Skip to main content

agent_client_protocol/role/
acp.rs

1use std::{fmt::Debug, future::Future, hash::Hash};
2
3#[cfg(feature = "unstable_protocol_v2")]
4use futures::{StreamExt as _, future};
5#[cfg(feature = "unstable_protocol_v2")]
6use serde::{Serialize, de::DeserializeOwned};
7
8#[cfg(feature = "unstable_protocol_v2")]
9use crate::DynConnectTo;
10use crate::jsonrpc::{Builder, handlers::NullHandler, run::NullRun};
11#[cfg(feature = "unstable_protocol_v2")]
12use crate::jsonrpc::{
13    TransportBatch, TransportBatchEntry, TransportFrame, V2Builder, is_response_only_shape,
14    raw_is_response_only_shape,
15};
16use crate::role::{HasPeer, RemoteStyle};
17#[cfg(not(feature = "unstable_protocol_v2"))]
18use crate::schema::InitializeProxyRequest;
19use crate::schema::METHOD_INITIALIZE_PROXY;
20#[cfg(feature = "unstable_protocol_v2")]
21use crate::schema::v1::RequestId;
22use crate::schema::v1::{InitializeRequest, SessionId};
23#[cfg(not(feature = "unstable_protocol_v2"))]
24use crate::schema::v1::{NewSessionRequest, NewSessionResponse};
25#[cfg(feature = "unstable_protocol_v2")]
26use crate::schema::{ProtocolVersion, v2};
27use crate::util::MatchDispatchFrom;
28#[cfg(feature = "unstable_protocol_v2")]
29use crate::{
30    Channel, RawJsonRpcError, RawJsonRpcMessage, RawJsonRpcParams,
31    RawJsonRpcResponse as RpcResponse,
32};
33use crate::{ConnectTo, ConnectionTo, Dispatch, HandleDispatchFrom, Handled, Role, RoleId};
34
35#[cfg(feature = "unstable_protocol_v2")]
36#[derive(serde::Deserialize)]
37struct NewSessionResponseEnvelope {
38    #[serde(rename = "sessionId")]
39    session_id: SessionId,
40}
41
42/// The client role - typically an IDE or CLI that controls an agent.
43///
44/// Clients send prompts and receive responses from agents.
45#[derive(Debug, Default, Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
46pub struct Client;
47
48impl Role for Client {
49    type Counterpart = Agent;
50
51    fn builder(self) -> Builder<Self> {
52        Builder::new(self).v1_client()
53    }
54
55    fn default_handle_dispatch_from(
56        &self,
57        message: Dispatch,
58        _connection: ConnectionTo<Client>,
59    ) -> impl Future<Output = Result<Handled<Dispatch>, crate::Error>> + Send {
60        std::future::ready(Ok(Handled::No {
61            message,
62            retry: false,
63        }))
64    }
65
66    fn role_id(&self) -> RoleId {
67        RoleId::from_singleton(self)
68    }
69
70    fn counterpart(&self) -> Self::Counterpart {
71        Agent
72    }
73}
74
75impl Client {
76    /// Create a connection builder for a client.
77    pub fn builder(self) -> Builder<Client, NullHandler, NullRun> {
78        <Self as Role>::builder(self)
79    }
80
81    /// Create a client builder that requires an ACP protocol v2 agent.
82    ///
83    /// If the agent negotiates v1 during initialization, the initialize
84    /// request resolves with an error so callers can choose an explicit v1
85    /// fallback path.
86    ///
87    /// Requires the `unstable_protocol_v2` crate feature.
88    #[cfg(feature = "unstable_protocol_v2")]
89    pub fn v2(self) -> V2Builder<Client, NullHandler, NullRun> {
90        self.builder().v2_client()
91    }
92
93    /// Create a connector that chooses between configured protocol implementations.
94    ///
95    /// Add implementation factories with [`ClientProtocolConnector::with_v1`]
96    /// and [`ClientProtocolConnector::with_v2`]. The resulting connector starts
97    /// the highest configured protocol implementation. If a v2 implementation
98    /// successfully negotiates v1 and a v1 implementation is configured, the
99    /// connector reuses the connection only when the v1 implementation's
100    /// complete `initialize` parameters match what the agent already saw;
101    /// otherwise it opens a fresh agent connection and restarts with v1.
102    ///
103    /// Requires the `unstable_protocol_v2` crate feature while protocol v2
104    /// stabilizes.
105    #[cfg(feature = "unstable_protocol_v2")]
106    #[must_use]
107    pub fn protocol_connector(self) -> ClientProtocolConnector {
108        ClientProtocolConnector::new()
109    }
110
111    /// Connect to `agent` and run `main_fn` with the [`ConnectionTo`].
112    /// Returns the result of `main_fn` (or an error if something goes wrong).
113    ///
114    /// Equivalent to `self.builder().connect_with(agent, main_fn)`.
115    pub async fn connect_with<R>(
116        self,
117        agent: impl ConnectTo<Client>,
118        main_fn: impl AsyncFnOnce(ConnectionTo<Agent>) -> Result<R, crate::Error>,
119    ) -> Result<R, crate::Error> {
120        self.builder().connect_with(agent, main_fn).await
121    }
122}
123
124/// Client connector that opens an agent connection with a configured protocol implementation.
125///
126/// Use [`Client::protocol_connector`] to start the builder, then add each
127/// supported protocol version independently. Implementations and the agent
128/// connection are provided as factories because fallback from v2 to v1 may
129/// require a fresh connection initialized by the v1 implementation.
130#[cfg(feature = "unstable_protocol_v2")]
131#[derive(Debug, Default)]
132pub struct ClientProtocolConnector {
133    v1: Option<DynConnectToFactory<Agent>>,
134    v2: Option<DynConnectToFactory<Agent>>,
135}
136
137#[cfg(feature = "unstable_protocol_v2")]
138impl ClientProtocolConnector {
139    /// Create an empty client protocol connector.
140    #[must_use]
141    pub fn new() -> Self {
142        Self::default()
143    }
144
145    /// Return this connector with an ACP v1 implementation factory configured.
146    #[must_use]
147    pub fn with_v1<C>(mut self, client: impl FnMut() -> C + Send + 'static) -> Self
148    where
149        C: ConnectTo<Agent>,
150    {
151        self.v1 = Some(DynConnectToFactory::new(client));
152        self
153    }
154
155    /// Return this connector with an ACP v2 implementation factory configured.
156    #[must_use]
157    pub fn with_v2<C>(mut self, client: impl FnMut() -> C + Send + 'static) -> Self
158    where
159        C: ConnectTo<Agent>,
160    {
161        self.v2 = Some(DynConnectToFactory::new(client));
162        self
163    }
164
165    /// Connect to an agent produced by `agent` using the highest configured
166    /// compatible protocol implementation.
167    pub async fn connect_to<C>(
168        mut self,
169        mut agent: impl FnMut() -> C + Send + 'static,
170    ) -> Result<(), crate::Error>
171    where
172        C: ConnectTo<Client>,
173    {
174        let supported = SupportedClientProtocols {
175            v1: self.v1.is_some(),
176            v2: self.v2.is_some(),
177        };
178        let Some(selected) = supported.highest_configured() else {
179            return Err(crate::Error::invalid_request()
180                .data("client protocol connector has no configured ACP protocol implementations"));
181        };
182
183        match selected {
184            ClientProtocol::V1 => {
185                let client = self
186                    .v1
187                    .as_mut()
188                    .expect("selected protocol is configured")
189                    .create();
190                connect_client_protocol(ClientProtocol::V1, client, agent()).await
191            }
192            ClientProtocol::V2 => {
193                let client = self
194                    .v2
195                    .as_mut()
196                    .expect("selected protocol is configured")
197                    .create();
198                let agent_connection = RunningProtocolPeer::new(agent());
199                let (client, initialize) =
200                    start_client_protocol(ClientProtocol::V2, client).await?;
201                // This normalization is only a probe for the connection-reuse
202                // optimization. A request that cannot be represented in v1
203                // can still be valid v2 traffic and must reach the agent.
204                let v2_initialize_as_v1 = normalize_v2_initialize_params_for_reuse(&initialize);
205                let (client, agent_connection, initialize_response) =
206                    send_initialize_and_receive(client, agent_connection, initialize).await?;
207
208                if initialize_response_negotiated_v1(&initialize_response)
209                    && let Some(v1) = self.v1.as_mut()
210                {
211                    let fallback_client = v1.create();
212                    let (fallback_client, fallback_initialize) =
213                        start_client_protocol(ClientProtocol::V1, fallback_client).await?;
214                    let v1_initialize =
215                        validated_initialize_params::<InitializeRequest>(&fallback_initialize)?;
216
217                    if v2_initialize_as_v1
218                        .as_ref()
219                        .is_ok_and(|v2_initialize| v2_initialize == &v1_initialize)
220                    {
221                        let fallback_response = initialize_response.with_id(
222                            initialize_request_id(&fallback_initialize)
223                                .expect("validated initialize request has an id"),
224                        );
225                        // The v2 implementation will never receive its initialize response once
226                        // the matching v1 implementation takes over this connection. Drop its
227                        // future and channels before running the v1 session so any resources it
228                        // owns are released promptly.
229                        drop(client);
230                        fallback_client.send(fallback_response)?;
231                        return pipe_protocol_peers_until_done(fallback_client, agent_connection)
232                            .await;
233                    }
234
235                    // Neither probe can continue on the replacement connection. Release both
236                    // client implementations and the original agent connection before starting
237                    // the real v1 session so they cannot retain resources for its lifetime.
238                    drop((
239                        client,
240                        fallback_client,
241                        agent_connection,
242                        initialize_response,
243                    ));
244                    return connect_client_protocol(ClientProtocol::V1, v1.create(), agent()).await;
245                }
246
247                client.send(initialize_response.into_message())?;
248                pipe_protocol_peers_until_done(client, agent_connection).await
249            }
250        }
251    }
252}
253
254#[cfg(feature = "unstable_protocol_v2")]
255struct DynConnectToFactory<R: Role> {
256    inner: Box<dyn FnMut() -> DynConnectTo<R> + Send>,
257}
258
259#[cfg(feature = "unstable_protocol_v2")]
260impl<R: Role> DynConnectToFactory<R> {
261    fn new<C>(mut factory: impl FnMut() -> C + Send + 'static) -> Self
262    where
263        C: ConnectTo<R>,
264    {
265        Self {
266            inner: Box::new(move || DynConnectTo::new(factory())),
267        }
268    }
269
270    fn create(&mut self) -> DynConnectTo<R> {
271        (self.inner)()
272    }
273}
274
275#[cfg(feature = "unstable_protocol_v2")]
276impl<R: Role> Debug for DynConnectToFactory<R> {
277    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
278        f.debug_struct("DynConnectToFactory")
279            .finish_non_exhaustive()
280    }
281}
282
283impl HasPeer<Client> for Client {
284    fn remote_style(&self, _peer: Client) -> RemoteStyle {
285        RemoteStyle::Counterpart
286    }
287}
288
289/// The agent role - typically an LLM that responds to prompts.
290///
291/// Agents receive prompts from clients and respond with answers,
292/// potentially invoking tools along the way.
293#[derive(Debug, Default, Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
294pub struct Agent;
295
296impl Role for Agent {
297    type Counterpart = Client;
298
299    fn builder(self) -> Builder<Self> {
300        Builder::new(self).v1_agent()
301    }
302
303    fn role_id(&self) -> RoleId {
304        RoleId::from_singleton(self)
305    }
306
307    fn counterpart(&self) -> Self::Counterpart {
308        Client
309    }
310
311    async fn default_handle_dispatch_from(
312        &self,
313        message: Dispatch,
314        connection: ConnectionTo<Agent>,
315    ) -> Result<Handled<Dispatch>, crate::Error> {
316        MatchDispatchFrom::new(message, &connection)
317            .if_dispatch_from(Agent, async |message: Dispatch| {
318                // Stable v1 session helpers install a dynamic handler after
319                // `session/new`. Retry session messages to close the race
320                // between the response and that registration.
321                //
322                // V2 uses typed handlers installed before the connection
323                // starts. Retrying an unhandled v2 message would retain it
324                // forever because no per-session dynamic handler is expected.
325                #[cfg(feature = "unstable_protocol_v2")]
326                let retry = message.has_session_id()
327                    && connection.acp_protocol_version()
328                        != Some(crate::schema::ProtocolVersion::V2);
329                #[cfg(not(feature = "unstable_protocol_v2"))]
330                let retry = message.has_session_id();
331                Ok(Handled::No { message, retry })
332            })
333            .await
334            .done()
335    }
336}
337
338impl Agent {
339    /// Create a connection builder for an agent.
340    pub fn builder(self) -> Builder<Agent, NullHandler, NullRun> {
341        <Self as Role>::builder(self)
342    }
343
344    /// Create an agent builder that uses the ACP protocol v2 API.
345    ///
346    /// This builder selects protocol v2 for clients requesting v2 or a newer
347    /// version with compatible initialization parameters. Use
348    /// [`Self::protocol_router`] with a v1 implementation for v1 clients.
349    ///
350    /// Requires the `unstable_protocol_v2` crate feature.
351    #[cfg(feature = "unstable_protocol_v2")]
352    pub fn v2(self) -> V2Builder<Agent, NullHandler, NullRun> {
353        self.builder().v2_agent()
354    }
355
356    /// Create a router that chooses between configured protocol implementations.
357    ///
358    /// Add implementations with [`AgentProtocolRouter::with_v1`] and
359    /// [`AgentProtocolRouter::with_v2`].
360    /// The resulting router reads the initial
361    /// `initialize` request and selects a version-specific implementation. An
362    /// endpoint exposing v2 selects the highest configured version at or below
363    /// the requested version. A v1-only endpoint retains v1's fallback-to-latest
364    /// rule. The router does not convert traffic after initialization.
365    ///
366    /// Requires the `unstable_protocol_v2` crate feature while protocol v2
367    /// stabilizes.
368    #[cfg(feature = "unstable_protocol_v2")]
369    #[must_use]
370    pub fn protocol_router(self) -> AgentProtocolRouter {
371        AgentProtocolRouter::new()
372    }
373}
374
375/// Agent component that routes each connection to a configured protocol implementation.
376///
377/// Use [`Agent::protocol_router`] to start the builder, then add each supported
378/// protocol version independently. The selected implementation owns the
379/// connection after the initial `initialize` negotiation.
380#[cfg(feature = "unstable_protocol_v2")]
381#[derive(Debug, Default)]
382pub struct AgentProtocolRouter {
383    v1: Option<DynConnectTo<Client>>,
384    v2: Option<DynConnectTo<Client>>,
385}
386
387#[cfg(feature = "unstable_protocol_v2")]
388impl AgentProtocolRouter {
389    /// Create an empty agent protocol router.
390    #[must_use]
391    pub fn new() -> Self {
392        Self::default()
393    }
394
395    /// Return this router with an ACP v1 implementation configured.
396    #[must_use]
397    pub fn with_v1(mut self, agent: impl ConnectTo<Client>) -> Self {
398        self.v1 = Some(DynConnectTo::new(agent));
399        self
400    }
401
402    /// Return this router with an ACP v2 implementation configured.
403    #[must_use]
404    pub fn with_v2(mut self, agent: impl ConnectTo<Client>) -> Self {
405        self.v2 = Some(DynConnectTo::new(agent));
406        self
407    }
408}
409
410#[cfg(feature = "unstable_protocol_v2")]
411impl ConnectTo<Client> for AgentProtocolRouter {
412    async fn connect_to(self, client: impl ConnectTo<Agent>) -> Result<(), crate::Error> {
413        let supported = SupportedProtocols {
414            v1: self.v1.is_some(),
415            v2: self.v2.is_some(),
416        };
417        let mut client = RunningProtocolPeer::new(client);
418        let (first_frame, client, selected) = loop {
419            let Some((mut frame, next_client)) = client.next_frame().await? else {
420                return Ok(());
421            };
422            let message = match initialize_message_mut(&mut frame) {
423                Ok(Some(message)) => message,
424                Ok(None) => {
425                    client = next_client;
426                    continue;
427                }
428                Err(error) => return reject_initialize(next_client, &frame, error).await,
429            };
430            let selected = match select_agent_protocol(message, supported) {
431                Ok(selected) => selected,
432                Err(error) => return reject_initialize(next_client, &frame, error).await,
433            };
434            break (frame, next_client, selected);
435        };
436        let Some(agent) = selected.take_agent(self) else {
437            let error = selected.unsupported_error(supported);
438            return reject_initialize(client, &first_frame, error).await;
439        };
440
441        let agent = RunningProtocolPeer::new(agent);
442        agent.send_frame(first_frame)?;
443        pipe_protocol_peers_until_done(client, agent).await
444    }
445}
446
447#[cfg(feature = "unstable_protocol_v2")]
448#[derive(Debug, Clone, Copy, PartialEq, Eq)]
449enum SelectedProtocol {
450    V1,
451    V2,
452}
453
454#[cfg(feature = "unstable_protocol_v2")]
455impl SelectedProtocol {
456    fn take_agent(self, agent: AgentProtocolRouter) -> Option<DynConnectTo<Client>> {
457        match self {
458            Self::V1 => agent.v1,
459            Self::V2 => agent.v2,
460        }
461    }
462
463    fn version(self) -> ProtocolVersion {
464        match self {
465            Self::V1 => ProtocolVersion::V1,
466            Self::V2 => ProtocolVersion::V2,
467        }
468    }
469
470    fn name(self) -> &'static str {
471        match self {
472            Self::V1 => "1",
473            Self::V2 => "2",
474        }
475    }
476
477    fn unsupported_error(self, supported: SupportedProtocols) -> crate::Error {
478        crate::Error::invalid_request().data(format!(
479            "ACP protocol version {} is not configured; this endpoint supports {}",
480            self.name(),
481            supported.description()
482        ))
483    }
484}
485
486#[cfg(feature = "unstable_protocol_v2")]
487#[derive(Debug, Clone, Copy, PartialEq, Eq)]
488struct SupportedProtocols {
489    v1: bool,
490    v2: bool,
491}
492
493#[cfg(feature = "unstable_protocol_v2")]
494impl SupportedProtocols {
495    fn select_for_agent(self, requested: ProtocolVersion) -> Option<SelectedProtocol> {
496        if self.v1 && !self.v2 {
497            // A v1-only endpoint retains the legacy fallback-to-latest rule.
498            return Some(SelectedProtocol::V1);
499        }
500        if self.v2 && requested >= ProtocolVersion::V2 {
501            return Some(SelectedProtocol::V2);
502        }
503
504        if self.v1 && requested >= ProtocolVersion::V1 {
505            return Some(SelectedProtocol::V1);
506        }
507
508        None
509    }
510
511    fn exact(self, requested: ProtocolVersion) -> Option<SelectedProtocol> {
512        if self.v1 && requested == ProtocolVersion::V1 {
513            Some(SelectedProtocol::V1)
514        } else if self.v2 && requested == ProtocolVersion::V2 {
515            Some(SelectedProtocol::V2)
516        } else {
517            None
518        }
519    }
520
521    fn versions(self) -> Vec<ProtocolVersion> {
522        [
523            self.v1.then_some(ProtocolVersion::V1),
524            self.v2.then_some(ProtocolVersion::V2),
525        ]
526        .into_iter()
527        .flatten()
528        .collect()
529    }
530
531    fn description(self) -> String {
532        match (self.v1, self.v2) {
533            (true, true) => "ACP protocol versions 1 and 2".into(),
534            (true, false) => "ACP protocol version 1".into(),
535            (false, true) => "ACP protocol version 2".into(),
536            (false, false) => "no ACP protocol versions".into(),
537        }
538    }
539}
540
541#[cfg(feature = "unstable_protocol_v2")]
542fn select_agent_protocol(
543    message: &mut RawJsonRpcMessage,
544    supported: SupportedProtocols,
545) -> Result<SelectedProtocol, crate::Error> {
546    let RawJsonRpcMessage::Request(request) = message else {
547        return Err(
548            crate::Error::invalid_request().data("first ACP message must be an initialize request")
549        );
550    };
551
552    if request.method.as_ref() != "initialize" {
553        return Err(crate::Error::invalid_request().data("first ACP request must be initialize"));
554    }
555
556    let Some(RawJsonRpcParams::Object(params)) = &mut request.params else {
557        return Err(invalid_initialize_protocol_version());
558    };
559    let Some(protocol_version) = params.get("protocolVersion") else {
560        return Err(invalid_initialize_protocol_version());
561    };
562
563    let requested = serde_json::from_value::<ProtocolVersion>(protocol_version.clone())
564        .map_err(|_| invalid_initialize_protocol_version())?;
565    let selected = select_configured_agent_protocol(requested, supported)?;
566    rewrite_initialize_params(params, requested, selected)?;
567
568    Ok(selected)
569}
570
571#[cfg(feature = "unstable_protocol_v2")]
572fn initialize_request_params(
573    message: &RawJsonRpcMessage,
574) -> Result<&serde_json::Map<String, serde_json::Value>, crate::Error> {
575    let RawJsonRpcMessage::Request(request) = message else {
576        return Err(
577            crate::Error::invalid_request().data("first ACP message must be an initialize request")
578        );
579    };
580
581    if request.method.as_ref() != "initialize" {
582        return Err(crate::Error::invalid_request().data("first ACP request must be initialize"));
583    }
584
585    let Some(RawJsonRpcParams::Object(params)) = &request.params else {
586        return Err(invalid_initialize_protocol_version());
587    };
588    if !params.contains_key("protocolVersion") {
589        return Err(invalid_initialize_protocol_version());
590    }
591    Ok(params)
592}
593
594#[cfg(feature = "unstable_protocol_v2")]
595fn validated_initialize_params<T: DeserializeOwned>(
596    message: &RawJsonRpcMessage,
597) -> Result<serde_json::Map<String, serde_json::Value>, crate::Error> {
598    let params = initialize_request_params(message)?;
599    parse_initialize_params::<T>(params)?;
600    Ok(params.clone())
601}
602
603#[cfg(feature = "unstable_protocol_v2")]
604fn normalize_v2_initialize_params_for_reuse(
605    message: &RawJsonRpcMessage,
606) -> Result<serde_json::Map<String, serde_json::Value>, crate::Error> {
607    let params = initialize_request_params(message)?;
608    let requested = params
609        .get("protocolVersion")
610        .cloned()
611        .ok_or_else(invalid_initialize_protocol_version)
612        .and_then(|version| {
613            serde_json::from_value::<ProtocolVersion>(version)
614                .map_err(|_| invalid_initialize_protocol_version())
615        })?;
616    if requested == ProtocolVersion::V1 {
617        parse_initialize_params::<InitializeRequest>(params)?;
618        return Ok(params.clone());
619    }
620    normalize_v2_initialize_params_for_v1(params, true)
621}
622
623#[cfg(feature = "unstable_protocol_v2")]
624fn rewrite_initialize_params(
625    params: &mut serde_json::Map<String, serde_json::Value>,
626    requested: ProtocolVersion,
627    selected: SelectedProtocol,
628) -> Result<(), crate::Error> {
629    // Validate exact-version initialization without replacing its raw
630    // parameters. Reserializing through the SDK's pinned schema would discard
631    // fields added by newer compatible peers.
632    if requested == selected.version() {
633        match selected {
634            SelectedProtocol::V1 => {
635                parse_initialize_params::<InitializeRequest>(params)?;
636            }
637            SelectedProtocol::V2 => {
638                parse_initialize_params::<v2::InitializeRequest>(params)?;
639            }
640        }
641        return Ok(());
642    }
643
644    match selected {
645        SelectedProtocol::V1 => {
646            if requested > ProtocolVersion::V1 {
647                *params = normalize_v2_initialize_params_for_v1(params, false)?;
648            } else {
649                let mut initialize = parse_initialize_params::<InitializeRequest>(params)?;
650                initialize.protocol_version = ProtocolVersion::V1;
651                *params = serialize_initialize_params(initialize)?;
652            }
653            Ok(())
654        }
655        SelectedProtocol::V2 => {
656            let mut initialize = parse_initialize_params::<v2::InitializeRequest>(params)?;
657            initialize.protocol_version = ProtocolVersion::V2;
658            *params = serialize_initialize_params(initialize)?;
659            Ok(())
660        }
661    }
662}
663
664#[cfg(feature = "unstable_protocol_v2")]
665fn normalize_v2_initialize_params_for_v1(
666    params: &serde_json::Map<String, serde_json::Value>,
667    require_lossless: bool,
668) -> Result<serde_json::Map<String, serde_json::Value>, crate::Error> {
669    if !require_lossless {
670        return crate::jsonrpc::normalize_v1_initialize_request(&serde_json::Value::Object(
671            params.clone(),
672        ));
673    }
674    // The connection-reuse probe additionally requires lossless v2
675    // canonicalization. Normal v1 routing above uses the same bootstrap as a
676    // standalone v1 agent, even when the v2 schema is not compiled in.
677    let initialize = parse_initialize_params::<v2::InitializeRequest>(params)?;
678    let mut target = serialize_initialize_params(initialize)?;
679    if require_lossless && target != *params {
680        return Err(invalid_initialize_params(
681            "v2 initialize parameters are not losslessly representable in v1",
682        ));
683    }
684
685    target.insert(
686        "protocolVersion".into(),
687        serde_json::to_value(ProtocolVersion::V1).map_err(crate::Error::into_internal_error)?,
688    );
689    let info = target
690        .remove("info")
691        .ok_or_else(|| invalid_initialize_params("v2 InitializeRequest.info is required"))?;
692    target.insert("clientInfo".into(), info);
693    let capabilities = target
694        .remove("capabilities")
695        .and_then(|capabilities| capabilities.as_object().cloned())
696        .ok_or_else(|| {
697            crate::util::internal_error("v2 initialize capabilities did not serialize as an object")
698        })?;
699    let mut capabilities = capabilities;
700    if let Some(auth) = capabilities
701        .get_mut("auth")
702        .and_then(serde_json::Value::as_object_mut)
703    {
704        let terminal = auth.remove("terminal");
705        if require_lossless
706            && terminal
707                .as_ref()
708                .and_then(serde_json::Value::as_object)
709                .is_some_and(|terminal| terminal.contains_key("_meta"))
710        {
711            return Err(invalid_initialize_params(
712                "v2 terminal authentication metadata is not representable in v1",
713            ));
714        }
715        auth.insert("terminal".into(), terminal.is_some().into());
716    }
717    capabilities.insert(
718        "session".into(),
719        serde_json::json!({ "configOptions": { "boolean": {} } }),
720    );
721    target.insert("clientCapabilities".into(), capabilities.into());
722
723    let initialize = parse_initialize_params::<InitializeRequest>(&target)?;
724    let normalized = serialize_initialize_params(initialize)?;
725    if require_lossless && !json_object_contains(&normalized, &target) {
726        return Err(invalid_initialize_params(
727            "v2 initialize parameters are not losslessly representable in v1",
728        ));
729    }
730    Ok(normalized)
731}
732
733#[cfg(all(test, feature = "unstable_protocol_v2"))]
734mod initialize_normalization_tests {
735    use super::*;
736
737    fn v2_initialize_params() -> serde_json::Map<String, serde_json::Value> {
738        let value = serde_json::to_value(v2::InitializeRequest::new(
739            ProtocolVersion::V2,
740            v2::Implementation::new("test-client", "1.0.0"),
741        ))
742        .expect("serialize v2 initialize request");
743        value
744            .as_object()
745            .expect("initialize params serialize as an object")
746            .clone()
747    }
748
749    #[test]
750    fn v2_tolerant_fields_are_canonicalized_before_v1_normalization() {
751        let mut params = v2_initialize_params();
752        params.insert(
753            "capabilities".into(),
754            serde_json::Value::String("malformed".into()),
755        );
756        params.insert(
757            "_meta".into(),
758            serde_json::Value::String("malformed".into()),
759        );
760
761        let normalized = normalize_v2_initialize_params_for_v1(&params, false)
762            .expect("tolerant v2 fields should normalize through their defaults");
763        let normalized = serde_json::Value::Object(normalized);
764
765        assert!(normalized.get("_meta").is_none());
766        assert_eq!(
767            normalized.pointer("/clientCapabilities/session/configOptions/boolean"),
768            Some(&serde_json::json!({}))
769        );
770    }
771
772    #[test]
773    fn noncanonical_v2_fields_disable_reuse_but_not_v1_routing() {
774        let mut params = v2_initialize_params();
775        params
776            .get_mut("info")
777            .and_then(serde_json::Value::as_object_mut)
778            .expect("v2 initialize info is an object")
779            .insert("buildCommit".into(), serde_json::json!("abc123"));
780
781        normalize_v2_initialize_params_for_v1(&params, false)
782            .expect("v1 routing may ignore parameters unavailable in v1");
783        normalize_v2_initialize_params_for_v1(&params, true)
784            .expect_err("connection reuse requires lossless normalization");
785    }
786
787    #[test]
788    fn v1_reuse_probe_preserves_raw_initialize_params() {
789        let mut params = v2_initialize_params();
790        params.insert(
791            "protocolVersion".into(),
792            serde_json::json!(ProtocolVersion::V1),
793        );
794        let message = RawJsonRpcMessage::request(
795            "initialize".into(),
796            serde_json::Value::Object(params.clone()),
797            RequestId::Number(1),
798        )
799        .expect("build initialize request");
800
801        let normalized = normalize_v2_initialize_params_for_reuse(&message)
802            .expect("v1-shaped initialize request should be valid");
803
804        assert_eq!(normalized, params);
805    }
806
807    #[test]
808    fn null_v2_terminal_marker_meta_is_omitted_before_v1_normalization() {
809        let mut params = v2_initialize_params();
810        params.insert(
811            "capabilities".into(),
812            serde_json::json!({
813                "auth": {
814                    "terminal": { "_meta": null }
815                }
816            }),
817        );
818
819        let normalized = normalize_v2_initialize_params_for_v1(&params, false)
820            .expect("null marker metadata is equivalent to omission");
821        let normalized = serde_json::Value::Object(normalized);
822
823        assert_eq!(
824            normalized.pointer("/clientCapabilities/auth/terminal"),
825            Some(&serde_json::Value::Bool(true))
826        );
827    }
828
829    #[test]
830    fn terminal_marker_metadata_disables_reuse_but_not_v1_routing() {
831        let mut params = v2_initialize_params();
832        params.insert(
833            "capabilities".into(),
834            serde_json::json!({
835                "auth": {
836                    "terminal": {
837                        "_meta": { "source": "test" }
838                    }
839                }
840            }),
841        );
842
843        normalize_v2_initialize_params_for_v1(&params, false)
844            .expect("v1 routing may discard terminal marker metadata");
845        normalize_v2_initialize_params_for_v1(&params, true)
846            .expect_err("connection reuse must preserve terminal marker metadata");
847    }
848}
849
850#[cfg(feature = "unstable_protocol_v2")]
851fn parse_initialize_params<T: DeserializeOwned>(
852    params: &serde_json::Map<String, serde_json::Value>,
853) -> Result<T, crate::Error> {
854    serde_json::from_value(serde_json::Value::Object(params.clone()))
855        .map_err(invalid_initialize_params)
856}
857
858#[cfg(feature = "unstable_protocol_v2")]
859fn serialize_initialize_params(
860    initialize: impl Serialize,
861) -> Result<serde_json::Map<String, serde_json::Value>, crate::Error> {
862    let value = serde_json::to_value(initialize).map_err(crate::Error::into_internal_error)?;
863    let serde_json::Value::Object(object) = value else {
864        return Err(crate::util::internal_error(
865            "initialize params did not serialize to an object",
866        ));
867    };
868    Ok(object)
869}
870
871#[cfg(feature = "unstable_protocol_v2")]
872fn json_object_contains(
873    actual: &serde_json::Map<String, serde_json::Value>,
874    expected: &serde_json::Map<String, serde_json::Value>,
875) -> bool {
876    fn contains(actual: &serde_json::Value, expected: &serde_json::Value) -> bool {
877        match (actual, expected) {
878            (serde_json::Value::Object(actual), serde_json::Value::Object(expected)) => expected
879                .iter()
880                .all(|(key, value)| actual.get(key).is_some_and(|item| contains(item, value))),
881            _ => actual == expected,
882        }
883    }
884
885    expected
886        .iter()
887        .all(|(key, value)| actual.get(key).is_some_and(|item| contains(item, value)))
888}
889
890#[cfg(feature = "unstable_protocol_v2")]
891fn select_configured_agent_protocol(
892    requested: ProtocolVersion,
893    supported: SupportedProtocols,
894) -> Result<SelectedProtocol, crate::Error> {
895    supported
896        .select_for_agent(requested)
897        .ok_or_else(|| crate::util::unsupported_protocol_version(requested, &supported.versions()))
898}
899
900#[cfg(feature = "unstable_protocol_v2")]
901fn invalid_initialize_protocol_version() -> crate::Error {
902    crate::Error::invalid_params()
903        .data("initialize.protocolVersion must be a valid ACP protocol version")
904}
905
906#[cfg(feature = "unstable_protocol_v2")]
907fn invalid_initialize_params(error: impl ToString) -> crate::Error {
908    crate::Error::invalid_params().data(format!("invalid initialize params: {}", error.to_string()))
909}
910
911#[cfg(feature = "unstable_protocol_v2")]
912fn send_initialize_error(
913    tx: &futures::channel::mpsc::UnboundedSender<TransportFrame>,
914    frame: &TransportFrame,
915    error: crate::Error,
916) -> Result<(), crate::Error> {
917    fn response_for_message(
918        entry: &RawJsonRpcMessage,
919        initialize_error: &crate::Error,
920    ) -> Option<RawJsonRpcMessage> {
921        match entry {
922            RawJsonRpcMessage::Request(request) => Some(RawJsonRpcMessage::response(
923                request.id.clone(),
924                Err(initialize_error.clone()),
925            )),
926            RawJsonRpcMessage::Notification(_) | RawJsonRpcMessage::Response(_) => None,
927        }
928    }
929
930    fn response_for_entry(
931        entry: &TransportBatchEntry,
932        initialize_error: &crate::Error,
933    ) -> Option<RawJsonRpcMessage> {
934        match entry {
935            TransportBatchEntry::Message(message) => {
936                response_for_message(message, initialize_error)
937            }
938            TransportBatchEntry::Malformed { raw, error } if !is_response_only_shape(raw) => Some(
939                RawJsonRpcMessage::response(RequestId::Null, Err(error.clone())),
940            ),
941            TransportBatchEntry::Malformed { .. } => None,
942        }
943    }
944
945    let response = match frame {
946        TransportFrame::Single(entry) => {
947            let Some(response) = response_for_message(entry, &error) else {
948                return Ok(());
949            };
950            TransportFrame::Single(response)
951        }
952        TransportFrame::Malformed { raw, error } if !raw_is_response_only_shape(raw) => {
953            TransportFrame::Single(RawJsonRpcMessage::response(
954                RequestId::Null,
955                Err(error.clone()),
956            ))
957        }
958        TransportFrame::Malformed { .. } => return Ok(()),
959        TransportFrame::Batch(batch) => {
960            let responses = batch
961                .entries()
962                .filter_map(|entry| response_for_entry(entry, &error))
963                .collect::<Vec<_>>();
964            let Some(responses) = TransportBatch::from_messages(responses) else {
965                return Ok(());
966            };
967            TransportFrame::Batch(responses)
968        }
969    };
970
971    tx.unbounded_send(response)
972        .map_err(crate::util::internal_error)
973}
974
975#[cfg(feature = "unstable_protocol_v2")]
976async fn reject_initialize(
977    client: RunningProtocolPeer,
978    frame: &TransportFrame,
979    error: crate::Error,
980) -> Result<(), crate::Error> {
981    let RunningProtocolPeer { mut rx, tx, driver } = client;
982    send_initialize_error(&tx, frame, error)?;
983    drop(tx);
984
985    let Some(mut driver) = driver.into_driver() else {
986        // The rejection has already been handed to the raw channel. There is
987        // no owned transport work or physical drain to await.
988        return Ok(());
989    };
990    if !driver.request_finish() {
991        // An opaque driver has no finite physical-finish contract. Preserve a
992        // ready error before cancelling it rather than wait for remote EOF.
993        return crate::util::run_until(driver, future::ready(Ok(()))).await;
994    }
995
996    let drain_incoming = async move {
997        // Later input has no protocol meaning once initialization is rejected.
998        // Keep draining it only so the transport can flush the queued rejection.
999        // Treating a malformed trailing frame as fatal would cancel that flush.
1000        while rx.next().await.is_some() {}
1001        Ok::<_, crate::Error>(())
1002    };
1003
1004    match future::select(driver, Box::pin(drain_incoming)).await {
1005        future::Either::Left((result, _)) => result,
1006        future::Either::Right((result, driver)) => {
1007            result?;
1008            driver.await
1009        }
1010    }
1011}
1012
1013#[cfg(feature = "unstable_protocol_v2")]
1014struct RunningProtocolPeer {
1015    rx: futures::channel::mpsc::UnboundedReceiver<TransportFrame>,
1016    tx: futures::channel::mpsc::UnboundedSender<TransportFrame>,
1017    driver: ProtocolPeerDriver,
1018}
1019
1020#[cfg(feature = "unstable_protocol_v2")]
1021enum ProtocolPeerDriver {
1022    Passive,
1023    Active(crate::ConnectionDriver),
1024    Completed {
1025        finish: Option<crate::component::FinishControl>,
1026    },
1027}
1028
1029#[cfg(feature = "unstable_protocol_v2")]
1030impl ProtocolPeerDriver {
1031    fn into_driver(self) -> Option<crate::ConnectionDriver> {
1032        match self {
1033            Self::Passive => None,
1034            Self::Active(driver) => Some(driver),
1035            // Conversion happens only when handing the peer to its final
1036            // bridge, never while reading its remaining queued frames. This
1037            // records actual owned completion, not a passive ready sentinel.
1038            Self::Completed { finish } => Some(match finish {
1039                Some(mut finish) => {
1040                    crate::ConnectionDriver::with_finish(future::ready(Ok(())), move || {
1041                        finish.request();
1042                    })
1043                }
1044                None => crate::ConnectionDriver::new(future::ready(Ok(()))),
1045            }),
1046        }
1047    }
1048}
1049
1050#[cfg(feature = "unstable_protocol_v2")]
1051impl RunningProtocolPeer {
1052    fn new<R: Role>(component: impl ConnectTo<R>) -> Self {
1053        let (Channel { rx, tx }, future) = component.into_channel_and_future();
1054        let driver = match future {
1055            None => ProtocolPeerDriver::Passive,
1056            Some(future) => ProtocolPeerDriver::Active(future),
1057        };
1058        Self { rx, tx, driver }
1059    }
1060
1061    async fn next_frame(self) -> Result<Option<(TransportFrame, Self)>, crate::Error> {
1062        let Self { mut rx, tx, driver } = self;
1063        let ProtocolPeerDriver::Active(mut future) = driver else {
1064            return Ok(rx
1065                .next()
1066                .await
1067                .map(|frame| (frame, Self { rx, tx, driver })));
1068        };
1069
1070        // Poll the owned driver first: a ready error must not be hidden by
1071        // an equally ready frame or clean channel EOF.
1072        match future::select(&mut future, Box::pin(rx.next())).await {
1073            future::Either::Right((Some(frame), _)) => Ok(Some((
1074                frame,
1075                Self {
1076                    rx,
1077                    tx,
1078                    driver: ProtocolPeerDriver::Active(future),
1079                },
1080            ))),
1081            future::Either::Right((None, _)) => {
1082                drop(tx);
1083                future.await?;
1084                Ok(None)
1085            }
1086            future::Either::Left((result, next_message)) => {
1087                result?;
1088                drop(next_message);
1089                // No more output may be accepted from an owned endpoint once
1090                // its driver completes, even if a sender escaped the component.
1091                rx.close();
1092                let Some(frame) = rx.next().await else {
1093                    return Ok(None);
1094                };
1095                Ok(Some((
1096                    frame,
1097                    Self {
1098                        rx,
1099                        tx,
1100                        driver: ProtocolPeerDriver::Completed {
1101                            finish: future.take_finish(),
1102                        },
1103                    },
1104                )))
1105            }
1106        }
1107    }
1108
1109    async fn next_message(self) -> Result<Option<(RawJsonRpcMessage, Self)>, crate::Error> {
1110        let Some((frame, peer)) = self.next_frame().await? else {
1111            return Ok(None);
1112        };
1113        Ok(Some((initialize_message(frame)?, peer)))
1114    }
1115
1116    fn send(&self, message: RawJsonRpcMessage) -> Result<(), crate::Error> {
1117        self.send_frame(TransportFrame::Single(message))
1118    }
1119
1120    fn send_frame(&self, frame: TransportFrame) -> Result<(), crate::Error> {
1121        self.tx
1122            .unbounded_send(frame)
1123            .map_err(crate::util::internal_error)
1124    }
1125}
1126
1127#[cfg(feature = "unstable_protocol_v2")]
1128fn initialize_message(frame: TransportFrame) -> Result<RawJsonRpcMessage, crate::Error> {
1129    match frame {
1130        TransportFrame::Single(message) => Ok(message),
1131        TransportFrame::Malformed { error, .. } => Err(error),
1132        TransportFrame::Batch(_) => Err(crate::Error::invalid_request()
1133            .data("ACP initialize request and response messages must be sent individually")),
1134    }
1135}
1136
1137#[cfg(feature = "unstable_protocol_v2")]
1138fn initialize_message_mut(
1139    frame: &mut TransportFrame,
1140) -> Result<Option<&mut RawJsonRpcMessage>, crate::Error> {
1141    match frame {
1142        TransportFrame::Single(RawJsonRpcMessage::Response(_)) => Ok(None),
1143        TransportFrame::Single(entry) => Ok(Some(entry)),
1144        TransportFrame::Malformed { raw, .. } if raw_is_response_only_shape(raw) => Ok(None),
1145        TransportFrame::Malformed { error, .. } => Err(error.clone()),
1146        TransportFrame::Batch(batch) => {
1147            for entry in batch.entries_mut() {
1148                match entry {
1149                    TransportBatchEntry::Message(RawJsonRpcMessage::Response(_)) => {}
1150                    TransportBatchEntry::Message(message) => return Ok(Some(message)),
1151                    TransportBatchEntry::Malformed { raw, .. } if is_response_only_shape(raw) => {}
1152                    TransportBatchEntry::Malformed { error, .. } => return Err(error.clone()),
1153                }
1154            }
1155            Ok(None)
1156        }
1157    }
1158}
1159
1160// Every protocol router uses the same ownership rule. Passive halves keep
1161// independent lifetimes; owned completion drains output and then either joins
1162// an opposed cooperative driver or cancels opaque work after polling errors.
1163#[cfg(feature = "unstable_protocol_v2")]
1164async fn pipe_protocol_peers_until_done(
1165    left: RunningProtocolPeer,
1166    right: RunningProtocolPeer,
1167) -> Result<(), crate::Error> {
1168    let mut left_driver = left.driver.into_driver();
1169    let mut right_driver = right.driver.into_driver();
1170    let left_passive = left_driver.is_none();
1171    let right_passive = right_driver.is_none();
1172    let left_finish = left_driver
1173        .as_mut()
1174        .and_then(crate::ConnectionDriver::take_finish);
1175    let right_finish = right_driver
1176        .as_mut()
1177        .and_then(crate::ConnectionDriver::take_finish);
1178    let (stop_left_tx, stop_left_rx) = futures::channel::oneshot::channel();
1179    let (stop_right_tx, stop_right_rx) = futures::channel::oneshot::channel();
1180    let stop = async |rx: futures::channel::oneshot::Receiver<()>| {
1181        if rx.await.is_err() {
1182            future::pending::<()>().await;
1183        }
1184    };
1185    let left_to_right = Box::pin(
1186        Channel {
1187            rx: left.rx,
1188            tx: right.tx,
1189        }
1190        .copy_with_driver_until(left_driver, stop(stop_left_rx)),
1191    );
1192    let right_to_left = Box::pin(
1193        Channel {
1194            rx: right.rx,
1195            tx: left.tx,
1196        }
1197        .copy_with_driver_until(right_driver, stop(stop_right_rx)),
1198    );
1199
1200    match future::select(left_to_right, right_to_left).await {
1201        future::Either::Left((result, right_to_left)) => {
1202            result?;
1203            if !left_passive {
1204                let _ = stop_right_tx.send(());
1205            }
1206            if left_passive || right_finish.is_some() {
1207                if !left_passive && let Some(mut finish) = right_finish {
1208                    finish.request();
1209                }
1210                right_to_left.await
1211            } else {
1212                // Without a cooperative finish hook, opposed work may remain
1213                // open indefinitely. Poll ready errors before cancelling it.
1214                crate::util::run_until(right_to_left, future::ready(Ok(()))).await
1215            }
1216        }
1217        future::Either::Right((result, left_to_right)) => {
1218            result?;
1219            if !right_passive {
1220                let _ = stop_left_tx.send(());
1221            }
1222            if right_passive || left_finish.is_some() {
1223                if !right_passive && let Some(mut finish) = left_finish {
1224                    finish.request();
1225                }
1226                left_to_right.await
1227            } else {
1228                crate::util::run_until(left_to_right, future::ready(Ok(()))).await
1229            }
1230        }
1231    }
1232}
1233
1234#[cfg(feature = "unstable_protocol_v2")]
1235#[derive(Debug)]
1236struct InitializeResponse {
1237    id: RequestId,
1238    result: Result<serde_json::Value, Box<RawJsonRpcError>>,
1239}
1240
1241#[cfg(feature = "unstable_protocol_v2")]
1242impl InitializeResponse {
1243    fn from_message(message: RawJsonRpcMessage) -> Result<Self, crate::Error> {
1244        match message {
1245            RawJsonRpcMessage::Response(RpcResponse::Result { id, result }) => Ok(Self {
1246                id,
1247                result: Ok(result),
1248            }),
1249            RawJsonRpcMessage::Response(RpcResponse::Error { id, error }) => Ok(Self {
1250                id,
1251                result: Err(error),
1252            }),
1253            message => Err(crate::Error::invalid_request().data(format!(
1254                "first ACP response must be an initialize response, got {message:?}",
1255            ))),
1256        }
1257    }
1258
1259    fn into_message(self) -> RawJsonRpcMessage {
1260        RawJsonRpcMessage::Response(RpcResponse::new(self.id, self.result))
1261    }
1262
1263    fn with_id(self, id: RequestId) -> RawJsonRpcMessage {
1264        RawJsonRpcMessage::Response(RpcResponse::new(id, self.result))
1265    }
1266
1267    fn protocol_version(&self) -> Option<ProtocolVersion> {
1268        serde_json::from_value(self.result.as_ref().ok()?.get("protocolVersion")?.clone()).ok()
1269    }
1270}
1271
1272#[cfg(all(test, feature = "unstable_protocol_v2"))]
1273mod raw_initialize_tests {
1274    use super::*;
1275    use serde_json::json;
1276
1277    #[test]
1278    fn initialize_error_forwarding_preserves_raw_fields() {
1279        for data in [
1280            None,
1281            Some(serde_json::Value::Null),
1282            Some(json!({"detail":"kept"})),
1283        ] {
1284            let mut error = json!({
1285                "code":-32000, "message":"peer", "extension":{"retry":true}
1286            });
1287            if let Some(data) = data {
1288                error["data"] = data;
1289            }
1290            let wire = json!({"jsonrpc":"2.0", "id":"original", "error":error});
1291            let response =
1292                InitializeResponse::from_message(serde_json::from_value(wire.clone()).unwrap())
1293                    .unwrap();
1294            assert_eq!(serde_json::to_value(response.into_message()).unwrap(), wire);
1295            let response =
1296                InitializeResponse::from_message(serde_json::from_value(wire).unwrap()).unwrap();
1297            let forwarded = response.with_id(RequestId::Str("replacement".into()));
1298            assert_eq!(
1299                serde_json::to_value(forwarded).unwrap(),
1300                json!({"jsonrpc":"2.0", "id":"replacement", "error":error})
1301            );
1302        }
1303    }
1304}
1305
1306#[cfg(feature = "unstable_protocol_v2")]
1307#[derive(Debug, Clone, Copy, PartialEq, Eq)]
1308enum ClientProtocol {
1309    V1,
1310    V2,
1311}
1312
1313#[cfg(feature = "unstable_protocol_v2")]
1314impl ClientProtocol {
1315    fn name(self) -> &'static str {
1316        match self {
1317            Self::V1 => "1",
1318            Self::V2 => "2",
1319        }
1320    }
1321}
1322
1323#[cfg(feature = "unstable_protocol_v2")]
1324#[derive(Debug, Clone, Copy, PartialEq, Eq)]
1325struct SupportedClientProtocols {
1326    v1: bool,
1327    v2: bool,
1328}
1329
1330#[cfg(feature = "unstable_protocol_v2")]
1331impl SupportedClientProtocols {
1332    fn highest_configured(self) -> Option<ClientProtocol> {
1333        if self.v2 {
1334            return Some(ClientProtocol::V2);
1335        }
1336
1337        if self.v1 {
1338            return Some(ClientProtocol::V1);
1339        }
1340
1341        None
1342    }
1343}
1344
1345#[cfg(feature = "unstable_protocol_v2")]
1346async fn start_client_protocol(
1347    protocol: ClientProtocol,
1348    client: DynConnectTo<Agent>,
1349) -> Result<(RunningProtocolPeer, RawJsonRpcMessage), crate::Error> {
1350    let client = RunningProtocolPeer::new(client);
1351    let Some((initialize, client)) = client.next_message().await? else {
1352        return Err(crate::Error::invalid_request().data(format!(
1353            "ACP protocol version {} client implementation ended before initialize",
1354            protocol.name()
1355        )));
1356    };
1357    ensure_client_initialize_request(protocol, &initialize)?;
1358    Ok((client, initialize))
1359}
1360
1361#[cfg(feature = "unstable_protocol_v2")]
1362async fn send_initialize_and_receive(
1363    client: RunningProtocolPeer,
1364    agent: RunningProtocolPeer,
1365    initialize: RawJsonRpcMessage,
1366) -> Result<(RunningProtocolPeer, RunningProtocolPeer, InitializeResponse), crate::Error> {
1367    agent.send(initialize)?;
1368    let Some((response, agent)) = agent.next_message().await? else {
1369        return Err(crate::Error::internal_error().data("agent closed before initialize response"));
1370    };
1371    let response = InitializeResponse::from_message(response)?;
1372    Ok((client, agent, response))
1373}
1374
1375#[cfg(feature = "unstable_protocol_v2")]
1376async fn initialize_client_protocol(
1377    protocol: ClientProtocol,
1378    client: DynConnectTo<Agent>,
1379    agent: impl ConnectTo<Client>,
1380) -> Result<(RunningProtocolPeer, RunningProtocolPeer, InitializeResponse), crate::Error> {
1381    let agent = RunningProtocolPeer::new(agent);
1382    let (client, initialize) = start_client_protocol(protocol, client).await?;
1383    send_initialize_and_receive(client, agent, initialize).await
1384}
1385
1386#[cfg(feature = "unstable_protocol_v2")]
1387async fn connect_client_protocol(
1388    protocol: ClientProtocol,
1389    client: DynConnectTo<Agent>,
1390    agent: impl ConnectTo<Client>,
1391) -> Result<(), crate::Error> {
1392    let (client, agent, initialize_response) =
1393        initialize_client_protocol(protocol, client, agent).await?;
1394    client.send(initialize_response.into_message())?;
1395    pipe_protocol_peers_until_done(client, agent).await
1396}
1397
1398#[cfg(feature = "unstable_protocol_v2")]
1399fn ensure_client_initialize_request(
1400    protocol: ClientProtocol,
1401    message: &RawJsonRpcMessage,
1402) -> Result<(), crate::Error> {
1403    let RawJsonRpcMessage::Request(request) = message else {
1404        return Err(crate::Error::invalid_request().data(format!(
1405            "ACP protocol version {} client implementation must send initialize first",
1406            protocol.name()
1407        )));
1408    };
1409
1410    if request.method.as_ref() != "initialize" {
1411        return Err(crate::Error::invalid_request().data(format!(
1412            "ACP protocol version {} client implementation must send initialize first",
1413            protocol.name()
1414        )));
1415    }
1416
1417    Ok(())
1418}
1419
1420#[cfg(feature = "unstable_protocol_v2")]
1421fn initialize_request_id(message: &RawJsonRpcMessage) -> Option<RequestId> {
1422    let RawJsonRpcMessage::Request(request) = message else {
1423        return None;
1424    };
1425    Some(request.id.clone())
1426}
1427
1428#[cfg(feature = "unstable_protocol_v2")]
1429fn initialize_response_negotiated_v1(response: &InitializeResponse) -> bool {
1430    response.protocol_version() == Some(ProtocolVersion::V1)
1431}
1432
1433impl HasPeer<Agent> for Agent {
1434    fn remote_style(&self, _peer: Agent) -> RemoteStyle {
1435        RemoteStyle::Counterpart
1436    }
1437}
1438
1439/// The proxy role - an intermediary that can intercept and modify messages.
1440///
1441/// Proxies sit between a client and an agent (or another proxy), and can:
1442/// - Add tools via MCP servers
1443/// - Filter or transform messages
1444/// - Inject additional context
1445///
1446/// Proxies connect to a [`Conductor`] which orchestrates the proxy chain.
1447#[derive(Debug, Default, Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
1448pub struct Proxy;
1449
1450impl Role for Proxy {
1451    type Counterpart = Conductor;
1452
1453    fn default_handle_dispatch_from(
1454        &self,
1455        message: crate::Dispatch,
1456        _connection: crate::ConnectionTo<Self>,
1457    ) -> impl Future<Output = Result<crate::Handled<crate::Dispatch>, crate::Error>> + Send {
1458        std::future::ready(Ok(Handled::No {
1459            message,
1460            retry: false,
1461        }))
1462    }
1463
1464    fn role_id(&self) -> RoleId {
1465        RoleId::from_singleton(self)
1466    }
1467
1468    fn counterpart(&self) -> Self::Counterpart {
1469        Conductor
1470    }
1471}
1472
1473impl Proxy {
1474    /// Create a stable protocol v1 connection builder for a proxy.
1475    ///
1476    /// Use `Proxy::v2` for a protocol-v2-only proxy with typed callbacks and
1477    /// wire validation. Protocol-routing infrastructure that deliberately
1478    /// selects a version itself can disable the guard with
1479    /// `Builder::without_acp_version_guard`.
1480    pub fn builder(self) -> Builder<Proxy, NullHandler, NullRun> {
1481        Builder::new(self)
1482    }
1483
1484    /// Create a proxy builder that uses the ACP protocol v2 API.
1485    ///
1486    /// This builder requires `_proxy/initialize` to select protocol v2.
1487    /// Fluent callbacks receive [`crate::V2ConnectionTo<Conductor>`], while
1488    /// low-level custom handlers and runners retain the protocol-neutral
1489    /// [`ConnectionTo`] interface.
1490    ///
1491    /// Requires the `unstable_protocol_v2` crate feature.
1492    #[cfg(feature = "unstable_protocol_v2")]
1493    pub fn v2(self) -> V2Builder<Proxy, NullHandler, NullRun> {
1494        self.builder().v2_proxy()
1495    }
1496
1497    /// Create a router that chooses between configured proxy implementations.
1498    ///
1499    /// Add implementations with [`ProxyProtocolRouter::with_v1`] and
1500    /// [`ProxyProtocolRouter::with_v2`]. The router reads the initial
1501    /// `_proxy/initialize` request, selects the implementation for that exact
1502    /// protocol version, and hands over the complete initial transport frame.
1503    /// It does not downgrade proxy traffic or convert later messages.
1504    ///
1505    /// Requires the `unstable_protocol_v2` crate feature while protocol v2
1506    /// stabilizes.
1507    #[cfg(feature = "unstable_protocol_v2")]
1508    #[must_use]
1509    pub fn protocol_router(self) -> ProxyProtocolRouter {
1510        ProxyProtocolRouter::new()
1511    }
1512}
1513
1514/// Proxy component that routes each connection to a configured protocol implementation.
1515///
1516/// Use [`Proxy::protocol_router`] to start the builder, then add stable-v1 and
1517/// draft-v2 proxy implementations independently. Unlike
1518/// [`AgentProtocolRouter`], this router requires an exact version match: the
1519/// conductor has already selected and canonicalized the wire protocol before
1520/// sending `_proxy/initialize` to a proxy.
1521#[cfg(feature = "unstable_protocol_v2")]
1522#[derive(Debug, Default)]
1523pub struct ProxyProtocolRouter {
1524    v1: Option<DynConnectTo<Conductor>>,
1525    v2: Option<DynConnectTo<Conductor>>,
1526}
1527
1528#[cfg(feature = "unstable_protocol_v2")]
1529impl ProxyProtocolRouter {
1530    /// Create an empty proxy protocol router.
1531    #[must_use]
1532    pub fn new() -> Self {
1533        Self::default()
1534    }
1535
1536    /// Return this router with a stable ACP v1 proxy implementation.
1537    #[must_use]
1538    pub fn with_v1(mut self, proxy: impl ConnectTo<Conductor>) -> Self {
1539        self.v1 = Some(DynConnectTo::new(proxy));
1540        self
1541    }
1542
1543    /// Return this router with a draft ACP v2 proxy implementation.
1544    #[must_use]
1545    pub fn with_v2(mut self, proxy: impl ConnectTo<Conductor>) -> Self {
1546        self.v2 = Some(DynConnectTo::new(proxy));
1547        self
1548    }
1549}
1550
1551#[cfg(feature = "unstable_protocol_v2")]
1552impl ConnectTo<Conductor> for ProxyProtocolRouter {
1553    async fn connect_to(self, conductor: impl ConnectTo<Proxy>) -> Result<(), crate::Error> {
1554        let supported = SupportedProtocols {
1555            v1: self.v1.is_some(),
1556            v2: self.v2.is_some(),
1557        };
1558        let mut conductor = RunningProtocolPeer::new(conductor);
1559        let (first_frame, conductor, selected) = loop {
1560            let Some((mut frame, next_conductor)) = conductor.next_frame().await? else {
1561                return Ok(());
1562            };
1563            let message = match initialize_message_mut(&mut frame) {
1564                Ok(Some(message)) => message,
1565                Ok(None) => {
1566                    conductor = next_conductor;
1567                    continue;
1568                }
1569                Err(error) => return reject_initialize(next_conductor, &frame, error).await,
1570            };
1571            let selected = match select_proxy_protocol(message, supported) {
1572                Ok(selected) => selected,
1573                Err(error) => return reject_initialize(next_conductor, &frame, error).await,
1574            };
1575            break (frame, next_conductor, selected);
1576        };
1577        let Some(proxy) = selected.take_proxy(self) else {
1578            let error = selected.unsupported_error(supported);
1579            return reject_initialize(conductor, &first_frame, error).await;
1580        };
1581
1582        let proxy = RunningProtocolPeer::new(proxy);
1583        proxy.send_frame(first_frame)?;
1584        pipe_protocol_peers_until_done(conductor, proxy).await
1585    }
1586}
1587
1588#[cfg(feature = "unstable_protocol_v2")]
1589impl SelectedProtocol {
1590    fn take_proxy(self, proxy: ProxyProtocolRouter) -> Option<DynConnectTo<Conductor>> {
1591        match self {
1592            Self::V1 => proxy.v1,
1593            Self::V2 => proxy.v2,
1594        }
1595    }
1596}
1597
1598#[cfg(feature = "unstable_protocol_v2")]
1599fn select_proxy_protocol(
1600    message: &RawJsonRpcMessage,
1601    supported: SupportedProtocols,
1602) -> Result<SelectedProtocol, crate::Error> {
1603    let RawJsonRpcMessage::Request(request) = message else {
1604        return Err(crate::Error::invalid_request()
1605            .data("first ACP proxy message must be an `_proxy/initialize` request"));
1606    };
1607
1608    if request.method.as_ref() != METHOD_INITIALIZE_PROXY {
1609        return Err(crate::Error::invalid_request()
1610            .data("first ACP proxy request must be `_proxy/initialize`"));
1611    }
1612
1613    let Some(RawJsonRpcParams::Object(params)) = &request.params else {
1614        return Err(invalid_initialize_protocol_version());
1615    };
1616    let Some(protocol_version) = params.get("protocolVersion") else {
1617        return Err(invalid_initialize_protocol_version());
1618    };
1619    let requested = serde_json::from_value::<ProtocolVersion>(protocol_version.clone())
1620        .map_err(|_| invalid_initialize_protocol_version())?;
1621    let selected = supported.exact(requested).ok_or_else(|| {
1622        crate::Error::invalid_request().data(format!(
1623            "unsupported ACP protocol version {requested}; this proxy supports {}",
1624            supported.description()
1625        ))
1626    })?;
1627
1628    match selected {
1629        SelectedProtocol::V1 => {
1630            parse_initialize_params::<crate::schema::InitializeProxyRequest>(params)?;
1631        }
1632        SelectedProtocol::V2 => {
1633            parse_initialize_params::<v2::InitializeProxyRequest>(params)?;
1634        }
1635    }
1636    Ok(selected)
1637}
1638
1639impl HasPeer<Proxy> for Proxy {
1640    fn remote_style(&self, _peer: Proxy) -> RemoteStyle {
1641        RemoteStyle::Counterpart
1642    }
1643}
1644
1645/// The conductor role - orchestrates proxy chains.
1646///
1647/// Conductors manage connections between clients, proxies, and agents,
1648/// routing messages through the appropriate proxy chain.
1649#[derive(Debug, Default, Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
1650pub struct Conductor;
1651
1652impl Role for Conductor {
1653    type Counterpart = Proxy;
1654
1655    fn role_id(&self) -> RoleId {
1656        RoleId::from_singleton(self)
1657    }
1658
1659    fn counterpart(&self) -> Self::Counterpart {
1660        Proxy
1661    }
1662
1663    async fn default_handle_dispatch_from(
1664        &self,
1665        message: Dispatch,
1666        cx: ConnectionTo<Conductor>,
1667    ) -> Result<Handled<Dispatch>, crate::Error> {
1668        #[cfg(not(feature = "unstable_protocol_v2"))]
1669        {
1670            MatchDispatchFrom::new(message, &cx)
1671                .if_request_from(Client, async |_req: InitializeRequest, responder| {
1672                    responder.respond_with_error(crate::Error::invalid_request().data(format!(
1673                        "proxies must be initialized with `{METHOD_INITIALIZE_PROXY}`"
1674                    )))
1675                })
1676                .await
1677                .if_request_from(
1678                    Client,
1679                    async |request: InitializeProxyRequest, responder| {
1680                        let InitializeProxyRequest { initialize } = request;
1681                        cx.send_ordered_request_to(Agent, initialize)
1682                            .forward_response_to(responder)
1683                    },
1684                )
1685                .await
1686                .if_request_from(Client, async |request: NewSessionRequest, responder| {
1687                    let sent = cx.send_ordered_request_to(Agent, request);
1688                    let sent = sent.forward_cancellation_from(responder.cancellation());
1689                    sent.on_receiving_result({
1690                        let cx = cx.clone();
1691                        async move |result| {
1692                            if let Ok(NewSessionResponse { session_id, .. }) = &result {
1693                                cx.add_dynamic_handler(ProxySessionMessages::new(
1694                                    session_id.clone(),
1695                                ))?
1696                                .detach();
1697                            }
1698                            responder.respond_with_result(result)
1699                        }
1700                    })
1701                })
1702                .await
1703                .if_dispatch_from(Client, async |message: Dispatch| {
1704                    cx.send_proxied_message_to(Agent, message)
1705                })
1706                .await
1707                .if_dispatch_from(Agent, async |message: Dispatch| {
1708                    cx.send_proxied_message_to(Client, message)
1709                })
1710                .await
1711                .done()
1712        }
1713
1714        #[cfg(feature = "unstable_protocol_v2")]
1715        {
1716            let message = match message {
1717                Dispatch::Request(request, responder) if request.method() == "initialize" => {
1718                    responder.respond_with_error(crate::Error::invalid_request().data(format!(
1719                        "proxies must be initialized with `{METHOD_INITIALIZE_PROXY}`"
1720                    )))?;
1721                    return Ok(Handled::Yes);
1722                }
1723                Dispatch::Request(mut request, responder)
1724                    if request.method() == METHOD_INITIALIZE_PROXY =>
1725                {
1726                    request.method = "initialize".to_string();
1727                    cx.send_ordered_request_to(Agent, request)
1728                        .forward_response_to(responder)?;
1729                    return Ok(Handled::Yes);
1730                }
1731                Dispatch::Request(request, responder) if request.method() == "session/new" => {
1732                    let sent = cx.send_ordered_request_to(Agent, request);
1733                    // The dynamic-handler hook below means we cannot use
1734                    // `forward_response_to`, so wire up cancellation forwarding
1735                    // explicitly to keep `session/new` cancellable like every
1736                    // other proxied request.
1737                    let sent = sent.forward_cancellation_from(responder.cancellation());
1738                    sent.on_receiving_result({
1739                        let cx = cx.clone();
1740                        async move |result| {
1741                            let result = result.and_then(|response| {
1742                                let envelope: NewSessionResponseEnvelope =
1743                                    crate::util::json_cast(response.clone())?;
1744                                cx.add_dynamic_handler(ProxySessionMessages::new(
1745                                    envelope.session_id,
1746                                ))?
1747                                .detach();
1748                                Ok(response)
1749                            });
1750                            responder.respond_with_result(result)
1751                        }
1752                    })?;
1753                    return Ok(Handled::Yes);
1754                }
1755                message => message,
1756            };
1757
1758            MatchDispatchFrom::new(message, &cx)
1759                .if_dispatch_from(Client, async |message: Dispatch| {
1760                    cx.send_proxied_message_to(Agent, message)
1761                })
1762                .await
1763                .if_dispatch_from(Agent, async |message: Dispatch| {
1764                    cx.send_proxied_message_to(Client, message)
1765                })
1766                .await
1767                .done()
1768        }
1769    }
1770}
1771
1772impl Conductor {
1773    /// Create a connection builder for a conductor.
1774    pub fn builder(self) -> Builder<Conductor, NullHandler, NullRun> {
1775        Builder::new(self)
1776    }
1777}
1778
1779impl HasPeer<Client> for Conductor {
1780    fn remote_style(&self, _peer: Client) -> RemoteStyle {
1781        RemoteStyle::Predecessor
1782    }
1783}
1784
1785impl HasPeer<Agent> for Conductor {
1786    fn remote_style(&self, _peer: Agent) -> RemoteStyle {
1787        RemoteStyle::Successor
1788    }
1789}
1790
1791/// Dynamic handler that proxies session messages from Agent to Client.
1792///
1793/// This is used internally to handle session message routing after a
1794/// `session.new` request has been forwarded.
1795pub(crate) struct ProxySessionMessages {
1796    session_id: SessionId,
1797}
1798
1799impl ProxySessionMessages {
1800    /// Create a new proxy handler for the given session.
1801    pub fn new(session_id: SessionId) -> Self {
1802        Self { session_id }
1803    }
1804}
1805
1806impl<Counterpart: Role> HandleDispatchFrom<Counterpart> for ProxySessionMessages
1807where
1808    Counterpart: HasPeer<Agent> + HasPeer<Client>,
1809{
1810    async fn handle_dispatch_from(
1811        &mut self,
1812        message: Dispatch,
1813        connection: ConnectionTo<Counterpart>,
1814    ) -> Result<Handled<Dispatch>, crate::Error> {
1815        MatchDispatchFrom::new(message, &connection)
1816            .if_dispatch_from(Agent, async |message| {
1817                // If this is for our session-id, proxy it to the client.
1818                if let Some(session_id) = message.get_session_id()?
1819                    && session_id == self.session_id
1820                {
1821                    connection.send_proxied_message_to(Client, message)?;
1822                    return Ok(Handled::Yes);
1823                }
1824
1825                // Otherwise, leave it alone.
1826                Ok(Handled::No {
1827                    message,
1828                    retry: false,
1829                })
1830            })
1831            .await
1832            .done()
1833    }
1834
1835    fn describe_chain(&self) -> impl std::fmt::Debug {
1836        format!("ProxySessionMessages({})", self.session_id)
1837    }
1838}
1839
1840#[cfg(all(test, feature = "unstable_protocol_v2"))]
1841mod lifetime_tests {
1842    use super::*;
1843    use crate::{ConnectionDriver, UntypedRole};
1844    use futures::FutureExt as _;
1845
1846    fn frame() -> TransportFrame {
1847        TransportFrame::parse_json(r#"{"jsonrpc":"2.0","method":"test/queued","params":{}}"#)
1848    }
1849
1850    #[tokio::test]
1851    async fn passive_initialization_waits_for_a_frame_not_driver_readiness() {
1852        let (channel, remote) = Channel::duplex();
1853        let peer = RunningProtocolPeer::new::<UntypedRole>(channel);
1854        let mut next = Box::pin(peer.next_frame());
1855        assert!(next.as_mut().now_or_never().is_none());
1856
1857        remote.tx.unbounded_send(frame()).unwrap();
1858        let (_, peer) = next.await.unwrap().expect("passive peer remains connected");
1859        assert!(matches!(peer.driver, ProtocolPeerDriver::Passive));
1860    }
1861
1862    #[tokio::test]
1863    async fn owned_peer_preserves_finish_metadata_through_active_and_completed_states() {
1864        let (Channel { rx, tx }, remote) = Channel::duplex();
1865        let (done_tx, done_rx) = futures::channel::oneshot::channel();
1866        let (finish_tx, finish_rx) = futures::channel::oneshot::channel();
1867        let peer = RunningProtocolPeer {
1868            rx,
1869            tx,
1870            driver: ProtocolPeerDriver::Active(ConnectionDriver::with_finish(
1871                async move {
1872                    done_rx.await.unwrap();
1873                    Ok(())
1874                },
1875                move || {
1876                    let _ = finish_tx.send(());
1877                },
1878            )),
1879        };
1880
1881        remote.tx.unbounded_send(frame()).unwrap();
1882        let (_, peer) = peer.next_frame().await.unwrap().unwrap();
1883        assert!(matches!(&peer.driver, ProtocolPeerDriver::Active(_)));
1884
1885        remote.tx.unbounded_send(frame()).unwrap();
1886        done_tx.send(()).unwrap();
1887        let (_, peer) = peer.next_frame().await.unwrap().unwrap();
1888        assert!(matches!(
1889            &peer.driver,
1890            ProtocolPeerDriver::Completed { finish: Some(_) }
1891        ));
1892        let mut driver = peer
1893            .driver
1894            .into_driver()
1895            .expect("completed owned work must not become passive");
1896        assert!(driver.request_finish());
1897        finish_rx.await.unwrap();
1898        driver.await.unwrap();
1899    }
1900
1901    #[tokio::test]
1902    async fn active_initialization_drains_accepted_frames_without_escaped_sender_eof() {
1903        let (Channel { rx, tx }, remote) = Channel::duplex();
1904        remote.tx.unbounded_send(frame()).unwrap();
1905        remote.tx.unbounded_send(frame()).unwrap();
1906        let mut polls = 0;
1907        let driver = ConnectionDriver::new(future::poll_fn(move |_| {
1908            polls += 1;
1909            assert_eq!(polls, 1, "the completed driver must never be re-polled");
1910            std::task::Poll::Ready(Ok(()))
1911        }));
1912        let peer = RunningProtocolPeer {
1913            rx,
1914            tx,
1915            driver: ProtocolPeerDriver::Active(driver),
1916        };
1917
1918        let (_, peer) = peer.next_frame().await.unwrap().unwrap();
1919        assert!(
1920            remote.tx.unbounded_send(frame()).is_err(),
1921            "active completion must reject new output from escaped handles"
1922        );
1923        let (_, peer) = peer.next_frame().await.unwrap().unwrap();
1924        assert!(peer.next_frame().await.unwrap().is_none());
1925    }
1926
1927    #[tokio::test]
1928    async fn ready_initialization_driver_error_beats_queued_frames() {
1929        let (Channel { rx, tx }, remote) = Channel::duplex();
1930        remote.tx.unbounded_send(frame()).unwrap();
1931        let error = crate::Error::internal_error().data("owned initialization failed");
1932        let peer = RunningProtocolPeer {
1933            rx,
1934            tx,
1935            driver: ProtocolPeerDriver::Active(ConnectionDriver::new(future::ready(Err(
1936                error.clone()
1937            )))),
1938        };
1939
1940        match peer.next_frame().await {
1941            Err(actual) => assert_eq!(actual, error),
1942            Ok(_) => panic!("ready driver error must not be hidden by a queued frame"),
1943        }
1944    }
1945
1946    struct QueuedFinalClient;
1947
1948    impl ConnectTo<Agent> for QueuedFinalClient {
1949        async fn connect_to(self, agent: impl ConnectTo<Client>) -> Result<(), crate::Error> {
1950            let (mut channel, driver) = agent.into_channel_and_future();
1951            let foreground = async move {
1952                channel
1953                    .tx
1954                    .unbounded_send(TransportFrame::Single(RawJsonRpcMessage::request(
1955                        "initialize".into(),
1956                        serde_json::json!({ "protocolVersion": 1, "clientCapabilities": {} }),
1957                        RequestId::Number(1),
1958                    )?))
1959                    .unwrap();
1960                assert!(channel.rx.next().await.is_some(), "initialize response");
1961                // Exceed the physical writer capacity so only concurrent polling
1962                // of the sink can make this bridge finish.
1963                for index in 0..3 {
1964                    channel
1965                        .tx
1966                        .unbounded_send(TransportFrame::Single(RawJsonRpcMessage::notification(
1967                            "test/final".into(),
1968                            serde_json::json!({ "index": index, "payload": "x".repeat(1024) }),
1969                        )?))
1970                        .unwrap();
1971                }
1972                Ok(())
1973            };
1974            match driver {
1975                Some(driver) => crate::util::run_until(driver, foreground).await,
1976                None => foreground.await,
1977            }
1978        }
1979    }
1980
1981    #[tokio::test]
1982    async fn connector_completion_flushes_byte_streams_without_remote_read_eof() {
1983        use tokio::io::{AsyncBufReadExt as _, AsyncWriteExt as _, BufReader};
1984        use tokio_util::compat::{TokioAsyncReadCompatExt as _, TokioAsyncWriteCompatExt as _};
1985
1986        let (writer, remote_reader) = tokio::io::duplex(64);
1987        let (mut remote_input, reader) = tokio::io::duplex(64);
1988        let physical = crate::ByteStreams::new(writer.compat_write(), reader.compat());
1989        let mut physical = Some(crate::DynConnectTo::<Client>::new(physical));
1990        let connector = tokio::spawn(
1991            ClientProtocolConnector::new()
1992                .with_v1(|| QueuedFinalClient)
1993                .connect_to(move || physical.take().expect("one physical connection")),
1994        );
1995        let mut lines = BufReader::new(remote_reader).lines();
1996
1997        tokio::time::timeout(std::time::Duration::from_secs(5), async {
1998            let initialize = lines
1999                .next_line()
2000                .await
2001                .unwrap()
2002                .expect("initialize request");
2003            let value: serde_json::Value = serde_json::from_str(&initialize).unwrap();
2004            assert_eq!(value["method"], "initialize");
2005            remote_input
2006                .write_all(b"{\"jsonrpc\":\"2.0\",\"id\":1,\"result\":{\"protocolVersion\":1}}\n")
2007                .await
2008                .unwrap();
2009            for index in 0..3 {
2010                let line = lines
2011                    .next_line()
2012                    .await
2013                    .unwrap()
2014                    .expect("accepted final output");
2015                let value: serde_json::Value = serde_json::from_str(&line).unwrap();
2016                assert_eq!(value["params"]["index"], index);
2017                assert_eq!(value["params"]["payload"].as_str().unwrap().len(), 1024);
2018            }
2019            connector.await.unwrap().unwrap();
2020            assert!(lines.next_line().await.unwrap().is_none());
2021        })
2022        .await
2023        .expect("physical flush must not wait for independent remote input EOF");
2024        drop(remote_input);
2025    }
2026
2027    #[derive(Default, Debug)]
2028    struct GatedLineSinkState {
2029        pending: Vec<String>,
2030        flushed: Vec<String>,
2031        closed: bool,
2032        dropped: bool,
2033    }
2034
2035    struct GatedLineSink {
2036        state: std::sync::Arc<std::sync::Mutex<GatedLineSinkState>>,
2037        release: futures::channel::oneshot::Receiver<()>,
2038        released: bool,
2039    }
2040
2041    impl futures::Sink<String> for GatedLineSink {
2042        type Error = std::io::Error;
2043
2044        fn poll_ready(
2045            self: std::pin::Pin<&mut Self>,
2046            _: &mut std::task::Context<'_>,
2047        ) -> std::task::Poll<Result<(), Self::Error>> {
2048            std::task::Poll::Ready(Ok(()))
2049        }
2050
2051        fn start_send(self: std::pin::Pin<&mut Self>, line: String) -> Result<(), Self::Error> {
2052            self.state.lock().unwrap().pending.push(line);
2053            Ok(())
2054        }
2055
2056        fn poll_flush(
2057            mut self: std::pin::Pin<&mut Self>,
2058            cx: &mut std::task::Context<'_>,
2059        ) -> std::task::Poll<Result<(), Self::Error>> {
2060            if !self.released {
2061                if std::pin::Pin::new(&mut self.release).poll(cx).is_pending() {
2062                    return std::task::Poll::Pending;
2063                }
2064                self.released = true;
2065            }
2066            let mut state = self.state.lock().unwrap();
2067            let pending = std::mem::take(&mut state.pending);
2068            state.flushed.extend(pending);
2069            std::task::Poll::Ready(Ok(()))
2070        }
2071
2072        fn poll_close(
2073            mut self: std::pin::Pin<&mut Self>,
2074            cx: &mut std::task::Context<'_>,
2075        ) -> std::task::Poll<Result<(), Self::Error>> {
2076            match self.as_mut().poll_flush(cx) {
2077                std::task::Poll::Ready(Ok(())) => {
2078                    self.state.lock().unwrap().closed = true;
2079                    std::task::Poll::Ready(Ok(()))
2080                }
2081                result => result,
2082            }
2083        }
2084    }
2085
2086    impl Drop for GatedLineSink {
2087        fn drop(&mut self) {
2088            self.state.lock().unwrap().dropped = true;
2089        }
2090    }
2091
2092    #[tokio::test]
2093    async fn foreground_completion_flushes_lines_with_already_normalized_remote_input() {
2094        let state = std::sync::Arc::new(std::sync::Mutex::new(GatedLineSinkState::default()));
2095        let (release_tx, release_rx) = futures::channel::oneshot::channel();
2096        let sink = GatedLineSink {
2097            state: state.clone(),
2098            release: release_rx,
2099            released: false,
2100        };
2101        let (remote_input, incoming) = futures::channel::mpsc::unbounded();
2102        remote_input
2103            .unbounded_send(Ok(
2104                r#"{"jsonrpc":"2.0","id":1,"result":{"protocolVersion":1}}"#.to_string(),
2105            ))
2106            .unwrap();
2107        remote_input
2108            .unbounded_send(Ok(
2109                r#"{"jsonrpc":"2.0","method":"test/queued","params":{}}"#.to_string(),
2110            ))
2111            .unwrap();
2112        let physical = RunningProtocolPeer::new::<Client>(crate::Lines::new(sink, incoming));
2113        // The real Lines driver reads both ready lines before this returns.
2114        let (initialize, physical) = physical.next_frame().await.unwrap().unwrap();
2115        assert_eq!(
2116            futures::Stream::size_hint(&physical.rx).0,
2117            1,
2118            "trailing input must already be in the original normalized queue"
2119        );
2120
2121        let (Channel { rx, tx }, mut local) = Channel::duplex();
2122        tx.unbounded_send(initialize).unwrap();
2123        let foreground = RunningProtocolPeer {
2124            rx,
2125            tx,
2126            driver: ProtocolPeerDriver::Active(ConnectionDriver::new(async move {
2127                assert!(local.rx.next().await.is_some(), "initialize response");
2128                for index in 0..3 {
2129                    local
2130                        .tx
2131                        .unbounded_send(TransportFrame::Single(RawJsonRpcMessage::notification(
2132                            "test/final".into(),
2133                            serde_json::json!({ "index": index, "payload": "x".repeat(1024) }),
2134                        )?))
2135                        .unwrap();
2136                }
2137                // This owns the foreground input receiver: completion closes it.
2138                drop(local);
2139                Ok(())
2140            })),
2141        };
2142        let mut bridge = Box::pin(pipe_protocol_peers_until_done(foreground, physical));
2143        let early_result = bridge.as_mut().now_or_never();
2144        assert!(
2145            !state.lock().unwrap().pending.is_empty(),
2146            "the real physical writer must accept output before the gated flush"
2147        );
2148        // Never use time to establish the race: only the sink gate controls drain.
2149        let _released = release_tx.send(());
2150        let result = match early_result {
2151            Some(result) => result,
2152            None => tokio::time::timeout(std::time::Duration::from_secs(1), bridge)
2153                .await
2154                .expect("physical drain must not require remote input EOF"),
2155        };
2156        let state = state.lock().unwrap();
2157        assert!(
2158            result.is_ok() && state.flushed.len() == 3 && state.closed,
2159            "accepted output must drain cleanly despite queued remote input: result={result:?}, pending={}, flushed={}, closed={}, dropped={}",
2160            state.pending.len(),
2161            state.flushed.len(),
2162            state.closed,
2163            state.dropped,
2164        );
2165        for (index, line) in state.flushed.iter().enumerate() {
2166            let value: serde_json::Value = serde_json::from_str(line).unwrap();
2167            assert_eq!(value["method"], "test/final");
2168            assert_eq!(value["params"]["index"], index);
2169            assert_eq!(value["params"]["payload"].as_str().unwrap().len(), 1024);
2170        }
2171        drop(remote_input);
2172    }
2173
2174    #[tokio::test]
2175    async fn foreground_completion_keeps_read_errors_during_lines_drain() {
2176        let state = std::sync::Arc::new(std::sync::Mutex::new(GatedLineSinkState::default()));
2177        let (_release_tx, release_rx) = futures::channel::oneshot::channel();
2178        let sink = GatedLineSink {
2179            state: state.clone(),
2180            release: release_rx,
2181            released: false,
2182        };
2183        let (remote_input, incoming) = futures::channel::mpsc::unbounded();
2184        let physical = RunningProtocolPeer::new::<Client>(crate::Lines::new(sink, incoming));
2185        let (Channel { rx, tx }, local) = Channel::duplex();
2186        local.tx.unbounded_send(frame()).unwrap();
2187        drop(local);
2188        let foreground = RunningProtocolPeer {
2189            rx,
2190            tx,
2191            driver: ProtocolPeerDriver::Active(ConnectionDriver::new(future::ready(Ok(())))),
2192        };
2193        let mut bridge = Box::pin(pipe_protocol_peers_until_done(foreground, physical));
2194        assert!(bridge.as_mut().now_or_never().is_none());
2195        assert_eq!(state.lock().unwrap().pending.len(), 1);
2196
2197        // Newly read successful input is irrelevant to the completed foreground,
2198        // but a genuine read failure must still cancel the blocked sink drain.
2199        remote_input
2200            .unbounded_send(Ok(r#"{"jsonrpc":"2.0","method":"late"}"#.to_string()))
2201            .unwrap();
2202        remote_input
2203            .unbounded_send(Err(std::io::Error::other("read failed after foreground")))
2204            .unwrap();
2205        let error = tokio::time::timeout(std::time::Duration::from_secs(1), bridge)
2206            .await
2207            .expect("read failure must not wait for the sink gate")
2208            .expect_err("real read error must win over foreground success");
2209        assert!(
2210            error
2211                .data
2212                .unwrap()
2213                .to_string()
2214                .contains("read failed after foreground"),
2215            "the read error must not become a receiver-gone forwarding error"
2216        );
2217        assert_eq!(state.lock().unwrap().flushed, Vec::<String>::new());
2218    }
2219
2220    #[tokio::test]
2221    async fn foreground_completion_cancels_opposed_opaque_work_in_both_directions() {
2222        for foreground_on_left in [true, false] {
2223            let (Channel { rx, tx }, foreground_remote) = Channel::duplex();
2224            foreground_remote.tx.unbounded_send(frame()).unwrap();
2225            let foreground = RunningProtocolPeer {
2226                rx,
2227                tx,
2228                driver: ProtocolPeerDriver::Active(ConnectionDriver::new(future::ready(Ok(())))),
2229            };
2230            let (Channel { rx, tx }, mut opposed_remote) = Channel::duplex();
2231            let (work_tx, work_rx) = futures::channel::oneshot::channel::<()>();
2232            let opposed = RunningProtocolPeer {
2233                rx,
2234                tx,
2235                driver: ProtocolPeerDriver::Active(ConnectionDriver::new(async move {
2236                    work_rx.await.map_err(crate::util::internal_error)?;
2237                    Ok(())
2238                })),
2239            };
2240            let mut bridge = Box::pin(if foreground_on_left {
2241                pipe_protocol_peers_until_done(foreground, opposed)
2242            } else {
2243                pipe_protocol_peers_until_done(opposed, foreground)
2244            });
2245
2246            assert_eq!(
2247                bridge.as_mut().now_or_never(),
2248                Some(Ok(())),
2249                "finite foreground must not join an opaque pending peer"
2250            );
2251            assert!(work_tx.is_canceled(), "opaque work must be dropped");
2252            assert!(opposed_remote.rx.next().await.is_some());
2253            assert!(opposed_remote.rx.next().await.is_none());
2254        }
2255    }
2256
2257    #[tokio::test]
2258    async fn foreground_completion_does_not_hide_opposed_ready_driver_error() {
2259        let (Channel { rx, tx }, _foreground_remote) = Channel::duplex();
2260        let foreground = RunningProtocolPeer {
2261            rx,
2262            tx,
2263            driver: ProtocolPeerDriver::Active(ConnectionDriver::new(future::ready(Ok(())))),
2264        };
2265        let (Channel { rx, tx }, _opposed_remote) = Channel::duplex();
2266        let error = crate::Error::internal_error().data("opposed driver failed");
2267        let opposed = RunningProtocolPeer {
2268            rx,
2269            tx,
2270            driver: ProtocolPeerDriver::Active(ConnectionDriver::new(future::ready(Err(
2271                error.clone()
2272            )))),
2273        };
2274
2275        assert_eq!(
2276            pipe_protocol_peers_until_done(foreground, opposed).await,
2277            Err(error)
2278        );
2279    }
2280
2281    #[tokio::test]
2282    async fn passive_protocol_bridge_preserves_the_reverse_half_after_eof() {
2283        let (left, mut remote_left) = Channel::duplex();
2284        let (right, mut remote_right) = Channel::duplex();
2285        let mut bridge = Box::pin(pipe_protocol_peers_until_done(
2286            RunningProtocolPeer::new::<UntypedRole>(left),
2287            RunningProtocolPeer::new::<UntypedRole>(right),
2288        ));
2289        remote_left.tx.close_channel();
2290        assert!(bridge.as_mut().now_or_never().is_none());
2291        assert!(remote_right.rx.next().await.is_none());
2292
2293        remote_right.tx.unbounded_send(frame()).unwrap();
2294        remote_right.tx.close_channel();
2295        bridge.await.unwrap();
2296        assert!(remote_left.rx.next().await.is_some());
2297        assert!(remote_left.rx.next().await.is_none());
2298    }
2299}