Skip to main content

agent_client_protocol/
jsonrpc.rs

1//! Core JSON-RPC server support.
2
3use agent_client_protocol_schema::v1::{JsonRpcMessage as VersionedJsonRpcMessage, SessionId};
4
5// Types re-exported from crate root
6use serde::{Deserialize, Serialize};
7use std::any::TypeId;
8use std::collections::HashMap;
9use std::fmt::Debug;
10use std::marker::PhantomData;
11use std::panic::Location;
12use std::pin::pin;
13use std::sync::{
14    Arc, Mutex, Weak,
15    atomic::{AtomicBool, AtomicU8, Ordering},
16};
17use uuid::Uuid;
18
19use futures::FutureExt;
20use futures::channel::{mpsc, oneshot};
21use futures::future::{self, BoxFuture, Either};
22use futures::{AsyncRead, AsyncWrite, StreamExt};
23
24pub(crate) mod close;
25mod dynamic_handler;
26pub(crate) mod handlers;
27mod incoming_actor;
28mod outgoing_actor;
29#[cfg(test)]
30mod prepared_request_tests;
31pub(crate) mod protocol;
32mod protocol_compat;
33mod raw_error;
34pub(crate) mod raw_id;
35pub(crate) mod run;
36mod task_actor;
37mod termination;
38mod transport_actor;
39pub(crate) mod v1_connection;
40#[cfg(feature = "unstable_protocol_v2")]
41pub(crate) mod v2_connection;
42
43use crate::jsonrpc::close::{ChainedClose, CloseCallback};
44pub use crate::jsonrpc::close::{HandleConnectionClose, NullClose};
45use crate::jsonrpc::dynamic_handler::DynamicHandlerMessage;
46pub use crate::jsonrpc::handlers::NullHandler;
47use crate::jsonrpc::handlers::{ChainedHandler, NamedHandler};
48use crate::jsonrpc::handlers::{MessageHandler, NotificationHandler, RequestHandler};
49use crate::jsonrpc::outgoing_actor::{OutgoingMessageTx, send_raw_message};
50use crate::jsonrpc::protocol::HandleId;
51#[cfg(feature = "unstable_protocol_v2")]
52pub use crate::jsonrpc::protocol::V2;
53
54#[cfg(feature = "unstable_mcp_over_acp")]
55use crate::Client;
56pub use crate::jsonrpc::protocol::{Protocol, Raw, V1};
57#[cfg(feature = "unstable_protocol_v2")]
58pub(crate) use crate::jsonrpc::protocol_compat::normalize_newer_v2_initialize_request;
59pub(crate) use crate::jsonrpc::protocol_compat::normalize_v1_initialize_request;
60#[cfg(feature = "unstable_protocol_v2")]
61pub use crate::jsonrpc::protocol_compat::{
62    InitializationVersionMismatch, InitializeOutcome, SelectedInitializeResponse,
63};
64use crate::jsonrpc::protocol_compat::{ProtocolCompat, ProtocolMode};
65pub use crate::jsonrpc::raw_error::{RawJsonRpcError, RawJsonRpcResponse};
66pub use crate::jsonrpc::raw_id::RawJsonRpcId;
67use crate::jsonrpc::run::SpawnedRun;
68use crate::jsonrpc::run::{ChainRun, NullRun, RunWithConnectionTo};
69use crate::jsonrpc::task_actor::{Task, TaskTx};
70pub use crate::jsonrpc::termination::{
71    CONNECTION_TERMINATED_REASON, ConnectionTermination, is_connection_terminated,
72};
73use crate::jsonrpc::termination::{
74    TerminationSource, TerminationSources, connection_terminated_error,
75};
76#[cfg(feature = "unstable_mcp_over_acp")]
77use crate::mcp_server::McpServer;
78use crate::role::HasPeer;
79use crate::role::Role;
80use crate::{Agent, ConnectTo, RoleId};
81
82/// One valid JSON-RPC message carried inside a [`TransportFrame`].
83///
84/// Raw messages are the same for every ACP protocol version: IDs are
85/// [`RawJsonRpcId`]s, and method params and response errors are not
86/// interpreted at the transport boundary.
87///
88/// Serializing adds the `"jsonrpc": "2.0"` member, and deserializing requires
89/// it. Requests and notifications keep only `id`, `method` and `params`:
90/// other members of the envelope are dropped. Error responses keep unknown
91/// members of the error object in [`RawJsonRpcError::extra`].
92#[derive(Debug, Clone, PartialEq, Eq)]
93#[allow(
94    clippy::exhaustive_enums,
95    reason = "JSON-RPC defines exactly these message kinds"
96)]
97pub enum RawJsonRpcMessage {
98    /// A JSON-RPC request with an id and expected response.
99    Request(RawJsonRpcRequest),
100    /// A JSON-RPC notification without a response.
101    Notification(RawJsonRpcNotification),
102    /// A JSON-RPC response to a prior request.
103    Response(RawJsonRpcResponse),
104}
105
106/// A JSON-RPC frame exchanged between protocol components and transports.
107///
108/// A frame is one JSON-RPC wire value. Malformed wire input, including a
109/// JSON-RPC batch array, which the SDK does not support, is represented
110/// explicitly; transport failures are reported by the future that drives the
111/// transport rather than sent through a [`Channel`].
112#[derive(Clone, Debug)]
113#[non_exhaustive]
114pub enum TransportFrame {
115    /// One valid JSON-RPC message.
116    Single(RawJsonRpcMessage),
117    /// One malformed or invalid wire value retained for relays.
118    Malformed {
119        /// The original wire representation.
120        raw: String,
121        /// The JSON-RPC error associated with the malformed value.
122        error: crate::Error,
123    },
124}
125
126pub(crate) fn is_response_only_shape(value: &serde_json::Value) -> bool {
127    value.as_object().is_some_and(|object| {
128        !object.contains_key("method")
129            && (object.contains_key("result") || object.contains_key("error"))
130    })
131}
132
133pub(crate) fn raw_is_response_only_shape(raw: &str) -> bool {
134    serde_json::from_str(raw).is_ok_and(|value| is_response_only_shape(&value))
135}
136
137impl TransportFrame {
138    fn inspect_messages(
139        &self,
140        observer: &mut impl FnMut(&RawJsonRpcMessage) -> Result<(), crate::Error>,
141    ) -> Result<(), crate::Error> {
142        match self {
143            Self::Single(message) => observer(message),
144            Self::Malformed { .. } => Ok(()),
145        }
146    }
147}
148
149/// A JSON-RPC request with uninterpreted params.
150///
151/// This is the request without the `jsonrpc` member, which
152/// [`RawJsonRpcMessage`] adds and checks. Other members of the request object
153/// are not kept: a relay that parses and re-serializes a request drops them.
154#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
155#[allow(
156    clippy::exhaustive_structs,
157    reason = "JSON-RPC defines exactly these request members"
158)]
159pub struct RawJsonRpcRequest {
160    /// The ID the response must carry.
161    pub id: RawJsonRpcId,
162    /// The method to invoke.
163    pub method: Arc<str>,
164    /// The params, if any. `None` omits the member.
165    #[serde(default, skip_serializing_if = "Option::is_none")]
166    pub params: Option<RawJsonRpcParams>,
167}
168
169/// A JSON-RPC notification with uninterpreted params.
170///
171/// This is the notification without the `jsonrpc` member, which
172/// [`RawJsonRpcMessage`] adds and checks. Other members of the notification
173/// object are not kept: a relay that parses and re-serializes a notification
174/// drops them.
175#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
176#[allow(
177    clippy::exhaustive_structs,
178    reason = "JSON-RPC defines exactly these notification members"
179)]
180pub struct RawJsonRpcNotification {
181    /// The notification method.
182    pub method: Arc<str>,
183    /// The params, if any. `None` omits the member.
184    #[serde(default, skip_serializing_if = "Option::is_none")]
185    pub params: Option<RawJsonRpcParams>,
186}
187
188impl From<RawJsonRpcRequest> for RawJsonRpcMessage {
189    fn from(request: RawJsonRpcRequest) -> Self {
190        Self::Request(request)
191    }
192}
193
194impl From<RawJsonRpcNotification> for RawJsonRpcMessage {
195    fn from(notification: RawJsonRpcNotification) -> Self {
196        Self::Notification(notification)
197    }
198}
199
200impl From<RawJsonRpcResponse> for RawJsonRpcMessage {
201    fn from(response: RawJsonRpcResponse) -> Self {
202        Self::Response(response)
203    }
204}
205
206/// Raw JSON-RPC request or notification parameters.
207///
208/// JSON-RPC params, when present, must be either an array or an object.
209#[derive(Debug, Clone, PartialEq, Eq)]
210#[allow(
211    clippy::exhaustive_enums,
212    reason = "JSON-RPC params are either an array or an object"
213)]
214pub enum RawJsonRpcParams {
215    /// Positional JSON-RPC params.
216    Array(Vec<serde_json::Value>),
217    /// Named JSON-RPC params.
218    Object(serde_json::Map<String, serde_json::Value>),
219}
220
221impl RawJsonRpcParams {
222    /// Convert a JSON value into JSON-RPC params.
223    pub fn from_value(value: serde_json::Value) -> Result<Option<Self>, crate::Error> {
224        match value {
225            serde_json::Value::Null => Ok(None),
226            serde_json::Value::Array(array) => Ok(Some(Self::Array(array))),
227            serde_json::Value::Object(object) => Ok(Some(Self::Object(object))),
228            _ => {
229                Err(crate::Error::invalid_params()
230                    .data("JSON-RPC params must be an object or array"))
231            }
232        }
233    }
234
235    /// Convert params back into a JSON value.
236    #[must_use]
237    pub fn into_value(self) -> serde_json::Value {
238        match self {
239            Self::Array(array) => serde_json::Value::Array(array),
240            Self::Object(object) => serde_json::Value::Object(object),
241        }
242    }
243}
244
245impl Serialize for RawJsonRpcParams {
246    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
247    where
248        S: serde::Serializer,
249    {
250        match self {
251            Self::Array(array) => array.serialize(serializer),
252            Self::Object(object) => object.serialize(serializer),
253        }
254    }
255}
256
257impl<'de> Deserialize<'de> for RawJsonRpcParams {
258    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
259    where
260        D: serde::Deserializer<'de>,
261    {
262        let value = serde_json::Value::deserialize(deserializer)?;
263        match value {
264            serde_json::Value::Array(array) => Ok(Self::Array(array)),
265            serde_json::Value::Object(object) => Ok(Self::Object(object)),
266            _ => Err(serde::de::Error::custom(
267                "JSON-RPC params must be an object or array",
268            )),
269        }
270    }
271}
272
273impl RawJsonRpcMessage {
274    /// Build a raw JSON-RPC request message.
275    pub fn request(
276        method: String,
277        params: serde_json::Value,
278        id: impl Into<RawJsonRpcId>,
279    ) -> Result<Self, crate::Error> {
280        Ok(Self::Request(RawJsonRpcRequest {
281            id: id.into(),
282            method: Arc::from(method),
283            params: RawJsonRpcParams::from_value(params)?,
284        }))
285    }
286
287    /// Build a raw JSON-RPC notification message.
288    pub fn notification(method: String, params: serde_json::Value) -> Result<Self, crate::Error> {
289        Ok(Self::Notification(RawJsonRpcNotification {
290            method: Arc::from(method),
291            params: RawJsonRpcParams::from_value(params)?,
292        }))
293    }
294
295    /// Build a JSON-RPC response from an ACP result.
296    ///
297    /// For other protocols, construct [`RawJsonRpcResponse`] directly so error
298    /// codes and fields are not interpreted as ACP.
299    #[must_use]
300    pub fn response(
301        id: impl Into<RawJsonRpcId>,
302        response: Result<serde_json::Value, crate::Error>,
303    ) -> Self {
304        Self::Response(RawJsonRpcResponse::new(
305            id,
306            response.map_err(|error| Box::new(error.into())),
307        ))
308    }
309
310    /// The response id, if this is a response.
311    #[must_use]
312    pub fn response_id(&self) -> Option<&RawJsonRpcId> {
313        match self {
314            Self::Response(
315                RawJsonRpcResponse::Result { id, .. } | RawJsonRpcResponse::Error { id, .. },
316            ) => Some(id),
317            Self::Request(_) | Self::Notification(_) => None,
318        }
319    }
320}
321
322impl Serialize for RawJsonRpcMessage {
323    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
324    where
325        S: serde::Serializer,
326    {
327        match self {
328            Self::Request(request) => {
329                VersionedJsonRpcMessage::wrap(request.clone()).serialize(serializer)
330            }
331            Self::Notification(notification) => {
332                VersionedJsonRpcMessage::wrap(notification.clone()).serialize(serializer)
333            }
334            Self::Response(response) => {
335                VersionedJsonRpcMessage::wrap(response.clone()).serialize(serializer)
336            }
337        }
338    }
339}
340
341impl<'de> Deserialize<'de> for RawJsonRpcMessage {
342    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
343    where
344        D: serde::Deserializer<'de>,
345    {
346        let value = serde_json::Value::deserialize(deserializer)?;
347        let Some(object) = value.as_object() else {
348            return Err(serde::de::Error::custom("invalid JSON-RPC message"));
349        };
350
351        let has_method = object.contains_key("method");
352        let has_id = object.contains_key("id");
353        let has_result = object.contains_key("result");
354        let has_error = object.contains_key("error");
355
356        if has_method && !has_result && !has_error {
357            if has_id {
358                let request =
359                    serde_json::from_value::<VersionedJsonRpcMessage<RawJsonRpcRequest>>(value)
360                        .map_err(serde::de::Error::custom)?
361                        .into_inner();
362                Ok(Self::Request(request))
363            } else {
364                let notification = serde_json::from_value::<
365                    VersionedJsonRpcMessage<RawJsonRpcNotification>,
366                >(value)
367                .map_err(serde::de::Error::custom)?
368                .into_inner();
369                Ok(Self::Notification(notification))
370            }
371        } else if !has_method && has_id && has_result != has_error {
372            let response =
373                serde_json::from_value::<VersionedJsonRpcMessage<RawJsonRpcResponse>>(value)
374                    .map_err(serde::de::Error::custom)?
375                    .into_inner();
376            Ok(Self::Response(response))
377        } else {
378            Err(serde::de::Error::custom("invalid JSON-RPC message"))
379        }
380    }
381}
382
383fn params_from_transport(params: Option<RawJsonRpcParams>) -> serde_json::Value {
384    params.map_or(serde_json::Value::Null, RawJsonRpcParams::into_value)
385}
386
387/// Handlers process incoming JSON-RPC messages on a connection.
388///
389/// When messages arrive, they flow through a chain of handlers. Each handler can
390/// either **claim** the message (handle it) or **decline** it (pass to the next handler).
391///
392/// # Message Flow
393///
394/// Messages flow through three layers of handlers in order:
395///
396/// ```text
397/// ┌─────────────────────────────────────────────────────────────────┐
398/// │                     Incoming Message                            │
399/// └─────────────────────────────────────────────────────────────────┘
400///                              │
401///                              ▼
402/// ┌─────────────────────────────────────────────────────────────────┐
403/// │  1. User Handlers (registered via on_receive_request, etc.)     │
404/// │     - Tried in registration order                               │
405/// │     - First handler to return Handled::Yes claims the message   │
406/// └─────────────────────────────────────────────────────────────────┘
407///                              │ Handled::No
408///                              ▼
409/// ┌─────────────────────────────────────────────────────────────────┐
410/// │  2. Dynamic Handlers (added at runtime)                         │
411/// │     - Used for session-specific message handling                │
412/// │     - Added via ConnectionTo::add_dynamic_handler             │
413/// └─────────────────────────────────────────────────────────────────┘
414///                              │ Handled::No
415///                              ▼
416/// ┌─────────────────────────────────────────────────────────────────┐
417/// │  3. Role Default Handler                                        │
418/// │     - Fallback based on the connection's Role                   │
419/// │     - Handles protocol-level messages (e.g., proxy forwarding)  │
420/// └─────────────────────────────────────────────────────────────────┘
421///                              │ Handled::No
422///                              ▼
423/// ┌─────────────────────────────────────────────────────────────────┐
424/// │  Unhandled: requests error, notifications ignored               │
425/// └─────────────────────────────────────────────────────────────────┘
426/// ```
427///
428/// # The `Handled` Return Value
429///
430/// Each handler returns [`Handled`] to indicate whether it processed the message:
431///
432/// - **`Handled::Yes`** - Message was handled. No further handlers are invoked.
433/// - **`Handled::No { message, retry }`** - Message was not handled. The message
434///   (possibly modified) is passed to the next handler in the chain.
435///
436/// For convenience, handlers can return `()` which is equivalent to `Handled::Yes`.
437///
438/// # The Retry Mechanism
439///
440/// The `retry` flag in `Handled::No` controls what happens when no handler claims a message:
441///
442/// - **`retry: false`** (default) - Send a "method not found" error
443///   response immediately for requests, or ignore notifications.
444/// - **`retry: true`** - Queue the message and retry it when new dynamic handlers are added.
445///
446/// This mechanism exists because of a timing issue with sessions: when a `session/new`
447/// response is being processed, the dynamic handler for that session hasn't been registered
448/// yet, but `session/update` notifications for that session may already be arriving.
449/// By setting `retry: true`, these early notifications are queued until the session's
450/// dynamic handler is added.
451///
452/// The queue holds at most 1024 messages per connection. Queuing one more gives up on the
453/// oldest: a request gets the "method not found" error, a notification is dropped with a
454/// log, and a response goes to its awaiter.
455///
456/// On v1 and unguarded connections, a client's role also retries unhandled messages from
457/// the agent that carry a `sessionId`. Its notifications are queued as above. Its requests
458/// are queued only while a dropped dynamic handler's removal is still pending, which covers
459/// a session handler being replaced, and are answered with "method not found" once no such
460/// removal is left. Any other unhandled session request is answered right away, so the
461/// agent never waits for a handler that nobody registers.
462///
463/// # Handler Registration
464///
465/// Most users register handlers using the builder methods on [`Builder`]:
466///
467/// ```
468/// # use agent_client_protocol::{Agent, Client, ConnectTo};
469/// # use agent_client_protocol::v1::schema::{AgentCapabilities, InitializeRequest, InitializeResponse};
470/// # use agent_client_protocol_test::StatusUpdate;
471/// # async fn example(transport: impl ConnectTo<Agent>) -> Result<(), agent_client_protocol::Error> {
472/// Agent.v1()
473///     .on_receive_request(async |req: InitializeRequest, responder, cx| {
474///         responder.respond(
475///             InitializeResponse::new(req.protocol_version)
476///                 .agent_capabilities(AgentCapabilities::new()),
477///         )
478///     }, agent_client_protocol::on_receive_request!())
479///     .on_receive_notification(async |notif: StatusUpdate, cx| {
480///         // Process notification
481///         Ok(())
482///     }, agent_client_protocol::on_receive_notification!())
483///     .connect_to(transport)
484///     .await?;
485/// # Ok(())
486/// # }
487/// ```
488///
489/// The type parameter on the closure determines which messages are dispatched to it.
490/// Messages that don't match the type are automatically passed to the next handler.
491///
492/// # Implementing Custom Handlers
493///
494/// For advanced use cases, you can implement [`HandleDispatchFrom`] directly:
495///
496/// ```no_run
497/// use agent_client_protocol::{
498///     Client, ConnectionTo, Dispatch, Error, HandleDispatchFrom, Handled,
499/// };
500///
501/// struct MyHandler;
502///
503/// impl HandleDispatchFrom<Client> for MyHandler {
504///     async fn handle_dispatch_from(
505///         &mut self,
506///         message: Dispatch,
507///         _connection: ConnectionTo<Client>,
508///     ) -> Result<Handled<Dispatch>, Error> {
509///         if message.method() == "my/custom/method" {
510///             // Handle it
511///             Ok(Handled::Yes)
512///         } else {
513///             // Pass to next handler
514///             Ok(Handled::No { message, retry: false })
515///         }
516///     }
517///
518///     fn describe_chain(&self) -> impl std::fmt::Debug {
519///         "MyHandler"
520///     }
521/// }
522/// ```
523///
524/// # Important: Handlers Must Not Block
525///
526/// The connection processes messages on a single async task. While a handler is running,
527/// no other messages can be processed. For expensive operations, use [`ConnectionTo::spawn`]
528/// to run work concurrently:
529///
530/// ```
531/// # use agent_client_protocol::{Client, Agent, ConnectTo};
532/// # use agent_client_protocol_test::{expensive_operation, ProcessComplete};
533/// # async fn example(transport: impl ConnectTo<Client>) -> Result<(), agent_client_protocol::Error> {
534/// # Client.v1().connect_with(transport, async |cx| {
535/// cx.spawn({
536///     let connection = cx.clone();
537///     async move {
538///         let result = expensive_operation("data").await?;
539///         connection.send_notification(ProcessComplete { result })?;
540///         Ok(())
541///     }
542/// })?;
543/// # Ok(())
544/// # }).await?;
545/// # Ok(())
546/// # }
547/// ```
548#[allow(async_fn_in_trait)]
549/// A handler for incoming JSON-RPC messages.
550///
551/// This trait is implemented by types that can process incoming messages on a connection.
552/// Handlers are registered with a [`Builder`] and are called in order until
553/// one claims the message.
554///
555/// The type parameter is the counterpart role that messages arrive from and
556/// that the supplied [`ConnectionTo`] addresses. An agent handler therefore
557/// implements `HandleDispatchFrom<Client>`, while a client handler implements
558/// `HandleDispatchFrom<Agent>`.
559pub trait HandleDispatchFrom<Counterpart: Role>: Send {
560    /// Attempt to claim an incoming dispatch (request, notification, or response).
561    ///
562    /// # Important: do not block
563    ///
564    /// The server will not process new messages until this handler returns.
565    /// You should avoid blocking in this callback unless you wish to block the server (e.g., for rate limiting).
566    /// The recommended approach to manage expensive operations is to the [`ConnectionTo::spawn`] method available on the message context.
567    ///
568    /// # Parameters
569    ///
570    /// * `message` - The incoming message to handle.
571    /// * `connection` - The connection, used to send messages and access connection state.
572    ///
573    /// # Returns
574    ///
575    /// * `Ok(Handled::Yes)` if the message was claimed. It will not be propagated further.
576    /// * `Ok(Handled::No(message))` if not; the (possibly changed) message will be passed to the remaining handlers.
577    /// * `Err` if processing fails. Requests receive an Error Response, response
578    ///   errors are routed to the local request awaiter, and notification errors
579    ///   are logged without a wire reply.
580    fn handle_dispatch_from(
581        &mut self,
582        message: Dispatch,
583        connection: ConnectionTo<Counterpart>,
584    ) -> impl Future<Output = Result<Handled<Dispatch>, crate::Error>> + Send;
585
586    /// Returns a debug description of the registered handlers for diagnostics.
587    fn describe_chain(&self) -> impl std::fmt::Debug;
588}
589
590impl<Counterpart: Role, H> HandleDispatchFrom<Counterpart> for &mut H
591where
592    H: HandleDispatchFrom<Counterpart>,
593{
594    fn handle_dispatch_from(
595        &mut self,
596        message: Dispatch,
597        cx: ConnectionTo<Counterpart>,
598    ) -> impl Future<Output = Result<Handled<Dispatch>, crate::Error>> + Send {
599        H::handle_dispatch_from(self, message, cx)
600    }
601
602    fn describe_chain(&self) -> impl std::fmt::Debug {
603        H::describe_chain(self)
604    }
605}
606
607/// A JSON-RPC connection that can act as either a server, client, or both.
608///
609/// [`Builder`] provides a builder-style API for creating JSON-RPC servers and clients.
610/// You start with a versioned entry point such as `Client.v1()` or `Client.v2()`, then add message
611/// handlers, and finally drive the connection with either [`connect_to`](Builder::connect_to)
612/// or [`connect_with`](Builder::connect_with), providing a component implementation
613/// (e.g., [`ByteStreams`] for byte streams).
614///
615/// # JSON-RPC Primer
616///
617/// JSON-RPC 2.0 has two fundamental message types:
618///
619/// * **Requests** - Messages that expect a response. They have an `id` field that gets
620///   echoed back in the response so the sender can correlate them.
621/// * **Notifications** - Fire-and-forget messages with no `id` field. The sender doesn't
622///   expect or receive a response.
623///
624/// # Type-Driven Message Dispatch
625///
626/// The handler registration methods use Rust's type system to determine which messages
627/// to handle. The type parameter you provide controls what gets dispatched to your handler:
628///
629/// ## Single Message Types
630///
631/// The simplest case - handle one specific message type:
632///
633/// ```no_run
634/// # use agent_client_protocol_test::*;
635/// # use agent_client_protocol::v1::schema::{InitializeRequest, InitializeResponse, SessionNotification};
636/// # async fn example() -> Result<(), agent_client_protocol::Error> {
637/// # let connection = mock_connection();
638/// connection
639///     .on_receive_request(async |req: InitializeRequest, responder, cx| {
640///         // Handle only InitializeRequest messages
641///         responder.respond(InitializeResponse::make())
642///     }, agent_client_protocol::on_receive_request!())
643///     .on_receive_notification(async |notif: SessionNotification, cx| {
644///         // Handle only SessionUpdate notifications
645///         Ok(())
646///     }, agent_client_protocol::on_receive_notification!())
647/// # .connect_to(agent_client_protocol_test::MockTransport).await?;
648/// # Ok(())
649/// # }
650/// ```
651///
652/// ## Enum Message Types
653///
654/// You can also handle multiple related messages with a single handler by defining an enum
655/// that implements the appropriate trait ([`JsonRpcRequest`] or [`JsonRpcNotification`]):
656///
657/// ```no_run
658/// # use agent_client_protocol_test::*;
659/// # use agent_client_protocol::{JsonRpcRequest, JsonRpcMessage, UntypedMessage};
660/// # use agent_client_protocol::v1::schema::{InitializeRequest, InitializeResponse, PromptRequest, PromptResponse};
661/// # async fn example() -> Result<(), agent_client_protocol::Error> {
662/// # let connection = mock_connection();
663/// // Define an enum for multiple request types
664/// #[derive(Debug, Clone)]
665/// enum MyRequests {
666///     Initialize(InitializeRequest),
667///     Prompt(PromptRequest),
668/// }
669///
670/// // Implement JsonRpcRequest for your enum
671/// # impl JsonRpcMessage for MyRequests {
672/// #     fn matches_method(_method: &str) -> bool { false }
673/// #     fn method(&self) -> &str { "myRequests" }
674/// #     fn to_untyped_message(&self) -> Result<UntypedMessage, agent_client_protocol::Error> { todo!() }
675/// #     fn parse_message(_method: &str, _params: &serde_json::Value) -> Result<Self, agent_client_protocol::Error> { Err(agent_client_protocol::Error::method_not_found()) }
676/// # }
677/// impl JsonRpcRequest for MyRequests { type Response = serde_json::Value; }
678///
679/// // Handle all variants in one place
680/// connection.on_receive_request(async |req: MyRequests, responder, cx| {
681///     match req {
682///         MyRequests::Initialize(init) => { responder.respond(serde_json::json!({})) }
683///         MyRequests::Prompt(prompt) => { responder.respond(serde_json::json!({})) }
684///     }
685/// }, agent_client_protocol::on_receive_request!())
686/// # .connect_to(agent_client_protocol_test::MockTransport).await?;
687/// # Ok(())
688/// # }
689/// ```
690///
691/// ## Mixed Message Types
692///
693/// To handle requests, notifications, and responses in one callback, use
694/// [`on_receive_dispatch`](Self::on_receive_dispatch):
695///
696/// ```no_run
697/// # use agent_client_protocol_test::*;
698/// # use agent_client_protocol::v1::Dispatch;
699/// # use agent_client_protocol::v1::schema::{InitializeRequest, InitializeResponse, SessionNotification};
700/// # async fn example() -> Result<(), agent_client_protocol::Error> {
701/// # let connection = mock_connection();
702/// // on_receive_dispatch receives requests, notifications, and responses
703/// connection.on_receive_dispatch(async |msg: Dispatch<InitializeRequest, SessionNotification>, _cx| {
704///     match msg {
705///         Dispatch::Request(req, responder) => {
706///             responder.respond(InitializeResponse::make())
707///         }
708///         Dispatch::Notification(notif) => {
709///             Ok(())
710///         }
711///         Dispatch::Response(result, router) => {
712///             // Forward response to its destination
713///             router.route_with_result(result)
714///         }
715///     }
716/// }, agent_client_protocol::on_receive_dispatch!())
717/// # .connect_to(agent_client_protocol_test::MockTransport).await?;
718/// # Ok(())
719/// # }
720/// ```
721///
722/// # Handler Registration
723///
724/// Register handlers using these methods (listed from most common to most flexible):
725///
726/// * [`on_receive_request`](Self::on_receive_request) - Handle JSON-RPC requests (messages expecting responses)
727/// * [`on_receive_notification`](Self::on_receive_notification) - Handle JSON-RPC notifications (fire-and-forget)
728/// * [`on_receive_dispatch`](Self::on_receive_dispatch) - Handle requests, notifications, and responses in one callback
729/// * [`with_handler`](Self::with_handler) - Low-level primitive for maximum flexibility
730///
731/// ## Handler Ordering
732///
733/// Handlers are tried in the order you register them. The first handler that claims a message
734/// (by matching its type) will process it. Subsequent handlers won't see that message:
735///
736/// ```no_run
737/// # use agent_client_protocol_test::*;
738/// # use agent_client_protocol::v1::schema::{InitializeRequest, InitializeResponse, PromptRequest, PromptResponse};
739/// # async fn example() -> Result<(), agent_client_protocol::Error> {
740/// # let connection = mock_connection();
741/// connection
742///     .on_receive_request(async |req: InitializeRequest, responder, cx| {
743///         // This runs first for InitializeRequest
744///         responder.respond(InitializeResponse::make())
745///     }, agent_client_protocol::on_receive_request!())
746///     .on_receive_request(async |req: PromptRequest, responder, cx| {
747///         // This runs first for PromptRequest
748///         responder.respond(PromptResponse::make())
749///     }, agent_client_protocol::on_receive_request!())
750///     // Unknown requests receive Method not found automatically; unhandled
751///     // notifications are ignored.
752/// # .connect_to(agent_client_protocol_test::MockTransport).await?;
753/// # Ok(())
754/// # }
755/// ```
756///
757/// # Event Loop and Concurrency
758///
759/// Understanding the event loop is critical for writing correct handlers.
760///
761/// ## The Event Loop
762///
763/// [`Builder`] runs all handler callbacks on a single async task - the event loop.
764/// While a handler is running, **the server cannot receive new messages**. This means
765/// any blocking or expensive work in your handlers will stall the entire connection.
766///
767/// To avoid blocking the event loop, use [`ConnectionTo::spawn`] to offload serious
768/// work to concurrent tasks:
769///
770/// ```no_run
771/// # use agent_client_protocol_test::*;
772/// # async fn example() -> Result<(), agent_client_protocol::Error> {
773/// # let connection = mock_connection();
774/// connection.on_receive_request(async |req: AnalyzeRequest, responder, cx| {
775///     // Clone cx for the spawned task
776///     cx.spawn({
777///         let connection = cx.clone();
778///         async move {
779///             let result = expensive_analysis(&req.data).await?;
780///             connection.send_notification(AnalysisComplete { result })?;
781///             Ok(())
782///         }
783///     })?;
784///
785///     // Respond immediately without blocking
786///     responder.respond(AnalysisStarted { job_id: 42 })
787/// }, agent_client_protocol::on_receive_request!())
788/// # .connect_to(agent_client_protocol_test::MockTransport).await?;
789/// # Ok(())
790/// # }
791/// ```
792///
793/// Note that the entire connection runs within one async task, so parallelism must be
794/// managed explicitly using [`spawn`](ConnectionTo::spawn).
795///
796/// ## The Connection Context
797///
798/// Handler callbacks receive a context object (`cx`) for interacting with the connection:
799///
800/// * **For request handlers** - [`Responder<R>`] provides [`respond`](Responder::respond)
801///   to send the response, plus methods to send other messages
802/// * **For notification handlers** - [`ConnectionTo`] provides methods to send messages
803///   and spawn tasks
804///
805/// Both context types support:
806/// * [`send_request`](ConnectionTo::send_request) - Send requests to the other side
807/// * [`send_notification`](ConnectionTo::send_notification) - Send notifications
808/// * [`spawn`](ConnectionTo::spawn) - Run tasks concurrently without blocking the event loop
809///
810/// The [`SentRequest`] returned by `send_request` provides methods like
811/// [`on_receiving_result`](SentRequest::on_receiving_result) that help you
812/// avoid accidentally blocking the event loop while waiting for responses.
813///
814/// # Driving the Connection
815///
816/// After adding handlers, you must drive the connection using one of two modes:
817///
818/// ## Server Mode: `connect_to()`
819///
820/// Use [`connect_to`](Self::connect_to) when you only need to respond to incoming messages:
821///
822/// ```no_run
823/// # use agent_client_protocol_test::*;
824/// # async fn example() -> Result<(), agent_client_protocol::Error> {
825/// # let connection = mock_connection();
826/// connection
827///     .on_receive_request(async |req: MyRequest, responder, cx| {
828///         responder.respond(MyResponse { status: "ok".into() })
829///     }, agent_client_protocol::on_receive_request!())
830///     .connect_to(MockTransport)  // Runs until connection closes or error occurs
831///     .await?;
832/// # Ok(())
833/// # }
834/// ```
835///
836/// The connection will process incoming messages and invoke your handlers until the
837/// connection is closed or an error occurs.
838///
839/// ## Client Mode: `connect_with()`
840///
841/// Use [`connect_with`](Self::connect_with) when you need to both handle incoming messages
842/// AND send your own requests/notifications:
843///
844/// ```no_run
845/// # use agent_client_protocol_test::*;
846/// # use agent_client_protocol::v1::schema::InitializeRequest;
847/// # async fn example() -> Result<(), agent_client_protocol::Error> {
848/// # let connection = mock_connection();
849/// connection
850///     .on_receive_request(async |req: MyRequest, responder, cx| {
851///         responder.respond(MyResponse { status: "ok".into() })
852///     }, agent_client_protocol::on_receive_request!())
853///     .connect_with(MockTransport, async |cx| {
854///         // You can send requests to the other side
855///         let response = cx.send_request(InitializeRequest::make())
856///             .block_task()
857///             .await?;
858///
859///         // And send notifications
860///         cx.send_notification(StatusUpdate { message: "ready".into() })?;
861///
862///         Ok(())
863///     })
864///     .await?;
865/// # Ok(())
866/// # }
867/// ```
868///
869/// The connection will serve incoming messages in the background while your client closure
870/// runs. When the closure returns, the connection shuts down.
871///
872/// # Example: Complete Agent
873///
874/// ```no_run
875/// # use agent_client_protocol::UntypedRole;
876/// # use agent_client_protocol::{Builder};
877/// # use agent_client_protocol::v1::schema::{InitializeRequest, InitializeResponse, PromptRequest, PromptResponse, SessionNotification};
878/// # async fn example(transport: impl agent_client_protocol::ConnectTo<UntypedRole>) -> Result<(), agent_client_protocol::Error> {
879///
880/// UntypedRole.builder()
881///     .name("my-agent")  // Optional: for debugging logs
882///     .on_receive_request(async |init: InitializeRequest, responder, cx| {
883///         let response: InitializeResponse = todo!();
884///         responder.respond(response)
885///     }, agent_client_protocol::on_receive_request!())
886///     .on_receive_request(async |prompt: PromptRequest, responder, cx| {
887///         // You can send notifications while processing a request
888///         let notif: SessionNotification = todo!();
889///         cx.send_notification(notif)?;
890///
891///         // Then respond to the request
892///         let response: PromptResponse = todo!();
893///         responder.respond(response)
894///     }, agent_client_protocol::on_receive_request!())
895///     .connect_to(transport)
896///     .await?;
897/// # Ok(())
898/// # }
899/// ```
900///
901/// # Protocol version
902///
903/// `P` is the [`Protocol`] marker of the request handles that typed handlers
904/// receive and of the connection type callbacks receive. `Client.v1()`,
905/// `Agent.v1()` and `Proxy.v1()` return [`V1`] builders
906/// ([`v1::Builder`](crate::v1::Builder)), the v2 builders use `V2`
907/// (`v2::Builder`), and [`Builder::new`] and other roles' builders use the
908/// version-neutral default, [`Raw`], with no protocol-version guard.
909#[must_use]
910#[derive(Debug)]
911pub struct Builder<Host: Role, Handler = NullHandler, Runner = NullRun, Close = NullClose, P = Raw>
912where
913    Handler: HandleDispatchFrom<Host::Counterpart>,
914    Runner: RunWithConnectionTo<Host::Counterpart>,
915    Close: HandleConnectionClose<Host::Counterpart>,
916    P: Protocol,
917{
918    /// My role.
919    host: Host,
920
921    /// Name of the connection, used in tracing logs.
922    name: Option<String>,
923
924    /// Handler for incoming messages.
925    handler: Handler,
926
927    /// Runner for background connection tasks.
928    runner: Runner,
929
930    /// Protocol version mode for the public API and wire compatibility layer.
931    protocol_mode: ProtocolMode,
932
933    /// Handler run when the incoming transport reaches clean EOF.
934    on_close: Close,
935
936    /// The protocol version of the callbacks: their connection type and the
937    /// request handles they receive.
938    protocol: PhantomData<fn() -> P>,
939}
940
941impl<Host: Role> Builder<Host, NullHandler, NullRun, NullClose> {
942    /// Create a new version-neutral connection builder for the given role.
943    /// This type follows a builder pattern; use other methods to configure and then invoke
944    /// [`Self::connect_to`] (to use as a server) or [`Self::connect_with`] to use as a client.
945    ///
946    /// The builder is raw: it applies no ACP protocol-version guard, even for
947    /// [`Client`](crate::Client), [`Agent`] and [`Proxy`](crate::Proxy), and
948    /// its callbacks receive [`Raw`] request handles. Messages are sent and received exactly as
949    /// handlers produce them. This suits protocol-routing infrastructure that
950    /// validates `initialize` itself and forwarders such as the conductor,
951    /// which leave version validation to the endpoints.
952    ///
953    /// Applications should use [`Client.v1()`](crate::Client::v1),
954    /// [`Agent.v1()`](crate::Agent::v1) or [`Proxy.v1()`](crate::Proxy::v1)
955    /// (or the `v2()` entry points with the `unstable_protocol_v2` feature),
956    /// which apply that version's wire guard.
957    pub fn new(role: Host) -> Self {
958        Self {
959            host: role,
960            name: None,
961            handler: NullHandler,
962            runner: NullRun,
963            protocol_mode: ProtocolMode::disabled(),
964            on_close: NullClose,
965            protocol: PhantomData,
966        }
967    }
968}
969
970impl<Host: Role, Handler> Builder<Host, Handler, NullRun, NullClose>
971where
972    Handler: HandleDispatchFrom<Host::Counterpart>,
973{
974    /// Create a new connection builder with the given handler.
975    ///
976    /// Like [`Builder::new`], the builder applies no protocol-version guard.
977    pub fn new_with(role: Host, handler: Handler) -> Self {
978        Self {
979            host: role,
980            name: None,
981            handler,
982            runner: NullRun,
983            protocol_mode: ProtocolMode::disabled(),
984            on_close: NullClose,
985            protocol: PhantomData,
986        }
987    }
988}
989
990impl<
991    Host: Role,
992    Handler: HandleDispatchFrom<Host::Counterpart>,
993    Runner: RunWithConnectionTo<Host::Counterpart>,
994    Close: HandleConnectionClose<Host::Counterpart>,
995    P: Protocol,
996> Builder<Host, Handler, Runner, Close, P>
997{
998    /// Set the "name" of this connection -- used only for debugging logs.
999    pub fn name(mut self, name: impl ToString) -> Self {
1000        self.name = Some(name.to_string());
1001        self
1002    }
1003
1004    /// The same builder with callbacks for protocol `Q`. Handlers already
1005    /// registered keep their own protocol.
1006    pub(crate) fn into_protocol<Q: Protocol>(self) -> Builder<Host, Handler, Runner, Close, Q> {
1007        Builder {
1008            host: self.host,
1009            name: self.name,
1010            handler: self.handler,
1011            runner: self.runner,
1012            protocol_mode: self.protocol_mode,
1013            on_close: self.on_close,
1014            protocol: PhantomData,
1015        }
1016    }
1017
1018    #[cfg(feature = "unstable_protocol_v2")]
1019    pub(crate) fn v2_agent(self) -> Builder<Host, Handler, Runner, Close, V2> {
1020        let mut builder = self.into_protocol();
1021        builder.protocol_mode = ProtocolMode::v2_agent();
1022        builder
1023    }
1024
1025    #[cfg(feature = "unstable_protocol_v2")]
1026    pub(crate) fn v2_client(self) -> Builder<Host, Handler, Runner, Close, V2> {
1027        let mut builder = self.into_protocol();
1028        builder.protocol_mode = ProtocolMode::v2_client();
1029        builder
1030    }
1031
1032    #[cfg(feature = "unstable_protocol_v2")]
1033    pub(crate) fn v2_proxy(self) -> Builder<Host, Handler, Runner, Close, V2> {
1034        let mut builder = self.into_protocol();
1035        builder.protocol_mode = ProtocolMode::v2_proxy();
1036        builder
1037    }
1038
1039    pub(crate) fn v1_agent(self) -> Builder<Host, Handler, Runner, Close, V1> {
1040        let mut builder = self.into_protocol();
1041        builder.protocol_mode = ProtocolMode::v1_agent();
1042        builder
1043    }
1044
1045    pub(crate) fn v1_client(self) -> Builder<Host, Handler, Runner, Close, V1> {
1046        let mut builder = self.into_protocol();
1047        builder.protocol_mode = ProtocolMode::v1_client();
1048        builder
1049    }
1050
1051    pub(crate) fn v1_proxy(self) -> Builder<Host, Handler, Runner, Close, V1> {
1052        let mut builder = self.into_protocol();
1053        builder.protocol_mode = ProtocolMode::v1_proxy();
1054        builder
1055    }
1056
1057    /// Record a request handler of type `Req`.
1058    ///
1059    /// A v1 agent that handles `initialize` itself rejects other requests
1060    /// until initialization completes, as v1 agents SHOULD. Agents without
1061    /// such a handler, including those with only catch-all handlers, keep
1062    /// accepting requests without initialization, because they could not
1063    /// complete one.
1064    fn note_request_handler<Req: JsonRpcRequest>(mut self) -> Self {
1065        // A handler for every method is a catch-all, not an initialize handler.
1066        const UNRELATED_METHOD: &str = "\0acp-sdk/not-a-method";
1067        if TypeId::of::<Host>() == TypeId::of::<Agent>()
1068            && Req::matches_method("initialize")
1069            && !Req::matches_method(UNRELATED_METHOD)
1070        {
1071            self.protocol_mode = self.protocol_mode.require_initialize();
1072        }
1073        self
1074    }
1075
1076    /// Merge another [`Builder`] into this one.
1077    ///
1078    /// Prefer [`Self::on_receive_request`] or [`Self::on_receive_notification`].
1079    /// This is a low-level method that is not intended for general use.
1080    pub fn with_connection_builder(
1081        self,
1082        other: Builder<
1083            Host,
1084            impl HandleDispatchFrom<Host::Counterpart>,
1085            impl RunWithConnectionTo<Host::Counterpart>,
1086            impl HandleConnectionClose<Host::Counterpart>,
1087            P,
1088        >,
1089    ) -> Builder<
1090        Host,
1091        impl HandleDispatchFrom<Host::Counterpart>,
1092        impl RunWithConnectionTo<Host::Counterpart>,
1093        impl HandleConnectionClose<Host::Counterpart>,
1094        P,
1095    > {
1096        let Builder {
1097            name: other_name,
1098            handler: other_handler,
1099            runner: other_runner,
1100            protocol_mode: other_protocol_mode,
1101            on_close: other_on_close,
1102            protocol: _,
1103            host: _,
1104        } = other;
1105        Builder {
1106            host: self.host,
1107            name: self.name,
1108            handler: ChainedHandler::new(
1109                self.handler,
1110                NamedHandler::new(other_name, other_handler),
1111            ),
1112            runner: ChainRun::new(self.runner, other_runner),
1113            protocol_mode: self.protocol_mode.merge(other_protocol_mode),
1114            on_close: ChainedClose::new(self.on_close, other_on_close),
1115            protocol: PhantomData,
1116        }
1117    }
1118
1119    /// Add a new [`HandleDispatchFrom`] to the chain.
1120    ///
1121    /// Prefer [`Self::on_receive_request`] or [`Self::on_receive_notification`].
1122    /// This is a low-level method that is not intended for general use.
1123    pub fn with_handler(
1124        self,
1125        handler: impl HandleDispatchFrom<Host::Counterpart>,
1126    ) -> Builder<Host, impl HandleDispatchFrom<Host::Counterpart>, Runner, Close, P> {
1127        Builder {
1128            host: self.host,
1129            name: self.name,
1130            handler: ChainedHandler::new(self.handler, handler),
1131            runner: self.runner,
1132            protocol_mode: self.protocol_mode,
1133            on_close: self.on_close,
1134            protocol: PhantomData,
1135        }
1136    }
1137
1138    /// Add a new [`RunWithConnectionTo`] to the chain.
1139    pub fn with_runner<Run1>(
1140        self,
1141        runner: Run1,
1142    ) -> Builder<Host, Handler, impl RunWithConnectionTo<Host::Counterpart>, Close, P>
1143    where
1144        Run1: RunWithConnectionTo<Host::Counterpart>,
1145    {
1146        Builder {
1147            host: self.host,
1148            name: self.name,
1149            handler: self.handler,
1150            runner: ChainRun::new(self.runner, runner),
1151            protocol_mode: self.protocol_mode,
1152            on_close: self.on_close,
1153            protocol: PhantomData,
1154        }
1155    }
1156
1157    /// Enqueue a task to run once the connection is actively serving traffic.
1158    #[track_caller]
1159    pub fn with_spawned<F, Fut>(
1160        self,
1161        task: F,
1162    ) -> Builder<Host, Handler, impl RunWithConnectionTo<Host::Counterpart>, Close, P>
1163    where
1164        F: FnOnce(P::Connection<Host::Counterpart>) -> Fut + Send,
1165        Fut: Future<Output = Result<(), crate::Error>> + Send,
1166    {
1167        let location = Location::caller();
1168        self.with_runner(SpawnedRun::<_, P>::new(location, task))
1169    }
1170
1171    /// Run a callback when the incoming transport reaches clean EOF.
1172    ///
1173    /// Each callback runs at most once and receives the connection context. A
1174    /// successful callback observes the close without otherwise changing the
1175    /// lifetime of [`connect_with`](Self::connect_with). Returning an error
1176    /// shuts down the connection and cancels a still-running `connect_with`
1177    /// future.
1178    ///
1179    /// Multiple callbacks run sequentially in registration order. All of them
1180    /// run even if an earlier callback fails, after which the first error is
1181    /// returned. Pending requests are failed before callbacks begin, while the
1182    /// selected connection context's `incoming_closed` future completes only
1183    /// after they finish. A callback must therefore not await that close
1184    /// future itself.
1185    ///
1186    /// This separation lets applications choose their cancellation policy. A
1187    /// callback can notify application-owned tasks and return `Ok(())` for
1188    /// graceful cleanup, or return an error to stop them immediately.
1189    ///
1190    /// Close callbacks do not run for transport or handler errors, foreground
1191    /// completion, or a dropped driver. Use [`ConnectionTo::terminated`] to
1192    /// observe every way the connection can end.
1193    ///
1194    /// ```
1195    /// # use agent_client_protocol::{Client, ConnectTo, Error};
1196    /// # async fn example(transport: impl ConnectTo<Client>) -> Result<(), Error> {
1197    /// Client.v1()
1198    ///     .on_close(async |_cx| {
1199    ///         Err(Error::internal_error().data("agent transport closed"))
1200    ///     })
1201    ///     .connect_with(transport, async |_cx| {
1202    ///         std::future::pending().await
1203    ///     })
1204    ///     .await?;
1205    /// # Ok(())
1206    /// # }
1207    /// ```
1208    pub fn on_close<F, Fut>(
1209        self,
1210        callback: F,
1211    ) -> Builder<Host, Handler, Runner, impl HandleConnectionClose<Host::Counterpart>, P>
1212    where
1213        F: FnOnce(P::Connection<Host::Counterpart>) -> Fut + Send,
1214        Fut: Future<Output = Result<(), crate::Error>> + Send,
1215    {
1216        Builder {
1217            host: self.host,
1218            name: self.name,
1219            handler: self.handler,
1220            runner: self.runner,
1221            protocol_mode: self.protocol_mode,
1222            on_close: ChainedClose::new(self.on_close, CloseCallback::<_, P>::new(callback)),
1223            protocol: PhantomData,
1224        }
1225    }
1226
1227    /// Register a handler for requests, notifications, and responses.
1228    ///
1229    /// Use this when you want to handle all JSON-RPC message kinds in one callback.
1230    /// Your handler receives a [`Dispatch<Req, Notif>`] with three variants:
1231    ///
1232    /// - `Dispatch::Request(request, responder)` - A request with its response context
1233    /// - `Dispatch::Notification(notification)` - A notification
1234    /// - `Dispatch::Response(result, router)` - A response to a request we sent
1235    ///
1236    /// # Example
1237    ///
1238    /// ```no_run
1239    /// # use agent_client_protocol_test::*;
1240    /// # use agent_client_protocol::v1::Dispatch;
1241    /// # async fn example() -> Result<(), agent_client_protocol::Error> {
1242    /// # let connection = mock_connection();
1243    /// connection.on_receive_dispatch(async |message: Dispatch<MyRequest, StatusUpdate>, _cx| {
1244    ///     match message {
1245    ///         Dispatch::Request(req, responder) => {
1246    ///             // Handle request and send response
1247    ///             responder.respond(MyResponse { status: "ok".into() })
1248    ///         }
1249    ///         Dispatch::Notification(notif) => {
1250    ///             // Handle notification (no response needed)
1251    ///             Ok(())
1252    ///         }
1253    ///         Dispatch::Response(result, router) => {
1254    ///             // Forward response to its destination
1255    ///             router.route_with_result(result)
1256    ///         }
1257    ///     }
1258    /// }, agent_client_protocol::on_receive_dispatch!())
1259    /// # .connect_to(agent_client_protocol_test::MockTransport).await?;
1260    /// # Ok(())
1261    /// # }
1262    /// ```
1263    ///
1264    /// For most use cases, prefer [`on_receive_request`](Self::on_receive_request) or
1265    /// [`on_receive_notification`](Self::on_receive_notification) which provide cleaner APIs
1266    /// for handling requests or notifications separately.
1267    ///
1268    /// # Ordering
1269    ///
1270    /// This callback runs inside the dispatch loop and blocks further message processing
1271    /// until it completes. See the [`ordering`](crate::concepts::ordering) module for details on
1272    /// ordering guarantees and how to avoid deadlocks.
1273    pub fn on_receive_dispatch<Req, Notif, F, T, ToFut>(
1274        self,
1275        op: F,
1276        to_future_hack: ToFut,
1277    ) -> Builder<Host, impl HandleDispatchFrom<Host::Counterpart>, Runner, Close, P>
1278    where
1279        Host::Counterpart: HasPeer<Host::Counterpart>,
1280        Req: JsonRpcRequest,
1281        Notif: JsonRpcNotification,
1282        F: AsyncFnMut(
1283                Dispatch<Req, Notif, P>,
1284                P::Connection<Host::Counterpart>,
1285            ) -> Result<T, crate::Error>
1286            + Send,
1287        T: IntoHandled<Dispatch<Req, Notif, P>>,
1288        ToFut: Fn(
1289                &mut F,
1290                Dispatch<Req, Notif, P>,
1291                P::Connection<Host::Counterpart>,
1292            ) -> crate::BoxFuture<'_, Result<T, crate::Error>>
1293            + Send
1294            + Sync,
1295    {
1296        let handler = MessageHandler::<_, _, P, _, _, _, _>::new(
1297            self.host.counterpart(),
1298            self.host.counterpart(),
1299            op,
1300            to_future_hack,
1301        );
1302        self.note_request_handler::<Req>().with_handler(handler)
1303    }
1304
1305    /// Register a handler for JSON-RPC requests of type `Req`.
1306    ///
1307    /// Your handler receives three arguments:
1308    /// 1. The request (type `Req`)
1309    /// 2. A [`Responder<Req::Response, P>`] for sending the response
1310    /// 3. The builder-selected connection context for the peer that sent the
1311    ///    request (`ConnectionTo` by default, or `v2::ConnectionTo` for a
1312    ///    `v2::Builder`)
1313    ///
1314    /// The request context allows you to:
1315    /// - Send the response with [`Responder::respond`]
1316    /// - Send notifications to the client with the context's
1317    ///   `send_notification` method
1318    /// - Send requests to the client with the context's `send_request` method
1319    ///
1320    /// # Example
1321    ///
1322    /// ```no_run
1323    /// # use agent_client_protocol::{Agent, ConnectTo};
1324    /// # use agent_client_protocol::v1::schema::{PromptRequest, PromptResponse, SessionNotification};
1325    /// # async fn example(transport: impl ConnectTo<Agent>) -> Result<(), agent_client_protocol::Error> {
1326    /// Agent.v1().on_receive_request(async |request: PromptRequest, responder, cx| {
1327    ///     // Send a notification while processing
1328    ///     let notif: SessionNotification = todo!();
1329    ///     cx.send_notification(notif)?;
1330    ///
1331    ///     // Send the response
1332    ///     let response: PromptResponse = todo!();
1333    ///     responder.respond(response)
1334    /// }, agent_client_protocol::on_receive_request!())
1335    /// .connect_to(transport)
1336    /// .await
1337    /// # }
1338    /// ```
1339    ///
1340    /// # Type Parameter
1341    ///
1342    /// `Req` can be either a single request type or an enum of multiple request types.
1343    /// See the [type-driven dispatch](Self#type-driven-message-dispatch) section for details.
1344    ///
1345    /// # Ordering
1346    ///
1347    /// This callback runs inside the dispatch loop and blocks further message processing
1348    /// until it completes. See the [`ordering`](crate::concepts::ordering) module for details on
1349    /// ordering guarantees and how to avoid deadlocks.
1350    pub fn on_receive_request<Req: JsonRpcRequest, F, T, ToFut>(
1351        self,
1352        op: F,
1353        to_future_hack: ToFut,
1354    ) -> Builder<Host, impl HandleDispatchFrom<Host::Counterpart>, Runner, Close, P>
1355    where
1356        Host::Counterpart: HasPeer<Host::Counterpart>,
1357        F: AsyncFnMut(
1358                Req,
1359                Responder<Req::Response, P>,
1360                P::Connection<Host::Counterpart>,
1361            ) -> Result<T, crate::Error>
1362            + Send,
1363        T: IntoHandled<(Req, Responder<Req::Response, P>)>,
1364        ToFut: Fn(
1365                &mut F,
1366                Req,
1367                Responder<Req::Response, P>,
1368                P::Connection<Host::Counterpart>,
1369            ) -> crate::BoxFuture<'_, Result<T, crate::Error>>
1370            + Send
1371            + Sync,
1372    {
1373        let handler = RequestHandler::<_, _, P, _, _, _>::new(
1374            self.host.counterpart(),
1375            self.host.counterpart(),
1376            op,
1377            to_future_hack,
1378        );
1379        self.note_request_handler::<Req>().with_handler(handler)
1380    }
1381
1382    /// Register a handler for JSON-RPC notifications of type `Notif`.
1383    ///
1384    /// Notifications are fire-and-forget messages that don't expect a response.
1385    /// Your handler receives:
1386    /// 1. The notification (type `Notif`)
1387    /// 2. The builder-selected connection context for sending messages to the
1388    ///    other side
1389    ///
1390    /// Unlike request handlers, you cannot send a response (notifications don't have IDs),
1391    /// but you can still send your own requests and notifications using the context.
1392    ///
1393    /// # Example
1394    ///
1395    /// ```no_run
1396    /// # use agent_client_protocol_test::*;
1397    /// # async fn example() -> Result<(), agent_client_protocol::Error> {
1398    /// # let connection = mock_connection();
1399    /// connection.on_receive_notification(async |notif: SessionUpdate, cx| {
1400    ///     // Process the notification
1401    ///     update_session_state(&notif)?;
1402    ///
1403    ///     // Optionally send a notification back
1404    ///     cx.send_notification(StatusUpdate {
1405    ///         message: "Acknowledged".into(),
1406    ///     })?;
1407    ///
1408    ///     Ok(())
1409    /// }, agent_client_protocol::on_receive_notification!())
1410    /// # .connect_to(agent_client_protocol_test::MockTransport).await?;
1411    /// # Ok(())
1412    /// # }
1413    /// ```
1414    ///
1415    /// # Type Parameter
1416    ///
1417    /// `Notif` can be either a single notification type or an enum of multiple notification types.
1418    /// See the [type-driven dispatch](Self#type-driven-message-dispatch) section for details.
1419    ///
1420    /// # Ordering
1421    ///
1422    /// This callback runs inside the dispatch loop and blocks further message processing
1423    /// until it completes. See the [`ordering`](crate::concepts::ordering) module for details on
1424    /// ordering guarantees and how to avoid deadlocks.
1425    pub fn on_receive_notification<Notif, F, T, ToFut>(
1426        self,
1427        op: F,
1428        to_future_hack: ToFut,
1429    ) -> Builder<Host, impl HandleDispatchFrom<Host::Counterpart>, Runner, Close, P>
1430    where
1431        Host::Counterpart: HasPeer<Host::Counterpart>,
1432        Notif: JsonRpcNotification,
1433        F: AsyncFnMut(Notif, P::Connection<Host::Counterpart>) -> Result<T, crate::Error> + Send,
1434        T: IntoHandled<(Notif, P::Connection<Host::Counterpart>)>,
1435        ToFut: Fn(
1436                &mut F,
1437                Notif,
1438                P::Connection<Host::Counterpart>,
1439            ) -> crate::BoxFuture<'_, Result<T, crate::Error>>
1440            + Send
1441            + Sync,
1442    {
1443        let handler = NotificationHandler::<_, _, P, _, _, _>::new(
1444            self.host.counterpart(),
1445            self.host.counterpart(),
1446            op,
1447            to_future_hack,
1448        );
1449        self.with_handler(handler)
1450    }
1451
1452    /// Register a handler for messages from a specific peer.
1453    ///
1454    /// This is similar to [`on_receive_dispatch`](Self::on_receive_dispatch), but allows
1455    /// specifying the source peer explicitly. This is useful when receiving messages
1456    /// from a peer that requires message transformation (e.g., unwrapping `SuccessorMessage`
1457    /// envelopes when receiving from an agent via a proxy).
1458    ///
1459    /// For the common case of receiving from the default counterpart, use
1460    /// [`on_receive_dispatch`](Self::on_receive_dispatch) instead.
1461    ///
1462    /// # Ordering
1463    ///
1464    /// This callback runs inside the dispatch loop and blocks further message processing
1465    /// until it completes. See the [`ordering`](crate::concepts::ordering) module for details on
1466    /// ordering guarantees and how to avoid deadlocks.
1467    pub fn on_receive_dispatch_from<
1468        Req: JsonRpcRequest,
1469        Notif: JsonRpcNotification,
1470        Peer: Role,
1471        F,
1472        T,
1473        ToFut,
1474    >(
1475        self,
1476        peer: Peer,
1477        op: F,
1478        to_future_hack: ToFut,
1479    ) -> Builder<Host, impl HandleDispatchFrom<Host::Counterpart>, Runner, Close, P>
1480    where
1481        Host::Counterpart: HasPeer<Peer>,
1482        F: AsyncFnMut(
1483                Dispatch<Req, Notif, P>,
1484                P::Connection<Host::Counterpart>,
1485            ) -> Result<T, crate::Error>
1486            + Send,
1487        T: IntoHandled<Dispatch<Req, Notif, P>>,
1488        ToFut: Fn(
1489                &mut F,
1490                Dispatch<Req, Notif, P>,
1491                P::Connection<Host::Counterpart>,
1492            ) -> crate::BoxFuture<'_, Result<T, crate::Error>>
1493            + Send
1494            + Sync,
1495    {
1496        let handler = MessageHandler::<_, _, P, _, _, _, _>::new(
1497            self.host.counterpart(),
1498            peer,
1499            op,
1500            to_future_hack,
1501        );
1502        self.note_request_handler::<Req>().with_handler(handler)
1503    }
1504
1505    /// Register a handler for JSON-RPC requests from a specific peer.
1506    ///
1507    /// This is similar to [`on_receive_request`](Self::on_receive_request), but allows
1508    /// specifying the source peer explicitly. This is useful when receiving messages
1509    /// from a peer that requires message transformation (e.g., unwrapping `SuccessorRequest`
1510    /// envelopes when receiving from an agent via a proxy).
1511    ///
1512    /// For the common case of receiving from the default counterpart, use
1513    /// [`on_receive_request`](Self::on_receive_request) instead.
1514    ///
1515    /// # Example
1516    ///
1517    /// ```ignore
1518    /// use agent_client_protocol::Agent;
1519    /// use agent_client_protocol::v1::schema::InitializeRequest;
1520    ///
1521    /// // Conductor receiving from agent direction - messages will be unwrapped from SuccessorMessage
1522    /// connection.on_receive_request_from(Agent, async |req: InitializeRequest, responder, cx| {
1523    ///     // Handle the request
1524    ///     responder.respond(InitializeResponse::make())
1525    /// })
1526    /// ```
1527    ///
1528    /// # Ordering
1529    ///
1530    /// This callback runs inside the dispatch loop and blocks further message processing
1531    /// until it completes. See the [`ordering`](crate::concepts::ordering) module for details on
1532    /// ordering guarantees and how to avoid deadlocks.
1533    pub fn on_receive_request_from<Req: JsonRpcRequest, Peer: Role, F, T, ToFut>(
1534        self,
1535        peer: Peer,
1536        op: F,
1537        to_future_hack: ToFut,
1538    ) -> Builder<Host, impl HandleDispatchFrom<Host::Counterpart>, Runner, Close, P>
1539    where
1540        Host::Counterpart: HasPeer<Peer>,
1541        F: AsyncFnMut(
1542                Req,
1543                Responder<Req::Response, P>,
1544                P::Connection<Host::Counterpart>,
1545            ) -> Result<T, crate::Error>
1546            + Send,
1547        T: IntoHandled<(Req, Responder<Req::Response, P>)>,
1548        ToFut: Fn(
1549                &mut F,
1550                Req,
1551                Responder<Req::Response, P>,
1552                P::Connection<Host::Counterpart>,
1553            ) -> crate::BoxFuture<'_, Result<T, crate::Error>>
1554            + Send
1555            + Sync,
1556    {
1557        let handler = RequestHandler::<_, _, P, _, _, _>::new(
1558            self.host.counterpart(),
1559            peer,
1560            op,
1561            to_future_hack,
1562        );
1563        self.note_request_handler::<Req>().with_handler(handler)
1564    }
1565
1566    /// Register a handler for JSON-RPC notifications from a specific peer.
1567    ///
1568    /// This is similar to [`on_receive_notification`](Self::on_receive_notification), but allows
1569    /// specifying the source peer explicitly. This is useful when receiving messages
1570    /// from a peer that requires message transformation (e.g., unwrapping `SuccessorNotification`
1571    /// envelopes when receiving from an agent via a proxy).
1572    ///
1573    /// For the common case of receiving from the default counterpart, use
1574    /// [`on_receive_notification`](Self::on_receive_notification) instead.
1575    ///
1576    /// # Ordering
1577    ///
1578    /// This callback runs inside the dispatch loop and blocks further message processing
1579    /// until it completes. See the [`ordering`](crate::concepts::ordering) module for details on
1580    /// ordering guarantees and how to avoid deadlocks.
1581    pub fn on_receive_notification_from<Notif: JsonRpcNotification, Peer: Role, F, T, ToFut>(
1582        self,
1583        peer: Peer,
1584        op: F,
1585        to_future_hack: ToFut,
1586    ) -> Builder<Host, impl HandleDispatchFrom<Host::Counterpart>, Runner, Close, P>
1587    where
1588        Host::Counterpart: HasPeer<Peer>,
1589        F: AsyncFnMut(Notif, P::Connection<Host::Counterpart>) -> Result<T, crate::Error> + Send,
1590        T: IntoHandled<(Notif, P::Connection<Host::Counterpart>)>,
1591        ToFut: Fn(
1592                &mut F,
1593                Notif,
1594                P::Connection<Host::Counterpart>,
1595            ) -> crate::BoxFuture<'_, Result<T, crate::Error>>
1596            + Send
1597            + Sync,
1598    {
1599        let handler = NotificationHandler::<_, _, P, _, _, _>::new(
1600            self.host.counterpart(),
1601            peer,
1602            op,
1603            to_future_hack,
1604        );
1605        self.with_handler(handler)
1606    }
1607
1608    /// Run in server mode with the provided transport.
1609    ///
1610    /// This drives the connection by continuously processing messages from the transport
1611    /// and dispatching them to your registered handlers. The connection will run until:
1612    /// - The transport closes (e.g., EOF on byte streams)
1613    /// - An error occurs
1614    ///
1615    /// Handler errors are normally contained: requests receive an Error Response,
1616    /// response-handler errors are routed to the pending local request, and
1617    /// notification errors are logged without a wire reply.
1618    ///
1619    /// On clean EOF, messages already accepted by the outgoing queue—including
1620    /// handler responses and close-callback notifications—are drained through
1621    /// the transport sink before this returns `Ok(())`.
1622    ///
1623    /// The transport boundary carries [`TransportFrame`] values. Physical stream adapters
1624    /// serialize and deserialize frames, while channel-based components relay them directly.
1625    ///
1626    /// Use this mode when you only need to respond to incoming messages and don't need
1627    /// to initiate your own requests. If you need to send requests to the other side,
1628    /// use [`connect_with`](Self::connect_with) instead.
1629    ///
1630    /// # Example: Byte Stream Transport
1631    ///
1632    /// ```no_run
1633    /// # use agent_client_protocol::UntypedRole;
1634    /// # use agent_client_protocol::{Builder};
1635    /// # use agent_client_protocol_test::*;
1636    /// # async fn example(transport: impl agent_client_protocol::ConnectTo<UntypedRole>) -> Result<(), agent_client_protocol::Error> {
1637    ///
1638    /// UntypedRole.builder()
1639    ///     .on_receive_request(async |req: MyRequest, responder, cx| {
1640    ///         responder.respond(MyResponse { status: "ok".into() })
1641    ///     }, agent_client_protocol::on_receive_request!())
1642    ///     .connect_to(transport)
1643    ///     .await?;
1644    /// # Ok(())
1645    /// # }
1646    /// ```
1647    pub async fn connect_to(
1648        self,
1649        transport: impl ConnectTo<Host> + 'static,
1650    ) -> Result<(), crate::Error> {
1651        let (_, future) = self.into_connection_and_future(transport, true, async move |cx| {
1652            cx.incoming_closed().await;
1653            Ok(())
1654        });
1655        future.await
1656    }
1657
1658    /// Run the connection until the provided closure completes.
1659    ///
1660    /// This drives the connection by:
1661    /// 1. Running your registered handlers in the background to process incoming messages
1662    /// 2. Executing your `main_fn` closure with the builder-selected connection
1663    ///    context for sending requests and notifications
1664    ///
1665    /// The connection stays active until your `main_fn` returns, then shuts down.
1666    /// Clean incoming EOF fails every pending request and makes future
1667    /// requests fail immediately. It does not cancel unrelated work in
1668    /// `main_fn`: that future may observe the context's `incoming_closed`
1669    /// future, or the builder can use [`on_close`](Self::on_close) to notify it
1670    /// or return an error and stop it.
1671    ///
1672    /// Use this mode when you need to initiate communication (send requests/notifications)
1673    /// in addition to responding to incoming messages. For server-only mode where you just
1674    /// respond to messages, use [`connect_to`](Self::connect_to) instead.
1675    ///
1676    /// # Example
1677    ///
1678    /// ```no_run
1679    /// # use agent_client_protocol::UntypedRole;
1680    /// # use agent_client_protocol::{Builder};
1681    /// # use agent_client_protocol::ByteStreams;
1682    /// # use agent_client_protocol::v1::schema::InitializeRequest;
1683    /// # use agent_client_protocol_test::*;
1684    /// # async fn example(transport: impl agent_client_protocol::ConnectTo<UntypedRole>) -> Result<(), agent_client_protocol::Error> {
1685    ///
1686    /// UntypedRole.builder()
1687    ///     .on_receive_request(async |req: MyRequest, responder, cx| {
1688    ///         // Handle incoming requests in the background
1689    ///         responder.respond(MyResponse { status: "ok".into() })
1690    ///     }, agent_client_protocol::on_receive_request!())
1691    ///     .connect_with(transport, async |cx| {
1692    ///         // Initialize the protocol
1693    ///         let init_response = cx.send_request(InitializeRequest::make())
1694    ///             .block_task()
1695    ///             .await?;
1696    ///
1697    ///         // Send more requests...
1698    ///         let result = cx.send_request(MyRequest {})
1699    ///             .block_task()
1700    ///             .await?;
1701    ///
1702    ///         // When this closure returns, the connection shuts down
1703    ///         Ok(())
1704    ///     })
1705    ///     .await?;
1706    /// # Ok(())
1707    /// # }
1708    /// ```
1709    ///
1710    /// # Parameters
1711    ///
1712    /// - `main_fn`: Your client logic. Receives the builder-selected connection
1713    ///   context for sending messages.
1714    ///
1715    /// # Errors
1716    ///
1717    /// Returns an error if a handler, background task, transport, or close
1718    /// callback fails, or if `main_fn` returns an error. Clean incoming EOF is
1719    /// observable through the context's `incoming_closed` future and is not
1720    /// itself an error in this mode.
1721    pub async fn connect_with<R>(
1722        self,
1723        transport: impl ConnectTo<Host> + 'static,
1724        main_fn: impl AsyncFnOnce(P::Connection<Host::Counterpart>) -> Result<R, crate::Error>,
1725    ) -> Result<R, crate::Error> {
1726        let (_, future) =
1727            self.into_connection_and_future(transport, false, async move |connection| {
1728                main_fn(protocol::connection_from_raw::<P, _>(connection)).await
1729            });
1730        future.await
1731    }
1732
1733    /// Helper that returns a [`ConnectionTo<R>`] and a future that runs this connection until `main_fn` returns.
1734    fn into_connection_and_future<R>(
1735        self,
1736        transport: impl ConnectTo<Host> + 'static,
1737        wait_owned_transport: bool,
1738        main_fn: impl AsyncFnOnce(ConnectionTo<Host::Counterpart>) -> Result<R, crate::Error>,
1739    ) -> (
1740        ConnectionTo<Host::Counterpart>,
1741        impl Future<Output = Result<R, crate::Error>>,
1742    ) {
1743        let Self {
1744            name,
1745            handler,
1746            runner,
1747            host: me,
1748            protocol_mode,
1749            on_close,
1750            protocol: _,
1751        } = self;
1752
1753        let (outgoing_tx, outgoing_rx) = mpsc::unbounded();
1754        let (new_task_tx, new_task_rx) = mpsc::unbounded();
1755        let (dynamic_handler_tx, dynamic_handler_rx) = mpsc::unbounded();
1756        let (foreground_succeeded_tx, foreground_succeeded) = completion_signal();
1757        let (foreground_done_tx, foreground_done) = completion_signal();
1758        let pending_replies = PendingReplies::default();
1759        let termination_sources = TerminationSources::default();
1760
1761        // Normalize the transport without losing ownership or finish metadata.
1762        let transport_component = crate::DynConnectTo::new(transport);
1763        let (transport_channel, mut transport_future) =
1764            transport_component.into_channel_and_future();
1765        let owned_transport = transport_future.is_some();
1766        let transport_finish = transport_future
1767            .as_mut()
1768            .and_then(crate::ConnectionDriver::take_finish);
1769        let (transport_completion_tx, transport_completion_rx) = oneshot::channel();
1770        let transport_completion = transport_completion_rx
1771            .map(|result| {
1772                result.unwrap_or_else(|error| {
1773                    Err(crate::util::internal_error(format!(
1774                        "transport task dropped before reporting completion: {error}"
1775                    )))
1776                })
1777            })
1778            .boxed()
1779            .shared();
1780
1781        let connection = ConnectionTo::new(
1782            me.counterpart(),
1783            outgoing_tx,
1784            new_task_tx,
1785            dynamic_handler_tx,
1786            transport_completion,
1787            pending_replies.registrar(),
1788            protocol_mode,
1789        );
1790        // Transport progress must outlive successful foreground completion.
1791        // Application tasks remain cancellable. The inherited close-callback
1792        // phase still polls them for cleanup, but physical drain alone does not.
1793        let transport_driver = if let Some(driver) = transport_future {
1794            let termination_sources = termination_sources.clone();
1795            async move {
1796                let result = driver.await;
1797                termination_sources.record(TerminationSource::Transport, &result);
1798                drop(transport_completion_tx.send(result.clone()));
1799                result
1800            }
1801            .boxed()
1802        } else {
1803            // Channel-only endpoints have no physical sink work to await.
1804            // Their protocol drain marker still orders accepted output.
1805            drop(transport_completion_tx.send(Ok(())));
1806            future::ready(Ok(())).boxed()
1807        };
1808
1809        // Destructure the channel endpoints
1810        let Channel {
1811            rx: mut transport_incoming_rx,
1812            tx: transport_outgoing_tx,
1813        } = transport_channel;
1814
1815        let transport_incoming = futures::stream::poll_fn({
1816            let mut completion = connection.transport_completion.clone();
1817            let mut completed = false;
1818            move |cx| {
1819                if owned_transport
1820                    && !completed
1821                    && let std::task::Poll::Ready(Ok(())) =
1822                        std::pin::Pin::new(&mut completion).poll(cx)
1823                {
1824                    // Owned completion closes the producer boundary, not the
1825                    // accepted buffer. The incoming actor still dispatches every
1826                    // accepted frame and completes its close callbacks in order.
1827                    transport_incoming_rx.close();
1828                    completed = true;
1829                }
1830                transport_incoming_rx.poll_next_unpin(cx)
1831            }
1832        });
1833
1834        let protocol_compat = ProtocolCompat::new(protocol_mode);
1835
1836        // Constructed outside the driver so dropping it unpolled still
1837        // publishes a termination.
1838        let termination_guard = TerminationGuard {
1839            incoming_closed: connection.incoming_closed.clone(),
1840            pending_replies: pending_replies.clone(),
1841        };
1842
1843        let future = crate::util::instrument_with_connection_name(name, {
1844            let connection = connection.clone();
1845            async move {
1846                let sources = &termination_sources;
1847                let driver = async {
1848                    let background = async {
1849                        let incoming = {
1850                            let pending_replies = pending_replies.clone();
1851                            let protocol_compat = protocol_compat.clone();
1852                            async {
1853                                let mut transport_incoming = std::pin::pin!(transport_incoming);
1854                                let incoming = incoming_actor::incoming_protocol_actor(
1855                                    me.counterpart(),
1856                                    &connection,
1857                                    transport_incoming.as_mut(),
1858                                    dynamic_handler_rx,
1859                                    pending_replies,
1860                                    incoming_actor::IncomingHandlers::new(
1861                                        handler,
1862                                        on_close,
1863                                        foreground_succeeded.clone(),
1864                                    ),
1865                                    protocol_compat,
1866                                );
1867                                // Success stops delivery, not physical I/O. An
1868                                // underway close callback finishes before sealing.
1869                                let result = run_incoming_until_foreground_succeeds(
1870                                    incoming,
1871                                    foreground_succeeded,
1872                                    connection.incoming_closed.clone(),
1873                                )
1874                                .await;
1875                                sources.record(TerminationSource::Handler, &result);
1876                                if result.is_err() {
1877                                    connection.request_shutdown();
1878                                }
1879                                result?;
1880                                // Keep the raw producer boundary alive while the
1881                                // physical driver drains. Discard without application
1882                                // delivery, rather than fail a late read-side send.
1883                                while transport_incoming.next().await.is_some() {}
1884                                Ok(())
1885                            }
1886                        };
1887                        let other_actors = async {
1888                            let result = futures::try_join!(
1889                                // A ready driver error is authoritative even if its
1890                                // closed channel would also make output forwarding fail.
1891                                transport_driver,
1892                                // Protocol layer: OutgoingMessage -> RawJsonRpcMessage
1893                                outgoing_actor::outgoing_protocol_actor(
1894                                    outgoing_rx,
1895                                    pending_replies,
1896                                    transport_outgoing_tx,
1897                                    protocol_compat,
1898                                    foreground_done,
1899                                )
1900                                .inspect(
1901                                    |result| sources.record(TerminationSource::Transport, result)
1902                                ),
1903                            );
1904                            // Signal before awaiting an underway close callback.
1905                            if result.is_err() {
1906                                connection.request_shutdown();
1907                            }
1908                            result?;
1909                            Ok(())
1910                        };
1911
1912                        // Keep close callbacks alive when another core actor fails.
1913                        // The outer coordination provides the same protection when
1914                        // EOF wakes an application task or the foreground.
1915                        run_until_connection_close(
1916                            incoming,
1917                            other_actors,
1918                            connection.incoming_closed.clone(),
1919                        )
1920                        .await
1921                    };
1922
1923                    run_until_connection_close(
1924                        finish_actor_error(background, &connection),
1925                        async {
1926                            let application = async {
1927                                futures::try_join!(
1928                                    finish_actor_error(
1929                                        task_actor::task_actor(new_task_rx, &connection)
1930                                            .inspect(|result| sources
1931                                                .record(TerminationSource::Handler, result),),
1932                                        &connection,
1933                                    ),
1934                                    finish_actor_error(
1935                                        runner
1936                                            .run_with_connection_to(connection.clone())
1937                                            .inspect(|result| sources
1938                                                .record(TerminationSource::Handler, result),),
1939                                        &connection,
1940                                    ),
1941                                )?;
1942                                Ok(())
1943                            };
1944                            let result = run_until_connection_close(
1945                                application,
1946                                async {
1947                                    let result = main_fn(connection.clone()).await;
1948                                    sources.record(TerminationSource::Foreground, &result);
1949                                    connection.request_shutdown();
1950                                    if result.is_ok() {
1951                                        // Stop new incoming delivery immediately,
1952                                        // including during the callback cleanup phase.
1953                                        let _ = foreground_succeeded_tx.send(());
1954                                    }
1955                                    // Shutdown cancels local consumers, not remote
1956                                    // requests. Do not add cancellation traffic while
1957                                    // dropping those consumers before the drain.
1958                                    connection.pending_replies.disarm_cancellations();
1959                                    // The application actor (including actual scoped
1960                                    // runners) remains polled while supervisors finish.
1961                                    // Only then may ordinary application tasks drop.
1962                                    connection.wait_protected_operations().await;
1963                                    result
1964                                },
1965                                connection.incoming_closed.clone(),
1966                            )
1967                            .await?;
1968                            let _ = foreground_done_tx.send(());
1969                            let drained = connection
1970                                .drain_outgoing(transport_finish, wait_owned_transport)
1971                                .await;
1972                            sources.record(TerminationSource::Transport, &drained);
1973                            drained?;
1974                            Ok(result)
1975                        },
1976                        connection.incoming_closed.clone(),
1977                    )
1978                    .await
1979                };
1980                // Every actor, handler, and task has completed or been
1981                // dropped by the time `driver` resolves.
1982                let result = driver.await;
1983                termination_guard.finish(match &result {
1984                    // A foreground that finishes after observing EOF still
1985                    // ended at a clean EOF boundary.
1986                    Ok(_) if connection.incoming_closed.is_closing() => {
1987                        ConnectionTermination::IncomingClosed
1988                    }
1989                    Ok(_) => ConnectionTermination::LocalShutdown,
1990                    Err(error) => sources.classify(error.clone()),
1991                });
1992                result
1993            }
1994        });
1995
1996        (connection, future)
1997    }
1998}
1999
2000/// Defer an actor error without dropping the other actors that drive owned
2001/// cleanup or an underway close callback. Physical output drain stays separate.
2002async fn finish_actor_error<R: Role>(
2003    actor: impl Future<Output = Result<(), crate::Error>>,
2004    connection: &ConnectionTo<R>,
2005) -> Result<(), crate::Error> {
2006    let result = actor.await;
2007    if result.is_err() {
2008        connection.request_shutdown();
2009        if connection.incoming_closed.is_closing() {
2010            connection.incoming_closed.closed().await;
2011        }
2012        connection.wait_protected_operations().await;
2013    }
2014    result
2015}
2016
2017#[cfg(feature = "unstable_mcp_over_acp")]
2018impl<
2019    Host: Role,
2020    Handler: HandleDispatchFrom<Host::Counterpart>,
2021    Runner: RunWithConnectionTo<Host::Counterpart>,
2022    Close: HandleConnectionClose<Host::Counterpart>,
2023> Builder<Host, Handler, Runner, Close, V1>
2024{
2025    /// Add an MCP server to protocol v1 session setup requests proxied through
2026    /// this connection.
2027    ///
2028    /// The same native MCP server declaration is added to new, load, and resume
2029    /// requests, plus fork requests when `unstable_session_fork` is enabled.
2030    /// Unrelated request fields are preserved exactly.
2031    ///
2032    /// Only applicable to proxies. Use the same method on `v2::Builder` to
2033    /// attach the server to protocol v2 setup requests.
2034    ///
2035    /// The declaration is an `acp` transport (`McpServer::Acp`). The agent must
2036    /// advertise `agentCapabilities.mcpCapabilities.acp`, or a compatibility
2037    /// proxy such as `McpOverAcpPolyfill` must sit between this proxy and the
2038    /// agent; otherwise a conforming agent rejects the setup request with
2039    /// `-32602`. This method does not check the agent's capabilities.
2040    pub fn with_mcp_server(
2041        self,
2042        mcp_server: McpServer<Host::Counterpart, impl RunWithConnectionTo<Host::Counterpart>>,
2043    ) -> Builder<
2044        Host,
2045        impl HandleDispatchFrom<Host::Counterpart>,
2046        impl RunWithConnectionTo<Host::Counterpart>,
2047        Close,
2048        V1,
2049    >
2050    where
2051        Host::Counterpart: HasPeer<Agent> + HasPeer<Client>,
2052    {
2053        let (handler, runner) = mcp_server.into_handler_and_runner();
2054        self.with_handler(handler).with_runner(runner)
2055    }
2056}
2057
2058#[cfg(all(feature = "unstable_mcp_over_acp", feature = "unstable_protocol_v2"))]
2059impl<
2060    Host: Role,
2061    Handler: HandleDispatchFrom<Host::Counterpart>,
2062    Runner: RunWithConnectionTo<Host::Counterpart>,
2063    Close: HandleConnectionClose<Host::Counterpart>,
2064> Builder<Host, Handler, Runner, Close, V2>
2065{
2066    /// Add an MCP server to protocol v2 session setup requests proxied through
2067    /// this connection.
2068    ///
2069    /// The same native MCP server declaration is added to new and resume
2070    /// requests, plus fork requests when `unstable_session_fork` is enabled.
2071    /// Unrelated request fields are preserved exactly.
2072    ///
2073    /// MCP server names must be unique within `mcpServers`. If a client's
2074    /// setup request already declares a server with this server's name, the
2075    /// proxy answers that request with `-32602` naming the conflicting server
2076    /// and does not forward it. The client's servers are never renamed or
2077    /// dropped.
2078    ///
2079    /// Only applicable to proxies.
2080    ///
2081    /// The declaration is an `acp` transport (`McpServer::Acp`). The downstream
2082    /// agent must advertise `session.mcp.acp`, or a compatibility proxy such as
2083    /// `McpOverAcpPolyfill` must sit between this proxy and the agent;
2084    /// otherwise a conforming agent rejects the setup request with `-32602`.
2085    /// This method does not check the agent's capabilities.
2086    pub fn with_mcp_server(
2087        self,
2088        mcp_server: McpServer<Host::Counterpart, impl RunWithConnectionTo<Host::Counterpart>>,
2089    ) -> Builder<
2090        Host,
2091        impl HandleDispatchFrom<Host::Counterpart>,
2092        impl RunWithConnectionTo<Host::Counterpart>,
2093        Close,
2094        V2,
2095    >
2096    where
2097        Host::Counterpart: HasPeer<Agent> + HasPeer<Client>,
2098    {
2099        let (handler, runner) = mcp_server.into_v2_handler_and_runner();
2100        self.with_handler(handler).with_runner(runner)
2101    }
2102}
2103
2104impl<R, H, Run, Close, P> ConnectTo<R::Counterpart> for Builder<R, H, Run, Close, P>
2105where
2106    R: Role,
2107    H: HandleDispatchFrom<R::Counterpart> + 'static,
2108    Run: RunWithConnectionTo<R::Counterpart> + 'static,
2109    Close: HandleConnectionClose<R::Counterpart> + 'static,
2110    P: Protocol,
2111{
2112    async fn connect_to(self, client: impl ConnectTo<R>) -> Result<(), crate::Error> {
2113        Builder::connect_to(self, client).await
2114    }
2115}
2116
2117/// The payload sent through the response oneshot channel.
2118///
2119/// Includes the response value and an optional ack channel for dispatch loop
2120/// synchronization.
2121pub(crate) struct ResponsePayload {
2122    /// The response result - either the JSON value or an error.
2123    pub(crate) result: Result<serde_json::Value, crate::Error>,
2124
2125    /// Optional acknowledgment channel for dispatch loop synchronization.
2126    ///
2127    /// When present, the receiver must send on this channel to signal that
2128    /// response processing is complete, allowing the dispatch loop to continue
2129    /// to the next message.
2130    ///
2131    /// This is present when ordered response consumption was selected before
2132    /// the response was routed during its original dispatch. Public callback
2133    /// consumption and framework-owned ordered blocking transforms can hold
2134    /// the dispatch loop; ordinary blocking consumers, local error paths, and
2135    /// responses routed later do not.
2136    pub(crate) ack_tx: Option<oneshot::Sender<()>>,
2137}
2138
2139type ResponseRouteHook =
2140    Box<dyn FnOnce(&str, &serde_json::Value) -> Result<(), crate::Error> + Send>;
2141
2142/// A prerequisite that must complete before an outgoing request is published
2143/// to the transport.
2144struct RequestReadiness {
2145    future: BoxFuture<'static, Result<(), crate::Error>>,
2146}
2147
2148impl RequestReadiness {
2149    fn new(future: impl Future<Output = Result<(), crate::Error>> + Send + 'static) -> Self {
2150        Self {
2151            future: future.boxed(),
2152        }
2153    }
2154}
2155
2156impl Future for RequestReadiness {
2157    type Output = Result<(), crate::Error>;
2158
2159    fn poll(
2160        mut self: std::pin::Pin<&mut Self>,
2161        cx: &mut std::task::Context<'_>,
2162    ) -> std::task::Poll<Self::Output> {
2163        self.future.as_mut().poll(cx)
2164    }
2165}
2166
2167impl Debug for RequestReadiness {
2168    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
2169        formatter
2170            .debug_struct("RequestReadiness")
2171            .finish_non_exhaustive()
2172    }
2173}
2174
2175impl std::fmt::Debug for ResponsePayload {
2176    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
2177        f.debug_struct("ResponsePayload")
2178            .field("result", &self.result)
2179            .field("ack_tx", &self.ack_tx.as_ref().map(|_| "..."))
2180            .finish()
2181    }
2182}
2183
2184#[derive(Clone, Debug, Default)]
2185struct ResponseOrdering {
2186    ordered: Arc<AtomicBool>,
2187}
2188
2189impl ResponseOrdering {
2190    fn mark_ordered(&self) {
2191        self.ordered.store(true, Ordering::Release);
2192    }
2193
2194    fn is_ordered(&self) -> bool {
2195        self.ordered.load(Ordering::Acquire)
2196    }
2197}
2198
2199struct PendingReply {
2200    id: RawJsonRpcId,
2201    method: String,
2202    role_id: RoleId,
2203    sender: oneshot::Sender<ResponsePayload>,
2204    cancellation_disarm: SentRequestCancellationDisarm,
2205    ordering: ResponseOrdering,
2206    response_route_hook: Option<ResponseRouteHook>,
2207    /// Set only by `v2::ConnectionTo::send_initialize`, whose typed outcome
2208    /// reports a lower-version initialize response as a success.
2209    accepts_lower_initialize_version: bool,
2210}
2211
2212impl PendingReply {
2213    fn fail(self, error: crate::Error) {
2214        self.cancellation_disarm.disarm();
2215        if self
2216            .sender
2217            .send(ResponsePayload {
2218                result: Err(error),
2219                ack_tx: None,
2220            })
2221            .is_err()
2222        {
2223            tracing::trace!(method = %self.method, "Pending request was already dropped");
2224        }
2225    }
2226
2227    fn fail_incoming_closed(self) {
2228        let error = incoming_transport_closed_error(&self.method);
2229        self.fail(error);
2230    }
2231}
2232
2233#[derive(Default)]
2234struct PendingRepliesInner {
2235    incoming_closed: bool,
2236    termination: Option<ConnectionTermination>,
2237    replies: HashMap<RawJsonRpcId, PendingReply>,
2238}
2239
2240#[derive(Clone, Default)]
2241struct PendingReplies {
2242    inner: Arc<Mutex<PendingRepliesInner>>,
2243}
2244
2245impl PendingReplies {
2246    fn registrar(&self) -> PendingRepliesRegistrar {
2247        PendingRepliesRegistrar {
2248            inner: Arc::downgrade(&self.inner),
2249        }
2250    }
2251
2252    fn contains(&self, id: &RawJsonRpcId) -> bool {
2253        self.inner
2254            .lock()
2255            .expect("pending replies mutex poisoned")
2256            .replies
2257            .contains_key(id)
2258    }
2259
2260    fn remove(&self, id: &RawJsonRpcId) -> Option<PendingReply> {
2261        self.inner
2262            .lock()
2263            .expect("pending replies mutex poisoned")
2264            .replies
2265            .remove(id)
2266    }
2267
2268    /// Atomically reject new subscriptions and fail every existing one.
2269    fn close_incoming(&self) -> usize {
2270        let replies = {
2271            let mut inner = self.inner.lock().expect("pending replies mutex poisoned");
2272            inner.incoming_closed = true;
2273            std::mem::take(&mut inner.replies)
2274        };
2275        let count = replies.len();
2276        for (_, reply) in replies {
2277            reply.fail_incoming_closed();
2278        }
2279        count
2280    }
2281
2282    /// Atomically reject new subscriptions and fail every remaining one
2283    /// because the driver has terminated.
2284    fn terminate(&self, termination: &ConnectionTermination) -> usize {
2285        let replies = {
2286            let mut inner = self.inner.lock().expect("pending replies mutex poisoned");
2287            inner.termination = Some(termination.clone());
2288            std::mem::take(&mut inner.replies)
2289        };
2290        let count = replies.len();
2291        for (_, reply) in replies {
2292            let error = connection_terminated_error(&reply.method, termination);
2293            reply.fail(error);
2294        }
2295        count
2296    }
2297}
2298
2299/// A non-owning handle used to register a request before it enters the
2300/// outgoing queue. Keeping this weak prevents escaped [`ConnectionTo`] clones
2301/// from extending the lifetime of response senders after the driver stops.
2302#[derive(Clone)]
2303struct PendingRepliesRegistrar {
2304    inner: Weak<Mutex<PendingRepliesInner>>,
2305}
2306
2307impl PendingRepliesRegistrar {
2308    fn disarm_cancellations(&self) {
2309        if let Some(inner) = self.inner.upgrade() {
2310            let inner = inner.lock().expect("pending replies mutex poisoned");
2311            for reply in inner.replies.values() {
2312                reply.cancellation_disarm.disarm();
2313            }
2314        }
2315    }
2316
2317    /// Register a response destination before the request becomes observable.
2318    ///
2319    /// Returns an error after failing `reply` when EOF has already made a
2320    /// response impossible or the connection driver is no longer running.
2321    fn subscribe(
2322        &self,
2323        id: RawJsonRpcId,
2324        reply: PendingReply,
2325        incoming_closed: &IncomingClosed,
2326    ) -> Result<(), crate::Error> {
2327        let Some(inner) = self.inner.upgrade() else {
2328            let error = if incoming_closed.is_closing() {
2329                incoming_transport_closed_error(&reply.method)
2330            } else if let Some(termination) = incoming_closed.termination() {
2331                connection_terminated_error(&reply.method, &termination)
2332            } else {
2333                crate::util::internal_error(format!(
2334                    "failed to send outgoing request `{}`: connection is no longer running",
2335                    reply.method
2336                ))
2337            };
2338            reply.fail(error.clone());
2339            return Err(error);
2340        };
2341
2342        let result = {
2343            let mut inner = inner.lock().expect("pending replies mutex poisoned");
2344            if inner.incoming_closed {
2345                Err((reply, None))
2346            } else if let Some(termination) = &inner.termination {
2347                Err((reply, Some(termination.clone())))
2348            } else {
2349                Ok(inner.replies.insert(id, reply))
2350            }
2351        };
2352
2353        match result {
2354            Err((rejected, termination)) => {
2355                let error = match termination {
2356                    Some(termination) => {
2357                        connection_terminated_error(&rejected.method, &termination)
2358                    }
2359                    None => incoming_transport_closed_error(&rejected.method),
2360                };
2361                rejected.fail(error.clone());
2362                Err(error)
2363            }
2364            Ok(replaced) => {
2365                if let Some(replaced) = replaced {
2366                    replaced.fail(
2367                        crate::Error::internal_error()
2368                            .data("outgoing request ID was reused before its response arrived"),
2369                    );
2370                }
2371                Ok(())
2372            }
2373        }
2374    }
2375
2376    fn remove(&self, id: &RawJsonRpcId) -> Option<PendingReply> {
2377        self.inner
2378            .upgrade()?
2379            .lock()
2380            .expect("pending replies mutex poisoned")
2381            .replies
2382            .remove(id)
2383    }
2384}
2385
2386impl Debug for PendingRepliesRegistrar {
2387    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
2388        formatter
2389            .debug_struct("PendingRepliesRegistrar")
2390            .field("is_connected", &(self.inner.strong_count() > 0))
2391            .finish()
2392    }
2393}
2394
2395/// A request-local marker that is set when the peer asks to cancel the request.
2396///
2397/// Request handlers can get this handle from [`Responder::cancellation`] and
2398/// use it from spawned work to stop long-running request processing
2399/// cooperatively.
2400#[derive(Clone)]
2401pub struct RequestCancellation {
2402    state: Arc<RequestCancellationState>,
2403}
2404
2405struct RequestCancellationState {
2406    cancelled: AtomicBool,
2407    signal_tx: Mutex<Option<oneshot::Sender<()>>>,
2408    signal_rx: future::Shared<BoxFuture<'static, ()>>,
2409}
2410
2411impl RequestCancellation {
2412    fn new() -> Self {
2413        let (signal_tx, signal_rx) = oneshot::channel();
2414        let signal_rx = signal_rx.map(|_| ()).boxed().shared();
2415        Self {
2416            state: Arc::new(RequestCancellationState {
2417                cancelled: AtomicBool::new(false),
2418                signal_tx: Mutex::new(Some(signal_tx)),
2419                signal_rx,
2420            }),
2421        }
2422    }
2423
2424    /// Wait until the peer sends `$/cancel_request` for this request.
2425    ///
2426    /// If cancellation was already requested, this returns immediately.
2427    pub async fn cancelled(&self) {
2428        self.state.signal_rx.clone().await;
2429    }
2430
2431    /// Run request work until it completes or the peer asks to cancel it.
2432    ///
2433    /// If cancellation is requested first, this returns
2434    /// [`Error::request_cancelled`]. This is a convenience for request handlers
2435    /// that want to respond with the normal result or the standard
2436    /// cancellation error.
2437    ///
2438    /// When cancellation wins, `future` is dropped: work stops at its next
2439    /// await point, partial results are lost, and any cleanup must happen in
2440    /// `Drop` implementations. Handlers that need to flush partial results or
2441    /// run async cleanup should instead watch [`cancelled`](Self::cancelled)
2442    /// or poll [`is_cancelled`](Self::is_cancelled) from inside the work.
2443    ///
2444    /// [`Error::request_cancelled`]: crate::Error::request_cancelled
2445    pub async fn run_until_cancelled<T>(
2446        &self,
2447        future: impl std::future::Future<Output = Result<T, crate::Error>>,
2448    ) -> Result<T, crate::Error> {
2449        if self.is_cancelled() {
2450            return Err(crate::Error::request_cancelled());
2451        }
2452
2453        match future::select(pin!(future), pin!(self.cancelled())).await {
2454            Either::Left((result, _)) => result,
2455            Either::Right(((), _)) => Err(crate::Error::request_cancelled()),
2456        }
2457    }
2458
2459    /// Returns whether the peer has already requested cancellation.
2460    #[must_use]
2461    pub fn is_cancelled(&self) -> bool {
2462        self.state.cancelled.load(Ordering::Acquire)
2463    }
2464
2465    fn cancel(&self) {
2466        if self.state.cancelled.swap(true, Ordering::AcqRel) {
2467            return;
2468        }
2469
2470        let signal_tx = self
2471            .state
2472            .signal_tx
2473            .lock()
2474            .expect("request cancellation signal mutex poisoned")
2475            .take();
2476
2477        // Complete the oneshot outside the lock: it wakes waiters, and
2478        // arbitrary waker code must not observe the lock held.
2479        if let Some(signal_tx) = signal_tx {
2480            let _ = signal_tx.send(());
2481        }
2482    }
2483}
2484
2485impl Debug for RequestCancellation {
2486    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
2487        formatter
2488            .debug_struct("RequestCancellation")
2489            .field("is_cancelled", &self.is_cancelled())
2490            .finish_non_exhaustive()
2491    }
2492}
2493
2494/// Per-request cancellation state tracked by [`RequestCancellationRegistry`].
2495///
2496/// The full [`RequestCancellation`] marker (with its wakeup machinery) is only
2497/// allocated once a handler asks for it via [`Responder::cancellation`]; until
2498/// then an incoming `$/cancel_request` just flips the entry to `Cancelled`.
2499/// This keeps the per-request cost of the registry to a single map entry.
2500#[derive(Debug)]
2501enum RequestCancellationEntry {
2502    /// The request is in flight; no marker handed out, no cancellation yet.
2503    Armed,
2504    /// `$/cancel_request` arrived before a marker was handed out.
2505    Cancelled,
2506    /// A marker was handed out via [`Responder::cancellation`].
2507    Marker(RequestCancellation),
2508}
2509
2510/// A registered request's cancellation state, tagged with the generation of
2511/// its registration.
2512///
2513/// The generation distinguishes a registration from earlier ones that used
2514/// the same request ID, so that when a (protocol-violating) peer reuses the
2515/// ID of a request that is still in flight, the stale request's responder can
2516/// neither remove nor observe the cancellation state of the newer request.
2517#[derive(Debug)]
2518struct RequestCancellationSlot {
2519    generation: u64,
2520    entry: RequestCancellationEntry,
2521}
2522
2523#[derive(Debug, Default)]
2524struct RequestCancellationRegistryInner {
2525    slots: HashMap<RawJsonRpcId, RequestCancellationSlot>,
2526    next_generation: u64,
2527}
2528
2529#[derive(Clone, Debug, Default)]
2530struct RequestCancellationRegistry {
2531    inner: Arc<Mutex<RequestCancellationRegistryInner>>,
2532}
2533
2534#[derive(Debug)]
2535struct ResponderCancellation {
2536    id: RawJsonRpcId,
2537    generation: u64,
2538    registry: RequestCancellationRegistry,
2539}
2540
2541impl RequestCancellationRegistry {
2542    fn new() -> Self {
2543        Self::default()
2544    }
2545
2546    fn register(&self, id: &RawJsonRpcId) -> ResponderCancellation {
2547        let generation = {
2548            let mut inner = self
2549                .inner
2550                .lock()
2551                .expect("request cancellation registry mutex poisoned");
2552            let generation = inner.next_generation;
2553            inner.next_generation += 1;
2554            if inner
2555                .slots
2556                .insert(
2557                    id.clone(),
2558                    RequestCancellationSlot {
2559                        generation,
2560                        entry: RequestCancellationEntry::Armed,
2561                    },
2562                )
2563                .is_some()
2564            {
2565                tracing::debug!(
2566                    ?id,
2567                    "peer reused the ID of a request that is still in flight"
2568                );
2569            }
2570            generation
2571        };
2572        ResponderCancellation {
2573            id: id.clone(),
2574            generation,
2575            registry: self.clone(),
2576        }
2577    }
2578
2579    /// Get the cancellation marker for a registered request, creating it on
2580    /// first use. Repeated calls return markers that share the same state.
2581    ///
2582    /// Exception: when the registration is stale (a protocol-violating peer
2583    /// reused this request ID and the slot now belongs to a newer request, or
2584    /// was already removed by it), every call returns a fresh *detached*
2585    /// marker. Detached markers can never fire, and detached markers from
2586    /// repeated calls do not share state with each other.
2587    fn marker(&self, id: &RawJsonRpcId, generation: u64) -> RequestCancellation {
2588        let mut inner = self
2589            .inner
2590            .lock()
2591            .expect("request cancellation registry mutex poisoned");
2592        let Some(slot) = inner.slots.get_mut(id) else {
2593            // The slot lives as long as the responder that owns it, so this
2594            // is only reachable if the peer reused this request ID and the
2595            // newer request's responder already removed the replacement slot.
2596            // Hand out a detached marker rather than panicking.
2597            return RequestCancellation::new();
2598        };
2599        if slot.generation != generation {
2600            // The peer reused this request ID while the request was still in
2601            // flight, and the slot now belongs to the newer request. Hand the
2602            // stale responder a detached marker instead of cross-wiring the
2603            // two requests' cancellation states.
2604            return RequestCancellation::new();
2605        }
2606        let entry = &mut slot.entry;
2607        match entry {
2608            RequestCancellationEntry::Marker(marker) => marker.clone(),
2609            RequestCancellationEntry::Armed => {
2610                let marker = RequestCancellation::new();
2611                *entry = RequestCancellationEntry::Marker(marker.clone());
2612                marker
2613            }
2614            RequestCancellationEntry::Cancelled => {
2615                // No one can be waiting on a marker that did not exist yet,
2616                // so firing it while holding the registry lock is fine.
2617                let marker = RequestCancellation::new();
2618                marker.cancel();
2619                *entry = RequestCancellationEntry::Marker(marker.clone());
2620                marker
2621            }
2622        }
2623    }
2624
2625    fn cancel_if_requested(&self, dispatch: &Dispatch) -> Result<bool, crate::Error> {
2626        let Some(request_id) = cancellation_request_id(dispatch)? else {
2627            return Ok(false);
2628        };
2629        Ok(self.cancel(&request_id))
2630    }
2631
2632    /// Mark whichever request currently owns `request_id` as cancelled.
2633    fn cancel(&self, request_id: &RawJsonRpcId) -> bool {
2634        let marker = {
2635            let mut inner = self
2636                .inner
2637                .lock()
2638                .expect("request cancellation registry mutex poisoned");
2639            let Some(slot) = inner.slots.get_mut(request_id) else {
2640                return false;
2641            };
2642            let entry = &mut slot.entry;
2643            match entry {
2644                RequestCancellationEntry::Marker(marker) => marker.clone(),
2645                RequestCancellationEntry::Cancelled => return true,
2646                RequestCancellationEntry::Armed => {
2647                    *entry = RequestCancellationEntry::Cancelled;
2648                    return true;
2649                }
2650            }
2651        };
2652
2653        // Fire the marker outside the registry lock: waking waiters runs
2654        // arbitrary waker code that must not observe the lock held.
2655        marker.cancel();
2656        true
2657    }
2658
2659    /// Remove the slot for `request_id`, but only if it still belongs to the
2660    /// registration identified by `generation`.
2661    fn remove(&self, request_id: &RawJsonRpcId, generation: u64) {
2662        let mut inner = self
2663            .inner
2664            .lock()
2665            .expect("request cancellation registry mutex poisoned");
2666        if inner
2667            .slots
2668            .get(request_id)
2669            .is_some_and(|slot| slot.generation == generation)
2670        {
2671            inner.slots.remove(request_id);
2672        }
2673    }
2674}
2675
2676impl ResponderCancellation {
2677    fn cancellation(&self) -> RequestCancellation {
2678        self.registry.marker(&self.id, self.generation)
2679    }
2680
2681    /// Whether the peer has cancelled this request, without allocating a
2682    /// marker.
2683    fn is_cancelled(&self) -> bool {
2684        let inner = self
2685            .registry
2686            .inner
2687            .lock()
2688            .expect("request cancellation registry mutex poisoned");
2689        let Some(slot) = inner.slots.get(&self.id) else {
2690            return false;
2691        };
2692        if slot.generation != self.generation {
2693            return false;
2694        }
2695        match &slot.entry {
2696            RequestCancellationEntry::Armed => false,
2697            RequestCancellationEntry::Cancelled => true,
2698            RequestCancellationEntry::Marker(marker) => marker.is_cancelled(),
2699        }
2700    }
2701}
2702
2703impl Drop for ResponderCancellation {
2704    fn drop(&mut self) {
2705        self.registry.remove(&self.id, self.generation);
2706    }
2707}
2708
2709fn cancellation_request_id(dispatch: &Dispatch) -> Result<Option<RawJsonRpcId>, crate::Error> {
2710    let Dispatch::Notification(message) = dispatch else {
2711        return Ok(None);
2712    };
2713    cancellation_request_id_from_message(message)
2714}
2715
2716fn cancellation_request_id_from_message(
2717    message: &UntypedMessage,
2718) -> Result<Option<RawJsonRpcId>, crate::Error> {
2719    let (method, params) = peel_successor_envelopes(&message.method, &message.params);
2720    if !crate::v1::schema::CancelRequestNotification::matches_method(method) {
2721        return Ok(None);
2722    }
2723
2724    let notification = crate::v1::schema::CancelRequestNotification::parse_message(method, params)?;
2725    Ok(Some(notification.request_id.into()))
2726}
2727
2728/// Peel any [`SuccessorMessage`] envelopes off a notification by reference,
2729/// returning the innermost method and params.
2730///
2731/// This only peeks at the envelope's `method`/`params` fields instead of
2732/// deserializing the envelope, for two reasons:
2733///
2734/// - It avoids deep-cloning the params of every wrapped notification on the
2735///   hot dispatch path just to inspect the inner method name.
2736/// - It is deliberately lenient: a malformed envelope is left as-is here and
2737///   flows on to the handler chain, which is responsible for reporting it.
2738///
2739/// [`SuccessorMessage`]: crate::schema::SuccessorMessage
2740fn peel_successor_envelopes<'message>(
2741    mut method: &'message str,
2742    mut params: &'message serde_json::Value,
2743) -> (&'message str, &'message serde_json::Value) {
2744    while crate::schema::SuccessorMessage::<UntypedMessage>::matches_method(method) {
2745        let Some(inner_method) = params.get("method").and_then(serde_json::Value::as_str) else {
2746            break;
2747        };
2748        method = inner_method;
2749        params = params.get("params").unwrap_or(&serde_json::Value::Null);
2750    }
2751    (method, params)
2752}
2753
2754/// Whether a notification is a `$/cancel_request`, even when it is still
2755/// wrapped in `proxy/successor` envelopes.
2756///
2757/// `$/cancel_request` is connection-scoped: its `requestId` was allocated on
2758/// the connection the notification arrived over and means nothing on any
2759/// other connection. Generic forwarding code (such as
2760/// [`ConnectionTo::send_proxied_message_to`]) uses this check to drop the raw
2761/// notification instead of tunneling it across a hop; the cancellation still
2762/// propagates because [`forward_response_to`](SentRequest::forward_response_to)
2763/// re-issues it with the forwarded request's own ID.
2764///
2765/// Checking a notification whose method is not the successor envelope is a
2766/// plain method-name comparison. Only successor-wrapped notifications pay for
2767/// a serialization to peel the envelope.
2768#[must_use]
2769pub fn is_cancel_request_notification<N: JsonRpcNotification>(notification: &N) -> bool {
2770    let method = notification.method();
2771    if crate::v1::schema::CancelRequestNotification::matches_method(method) {
2772        return true;
2773    }
2774    if !crate::schema::SuccessorMessage::<UntypedMessage>::matches_method(method) {
2775        return false;
2776    }
2777
2778    match notification.to_untyped_message() {
2779        Ok(untyped) => {
2780            let (method, _params) = peel_successor_envelopes(&untyped.method, &untyped.params);
2781            crate::v1::schema::CancelRequestNotification::matches_method(method)
2782        }
2783        Err(error) => {
2784            tracing::debug!(
2785                error = ?crate::util::redacted_error(error.code),
2786                "failed to inspect successor-wrapped notification for cancellation"
2787            );
2788            false
2789        }
2790    }
2791}
2792
2793/// Where the response to an incoming request is sent.
2794///
2795/// A request is answered at most once; a duplicate completion is ignored.
2796///
2797/// While a handler attempt for the request is running (see
2798/// [`begin_handler_attempt`](Self::begin_handler_attempt)), the automatic
2799/// error for a dropped [`Responder`] is held back: if the handler then
2800/// returns an error, that error is the response, not the generic fallback.
2801#[derive(Clone, Debug, Default)]
2802struct ResponseDestination {
2803    state: Arc<ResponseDestinationState>,
2804}
2805
2806#[derive(Debug, Default)]
2807struct ResponseDestinationState {
2808    completed: AtomicBool,
2809    attempts: Mutex<HandlerAttempts>,
2810}
2811
2812#[derive(Debug, Default)]
2813struct HandlerAttempts {
2814    active: usize,
2815    deferred: Option<AbandonedResponse>,
2816}
2817
2818/// The automatic error response for a [`Responder`] dropped without replying.
2819#[derive(Debug)]
2820struct AbandonedResponse {
2821    message_tx: OutgoingMessageTx,
2822    connection_state: IncomingClosed,
2823    message: OutgoingMessage,
2824}
2825
2826impl AbandonedResponse {
2827    fn send(self) {
2828        // Re-check: shutdown may have begun while the response was deferred.
2829        if self.connection_state.is_shutting_down() {
2830            tracing::debug!("not answering dropped responder: connection is shutting down");
2831            return;
2832        }
2833        if let Err(error) = send_raw_message(&self.message_tx, self.message) {
2834            tracing::debug!(
2835                error = ?crate::util::redacted_error(error.code),
2836                "could not answer request whose responder was dropped"
2837            );
2838        }
2839    }
2840}
2841
2842/// A running handler attempt for one request; see [`ResponseDestination`].
2843struct HandlerAttempt {
2844    destination: ResponseDestination,
2845}
2846
2847impl Drop for HandlerAttempt {
2848    fn drop(&mut self) {
2849        let deferred = {
2850            let mut attempts = self.destination.attempts();
2851            attempts.active -= 1;
2852            if attempts.active == 0 {
2853                attempts.deferred.take()
2854            } else {
2855                None
2856            }
2857        };
2858        if let Some(deferred) = deferred {
2859            deferred.send();
2860        }
2861    }
2862}
2863
2864impl ResponseDestination {
2865    fn individual() -> Self {
2866        Self::default()
2867    }
2868
2869    fn attempts(&self) -> std::sync::MutexGuard<'_, HandlerAttempts> {
2870        self.state
2871            .attempts
2872            .lock()
2873            .expect("response destination mutex poisoned")
2874    }
2875
2876    fn is_completed(&self) -> bool {
2877        self.state.completed.load(Ordering::Acquire)
2878    }
2879
2880    fn complete(self, response: RawJsonRpcMessage) -> Option<TransportFrame> {
2881        if self.state.completed.swap(true, Ordering::AcqRel) {
2882            tracing::warn!("Ignoring duplicate completion of JSON-RPC request");
2883            return None;
2884        }
2885
2886        Some(TransportFrame::Single(response))
2887    }
2888
2889    fn begin_handler_attempt(&self) -> HandlerAttempt {
2890        self.attempts().active += 1;
2891        HandlerAttempt {
2892            destination: self.clone(),
2893        }
2894    }
2895
2896    /// Send, or defer until the running handler attempt ends, the automatic
2897    /// error for a dropped responder.
2898    fn abandon(&self, response: AbandonedResponse) {
2899        {
2900            let mut attempts = self.attempts();
2901            if attempts.active > 0 {
2902                attempts.deferred = Some(response);
2903                return;
2904            }
2905        }
2906        response.send();
2907    }
2908
2909    /// A handler error answers the request; drop any deferred fallback.
2910    fn discard_abandoned(&self) {
2911        self.attempts().deferred = None;
2912    }
2913}
2914
2915#[derive(Clone, Debug)]
2916struct RequestReplyTarget {
2917    id: RawJsonRpcId,
2918    method: String,
2919    destination: ResponseDestination,
2920}
2921
2922#[derive(Clone)]
2923struct ResponseReplyTarget {
2924    id: RawJsonRpcId,
2925    method: String,
2926    sender: Arc<Mutex<Option<oneshot::Sender<ResponsePayload>>>>,
2927    ordering: ResponseOrdering,
2928    dispatch: ResponseDispatch,
2929}
2930
2931impl ResponseReplyTarget {
2932    fn route(self, result: Result<serde_json::Value, crate::Error>) {
2933        let sender = self
2934            .sender
2935            .lock()
2936            .expect("response reply mutex poisoned")
2937            .take();
2938        let Some(sender) = sender else {
2939            tracing::debug!(
2940                method = %self.method,
2941                id = ?self.id,
2942                "response was already routed to its local awaiter"
2943            );
2944            return;
2945        };
2946
2947        let ack_tx = self.dispatch.acknowledgment(&self.ordering);
2948        if sender.send(ResponsePayload { result, ack_tx }).is_err() {
2949            tracing::debug!(
2950                method = %self.method,
2951                id = ?self.id,
2952                "dropped response because local receiver was gone"
2953            );
2954        }
2955    }
2956}
2957
2958#[derive(Clone, Default)]
2959struct ResponseDispatch {
2960    state: Arc<Mutex<ResponseDispatchState>>,
2961}
2962
2963#[derive(Default)]
2964struct ResponseDispatchState {
2965    complete: bool,
2966    ack_rx: Option<oneshot::Receiver<()>>,
2967}
2968
2969impl ResponseDispatch {
2970    fn acknowledgment(&self, ordering: &ResponseOrdering) -> Option<oneshot::Sender<()>> {
2971        if !ordering.is_ordered() {
2972            return None;
2973        }
2974
2975        let mut state = self.state.lock().expect("response dispatch mutex poisoned");
2976        if state.complete {
2977            return None;
2978        }
2979
2980        let (ack_tx, ack_rx) = oneshot::channel();
2981        let previous_ack = state.ack_rx.replace(ack_rx);
2982        debug_assert!(
2983            previous_ack.is_none(),
2984            "a response dispatch can only be routed once"
2985        );
2986        Some(ack_tx)
2987    }
2988
2989    fn complete(&self) -> Option<oneshot::Receiver<()>> {
2990        let mut state = self.state.lock().expect("response dispatch mutex poisoned");
2991        state.complete = true;
2992        state.ack_rx.take()
2993    }
2994}
2995
2996enum HandlerErrorTarget {
2997    Request(RequestReplyTarget),
2998    Response(ResponseReplyTarget),
2999}
3000
3001impl HandlerErrorTarget {
3002    fn begin_handler_attempt(&self) -> Option<HandlerAttempt> {
3003        match self {
3004            Self::Request(target) => Some(target.destination.begin_handler_attempt()),
3005            Self::Response(_) => None,
3006        }
3007    }
3008}
3009
3010#[derive(Debug)]
3011enum OutgoingMessage {
3012    /// Close the outgoing application queue and acknowledge after every
3013    /// already-accepted message has entered the raw transport queue.
3014    CloseAfterDraining { done: oneshot::Sender<()> },
3015
3016    /// Send a request to the server.
3017    Request {
3018        /// id assigned to this request (generated by sender)
3019        id: RawJsonRpcId,
3020
3021        /// the original method
3022        method: String,
3023
3024        /// The logical message before peer-direction wrapping.
3025        untyped: UntypedMessage,
3026
3027        /// How to transform the logical message for its target peer.
3028        remote_style: crate::role::RemoteStyle,
3029
3030        /// Optional prerequisite that must finish before the request becomes
3031        /// visible on the transport.
3032        readiness: Option<RequestReadiness>,
3033    },
3034
3035    /// Send a notification to the server.
3036    Notification {
3037        /// the message to send; this may have a distinct method
3038        /// depending on the peer
3039        untyped: UntypedMessage,
3040    },
3041
3042    /// Send a response to a message from the server
3043    Response {
3044        id: RawJsonRpcId,
3045
3046        /// Method of the incoming request this response completes.
3047        method: String,
3048
3049        response: Result<serde_json::Value, crate::Error>,
3050
3051        destination: ResponseDestination,
3052    },
3053
3054    /// Send an Error Response that cannot be correlated to a request ID.
3055    UncorrelatedErrorResponse {
3056        error: crate::Error,
3057        destination: ResponseDestination,
3058    },
3059}
3060
3061impl OutgoingMessage {
3062    /// Kind, method, and ID only; params, results, and errors can carry
3063    /// credentials and must not appear in ordinary logs.
3064    fn log_summary(&self) -> (&'static str, Option<&str>, Option<&RawJsonRpcId>) {
3065        match self {
3066            Self::CloseAfterDraining { .. } => ("close_after_draining", None, None),
3067            Self::Request { id, method, .. } => ("request", Some(method), Some(id)),
3068            Self::Notification { untyped } => ("notification", Some(&untyped.method), None),
3069            Self::Response { id, method, .. } => ("response", Some(method), Some(id)),
3070            Self::UncorrelatedErrorResponse { .. } => ("uncorrelated_error_response", None, None),
3071        }
3072    }
3073}
3074
3075/// Return type from JrHandler; indicates whether the request was handled or not.
3076#[must_use]
3077#[derive(Debug)]
3078pub enum Handled<T> {
3079    /// The message was handled
3080    Yes,
3081
3082    /// The message was not handled; returns the original value.
3083    ///
3084    /// If `retry` is true,
3085    No {
3086        /// The message to be passed to subsequent handlers
3087        /// (typically the original message, but it may have been
3088        /// mutated.)
3089        message: T,
3090
3091        /// If true, request the message to be queued and retried with
3092        /// dynamic handlers as they are added.
3093        ///
3094        /// This is used for managing session updates since the dynamic
3095        /// handler for a session cannot be added until the response to the
3096        /// new session request has been processed and there may be updates
3097        /// that get processed at the same time.
3098        retry: bool,
3099    },
3100}
3101
3102/// Trait for converting handler return values into [`Handled`].
3103///
3104/// This trait allows handlers to return either `()` (which becomes `Handled::Yes`)
3105/// or an explicit `Handled<T>` value for more control over handler propagation.
3106pub trait IntoHandled<T> {
3107    /// Convert this value into a `Handled<T>`.
3108    fn into_handled(self) -> Handled<T>;
3109}
3110
3111impl<T> IntoHandled<T> for () {
3112    fn into_handled(self) -> Handled<T> {
3113        Handled::Yes
3114    }
3115}
3116
3117impl<T> IntoHandled<T> for Handled<T> {
3118    fn into_handled(self) -> Handled<T> {
3119        self
3120    }
3121}
3122
3123/// Connection context for sending messages and spawning tasks.
3124///
3125/// This is the primary handle for interacting with the JSON-RPC connection from
3126/// within handler callbacks. You can use it to:
3127///
3128/// * Send requests and notifications to the other side
3129/// * Spawn concurrent tasks that run alongside the connection
3130/// * Respond to requests (via [`Responder`] which wraps this)
3131///
3132/// # Cloning
3133///
3134/// `ConnectionTo` is cheaply cloneable - all clones refer to the same underlying connection.
3135/// This makes it easy to share across async tasks.
3136///
3137/// # Event Loop and Concurrency
3138///
3139/// Handler callbacks run on the event loop, which means the connection cannot process new
3140/// messages while your handler is running. Use [`spawn`](Self::spawn) to offload any
3141/// expensive or blocking work to concurrent tasks.
3142///
3143/// See the [Event Loop and Concurrency](Builder#event-loop-and-concurrency) section
3144/// for more details.
3145#[derive(Clone, Debug)]
3146pub struct ConnectionTo<Counterpart: Role> {
3147    counterpart: Counterpart,
3148    message_tx: OutgoingMessageTx,
3149    task_tx: TaskTx,
3150    dynamic_handler_tx: mpsc::UnboundedSender<DynamicHandlerMessage<Counterpart>>,
3151    transport_completion: SharedTransportCompletion,
3152    pending_replies: PendingRepliesRegistrar,
3153    protocol_mode: ProtocolMode,
3154    incoming_closed: IncomingClosed,
3155    protected_operations: Arc<Mutex<ProtectedOperations>>,
3156    runner_error_scope: Option<run::RunnerErrorScope>,
3157}
3158
3159type SharedTransportCompletion = future::Shared<BoxFuture<'static, Result<(), crate::Error>>>;
3160
3161type SharedCompletionSignal = future::Shared<BoxFuture<'static, ()>>;
3162
3163#[derive(Default)]
3164struct ProtectedOperations {
3165    pending: Vec<oneshot::Receiver<()>>,
3166    joining: Option<SharedCompletionSignal>,
3167}
3168
3169impl Debug for ProtectedOperations {
3170    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
3171        formatter
3172            .debug_struct("ProtectedOperations")
3173            .field("pending", &self.pending.len())
3174            .field("joining", &self.joining.is_some())
3175            .finish_non_exhaustive()
3176    }
3177}
3178
3179fn completion_signal() -> (oneshot::Sender<()>, SharedCompletionSignal) {
3180    let (tx, rx) = oneshot::channel();
3181    let signal = async move {
3182        // Dropping a sender (e.g. foreground failure) is not success.
3183        if rx.await.is_err() {
3184            future::pending::<()>().await;
3185        }
3186    }
3187    .boxed()
3188    .shared();
3189    (tx, signal)
3190}
3191
3192#[derive(Clone)]
3193struct IncomingClosed {
3194    state: Arc<IncomingClosedState>,
3195}
3196
3197struct IncomingClosedState {
3198    closing: AtomicBool,
3199    closed: AtomicBool,
3200    signal_tx: Mutex<Option<oneshot::Sender<()>>>,
3201    signal_rx: future::Shared<BoxFuture<'static, ()>>,
3202    shutdown_tx: Mutex<Option<oneshot::Sender<()>>>,
3203    #[cfg(any(feature = "unstable_mcp_over_acp", test))]
3204    shutdown_rx: SharedCompletionSignal,
3205    termination: Mutex<Option<ConnectionTermination>>,
3206    terminated_tx: Mutex<Option<oneshot::Sender<()>>>,
3207    terminated_rx: SharedCompletionSignal,
3208}
3209
3210impl IncomingClosed {
3211    fn new() -> Self {
3212        let (signal_tx, signal_rx) = oneshot::channel();
3213        let (shutdown_tx, shutdown_rx) = oneshot::channel();
3214        #[cfg(not(any(feature = "unstable_mcp_over_acp", test)))]
3215        drop(shutdown_rx);
3216        let (terminated_tx, terminated_rx) = oneshot::channel();
3217        Self {
3218            state: Arc::new(IncomingClosedState {
3219                closing: AtomicBool::new(false),
3220                closed: AtomicBool::new(false),
3221                signal_tx: Mutex::new(Some(signal_tx)),
3222                signal_rx: signal_rx.map(|_| ()).boxed().shared(),
3223                shutdown_tx: Mutex::new(Some(shutdown_tx)),
3224                #[cfg(any(feature = "unstable_mcp_over_acp", test))]
3225                shutdown_rx: shutdown_rx.map(|_| ()).boxed().shared(),
3226                termination: Mutex::new(None),
3227                terminated_tx: Mutex::new(Some(terminated_tx)),
3228                terminated_rx: terminated_rx.map(|_| ()).boxed().shared(),
3229            }),
3230        }
3231    }
3232
3233    fn begin_close(&self) {
3234        self.state.closing.store(true, Ordering::Release);
3235        self.request_shutdown();
3236    }
3237
3238    fn request_shutdown(&self) {
3239        if let Some(tx) = self
3240            .state
3241            .shutdown_tx
3242            .lock()
3243            .expect("shutdown signal mutex poisoned")
3244            .take()
3245        {
3246            let _ = tx.send(());
3247        }
3248    }
3249
3250    /// Whether the connection has begun shutting down (incoming EOF, a
3251    /// finished foreground, an actor error) or has already terminated.
3252    fn is_shutting_down(&self) -> bool {
3253        self.state
3254            .shutdown_tx
3255            .lock()
3256            .expect("shutdown signal mutex poisoned")
3257            .is_none()
3258            || self.termination().is_some()
3259    }
3260
3261    fn finish_close(&self) {
3262        self.state.closed.store(true, Ordering::Release);
3263        let signal_tx = self
3264            .state
3265            .signal_tx
3266            .lock()
3267            .expect("incoming-close signal mutex poisoned")
3268            .take();
3269
3270        if let Some(signal_tx) = signal_tx {
3271            let _ = signal_tx.send(());
3272        }
3273    }
3274
3275    async fn closed(&self) {
3276        self.state.signal_rx.clone().await;
3277    }
3278
3279    fn is_closed(&self) -> bool {
3280        self.state.closed.load(Ordering::Acquire)
3281    }
3282
3283    fn is_closing(&self) -> bool {
3284        self.state.closing.load(Ordering::Acquire)
3285    }
3286
3287    /// Record the termination reason. Returns `false` if one was already set.
3288    fn set_termination(&self, termination: ConnectionTermination) -> bool {
3289        let mut slot = self
3290            .state
3291            .termination
3292            .lock()
3293            .expect("termination mutex poisoned");
3294        if slot.is_some() {
3295            return false;
3296        }
3297        *slot = Some(termination);
3298        true
3299    }
3300
3301    fn signal_terminated(&self) {
3302        if let Some(tx) = self
3303            .state
3304            .terminated_tx
3305            .lock()
3306            .expect("terminated signal mutex poisoned")
3307            .take()
3308        {
3309            let _ = tx.send(());
3310        }
3311    }
3312
3313    fn termination(&self) -> Option<ConnectionTermination> {
3314        self.state
3315            .termination
3316            .lock()
3317            .expect("termination mutex poisoned")
3318            .clone()
3319    }
3320
3321    async fn terminated(&self) -> ConnectionTermination {
3322        self.state.terminated_rx.clone().await;
3323        self.termination()
3324            .expect("termination is recorded before it is signalled")
3325    }
3326}
3327
3328/// Publishes the connection's termination exactly once: when the driver
3329/// completes via [`Self::finish`], or synchronously when the driver future is
3330/// dropped first.
3331struct TerminationGuard {
3332    incoming_closed: IncomingClosed,
3333    pending_replies: PendingReplies,
3334}
3335
3336impl TerminationGuard {
3337    fn finish(self, termination: ConnectionTermination) {
3338        self.publish(termination);
3339    }
3340
3341    fn publish(&self, termination: ConnectionTermination) {
3342        if !self.incoming_closed.set_termination(termination.clone()) {
3343            return;
3344        }
3345        // Record first, so a requester woken by its failure observes the
3346        // termination; then fail requests before waking `terminated()`.
3347        let pending_reply_count = self.pending_replies.terminate(&termination);
3348        tracing::debug!(
3349            termination = termination.kind(),
3350            pending_reply_count,
3351            "Connection terminated"
3352        );
3353        self.incoming_closed.signal_terminated();
3354    }
3355}
3356
3357impl Drop for TerminationGuard {
3358    fn drop(&mut self) {
3359        // No-op after `finish`.
3360        self.publish(ConnectionTermination::Dropped);
3361    }
3362}
3363
3364impl Debug for IncomingClosed {
3365    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
3366        formatter
3367            .debug_struct("IncomingClosed")
3368            .field("is_closing", &self.is_closing())
3369            .field("is_closed", &self.is_closed())
3370            .finish_non_exhaustive()
3371    }
3372}
3373
3374/// Stable discriminator stored in the `data.reason` field of errors produced
3375/// when the incoming transport reaches clean EOF before a request receives its
3376/// response.
3377///
3378/// The `data` field is informational. Use [`is_incoming_transport_closed`] to
3379/// detect these errors: a peer can send this value in its own error data.
3380pub const INCOMING_TRANSPORT_CLOSED_REASON: &str = "incoming_transport_closed";
3381
3382/// Return whether `error` reports that the incoming transport reached clean
3383/// EOF before a request received its response.
3384///
3385/// Only errors the SDK produced for its own connection match. An error
3386/// response received from a peer never does, even when its `data.reason` is
3387/// [`INCOMING_TRANSPORT_CLOSED_REASON`], and neither does an error built from
3388/// the same code, message, and data.
3389#[must_use]
3390pub fn is_incoming_transport_closed(error: &crate::Error) -> bool {
3391    error.origin == crate::error::ErrorOrigin::IncomingTransportClosed
3392}
3393
3394fn incoming_transport_closed_error(method: &str) -> crate::Error {
3395    let mut error = crate::Error::internal_error();
3396    error.message = "Incoming transport closed".to_string();
3397    error.origin = crate::error::ErrorOrigin::IncomingTransportClosed;
3398    error.data(serde_json::json!({
3399        "reason": INCOMING_TRANSPORT_CLOSED_REASON,
3400        "method": method,
3401    }))
3402}
3403
3404/// Unlike the cleanup coordinator below, check success before polling delivery:
3405/// an already-ready success must not resume a message handler into another
3406/// dispatch.
3407fn run_incoming_until_foreground_succeeds(
3408    incoming: impl Future<Output = Result<(), crate::Error>>,
3409    foreground_succeeded: SharedCompletionSignal,
3410    incoming_closed: IncomingClosed,
3411) -> impl Future<Output = Result<(), crate::Error>> {
3412    let mut incoming = Box::pin(incoming);
3413    future::poll_fn(move |cx| {
3414        if foreground_succeeded.clone().poll_unpin(cx).is_ready() && !incoming_closed.is_closing() {
3415            return std::task::Poll::Ready(Ok(()));
3416        }
3417        // A close callback already underway is protected. Its result is polled
3418        // before stopping, so callback errors retain their existing precedence.
3419        incoming.as_mut().poll(cx)
3420    })
3421}
3422
3423/// Run the connection background alongside its foreground while ensuring that
3424/// a foreground woken by incoming EOF cannot cancel close callbacks midway.
3425fn run_until_connection_close<R>(
3426    background: impl Future<Output = Result<(), crate::Error>>,
3427    foreground: impl Future<Output = Result<R, crate::Error>>,
3428    incoming_closed: IncomingClosed,
3429) -> impl Future<Output = Result<R, crate::Error>> {
3430    // Box these before constructing the returned future. Keeping the generic
3431    // connection actors directly in this async state would substantially grow
3432    // every `connect_*` future.
3433    let background = Box::pin(background);
3434    let foreground = Box::pin(foreground);
3435
3436    async move {
3437        match future::select(background, foreground).await {
3438            Either::Left((background_result, foreground)) => {
3439                background_result?;
3440                foreground.await
3441            }
3442            Either::Right((foreground_result, background)) => {
3443                if !incoming_closed.is_closing() {
3444                    return foreground_result;
3445                }
3446
3447                match future::select(background, Box::pin(incoming_closed.closed())).await {
3448                    Either::Left((background_result, _)) => {
3449                        background_result?;
3450                        foreground_result
3451                    }
3452                    Either::Right(((), background)) => {
3453                        // Poll the background first once more so an error returned
3454                        // by the just-finished close callback wins over the ready
3455                        // foreground result.
3456                        crate::util::run_until(background, future::ready(foreground_result)).await
3457                    }
3458                }
3459            }
3460        }
3461    }
3462}
3463
3464impl<Counterpart: Role> ConnectionTo<Counterpart> {
3465    fn new(
3466        counterpart: Counterpart,
3467        message_tx: mpsc::UnboundedSender<OutgoingMessage>,
3468        task_tx: mpsc::UnboundedSender<Task>,
3469        dynamic_handler_tx: mpsc::UnboundedSender<DynamicHandlerMessage<Counterpart>>,
3470        transport_completion: SharedTransportCompletion,
3471        pending_replies: PendingRepliesRegistrar,
3472        protocol_mode: ProtocolMode,
3473    ) -> Self {
3474        Self {
3475            counterpart,
3476            message_tx,
3477            task_tx,
3478            dynamic_handler_tx,
3479            transport_completion,
3480            pending_replies,
3481            protocol_mode,
3482            incoming_closed: IncomingClosed::new(),
3483            protected_operations: Arc::default(),
3484            runner_error_scope: None,
3485        }
3486    }
3487
3488    pub(crate) fn with_runner_error_scope(mut self, scope: run::RunnerErrorScope) -> Self {
3489        self.runner_error_scope = Some(scope);
3490        self
3491    }
3492
3493    pub(crate) fn finish_runner_error(
3494        &self,
3495        error: crate::Error,
3496    ) -> impl Future<Output = ()> + Send + '_ {
3497        if let Some(scope) = &self.runner_error_scope {
3498            Either::Left(scope.finish(error))
3499        } else {
3500            // Unscoped runners belong to the connection itself.
3501            self.request_shutdown();
3502            Either::Right(self.wait_protected_operations())
3503        }
3504    }
3505
3506    /// Spawn only a connection-owned supervisor whose async cleanup must finish
3507    /// before the driver returns. Ordinary application work must use `spawn`.
3508    #[cfg(any(feature = "unstable_mcp_over_acp", test))]
3509    #[track_caller]
3510    pub(crate) fn spawn_protected(
3511        &self,
3512        task: impl IntoFuture<Output = Result<(), crate::Error>, IntoFuture: Send + 'static>,
3513    ) -> Result<(), crate::Error> {
3514        let mut state = self
3515            .protected_operations
3516            .lock()
3517            .expect("protected operations mutex poisoned");
3518        if state.joining.is_some() {
3519            return Err(crate::Error::request_cancelled());
3520        }
3521        // Reap completed acknowledgments at admission, rather than retaining
3522        // every operation for the entire lifetime of the connection.
3523        state
3524            .pending
3525            .retain_mut(|done| matches!(done.try_recv(), Ok(None)));
3526        let (done_tx, done_rx) = oneshot::channel();
3527        let task = task.into_future();
3528        self.spawn(async move {
3529            let result = task.await;
3530            let _ = done_tx.send(());
3531            result
3532        })?;
3533        state.pending.push(done_rx);
3534        Ok(())
3535    }
3536
3537    pub(crate) async fn wait_protected_operations(&self) {
3538        let joining = {
3539            let mut state = self
3540                .protected_operations
3541                .lock()
3542                .expect("protected operations mutex poisoned");
3543            if state.joining.is_none() {
3544                let operations = std::mem::take(&mut state.pending);
3545                state.joining = Some(
3546                    async move {
3547                        for operation in operations {
3548                            let _ = operation.await;
3549                        }
3550                    }
3551                    .boxed()
3552                    .shared(),
3553                );
3554            }
3555            state.joining.as_ref().expect("join initialized").clone()
3556        };
3557        joining.await;
3558    }
3559
3560    pub(crate) fn request_shutdown(&self) {
3561        self.incoming_closed.request_shutdown();
3562    }
3563
3564    /// Early cancellation for owned native work, before close callbacks or drain.
3565    #[cfg(any(feature = "unstable_mcp_over_acp", test))]
3566    pub(crate) async fn shutdown_requested(&self) {
3567        self.incoming_closed.state.shutdown_rx.clone().await;
3568    }
3569
3570    /// Return the counterpart role this connection is talking to.
3571    pub fn counterpart(&self) -> Counterpart {
3572        self.counterpart.clone()
3573    }
3574
3575    /// Wait until the incoming transport reaches clean EOF.
3576    ///
3577    /// Transport closure means that no more messages or responses can arrive.
3578    /// Pending requests are failed first; this completes after registered
3579    /// [`Builder::on_close`] callbacks finish.
3580    /// It does not automatically cancel the future passed to
3581    /// [`Builder::connect_with`]; use [`Builder::on_close`] when the connection
3582    /// should run application-specific cleanup or terminate that future.
3583    ///
3584    /// This never completes if the connection ends some other way; use
3585    /// [`terminated`](Self::terminated) to observe every termination.
3586    pub async fn incoming_closed(&self) {
3587        self.incoming_closed.closed().await;
3588    }
3589
3590    /// Return whether clean incoming-EOF processing has completed.
3591    ///
3592    /// This remains `false` while [`Builder::on_close`] callbacks are running.
3593    #[must_use]
3594    pub fn is_incoming_closed(&self) -> bool {
3595        self.incoming_closed.is_closed()
3596    }
3597
3598    /// Wait until the connection driver has terminated, for any reason.
3599    ///
3600    /// This is the universal terminal event for a connection driven by
3601    /// [`Builder::connect_to`] or [`Builder::connect_with`]. Unlike
3602    /// [`incoming_closed`](Self::incoming_closed), it also completes after
3603    /// transport errors, handler or close-callback errors, foreground
3604    /// completion, and a dropped driver future.
3605    ///
3606    /// The termination is published once, synchronously, as the driver future
3607    /// returns (or is dropped). At that point:
3608    ///
3609    /// - No further handler, close callback, runner, or
3610    ///   [`spawn`](Self::spawn)ed task of this connection will run; each has
3611    ///   completed or been dropped.
3612    /// - Every request still awaiting a response has been failed. Requests
3613    ///   failed earlier by clean incoming EOF satisfy
3614    ///   [`is_incoming_transport_closed`](crate::is_incoming_transport_closed);
3615    ///   the rest satisfy [`is_connection_terminated`](crate::is_connection_terminated),
3616    ///   as do requests sent afterwards.
3617    ///
3618    /// Do not await this from inside the connection's own handlers, close
3619    /// callbacks, or tasks: termination happens only after they finish.
3620    ///
3621    /// Termination does not reap child processes owned by the transport, nor
3622    /// drain application-owned queues.
3623    pub async fn terminated(&self) -> ConnectionTermination {
3624        self.incoming_closed.terminated().await
3625    }
3626
3627    /// Return the termination reason if the connection driver has terminated.
3628    ///
3629    /// See [`terminated`](Self::terminated).
3630    #[must_use]
3631    pub fn termination(&self) -> Option<ConnectionTermination> {
3632        self.incoming_closed.termination()
3633    }
3634
3635    /// Stop accepting outgoing messages, drain routable output through the
3636    /// protocol actor, and finish cooperative physical sinks. Reactive serving also
3637    /// joins owned transport work after incoming EOF.
3638    async fn drain_outgoing(
3639        &self,
3640        finish: Option<crate::component::FinishControl>,
3641        wait_owned_transport: bool,
3642    ) -> Result<(), crate::Error> {
3643        let (done_tx, done_rx) = oneshot::channel();
3644        let marker_result = send_raw_message(
3645            &self.message_tx,
3646            OutgoingMessage::CloseAfterDraining { done: done_tx },
3647        );
3648        let marker_result = match marker_result {
3649            Ok(()) => done_rx.await.map_err(|error| {
3650                crate::util::internal_error(format!(
3651                    "outgoing drain marker was dropped before completion: {error}"
3652                ))
3653            }),
3654            Err(error) => Err(error),
3655        };
3656
3657        let physical_finish = finish.is_some();
3658        if let Some(mut finish) = finish {
3659            // Only finish the physical sink after the protocol actor has handed
3660            // off its accepted output. Closing it earlier races the drain.
3661            finish.request();
3662        }
3663        if physical_finish || wait_owned_transport {
3664            // Cooperative completion proves physical sink drain. Reactive serving
3665            // also waits for owned work after EOF (e.g. child exit status).
3666            self.transport_completion.clone().await?;
3667        }
3668        // Opaque application drivers have no physical finish contract. Keep
3669        // polling their errors in the background, but do not globally join work
3670        // which may intentionally run forever after the foreground returns.
3671        marker_result
3672    }
3673
3674    fn is_incoming_closing(&self) -> bool {
3675        self.incoming_closed.is_closing()
3676    }
3677
3678    pub(super) fn begin_incoming_close(&self) {
3679        self.incoming_closed.begin_close();
3680    }
3681
3682    pub(super) fn finish_incoming_close(&self) {
3683        self.incoming_closed.finish_close();
3684    }
3685
3686    /// Spawns a task that will run so long as the JSON-RPC connection is being served.
3687    ///
3688    /// This is the primary mechanism for offloading expensive work from handler callbacks
3689    /// to avoid blocking the event loop. Spawned tasks run concurrently with the connection,
3690    /// allowing the server to continue processing messages.
3691    ///
3692    /// # Event Loop
3693    ///
3694    /// Handler callbacks run on the event loop, which cannot process new messages while
3695    /// your handler is running. Use `spawn` for any expensive operations:
3696    ///
3697    /// ```no_run
3698    /// # use agent_client_protocol_test::*;
3699    /// # async fn example() -> Result<(), agent_client_protocol::Error> {
3700    /// # let connection = mock_connection();
3701    /// connection.on_receive_request(async |req: ProcessRequest, responder, cx| {
3702    ///     // Clone cx for the spawned task
3703    ///     cx.spawn({
3704    ///         let connection = cx.clone();
3705    ///         async move {
3706    ///             let result = expensive_operation(&req.data).await?;
3707    ///             connection.send_notification(ProcessComplete { result })?;
3708    ///             Ok(())
3709    ///         }
3710    ///     })?;
3711    ///
3712    ///     // Respond immediately
3713    ///     responder.respond(ProcessResponse { result: "started".into() })
3714    /// }, agent_client_protocol::on_receive_request!())
3715    /// # .connect_to(agent_client_protocol_test::MockTransport).await?;
3716    /// # Ok(())
3717    /// # }
3718    /// ```
3719    ///
3720    /// # Errors
3721    ///
3722    /// If the spawned task returns an error, the entire server will shut down.
3723    #[track_caller]
3724    pub fn spawn(
3725        &self,
3726        task: impl IntoFuture<Output = Result<(), crate::Error>, IntoFuture: Send + 'static>,
3727    ) -> Result<(), crate::Error> {
3728        let location = std::panic::Location::caller();
3729        let task = task.into_future();
3730        Task::new(location, task).spawn(&self.task_tx)
3731    }
3732
3733    /// Spawn a JSON-RPC connection in the background and return the
3734    /// connection context its builder selects.
3735    ///
3736    /// This is useful for creating multiple connections that communicate with each other,
3737    /// such as implementing proxy patterns or connecting to multiple backend services.
3738    ///
3739    /// # Arguments
3740    ///
3741    /// - `builder`: The connection builder with handlers configured
3742    /// - `transport`: The transport component to connect to
3743    ///
3744    /// # Returns
3745    ///
3746    /// The connection type that `builder`'s callbacks receive: a raw
3747    /// [`ConnectionTo`] for a [`Raw`] builder such as [`Builder::new`], a
3748    /// [`v1::ConnectionTo`](crate::v1::ConnectionTo) for
3749    /// [`Client.v1()`](crate::Client::v1), and a `v2::ConnectionTo` for
3750    /// `Client.v2()`. Convert a typed connection back with `into_raw`.
3751    ///
3752    /// # Example: Proxying to a backend connection
3753    ///
3754    /// ```
3755    /// # use agent_client_protocol::UntypedRole;
3756    /// # use agent_client_protocol::{Builder, ConnectionTo};
3757    /// # use agent_client_protocol_test::*;
3758    /// # async fn example(cx: ConnectionTo<UntypedRole>) -> Result<(), agent_client_protocol::Error> {
3759    /// // Set up a backend connection builder
3760    /// let backend = UntypedRole.builder()
3761    ///     .on_receive_request(async |req: MyRequest, responder, _cx| {
3762    ///         responder.respond(MyResponse { status: "ok".into() })
3763    ///     }, agent_client_protocol::on_receive_request!());
3764    ///
3765    /// // Spawn it and get a context to send requests to it
3766    /// let backend_connection = cx.spawn_connection(backend, MockTransport)?;
3767    ///
3768    /// // Now you can forward requests to the backend
3769    /// let response = backend_connection.send_request(MyRequest {}).block_task().await?;
3770    /// # Ok(())
3771    /// # }
3772    /// ```
3773    ///
3774    /// # Example: Spawning a v1 client
3775    ///
3776    /// ```no_run
3777    /// # use agent_client_protocol::{
3778    /// #     Agent, Client, ConnectTo, ConnectionTo, Error, UntypedRole, v1,
3779    /// # };
3780    /// # fn example(
3781    /// #     connection: ConnectionTo<UntypedRole>,
3782    /// #     transport: impl ConnectTo<Client> + 'static,
3783    /// # ) -> Result<(), Error> {
3784    /// let child: v1::ConnectionTo<Agent> = connection.spawn_connection(Client.v1(), transport)?;
3785    /// # drop(child);
3786    /// # Ok(())
3787    /// # }
3788    /// ```
3789    #[track_caller]
3790    pub fn spawn_connection<R: Role, P: Protocol>(
3791        &self,
3792        builder: Builder<
3793            R,
3794            impl HandleDispatchFrom<R::Counterpart> + 'static,
3795            impl RunWithConnectionTo<R::Counterpart> + 'static,
3796            impl HandleConnectionClose<R::Counterpart> + 'static,
3797            P,
3798        >,
3799        transport: impl ConnectTo<R> + 'static,
3800    ) -> Result<P::Connection<R::Counterpart>, crate::Error> {
3801        let (connection, future) =
3802            builder.into_connection_and_future(transport, false, |_| std::future::pending());
3803        Task::new(std::panic::Location::caller(), future).spawn(&self.task_tx)?;
3804        Ok(protocol::connection_from_raw::<P, _>(connection))
3805    }
3806
3807    /// The ACP protocol version of this connection's wire guard, or `None`
3808    /// when its builder applies no guard, as with [`Builder::new`].
3809    ///
3810    /// `Client.v1()`, `Agent.v1()` and `Proxy.v1()` connections report
3811    /// [`ProtocolVersion::V1`](crate::schema::ProtocolVersion::V1), and the
3812    /// `v2()` entry points report `ProtocolVersion::V2`.
3813    #[must_use]
3814    pub fn protocol_version(&self) -> Option<crate::schema::ProtocolVersion> {
3815        self.protocol_mode.api_protocol_version()
3816    }
3817
3818    /// The same connection, typed for protocol `P`.
3819    ///
3820    /// `into_protocol::<V1>()` returns a [`v1::ConnectionTo`](crate::v1::ConnectionTo)
3821    /// (and, with `unstable_protocol_v2`, `into_protocol::<V2>()` a
3822    /// `v2::ConnectionTo`), whose sends return request handles typed for that
3823    /// version and which carries that version's session helpers. Only the
3824    /// Rust API changes: the connection's wire guard is the one its builder
3825    /// selected. The versioned connections convert back with `into_raw`.
3826    ///
3827    /// # Errors
3828    ///
3829    /// Returns an Internal error (`-32603`) when the builder's guard selected
3830    /// a different version than `P`, for example `into_protocol::<V2>()` on a
3831    /// connection from `Client.v1()`. The mismatch is a local programming
3832    /// error, not a malformed message from the peer. A
3833    /// connection whose builder applies no guard (see
3834    /// [`protocol_version`](Self::protocol_version)) converts to every
3835    /// protocol, and `into_protocol::<Raw>()` never fails.
3836    ///
3837    /// ```
3838    /// # use agent_client_protocol::{Agent, ConnectionTo, Error, V1, v1};
3839    /// fn sessions(connection: ConnectionTo<Agent>) -> Result<v1::ConnectionTo<Agent>, Error> {
3840    ///     connection.into_protocol::<V1>()
3841    /// }
3842    /// ```
3843    pub fn into_protocol<P: Protocol>(self) -> Result<P::Connection<Counterpart>, crate::Error> {
3844        if let (Some(guard), Some(requested)) =
3845            (self.protocol_version(), protocol::protocol_version::<P>())
3846            && guard != requested
3847        {
3848            return Err(crate::Error::internal_error().data(format!(
3849                "cannot use a connection guarded for ACP protocol version {guard} \
3850                 as a protocol version {requested} connection",
3851            )));
3852        }
3853        Ok(protocol::connection_from_raw::<P, _>(self))
3854    }
3855
3856    /// Send a request/notification and forward the response appropriately.
3857    ///
3858    /// The request context's response type matches the request's response type,
3859    /// enabling type-safe message forwarding.
3860    pub fn send_proxied_message<
3861        Req: JsonRpcRequest<Response: Send>,
3862        Notif: JsonRpcNotification,
3863        Q: Protocol,
3864    >(
3865        &self,
3866        message: Dispatch<Req, Notif, Q>,
3867    ) -> Result<(), crate::Error>
3868    where
3869        Counterpart: HasPeer<Counterpart>,
3870    {
3871        self.send_proxied_message_to(self.counterpart(), message)
3872    }
3873
3874    /// Send a request/notification and forward the response appropriately.
3875    ///
3876    /// The request context's response type matches the request's response type,
3877    /// enabling type-safe message forwarding.
3878    ///
3879    /// `$/cancel_request` notifications are *not* forwarded: their `requestId`
3880    /// refers to a request on the connection they arrived over and would be
3881    /// meaningless to `peer`. Cancellation instead propagates hop by hop,
3882    /// because the responders passed to
3883    /// [`forward_response_to`](SentRequest::forward_response_to) observe it
3884    /// and re-issue the cancellation with the forwarded request's own ID.
3885    pub fn send_proxied_message_to<
3886        Peer: Role,
3887        Req: JsonRpcRequest<Response: Send>,
3888        Notif: JsonRpcNotification,
3889        Q: Protocol,
3890    >(
3891        &self,
3892        peer: Peer,
3893        message: Dispatch<Req, Notif, Q>,
3894    ) -> Result<(), crate::Error>
3895    where
3896        Counterpart: HasPeer<Peer>,
3897    {
3898        match message {
3899            Dispatch::Request(request, responder) => self
3900                .send_ordered_request_to::<Raw, _, _>(peer, request)
3901                .forward_response_to(responder),
3902            Dispatch::Notification(notification) => {
3903                // `$/cancel_request` is connection-scoped: its `requestId` was
3904                // allocated on the connection the notification arrived over
3905                // and means nothing to `peer`. The cancellation has already
3906                // been recorded on this connection's responder markers, and
3907                // `forward_response_to` re-issues it for the forwarded request
3908                // with the correct per-hop ID, so drop the raw notification
3909                // instead of tunneling a meaningless ID across the hop.
3910                if is_cancel_request_notification(&notification) {
3911                    tracing::debug!(
3912                        "not forwarding hop-scoped `$/cancel_request` notification across proxy hop"
3913                    );
3914                    return Ok(());
3915                }
3916                self.send_notification_to(peer, notification)
3917            }
3918            Dispatch::Response(result, router) => {
3919                // Responses are forwarded directly to their destination
3920                router.route_with_result(result)
3921            }
3922        }
3923    }
3924
3925    /// Send an outgoing request and return a [`SentRequest`] for handling the reply.
3926    ///
3927    /// The returned [`SentRequest`] makes the response-consumption mode explicit:
3928    ///
3929    /// * [`on_receiving_result`](SentRequest::on_receiving_result) - Register a callback and
3930    ///   return immediately. If registered before the response is routed during its original
3931    ///   dispatch, the loop waits for the callback to complete.
3932    /// * [`block_task`](SentRequest::block_task) - Wait on the current task until the response
3933    ///   arrives. This is only safe when that task already runs outside the dispatch loop.
3934    ///
3935    /// For callback ordering selected before publication, use
3936    /// [`prepare_request`](Self::prepare_request) instead. Even an immediately
3937    /// chained callback can race with a fast response on a concurrent connection.
3938    ///
3939    /// # Anti-Footgun Design
3940    ///
3941    /// The API intentionally makes it difficult to block on the result directly to prevent
3942    /// the common mistake of blocking the event loop while waiting for a response:
3943    ///
3944    /// ```compile_fail
3945    /// # use agent_client_protocol_test::*;
3946    /// # async fn example(cx: agent_client_protocol::ConnectionTo<agent_client_protocol::UntypedRole>) -> Result<(), agent_client_protocol::Error> {
3947    /// // ❌ This doesn't compile - prevents blocking the event loop
3948    /// let response = cx.send_request(MyRequest {}).await?;
3949    /// # Ok(())
3950    /// # }
3951    /// ```
3952    ///
3953    /// ```no_run
3954    /// # use agent_client_protocol_test::*;
3955    /// # async fn example(cx: agent_client_protocol::ConnectionTo<agent_client_protocol::UntypedRole>) -> Result<(), agent_client_protocol::Error> {
3956    /// // ✅ Option 1: Register an ordered callback (safe in handlers)
3957    /// cx.send_request(MyRequest {})
3958    ///     .on_receiving_result(async |result| {
3959    ///         // Handle the response
3960    ///         Ok(())
3961    ///     })?;
3962    ///
3963    /// // ✅ Option 2: Block in spawned task (safe because task is concurrent)
3964    /// cx.spawn({
3965    ///     let cx = cx.clone();
3966    ///     async move {
3967    ///         let response = cx.send_request(MyRequest {})
3968    ///             .block_task()
3969    ///             .await?;
3970    ///         // Process response...
3971    ///         Ok(())
3972    ///     }
3973    /// })?;
3974    /// # Ok(())
3975    /// # }
3976    /// ```
3977    /// Send an outgoing request to the default counterpart peer.
3978    ///
3979    /// This is a convenience method that sends to the counterpart role `R`.
3980    /// For explicit control over the target peer, use [`send_request_to`](Self::send_request_to).
3981    pub fn send_request<Req: JsonRpcRequest>(&self, request: Req) -> SentRequest<Req::Response>
3982    where
3983        Counterpart: HasPeer<Counterpart>,
3984    {
3985        self.send_request_to(self.counterpart.clone(), request)
3986    }
3987
3988    /// Send an outgoing request to a specific peer.
3989    ///
3990    /// The message will be transformed according to the [`HasPeer`](crate::role::HasPeer)
3991    /// implementation before being sent.
3992    pub fn send_request_to<Peer: Role, Req: JsonRpcRequest>(
3993        &self,
3994        peer: Peer,
3995        request: Req,
3996    ) -> SentRequest<Req::Response>
3997    where
3998        Counterpart: HasPeer<Peer>,
3999    {
4000        self.send_request_to_with_options(peer, request, false, None, None)
4001    }
4002
4003    /// Prepare a request without sending it until response handling is selected.
4004    ///
4005    /// Unlike [`send_request`](Self::send_request), this does not register a
4006    /// pending reply or enqueue the request. A consuming method on the returned
4007    /// [`PreparedRequest`] publishes it synchronously. Callback-style methods
4008    /// select ordered consumption before publication, closing the race with a
4009    /// fast peer response even when the connection runs on another task.
4010    ///
4011    /// Dropping the prepared request sends nothing. See [`PreparedRequest`] for
4012    /// the available consumption modes and their error and cancellation behavior.
4013    ///
4014    /// # Examples
4015    ///
4016    /// ```no_run
4017    /// # use agent_client_protocol::{ConnectionTo, Error, UntypedRole};
4018    /// # use agent_client_protocol_test::MyRequest;
4019    /// # fn example(connection: ConnectionTo<UntypedRole>) -> Result<(), Error> {
4020    /// connection.prepare_request(MyRequest {}).on_receiving_result(async |result| {
4021    ///     let response = result?;
4022    ///     // Apply bounded response work before later inbound messages.
4023    ///     Ok(())
4024    /// })?;
4025    /// # Ok(())
4026    /// # }
4027    /// ```
4028    pub fn prepare_request<Req: JsonRpcRequest>(
4029        &self,
4030        request: Req,
4031    ) -> PreparedRequest<Req::Response>
4032    where
4033        Counterpart: HasPeer<Counterpart>,
4034    {
4035        self.prepare_request_to(self.counterpart.clone(), request)
4036    }
4037
4038    /// Prepare a request to a specific peer without sending it.
4039    ///
4040    /// The request is serialized now. A consuming method synchronously registers
4041    /// its pending reply and enqueues the request; peer transformation and
4042    /// transmission happen later in the connection driver.
4043    /// See [`prepare_request`](Self::prepare_request).
4044    pub fn prepare_request_to<Peer: Role, Req: JsonRpcRequest>(
4045        &self,
4046        peer: Peer,
4047        request: Req,
4048    ) -> PreparedRequest<Req::Response>
4049    where
4050        Counterpart: HasPeer<Peer>,
4051    {
4052        self.prepare_request_to_with_options(peer, request, None, None)
4053    }
4054
4055    /// Send a request and run a synchronous side effect when its valid success
4056    /// response is routed, after `before_send` completes and independently
4057    /// from how the returned request is eventually consumed.
4058    #[cfg(feature = "unstable_protocol_v2")]
4059    pub(crate) fn send_request_to_with_response_hook_after<
4060        P: Protocol,
4061        Peer: Role,
4062        Req: JsonRpcRequest,
4063        BeforeSend: Future<Output = Result<(), crate::Error>> + Send + 'static,
4064    >(
4065        &self,
4066        peer: Peer,
4067        request: Req,
4068        before_send: BeforeSend,
4069        response_hook: impl FnOnce(&Req::Response) -> Result<(), crate::Error> + Send + 'static,
4070    ) -> SentRequest<Req::Response, P>
4071    where
4072        Counterpart: HasPeer<Peer>,
4073    {
4074        let hook: ResponseRouteHook = Box::new(move |method, value| {
4075            let response = Req::Response::from_json(method, value.clone())?;
4076            response_hook(&response)
4077        });
4078        self.send_request_to_with_options(
4079            peer,
4080            request,
4081            false,
4082            Some(RequestReadiness::new(before_send)),
4083            Some(hook),
4084        )
4085    }
4086
4087    /// Send an ordered request with readiness and valid-success hooks.
4088    #[cfg(feature = "unstable_protocol_v2")]
4089    pub(crate) fn send_ordered_request_to_with_response_hook_after<
4090        P: Protocol,
4091        Peer: Role,
4092        Req: JsonRpcRequest,
4093        BeforeSend: Future<Output = Result<(), crate::Error>> + Send + 'static,
4094    >(
4095        &self,
4096        peer: Peer,
4097        request: Req,
4098        before_send: BeforeSend,
4099        response_hook: impl FnOnce(&Req::Response) -> Result<(), crate::Error> + Send + 'static,
4100    ) -> SentRequest<Req::Response, P>
4101    where
4102        Counterpart: HasPeer<Peer>,
4103    {
4104        let hook: ResponseRouteHook = Box::new(move |method, value| {
4105            let response = Req::Response::from_json(method, value.clone())?;
4106            response_hook(&response)
4107        });
4108        self.send_request_to_with_options(
4109            peer,
4110            request,
4111            true,
4112            Some(RequestReadiness::new(before_send)),
4113            Some(hook),
4114        )
4115    }
4116
4117    /// Send a request whose callback must run before later inbound messages.
4118    ///
4119    /// The ordering marker is installed before the request enters the outgoing
4120    /// queue, closing the race between a fast peer response and the immediate
4121    /// [`SentRequest::on_receiving_result`] call. Callers must consume the
4122    /// returned request with a callback-style method without yielding.
4123    pub(crate) fn send_ordered_request_to<P: Protocol, Peer: Role, Req: JsonRpcRequest>(
4124        &self,
4125        peer: Peer,
4126        request: Req,
4127    ) -> SentRequest<Req::Response, P>
4128    where
4129        Counterpart: HasPeer<Peer>,
4130    {
4131        self.send_request_to_with_options(peer, request, true, None, None)
4132    }
4133
4134    /// Send an ordered request after `before_send` completes successfully.
4135    ///
4136    /// The ordering marker and readiness prerequisite are both registered
4137    /// before the request enters the outgoing queue. This is used by framework
4138    /// setup paths that must acknowledge local routing before the peer can
4139    /// observe the request.
4140    pub(crate) fn send_ordered_request_to_after<
4141        P: Protocol,
4142        Peer: Role,
4143        Req: JsonRpcRequest,
4144        BeforeSend: Future<Output = Result<(), crate::Error>> + Send + 'static,
4145    >(
4146        &self,
4147        peer: Peer,
4148        request: Req,
4149        before_send: BeforeSend,
4150    ) -> SentRequest<Req::Response, P>
4151    where
4152        Counterpart: HasPeer<Peer>,
4153    {
4154        self.send_request_to_with_options(
4155            peer,
4156            request,
4157            true,
4158            Some(RequestReadiness::new(before_send)),
4159            None,
4160        )
4161    }
4162
4163    fn send_request_to_with_options<P: Protocol, Peer: Role, Req: JsonRpcRequest>(
4164        &self,
4165        peer: Peer,
4166        request: Req,
4167        ordered: bool,
4168        readiness: Option<RequestReadiness>,
4169        response_route_hook: Option<ResponseRouteHook>,
4170    ) -> SentRequest<Req::Response, P>
4171    where
4172        Counterpart: HasPeer<Peer>,
4173    {
4174        self.prepare_request_to_with_options(peer, request, readiness, response_route_hook)
4175            .into_sent_request(ordered)
4176    }
4177
4178    fn prepare_request_to_with_options<P: Protocol, Peer: Role, Req: JsonRpcRequest>(
4179        &self,
4180        peer: Peer,
4181        request: Req,
4182        readiness: Option<RequestReadiness>,
4183        response_route_hook: Option<ResponseRouteHook>,
4184    ) -> PreparedRequest<Req::Response, P>
4185    where
4186        Counterpart: HasPeer<Peer>,
4187    {
4188        let method = request.method().to_string();
4189        let id = HandleId::<P>::generated(Uuid::new_v4().to_string());
4190        let (response_tx, response_rx) = oneshot::channel();
4191        let response_ordering = ResponseOrdering::default();
4192        let role_id = peer.role_id();
4193        let remote_style = self.counterpart.remote_style(peer);
4194        let cancellation = SentRequestCancellation::new(self.message_tx.clone(), id.raw().clone());
4195        let pending_reply = PendingReply {
4196            id: id.raw().clone(),
4197            method: method.clone(),
4198            role_id,
4199            sender: response_tx,
4200            cancellation_disarm: cancellation.disarm_handle(),
4201            ordering: response_ordering.clone(),
4202            response_route_hook,
4203            accepts_lower_initialize_version: false,
4204        };
4205        let message = if self.is_incoming_closing() {
4206            Err(incoming_transport_closed_error(&method))
4207        } else {
4208            request
4209                .to_untyped_message()
4210                .map(|untyped| OutgoingMessage::Request {
4211                    id: id.raw().clone(),
4212                    method: method.clone(),
4213                    untyped,
4214                    remote_style,
4215                    readiness,
4216                })
4217                .map_err(|error| {
4218                    crate::util::internal_error(format!(
4219                        "failed to create untyped request for `{method}`: {error}"
4220                    ))
4221                })
4222        };
4223        let sent = SentRequest::new(
4224            id,
4225            method.clone(),
4226            self.task_tx.clone(),
4227            response_rx,
4228            cancellation,
4229            response_ordering,
4230        )
4231        .map(move |json| <Req::Response>::from_json(&method, json));
4232        PreparedRequest {
4233            sent,
4234            publication: RequestPublication {
4235                message,
4236                pending_reply,
4237                message_tx: self.message_tx.clone(),
4238                pending_replies: self.pending_replies.clone(),
4239                incoming_closed: self.incoming_closed.clone(),
4240            },
4241        }
4242    }
4243
4244    /// Send an outgoing notification to the default counterpart peer (no reply expected).
4245    ///
4246    /// Notifications are fire-and-forget messages that don't have IDs and don't expect responses.
4247    /// This method sends the notification immediately and returns.
4248    ///
4249    /// This is a convenience method that sends to the counterpart role `R`.
4250    /// For explicit control over the target peer, use [`send_notification_to`](Self::send_notification_to).
4251    ///
4252    /// ```no_run
4253    /// # use agent_client_protocol_test::*;
4254    /// # async fn example(cx: agent_client_protocol::ConnectionTo<agent_client_protocol::Agent>) -> Result<(), agent_client_protocol::Error> {
4255    /// cx.send_notification(StatusUpdate {
4256    ///     message: "Processing...".into(),
4257    /// })?;
4258    /// # Ok(())
4259    /// # }
4260    /// ```
4261    pub fn send_notification<N: JsonRpcNotification>(
4262        &self,
4263        notification: N,
4264    ) -> Result<(), crate::Error>
4265    where
4266        Counterpart: HasPeer<Counterpart>,
4267    {
4268        self.send_notification_to(self.counterpart.clone(), notification)
4269    }
4270
4271    /// Send an outgoing notification to a specific peer (no reply expected).
4272    ///
4273    /// The message will be transformed according to the [`HasPeer`](crate::role::HasPeer)
4274    /// implementation before being sent.
4275    pub fn send_notification_to<Peer: Role, N: JsonRpcNotification>(
4276        &self,
4277        peer: Peer,
4278        notification: N,
4279    ) -> Result<(), crate::Error>
4280    where
4281        Counterpart: HasPeer<Peer>,
4282    {
4283        let remote_style = self.counterpart.remote_style(peer);
4284        tracing::debug!(
4285            role = std::any::type_name::<Counterpart>(),
4286            peer = std::any::type_name::<Peer>(),
4287            notification_type = std::any::type_name::<N>(),
4288            ?remote_style,
4289            original_method = notification.method(),
4290            "send_notification_to"
4291        );
4292        let transformed = remote_style.transform_outgoing_message(notification)?;
4293        tracing::debug!(
4294            transformed_method = %transformed.method,
4295            "send_notification_to transformed"
4296        );
4297        send_raw_message(
4298            &self.message_tx,
4299            OutgoingMessage::Notification {
4300                untyped: transformed,
4301            },
4302        )
4303    }
4304
4305    /// Send a `$/cancel_request` notification for an arbitrary request ID to
4306    /// the default counterpart peer.
4307    ///
4308    /// Prefer [`SentRequest::cancel`] when you have the request handle: it
4309    /// already knows the correct request ID. Use this low-level method only
4310    /// when implementing custom routing with a request ID that is valid on
4311    /// this connection.
4312    pub fn send_cancel_request(
4313        &self,
4314        request_id: impl Into<RawJsonRpcId>,
4315    ) -> Result<(), crate::Error>
4316    where
4317        Counterpart: HasPeer<Counterpart>,
4318    {
4319        self.send_cancel_request_to(self.counterpart.clone(), request_id)
4320    }
4321
4322    /// Send a `$/cancel_request` notification for an arbitrary request ID to a
4323    /// specific peer.
4324    ///
4325    /// `$/cancel_request` is hop-local, so the notification is always sent
4326    /// plain, never wrapped in `proxy/successor`, even to a proxy's
4327    /// successor. To cancel a request sent to a successor, pass the ID of the
4328    /// `proxy/successor` request on this connection.
4329    ///
4330    /// Prefer [`SentRequest::cancel`] when you have the request handle: it
4331    /// already knows the correct request ID. Use this low-level method only
4332    /// when implementing custom routing with a request ID that is valid on
4333    /// this connection.
4334    pub fn send_cancel_request_to<Peer: Role>(
4335        &self,
4336        _peer: Peer,
4337        request_id: impl Into<RawJsonRpcId>,
4338    ) -> Result<(), crate::Error>
4339    where
4340        Counterpart: HasPeer<Peer>,
4341    {
4342        let untyped = crate::v1::schema::CancelRequestNotification::new(
4343            crate::v1::schema::RequestId::from(request_id.into()),
4344        )
4345        .to_untyped_message()?;
4346        send_raw_message(&self.message_tx, OutgoingMessage::Notification { untyped })
4347    }
4348
4349    /// Register a dynamic message handler, used to intercept messages specific to a particular session
4350    /// or some similar modal thing.
4351    ///
4352    /// Dynamic message handlers run after the handlers registered on [`Builder`] and before the
4353    /// role's default handler. They receive messages that the builder handlers decline.
4354    ///
4355    /// The handler will stay registered until the returned registration guard is dropped.
4356    pub fn add_dynamic_handler(
4357        &self,
4358        handler: impl HandleDispatchFrom<Counterpart> + 'static,
4359    ) -> Result<DynamicHandlerGuard<Counterpart>, crate::Error> {
4360        let uuid = Uuid::new_v4();
4361        let active = Arc::new(AtomicBool::new(true));
4362        self.dynamic_handler_tx
4363            .unbounded_send(DynamicHandlerMessage::AddDynamicHandler(
4364                uuid,
4365                Box::new(GuardedDynamicHandler {
4366                    active: active.clone(),
4367                    handler,
4368                }),
4369                active.clone(),
4370            ))
4371            .map_err(crate::util::internal_error)?;
4372
4373        Ok(DynamicHandlerGuard::new(uuid, active, self.clone()))
4374    }
4375
4376    /// Wait until every dynamic-handler update queued before this call has
4377    /// been applied by the incoming protocol actor.
4378    pub(crate) fn dynamic_handler_barrier(&self) -> BoxFuture<'static, Result<(), crate::Error>> {
4379        let (acknowledgment_tx, acknowledgment_rx) = oneshot::channel();
4380        if let Err(error) =
4381            self.dynamic_handler_tx
4382                .unbounded_send(DynamicHandlerMessage::AcknowledgedBarrier(
4383                    acknowledgment_tx,
4384                ))
4385        {
4386            return future::ready(Err(crate::Error::into_internal_error(error))).boxed();
4387        }
4388
4389        async move {
4390            acknowledgment_rx.await.map_err(|error| {
4391                crate::util::internal_error(format!(
4392                    "dynamic-handler barrier was dropped before acknowledgment: {error}"
4393                ))
4394            })
4395        }
4396        .boxed()
4397    }
4398
4399    fn remove_dynamic_handler(&self, uuid: Uuid) {
4400        // Ignore errors
4401        drop(
4402            self.dynamic_handler_tx
4403                .unbounded_send(DynamicHandlerMessage::RemoveDynamicHandler(uuid)),
4404        );
4405    }
4406}
4407
4408struct GuardedDynamicHandler<Handler> {
4409    active: Arc<AtomicBool>,
4410    handler: Handler,
4411}
4412
4413impl<Counterpart, Handler> HandleDispatchFrom<Counterpart> for GuardedDynamicHandler<Handler>
4414where
4415    Counterpart: Role,
4416    Handler: HandleDispatchFrom<Counterpart>,
4417{
4418    async fn handle_dispatch_from(
4419        &mut self,
4420        message: Dispatch,
4421        connection: ConnectionTo<Counterpart>,
4422    ) -> Result<Handled<Dispatch>, crate::Error> {
4423        if !self.active.load(Ordering::Acquire) {
4424            return Ok(Handled::No {
4425                message,
4426                retry: false,
4427            });
4428        }
4429        self.handler.handle_dispatch_from(message, connection).await
4430    }
4431
4432    fn describe_chain(&self) -> impl Debug {
4433        self.handler.describe_chain()
4434    }
4435}
4436
4437/// A guard that keeps a dynamic message handler registered.
4438///
4439/// Dropping the guard immediately deactivates the handler and queues its
4440/// removal from the connection. Use [`detach`](Self::detach) to keep the
4441/// handler registered for the remaining lifetime of the connection.
4442#[must_use = "dropping this guard unregisters the dynamic handler"]
4443#[derive(Debug)]
4444pub struct DynamicHandlerGuard<R: Role> {
4445    uuid: Option<Uuid>,
4446    active: Arc<AtomicBool>,
4447    cx: ConnectionTo<R>,
4448    cleanup: Option<Arc<dyn DynamicHandlerCleanup>>,
4449}
4450
4451/// Private registration-local cleanup, independent of connection task admission.
4452pub(crate) trait DynamicHandlerCleanup: std::fmt::Debug + Send + Sync {
4453    fn close(&self);
4454    fn wait(&self) -> futures::future::BoxFuture<'static, ()>;
4455}
4456
4457impl<R: Role> DynamicHandlerGuard<R> {
4458    fn new(uuid: Uuid, active: Arc<AtomicBool>, cx: ConnectionTo<R>) -> Self {
4459        Self {
4460            uuid: Some(uuid),
4461            active,
4462            cx,
4463            cleanup: None,
4464        }
4465    }
4466
4467    #[cfg(feature = "unstable_mcp_over_acp")]
4468    pub(crate) fn with_cleanup(mut self, cleanup: Arc<dyn DynamicHandlerCleanup>) -> Self {
4469        self.cleanup = Some(cleanup);
4470        self
4471    }
4472
4473    pub(crate) fn cleanup(&self) -> Option<Arc<dyn DynamicHandlerCleanup>> {
4474        self.cleanup.clone()
4475    }
4476
4477    /// Stop the handler from claiming messages without queuing its removal.
4478    ///
4479    /// Until the guard is dropped, the incoming actor keeps unclaimed session
4480    /// requests parked for a replacement handler.
4481    pub(crate) fn deactivate(&self) {
4482        self.active.store(false, Ordering::Release);
4483    }
4484
4485    /// Keep the dynamic handler registered after this guard is dropped.
4486    ///
4487    /// The handler remains registered until the connection itself shuts down.
4488    /// Unlike leaking the guard, detaching does not retain an extra
4489    /// [`ConnectionTo`] handle.
4490    pub fn detach(mut self) {
4491        self.uuid = None;
4492    }
4493}
4494
4495impl<R: Role> Drop for DynamicHandlerGuard<R> {
4496    fn drop(&mut self) {
4497        if let Some(uuid) = self.uuid {
4498            self.active.store(false, Ordering::Release);
4499            if let Some(cleanup) = &self.cleanup {
4500                cleanup.close();
4501            }
4502            self.cx.remove_dynamic_handler(uuid);
4503        }
4504    }
4505}
4506
4507/// The context to respond to an incoming request.
4508///
4509/// This context is provided to request handlers and serves a dual role:
4510///
4511/// 1. **Respond to the request** - Use [`respond`](Self::respond) or
4512///    [`respond_with_result`](Self::respond_with_result) to send the response
4513/// 2. **Send other messages** - Use the [`ConnectionTo`] parameter passed to your
4514///    handler, which provides [`send_request`](`ConnectionTo::send_request`),
4515///    [`send_notification`](`ConnectionTo::send_notification`), and
4516///    [`spawn`](`ConnectionTo::spawn`)
4517///
4518/// # Example
4519///
4520/// ```no_run
4521/// # use agent_client_protocol_test::*;
4522/// # async fn example() -> Result<(), agent_client_protocol::Error> {
4523/// # let connection = mock_connection();
4524/// connection.on_receive_request(async |req: ProcessRequest, responder, cx| {
4525///     // Send a notification while processing
4526///     cx.send_notification(StatusUpdate {
4527///         message: "processing".into(),
4528///     })?;
4529///
4530///     // Do some work...
4531///     let result = process(&req.data)?;
4532///
4533///     // Respond to the request
4534///     responder.respond(ProcessResponse { result })
4535/// }, agent_client_protocol::on_receive_request!())
4536/// # .connect_to(agent_client_protocol_test::MockTransport).await?;
4537/// # Ok(())
4538/// # }
4539/// ```
4540///
4541/// # Event Loop Considerations
4542///
4543/// Like all handlers, request handlers run on the event loop. Use
4544/// [`spawn`](ConnectionTo::spawn) for expensive operations to avoid blocking
4545/// the connection.
4546///
4547/// See the [Event Loop and Concurrency](Builder#event-loop-and-concurrency)
4548/// section for more details.
4549///
4550/// # Drop behavior
4551///
4552/// A request must be answered. If a responder for a request is dropped
4553/// without a response having been sent through it, the SDK answers the
4554/// request itself with an error, so the peer is never left waiting:
4555///
4556/// - `-32800` ([`Error::request_cancelled`]) if the peer had cancelled the
4557///   request with `$/cancel_request`;
4558/// - otherwise `-32603` ([`Error::internal_error`]), with data naming the
4559///   method whose responder was dropped.
4560///
4561/// No automatic response is sent when:
4562///
4563/// - a response was already sent through the responder (any of the
4564///   `respond*` methods), or the responder was handed to something that
4565///   answers through it (such as
4566///   [`SentRequest::forward_response_to`]) — only dropping the responder
4567///   itself without responding triggers the fallback;
4568/// - the responder was moved into another owner (another task, a wrapped
4569///   responder from [`cast`](Self::cast) or [`wrap_params`](Self::wrap_params),
4570///   a retried [`Dispatch`]); that owner's eventual drop is what counts;
4571/// - the connection is shutting down or has terminated.
4572///
4573/// When a handler drops the responder and then returns an error, that error
4574/// is the response instead of the generic fallback.
4575///
4576/// The automatic error is an ordinary error response: for example, a dropped
4577/// `initialize` responder fails initialization just as an explicit error
4578/// would.
4579///
4580/// # Protocol version
4581///
4582/// `P` is the [`Protocol`] marker and selects the type [`id`](Self::id)
4583/// returns. Typed handlers on `Client.v1()`, `Agent.v1()` and
4584/// `Proxy.v1()` receive [`V1`] responders and those on the v2 builders
4585/// receive `V2` responders. The version-neutral raw dispatch layer
4586/// ([`HandleDispatchFrom`], [`MatchDispatch`](crate::util::MatchDispatch))
4587/// uses the default, [`Raw`]. [`raw_id`](Self::raw_id) returns the wire ID
4588/// whatever the marker; [`into_protocol`](Responder::into_protocol) and
4589/// [`into_raw`](Self::into_raw) convert between markers.
4590///
4591/// [`Error::request_cancelled`]: crate::Error::request_cancelled
4592/// [`Error::internal_error`]: crate::Error::internal_error
4593#[must_use]
4594pub struct Responder<T: JsonRpcResponse = serde_json::Value, P: Protocol = Raw> {
4595    /// The method of the request.
4596    method: String,
4597
4598    /// The `id` of the message we are replying to.
4599    id: HandleId<P>,
4600
4601    /// Where the response is sent; it is completed at most once.
4602    destination: ResponseDestination,
4603
4604    /// Function to send the response to its destination.
4605    ///
4606    /// For incoming requests: serializes to JSON and sends over the wire.
4607    /// For incoming responses: sends to the waiting oneshot channel.
4608    send_fn: Box<dyn FnOnce(Result<T, crate::Error>) -> Result<(), crate::Error> + Send>,
4609
4610    /// Answers the request with an error if no response is sent. Owns the
4611    /// request-local cancellation state, so it can still observe a
4612    /// cancellation when it fires.
4613    drop_guard: ResponderDropGuard,
4614}
4615
4616/// Sends the automatic error for a [`Responder`] dropped without responding.
4617struct ResponderDropGuard {
4618    message_tx: OutgoingMessageTx,
4619    connection_state: IncomingClosed,
4620    /// The method of the request as received, used to route the response.
4621    wire_method: String,
4622    id: RawJsonRpcId,
4623    destination: ResponseDestination,
4624    /// Request-local cancellation state.
4625    cancellation: ResponderCancellation,
4626    armed: bool,
4627}
4628
4629impl ResponderDropGuard {
4630    fn disarm(&mut self) {
4631        self.armed = false;
4632    }
4633}
4634
4635impl Drop for ResponderDropGuard {
4636    fn drop(&mut self) {
4637        if !self.armed || self.destination.is_completed() {
4638            return;
4639        }
4640        if self.connection_state.is_shutting_down() {
4641            tracing::debug!(
4642                id = ?self.id,
4643                method = %self.wire_method,
4644                "responder dropped while the connection is shutting down; not answering"
4645            );
4646            return;
4647        }
4648
4649        let error = if self.cancellation.is_cancelled() {
4650            crate::Error::request_cancelled()
4651        } else {
4652            crate::Error::internal_error().data(format!(
4653                "request handler dropped its responder for `{}` without responding",
4654                self.wire_method
4655            ))
4656        };
4657        tracing::warn!(
4658            id = ?self.id,
4659            method = %self.wire_method,
4660            code = ?error.code,
4661            "Responder dropped without responding; answering with an error"
4662        );
4663        self.destination.abandon(AbandonedResponse {
4664            message_tx: self.message_tx.clone(),
4665            connection_state: self.connection_state.clone(),
4666            message: OutgoingMessage::Response {
4667                id: self.id.clone(),
4668                method: self.wire_method.clone(),
4669                response: Err(error),
4670                destination: self.destination.clone(),
4671            },
4672        });
4673    }
4674}
4675
4676impl<T: JsonRpcResponse, P: Protocol> std::fmt::Debug for Responder<T, P> {
4677    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
4678        f.debug_struct("Responder")
4679            .field("method", &self.method)
4680            .field("id", &self.id)
4681            .field("response_type", &std::any::type_name::<T>())
4682            .finish_non_exhaustive()
4683    }
4684}
4685
4686impl Responder<serde_json::Value, Raw> {
4687    /// Create a new request context for an incoming request.
4688    ///
4689    /// The response will be serialized to JSON and sent over the wire.
4690    fn new(
4691        message_tx: OutgoingMessageTx,
4692        connection_state: IncomingClosed,
4693        method: String,
4694        id: RawJsonRpcId,
4695        cancellation_registry: &RequestCancellationRegistry,
4696        destination: ResponseDestination,
4697    ) -> Self {
4698        let id_clone = id.clone();
4699        let method_clone = method.clone();
4700        let cancellation = cancellation_registry.register(&id);
4701        let send_destination = destination.clone();
4702        let drop_guard = ResponderDropGuard {
4703            message_tx: message_tx.clone(),
4704            connection_state,
4705            wire_method: method.clone(),
4706            id: id.clone(),
4707            destination: destination.clone(),
4708            cancellation,
4709            armed: true,
4710        };
4711        Self {
4712            method,
4713            id: HandleId::raw_layer(id),
4714            destination,
4715            send_fn: Box::new(move |response: Result<serde_json::Value, crate::Error>| {
4716                send_raw_message(
4717                    &message_tx,
4718                    OutgoingMessage::Response {
4719                        id: id_clone,
4720                        method: method_clone,
4721                        response,
4722                        destination: send_destination,
4723                    },
4724                )
4725            }),
4726            drop_guard,
4727        }
4728    }
4729}
4730
4731impl<P: Protocol> Responder<serde_json::Value, P> {
4732    /// Cast this request context to a different response type.
4733    ///
4734    /// The provided type `T` will be serialized to JSON before sending.
4735    pub fn cast<T: JsonRpcResponse>(self) -> Responder<T, P> {
4736        self.wrap_params(move |method, value| match value {
4737            Ok(value) => T::into_json(value, method),
4738            Err(e) => Err(e),
4739        })
4740    }
4741}
4742
4743impl<T: JsonRpcResponse, P: Protocol> Responder<T, P> {
4744    /// Method of the incoming request
4745    #[must_use]
4746    pub fn method(&self) -> &str {
4747        &self.method
4748    }
4749
4750    /// ID of the incoming request, as this responder's protocol version
4751    /// types it.
4752    ///
4753    /// On a [`V2`](crate::Protocol) responder this is never `null`: the SDK
4754    /// rejects v2 requests with a `null` ID before dispatch.
4755    #[must_use]
4756    pub fn id(&self) -> &P::RequestId {
4757        self.id.typed()
4758    }
4759
4760    /// ID of the incoming request exactly as it appeared on the wire.
4761    ///
4762    /// This is the same ID as [`id`](Self::id), in the version-neutral form
4763    /// used by routers, forwarders and tracing.
4764    #[must_use]
4765    pub fn raw_id(&self) -> &RawJsonRpcId {
4766        self.id.raw()
4767    }
4768
4769    /// Returns the cancellation marker for this request.
4770    ///
4771    /// The marker is set when the peer sends `$/cancel_request` for this
4772    /// request's JSON-RPC ID. Cancellation is cooperative: handlers should use
4773    /// the marker to stop long-running work and then respond. Unless the
4774    /// method defines a successful response for cancellation (for example a
4775    /// permission request's `cancelled` outcome), that response is
4776    /// [`Error::request_cancelled`] (`-32800`) rather than partial data.
4777    /// Dropping the responder of a cancelled request sends that error
4778    /// automatically.
4779    ///
4780    /// [`Error::request_cancelled`]: crate::Error::request_cancelled
4781    #[must_use]
4782    pub fn cancellation(&self) -> RequestCancellation {
4783        self.drop_guard.cancellation.cancellation()
4784    }
4785
4786    /// Convert to a `Responder` that expects a JSON value
4787    /// and which checks (dynamically) that the JSON value it receives
4788    /// can be converted to `T`.
4789    pub fn erase_to_json(self) -> Responder<serde_json::Value, P> {
4790        self.wrap_params(|method, value| T::from_json(method, value?))
4791    }
4792
4793    /// Return a new Responder with a different method name.
4794    pub fn wrap_method(self, method: String) -> Responder<T, P> {
4795        Responder {
4796            method,
4797            id: self.id,
4798            destination: self.destination,
4799            send_fn: self.send_fn,
4800            drop_guard: self.drop_guard,
4801        }
4802    }
4803
4804    /// Return a new Responder that expects a response of type U.
4805    ///
4806    /// `wrap_fn` will be invoked with the method name and the result to transform
4807    /// type `U` into type `T` before sending.
4808    pub fn wrap_params<U: JsonRpcResponse>(
4809        self,
4810        wrap_fn: impl FnOnce(&str, Result<U, crate::Error>) -> Result<T, crate::Error> + Send + 'static,
4811    ) -> Responder<U, P> {
4812        let method = self.method.clone();
4813        let send_fn = self.send_fn;
4814        Responder {
4815            method: self.method,
4816            id: self.id,
4817            destination: self.destination,
4818            send_fn: Box::new(move |input: Result<U, crate::Error>| {
4819                let t_value = wrap_fn(&method, input);
4820                send_fn(t_value)
4821            }),
4822            drop_guard: self.drop_guard,
4823        }
4824    }
4825
4826    /// Respond to the JSON-RPC request with either a value (`Ok`) or an error (`Err`).
4827    pub fn respond_with_result(
4828        mut self,
4829        response: Result<T, crate::Error>,
4830    ) -> Result<(), crate::Error> {
4831        tracing::debug!(id = ?self.id, "respond called");
4832        self.drop_guard.disarm();
4833        (self.send_fn)(response)
4834    }
4835
4836    /// Respond to the JSON-RPC request with a value.
4837    pub fn respond(self, response: T) -> Result<(), crate::Error> {
4838        self.respond_with_result(Ok(response))
4839    }
4840
4841    /// Respond to the JSON-RPC request with an internal error containing a message.
4842    pub fn respond_with_internal_error(self, message: impl ToString) -> Result<(), crate::Error> {
4843        self.respond_with_error(crate::util::internal_error(message))
4844    }
4845
4846    /// Respond to the JSON-RPC request with an error.
4847    pub fn respond_with_error(self, error: crate::Error) -> Result<(), crate::Error> {
4848        tracing::debug!(id = ?self.id, error = ?crate::util::redacted_error(error.code), "respond_with_error called");
4849        self.respond_with_result(Err(error))
4850    }
4851
4852    /// The same responder for protocol `Q`. Fails only if the request ID is
4853    /// `null` and `Q` does not allow it, which cannot happen for requests
4854    /// received on a connection that speaks `Q`.
4855    pub(crate) fn retag<Q: Protocol>(self) -> Result<Responder<T, Q>, crate::Error> {
4856        Ok(Responder {
4857            method: self.method,
4858            id: self.id.into_protocol()?,
4859            destination: self.destination,
4860            send_fn: self.send_fn,
4861            drop_guard: self.drop_guard,
4862        })
4863    }
4864
4865    /// The same responder with a version-neutral [`Raw`] request ID, for
4866    /// code that is not tied to one protocol version, such as a
4867    /// [`HandleDispatchFrom`] implementation. This never fails.
4868    pub fn into_raw(self) -> Responder<T, Raw> {
4869        Responder {
4870            method: self.method,
4871            id: self.id.into_raw(),
4872            destination: self.destination,
4873            send_fn: self.send_fn,
4874            drop_guard: self.drop_guard,
4875        }
4876    }
4877
4878    fn reply_target(&self) -> RequestReplyTarget {
4879        RequestReplyTarget {
4880            id: self.id.raw().clone(),
4881            method: self.method.clone(),
4882            destination: self.destination.clone(),
4883        }
4884    }
4885}
4886
4887impl<T: JsonRpcResponse> Responder<T, Raw> {
4888    /// Convert this version-neutral responder to protocol `Q`'s responder,
4889    /// for example to pass a responder from a [`HandleDispatchFrom`] or
4890    /// [`MatchDispatch`](crate::util::MatchDispatch) handler to a v2 session
4891    /// helper.
4892    ///
4893    /// Only the request ID's type changes. It fails only when `Q` does not
4894    /// allow the ID: protocol v2 rejects `null` with the Invalid Request
4895    /// error (`-32600`) that a v2 connection sends for a `null` request ID.
4896    /// That cannot happen for a request received on a v2 connection, which
4897    /// rejects such requests before dispatch. On failure the responder is
4898    /// dropped, which answers the request with an error.
4899    ///
4900    /// ```
4901    /// # use agent_client_protocol::{Error, Responder, V1, v1::schema};
4902    /// # fn example(responder: Responder<schema::PromptResponse>) -> Result<(), Error> {
4903    /// let responder: agent_client_protocol::v1::Responder<schema::PromptResponse> =
4904    ///     responder.into_protocol::<V1>()?;
4905    /// let id: &schema::RequestId = responder.id();
4906    /// # let _ = id;
4907    /// # Ok(())
4908    /// # }
4909    /// ```
4910    pub fn into_protocol<Q: Protocol>(self) -> Result<Responder<T, Q>, crate::Error> {
4911        self.retag()
4912    }
4913}
4914
4915/// Context for handling an incoming JSON-RPC response.
4916///
4917/// This is the response-side counterpart to [`Responder`]. While `Responder` handles
4918/// incoming requests (where you send a response over the wire), `ResponseRouter` handles
4919/// incoming responses (where you route the response to a local task waiting for it).
4920///
4921/// Both are fundamentally "sinks" that push the message through a `send_fn`, but they
4922/// represent different points in the message lifecycle and carry different metadata.
4923///
4924/// # Drop Behavior
4925///
4926/// Dropping a `ResponseRouter` without routing the response (for example, from a
4927/// dispatch handler that claims a [`Dispatch::Response`]) discards the
4928/// response: the local awaiter observes the response as never received. The
4929/// request still counts as settled: routing a response this far disarms the
4930/// originating [`SentRequest`]'s drop-time auto-cancellation even if the router
4931/// is never invoked, since the peer has already answered.
4932///
4933/// `P` is the [`Protocol`] marker that types [`id`](Self::id); see
4934/// [`Responder`].
4935#[must_use]
4936pub struct ResponseRouter<T: JsonRpcResponse = serde_json::Value, P: Protocol = Raw> {
4937    /// The method of the original request.
4938    method: String,
4939
4940    /// The `id` of the original request.
4941    id: HandleId<P>,
4942
4943    /// The RoleId to which the original request was sent
4944    /// (and hence from which the reply is expected).
4945    role_id: RoleId,
4946
4947    /// Function to send the response to the waiting task.
4948    send_fn: Box<dyn FnOnce(Result<T, crate::Error>) -> Result<(), crate::Error> + Send>,
4949
4950    /// Shared route used to deliver a dispatch-handler error to the same waiter.
4951    reply_target: ResponseReplyTarget,
4952}
4953
4954impl<T: JsonRpcResponse, P: Protocol> std::fmt::Debug for ResponseRouter<T, P> {
4955    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
4956        f.debug_struct("ResponseRouter")
4957            .field("method", &self.method)
4958            .field("id", &self.id)
4959            .field("response_type", &std::any::type_name::<T>())
4960            .finish_non_exhaustive()
4961    }
4962}
4963
4964impl ResponseRouter<serde_json::Value> {
4965    /// Create a new response context for routing a response to a local awaiter.
4966    ///
4967    /// When [`route_with_result`](Self::route_with_result) is called, the response is sent through the oneshot
4968    /// channel to the code that originally sent the request. If that receiver was
4969    /// dropped, the response is discarded because there is no local awaiter left.
4970    fn new(pending_reply: PendingReply, dispatch: ResponseDispatch) -> Self {
4971        let PendingReply {
4972            id,
4973            method,
4974            role_id,
4975            sender,
4976            cancellation_disarm,
4977            ordering,
4978            response_route_hook,
4979            accepts_lower_initialize_version: _,
4980        } = pending_reply;
4981        let reply_target = ResponseReplyTarget {
4982            id: id.clone(),
4983            method: method.clone(),
4984            sender: Arc::new(Mutex::new(Some(sender))),
4985            ordering,
4986            dispatch,
4987        };
4988        let send_target = reply_target.clone();
4989        // A response for the request reached this router, so the request is
4990        // settled from the peer's perspective and a `$/cancel_request` could
4991        // only ever be redundant. Disarm immediately so handlers may retain
4992        // the router without leaving auto-cancellation armed.
4993        cancellation_disarm.disarm();
4994        let hook_method = method.clone();
4995        Self {
4996            method,
4997            id: HandleId::raw_layer(id),
4998            role_id,
4999            send_fn: Box::new(move |response: Result<serde_json::Value, crate::Error>| {
5000                let response = match response {
5001                    Ok(value) => match response_route_hook {
5002                        Some(hook) => hook(&hook_method, &value).map(|()| value),
5003                        None => Ok(value),
5004                    },
5005                    Err(error) => Err(error),
5006                };
5007                send_target.route(response);
5008                Ok(())
5009            }),
5010            reply_target,
5011        }
5012    }
5013}
5014
5015impl<P: Protocol> ResponseRouter<serde_json::Value, P> {
5016    /// Cast this response context to a different response type.
5017    ///
5018    /// The provided type `T` will be serialized to JSON before sending.
5019    pub fn cast<T: JsonRpcResponse>(self) -> ResponseRouter<T, P> {
5020        self.wrap_params(move |method, value| match value {
5021            Ok(value) => T::into_json(value, method),
5022            Err(e) => Err(e),
5023        })
5024    }
5025}
5026
5027impl<T: JsonRpcResponse, P: Protocol> ResponseRouter<T, P> {
5028    /// Method of the original request
5029    #[must_use]
5030    pub fn method(&self) -> &str {
5031        &self.method
5032    }
5033
5034    /// ID of the original request, as this router's protocol version types
5035    /// it.
5036    #[must_use]
5037    pub fn id(&self) -> &P::RequestId {
5038        self.id.typed()
5039    }
5040
5041    /// ID of the original request exactly as it was sent on the wire.
5042    ///
5043    /// This is the same ID as [`id`](Self::id), in the version-neutral form
5044    /// used by routers, forwarders and tracing.
5045    #[must_use]
5046    pub fn raw_id(&self) -> &RawJsonRpcId {
5047        self.id.raw()
5048    }
5049
5050    /// The peer to which the original request was sent.
5051    ///
5052    /// This is the peer from which we expect to receive the response.
5053    #[must_use]
5054    pub fn role_id(&self) -> RoleId {
5055        self.role_id.clone()
5056    }
5057
5058    /// Convert to a `ResponseRouter` that expects a JSON value
5059    /// and which checks (dynamically) that the JSON value it receives
5060    /// can be converted to `T`.
5061    pub fn erase_to_json(self) -> ResponseRouter<serde_json::Value, P> {
5062        self.wrap_params(|method, value| T::from_json(method, value?))
5063    }
5064
5065    /// The same router for protocol `Q`. Fails only if the request ID is not
5066    /// allowed by `Q`, which cannot happen because the SDK never assigns
5067    /// `null` to requests it sends.
5068    pub(crate) fn retag<Q: Protocol>(self) -> Result<ResponseRouter<T, Q>, crate::Error> {
5069        Ok(ResponseRouter {
5070            method: self.method,
5071            id: self.id.into_protocol()?,
5072            role_id: self.role_id,
5073            send_fn: self.send_fn,
5074            reply_target: self.reply_target,
5075        })
5076    }
5077
5078    /// The same router with a version-neutral [`Raw`] request ID. This never
5079    /// fails.
5080    pub fn into_raw(self) -> ResponseRouter<T, Raw> {
5081        ResponseRouter {
5082            method: self.method,
5083            id: self.id.into_raw(),
5084            role_id: self.role_id,
5085            send_fn: self.send_fn,
5086            reply_target: self.reply_target,
5087        }
5088    }
5089
5090    /// Return a new ResponseRouter that expects a response of type U.
5091    ///
5092    /// `wrap_fn` will be invoked with the method name and the result to transform
5093    /// type `U` into type `T` before sending.
5094    fn wrap_params<U: JsonRpcResponse>(
5095        self,
5096        wrap_fn: impl FnOnce(&str, Result<U, crate::Error>) -> Result<T, crate::Error> + Send + 'static,
5097    ) -> ResponseRouter<U, P> {
5098        let method = self.method.clone();
5099        ResponseRouter {
5100            method: self.method,
5101            id: self.id,
5102            role_id: self.role_id,
5103            send_fn: Box::new(move |input: Result<U, crate::Error>| {
5104                let t_value = wrap_fn(&method, input);
5105                (self.send_fn)(t_value)
5106            }),
5107            reply_target: self.reply_target,
5108        }
5109    }
5110
5111    /// Route the response result to the waiting task.
5112    pub fn route_with_result(self, response: Result<T, crate::Error>) -> Result<(), crate::Error> {
5113        tracing::debug!(id = ?self.id, "response routed to awaiter");
5114        (self.send_fn)(response)
5115    }
5116
5117    /// Route a successful response value to the waiting task.
5118    pub fn route(self, response: T) -> Result<(), crate::Error> {
5119        self.route_with_result(Ok(response))
5120    }
5121
5122    /// Route an internal error to the waiting task.
5123    pub fn route_with_internal_error(self, message: impl ToString) -> Result<(), crate::Error> {
5124        self.route_with_error(crate::util::internal_error(message))
5125    }
5126
5127    /// Route an error response to the waiting task.
5128    pub fn route_with_error(self, error: crate::Error) -> Result<(), crate::Error> {
5129        tracing::debug!(id = ?self.id, error = ?crate::util::redacted_error(error.code), "error routed to awaiter");
5130        self.route_with_result(Err(error))
5131    }
5132}
5133
5134impl<T: JsonRpcResponse> ResponseRouter<T, Raw> {
5135    /// Convert this version-neutral router to protocol `Q`'s router.
5136    ///
5137    /// Only the request ID's type changes. The SDK never assigns `null` to
5138    /// requests it sends, so this succeeds for every marker.
5139    pub fn into_protocol<Q: Protocol>(self) -> Result<ResponseRouter<T, Q>, crate::Error> {
5140        self.retag()
5141    }
5142}
5143
5144/// Common bounds for any JSON-RPC message.
5145///
5146/// # Derive Macro
5147///
5148/// For simple message types, you can use the `JsonRpcRequest` or `JsonRpcNotification` derive macros
5149/// which will implement both `JsonRpcMessage` and the respective trait. See [`JsonRpcRequest`] and
5150/// [`JsonRpcNotification`] for examples.
5151pub trait JsonRpcMessage: 'static + Debug + Sized + Send + Clone {
5152    /// Check if this message type matches the given method name.
5153    fn matches_method(method: &str) -> bool;
5154
5155    /// The method name for the message.
5156    fn method(&self) -> &str;
5157
5158    /// Convert this message into an untyped message.
5159    fn to_untyped_message(&self) -> Result<UntypedMessage, crate::Error>;
5160
5161    /// Parse this type from a method name and the message's JSON `params`.
5162    ///
5163    /// Returns an error if the method doesn't match or deserialization fails.
5164    /// Callers should use `matches_method` first to check if this type handles the method.
5165    /// A message without `params` is parsed from [`serde_json::Value::Null`].
5166    fn parse_message(method: &str, params: &serde_json::Value) -> Result<Self, crate::Error>;
5167}
5168
5169/// Defines the "payload" of a successful response to a JSON-RPC request.
5170///
5171/// # Derive Macro
5172///
5173/// Use `#[derive(JsonRpcResponse)]` to automatically implement this trait:
5174///
5175/// ```ignore
5176/// use agent_client_protocol::JsonRpcResponse;
5177/// use serde::{Serialize, Deserialize};
5178///
5179/// #[derive(Debug, Clone, Serialize, Deserialize, JsonRpcResponse)]
5180/// struct HelloResponse {
5181///     greeting: String,
5182/// }
5183/// ```
5184pub trait JsonRpcResponse: 'static + Debug + Sized + Send + Clone {
5185    /// Convert this message into a JSON value.
5186    fn into_json(self, method: &str) -> Result<serde_json::Value, crate::Error>;
5187
5188    /// Parse a JSON value into the response type.
5189    fn from_json(method: &str, value: serde_json::Value) -> Result<Self, crate::Error>;
5190}
5191
5192impl JsonRpcResponse for serde_json::Value {
5193    fn from_json(_method: &str, value: serde_json::Value) -> Result<Self, crate::Error> {
5194        Ok(value)
5195    }
5196
5197    fn into_json(self, _method: &str) -> Result<serde_json::Value, crate::Error> {
5198        Ok(self)
5199    }
5200}
5201
5202/// A struct that represents a notification (JSON-RPC message that does not expect a response).
5203///
5204/// # Derive Macro
5205///
5206/// Use `#[derive(JsonRpcNotification)]` to automatically implement both `JsonRpcMessage` and `JsonRpcNotification`:
5207///
5208/// ```ignore
5209/// use agent_client_protocol::JsonRpcNotification;
5210/// use serde::{Serialize, Deserialize};
5211///
5212/// #[derive(Debug, Clone, Serialize, Deserialize, JsonRpcNotification)]
5213/// #[notification(method = "_ping")]
5214/// struct PingNotification {
5215///     timestamp: u64,
5216/// }
5217/// ```
5218pub trait JsonRpcNotification: JsonRpcMessage {}
5219
5220/// A struct that represents a request (JSON-RPC message expecting a response).
5221///
5222/// # Derive Macro
5223///
5224/// Use `#[derive(JsonRpcRequest)]` to automatically implement both `JsonRpcMessage` and `JsonRpcRequest`:
5225///
5226/// ```ignore
5227/// use agent_client_protocol::{JsonRpcRequest, JsonRpcResponse};
5228/// use serde::{Serialize, Deserialize};
5229///
5230/// #[derive(Debug, Clone, Serialize, Deserialize, JsonRpcRequest)]
5231/// #[request(method = "_hello", response = HelloResponse)]
5232/// struct HelloRequest {
5233///     name: String,
5234/// }
5235///
5236/// #[derive(Debug, Clone, Serialize, Deserialize, JsonRpcResponse)]
5237/// struct HelloResponse {
5238///     greeting: String,
5239/// }
5240/// ```
5241pub trait JsonRpcRequest: JsonRpcMessage {
5242    /// The type of data expected in response.
5243    type Response: JsonRpcResponse;
5244}
5245
5246/// An incoming request, notification, or response being dispatched through handlers.
5247/// Requests include the context used to answer them; responses include the context
5248/// used to route them to the local requester.
5249///
5250/// Type parameters allow specifying the concrete request and notification types.
5251/// By default, both are `UntypedMessage` for dynamic dispatch.
5252/// The request context's response type matches the request's response type.
5253///
5254/// `P` is the [`Protocol`] marker of the request handles. Typed handlers on
5255/// v1 builders receive `Dispatch<Req, Notif, V1>` and those on v2 builders
5256/// `Dispatch<Req, Notif, V2>`. The version-neutral [`HandleDispatchFrom`]
5257/// layer uses the default, [`Raw`], whose `id()` is the wire ID.
5258#[derive(Debug)]
5259pub enum Dispatch<
5260    Req: JsonRpcRequest = UntypedMessage,
5261    Notif: JsonRpcNotification = UntypedMessage,
5262    P: Protocol = Raw,
5263> {
5264    /// Incoming request and the context where the response should be sent.
5265    Request(Req, Responder<Req::Response, P>),
5266
5267    /// Incoming notification.
5268    Notification(Notif),
5269
5270    /// Incoming response to a request we sent.
5271    ///
5272    /// The first field is the response result (success or error from the remote).
5273    /// The second field is the context for forwarding the response to its destination
5274    /// (typically a waiting oneshot channel).
5275    Response(
5276        Result<Req::Response, crate::Error>,
5277        ResponseRouter<Req::Response, P>,
5278    ),
5279}
5280
5281impl<Req: JsonRpcRequest, Notif: JsonRpcNotification, P: Protocol> Dispatch<Req, Notif, P> {
5282    /// Map the request and notification types to new types.
5283    ///
5284    /// Note: Response variants are passed through unchanged since they don't
5285    /// contain a parseable message payload.
5286    pub fn map<Req1, Notif1>(
5287        self,
5288        map_request: impl FnOnce(
5289            Req,
5290            Responder<Req::Response, P>,
5291        ) -> (Req1, Responder<Req1::Response, P>),
5292        map_notification: impl FnOnce(Notif) -> Notif1,
5293    ) -> Dispatch<Req1, Notif1, P>
5294    where
5295        Req1: JsonRpcRequest<Response = Req::Response>,
5296        Notif1: JsonRpcNotification,
5297    {
5298        match self {
5299            Dispatch::Request(request, responder) => {
5300                let (new_request, new_responder) = map_request(request, responder);
5301                Dispatch::Request(new_request, new_responder)
5302            }
5303            Dispatch::Notification(notification) => {
5304                let new_notification = map_notification(notification);
5305                Dispatch::Notification(new_notification)
5306            }
5307            Dispatch::Response(result, router) => Dispatch::Response(result, router),
5308        }
5309    }
5310
5311    /// Convert the message in self to an untyped message.
5312    ///
5313    /// Note: Response variants don't have an untyped message representation.
5314    /// This returns an error for Response variants.
5315    pub fn to_untyped_message(&self) -> Result<UntypedMessage, crate::Error> {
5316        match self {
5317            Dispatch::Request(request, _) => request.to_untyped_message(),
5318            Dispatch::Notification(notification) => notification.to_untyped_message(),
5319            Dispatch::Response(_, _) => Err(crate::util::internal_error(
5320                "Response variant has no untyped message representation",
5321            )),
5322        }
5323    }
5324
5325    /// Convert self to an untyped message context.
5326    ///
5327    /// Note: Response variants cannot be converted. This returns an error for Response variants.
5328    pub fn into_untyped_dispatch(
5329        self,
5330    ) -> Result<Dispatch<UntypedMessage, UntypedMessage, P>, crate::Error> {
5331        match self {
5332            Dispatch::Request(request, responder) => Ok(Dispatch::Request(
5333                request.to_untyped_message()?,
5334                responder.erase_to_json(),
5335            )),
5336            Dispatch::Notification(notification) => {
5337                Ok(Dispatch::Notification(notification.to_untyped_message()?))
5338            }
5339            Dispatch::Response(_, _) => Err(crate::util::internal_error(
5340                "cannot convert Response variant to untyped message context",
5341            )),
5342        }
5343    }
5344
5345    /// Returns the request ID if this is a request or response, None if notification.
5346    ///
5347    /// See [`Responder::id`] and [`ResponseRouter::id`].
5348    pub fn id(&self) -> Option<&P::RequestId> {
5349        match self {
5350            Dispatch::Request(_, cx) => Some(cx.id()),
5351            Dispatch::Notification(_) => None,
5352            Dispatch::Response(_, cx) => Some(cx.id()),
5353        }
5354    }
5355
5356    /// Returns the request ID exactly as it appears on the wire if this is a
5357    /// request or response, None if notification.
5358    ///
5359    /// See [`Responder::raw_id`] and [`ResponseRouter::raw_id`].
5360    #[must_use]
5361    pub fn raw_id(&self) -> Option<&RawJsonRpcId> {
5362        match self {
5363            Dispatch::Request(_, cx) => Some(cx.raw_id()),
5364            Dispatch::Notification(_) => None,
5365            Dispatch::Response(_, cx) => Some(cx.raw_id()),
5366        }
5367    }
5368
5369    /// The same dispatch with request handles for protocol `Q`.
5370    ///
5371    /// Fails only for a request whose `null` ID `Q` does not allow, which
5372    /// cannot happen for a message received on a connection that speaks `Q`.
5373    pub(crate) fn retag<Q: Protocol>(self) -> Result<Dispatch<Req, Notif, Q>, crate::Error> {
5374        Ok(match self {
5375            Dispatch::Request(request, responder) => Dispatch::Request(request, responder.retag()?),
5376            Dispatch::Notification(notification) => Dispatch::Notification(notification),
5377            Dispatch::Response(result, router) => Dispatch::Response(result, router.retag()?),
5378        })
5379    }
5380
5381    /// The same dispatch with version-neutral [`Raw`] request handles, for
5382    /// example to hand a typed dispatch to a [`HandleDispatchFrom`]
5383    /// implementation. This never fails.
5384    pub fn into_raw(self) -> Dispatch<Req, Notif, Raw> {
5385        match self {
5386            Dispatch::Request(request, responder) => {
5387                Dispatch::Request(request, responder.into_raw())
5388            }
5389            Dispatch::Notification(notification) => Dispatch::Notification(notification),
5390            Dispatch::Response(result, router) => Dispatch::Response(result, router.into_raw()),
5391        }
5392    }
5393
5394    fn handler_error_target(&self) -> Option<HandlerErrorTarget> {
5395        match self {
5396            Dispatch::Request(_, responder) => {
5397                Some(HandlerErrorTarget::Request(responder.reply_target()))
5398            }
5399            Dispatch::Notification(_) => None,
5400            Dispatch::Response(_, router) => {
5401                Some(HandlerErrorTarget::Response(router.reply_target.clone()))
5402            }
5403        }
5404    }
5405
5406    /// Returns the method of the message.
5407    ///
5408    /// For requests and notifications, this is the method from the message payload.
5409    /// For responses, this is the method of the original request.
5410    pub fn method(&self) -> &str {
5411        match self {
5412            Dispatch::Request(msg, _) => msg.method(),
5413            Dispatch::Notification(msg) => msg.method(),
5414            Dispatch::Response(_, cx) => cx.method(),
5415        }
5416    }
5417}
5418
5419impl<Req: JsonRpcRequest, Notif: JsonRpcNotification> Dispatch<Req, Notif, Raw> {
5420    /// Convert this version-neutral dispatch to protocol `Q`'s request
5421    /// handles.
5422    ///
5423    /// See [`Responder::into_protocol`]: this fails only for a request whose
5424    /// `null` ID `Q` does not allow, with the Invalid Request error
5425    /// (`-32600`).
5426    pub fn into_protocol<Q: Protocol>(self) -> Result<Dispatch<Req, Notif, Q>, crate::Error> {
5427        self.retag()
5428    }
5429}
5430
5431impl Dispatch {
5432    /// Attempts to parse `self` into a typed message context.
5433    ///
5434    /// # Returns
5435    ///
5436    /// * `Ok(Ok(typed))` if this dispatch matches the requested type for its variant
5437    /// * `Ok(Err(self))` if it does not match the requested type for its variant
5438    /// * `Err` if its method matches the requested type but parsing fails
5439    #[tracing::instrument(skip(self), fields(Request = ?std::any::type_name::<Req>(), Notif = ?std::any::type_name::<Notif>()), level = "trace", ret)]
5440    pub(crate) fn into_typed_dispatch<Req: JsonRpcRequest, Notif: JsonRpcNotification>(
5441        self,
5442    ) -> Result<Result<Dispatch<Req, Notif>, Dispatch>, crate::Error> {
5443        tracing::debug!(
5444            method = self.method(),
5445            id = ?self.raw_id(),
5446            "into_typed_dispatch"
5447        );
5448        match self {
5449            Dispatch::Request(message, responder) => {
5450                if Req::matches_method(&message.method) {
5451                    match Req::parse_message(&message.method, &message.params) {
5452                        Ok(req) => {
5453                            tracing::trace!("parsed ok");
5454                            Ok(Ok(Dispatch::Request(req, responder.cast())))
5455                        }
5456                        Err(err) => {
5457                            tracing::trace!(err = ?crate::util::redacted_error(err.code), "parse error");
5458                            Err(err)
5459                        }
5460                    }
5461                } else {
5462                    tracing::trace!("method doesn't match");
5463                    Ok(Err(Dispatch::Request(message, responder)))
5464                }
5465            }
5466
5467            Dispatch::Notification(message) => {
5468                if Notif::matches_method(&message.method) {
5469                    match Notif::parse_message(&message.method, &message.params) {
5470                        Ok(notif) => {
5471                            tracing::trace!("parse ok");
5472                            Ok(Ok(Dispatch::Notification(notif)))
5473                        }
5474                        Err(err) => {
5475                            tracing::trace!(err = ?crate::util::redacted_error(err.code), "parse error");
5476                            Err(err)
5477                        }
5478                    }
5479                } else {
5480                    tracing::trace!("method doesn't match");
5481                    Ok(Err(Dispatch::Notification(message)))
5482                }
5483            }
5484
5485            Dispatch::Response(result, cx) => {
5486                let method = cx.method();
5487                if Req::matches_method(method) {
5488                    // Parse the response result
5489                    let typed_result = match result {
5490                        Ok(value) => {
5491                            match <Req::Response as JsonRpcResponse>::from_json(method, value) {
5492                                Ok(parsed) => {
5493                                    tracing::trace!("parse ok");
5494                                    Ok(parsed)
5495                                }
5496                                Err(err) => {
5497                                    tracing::trace!(err = ?crate::util::redacted_error(err.code), "parse error");
5498                                    return Err(err);
5499                                }
5500                            }
5501                        }
5502                        Err(err) => {
5503                            tracing::trace!("error, passthrough");
5504                            Err(err)
5505                        }
5506                    };
5507                    Ok(Ok(Dispatch::Response(typed_result, cx.cast())))
5508                } else {
5509                    tracing::trace!("method doesn't match");
5510                    Ok(Err(Dispatch::Response(result, cx)))
5511                }
5512            }
5513        }
5514    }
5515
5516    /// True if this message has a field with the given name.
5517    ///
5518    /// Returns `false` for Response variants.
5519    #[must_use]
5520    pub fn has_field(&self, field_name: &str) -> bool {
5521        self.message()
5522            .and_then(|m| m.params().get(field_name))
5523            .is_some()
5524    }
5525
5526    /// Returns true if this message has a session-id field.
5527    ///
5528    /// Returns `false` for Response variants.
5529    pub(crate) fn has_session_id(&self) -> bool {
5530        self.has_field("sessionId")
5531    }
5532
5533    /// Extract the ACP session-id from this message (if any).
5534    ///
5535    /// Returns `Ok(None)` for Response variants.
5536    pub(crate) fn get_session_id(&self) -> Result<Option<SessionId>, crate::Error> {
5537        let Some(message) = self.message() else {
5538            return Ok(None);
5539        };
5540        let Some(value) = message.params().get("sessionId") else {
5541            return Ok(None);
5542        };
5543        let session_id = serde_json::from_value(value.clone())?;
5544        Ok(Some(session_id))
5545    }
5546
5547    /// Try to parse this as a notification of the given type.
5548    ///
5549    /// # Returns
5550    ///
5551    /// * `Ok(Ok(typed))` if this is a notification of the requested type
5552    /// * `Ok(Err(self))` if this is not a matching notification
5553    /// * `Err` if its method matches the requested type but parsing fails
5554    pub fn into_notification<N: JsonRpcNotification>(
5555        self,
5556    ) -> Result<Result<N, Dispatch>, crate::Error> {
5557        match self {
5558            Dispatch::Notification(msg) => {
5559                if !N::matches_method(&msg.method) {
5560                    return Ok(Err(Dispatch::Notification(msg)));
5561                }
5562                match N::parse_message(&msg.method, &msg.params) {
5563                    Ok(n) => Ok(Ok(n)),
5564                    Err(err) => Err(err),
5565                }
5566            }
5567            Dispatch::Request(..) | Dispatch::Response(..) => Ok(Err(self)),
5568        }
5569    }
5570
5571    /// Try to parse this as a request of the given type.
5572    ///
5573    /// # Returns
5574    ///
5575    /// * `Ok(Ok(typed))` if this is a request of the requested type
5576    /// * `Ok(Err(self))` if this is not a matching request
5577    /// * `Err` if its method matches the requested type but parsing fails
5578    pub fn into_request<Req: JsonRpcRequest>(
5579        self,
5580    ) -> Result<Result<(Req, Responder<Req::Response>), Dispatch>, crate::Error> {
5581        match self {
5582            Dispatch::Request(msg, responder) => {
5583                if !Req::matches_method(&msg.method) {
5584                    return Ok(Err(Dispatch::Request(msg, responder)));
5585                }
5586                match Req::parse_message(&msg.method, &msg.params) {
5587                    Ok(req) => Ok(Ok((req, responder.cast()))),
5588                    Err(err) => Err(err),
5589                }
5590            }
5591            Dispatch::Notification(..) | Dispatch::Response(..) => Ok(Err(self)),
5592        }
5593    }
5594}
5595
5596impl<M: JsonRpcRequest + JsonRpcNotification> Dispatch<M, M> {
5597    /// Returns the message payload for requests and notifications.
5598    ///
5599    /// Returns `None` for Response variants since they don't contain a message payload.
5600    pub fn message(&self) -> Option<&M> {
5601        match self {
5602            Dispatch::Request(msg, _) | Dispatch::Notification(msg) => Some(msg),
5603            Dispatch::Response(_, _) => None,
5604        }
5605    }
5606
5607    /// Map the request/notification message.
5608    ///
5609    /// Response variants pass through unchanged.
5610    pub(crate) fn try_map_message(
5611        self,
5612        map_message: impl FnOnce(M) -> Result<M, crate::Error>,
5613    ) -> Result<Dispatch<M, M>, crate::Error> {
5614        match self {
5615            Dispatch::Request(request, cx) => Ok(Dispatch::Request(map_message(request)?, cx)),
5616            Dispatch::Notification(notification) => {
5617                Ok(Dispatch::<M, M>::Notification(map_message(notification)?))
5618            }
5619            Dispatch::Response(result, cx) => Ok(Dispatch::Response(result, cx)),
5620        }
5621    }
5622}
5623
5624/// An incoming JSON message without any typing. Can be a request or a notification.
5625#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
5626pub struct UntypedMessage {
5627    /// The JSON-RPC method name
5628    pub method: String,
5629    /// The JSON-RPC parameters as a raw JSON value
5630    pub params: serde_json::Value,
5631}
5632
5633impl UntypedMessage {
5634    /// Returns an untyped message with the given method and parameters.
5635    pub fn new(method: &str, params: impl Serialize) -> Result<Self, crate::Error> {
5636        let params = serde_json::to_value(params)?;
5637        Ok(Self {
5638            method: method.to_string(),
5639            params,
5640        })
5641    }
5642
5643    /// Returns the method name
5644    #[must_use]
5645    pub fn method(&self) -> &str {
5646        &self.method
5647    }
5648
5649    /// Returns the parameters as a JSON value
5650    #[must_use]
5651    pub fn params(&self) -> &serde_json::Value {
5652        &self.params
5653    }
5654
5655    /// Consumes this message and returns the method and params
5656    #[must_use]
5657    pub fn into_parts(self) -> (String, serde_json::Value) {
5658        (self.method, self.params)
5659    }
5660
5661    /// Convert `self` to a raw JSON-RPC message.
5662    pub(crate) fn into_raw_jsonrpc_message(
5663        self,
5664        id: Option<RawJsonRpcId>,
5665    ) -> Result<RawJsonRpcMessage, crate::Error> {
5666        let Self { method, params } = self;
5667        match id {
5668            Some(id) => RawJsonRpcMessage::request(method, params, id),
5669            None => RawJsonRpcMessage::notification(method, params),
5670        }
5671    }
5672}
5673
5674impl JsonRpcMessage for UntypedMessage {
5675    fn matches_method(_method: &str) -> bool {
5676        // UntypedMessage matches any method - it's the untyped fallback
5677        true
5678    }
5679
5680    fn method(&self) -> &str {
5681        &self.method
5682    }
5683
5684    fn to_untyped_message(&self) -> Result<UntypedMessage, crate::Error> {
5685        Ok(self.clone())
5686    }
5687
5688    fn parse_message(method: &str, params: &serde_json::Value) -> Result<Self, crate::Error> {
5689        UntypedMessage::new(method, params)
5690    }
5691}
5692
5693impl JsonRpcRequest for UntypedMessage {
5694    type Response = serde_json::Value;
5695}
5696
5697impl JsonRpcNotification for UntypedMessage {}
5698
5699/// Represents a pending response of type `R` from an outgoing request.
5700///
5701/// Returned by [`ConnectionTo::send_request`], this type provides explicit response-consumption
5702/// modes. The API is intentionally designed to make it difficult to accidentally wait for a
5703/// response inside the dispatch loop.
5704///
5705/// # Anti-Footgun Design
5706///
5707/// You cannot directly `.await` a `SentRequest`. Instead, you must choose how to handle
5708/// the response:
5709///
5710/// ## Option 1: Register an Ordered Callback (Safe in Handlers)
5711///
5712/// Calling [`on_receiving_result`](Self::on_receiving_result) registers the callback and returns
5713/// immediately. When ordered consumption is selected before the response is routed during its
5714/// original dispatch, the loop waits for the callback to complete before processing the next
5715/// message:
5716///
5717/// ```no_run
5718/// # use agent_client_protocol_test::*;
5719/// # async fn example(cx: agent_client_protocol::ConnectionTo<agent_client_protocol::UntypedRole>) -> Result<(), agent_client_protocol::Error> {
5720/// cx.send_request(MyRequest {})
5721///     .on_receiving_result(async |result| {
5722///         match result {
5723///             Ok(response) => {
5724///                 // Handle successful response
5725///                 Ok(())
5726///             }
5727///             Err(error) => {
5728///                 // Handle error
5729///                 Err(error)
5730///             }
5731///         }
5732///     })?;
5733/// # Ok(())
5734/// # }
5735/// ```
5736///
5737/// ## Option 2: Wait Outside the Dispatch Loop
5738///
5739/// Use [`block_task`](Self::block_task) only when the current task already runs outside the
5740/// dispatch loop—for example, in the foreground future passed to `connect_with` or in a task
5741/// created with [`ConnectionTo::spawn`]. Never await it in a handler:
5742///
5743/// ```no_run
5744/// # use agent_client_protocol_test::*;
5745/// # async fn example(cx: agent_client_protocol::ConnectionTo<agent_client_protocol::UntypedRole>) -> Result<(), agent_client_protocol::Error> {
5746/// // ✅ Safe: Spawned task runs concurrently
5747/// cx.spawn({
5748///     let cx = cx.clone();
5749///     async move {
5750///         let response = cx.send_request(MyRequest {})
5751///             .block_task()
5752///             .await?;
5753///         // Process response...
5754///         Ok(())
5755///     }
5756/// })?;
5757/// # Ok(())
5758/// # }
5759/// ```
5760///
5761/// ```no_run
5762/// # use agent_client_protocol_test::*;
5763/// # async fn example() -> Result<(), agent_client_protocol::Error> {
5764/// # let connection = mock_connection();
5765/// // ❌ NEVER do this in a handler - blocks the event loop!
5766/// connection.on_receive_request(async |req: MyRequest, responder, cx| {
5767///     let response = cx.send_request(MyRequest {})
5768///         .block_task()  // This will deadlock!
5769///         .await?;
5770///     responder.respond(response)
5771/// }, agent_client_protocol::on_receive_request!())
5772/// # .connect_to(agent_client_protocol_test::MockTransport).await?;
5773/// # Ok(())
5774/// # }
5775/// ```
5776///
5777/// # Why This Design?
5778///
5779/// If you block the event loop while waiting for a response, the connection cannot process
5780/// the incoming response message, creating a deadlock. This API design prevents that footgun
5781/// by making blocking explicit and encouraging non-blocking patterns.
5782///
5783/// # Drop Behavior
5784///
5785/// By default, dropping a `SentRequest` before the SDK has received the
5786/// response sends a `$/cancel_request` notification asking the peer to cancel
5787/// the request, then discards the response when it arrives. Requests whose
5788/// eventual response should be ignored, but which should keep running on the
5789/// peer, should use [`detach`](Self::detach) instead.
5790///
5791/// # Incoming Transport EOF
5792///
5793/// If the incoming transport reaches clean EOF before the response arrives, every
5794/// consumption mode receives an error with the message `Incoming transport
5795/// closed` and data containing
5796/// `{"reason":"incoming_transport_closed","method":"..."}`. Requests made
5797/// after incoming EOF fail immediately with the same error. Use
5798/// [`is_incoming_transport_closed`] to identify it.
5799///
5800/// # Protocol version
5801///
5802/// `P` is the [`Protocol`] marker of the connection handle that sent the
5803/// request: [`Raw`] for the version-neutral [`ConnectionTo`], `V2` for
5804/// `v2::ConnectionTo`. It selects the type [`id`](Self::id) returns.
5805#[must_use = "dropping a SentRequest asks the peer to cancel the request and \
5806              discards the response; consume it with `block_task`, \
5807              `on_receiving_result`, `forward_response_to`, or `detach`"]
5808pub struct SentRequest<T, P: Protocol = Raw> {
5809    id: HandleId<P>,
5810    method: String,
5811    task_tx: TaskTx,
5812    response_rx: oneshot::Receiver<ResponsePayload>,
5813    /// Converts the raw response into the typed result. It receives error
5814    /// responses too, so [`SentRequest::recover_error`] can map an error to a
5815    /// success.
5816    to_result:
5817        Box<dyn FnOnce(Result<serde_json::Value, crate::Error>) -> Result<T, crate::Error> + Send>,
5818    cancellation: SentRequestCancellation,
5819    response_ordering: ResponseOrdering,
5820    /// Cancellation markers of other (incoming) requests whose cancellation
5821    /// should be forwarded to this request. See
5822    /// [`forward_cancellation_from`](Self::forward_cancellation_from).
5823    cancellation_sources: Vec<RequestCancellation>,
5824}
5825
5826/// A request that has not been published to its connection.
5827///
5828/// Created by [`ConnectionTo::prepare_request`] or
5829/// [`ConnectionTo::prepare_request_to`]. Preparation serializes the request but
5830/// does not register a pending reply or enqueue outgoing traffic. Dropping this
5831/// value sends neither the request nor a cancellation notification.
5832///
5833/// A consuming method publishes the request synchronously:
5834///
5835/// - [`on_receiving_result`](Self::on_receiving_result),
5836///   [`on_receiving_ok_result`](Self::on_receiving_ok_result), and
5837///   [`forward_response_to`](Self::forward_response_to) register ordered response
5838///   handling before publication. When a peer response is routed during its
5839///   original dispatch, later inbound messages wait for that handling to finish.
5840/// - [`block_task`](Self::block_task) publishes immediately and returns an
5841///   unordered response future. Publication does not wait for its first poll.
5842/// - [`detach`](Self::detach) publishes immediately and discards the response.
5843///
5844/// Ordered callbacks must do bounded work and must not await later inbound
5845/// traffic on the same connection. EOF failures and responses routed through a
5846/// retained [`ResponseRouter`] after their original dispatch have no ordering
5847/// barrier. See [`crate::concepts::ordering`].
5848///
5849/// # Errors
5850///
5851/// Preparation and publication failures are delivered to the selected response
5852/// consumer. Callback-style methods return an error if their task cannot be
5853/// registered; in that case the request is not published. [`detach`](Self::detach)
5854/// returns preparation or publication errors directly because it has no response
5855/// consumer. A callback returning an error terminates the connection.
5856///
5857/// `P` is the [`Protocol`] marker that types [`id`](Self::id); see
5858/// [`SentRequest`].
5859#[must_use = "a prepared request is not sent until consumed with `block_task`, \
5860              `on_receiving_result`, `forward_response_to`, or `detach`"]
5861pub struct PreparedRequest<T, P: Protocol = Raw> {
5862    sent: SentRequest<T, P>,
5863    publication: RequestPublication,
5864}
5865
5866struct RequestPublication {
5867    message: Result<OutgoingMessage, crate::Error>,
5868    pending_reply: PendingReply,
5869    message_tx: OutgoingMessageTx,
5870    pending_replies: PendingRepliesRegistrar,
5871    incoming_closed: IncomingClosed,
5872}
5873
5874impl RequestPublication {
5875    fn publish(self) -> Result<(), crate::Error> {
5876        let message = if self.incoming_closed.is_closing() {
5877            Err(incoming_transport_closed_error(&self.pending_reply.method))
5878        } else {
5879            self.message
5880        };
5881        let message = match message {
5882            Ok(message) => message,
5883            Err(error) => {
5884                self.pending_reply.fail(error.clone());
5885                return Err(error);
5886            }
5887        };
5888        let OutgoingMessage::Request { id, method, .. } = &message else {
5889            unreachable!();
5890        };
5891        let id = id.clone();
5892        let method = method.clone();
5893        let cancellation_disarm = self.pending_reply.cancellation_disarm.clone();
5894        // Register before enqueueing so incoming EOF can fail every observable
5895        // request before close callbacks begin. The outgoing actor checks that
5896        // the registration still exists before sending the request.
5897        self.pending_replies
5898            .subscribe(id.clone(), self.pending_reply, &self.incoming_closed)?;
5899        if self.message_tx.unbounded_send(message).is_err() {
5900            let error = if self.incoming_closed.is_closing() {
5901                incoming_transport_closed_error(&method)
5902            } else {
5903                crate::util::internal_error(format!("failed to send outgoing request `{method}`"))
5904            };
5905            if let Some(pending_reply) = self.pending_replies.remove(&id) {
5906                pending_reply.fail(error.clone());
5907            }
5908            return Err(error);
5909        }
5910        // An escaped cancellation handle must not enqueue cancellation before
5911        // the request. A fast response or EOF may already have disarmed it.
5912        cancellation_disarm.arm();
5913        Ok(())
5914    }
5915}
5916
5917impl<T: Debug, P: Protocol> Debug for PreparedRequest<T, P> {
5918    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
5919        formatter
5920            .debug_struct("PreparedRequest")
5921            .field("request", &self.sent)
5922            .finish_non_exhaustive()
5923    }
5924}
5925
5926impl<T, P: Protocol> PreparedRequest<T, P> {
5927    /// The same prepared request with a version-neutral [`Raw`] request ID,
5928    /// for code that is not tied to one protocol version. This never fails
5929    /// and does not publish the request.
5930    pub fn into_raw(self) -> PreparedRequest<T, Raw> {
5931        PreparedRequest {
5932            sent: self.sent.into_raw(),
5933            publication: self.publication,
5934        }
5935    }
5936}
5937
5938impl<T> PreparedRequest<T, Raw> {
5939    /// Convert this version-neutral prepared request to protocol `Q`'s
5940    /// handle without publishing it.
5941    ///
5942    /// Only the request ID's type changes. The SDK never assigns `null` to
5943    /// requests it sends, so this succeeds for every marker.
5944    pub fn into_protocol<Q: Protocol>(self) -> Result<PreparedRequest<T, Q>, crate::Error> {
5945        Ok(PreparedRequest {
5946            sent: self.sent.into_protocol()?,
5947            publication: self.publication,
5948        })
5949    }
5950}
5951
5952impl<T, P: Protocol> PreparedRequest<T, P> {
5953    /// The ID reserved for this request, which has not been sent yet, as
5954    /// this request's protocol version types it.
5955    ///
5956    /// It is the ID written to the wire when the request is published. The
5957    /// SDK never assigns `null`.
5958    #[must_use]
5959    pub fn id(&self) -> &P::RequestId {
5960        self.sent.id()
5961    }
5962
5963    /// The ID reserved for this request in the version-neutral form used by
5964    /// routers, forwarders and tracing. It is the same ID as
5965    /// [`id`](Self::id).
5966    #[must_use]
5967    pub fn raw_id(&self) -> &RawJsonRpcId {
5968        self.sent.raw_id()
5969    }
5970
5971    /// The method of the prepared request.
5972    #[must_use]
5973    pub fn method(&self) -> &str {
5974        self.sent.method()
5975    }
5976
5977    /// Retain explicit cancellation control without publishing this request.
5978    ///
5979    /// The handle can outlive consumption of this request by an ordered callback
5980    /// or response future. Calling it before publication is a no-op, not a
5981    /// cancellation to apply when the request is later published. Dropping the
5982    /// handle does not cancel. See [`RequestCancellationHandle`].
5983    #[must_use]
5984    pub fn cancellation_handle(&self) -> RequestCancellationHandle {
5985        self.sent.cancellation_handle()
5986    }
5987
5988    /// Map a successful response without publishing the request.
5989    ///
5990    /// The mapper has the same contract as [`SentRequest::map`].
5991    pub fn map<U>(
5992        self,
5993        map_fn: impl FnOnce(T) -> Result<U, crate::Error> + 'static + Send,
5994    ) -> PreparedRequest<U, P>
5995    where
5996        T: 'static,
5997    {
5998        PreparedRequest {
5999            sent: self.sent.map(map_fn),
6000            publication: self.publication,
6001        }
6002    }
6003
6004    /// Register a cancellation source without publishing the request.
6005    ///
6006    /// After publication, cancellation is forwarded while awaiting the response,
6007    /// as described by [`SentRequest::forward_cancellation_from`].
6008    pub fn forward_cancellation_from(mut self, source: RequestCancellation) -> Self {
6009        self.sent = self.sent.forward_cancellation_from(source);
6010        self
6011    }
6012
6013    /// Publish now and return an unordered future for the response.
6014    ///
6015    /// The request is enqueued during this call, not when the future is first
6016    /// polled. Dropping that future asks the peer to cancel a still-outstanding
6017    /// request. Await it only outside the dispatch loop; awaiting it in an
6018    /// incoming handler deadlocks just like [`SentRequest::block_task`].
6019    ///
6020    /// # Errors
6021    ///
6022    /// The returned future delivers preparation, publication, and response errors.
6023    pub fn block_task(self) -> impl Future<Output = Result<T, crate::Error>> {
6024        self.into_sent_request(false).block_task()
6025    }
6026
6027    /// Publish now and discard the eventual response without cancelling.
6028    ///
6029    /// # Errors
6030    ///
6031    /// Returns immediate preparation or enqueue failures. Later local
6032    /// transformation errors and peer response errors are discarded along with
6033    /// successful responses. Transport failures still propagate through the
6034    /// connection future. Retained [`RequestCancellationHandle`] values can
6035    /// still explicitly cancel the request while it remains pending.
6036    pub fn detach(self) -> Result<(), crate::Error> {
6037        let result = self.publication.publish();
6038        self.sent.detach();
6039        result
6040    }
6041
6042    /// Register an ordered callback, then publish the request.
6043    ///
6044    /// Ordering is selected before publication, even when the connection runs
6045    /// concurrently. See [`PreparedRequest`] for barrier limits and deadlock risks.
6046    ///
6047    /// # Errors
6048    ///
6049    /// Returns an error if the callback task cannot be registered, without
6050    /// publishing the request. Preparation and publication errors are delivered
6051    /// to the callback. Returning an error from the callback ends the connection.
6052    #[track_caller]
6053    pub fn on_receiving_result<F>(
6054        self,
6055        task: impl FnOnce(Result<T, crate::Error>) -> F + 'static + Send,
6056    ) -> Result<(), crate::Error>
6057    where
6058        T: 'static,
6059        F: Future<Output = Result<(), crate::Error>> + 'static + Send,
6060    {
6061        self.consume_with(move |response| match response {
6062            Ok(result) => Either::Left(task(result)),
6063            Err(error) => Either::Right(future::ready(Err(error))),
6064        })
6065    }
6066
6067    /// Register an ordered success callback, then publish the request.
6068    ///
6069    /// Errors are forwarded to `responder`, as with
6070    /// [`SentRequest::on_receiving_ok_result`].
6071    ///
6072    /// # Errors
6073    ///
6074    /// Returns a task-registration error without publishing the request.
6075    /// Preparation, publication, and response errors are forwarded to `responder`.
6076    /// Returning an error from the callback ends the connection.
6077    #[track_caller]
6078    pub fn on_receiving_ok_result<F, Q: Protocol>(
6079        self,
6080        responder: Responder<T, Q>,
6081        task: impl FnOnce(T, Responder<T, Q>) -> F + 'static + Send,
6082    ) -> Result<(), crate::Error>
6083    where
6084        T: JsonRpcResponse,
6085        F: Future<Output = Result<(), crate::Error>> + 'static + Send,
6086    {
6087        self.on_receiving_result(async move |result| match result {
6088            Ok(value) => task(value, responder).await,
6089            Err(error) => responder.respond_with_error(error),
6090        })
6091    }
6092
6093    /// Register ordered response forwarding, then publish the request.
6094    ///
6095    /// Cancellation and response errors propagate as with
6096    /// [`SentRequest::forward_response_to`].
6097    ///
6098    /// # Errors
6099    ///
6100    /// Returns a task-registration error without publishing the request.
6101    /// Preparation, publication, and response errors are forwarded to `responder`.
6102    #[track_caller]
6103    pub fn forward_response_to<Q: Protocol>(
6104        self,
6105        responder: Responder<T, Q>,
6106    ) -> Result<(), crate::Error>
6107    where
6108        T: JsonRpcResponse,
6109    {
6110        self.forward_cancellation_from(responder.cancellation())
6111            .consume_with(async move |response| {
6112                responder.respond_with_result(response.unwrap_or_else(Err))
6113            })
6114    }
6115
6116    fn into_sent_request(self, ordered: bool) -> SentRequest<T, P> {
6117        if ordered {
6118            self.sent.response_ordering.mark_ordered();
6119        }
6120        // Publication errors also settle the response channel.
6121        drop(self.publication.publish());
6122        self.sent
6123    }
6124
6125    #[track_caller]
6126    fn consume_with<F>(
6127        self,
6128        handle: impl FnOnce(Result<Result<T, crate::Error>, crate::Error>) -> F + 'static + Send,
6129    ) -> Result<(), crate::Error>
6130    where
6131        T: 'static,
6132        F: Future<Output = Result<(), crate::Error>> + 'static + Send,
6133    {
6134        let published_tx = self.register_consumer(handle)?;
6135        drop(self.publication.publish());
6136        // Keep the cancellation guard here until publication completes. Even
6137        // destruction of the registered task cannot cancel before enqueueing.
6138        drop(published_tx.send(self.sent));
6139        Ok(())
6140    }
6141
6142    #[track_caller]
6143    fn register_consumer<F>(
6144        &self,
6145        handle: impl FnOnce(Result<Result<T, crate::Error>, crate::Error>) -> F + 'static + Send,
6146    ) -> Result<oneshot::Sender<SentRequest<T, P>>, crate::Error>
6147    where
6148        T: 'static,
6149        F: Future<Output = Result<(), crate::Error>> + 'static + Send,
6150    {
6151        self.sent.response_ordering.mark_ordered();
6152        let (published_tx, published_rx) = oneshot::channel::<SentRequest<T, P>>();
6153        Task::new(Location::caller(), async move {
6154            match published_rx.await {
6155                Ok(sent) => sent.handle_response(handle).await,
6156                // Publication was abandoned before the consumer took ownership.
6157                Err(_) => Ok(()),
6158            }
6159        })
6160        .spawn(&self.sent.task_tx)?;
6161        Ok(published_tx)
6162    }
6163}
6164
6165#[derive(Clone, Debug)]
6166pub(crate) struct SentRequestCancellationDisarm {
6167    state: Arc<AtomicU8>,
6168}
6169
6170#[repr(u8)]
6171enum OutgoingCancellationState {
6172    Unpublished,
6173    Armed,
6174    Disarmed,
6175}
6176
6177impl SentRequestCancellationDisarm {
6178    fn new() -> Self {
6179        Self {
6180            state: Arc::new(AtomicU8::new(OutgoingCancellationState::Unpublished as u8)),
6181        }
6182    }
6183
6184    fn disarm(&self) {
6185        self.state
6186            .store(OutgoingCancellationState::Disarmed as u8, Ordering::Release);
6187    }
6188
6189    fn arm(&self) -> bool {
6190        self.state
6191            .compare_exchange(
6192                OutgoingCancellationState::Unpublished as u8,
6193                OutgoingCancellationState::Armed as u8,
6194                Ordering::AcqRel,
6195                Ordering::Acquire,
6196            )
6197            .is_ok()
6198    }
6199
6200    fn take_armed(&self) -> bool {
6201        self.state
6202            .compare_exchange(
6203                OutgoingCancellationState::Armed as u8,
6204                OutgoingCancellationState::Disarmed as u8,
6205                Ordering::AcqRel,
6206                Ordering::Acquire,
6207            )
6208            .is_ok()
6209    }
6210
6211    fn is_armed(&self) -> bool {
6212        self.state.load(Ordering::Acquire) == OutgoingCancellationState::Armed as u8
6213    }
6214}
6215
6216/// Explicit cancellation control for one outgoing request.
6217///
6218/// Obtain this handle from [`PreparedRequest::cancellation_handle`] or
6219/// [`SentRequest::cancellation_handle`] before consuming the request. It retains
6220/// neither the response consumer nor application callback, so cancellation can
6221/// be requested without discarding the eventual response. It does not keep the
6222/// connection driver or response consumer alive, or abort local callback work.
6223///
6224/// Clones share the request's cancellation state with [`SentRequest::cancel`],
6225/// forwarded cancellation, and request-drop automatic cancellation. At most one
6226/// cancellation notification is attempted, on the original request's
6227/// connection. It is a plain `$/cancel_request`, even for a request sent to a
6228/// successor in a `proxy/successor` envelope. Dropping this handle neither cancels the request nor disables its
6229/// automatic cancellation.
6230///
6231/// Cancellation before publication is a no-op and is not remembered for later
6232/// publication. A call racing publication may also be a no-op; call after the
6233/// publishing method returns to target a pending request. Once the SDK routes a
6234/// response or fails the request, subsequent cancellation calls are no-ops,
6235/// even if its callback has not yet run. Detaching the request suppresses only
6236/// automatic cancellation; retained handles can still explicitly cancel it.
6237///
6238/// This controls outgoing requests, unlike [`RequestCancellation`], which
6239/// observes a peer's cancellation of an incoming request.
6240#[derive(Clone)]
6241pub struct RequestCancellationHandle {
6242    message_tx: OutgoingMessageTx,
6243    request_id: RawJsonRpcId,
6244    disarm: SentRequestCancellationDisarm,
6245}
6246
6247impl RequestCancellationHandle {
6248    /// Ask the peer to cancel this request without discarding its response.
6249    ///
6250    /// Cancellation is cooperative: the peer may respond normally or with a
6251    /// cancellation error. Repeated calls, including calls through other clones
6252    /// or the original request, return `Ok(())` without sending another
6253    /// notification. Calls before publication or after settlement are no-ops.
6254    /// An attempt begun before settlement may still enqueue afterward.
6255    ///
6256    /// `Ok(())` means this call encountered no immediate error, not that a
6257    /// notification was sent or the peer stopped work. This method does not
6258    /// wait for transmission or acknowledgment.
6259    ///
6260    /// # Errors
6261    ///
6262    /// Only the call that attempts to send reports serialization or enqueue
6263    /// failure. A failed send is not retried by later cancellation calls.
6264    pub fn cancel(&self) -> Result<(), crate::Error> {
6265        if !self.disarm.take_armed() {
6266            return Ok(());
6267        }
6268
6269        // Build the notification lazily: most requests are never cancelled,
6270        // so this avoids serializing a notification per outgoing request.
6271        //
6272        // `$/cancel_request` is hop-local and always sent plain, even for a
6273        // request sent to a successor: its `requestId` names the
6274        // `proxy/successor` request on this connection.
6275        let untyped = crate::v1::schema::CancelRequestNotification::new(self.request_id.clone())
6276            .to_untyped_message()?;
6277
6278        send_raw_message(&self.message_tx, OutgoingMessage::Notification { untyped })
6279    }
6280}
6281
6282impl Debug for RequestCancellationHandle {
6283    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
6284        formatter
6285            .debug_struct("RequestCancellationHandle")
6286            .field("request_id", &self.request_id)
6287            .field("armed", &self.disarm.is_armed())
6288            .finish_non_exhaustive()
6289    }
6290}
6291
6292#[derive(Debug)]
6293struct SentRequestCancellation {
6294    handle: RequestCancellationHandle,
6295    cancel_on_drop: bool,
6296}
6297
6298impl SentRequestCancellation {
6299    fn new(message_tx: OutgoingMessageTx, request_id: RawJsonRpcId) -> Self {
6300        Self {
6301            handle: RequestCancellationHandle {
6302                message_tx,
6303                request_id,
6304                disarm: SentRequestCancellationDisarm::new(),
6305            },
6306            cancel_on_drop: true,
6307        }
6308    }
6309
6310    fn disarm(&self) {
6311        self.handle.disarm.disarm();
6312    }
6313
6314    fn disarm_handle(&self) -> SentRequestCancellationDisarm {
6315        self.handle.disarm.clone()
6316    }
6317
6318    fn send(&self) -> Result<(), crate::Error> {
6319        self.handle.cancel()
6320    }
6321}
6322
6323impl Drop for SentRequestCancellation {
6324    fn drop(&mut self) {
6325        if !self.cancel_on_drop {
6326            return;
6327        }
6328        if let Err(error) = self.send() {
6329            tracing::debug!(
6330                error = ?crate::util::redacted_error(error.code),
6331                "failed to auto-cancel dropped request"
6332            );
6333        }
6334    }
6335}
6336
6337/// Await the response payload for an outgoing request, watching `sources` for
6338/// cancellation of the upstream requests it was registered with.
6339///
6340/// When any source reports cancellation, a `$/cancel_request` is forwarded to
6341/// the outgoing request (at most once, shared with [`SentRequest::cancel`] and
6342/// drop-time auto-cancellation), and the response is *still* awaited: the peer
6343/// always answers, with normal data or a cancellation error.
6344///
6345/// Watching is deliberately bounded by response arrival so that completed
6346/// requests do not leak waiters on markers that will never fire.
6347async fn await_response_forwarding_cancellation(
6348    response_rx: oneshot::Receiver<ResponsePayload>,
6349    cancellation: &SentRequestCancellation,
6350    sources: &[RequestCancellation],
6351) -> Result<ResponsePayload, oneshot::Canceled> {
6352    // Failing to forward the cancellation must not abort the wait: the
6353    // response (normal data or a cancellation error) may still arrive and
6354    // must still be processed.
6355    let forward_cancellation = || {
6356        if let Err(error) = cancellation.send() {
6357            tracing::debug!(
6358                error = ?crate::util::redacted_error(error.code),
6359                "failed to forward cancellation to downstream request"
6360            );
6361        }
6362    };
6363
6364    let response = if sources.is_empty() {
6365        response_rx.await
6366    } else if sources.iter().any(RequestCancellation::is_cancelled) {
6367        forward_cancellation();
6368        response_rx.await
6369    } else {
6370        let cancelled = sources.iter().map(|source| source.state.signal_rx.clone());
6371        match future::select(future::select_all(cancelled), response_rx).await {
6372            Either::Left((_, response_rx)) => {
6373                forward_cancellation();
6374                response_rx.await
6375            }
6376            Either::Right((response, _)) => response,
6377        }
6378    };
6379
6380    cancellation.disarm();
6381    response
6382}
6383
6384impl<T: Debug, P: Protocol> Debug for SentRequest<T, P> {
6385    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
6386        let mut debug = f.debug_struct("SentRequest");
6387        debug
6388            .field("id", &self.id)
6389            .field("method", &self.method)
6390            .field("task_tx", &self.task_tx)
6391            .field("response_rx", &self.response_rx);
6392        debug
6393            .field("cancellation", &self.cancellation)
6394            .field("cancellation_sources", &self.cancellation_sources);
6395        debug.finish_non_exhaustive()
6396    }
6397}
6398
6399impl<P: Protocol> SentRequest<serde_json::Value, P> {
6400    fn new(
6401        id: HandleId<P>,
6402        method: String,
6403        task_tx: mpsc::UnboundedSender<Task>,
6404        response_rx: oneshot::Receiver<ResponsePayload>,
6405        cancellation: SentRequestCancellation,
6406        response_ordering: ResponseOrdering,
6407    ) -> Self {
6408        Self {
6409            id,
6410            method,
6411            response_rx,
6412            task_tx,
6413            to_result: Box::new(|result| result),
6414            cancellation,
6415            response_ordering,
6416            cancellation_sources: Vec::new(),
6417        }
6418    }
6419}
6420
6421impl<T, P: Protocol> SentRequest<T, P> {
6422    /// Detach this request handle without waiting for its response.
6423    ///
6424    /// The response will be discarded when it arrives. This also disables the
6425    /// drop-time automatic cancellation described in
6426    /// [Drop Behavior](Self#drop-behavior), so use it for requests whose
6427    /// eventual response should be ignored, but which should keep running on
6428    /// the peer. The peer is still expected to answer the JSON-RPC request
6429    /// eventually; use a notification instead when no response is expected at
6430    /// all.
6431    ///
6432    /// To ask the peer to stop the request, call `cancel` instead, or drop the
6433    /// handle while automatic cancellation is enabled. A retained
6434    /// [`RequestCancellationHandle`] can still explicitly cancel the detached
6435    /// request until the SDK receives its response or fails it.
6436    pub fn detach(mut self) {
6437        self.cancellation.cancel_on_drop = false;
6438    }
6439
6440    /// Send a `$/cancel_request` notification for this outgoing request.
6441    ///
6442    /// The notification goes over the connection the original request used,
6443    /// with that request's ID, so this is the preferred way to cancel a
6444    /// [`SentRequest`] when the request handle is still available. It is sent
6445    /// plain, never wrapped in `proxy/successor`: for a request sent to a
6446    /// successor, the ID names the `proxy/successor` request on this
6447    /// connection, as the proxy chains RFD specifies.
6448    ///
6449    /// At most one cancellation attempt is made per request, shared with
6450    /// retained handles, forwarded cancellation, and automatic cancellation
6451    /// described in [Drop Behavior](Self#drop-behavior). Later calls return
6452    /// `Ok(())` without another attempt, including when the first attempt failed.
6453    /// Once the SDK has routed the response, a new call is a no-op; an attempt
6454    /// begun before settlement may still enqueue afterward.
6455    ///
6456    /// `Ok(())` means this call encountered no immediate error, not that a
6457    /// notification was sent or the peer stopped work.
6458    ///
6459    /// Errors are only reported by the call that attempts to send the
6460    /// notification.
6461    pub fn cancel(&self) -> Result<(), crate::Error> {
6462        self.cancellation.send()
6463    }
6464
6465    /// Obtain a handle that remains usable after this request is consumed.
6466    ///
6467    /// The handle shares the once-only cancellation state used by
6468    /// [`cancel`](Self::cancel), response routing, and automatic request-drop
6469    /// cancellation, but has no cancel-on-drop behavior of its own.
6470    #[must_use]
6471    pub fn cancellation_handle(&self) -> RequestCancellationHandle {
6472        self.cancellation.handle.clone()
6473    }
6474
6475    /// Forward cancellation of another request to this one.
6476    ///
6477    /// When the request that `source` belongs to is cancelled by its peer,
6478    /// a `$/cancel_request` for *this* request is sent to its peer, plain and
6479    /// with this request's ID, as [`cancel`](Self::cancel) sends it. The response is still awaited
6480    /// and delivered as usual (normal data or a cancellation error), so this
6481    /// composes with [`block_task`](Self::block_task) and
6482    /// [`on_receiving_result`](Self::on_receiving_result).
6483    ///
6484    /// This is the building block for proxies that forward a request with
6485    /// custom logic instead of [`forward_response_to`](Self::forward_response_to)
6486    /// (which wires this up automatically from its responder). Without it,
6487    /// custom forwarding *absorbs* cancellation: the upstream marker is still
6488    /// set, but nothing is sent downstream.
6489    ///
6490    /// ```
6491    /// # use agent_client_protocol::{ConnectionTo, Error, Responder, UntypedRole};
6492    /// # use agent_client_protocol_test::{MyRequest, MyResponse};
6493    /// # async fn example(request: MyRequest, responder: Responder<MyResponse>, backend: ConnectionTo<UntypedRole>) -> Result<(), Error> {
6494    /// backend
6495    ///     .send_request(request)
6496    ///     .forward_cancellation_from(responder.cancellation())
6497    ///     .on_receiving_result(async move |result| {
6498    ///         // Custom result handling, e.g. bookkeeping or rewriting.
6499    ///         responder.respond_with_result(result)
6500    ///     })?;
6501    /// # Ok(())
6502    /// # }
6503    /// ```
6504    ///
6505    /// May be called multiple times; cancellation of any registered source
6506    /// triggers the forwarding (at most one `$/cancel_request` is ever sent
6507    /// per request). Sources are observed while the response is being
6508    /// awaited — that is, once the handle is consumed with
6509    /// [`block_task`](Self::block_task),
6510    /// [`on_receiving_result`](Self::on_receiving_result), or
6511    /// [`forward_response_to`](Self::forward_response_to); a source that was
6512    /// already cancelled by then is honored immediately.
6513    pub fn forward_cancellation_from(mut self, source: RequestCancellation) -> Self {
6514        self.cancellation_sources.push(source);
6515        self
6516    }
6517}
6518
6519impl<T, P: Protocol> SentRequest<T, P> {
6520    /// The same request handle with a version-neutral [`Raw`] request ID, for
6521    /// code that is not tied to one protocol version. This never fails.
6522    pub fn into_raw(self) -> SentRequest<T, Raw> {
6523        self.retag_id(HandleId::into_raw)
6524    }
6525
6526    fn retag_id<Q: Protocol>(
6527        self,
6528        retag: impl FnOnce(HandleId<P>) -> HandleId<Q>,
6529    ) -> SentRequest<T, Q> {
6530        SentRequest {
6531            id: retag(self.id),
6532            method: self.method,
6533            task_tx: self.task_tx,
6534            response_rx: self.response_rx,
6535            to_result: self.to_result,
6536            cancellation: self.cancellation,
6537            response_ordering: self.response_ordering,
6538            cancellation_sources: self.cancellation_sources,
6539        }
6540    }
6541}
6542
6543impl<T> SentRequest<T, Raw> {
6544    /// Convert this version-neutral request handle to protocol `Q`'s handle.
6545    ///
6546    /// Only the request ID's type changes. The SDK never assigns `null` to
6547    /// requests it sends, so this succeeds for every marker.
6548    pub fn into_protocol<Q: Protocol>(self) -> Result<SentRequest<T, Q>, crate::Error> {
6549        let id = self.id.clone().into_protocol()?;
6550        Ok(self.retag_id(|_| id))
6551    }
6552}
6553
6554impl<T, P: Protocol> SentRequest<T, P> {
6555    /// The id of the outgoing request, as this request's protocol version
6556    /// types it. It is the ID written to the wire; the SDK never assigns
6557    /// `null`.
6558    #[must_use]
6559    pub fn id(&self) -> &P::RequestId {
6560        self.id.typed()
6561    }
6562
6563    /// The id of the outgoing request in the version-neutral form used by
6564    /// routers, forwarders and tracing. It is the same ID as
6565    /// [`id`](Self::id).
6566    #[must_use]
6567    pub fn raw_id(&self) -> &RawJsonRpcId {
6568        self.id.raw()
6569    }
6570
6571    /// The method of the request this is in response to.
6572    #[must_use]
6573    pub fn method(&self) -> &str {
6574        &self.method
6575    }
6576
6577    /// Map a JSON-RPC error response to a result of the same type.
6578    ///
6579    /// The recovery function only sees error responses from the peer, not
6580    /// local failures to decode a successful response. Return `Ok` to treat
6581    /// the error as success, or `Err` to keep failing.
6582    #[cfg_attr(not(feature = "unstable_protocol_v2"), allow(dead_code))]
6583    pub(crate) fn recover_error(
6584        self,
6585        recover: impl FnOnce(crate::Error) -> Result<T, crate::Error> + 'static + Send,
6586    ) -> SentRequest<T, P>
6587    where
6588        T: 'static,
6589    {
6590        let to_result = self.to_result;
6591        SentRequest {
6592            to_result: Box::new(move |result| match result {
6593                Ok(value) => to_result(Ok(value)),
6594                Err(error) => recover(error),
6595            }),
6596            ..self
6597        }
6598    }
6599
6600    /// Map a successful JSON-RPC response into an application type.
6601    ///
6602    /// The mapped type does not need to implement [`JsonRpcResponse`]. The
6603    /// mapper runs at most once and may consume captured state. JSON-RPC error
6604    /// responses bypass the mapper. The mapped type may carry a non-`'static`
6605    /// lifetime when it is consumed with [`block_task`](Self::block_task);
6606    /// callback-style consumption still requires a `'static` mapped type
6607    /// because its work is spawned onto the connection.
6608    pub fn map<U>(
6609        self,
6610        map_fn: impl FnOnce(T) -> Result<U, crate::Error> + 'static + Send,
6611    ) -> SentRequest<U, P>
6612    where
6613        T: 'static,
6614    {
6615        SentRequest {
6616            id: self.id,
6617            method: self.method,
6618            response_rx: self.response_rx,
6619            task_tx: self.task_tx,
6620            to_result: Box::new(move |result| map_fn((self.to_result)(result)?)),
6621            cancellation: self.cancellation,
6622            response_ordering: self.response_ordering,
6623            cancellation_sources: self.cancellation_sources,
6624        }
6625    }
6626
6627    /// Forward the response (success or error) to a request context when it arrives.
6628    ///
6629    /// This is a convenience method for proxying messages between connections. When the
6630    /// response arrives, it will be automatically sent to the provided request context,
6631    /// whether it's a successful response or an error.
6632    ///
6633    /// # Example: Proxying requests
6634    ///
6635    /// ```
6636    /// # use agent_client_protocol::UntypedRole;
6637    /// # use agent_client_protocol::{Builder, ConnectionTo};
6638    /// # use agent_client_protocol_test::*;
6639    /// # async fn example(cx: ConnectionTo<UntypedRole>) -> Result<(), agent_client_protocol::Error> {
6640    /// // Set up backend connection builder
6641    /// let backend = UntypedRole.builder()
6642    ///     .on_receive_request(async |req: MyRequest, responder, cx| {
6643    ///         responder.respond(MyResponse { status: "ok".into() })
6644    ///     }, agent_client_protocol::on_receive_request!());
6645    ///
6646    /// // Spawn backend and get a context to send to it
6647    /// let backend_connection = cx.spawn_connection(backend, MockTransport)?;
6648    ///
6649    /// // Set up proxy that forwards requests to backend
6650    /// UntypedRole.builder()
6651    ///     .on_receive_request({
6652    ///         let backend_connection = backend_connection.clone();
6653    ///         async move |req: MyRequest, responder, cx| {
6654    ///             // Forward the request to backend and proxy the response back
6655    ///             backend_connection.send_request(req)
6656    ///                 .forward_response_to(responder)?;
6657    ///             Ok(())
6658    ///         }
6659    ///     }, agent_client_protocol::on_receive_request!());
6660    /// # Ok(())
6661    /// # }
6662    /// ```
6663    ///
6664    /// # Type Safety
6665    ///
6666    /// The request context's response type must match the request's response type,
6667    /// ensuring type-safe message forwarding.
6668    ///
6669    /// # When to Use
6670    ///
6671    /// Use this when:
6672    /// - You're implementing a proxy or gateway pattern
6673    /// - You want to forward responses without processing them
6674    /// - The response types match between the outgoing request and incoming request
6675    ///
6676    /// This is equivalent to calling `on_receiving_result` and manually forwarding
6677    /// the result, with two proxy-specific additions:
6678    ///
6679    /// - If the pending response cannot be delivered, the incoming request is
6680    ///   answered with an internal error instead of being left unanswered.
6681    ///   Known clean incoming EOF is delivered like any other response
6682    ///   error; an unexpected response-channel loss is forwarded as an outer
6683    ///   consumption error.
6684    /// - When the peer cancels the incoming request, the cancellation is
6685    ///   forwarded to the outgoing request, and the downstream response
6686    ///   (normal data or a cancellation error) is still forwarded back. This is
6687    ///   equivalent to registering the responder's marker with
6688    ///   `forward_cancellation_from`.
6689    #[track_caller]
6690    pub fn forward_response_to<Q: Protocol>(
6691        self,
6692        responder: Responder<T, Q>,
6693    ) -> Result<(), crate::Error>
6694    where
6695        T: JsonRpcResponse,
6696    {
6697        let this = self.forward_cancellation_from(responder.cancellation());
6698
6699        this.consume_with(async move |response| {
6700            // An unexpected response-channel loss (outer `Err`) is forwarded
6701            // as an error: the incoming request must not be left unanswered.
6702            responder.respond_with_result(response.unwrap_or_else(Err))
6703        })
6704    }
6705
6706    /// Spawn the response-consumption task shared by
6707    /// [`on_receiving_result`](Self::on_receiving_result) and
6708    /// [`forward_response_to`](Self::forward_response_to).
6709    ///
6710    /// The task awaits the response (forwarding cancellation from registered
6711    /// sources while waiting, converts the payload, and invokes `handle` with
6712    /// the typed result (`Ok(Result<T, _>)`). The dispatch loop's ack, if any,
6713    /// is sent after `handle` completes.
6714    ///
6715    /// Clean incoming EOF is delivered as `Ok(Err(error))`, just like
6716    /// a peer response error, so callback-style consumers still run. If the
6717    /// response channel disappears for another reason, `handle` receives an
6718    /// outer `Err` describing that unexpected loss; there is no ack then.
6719    #[track_caller]
6720    fn consume_with<F>(
6721        self,
6722        handle: impl FnOnce(Result<Result<T, crate::Error>, crate::Error>) -> F + 'static + Send,
6723    ) -> Result<(), crate::Error>
6724    where
6725        T: 'static,
6726        F: Future<Output = Result<(), crate::Error>> + 'static + Send,
6727    {
6728        self.response_ordering.mark_ordered();
6729        let task_tx = self.task_tx.clone();
6730        Task::new(Location::caller(), self.handle_response(handle)).spawn(&task_tx)
6731    }
6732
6733    fn handle_response<F>(
6734        self,
6735        handle: impl FnOnce(Result<Result<T, crate::Error>, crate::Error>) -> F + 'static + Send,
6736    ) -> impl Future<Output = Result<(), crate::Error>> + Send
6737    where
6738        T: 'static,
6739        F: Future<Output = Result<(), crate::Error>> + 'static + Send,
6740    {
6741        let method = self.method;
6742        let response_rx = self.response_rx;
6743        let to_result = self.to_result;
6744        let cancellation = self.cancellation;
6745        let cancellation_sources = self.cancellation_sources;
6746        async move {
6747            let response = await_response_forwarding_cancellation(
6748                response_rx,
6749                &cancellation,
6750                &cancellation_sources,
6751            )
6752            .await;
6753
6754            match response {
6755                Ok(ResponsePayload { result, ack_tx }) => {
6756                    let typed_result = to_result(result);
6757
6758                    let outcome = handle(Ok(typed_result)).await;
6759
6760                    // Ack AFTER the handler completes - this is the key
6761                    // difference from block_task. The dispatch loop waits for
6762                    // this ack.
6763                    if let Some(tx) = ack_tx {
6764                        let _ = tx.send(());
6765                    }
6766
6767                    outcome
6768                }
6769                Err(err) => {
6770                    handle(Err(crate::util::internal_error(format!(
6771                        "response to `{method}` never received: {err}"
6772                    ))))
6773                    .await
6774                }
6775            }
6776        }
6777    }
6778
6779    /// Block the current task until the response is received.
6780    ///
6781    /// **Warning:** This method blocks the current async task. It is safe only when that task
6782    /// already runs outside the dispatch loop, such as the foreground future passed to
6783    /// `connect_with` or a task created with [`ConnectionTo::spawn`]. Using it directly in a
6784    /// handler callback will deadlock the connection.
6785    ///
6786    /// # Safe Usage (outside the dispatch loop)
6787    ///
6788    /// ```no_run
6789    /// # use agent_client_protocol_test::*;
6790    /// # async fn example() -> Result<(), agent_client_protocol::Error> {
6791    /// # let connection = mock_connection();
6792    /// connection.on_receive_request(async |req: MyRequest, responder, cx| {
6793    ///     // Spawn a task to handle the request
6794    ///     cx.spawn({
6795    ///         let connection = cx.clone();
6796    ///         async move {
6797    ///             // Safe: We're in a spawned task, not blocking the event loop
6798    ///             let response = connection.send_request(OtherRequest {})
6799    ///                 .block_task()
6800    ///                 .await?;
6801    ///
6802    ///             // Process the response...
6803    ///             Ok(())
6804    ///         }
6805    ///     })?;
6806    ///
6807    ///     // Respond immediately
6808    ///     responder.respond(MyResponse { status: "ok".into() })
6809    /// }, agent_client_protocol::on_receive_request!())
6810    /// # .connect_to(agent_client_protocol_test::MockTransport).await?;
6811    /// # Ok(())
6812    /// # }
6813    /// ```
6814    ///
6815    /// # Unsafe Usage (in handlers - will deadlock!)
6816    ///
6817    /// ```no_run
6818    /// # use agent_client_protocol_test::*;
6819    /// # async fn example() -> Result<(), agent_client_protocol::Error> {
6820    /// # let connection = mock_connection();
6821    /// connection.on_receive_request(async |req: MyRequest, responder, cx| {
6822    ///     // ❌ DEADLOCK: Handler blocks event loop, which can't process the response
6823    ///     let response = cx.send_request(OtherRequest {})
6824    ///         .block_task()
6825    ///         .await?;
6826    ///
6827    ///     responder.respond(MyResponse { status: response.value })
6828    /// }, agent_client_protocol::on_receive_request!())
6829    /// # .connect_to(agent_client_protocol_test::MockTransport).await?;
6830    /// # Ok(())
6831    /// # }
6832    /// ```
6833    ///
6834    /// # When to Use
6835    ///
6836    /// Use this method when:
6837    /// - Your current task already runs outside the dispatch loop
6838    /// - You need the response value to proceed with your logic
6839    /// - Linear control flow is more natural than callbacks
6840    ///
6841    /// For handler callbacks, use [`on_receiving_result`](Self::on_receiving_result) instead.
6842    pub async fn block_task(self) -> Result<T, crate::Error> {
6843        let response = await_response_forwarding_cancellation(
6844            self.response_rx,
6845            &self.cancellation,
6846            &self.cancellation_sources,
6847        )
6848        .await;
6849
6850        match response {
6851            Ok(ResponsePayload { result, ack_tx }) => {
6852                // Blocking consumers ack before converting or returning the
6853                // value, so dispatch can continue while the caller processes it.
6854                if let Some(tx) = ack_tx {
6855                    let _ = tx.send(());
6856                }
6857                (self.to_result)(result)
6858            }
6859            Err(err) => Err(crate::util::internal_error(format!(
6860                "response to `{}` never received: {}",
6861                self.method, err
6862            ))),
6863        }
6864    }
6865
6866    /// Block the current task and transform the typed result before releasing
6867    /// the ordered-response barrier.
6868    ///
6869    /// Framework lifecycle code uses this when success transfers local state
6870    /// to the returned value while an error must drop that state before later
6871    /// messages from the same transport frame are dispatched. The synchronous
6872    /// transform must not wait for additional connection traffic.
6873    pub(crate) async fn block_task_with_ordered_result<U>(
6874        self,
6875        transform: impl FnOnce(Result<T, crate::Error>) -> Result<U, crate::Error>,
6876    ) -> Result<U, crate::Error> {
6877        let response = await_response_forwarding_cancellation(
6878            self.response_rx,
6879            &self.cancellation,
6880            &self.cancellation_sources,
6881        )
6882        .await;
6883
6884        let (result, ack_tx) = match response {
6885            Ok(ResponsePayload { result, ack_tx }) => ((self.to_result)(result), ack_tx),
6886            Err(error) => (
6887                Err(crate::util::internal_error(format!(
6888                    "response to `{}` never received: {error}",
6889                    self.method
6890                ))),
6891                None,
6892            ),
6893        };
6894
6895        let outcome = transform(result);
6896        if let Some(acknowledgment) = ack_tx {
6897            let _ = acknowledgment.send(());
6898        }
6899        outcome
6900    }
6901
6902    /// Schedule an async task to run when a successful response is received.
6903    ///
6904    /// This is a convenience wrapper around [`on_receiving_result`](Self::on_receiving_result)
6905    /// for the common pattern of forwarding errors to a request context while only processing
6906    /// successful responses.
6907    ///
6908    /// # Behavior
6909    ///
6910    /// - If the response is `Ok(value)`, your task receives the value and the request context
6911    /// - If the response is `Err(error)`, the error is automatically sent to `responder`
6912    ///   and your task is not called
6913    ///
6914    /// # Example: Chaining requests
6915    ///
6916    /// ```no_run
6917    /// # use agent_client_protocol_test::*;
6918    /// # async fn example() -> Result<(), agent_client_protocol::Error> {
6919    /// # let connection = mock_connection();
6920    /// connection.on_receive_request(async |req: ValidateRequest, responder, cx| {
6921    ///     // Send initial request
6922    ///     cx.send_request(ValidateRequest { data: req.data.clone() })
6923    ///         .on_receiving_ok_result(responder, async |validation, responder| {
6924    ///             // Only runs if validation succeeded
6925    ///             if validation.is_valid {
6926    ///                 // Respond to original request
6927    ///                 responder.respond(ValidateResponse { is_valid: true, error: None })
6928    ///             } else {
6929    ///                 responder.respond_with_error(agent_client_protocol::Error::internal_error().data("validation failed"))
6930    ///             }
6931    ///         })?;
6932    ///
6933    ///     Ok(())
6934    /// }, agent_client_protocol::on_receive_request!())
6935    /// # .connect_to(agent_client_protocol_test::MockTransport).await?;
6936    /// # Ok(())
6937    /// # }
6938    /// ```
6939    ///
6940    /// # Ordering
6941    ///
6942    /// Like [`on_receiving_result`](Self::on_receiving_result), response handling holds the
6943    /// dispatch loop through callback completion when ordered consumption is selected before a
6944    /// peer response is routed during its original dispatch. Pending-request failures delivered
6945    /// without an incoming response and delayed routes do not carry that barrier. The callback
6946    /// must not await later inbound traffic on the same connection. See the
6947    /// [`ordering`](crate::concepts::ordering) module for details.
6948    ///
6949    /// # When to Use
6950    ///
6951    /// Use this when:
6952    /// - You need to respond to a request based on another request's result
6953    /// - You want errors to automatically propagate to the request context
6954    /// - You only care about the success case
6955    ///
6956    /// For more control over error handling, use [`on_receiving_result`](Self::on_receiving_result).
6957    #[track_caller]
6958    pub fn on_receiving_ok_result<F, Q: Protocol>(
6959        self,
6960        responder: Responder<T, Q>,
6961        task: impl FnOnce(T, Responder<T, Q>) -> F + 'static + Send,
6962    ) -> Result<(), crate::Error>
6963    where
6964        F: Future<Output = Result<(), crate::Error>> + 'static + Send,
6965        T: JsonRpcResponse,
6966    {
6967        self.on_receiving_result(async move |result| match result {
6968            Ok(value) => task(value, responder).await,
6969            Err(err) => responder.respond_with_error(err),
6970        })
6971    }
6972
6973    /// Register an async callback to run when the response is received.
6974    ///
6975    /// This is the recommended way to select response handling from inside a handler because
6976    /// registration returns immediately. The response-consumption task waits concurrently for
6977    /// the response; once the response is dispatched, the ordered callback may hold the dispatch
6978    /// loop until it completes.
6979    ///
6980    /// # Example: Handle response in callback
6981    ///
6982    /// ```no_run
6983    /// # use agent_client_protocol_test::*;
6984    /// # async fn example() -> Result<(), agent_client_protocol::Error> {
6985    /// # let connection = mock_connection();
6986    /// connection.on_receive_request(async |req: MyRequest, responder, cx| {
6987    ///     // Send a request and schedule a callback for the response
6988    ///     cx.send_request(QueryRequest { id: 22 })
6989    ///         .on_receiving_result({
6990    ///             let connection = cx.clone();
6991    ///             async move |result| {
6992    ///                 match result {
6993    ///                     Ok(response) => {
6994    ///                         println!("Got response: {:?}", response);
6995    ///                         // Can send more messages here
6996    ///                         connection.send_notification(QueryComplete {})?;
6997    ///                         Ok(())
6998    ///                 }
6999    ///                     Err(error) => {
7000    ///                         eprintln!("Request failed: {}", error);
7001    ///                         Err(error)
7002    ///                     }
7003    ///                 }
7004    ///             }
7005    ///         })?;
7006    ///
7007    ///     // Handler continues immediately after registering the callback
7008    ///     responder.respond(MyResponse { status: "processing".into() })
7009    /// }, agent_client_protocol::on_receive_request!())
7010    /// # .connect_to(agent_client_protocol_test::MockTransport).await?;
7011    /// # Ok(())
7012    /// # }
7013    /// ```
7014    ///
7015    /// # Ordering
7016    ///
7017    /// When ordered consumption is selected before a peer response is routed during its original
7018    /// dispatch, the callback runs in a connection-managed task and the dispatch loop waits for
7019    /// it to complete before processing the next message.
7020    ///
7021    /// The barrier does not apply when the pending request is failed without an incoming response,
7022    /// such as on EOF. If the response was already routed, or an interceptor routes a retained
7023    /// [`ResponseRouter`] after its original dispatch, the callback still runs but cannot
7024    /// retroactively block messages that were already released.
7025    ///
7026    /// While the barrier is held, the callback must not await a later response, notification, or
7027    /// other inbound traffic on the same connection: that traffic cannot be dispatched until the
7028    /// callback completes. Spawn follow-up work with [`ConnectionTo::spawn`] and return, or use
7029    /// [`block_task`](Self::block_task) from a task already outside the dispatch loop.
7030    ///
7031    /// This differs from [`block_task`](Self::block_task), which does not select ordered
7032    /// consumption: dispatch remains free while the caller processes the delivered response.
7033    ///
7034    /// See the [`ordering`](crate::concepts::ordering) module for details on ordering guarantees
7035    /// and how to avoid deadlocks.
7036    ///
7037    /// # Error Handling
7038    ///
7039    /// If the scheduled task returns `Err`, the entire server will shut down. Make sure to handle
7040    /// errors appropriately within your task.
7041    ///
7042    /// # When to Use
7043    ///
7044    /// Use this method when:
7045    /// - You need to register response handling from a handler callback
7046    /// - You want a peer response callback to complete before later messages are dispatched
7047    /// - The callback performs bounded work that does not depend on later inbound traffic
7048    ///
7049    /// When already outside the dispatch loop and you do not need ordering guarantees, consider
7050    /// [`block_task`](Self::block_task).
7051    #[track_caller]
7052    pub fn on_receiving_result<F>(
7053        self,
7054        task: impl FnOnce(Result<T, crate::Error>) -> F + 'static + Send,
7055    ) -> Result<(), crate::Error>
7056    where
7057        T: 'static,
7058        F: Future<Output = Result<(), crate::Error>> + 'static + Send,
7059    {
7060        self.consume_with(move |response| match response {
7061            // Invoke the callback before constructing its future so the
7062            // response value does not need to be `Send` across an await.
7063            Ok(result) => Either::Left(task(result)),
7064            // A response that was never delivered fails the consuming
7065            // task instead of invoking the callback.
7066            Err(err) => Either::Right(future::ready(Err(err))),
7067        })
7068    }
7069}
7070
7071// ============================================================================
7072// IntoJrConnectionTransport Implementations
7073// ============================================================================
7074
7075/// A component that communicates over line streams.
7076///
7077/// `Lines` implements the [`ConnectTo`] trait for any pair of line-based streams
7078/// (a `Stream<Item = io::Result<String>>` for incoming and a `Sink<String>` for outgoing),
7079/// handling serialization of JSON-RPC messages to/from newline-delimited JSON.
7080/// Each line carries one JSON-RPC message. JSON-RPC batch arrays are not supported: an
7081/// incoming array is answered with a single invalid-request (`-32600`) error with `id: null`.
7082///
7083/// This is a lower-level primitive than [`ByteStreams`] that enables interception and
7084/// transformation of individual lines before they are parsed or after they are serialized.
7085/// This is particularly useful for debugging, logging, or implementing custom line-based
7086/// protocols.
7087///
7088/// # Use Cases
7089///
7090/// - **Line-by-line logging**: Intercept and log each line before parsing
7091/// - **Custom protocols**: Transform lines before/after JSON-RPC processing
7092/// - **Debugging**: Inspect raw message strings
7093/// - **Line filtering**: Skip or modify specific messages
7094///
7095/// Most users should use [`ByteStreams`] instead, which provides a simpler interface
7096/// for byte-based I/O.
7097///
7098/// [`ConnectTo`]: crate::ConnectTo
7099#[derive(Debug)]
7100pub struct Lines<OutgoingSink, IncomingStream> {
7101    outgoing: OutgoingSink,
7102    incoming: IncomingStream,
7103}
7104
7105impl<OutgoingSink, IncomingStream> Lines<OutgoingSink, IncomingStream>
7106where
7107    OutgoingSink: futures::Sink<String, Error = std::io::Error> + Send + 'static,
7108    IncomingStream: futures::Stream<Item = std::io::Result<String>> + Send + 'static,
7109{
7110    /// Create a new line stream transport.
7111    pub fn new(outgoing: OutgoingSink, incoming: IncomingStream) -> Self {
7112        Self { outgoing, incoming }
7113    }
7114
7115    fn into_channel_transport(self) -> (Channel, crate::ConnectionDriver) {
7116        let Self { outgoing, incoming } = self;
7117        let (channel_for_caller, channel_for_lines) = Channel::duplex();
7118        let Channel { mut rx, tx } = channel_for_lines;
7119        let (finish_tx, finish_rx) = oneshot::channel();
7120        let finish = async move {
7121            // Losing a finish handle is not a shutdown request.
7122            if finish_rx.await.is_err() {
7123                future::pending::<()>().await;
7124            }
7125        }
7126        .boxed()
7127        .shared();
7128        let outgoing_frames = futures::stream::poll_fn({
7129            let mut finish = finish.clone();
7130            let mut finishing = false;
7131            move |cx| {
7132                if !finishing && std::pin::Pin::new(&mut finish).poll(cx).is_ready() {
7133                    rx.close();
7134                    finishing = true;
7135                }
7136                rx.poll_next_unpin(cx)
7137            }
7138        });
7139        let discard_incoming = Arc::new(AtomicBool::new(false));
7140        let incoming = incoming.filter_map({
7141            let discard_incoming = discard_incoming.clone();
7142            move |item| {
7143                let discard = discard_incoming.load(Ordering::Acquire);
7144                future::ready((!discard || item.is_err()).then_some(item))
7145            }
7146        });
7147        let outgoing = transport_actor::transport_outgoing_lines_actor(outgoing_frames, outgoing)
7148            .boxed()
7149            .shared();
7150        let serve_self = Box::pin({
7151            let outgoing = outgoing.clone();
7152            async move {
7153                futures::try_join!(
7154                    outgoing,
7155                    transport_actor::transport_incoming_lines_actor(incoming, tx),
7156                )?;
7157                Ok(())
7158            }
7159        });
7160        let server_future = crate::ConnectionDriver::with_finish(
7161            async move {
7162                match future::select(finish, serve_self).await {
7163                    Either::Left(((), serve_self)) => {
7164                        discard_incoming.store(true, Ordering::Release);
7165                        // Keep reading while flushing, but do not require remote
7166                        // read EOF. Poll incoming errors before clean sink drain.
7167                        match future::select(serve_self, outgoing).await {
7168                            Either::Left((result, _)) | Either::Right((result, _)) => result,
7169                        }
7170                    }
7171                    Either::Right((result, _)) => result,
7172                }
7173            },
7174            move || {
7175                let _ = finish_tx.send(());
7176            },
7177        );
7178
7179        (channel_for_caller, server_future)
7180    }
7181}
7182
7183impl<OutgoingSink, IncomingStream, R: Role> ConnectTo<R> for Lines<OutgoingSink, IncomingStream>
7184where
7185    OutgoingSink: futures::Sink<String, Error = std::io::Error> + Send + 'static,
7186    IncomingStream: futures::Stream<Item = std::io::Result<String>> + Send + 'static,
7187{
7188    async fn connect_to(self, client: impl ConnectTo<R::Counterpart>) -> Result<(), crate::Error> {
7189        let (channel, mut serve_self) = self.into_channel_transport();
7190        let mut finish = serve_self
7191            .take_finish()
7192            .expect("built-in Lines transport supports explicit finishing");
7193        let client_future = Box::pin(ConnectTo::<R>::connect_to(channel, client));
7194
7195        match futures::future::select(client_future, serve_self).await {
7196            Either::Left((result, serve_self)) => {
7197                result?;
7198                // The local bridge has transferred all accepted client output.
7199                // Finish the physical sink without waiting for remote read EOF.
7200                finish.request();
7201                serve_self.await
7202            }
7203            Either::Right((result, _)) => result,
7204        }
7205    }
7206
7207    fn into_channel_and_future(self) -> (Channel, Option<crate::ConnectionDriver>) {
7208        let (channel, driver) = self.into_channel_transport();
7209        (channel, Some(driver))
7210    }
7211}
7212
7213/// A component that communicates over byte streams (stdin/stdout, sockets, pipes, etc.).
7214///
7215/// `ByteStreams` implements the [`ConnectTo`] trait for any pair of `AsyncRead` and `AsyncWrite`
7216/// streams, handling serialization of JSON-RPC messages to/from newline-delimited JSON.
7217/// This is the standard way to communicate with external processes or network connections.
7218///
7219/// # Use Cases
7220///
7221/// - **Stdio communication**: Connect to agents or proxies via stdin/stdout
7222/// - **Network sockets**: TCP, Unix domain sockets, or other stream-based protocols
7223/// - **Named pipes**: Cross-process communication on the same machine
7224/// - **File I/O**: Reading from and writing to file descriptors
7225///
7226/// # Example
7227///
7228/// Connecting to an agent via stdio:
7229///
7230/// ```no_run
7231/// use agent_client_protocol::UntypedRole;
7232/// # use agent_client_protocol::{ByteStreams};
7233/// use tokio_util::compat::{TokioAsyncReadCompatExt, TokioAsyncWriteCompatExt};
7234///
7235/// # async fn example() -> Result<(), agent_client_protocol::Error> {
7236/// let component = ByteStreams::new(
7237///     tokio::io::stdout().compat_write(),
7238///     tokio::io::stdin().compat(),
7239/// );
7240///
7241/// // Use as a component in a connection
7242/// agent_client_protocol::UntypedRole.builder()
7243///     .name("my-client")
7244///     .connect_to(component)
7245///     .await?;
7246/// # Ok(())
7247/// # }
7248/// ```
7249///
7250/// [`ConnectTo`]: crate::ConnectTo
7251#[derive(Debug)]
7252pub struct ByteStreams<OB, IB> {
7253    outgoing: OB,
7254    incoming: IB,
7255}
7256
7257impl<OB, IB> ByteStreams<OB, IB>
7258where
7259    OB: AsyncWrite + Send + 'static,
7260    IB: AsyncRead + Send + 'static,
7261{
7262    /// Create a new byte stream transport.
7263    pub fn new(outgoing: OB, incoming: IB) -> Self {
7264        Self { outgoing, incoming }
7265    }
7266
7267    fn into_lines(
7268        self,
7269    ) -> Lines<
7270        impl futures::Sink<String, Error = std::io::Error> + Send + 'static,
7271        impl futures::Stream<Item = std::io::Result<String>> + Send + 'static,
7272    > {
7273        use futures::AsyncBufReadExt;
7274        use futures::io::BufReader;
7275        let Self { outgoing, incoming } = self;
7276
7277        let incoming_lines = Box::pin(BufReader::new(incoming).lines());
7278        let outgoing_lines = transport_actor::LineWriter::new(outgoing);
7279
7280        Lines::new(outgoing_lines, incoming_lines)
7281    }
7282}
7283
7284#[cfg(any(
7285    all(
7286        any(feature = "process", feature = "stdio"),
7287        not(target_family = "wasm")
7288    ),
7289    test
7290))]
7291pub(crate) async fn write_line<W>(writer: &mut W, line: String) -> std::io::Result<()>
7292where
7293    W: AsyncWrite + Unpin + ?Sized,
7294{
7295    use futures::AsyncWriteExt as _;
7296
7297    let mut bytes = line.into_bytes();
7298    bytes.push(b'\n');
7299    writer.write_all(&bytes).await?;
7300    writer.flush().await
7301}
7302
7303impl<OB, IB, R: Role> ConnectTo<R> for ByteStreams<OB, IB>
7304where
7305    OB: AsyncWrite + Send + 'static,
7306    IB: AsyncRead + Send + 'static,
7307{
7308    async fn connect_to(self, client: impl ConnectTo<R::Counterpart>) -> Result<(), crate::Error> {
7309        ConnectTo::<R>::connect_to(self.into_lines(), client).await
7310    }
7311
7312    fn into_channel_and_future(self) -> (Channel, Option<crate::ConnectionDriver>) {
7313        ConnectTo::<R>::into_channel_and_future(self.into_lines())
7314    }
7315}
7316
7317/// A channel endpoint representing one side of a bidirectional JSON-RPC transport.
7318///
7319/// A channel carries complete TransportFrame values across in-process components and
7320/// transport adapters. Malformed wire input is an
7321/// explicit frame; failures while driving a physical transport are returned by that
7322/// transport's future.
7323///
7324/// The unbounded [`futures::channel::mpsc`] pair is the intended contract:
7325/// frames are never dropped and sends never wait, so the transport or
7326/// component on each side applies its own backpressure. Build a channel from
7327/// an existing pair with [`Channel::new`] and take it apart with
7328/// [`Channel::into_parts`] or by destructuring with `..`.
7329///
7330/// # Example
7331///
7332/// ```no_run
7333/// # use agent_client_protocol::UntypedRole;
7334/// # use agent_client_protocol::Channel;
7335/// # async fn example() -> Result<(), agent_client_protocol::Error> {
7336/// let (channel_a, _channel_b) = Channel::duplex();
7337///
7338/// UntypedRole.builder()
7339///     .name("connection-a")
7340///     .connect_to(channel_a)
7341///     .await?;
7342/// # Ok(())
7343/// # }
7344/// ```
7345#[derive(Debug)]
7346#[non_exhaustive]
7347pub struct Channel {
7348    /// Receives frames from the counterpart.
7349    pub rx: mpsc::UnboundedReceiver<TransportFrame>,
7350    /// Sends frames to the counterpart.
7351    pub tx: mpsc::UnboundedSender<TransportFrame>,
7352}
7353
7354impl Channel {
7355    /// Create a channel endpoint that receives frames on `rx` and sends them
7356    /// on `tx`.
7357    #[must_use]
7358    pub fn new(
7359        rx: mpsc::UnboundedReceiver<TransportFrame>,
7360        tx: mpsc::UnboundedSender<TransportFrame>,
7361    ) -> Self {
7362        Self { rx, tx }
7363    }
7364
7365    /// Split this endpoint into its receiving and sending halves, `(rx, tx)`.
7366    #[must_use]
7367    pub fn into_parts(
7368        self,
7369    ) -> (
7370        mpsc::UnboundedReceiver<TransportFrame>,
7371        mpsc::UnboundedSender<TransportFrame>,
7372    ) {
7373        (self.rx, self.tx)
7374    }
7375
7376    /// Create a pair of connected channel endpoints.
7377    ///
7378    /// Frames sent through either endpoint are received by the other endpoint.
7379    #[must_use]
7380    pub fn duplex() -> (Self, Self) {
7381        let (a_tx, b_rx) = mpsc::unbounded();
7382        let (b_tx, a_rx) = mpsc::unbounded();
7383
7384        (Self { rx: a_rx, tx: a_tx }, Self { rx: b_rx, tx: b_tx })
7385    }
7386
7387    /// Copy frames from `rx` to `tx` until the input closes.
7388    ///
7389    /// # Errors
7390    ///
7391    /// Returns an error if the receiving endpoint closes before the input.
7392    pub(crate) async fn copy(mut self) -> Result<(), crate::Error> {
7393        while let Some(frame) = self.rx.next().await {
7394            self.tx
7395                .unbounded_send(frame)
7396                .map_err(crate::util::internal_error)?;
7397        }
7398        Ok(())
7399    }
7400
7401    /// Copy output concurrently with its owning driver, then drain accepted frames.
7402    /// Passive endpoints instead retain the channel's independent half-close lifetime.
7403    pub(crate) async fn copy_with_driver(
7404        self,
7405        driver: Option<crate::ConnectionDriver>,
7406    ) -> Result<(), crate::Error> {
7407        self.copy_with_driver_until(driver, future::pending()).await
7408    }
7409
7410    /// After the destination's owned foreground finishes, keep driving source
7411    /// errors and sink work, but never deliver queued or new input to it.
7412    pub(crate) async fn copy_with_driver_until(
7413        mut self,
7414        mut driver: Option<crate::ConnectionDriver>,
7415        stop_delivery: impl Future<Output = ()>,
7416    ) -> Result<(), crate::Error> {
7417        let mut stop_delivery = pin!(stop_delivery);
7418        let mut delivering = true;
7419        let mut done = false;
7420        loop {
7421            let event = future::poll_fn(|cx| {
7422                if delivering && stop_delivery.as_mut().poll(cx).is_ready() {
7423                    delivering = false;
7424                }
7425                // Driver errors remain authoritative even when stop or EOF is ready.
7426                if !done
7427                    && let Some(driver) = driver.as_mut()
7428                    && let std::task::Poll::Ready(result) = std::pin::Pin::new(driver).poll(cx)
7429                {
7430                    return std::task::Poll::Ready(Either::Left(result));
7431                }
7432                if !delivering && driver.is_none() {
7433                    return std::task::Poll::Ready(Either::Right(None));
7434                }
7435                self.rx.poll_next_unpin(cx).map(Either::Right)
7436            })
7437            .await;
7438            let frame = match event {
7439                Either::Left(result) => {
7440                    result?;
7441                    done = true;
7442                    self.rx.close();
7443                    continue;
7444                }
7445                Either::Right(frame) => frame,
7446            };
7447            let Some(frame) = frame else {
7448                break;
7449            };
7450            if delivering {
7451                self.tx
7452                    .unbounded_send(frame)
7453                    .map_err(crate::util::internal_error)?;
7454            }
7455        }
7456        // Propagate this half-close before waiting for a still-running driver.
7457        drop(self);
7458        if !done && let Some(driver) = driver {
7459            driver.await?;
7460        }
7461        Ok(())
7462    }
7463
7464    /// Bridge two endpoints while inspecting every valid message.
7465    ///
7466    /// Observers are invoked in source order for each valid message. The
7467    /// original frame is forwarded unchanged after inspection.
7468    ///
7469    /// # Errors
7470    ///
7471    /// Returns an observer error or an error if a destination closes before its
7472    /// source.
7473    pub async fn bridge_with_inspection(
7474        left: Self,
7475        right: Self,
7476        mut left_to_right: impl FnMut(&RawJsonRpcMessage) -> Result<(), crate::Error> + Send,
7477        mut right_to_left: impl FnMut(&RawJsonRpcMessage) -> Result<(), crate::Error> + Send,
7478    ) -> Result<(), crate::Error> {
7479        let Self {
7480            rx: mut left_rx,
7481            tx: left_tx,
7482        } = left;
7483        let Self {
7484            rx: mut right_rx,
7485            tx: right_tx,
7486        } = right;
7487
7488        let left_to_right = async move {
7489            while let Some(frame) = left_rx.next().await {
7490                frame.inspect_messages(&mut left_to_right)?;
7491                right_tx
7492                    .unbounded_send(frame)
7493                    .map_err(crate::util::internal_error)?;
7494            }
7495            Ok::<(), crate::Error>(())
7496        };
7497        let right_to_left = async move {
7498            while let Some(frame) = right_rx.next().await {
7499                frame.inspect_messages(&mut right_to_left)?;
7500                left_tx
7501                    .unbounded_send(frame)
7502                    .map_err(crate::util::internal_error)?;
7503            }
7504            Ok::<(), crate::Error>(())
7505        };
7506
7507        futures::try_join!(left_to_right, right_to_left)?;
7508        Ok(())
7509    }
7510}
7511
7512impl<R: Role> ConnectTo<R> for Channel {
7513    async fn connect_to(self, client: impl ConnectTo<R::Counterpart>) -> Result<(), crate::Error> {
7514        let (client_channel, client_future) = client.into_channel_and_future();
7515
7516        let passive = client_future.is_none();
7517        let outgoing = Box::pin(
7518            Channel {
7519                rx: client_channel.rx,
7520                tx: self.tx,
7521            }
7522            .copy_with_driver(client_future),
7523        );
7524        let incoming = Box::pin(
7525            Channel {
7526                rx: self.rx,
7527                tx: client_channel.tx,
7528            }
7529            .copy(),
7530        );
7531        if passive {
7532            futures::try_join!(outgoing, incoming)?;
7533            return Ok(());
7534        }
7535
7536        match future::select(outgoing, incoming).await {
7537            Either::Left((result, _)) => result,
7538            Either::Right((result, outgoing)) => {
7539                result?;
7540                outgoing.await
7541            }
7542        }
7543    }
7544
7545    fn into_channel_and_future(self) -> (Channel, Option<crate::ConnectionDriver>) {
7546        (self, None)
7547    }
7548}
7549
7550#[cfg(test)]
7551mod tests {
7552    use super::*;
7553    use crate::Client;
7554    #[cfg(feature = "unstable_protocol_v2")]
7555    use crate::Proxy;
7556
7557    #[test]
7558    fn raw_notification_round_trips_and_drops_unknown_envelope_members() {
7559        let message: RawJsonRpcMessage = serde_json::from_value(serde_json::json!({
7560            "jsonrpc": "2.0",
7561            "method": "_vendor/ping",
7562            "params": { "n": 1 },
7563            "extra": true,
7564        }))
7565        .unwrap();
7566        let mut params = serde_json::Map::new();
7567        params.insert("n".into(), 1.into());
7568        assert_eq!(
7569            message,
7570            RawJsonRpcMessage::Notification(RawJsonRpcNotification {
7571                method: "_vendor/ping".into(),
7572                params: Some(RawJsonRpcParams::Object(params)),
7573            })
7574        );
7575        assert_eq!(
7576            serde_json::to_value(&message).unwrap(),
7577            serde_json::json!({
7578                "jsonrpc": "2.0",
7579                "method": "_vendor/ping",
7580                "params": { "n": 1 },
7581            })
7582        );
7583
7584        let without_params =
7585            RawJsonRpcMessage::notification("ping".into(), serde_json::Value::Null).unwrap();
7586        assert_eq!(
7587            serde_json::to_value(&without_params).unwrap(),
7588            serde_json::json!({ "jsonrpc": "2.0", "method": "ping" })
7589        );
7590    }
7591
7592    #[test]
7593    fn protected_cleanup_keeps_scoped_runners_polled_on_every_shutdown_path() {
7594        #[derive(Clone, Copy, Debug)]
7595        enum Stop {
7596            ForegroundSuccess,
7597            ForegroundError,
7598            InputEof,
7599            TransportError,
7600            TaskError,
7601            RunnerError,
7602            SupervisorError,
7603        }
7604
7605        struct Dropped(Arc<AtomicBool>);
7606        impl Drop for Dropped {
7607            fn drop(&mut self) {
7608                self.0.store(true, Ordering::Release);
7609            }
7610        }
7611
7612        for stop in [
7613            Stop::ForegroundSuccess,
7614            Stop::ForegroundError,
7615            Stop::InputEof,
7616            Stop::TransportError,
7617            Stop::TaskError,
7618            Stop::RunnerError,
7619            Stop::SupervisorError,
7620        ] {
7621            let cleaned = Arc::new(AtomicBool::new(false));
7622            let disposable_dropped = Arc::new(AtomicBool::new(false));
7623            let close_finished = Arc::new(AtomicBool::new(false));
7624            let (cleanup_tx, cleanup_rx) = oneshot::channel::<()>();
7625            let (scoped_done_tx, scoped_done_rx) = completion_signal();
7626            let (stop_tx, stop_rx) = oneshot::channel::<()>();
7627            let stop_signal = stop_rx.map(|_| ()).boxed().shared();
7628            let (incoming_tx, incoming_rx) = mpsc::unbounded();
7629            let outgoing = futures::sink::unfold((), |(), _line: String| {
7630                future::ready(Ok::<_, std::io::Error>(()))
7631            });
7632            let builder = Client
7633                .v1()
7634                .with_spawned({
7635                    let cleaned = cleaned.clone();
7636                    async move |cx: crate::v1::ConnectionTo<Agent>| {
7637                        cx.raw_connection().shutdown_requested().await;
7638                        // This stands in for the actual scoped native operation:
7639                        // its async cleanup only advances if this runner is polled.
7640                        cleanup_rx.await.unwrap();
7641                        cleaned.store(true, Ordering::Release);
7642                        let _ = scoped_done_tx.send(());
7643                        Ok(())
7644                    }
7645                })
7646                .with_spawned({
7647                    let stop_signal = stop_signal.clone();
7648                    async move |_cx| {
7649                        stop_signal.await;
7650                        if matches!(stop, Stop::RunnerError) {
7651                            Err(crate::Error::internal_error().data("runner failure"))
7652                        } else {
7653                            future::pending().await
7654                        }
7655                    }
7656                })
7657                .on_close({
7658                    let close_finished = close_finished.clone();
7659                    let scoped_done = scoped_done_rx.clone();
7660                    async move |cx: crate::v1::ConnectionTo<Agent>| {
7661                        // EOF cancellation must precede, not await, close callbacks.
7662                        cx.raw_connection().shutdown_requested().await;
7663                        assert!(!cx.is_incoming_closed());
7664                        scoped_done.await;
7665                        close_finished.store(true, Ordering::Release);
7666                        Ok(())
7667                    }
7668                });
7669            let (connection, driver) =
7670                builder.into_connection_and_future(Lines::new(outgoing, incoming_rx), false, {
7671                    let stop_signal = stop_signal.clone();
7672                    async move |cx| {
7673                        if matches!(stop, Stop::InputEof) {
7674                            cx.incoming_closed().await;
7675                            return Ok(());
7676                        }
7677                        stop_signal.await;
7678                        match stop {
7679                            Stop::ForegroundSuccess | Stop::SupervisorError => Ok(()),
7680                            Stop::ForegroundError => {
7681                                Err(crate::Error::internal_error().data("foreground failure"))
7682                            }
7683                            _ => future::pending().await,
7684                        }
7685                    }
7686                });
7687            let disposable = Dropped(disposable_dropped.clone());
7688            connection
7689                .spawn(async move {
7690                    let _disposable = disposable;
7691                    future::pending().await
7692                })
7693                .unwrap();
7694            connection
7695                .spawn({
7696                    let stop_signal = stop_signal.clone();
7697                    async move {
7698                        stop_signal.await;
7699                        if matches!(stop, Stop::TaskError) {
7700                            Err(crate::Error::internal_error().data("task failure"))
7701                        } else {
7702                            future::pending().await
7703                        }
7704                    }
7705                })
7706                .unwrap();
7707            connection
7708                .spawn_protected({
7709                    let connection = connection.clone();
7710                    async move {
7711                        connection.shutdown_requested().await;
7712                        scoped_done_rx.await;
7713                        if matches!(stop, Stop::SupervisorError) {
7714                            Err(crate::Error::internal_error().data("supervisor failure"))
7715                        } else {
7716                            Ok(())
7717                        }
7718                    }
7719                })
7720                .unwrap();
7721            let mut driver = Box::pin(driver);
7722            assert!(driver.as_mut().now_or_never().is_none(), "{stop:?}");
7723            assert!(connection.shutdown_requested().now_or_never().is_none());
7724            let _ = stop_tx.send(());
7725            let incoming_tx = match stop {
7726                Stop::InputEof => {
7727                    drop(incoming_tx);
7728                    None
7729                }
7730                Stop::TransportError => {
7731                    incoming_tx
7732                        .unbounded_send(Err(std::io::Error::other("transport failure")))
7733                        .unwrap();
7734                    Some(incoming_tx)
7735                }
7736                _ => Some(incoming_tx),
7737            };
7738            for _ in 0..10 {
7739                assert!(driver.as_mut().now_or_never().is_none(), "{stop:?}");
7740                if connection.shutdown_requested().now_or_never().is_some() {
7741                    break;
7742                }
7743            }
7744            assert!(
7745                connection.shutdown_requested().now_or_never().is_some(),
7746                "{stop:?}"
7747            );
7748            assert!(!cleaned.load(Ordering::Acquire), "{stop:?}");
7749            assert!(!disposable_dropped.load(Ordering::Acquire), "{stop:?}");
7750            cleanup_tx.send(()).unwrap();
7751            // Task acknowledgments may wake an actor already polled in this turn.
7752            // Bound the probe so a broken scoped-runner join fails, not hangs.
7753            let mut result = None;
7754            for _ in 0..10 {
7755                result = driver.as_mut().now_or_never();
7756                if result.is_some() {
7757                    break;
7758                }
7759            }
7760            let result =
7761                result.unwrap_or_else(|| panic!("driver did not finish owned cleanup: {stop:?}"));
7762            match stop {
7763                Stop::ForegroundSuccess | Stop::InputEof => result.unwrap(),
7764                _ => {
7765                    let error = result.expect_err("shutdown must preserve the first error");
7766                    let expected = match stop {
7767                        Stop::ForegroundError => "foreground failure",
7768                        Stop::TransportError => "transport failure",
7769                        Stop::TaskError => "task failure",
7770                        Stop::RunnerError => "runner failure",
7771                        Stop::SupervisorError => "supervisor failure",
7772                        _ => unreachable!(),
7773                    };
7774                    assert!(
7775                        error.data.unwrap().to_string().contains(expected),
7776                        "{stop:?}"
7777                    );
7778                }
7779            }
7780            assert!(cleaned.load(Ordering::Acquire), "{stop:?}");
7781            assert!(disposable_dropped.load(Ordering::Acquire), "{stop:?}");
7782            assert_eq!(
7783                close_finished.load(Ordering::Acquire),
7784                matches!(stop, Stop::InputEof),
7785                "{stop:?}",
7786            );
7787            assert!(connection.spawn_protected(async { Ok(()) }).is_err());
7788            drop(incoming_tx);
7789        }
7790    }
7791
7792    #[test]
7793    fn protected_operation_acknowledgments_are_reaped_and_join_seals_registration() {
7794        let (connection, _message_rx, _pending_replies) = connection_for_response_hook_tests();
7795        // The helper drops its task receiver, so use a live receiver for this probe.
7796        let (task_tx, mut task_rx) = mpsc::unbounded();
7797        let connection = ConnectionTo {
7798            task_tx,
7799            ..connection
7800        };
7801        for _ in 0..100 {
7802            connection.spawn_protected(async { Ok(()) }).unwrap();
7803            assert_eq!(
7804                connection
7805                    .protected_operations
7806                    .lock()
7807                    .unwrap()
7808                    .pending
7809                    .len(),
7810                1
7811            );
7812            let task = task_rx.next().now_or_never().unwrap().unwrap();
7813            futures::executor::block_on(task.run_for_test()).unwrap();
7814        }
7815        assert!(
7816            connection
7817                .wait_protected_operations()
7818                .now_or_never()
7819                .is_some()
7820        );
7821        assert!(
7822            connection
7823                .wait_protected_operations()
7824                .now_or_never()
7825                .is_some()
7826        );
7827        assert!(
7828            connection
7829                .protected_operations
7830                .lock()
7831                .unwrap()
7832                .pending
7833                .is_empty()
7834        );
7835        assert!(connection.spawn_protected(async { Ok(()) }).is_err());
7836        assert!(task_rx.next().now_or_never().is_none());
7837    }
7838
7839    #[test]
7840    fn dropping_unused_finish_signal_preserves_physical_half_closes() {
7841        let outgoing = futures::sink::unfold((), |(), _line: String| {
7842            future::ready(Ok::<_, std::io::Error>(()))
7843        });
7844        let (incoming_tx, incoming_rx) = mpsc::unbounded();
7845        let (Channel { mut rx, tx }, mut driver) =
7846            Lines::new(outgoing, incoming_rx).into_channel_transport();
7847
7848        drop(
7849            driver
7850                .take_finish()
7851                .expect("built-in Lines driver is finishable"),
7852        );
7853        drop(tx);
7854        assert!((&mut driver).now_or_never().is_none());
7855        incoming_tx
7856            .unbounded_send(Ok(
7857                r#"{"jsonrpc":"2.0","method":"test/after-output-eof"}"#.into()
7858            ))
7859            .unwrap();
7860        assert!((&mut driver).now_or_never().is_none());
7861        assert!(rx.next().now_or_never().unwrap().is_some());
7862
7863        drop(incoming_tx);
7864        futures::executor::block_on(driver).unwrap();
7865        assert!(rx.next().now_or_never().unwrap().is_none());
7866    }
7867
7868    #[test]
7869    fn explicit_physical_finish_does_not_hide_a_ready_read_error() {
7870        let outgoing = futures::sink::unfold((), |(), _line: String| {
7871            future::ready(Ok::<_, std::io::Error>(()))
7872        });
7873        let incoming = futures::stream::iter([Err(std::io::Error::other("finish read failed"))]);
7874        let (_channel, mut driver) = Lines::new(outgoing, incoming).into_channel_transport();
7875        assert!(driver.request_finish());
7876
7877        let error = futures::executor::block_on(driver).unwrap_err();
7878        assert_eq!(
7879            error
7880                .data
7881                .and_then(|value| value.as_str().map(str::to_owned)),
7882            Some("finish read failed".into())
7883        );
7884    }
7885
7886    #[cfg(feature = "unstable_protocol_v2")]
7887    fn connection_with_task_receiver() -> (
7888        ConnectionTo<crate::role::UntypedRole>,
7889        mpsc::UnboundedReceiver<Task>,
7890    ) {
7891        let (message_tx, _message_rx) = mpsc::unbounded();
7892        let (task_tx, task_rx) = mpsc::unbounded();
7893        let (dynamic_handler_tx, _dynamic_handler_rx) = mpsc::unbounded();
7894        let transport_completion: SharedTransportCompletion =
7895            future::ready(Ok::<(), crate::Error>(())).boxed().shared();
7896        let pending_replies = PendingReplies::default();
7897
7898        (
7899            ConnectionTo::new(
7900                crate::role::UntypedRole,
7901                message_tx,
7902                task_tx,
7903                dynamic_handler_tx,
7904                transport_completion,
7905                pending_replies.registrar(),
7906                ProtocolMode::disabled(),
7907            ),
7908            task_rx,
7909        )
7910    }
7911
7912    #[cfg(feature = "unstable_protocol_v2")]
7913    #[test]
7914    fn v2_builder_exposes_typed_context_to_user_callbacks() {
7915        fn assert_v2_context(_connection: &crate::v2::ConnectionTo<Agent>) {}
7916
7917        let _builder = Client
7918            .v2()
7919            .on_receive_request(
7920                async |_request: UntypedMessage, _responder, connection| {
7921                    assert_v2_context(&connection);
7922                    Ok(())
7923                },
7924                crate::on_receive_request!(),
7925            )
7926            .on_receive_notification(
7927                async |_notification: UntypedMessage, connection| {
7928                    assert_v2_context(&connection);
7929                    Ok(())
7930                },
7931                crate::on_receive_notification!(),
7932            )
7933            .on_receive_dispatch(
7934                async |_dispatch: Dispatch<UntypedMessage, UntypedMessage, V2>, connection| {
7935                    assert_v2_context(&connection);
7936                    Ok(())
7937                },
7938                crate::on_receive_dispatch!(),
7939            )
7940            .on_receive_request_from(
7941                Agent,
7942                async |_request: UntypedMessage, _responder, connection| {
7943                    assert_v2_context(&connection);
7944                    Ok(())
7945                },
7946                crate::on_receive_request!(),
7947            )
7948            .on_receive_notification_from(
7949                Agent,
7950                async |_notification: UntypedMessage, connection| {
7951                    assert_v2_context(&connection);
7952                    Ok(())
7953                },
7954                crate::on_receive_notification!(),
7955            )
7956            .on_receive_dispatch_from(
7957                Agent,
7958                async |_dispatch: Dispatch<UntypedMessage, UntypedMessage, V2>, connection| {
7959                    assert_v2_context(&connection);
7960                    Ok(())
7961                },
7962                crate::on_receive_dispatch!(),
7963            )
7964            .with_spawned(async |connection| {
7965                assert_v2_context(&connection);
7966                Ok(())
7967            })
7968            .on_close(async |connection| {
7969                assert_v2_context(&connection);
7970                Ok(())
7971            });
7972    }
7973
7974    #[cfg(feature = "unstable_protocol_v2")]
7975    #[test]
7976    fn proxy_builders_negotiate_proxy_initialize_like_agents() -> Result<(), crate::Error> {
7977        use crate::schema::ProtocolVersion;
7978
7979        // A v1 proxy selects v1 for a newer request.
7980        let mode = Proxy.v1().protocol_mode;
7981        assert_eq!(mode.api_protocol_version(), Some(ProtocolVersion::V1));
7982        let request = ProtocolCompat::new(mode).incoming_message(UntypedMessage::new(
7983            "proxy/initialize",
7984            serde_json::json!({ "protocolVersion": ProtocolVersion::V2 }),
7985        )?)?;
7986        assert_eq!(request.params()["protocolVersion"], serde_json::json!(1));
7987
7988        // A v2 proxy supports no version at or below v1.
7989        let mode = Proxy.v2().protocol_mode;
7990        assert_eq!(mode.api_protocol_version(), Some(ProtocolVersion::V2));
7991        let error = ProtocolCompat::new(mode)
7992            .incoming_message(UntypedMessage::new(
7993                "proxy/initialize",
7994                serde_json::json!({ "protocolVersion": ProtocolVersion::V1 }),
7995            )?)
7996            .expect_err("a v2 proxy builder must reject v1");
7997        assert_eq!(i32::from(error.code), -32022);
7998
7999        Ok(())
8000    }
8001
8002    #[cfg(feature = "unstable_protocol_v2")]
8003    #[test]
8004    fn v2_proxy_rejects_explicitly_prewrapped_initialize_request() {
8005        let (message_tx, message_rx) = mpsc::unbounded();
8006        let (task_tx, _task_rx) = mpsc::unbounded();
8007        let (dynamic_handler_tx, _dynamic_handler_rx) = mpsc::unbounded();
8008        let transport_completion: SharedTransportCompletion =
8009            future::ready(Ok::<(), crate::Error>(())).boxed().shared();
8010        let pending_replies = PendingReplies::default();
8011        let connection = ConnectionTo::new(
8012            crate::Conductor,
8013            message_tx,
8014            task_tx,
8015            dynamic_handler_tx,
8016            transport_completion,
8017            pending_replies.registrar(),
8018            ProtocolMode::v2_proxy(),
8019        );
8020
8021        let request = crate::schema::SuccessorMessage::new(
8022            UntypedMessage::new(
8023                "initialize",
8024                serde_json::json!({ "protocolVersion": crate::schema::ProtocolVersion::V1 }),
8025            )
8026            .expect("test initialize request should serialize"),
8027        );
8028        let sent = connection.send_request_to(Agent, request);
8029
8030        let (transport_tx, mut transport_rx) = mpsc::unbounded();
8031        let mut actor = Box::pin(outgoing_actor::outgoing_protocol_actor(
8032            message_rx,
8033            pending_replies,
8034            transport_tx,
8035            ProtocolCompat::new(ProtocolMode::v2_proxy()),
8036            future::pending::<()>().boxed().shared(),
8037        ));
8038        assert!(
8039            actor.as_mut().now_or_never().is_none(),
8040            "the outgoing actor should continue after rejecting the request"
8041        );
8042        assert!(
8043            transport_rx.next().now_or_never().is_none(),
8044            "an explicitly prewrapped initialize must not reach the transport"
8045        );
8046
8047        let error = futures::executor::block_on(sent.block_task())
8048            .expect_err("connection routing must own successor wrapping");
8049        let data = error
8050            .data
8051            .as_ref()
8052            .and_then(|data| data.as_str())
8053            .unwrap_or_default();
8054        assert!(data.contains("logical `initialize`"), "{error:?}");
8055        assert!(data.contains("proxy/successor"), "{error:?}");
8056    }
8057
8058    #[cfg(feature = "unstable_protocol_v2")]
8059    #[test]
8060    fn v2_proxy_builder_exposes_typed_context_to_user_callbacks() {
8061        fn assert_v2_context(_connection: &crate::v2::ConnectionTo<crate::Conductor>) {}
8062
8063        let _builder = Proxy
8064            .v2()
8065            .on_receive_request_from(
8066                Client,
8067                async |_request: UntypedMessage, _responder, connection| {
8068                    assert_v2_context(&connection);
8069                    Ok(())
8070                },
8071                crate::on_receive_request!(),
8072            )
8073            .on_receive_notification_from(
8074                Agent,
8075                async |_notification: UntypedMessage, connection| {
8076                    assert_v2_context(&connection);
8077                    Ok(())
8078                },
8079                crate::on_receive_notification!(),
8080            )
8081            .on_receive_dispatch_from(
8082                Client,
8083                async |_dispatch: Dispatch<UntypedMessage, UntypedMessage, V2>, connection| {
8084                    assert_v2_context(&connection);
8085                    Ok(())
8086                },
8087                crate::on_receive_dispatch!(),
8088            )
8089            .with_spawned(async |connection| {
8090                assert_v2_context(&connection);
8091                Ok(())
8092            })
8093            .on_close(async |connection| {
8094                assert_v2_context(&connection);
8095                Ok(())
8096            });
8097    }
8098
8099    #[cfg(feature = "unstable_protocol_v2")]
8100    #[test]
8101    fn raw_connection_spawns_v2_builder_with_typed_child_callback() {
8102        let (parent, mut task_rx) = connection_with_task_receiver();
8103        let (transport, _peer) = Channel::duplex();
8104        let (callback_tx, callback_rx) = oneshot::channel();
8105
8106        let child: crate::v2::ConnectionTo<Agent> = parent
8107            .spawn_connection(
8108                Client.v2().with_spawned(
8109                    async move |_connection: crate::v2::ConnectionTo<Agent>| {
8110                        callback_tx.send(()).map_err(|()| {
8111                            crate::util::internal_error("typed child callback receiver was dropped")
8112                        })
8113                    },
8114                ),
8115                transport,
8116            )
8117            .expect("v2 child connection should be spawned");
8118
8119        let task = futures::FutureExt::now_or_never(futures::StreamExt::next(&mut task_rx))
8120            .expect("child connection task should already be queued")
8121            .expect("parent task queue should remain open");
8122        futures::executor::block_on(async {
8123            match future::select(Box::pin(task.run_for_test()), Box::pin(callback_rx)).await {
8124                Either::Right((Ok(()), child_task)) => drop(child_task),
8125                Either::Right((Err(error), _)) => {
8126                    panic!("typed child callback sender was dropped: {error}")
8127                }
8128                Either::Left((result, _)) => {
8129                    panic!("child connection stopped before its typed callback ran: {result:?}")
8130                }
8131            }
8132        });
8133
8134        drop(child);
8135    }
8136
8137    #[cfg(feature = "unstable_protocol_v2")]
8138    #[test]
8139    fn raw_connection_returns_raw_context_for_raw_builder() {
8140        let (parent, mut task_rx) = connection_with_task_receiver();
8141        let (transport, _peer) = Channel::duplex();
8142
8143        let child: ConnectionTo<Agent> = parent
8144            .spawn_connection(Builder::new(Client), transport)
8145            .expect("raw child connection should be spawned");
8146        assert_eq!(child.protocol_version(), None);
8147
8148        let child_task = futures::FutureExt::now_or_never(futures::StreamExt::next(&mut task_rx))
8149            .expect("child connection task should already be queued")
8150            .expect("parent task queue should remain open");
8151
8152        drop((child, child_task));
8153    }
8154
8155    fn connection_with_dynamic_handler_receiver() -> (
8156        ConnectionTo<crate::role::UntypedRole>,
8157        mpsc::UnboundedReceiver<DynamicHandlerMessage<crate::role::UntypedRole>>,
8158    ) {
8159        let (message_tx, _message_rx) = mpsc::unbounded();
8160        let (task_tx, _task_rx) = mpsc::unbounded();
8161        let (dynamic_handler_tx, dynamic_handler_rx) = mpsc::unbounded();
8162        let transport_completion: SharedTransportCompletion =
8163            future::ready(Ok::<(), crate::Error>(())).boxed().shared();
8164        let pending_replies = PendingReplies::default();
8165
8166        (
8167            ConnectionTo::new(
8168                crate::role::UntypedRole,
8169                message_tx,
8170                task_tx,
8171                dynamic_handler_tx,
8172                transport_completion,
8173                pending_replies.registrar(),
8174                ProtocolMode::disabled(),
8175            ),
8176            dynamic_handler_rx,
8177        )
8178    }
8179
8180    struct ClaimingDynamicHandler;
8181
8182    impl HandleDispatchFrom<crate::role::UntypedRole> for ClaimingDynamicHandler {
8183        fn handle_dispatch_from(
8184            &mut self,
8185            _message: Dispatch,
8186            _connection: ConnectionTo<crate::role::UntypedRole>,
8187        ) -> impl Future<Output = Result<Handled<Dispatch>, crate::Error>> + Send {
8188            future::ready(Ok(Handled::Yes))
8189        }
8190
8191        fn describe_chain(&self) -> impl Debug {
8192            "ClaimingDynamicHandler"
8193        }
8194    }
8195
8196    fn connection_for_response_hook_tests() -> (
8197        ConnectionTo<crate::role::UntypedRole>,
8198        mpsc::UnboundedReceiver<OutgoingMessage>,
8199        PendingReplies,
8200    ) {
8201        let (message_tx, message_rx) = mpsc::unbounded();
8202        let (task_tx, _task_rx) = mpsc::unbounded();
8203        let (dynamic_handler_tx, _dynamic_handler_rx) = mpsc::unbounded();
8204        let transport_completion: SharedTransportCompletion =
8205            future::ready(Ok::<(), crate::Error>(())).boxed().shared();
8206        let pending_replies = PendingReplies::default();
8207
8208        (
8209            ConnectionTo::new(
8210                crate::role::UntypedRole,
8211                message_tx,
8212                task_tx,
8213                dynamic_handler_tx,
8214                transport_completion,
8215                pending_replies.registrar(),
8216                ProtocolMode::disabled(),
8217            ),
8218            message_rx,
8219            pending_replies,
8220        )
8221    }
8222
8223    #[cfg(feature = "unstable_protocol_v2")]
8224    fn route_test_response(
8225        request_id: RawJsonRpcId,
8226        pending_replies: &PendingReplies,
8227        result: Result<serde_json::Value, crate::Error>,
8228    ) {
8229        let pending_reply = pending_replies
8230            .remove(&request_id)
8231            .expect("the request should have a pending reply");
8232        let (dispatch, _) = incoming_actor::dispatch_from_response(pending_reply, result);
8233        let Dispatch::Response(result, router) = dispatch else {
8234            panic!("expected a response dispatch");
8235        };
8236        router
8237            .route_with_result(result)
8238            .expect("response should route to the pending request");
8239    }
8240
8241    #[cfg(feature = "unstable_protocol_v2")]
8242    #[test]
8243    fn handle_ids_are_typed_by_protocol_and_match_the_wire_id() {
8244        let (connection, _message_rx, pending_replies) = connection_for_response_hook_tests();
8245        let raw_sent = connection.send_request_to(
8246            crate::role::UntypedRole,
8247            UntypedMessage::new("probe", serde_json::json!({}))
8248                .expect("test request should serialize"),
8249        );
8250        let raw_sent_id: &RawJsonRpcId = raw_sent.id();
8251        assert_eq!(raw_sent_id, raw_sent.raw_id());
8252        raw_sent.detach();
8253
8254        let v2_connection = crate::v2::ConnectionTo { inner: connection };
8255        let sent = v2_connection.send_request_to(
8256            crate::role::UntypedRole,
8257            UntypedMessage::new("probe", serde_json::json!({}))
8258                .expect("test request should serialize"),
8259        );
8260        let v2_id: crate::v2::schema::RequestId = sent.id().clone();
8261        let raw_id = sent.raw_id().clone();
8262        assert_eq!(
8263            serde_json::to_value(&v2_id).expect("v2 id should serialize"),
8264            serde_json::to_value(&raw_id).expect("raw id should serialize"),
8265        );
8266
8267        let pending_reply = pending_replies
8268            .remove(&raw_id)
8269            .expect("the request should have a pending reply");
8270        let (dispatch, _) =
8271            incoming_actor::dispatch_from_response(pending_reply, Ok(serde_json::json!({})));
8272        assert_eq!(dispatch.raw_id(), Some(&raw_id));
8273        let dispatch = dispatch
8274            .into_protocol::<V2>()
8275            .expect("SDK-assigned ids are valid v2 ids");
8276        assert_eq!(dispatch.id(), Some(&v2_id));
8277        let Dispatch::Response(_, router) = dispatch else {
8278            panic!("expected a response dispatch");
8279        };
8280        assert_eq!(router.id(), &v2_id);
8281        assert_eq!(router.raw_id(), &raw_id);
8282        sent.detach();
8283    }
8284
8285    #[cfg(feature = "unstable_protocol_v2")]
8286    #[test]
8287    fn response_hook_runs_when_success_is_routed_before_consumption() {
8288        let (connection, _message_rx, pending_replies) = connection_for_response_hook_tests();
8289        let hook_ran = Arc::new(AtomicBool::new(false));
8290        let sent = connection.send_request_to_with_response_hook_after::<Raw, _, _, _>(
8291            crate::role::UntypedRole,
8292            UntypedMessage::new("hooked", serde_json::json!({}))
8293                .expect("test request should serialize"),
8294            future::ready(Ok(())),
8295            {
8296                let hook_ran = hook_ran.clone();
8297                move |response| {
8298                    assert_eq!(response, &serde_json::json!({"ok": true}));
8299                    hook_ran.store(true, Ordering::Release);
8300                    Ok(())
8301                }
8302            },
8303        );
8304        let request_id = sent.raw_id().clone();
8305
8306        route_test_response(
8307            request_id,
8308            &pending_replies,
8309            Ok(serde_json::json!({"ok": true})),
8310        );
8311
8312        assert!(hook_ran.load(Ordering::Acquire));
8313        assert_eq!(
8314            futures::executor::block_on(sent.block_task())
8315                .expect("routed response should remain consumable"),
8316            serde_json::json!({"ok": true})
8317        );
8318    }
8319
8320    #[cfg(feature = "unstable_protocol_v2")]
8321    #[test]
8322    fn response_hook_skips_errors_but_outlives_a_dropped_consumer() {
8323        let (connection, _message_rx, pending_replies) = connection_for_response_hook_tests();
8324        let peer_error_hook_ran = Arc::new(AtomicBool::new(false));
8325        let peer_error = connection.send_request_to_with_response_hook_after::<Raw, _, _, _>(
8326            crate::role::UntypedRole,
8327            UntypedMessage::new("peer-error", serde_json::json!({}))
8328                .expect("test request should serialize"),
8329            future::ready(Ok(())),
8330            {
8331                let hook_ran = peer_error_hook_ran.clone();
8332                move |_| {
8333                    hook_ran.store(true, Ordering::Release);
8334                    Ok(())
8335                }
8336            },
8337        );
8338        let peer_error_id = peer_error.raw_id().clone();
8339        route_test_response(
8340            peer_error_id,
8341            &pending_replies,
8342            Err(crate::Error::invalid_request()),
8343        );
8344        assert!(
8345            futures::executor::block_on(peer_error.block_task()).is_err(),
8346            "the peer error should reach the consumer"
8347        );
8348        assert!(!peer_error_hook_ran.load(Ordering::Acquire));
8349
8350        let dropped_hook_ran = Arc::new(AtomicBool::new(false));
8351        let dropped = connection.send_request_to_with_response_hook_after::<Raw, _, _, _>(
8352            crate::role::UntypedRole,
8353            UntypedMessage::new("dropped", serde_json::json!({}))
8354                .expect("test request should serialize"),
8355            future::ready(Ok(())),
8356            {
8357                let hook_ran = dropped_hook_ran.clone();
8358                move |_| {
8359                    hook_ran.store(true, Ordering::Release);
8360                    Ok(())
8361                }
8362            },
8363        );
8364        let dropped_id = dropped.raw_id().clone();
8365        drop(dropped);
8366        route_test_response(
8367            dropped_id,
8368            &pending_replies,
8369            Ok(serde_json::json!({"ok": true})),
8370        );
8371        assert!(dropped_hook_ran.load(Ordering::Acquire));
8372    }
8373
8374    #[cfg(feature = "unstable_protocol_v2")]
8375    #[test]
8376    fn response_hook_failure_replaces_the_success_result() {
8377        let (connection, _message_rx, pending_replies) = connection_for_response_hook_tests();
8378        let sent = connection.send_request_to_with_response_hook_after::<Raw, _, _, _>(
8379            crate::role::UntypedRole,
8380            UntypedMessage::new("hook-failure", serde_json::json!({}))
8381                .expect("test request should serialize"),
8382            future::ready(Ok(())),
8383            |_| Err(crate::Error::internal_error().data("response hook failed")),
8384        );
8385        let request_id = sent.raw_id().clone();
8386        route_test_response(
8387            request_id,
8388            &pending_replies,
8389            Ok(serde_json::json!({"ok": true})),
8390        );
8391
8392        let error = futures::executor::block_on(sent.block_task())
8393            .expect_err("the hook failure should replace the successful response");
8394        assert_eq!(error.code, crate::ErrorCode::INTERNAL_ERROR);
8395        assert_eq!(error.data, Some(serde_json::json!("response hook failed")));
8396    }
8397
8398    #[test]
8399    fn ordered_request_waits_for_readiness_before_publication() {
8400        let (connection, message_rx, pending_replies) = connection_for_response_hook_tests();
8401        let (ready_tx, ready_rx) = oneshot::channel();
8402        let sent = connection.send_ordered_request_to_after::<Raw, _, _, _>(
8403            crate::role::UntypedRole,
8404            UntypedMessage::new("after-ready", serde_json::json!({}))
8405                .expect("test request should serialize"),
8406            async move { ready_rx.await.map_err(crate::Error::into_internal_error) },
8407        );
8408
8409        let (transport_tx, mut transport_rx) = mpsc::unbounded();
8410        let mut actor = Box::pin(outgoing_actor::outgoing_protocol_actor(
8411            message_rx,
8412            pending_replies,
8413            transport_tx,
8414            ProtocolCompat::new(ProtocolMode::disabled()),
8415            future::pending::<()>().boxed().shared(),
8416        ));
8417
8418        assert!(
8419            actor.as_mut().now_or_never().is_none(),
8420            "the outgoing actor should wait for readiness"
8421        );
8422        assert!(
8423            transport_rx.next().now_or_never().is_none(),
8424            "the request must not be published before readiness"
8425        );
8426
8427        ready_tx
8428            .send(())
8429            .expect("the readiness receiver should remain active");
8430        assert!(
8431            actor.as_mut().now_or_never().is_none(),
8432            "the outgoing actor should continue serving after publication"
8433        );
8434        let frame = transport_rx
8435            .next()
8436            .now_or_never()
8437            .expect("the ready request should be published")
8438            .expect("the transport queue should remain open");
8439        assert!(matches!(
8440            frame,
8441            TransportFrame::Single(RawJsonRpcMessage::Request(_))
8442        ));
8443
8444        drop(sent);
8445    }
8446
8447    #[test]
8448    fn foreground_finish_settles_unready_requests_and_preserves_ready_output_fifo() {
8449        let (connection, message_rx, pending_replies) = connection_for_response_hook_tests();
8450        let unready = connection.send_ordered_request_to_after::<Raw, _, _, _>(
8451            crate::role::UntypedRole,
8452            UntypedMessage::new("unready", serde_json::json!({})).unwrap(),
8453            future::pending(),
8454        );
8455        let unready_id = unready.raw_id().clone();
8456        send_raw_message(
8457            &connection.message_tx,
8458            OutgoingMessage::Notification {
8459                untyped: UntypedMessage::new("first", serde_json::json!({})).unwrap(),
8460            },
8461        )
8462        .unwrap();
8463        let ready = connection.send_ordered_request_to_after::<Raw, _, _, _>(
8464            crate::role::UntypedRole,
8465            UntypedMessage::new("ready", serde_json::json!({})).unwrap(),
8466            future::ready(Ok(())),
8467        );
8468        let unready_after = connection.send_ordered_request_to_after::<Raw, _, _, _>(
8469            crate::role::UntypedRole,
8470            UntypedMessage::new("unready-after", serde_json::json!({})).unwrap(),
8471            future::pending(),
8472        );
8473        let unready_after_id = unready_after.raw_id().clone();
8474        send_raw_message(
8475            &connection.message_tx,
8476            OutgoingMessage::Notification {
8477                untyped: UntypedMessage::new("last", serde_json::json!({})).unwrap(),
8478            },
8479        )
8480        .unwrap();
8481        let (done_tx, done_rx) = oneshot::channel();
8482        send_raw_message(
8483            &connection.message_tx,
8484            OutgoingMessage::CloseAfterDraining { done: done_tx },
8485        )
8486        .unwrap();
8487        let (transport_tx, transport_rx) = mpsc::unbounded();
8488        futures::executor::block_on(outgoing_actor::outgoing_protocol_actor(
8489            message_rx,
8490            pending_replies.clone(),
8491            transport_tx,
8492            ProtocolCompat::new(ProtocolMode::disabled()),
8493            future::ready(()).boxed().shared(),
8494        ))
8495        .unwrap();
8496        futures::executor::block_on(done_rx).unwrap();
8497        let error = futures::executor::block_on(unready.block_task())
8498            .expect_err("an unresolved gate must explicitly fail its consumer");
8499        assert!(
8500            error
8501                .data
8502                .unwrap()
8503                .to_string()
8504                .contains("foreground completed before outgoing request readiness")
8505        );
8506        assert!(!pending_replies.contains(&unready_id));
8507        let error = futures::executor::block_on(unready_after.block_task())
8508            .expect_err("each unresolved gate must fail without repolling a consumed signal");
8509        assert!(
8510            error
8511                .data
8512                .unwrap()
8513                .to_string()
8514                .contains("foreground completed before outgoing request readiness")
8515        );
8516        assert!(!pending_replies.contains(&unready_after_id));
8517        assert!(pending_replies.contains(ready.raw_id()));
8518        let frames = futures::executor::block_on(transport_rx.collect::<Vec<_>>());
8519        let methods = frames
8520            .into_iter()
8521            .map(|frame| match frame {
8522                TransportFrame::Single(RawJsonRpcMessage::Notification(message)) => {
8523                    message.method.to_string()
8524                }
8525                TransportFrame::Single(RawJsonRpcMessage::Request(message)) => {
8526                    message.method.to_string()
8527                }
8528                _ => panic!("expected ready request/notification output"),
8529            })
8530            .collect::<Vec<_>>();
8531        assert_eq!(methods, ["first", "ready", "last"]);
8532    }
8533
8534    #[test]
8535    fn ordered_blocking_transform_precedes_response_acknowledgment() {
8536        let (connection, _message_rx, pending_replies) = connection_for_response_hook_tests();
8537        let sent = connection.send_ordered_request_to::<Raw, _, _>(
8538            crate::role::UntypedRole,
8539            UntypedMessage::new("ordered-transform", serde_json::json!({}))
8540                .expect("test request should serialize"),
8541        );
8542        let request_id = sent.raw_id().clone();
8543        let pending_reply = pending_replies
8544            .remove(&request_id)
8545            .expect("the request should have a pending reply");
8546        let (dispatch, response_dispatch) = incoming_actor::dispatch_from_response(
8547            pending_reply,
8548            Err(crate::Error::invalid_params()),
8549        );
8550        let Dispatch::Response(result, router) = dispatch else {
8551            panic!("expected a response dispatch");
8552        };
8553        router
8554            .route_with_result(result)
8555            .expect("response should route to the pending request");
8556        let acknowledgment = response_dispatch
8557            .complete()
8558            .expect("an ordered response should wait for acknowledgment");
8559        let acknowledgment = Arc::new(Mutex::new(Some(acknowledgment)));
8560        let acknowledgment_probe = acknowledgment.clone();
8561
8562        let error =
8563            futures::executor::block_on(sent.block_task_with_ordered_result(move |result| {
8564                assert_eq!(
8565                    acknowledgment_probe
8566                        .lock()
8567                        .expect("acknowledgment mutex poisoned")
8568                        .as_mut()
8569                        .expect("acknowledgment receiver should remain available")
8570                        .try_recv()
8571                        .expect("acknowledgment sender should remain open"),
8572                    None,
8573                    "the ordered response was acknowledged before its transform"
8574                );
8575                result
8576            }))
8577            .expect_err("the peer error should survive the ordered transform");
8578        assert_eq!(error.code, crate::ErrorCode::INVALID_PARAMS);
8579
8580        let acknowledgment = acknowledgment
8581            .lock()
8582            .expect("acknowledgment mutex poisoned")
8583            .take()
8584            .expect("acknowledgment receiver should remain available");
8585        futures::executor::block_on(acknowledgment)
8586            .expect("the transform should release the ordered response");
8587    }
8588
8589    #[cfg(feature = "unstable_protocol_v2")]
8590    #[test]
8591    fn outgoing_request_readiness_failure_rejects_without_publication() {
8592        let (connection, message_rx, pending_replies) = connection_for_response_hook_tests();
8593        let hook_ran = Arc::new(AtomicBool::new(false));
8594        let sent = connection.send_request_to_with_response_hook_after::<Raw, _, _, _>(
8595            crate::role::UntypedRole,
8596            UntypedMessage::new("never-published", serde_json::json!({}))
8597                .expect("test request should serialize"),
8598            future::ready(Err(crate::Error::internal_error().data("readiness failed"))),
8599            {
8600                let hook_ran = hook_ran.clone();
8601                move |_| {
8602                    hook_ran.store(true, Ordering::Release);
8603                    Ok(())
8604                }
8605            },
8606        );
8607
8608        let (transport_tx, mut transport_rx) = mpsc::unbounded();
8609        let mut actor = Box::pin(outgoing_actor::outgoing_protocol_actor(
8610            message_rx,
8611            pending_replies,
8612            transport_tx,
8613            ProtocolCompat::new(ProtocolMode::disabled()),
8614            future::pending::<()>().boxed().shared(),
8615        ));
8616
8617        assert!(
8618            actor.as_mut().now_or_never().is_none(),
8619            "the outgoing actor should continue serving after rejecting the request"
8620        );
8621        assert!(
8622            transport_rx.next().now_or_never().is_none(),
8623            "a request whose readiness failed must not be published"
8624        );
8625        let error = futures::executor::block_on(sent.block_task())
8626            .expect_err("the readiness error should reach the request consumer");
8627        assert_eq!(error.code, crate::ErrorCode::INTERNAL_ERROR);
8628        assert_eq!(error.data, Some(serde_json::json!("readiness failed")));
8629        assert!(!hook_ran.load(Ordering::Acquire));
8630    }
8631
8632    #[test]
8633    fn ordered_request_is_marked_before_entering_outgoing_queue() {
8634        let (message_tx, mut message_rx) = mpsc::unbounded();
8635        let (task_tx, mut task_rx) = mpsc::unbounded();
8636        let (dynamic_handler_tx, _dynamic_handler_rx) = mpsc::unbounded();
8637        let transport_completion: SharedTransportCompletion =
8638            future::ready(Ok::<(), crate::Error>(())).boxed().shared();
8639        let pending_replies = PendingReplies::default();
8640        let connection = ConnectionTo::new(
8641            crate::role::UntypedRole,
8642            message_tx,
8643            task_tx,
8644            dynamic_handler_tx,
8645            transport_completion,
8646            pending_replies.registrar(),
8647            ProtocolMode::disabled(),
8648        );
8649
8650        let sent = connection.send_ordered_request_to::<Raw, _, _>(
8651            crate::role::UntypedRole,
8652            UntypedMessage::new("ordered", serde_json::json!({}))
8653                .expect("test request should serialize"),
8654        );
8655        let request_id = sent.raw_id().clone();
8656        let message = futures::FutureExt::now_or_never(futures::StreamExt::next(&mut message_rx))
8657            .expect("outgoing request should already be queued")
8658            .expect("outgoing request queue should remain open");
8659        let OutgoingMessage::Request { id, .. } = message else {
8660            panic!("expected an outgoing request");
8661        };
8662        assert_eq!(id, request_id);
8663
8664        let pending_reply = pending_replies
8665            .remove(&request_id)
8666            .expect("the request should have a pending reply");
8667        assert!(
8668            pending_reply.ordering.is_ordered(),
8669            "the response ordering barrier must be installed before publication"
8670        );
8671
8672        // Route the response before the callback is registered. The pre-set
8673        // ordering marker must hold dispatch until the callback task is
8674        // subsequently installed and completes.
8675        let (dispatch, response_dispatch) = incoming_actor::dispatch_from_response(
8676            pending_reply,
8677            Ok(serde_json::json!({"ok": true})),
8678        );
8679        let Dispatch::Response(result, router) = dispatch else {
8680            panic!("expected a response dispatch");
8681        };
8682        router
8683            .route_with_result(result)
8684            .expect("response should route to the pending request");
8685        let acknowledgment = response_dispatch
8686            .complete()
8687            .expect("an ordered response should require acknowledgment");
8688
8689        let callback_ran = Arc::new(AtomicBool::new(false));
8690        sent.on_receiving_result({
8691            let callback_ran = callback_ran.clone();
8692            async move |result| {
8693                assert_eq!(result?, serde_json::json!({"ok": true}));
8694                callback_ran.store(true, Ordering::Release);
8695                Ok(())
8696            }
8697        })
8698        .expect("ordered callback should be scheduled");
8699
8700        let task = futures::FutureExt::now_or_never(futures::StreamExt::next(&mut task_rx))
8701            .expect("callback task should already be queued")
8702            .expect("callback task queue should remain open");
8703        futures::executor::block_on(task.run_for_test()).expect("callback task should succeed");
8704        futures::executor::block_on(acknowledgment)
8705            .expect("callback completion should acknowledge dispatch");
8706        assert!(callback_ran.load(Ordering::Acquire));
8707    }
8708
8709    fn next_dynamic_handler_message<Counterpart: Role>(
8710        receiver: &mut mpsc::UnboundedReceiver<DynamicHandlerMessage<Counterpart>>,
8711    ) -> Option<DynamicHandlerMessage<Counterpart>> {
8712        futures::FutureExt::now_or_never(futures::StreamExt::next(receiver))
8713            .expect("dynamic-handler receiver should be ready")
8714    }
8715
8716    #[cfg(feature = "unstable_protocol_v2")]
8717    #[test]
8718    fn v2_dynamic_handler_guard_registers_and_removes_handler() {
8719        let (message_tx, _message_rx) = mpsc::unbounded();
8720        let (task_tx, _task_rx) = mpsc::unbounded();
8721        let (dynamic_handler_tx, mut dynamic_handler_rx) = mpsc::unbounded();
8722        let transport_completion: SharedTransportCompletion =
8723            future::ready(Ok::<(), crate::Error>(())).boxed().shared();
8724        let pending_replies = PendingReplies::default();
8725        let connection = crate::v2::ConnectionTo {
8726            inner: ConnectionTo::new(
8727                Agent,
8728                message_tx,
8729                task_tx,
8730                dynamic_handler_tx,
8731                transport_completion,
8732                pending_replies.registrar(),
8733                ProtocolMode::v2_client(),
8734            ),
8735        };
8736
8737        let guard = connection
8738            .add_dynamic_handler(NullHandler)
8739            .expect("v2 dynamic handler should register");
8740        let added_uuid = match next_dynamic_handler_message(&mut dynamic_handler_rx) {
8741            Some(DynamicHandlerMessage::AddDynamicHandler(uuid, _, _)) => uuid,
8742            other => panic!("expected v2 handler registration, got {other:?}"),
8743        };
8744
8745        drop(guard);
8746
8747        match next_dynamic_handler_message(&mut dynamic_handler_rx) {
8748            Some(DynamicHandlerMessage::RemoveDynamicHandler(uuid)) => {
8749                assert_eq!(uuid, added_uuid);
8750            }
8751            other => panic!("expected v2 handler removal, got {other:?}"),
8752        }
8753    }
8754
8755    #[test]
8756    fn dropping_dynamic_handler_guard_unregisters_handler() {
8757        let (connection, mut receiver) = connection_with_dynamic_handler_receiver();
8758        let guard = connection.add_dynamic_handler(NullHandler).unwrap();
8759
8760        let added_uuid = match next_dynamic_handler_message(&mut receiver) {
8761            Some(DynamicHandlerMessage::AddDynamicHandler(uuid, _, _)) => uuid,
8762            other => panic!("expected handler registration, got {other:?}"),
8763        };
8764
8765        drop(guard);
8766
8767        match next_dynamic_handler_message(&mut receiver) {
8768            Some(DynamicHandlerMessage::RemoveDynamicHandler(uuid)) => {
8769                assert_eq!(uuid, added_uuid);
8770            }
8771            other => panic!("expected handler removal, got {other:?}"),
8772        }
8773    }
8774
8775    #[test]
8776    fn dropping_dynamic_handler_guard_deactivates_queued_handler_immediately() {
8777        let (connection, mut receiver) = connection_with_dynamic_handler_receiver();
8778        let guard = connection
8779            .add_dynamic_handler(ClaimingDynamicHandler)
8780            .expect("dynamic handler should register");
8781        let mut handler = match next_dynamic_handler_message(&mut receiver) {
8782            Some(DynamicHandlerMessage::AddDynamicHandler(_, handler, _)) => handler,
8783            other => panic!("expected handler registration, got {other:?}"),
8784        };
8785
8786        drop(guard);
8787
8788        let message = Dispatch::Notification(
8789            UntypedMessage::new("stale", serde_json::json!({}))
8790                .expect("test notification should serialize"),
8791        );
8792        let handled =
8793            futures::executor::block_on(handler.dyn_handle_dispatch_from(message, connection))
8794                .expect("inactive handler should decline cleanly");
8795        assert!(matches!(handled, Handled::No { retry: false, .. }));
8796    }
8797
8798    #[test]
8799    fn dynamic_handler_barrier_acknowledges_prior_messages() {
8800        let (connection, mut receiver) = connection_with_dynamic_handler_receiver();
8801        let _guard = connection.add_dynamic_handler(NullHandler).unwrap();
8802        let mut barrier = Box::pin(connection.dynamic_handler_barrier());
8803
8804        assert!(matches!(
8805            next_dynamic_handler_message(&mut receiver),
8806            Some(DynamicHandlerMessage::AddDynamicHandler(_, _, _))
8807        ));
8808        assert!(
8809            barrier.as_mut().now_or_never().is_none(),
8810            "the barrier must wait for the incoming actor"
8811        );
8812
8813        let acknowledgment = match next_dynamic_handler_message(&mut receiver) {
8814            Some(DynamicHandlerMessage::AcknowledgedBarrier(acknowledgment)) => acknowledgment,
8815            other => panic!("expected acknowledged barrier, got {other:?}"),
8816        };
8817        acknowledgment
8818            .send(())
8819            .expect("the barrier receiver should remain active");
8820        futures::executor::block_on(barrier)
8821            .expect("the acknowledged dynamic-handler barrier should complete");
8822    }
8823
8824    #[test]
8825    fn detaching_dynamic_handler_guard_does_not_leak_connection() {
8826        let (connection, mut receiver) = connection_with_dynamic_handler_receiver();
8827        let guard = connection.add_dynamic_handler(NullHandler).unwrap();
8828
8829        assert!(matches!(
8830            next_dynamic_handler_message(&mut receiver),
8831            Some(DynamicHandlerMessage::AddDynamicHandler(_, _, _))
8832        ));
8833
8834        drop(connection);
8835        guard.detach();
8836
8837        assert!(
8838            next_dynamic_handler_message(&mut receiver).is_none(),
8839            "detach should retain the handler without retaining a connection sender"
8840        );
8841    }
8842
8843    #[tokio::test]
8844    async fn write_line_flushes_buffered_writers() {
8845        let mut writer =
8846            futures::io::BufWriter::with_capacity(4096, futures::io::Cursor::new(Vec::new()));
8847
8848        write_line(&mut writer, "message".into()).await.unwrap();
8849
8850        assert_eq!(writer.into_inner().into_inner(), b"message\n");
8851    }
8852
8853    #[test]
8854    fn peel_successor_envelopes_returns_plain_messages_unchanged() {
8855        let params = serde_json::json!({ "key": "value" });
8856        let (method, peeled) = peel_successor_envelopes("session/update", &params);
8857        assert_eq!(method, "session/update");
8858        assert_eq!(peeled, &params);
8859    }
8860
8861    #[test]
8862    fn peel_successor_envelopes_unwraps_nested_envelopes() {
8863        let params = serde_json::json!({
8864            "method": "proxy/successor",
8865            "params": {
8866                "method": "$/cancel_request",
8867                "params": { "requestId": "req-1" }
8868            }
8869        });
8870        let (method, peeled) = peel_successor_envelopes("proxy/successor", &params);
8871        assert_eq!(method, "$/cancel_request");
8872        assert_eq!(peeled, &serde_json::json!({ "requestId": "req-1" }));
8873    }
8874
8875    #[test]
8876    fn peel_successor_envelopes_leaves_malformed_envelopes_intact() {
8877        // No string `method` field: the envelope cannot be peeled, so the
8878        // message is returned as-is for the handler chain to deal with.
8879        let params = serde_json::json!({ "unexpected": true });
8880        let (method, peeled) = peel_successor_envelopes("proxy/successor", &params);
8881        assert_eq!(method, "proxy/successor");
8882        assert_eq!(peeled, &params);
8883    }
8884
8885    mod cancel_request {
8886        use super::super::*;
8887
8888        fn notification(method: &str, params: serde_json::Value) -> UntypedMessage {
8889            UntypedMessage::new(method, params).expect("well-formed JSON")
8890        }
8891
8892        #[test]
8893        fn cancellation_request_id_is_extracted_from_wrapped_notifications() {
8894            let message = notification(
8895                "proxy/successor",
8896                serde_json::json!({
8897                    "method": "$/cancel_request",
8898                    "params": { "requestId": "req-1" }
8899                }),
8900            );
8901            let request_id = cancellation_request_id_from_message(&message)
8902                .expect("wrapped cancel should parse");
8903            assert_eq!(request_id, Some(RawJsonRpcId::Str("req-1".into())));
8904        }
8905
8906        #[test]
8907        fn malformed_successor_envelope_is_not_treated_as_cancellation() {
8908            // The envelope cannot be peeled; the message must flow on to the
8909            // handler chain instead of erroring the dispatch.
8910            let message = notification("proxy/successor", serde_json::json!({ "bogus": true }));
8911            let request_id = cancellation_request_id_from_message(&message)
8912                .expect("malformed envelope should be left to the handler chain");
8913            assert_eq!(request_id, None);
8914        }
8915
8916        #[test]
8917        fn cancel_request_notifications_are_detected_even_when_wrapped() {
8918            let plain = notification("$/cancel_request", serde_json::json!({ "requestId": 1 }));
8919            assert!(is_cancel_request_notification(&plain));
8920
8921            let wrapped = notification(
8922                "proxy/successor",
8923                serde_json::json!({
8924                    "method": "$/cancel_request",
8925                    "params": { "requestId": 1 }
8926                }),
8927            );
8928            assert!(is_cancel_request_notification(&wrapped));
8929
8930            let other_wrapped = notification(
8931                "proxy/successor",
8932                serde_json::json!({
8933                    "method": "session/update",
8934                    "params": {}
8935                }),
8936            );
8937            assert!(!is_cancel_request_notification(&other_wrapped));
8938
8939            let malformed_envelope =
8940                notification("proxy/successor", serde_json::json!({ "bogus": true }));
8941            assert!(!is_cancel_request_notification(&malformed_envelope));
8942        }
8943
8944        #[test]
8945        fn malformed_cancel_request_params_error() {
8946            let message = notification(
8947                "$/cancel_request",
8948                serde_json::json!({ "requestId": { "not": "an id" } }),
8949            );
8950            cancellation_request_id_from_message(&message)
8951                .expect_err("malformed cancel params should error");
8952        }
8953
8954        #[test]
8955        fn registry_marks_and_removes_requests() {
8956            let registry = RequestCancellationRegistry::new();
8957            let id = RawJsonRpcId::Str("req-1".into());
8958
8959            let responder_cancellation = registry.register(&id);
8960            let marker = responder_cancellation.cancellation();
8961            assert!(!marker.is_cancelled());
8962
8963            assert!(registry.cancel(&id));
8964            assert!(marker.is_cancelled());
8965            assert!(responder_cancellation.cancellation().is_cancelled());
8966
8967            drop(responder_cancellation);
8968            assert!(!registry.cancel(&id), "slot should be removed on drop");
8969        }
8970
8971        #[test]
8972        fn reused_request_id_does_not_cross_wire_cancellation_state() {
8973            let registry = RequestCancellationRegistry::new();
8974            let id = RawJsonRpcId::Str("dup".into());
8975
8976            // A protocol-violating peer reuses an in-flight request ID.
8977            let first = registry.register(&id);
8978            let first_marker = first.cancellation();
8979            let second = registry.register(&id);
8980            let second_marker = second.cancellation();
8981
8982            // A cancellation targets whichever request currently owns the ID.
8983            assert!(registry.cancel(&id));
8984            assert!(second_marker.is_cancelled());
8985            assert!(
8986                !first_marker.is_cancelled(),
8987                "the stale request must not observe the newer request's cancellation"
8988            );
8989
8990            // The stale responder must hand out detached markers, not the
8991            // newer request's marker.
8992            assert!(!first.cancellation().is_cancelled());
8993
8994            // Dropping the stale responder must not remove the newer
8995            // request's slot.
8996            drop(first);
8997            assert!(registry.cancel(&id), "newer slot should still be present");
8998
8999            drop(second);
9000            assert!(!registry.cancel(&id), "slot should be removed on drop");
9001        }
9002    }
9003}