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#[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 pub fn builder(self) -> Builder<Client, NullHandler, NullRun> {
78 <Self as Role>::builder(self)
79 }
80
81 #[cfg(feature = "unstable_protocol_v2")]
89 pub fn v2(self) -> V2Builder<Client, NullHandler, NullRun> {
90 self.builder().v2_client()
91 }
92
93 #[cfg(feature = "unstable_protocol_v2")]
106 #[must_use]
107 pub fn protocol_connector(self) -> ClientProtocolConnector {
108 ClientProtocolConnector::new()
109 }
110
111 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#[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 #[must_use]
141 pub fn new() -> Self {
142 Self::default()
143 }
144
145 #[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 #[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 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 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 drop(client);
230 fallback_client.send(fallback_response)?;
231 return pipe_protocol_peers_until_done(fallback_client, agent_connection)
232 .await;
233 }
234
235 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#[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 #[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 pub fn builder(self) -> Builder<Agent, NullHandler, NullRun> {
341 <Self as Role>::builder(self)
342 }
343
344 #[cfg(feature = "unstable_protocol_v2")]
352 pub fn v2(self) -> V2Builder<Agent, NullHandler, NullRun> {
353 self.builder().v2_agent()
354 }
355
356 #[cfg(feature = "unstable_protocol_v2")]
369 #[must_use]
370 pub fn protocol_router(self) -> AgentProtocolRouter {
371 AgentProtocolRouter::new()
372 }
373}
374
375#[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 #[must_use]
391 pub fn new() -> Self {
392 Self::default()
393 }
394
395 #[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 #[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 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 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 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(¶ms, 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(¶ms, false)
782 .expect("v1 routing may ignore parameters unavailable in v1");
783 normalize_v2_initialize_params_for_v1(¶ms, 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(¶ms, 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(¶ms, false)
844 .expect("v1 routing may discard terminal marker metadata");
845 normalize_v2_initialize_params_for_v1(¶ms, 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 return Ok(());
989 };
990 if !driver.request_finish() {
991 return crate::util::run_until(driver, future::ready(Ok(()))).await;
994 }
995
996 let drain_incoming = async move {
997 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 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 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 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#[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 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#[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 pub fn builder(self) -> Builder<Proxy, NullHandler, NullRun> {
1481 Builder::new(self)
1482 }
1483
1484 #[cfg(feature = "unstable_protocol_v2")]
1493 pub fn v2(self) -> V2Builder<Proxy, NullHandler, NullRun> {
1494 self.builder().v2_proxy()
1495 }
1496
1497 #[cfg(feature = "unstable_protocol_v2")]
1508 #[must_use]
1509 pub fn protocol_router(self) -> ProxyProtocolRouter {
1510 ProxyProtocolRouter::new()
1511 }
1512}
1513
1514#[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 #[must_use]
1532 pub fn new() -> Self {
1533 Self::default()
1534 }
1535
1536 #[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 #[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#[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 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 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
1791pub(crate) struct ProxySessionMessages {
1796 session_id: SessionId,
1797}
1798
1799impl ProxySessionMessages {
1800 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 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 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 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 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 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 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 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}