Skip to main content

agent_client_protocol_conductor/
conductor.rs

1//! # Conductor: ACP Proxy Chain Orchestrator
2//!
3//! This module implements the Conductor conductor, which orchestrates a chain of
4//! proxy components that sit between an editor and an agent, transforming the
5//! Agent-Client Protocol (ACP) stream bidirectionally.
6//!
7//! ## Architecture Overview
8//!
9//! The conductor builds and manages a chain of components:
10//!
11//! ```text
12//! Editor <-ACP-> [Component 0] <-ACP-> [Component 1] <-ACP-> ... <-ACP-> Agent
13//! ```
14//!
15//! Each component receives ACP messages, can transform them, and forwards them
16//! to the next component in the chain. The conductor:
17//!
18//! 1. Spawns each component as a subprocess
19//! 2. Establishes bidirectional JSON-RPC connections with each component
20//! 3. Routes messages between editor, components, and agent
21//! 4. Distinguishes proxy vs agent components via distinct request types
22//!
23//! ## Recursive Chain Building
24//!
25//! The chain is built recursively through the `_proxy/successor` envelope:
26//!
27//! 1. Editor connects to Component 0 via the conductor
28//! 2. When Component 0 wants to communicate with its successor, it sends
29//!    a `_proxy/successor` request or notification containing the inner method
30//!    and params
31//! 3. The conductor unwraps the inner message and forwards it to Component 1
32//! 4. Component 1 does the same for Component 2, and so on
33//! 5. The last component receives the unwrapped ACP message directly
34//!
35//! This allows each component to be written as if it's talking to a single successor,
36//! without knowing about the full chain.
37//!
38//! ## Proxy vs Agent Initialization
39//!
40//! Components discover whether they're a proxy or agent via the initialization request they receive:
41//!
42//! - **Proxy components**: Receive `InitializeProxyRequest` (`_proxy/initialize` method)
43//! - **Agent component**: Receives standard `InitializeRequest` (`initialize` method)
44//!
45//! The conductor sends `InitializeProxyRequest` to all proxy components in the chain,
46//! and `InitializeRequest` only to the final agent component. This allows proxies to
47//! know they should forward messages to a successor, while agents know they are the
48//! terminal component
49//!
50//! ## Message Routing
51//!
52//! The conductor runs an event loop processing messages from:
53//!
54//! - **Editor to first component**: Standard ACP messages
55//! - **Component to successor**: Via the `_proxy/successor` envelope
56//! - **Component responses**: Via futures channels back to requesters
57//!
58//! The message flow ensures bidirectional communication while maintaining the
59//! abstraction that each component only knows about its immediate successor.
60//!
61//! ## Lazy Component Initialization
62//!
63//! Components are instantiated lazily when the first `initialize` request is received
64//! from the editor. This enables dynamic proxy chain construction based on client capabilities.
65//!
66//! ### Fixed Chains
67//!
68//! Use [`ProxiesAndAgent`] to assemble a conductor that presents as an agent:
69//!
70//! ```ignore
71//! use agent_client_protocol_conductor::{ConductorImpl, ProxiesAndAgent};
72//!
73//! let conductor = ConductorImpl::new_agent(
74//!     "my-conductor",
75//!     ProxiesAndAgent::new(agent)
76//!         .proxy(proxy1)
77//!         .proxy(proxy2),
78//! );
79//! ```
80//!
81//! A conductor that presents as a proxy takes only its internal proxies; its
82//! final successor is supplied when the conductor is connected:
83//!
84//! ```ignore
85//! use agent_client_protocol_conductor::ConductorImpl;
86//!
87//! let conductor = ConductorImpl::new_proxy("my-proxy-conductor", vec![proxy]);
88//! ```
89//!
90//! ### Dynamic Chain Selection
91//!
92//! Both constructors also accept an instantiator closure. The closure receives
93//! the `InitializeRequest` and returns the possibly modified request together
94//! with type-erased connectors for the selected chain:
95//!
96//! ```ignore
97//! use agent_client_protocol::{Client, Conductor, DynConnectTo};
98//! use agent_client_protocol_conductor::ConductorImpl;
99//!
100//! let conductor = ConductorImpl::new_agent("my-conductor", |init_req| async move {
101//!     let mut proxies: Vec<DynConnectTo<Conductor>> = Vec::new();
102//!     if has_auth_capability(&init_req) {
103//!         proxies.push(DynConnectTo::new(make_auth_proxy()));
104//!     }
105//!
106//!     let agent: DynConnectTo<Client> = DynConnectTo::new(make_agent());
107//!     Ok((init_req, proxies, agent))
108//! });
109//! ```
110
111use std::sync::Arc;
112
113#[cfg(feature = "unstable_protocol_v2")]
114use agent_client_protocol::UntypedMessage;
115#[cfg(feature = "unstable_protocol_v2")]
116use agent_client_protocol::schema::ProtocolVersion;
117#[cfg(feature = "unstable_protocol_v2")]
118use agent_client_protocol::schema::v2;
119use agent_client_protocol::{
120    Agent, BoxFuture, Client, Conductor, ConnectTo, Dispatch, DynConnectTo, Error, JsonRpcMessage,
121    Proxy, Role, RunWithConnectionTo, role::HasPeer, util::MatchDispatch,
122};
123use agent_client_protocol::{
124    Builder, ConnectionTo, JsonRpcNotification, JsonRpcRequest, SentRequest,
125};
126use agent_client_protocol::{
127    HandleDispatchFrom,
128    schema::{InitializeProxyRequest, v1::InitializeRequest},
129    util::MatchDispatchFrom,
130};
131use agent_client_protocol::{Handled, schema::SuccessorMessage};
132use futures::{
133    SinkExt, StreamExt,
134    channel::mpsc::{self},
135};
136use tracing::{debug, info};
137
138#[cfg(feature = "unstable_protocol_v2")]
139#[derive(Debug, Clone, Copy, PartialEq, Eq)]
140enum InitializeProtocol {
141    V1,
142    V2,
143}
144
145#[cfg(feature = "unstable_protocol_v2")]
146impl InitializeProtocol {
147    fn from_request(
148        request: &agent_client_protocol::UntypedMessage,
149    ) -> Result<InitializeProtocolSelection, Error> {
150        let requested = request
151            .params()
152            .get("protocolVersion")
153            .cloned()
154            .ok_or_else(invalid_initialize_protocol_version)
155            .and_then(|version| {
156                serde_json::from_value::<ProtocolVersion>(version)
157                    .map_err(|_| invalid_initialize_protocol_version())
158            })?;
159
160        let protocol = if requested >= ProtocolVersion::V2 {
161            Self::V2
162        } else if requested == ProtocolVersion::V1 {
163            Self::V1
164        } else {
165            return Err(Error::invalid_request()
166                .data(format!("unsupported ACP protocol version {requested}")));
167        };
168
169        Ok(InitializeProtocolSelection {
170            requested,
171            protocol,
172        })
173    }
174
175    fn version(self) -> ProtocolVersion {
176        match self {
177            Self::V1 => ProtocolVersion::V1,
178            Self::V2 => ProtocolVersion::V2,
179        }
180    }
181}
182
183#[cfg(feature = "unstable_protocol_v2")]
184#[derive(Debug, Clone, Copy, PartialEq, Eq)]
185struct InitializeProtocolSelection {
186    requested: ProtocolVersion,
187    protocol: InitializeProtocol,
188}
189
190#[cfg(feature = "unstable_protocol_v2")]
191fn invalid_initialize_protocol_version() -> Error {
192    Error::invalid_params().data("initialize.protocolVersion must be a valid ACP protocol version")
193}
194
195#[cfg(feature = "unstable_protocol_v2")]
196fn forwarded_initialize_request<Request>(
197    raw_request: &UntypedMessage,
198    selection: InitializeProtocolSelection,
199    original_request: &Request,
200    modified_request: Request,
201) -> Result<UntypedMessage, Error>
202where
203    Request: JsonRpcRequest + PartialEq,
204{
205    if modified_request == *original_request && selection.requested == selection.protocol.version()
206    {
207        Ok(UntypedMessage {
208            method: "initialize".to_string(),
209            params: raw_request.params().clone(),
210        })
211    } else {
212        modified_request.to_untyped_message()
213    }
214}
215
216/// The conductor manages the proxy chain lifecycle and message routing.
217///
218/// It maintains connections to all components in the chain and routes messages
219/// bidirectionally between the editor, components, and agent.
220///
221#[derive(Debug)]
222pub struct ConductorImpl<Host: ConductorHostRole> {
223    host: Host,
224    name: String,
225    instantiator: Host::Instantiator,
226    trace_writer: Option<crate::trace::TraceWriter>,
227}
228
229impl<Host: ConductorHostRole> ConductorImpl<Host> {
230    pub fn new(host: Host, name: impl ToString, instantiator: Host::Instantiator) -> Self {
231        ConductorImpl {
232            name: name.to_string(),
233            host,
234            instantiator,
235            trace_writer: None,
236        }
237    }
238}
239
240impl ConductorImpl<Agent> {
241    /// Create a conductor in agent mode (the last component is an agent).
242    pub fn new_agent(
243        name: impl ToString,
244        instantiator: impl InstantiateProxiesAndAgent + 'static,
245    ) -> Self {
246        ConductorImpl::new(Agent, name, Box::new(instantiator))
247    }
248}
249
250impl ConductorImpl<Proxy> {
251    /// Create a conductor in proxy mode (forwards to another conductor).
252    pub fn new_proxy(name: impl ToString, instantiator: impl InstantiateProxies + 'static) -> Self {
253        ConductorImpl::new(Proxy, name, Box::new(instantiator))
254    }
255}
256
257impl<Host: ConductorHostRole> ConductorImpl<Host> {
258    /// Enable trace logging to a custom destination.
259    ///
260    /// Use `agent-client-protocol-trace-viewer` to view the trace as an interactive sequence diagram.
261    #[must_use]
262    pub fn trace_to(mut self, dest: impl crate::trace::WriteEvent) -> Self {
263        self.trace_writer = Some(crate::trace::TraceWriter::new(dest));
264        self
265    }
266
267    /// Enable trace logging to a file path.
268    ///
269    /// Events will be written as newline-delimited JSON (`.jsons` format).
270    /// Use `agent-client-protocol-trace-viewer` to view the trace as an interactive sequence diagram.
271    pub fn trace_to_path(mut self, path: impl AsRef<std::path::Path>) -> std::io::Result<Self> {
272        self.trace_writer = Some(crate::trace::TraceWriter::from_path(path)?);
273        Ok(self)
274    }
275
276    /// Enable trace logging with an existing TraceWriter.
277    #[must_use]
278    pub fn with_trace_writer(mut self, writer: crate::trace::TraceWriter) -> Self {
279        self.trace_writer = Some(writer);
280        self
281    }
282
283    /// Run the conductor with a transport.
284    pub async fn run(
285        self,
286        transport: impl ConnectTo<Host>,
287    ) -> Result<(), agent_client_protocol::Error> {
288        let (conductor_tx, conductor_rx) = mpsc::channel(128 /* chosen arbitrarily */);
289
290        // Set up tracing if enabled - spawn writer task and get handle
291        let trace_handle;
292        let trace_future: BoxFuture<'static, Result<(), agent_client_protocol::Error>>;
293        if let Some((h, f)) = self.trace_writer.map(super::trace::TraceWriter::spawn) {
294            trace_handle = Some(h);
295            trace_future = Box::pin(f);
296        } else {
297            trace_handle = None;
298            trace_future = Box::pin(std::future::ready(Ok(())));
299        }
300
301        let runner = ConductorRunner {
302            conductor_rx,
303            conductor_tx: conductor_tx.clone(),
304            #[cfg(not(feature = "unstable_protocol_v2"))]
305            instantiator: Some(self.instantiator),
306            #[cfg(feature = "unstable_protocol_v2")]
307            initialization: InitializationState::Pending(self.instantiator),
308            proxies: Vec::default(),
309            successor: Arc::new(agent_client_protocol::util::internal_error(
310                "successor not initialized",
311            )),
312            trace_handle,
313            host: self.host.clone(),
314        };
315
316        let connection = Builder::new_with(
317            self.host.clone(),
318            ConductorMessageHandler {
319                conductor_tx,
320                host: self.host.clone(),
321            },
322        );
323        #[cfg(feature = "unstable_protocol_v2")]
324        let connection = connection.without_acp_version_guard();
325
326        connection
327            .name(self.name)
328            .with_runner(runner)
329            .with_spawned(|_cx| trace_future)
330            .connect_to(transport)
331            .await
332    }
333
334    async fn incoming_message_from_client(
335        conductor_tx: &mut mpsc::Sender<ConductorMessage>,
336        message: Dispatch,
337    ) -> Result<(), agent_client_protocol::Error> {
338        conductor_tx
339            .send(ConductorMessage::LeftToRight {
340                target_component_index: 0,
341                message,
342            })
343            .await
344            .map_err(agent_client_protocol::util::internal_error)
345    }
346
347    async fn incoming_message_from_agent(
348        conductor_tx: &mut mpsc::Sender<ConductorMessage>,
349        message: Dispatch,
350    ) -> Result<(), agent_client_protocol::Error> {
351        conductor_tx
352            .send(ConductorMessage::RightToLeft {
353                source_component_index: SourceComponentIndex::Successor,
354                message,
355            })
356            .await
357            .map_err(agent_client_protocol::util::internal_error)
358    }
359}
360
361impl<Host: ConductorHostRole> ConnectTo<Host::Counterpart> for ConductorImpl<Host> {
362    async fn connect_to(
363        self,
364        client: impl ConnectTo<Host>,
365    ) -> Result<(), agent_client_protocol::Error> {
366        self.run(client).await
367    }
368}
369
370struct ConductorMessageHandler<Host: ConductorHostRole> {
371    conductor_tx: mpsc::Sender<ConductorMessage>,
372    host: Host,
373}
374
375impl<Host: ConductorHostRole> HandleDispatchFrom<Host::Counterpart>
376    for ConductorMessageHandler<Host>
377{
378    async fn handle_dispatch_from(
379        &mut self,
380        message: Dispatch,
381        connection: agent_client_protocol::ConnectionTo<Host::Counterpart>,
382    ) -> Result<agent_client_protocol::Handled<Dispatch>, agent_client_protocol::Error> {
383        self.host
384            .handle_dispatch(message, connection, &mut self.conductor_tx)
385            .await
386    }
387
388    fn describe_chain(&self) -> impl std::fmt::Debug {
389        "ConductorMessageHandler"
390    }
391}
392
393/// The conductor manages the proxy chain lifecycle and message routing.
394///
395/// It maintains connections to all components in the chain and routes messages
396/// bidirectionally between the editor, components, and agent.
397///
398pub struct ConductorRunner<Host>
399where
400    Host: ConductorHostRole,
401{
402    conductor_rx: mpsc::Receiver<ConductorMessage>,
403
404    conductor_tx: mpsc::Sender<ConductorMessage>,
405
406    /// The instantiator for lazy initialization.
407    /// Set to None after components are instantiated.
408    #[cfg(not(feature = "unstable_protocol_v2"))]
409    instantiator: Option<Host::Instantiator>,
410
411    /// The explicit initialization lifecycle used by the multi-version router.
412    #[cfg(feature = "unstable_protocol_v2")]
413    initialization: InitializationState<Host::Instantiator>,
414
415    /// The chain of proxies before the agent (if any).
416    ///
417    /// Populated lazily when the first Initialize request is received.
418    proxies: Vec<ConnectionTo<Proxy>>,
419
420    /// If the conductor is operating in agent mode, this will direct messages to the agent.
421    /// If the conductor is operating in proxy mode, this will direct messages to the successor.
422    /// Populated lazily when the first Initialize request is received; the initial value just returns errors.
423    successor: Arc<dyn ConductorSuccessor<Host>>,
424
425    /// Optional trace handle for sequence diagram visualization.
426    trace_handle: Option<crate::trace::TraceHandle>,
427
428    /// Defines what sort of link we have
429    host: Host,
430}
431
432#[cfg(feature = "unstable_protocol_v2")]
433enum InitializationState<Instantiator> {
434    Pending(Instantiator),
435    Initializing,
436    Ready,
437    Failed(Error),
438}
439
440impl<Host> std::fmt::Debug for ConductorRunner<Host>
441where
442    Host: ConductorHostRole,
443{
444    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
445        f.debug_struct("ConductorRunner")
446            .field("conductor_rx", &self.conductor_rx)
447            .field("conductor_tx", &self.conductor_tx)
448            .field("proxies", &self.proxies)
449            .field("trace_handle", &self.trace_handle)
450            .field("host", &self.host)
451            .finish_non_exhaustive()
452    }
453}
454
455impl<Host> RunWithConnectionTo<Host::Counterpart> for ConductorRunner<Host>
456where
457    Host: ConductorHostRole,
458{
459    async fn run_with_connection_to(
460        mut self,
461        connection: ConnectionTo<Host::Counterpart>,
462    ) -> Result<(), agent_client_protocol::Error> {
463        // Components are now spawned lazily in forward_initialize_request
464        // when the first Initialize request is received.
465
466        // This is the "central actor" of the conductor. Most other things forward messages
467        // via `conductor_tx` into this loop. This lets us serialize the conductor's activity.
468        while let Some(message) = self.conductor_rx.next().await {
469            self.handle_conductor_message(connection.clone(), message)
470                .await?;
471        }
472        Ok(())
473    }
474}
475
476impl<Host> ConductorRunner<Host>
477where
478    Host: ConductorHostRole,
479{
480    /// Recursively spawns components and builds the proxy chain.
481    ///
482    /// This function implements the recursive chain building pattern:
483    /// 1. Pop the next component from the `providers` list
484    /// 2. Create the component (either spawn subprocess or use mock)
485    /// 3. Set up JSON-RPC connection and message handlers
486    /// 4. Recursively call itself to spawn the next component
487    /// 5. When no components remain, continue in the central runner's message-routing loop
488    ///
489    /// Central message handling logic for the conductor.
490    /// The conductor routes all [`ConductorMessage`] messages through to this function.
491    /// Each message corresponds to a request or notification from one component to another.
492    /// The conductor ferries messages from one place to another, sometimes making modifications along the way.
493    /// Note that *responses to requests* are sent *directly* without going through this loop.
494    ///
495    /// The names we use are
496    ///
497    /// * The *client* is the originator of all ACP traffic, typically an editor or GUI.
498    /// * Then there is a sequence of *components* consisting of:
499    ///     * Zero or more *proxies*, which receive messages and forward them to the next component in the chain.
500    ///     * And finally the *agent*, which is the final component in the chain and handles the actual work.
501    ///
502    /// For the most part, we pass messages through the chain without modification. The initialization
503    /// handshake is the exception:
504    ///
505    /// * We send `InitializeProxyRequest` to proxy components and `InitializeRequest` to the agent component.
506    async fn handle_conductor_message(
507        &mut self,
508        client: ConnectionTo<Host::Counterpart>,
509        message: ConductorMessage,
510    ) -> Result<(), agent_client_protocol::Error> {
511        tracing::debug!(?message, "handle_conductor_message");
512
513        match message {
514            ConductorMessage::LeftToRight {
515                target_component_index,
516                message,
517            } => {
518                // Tracing happens inside forward_client_to_agent_message, after initialization,
519                // so that component_name() has access to the populated proxies list.
520                self.forward_client_to_agent_message(target_component_index, message, client)
521                    .await
522            }
523
524            ConductorMessage::RightToLeft {
525                source_component_index,
526                message,
527            } => {
528                tracing::debug!(
529                    ?source_component_index,
530                    message_method = ?message.method(),
531                    "Conductor: AgentToClient received"
532                );
533                self.send_message_to_predecessor_of(client, source_component_index, message)
534            }
535        }
536    }
537
538    /// Send a message (request or notification) to the predecessor of the given component.
539    ///
540    /// This is a bit subtle because the relationship of the conductor
541    /// is different depending on who will be receiving the message:
542    /// * If the message is going to the conductor's client, then no changes
543    ///   are needed, as the conductor is sending an agent-to-client message and
544    ///   the conductor is acting as the agent.
545    /// * If the message is going to a proxy component, then we have to wrap
546    ///   it in a "from successor" wrapper, because the conductor is the
547    ///   proxy's client.
548    fn send_message_to_predecessor_of<Req: JsonRpcRequest, N: JsonRpcNotification>(
549        &mut self,
550        client: ConnectionTo<Host::Counterpart>,
551        source_component_index: SourceComponentIndex,
552        message: Dispatch<Req, N>,
553    ) -> Result<(), agent_client_protocol::Error>
554    where
555        Req::Response: Send,
556    {
557        let source_component_index = match source_component_index {
558            SourceComponentIndex::Successor => self.proxies.len(),
559            SourceComponentIndex::Proxy(index) => index,
560        };
561
562        match message {
563            Dispatch::Request(request, responder) => self
564                .send_request_to_predecessor_of(client, source_component_index, request)
565                .forward_response_to(responder),
566            Dispatch::Notification(notification) => {
567                // `$/cancel_request` is connection-scoped: its `requestId` was
568                // allocated on the connection the notification arrived over
569                // and means nothing on the predecessor's connection. The SDK
570                // already propagates the cancellation hop by hop through the
571                // `forward_response_to` calls above, so drop the raw
572                // notification instead of tunneling a meaningless ID.
573                if agent_client_protocol::is_cancel_request_notification(&notification) {
574                    tracing::debug!(
575                        "not forwarding hop-scoped `$/cancel_request` notification to predecessor"
576                    );
577                    return Ok(());
578                }
579                self.send_notification_to_predecessor_of(
580                    client,
581                    source_component_index,
582                    notification,
583                )
584            }
585            Dispatch::Response(result, router) => router.route_with_result(result),
586        }
587    }
588
589    fn send_request_to_predecessor_of<Req: JsonRpcRequest>(
590        &mut self,
591        client_connection: ConnectionTo<Host::Counterpart>,
592        source_component_index: usize,
593        request: Req,
594    ) -> SentRequest<Req::Response> {
595        if source_component_index == 0 {
596            client_connection.send_request_to(Client, request)
597        } else {
598            self.proxies[source_component_index - 1].send_request(SuccessorMessage {
599                message: request,
600                meta: None,
601            })
602        }
603    }
604
605    /// Send a notification to the predecessor of the given component.
606    ///
607    /// This is a bit subtle because the relationship of the conductor
608    /// is different depending on who will be receiving the message:
609    /// * If the notification is going to the conductor's client, then no changes
610    ///   are needed, as the conductor is sending an agent-to-client message and
611    ///   the conductor is acting as the agent.
612    /// * If the notification is going to a proxy component, then we have to wrap
613    ///   it in a "from successor" wrapper, because the conductor is the
614    ///   proxy's client.
615    fn send_notification_to_predecessor_of<N: JsonRpcNotification>(
616        &mut self,
617        client: ConnectionTo<Host::Counterpart>,
618        source_component_index: usize,
619        notification: N,
620    ) -> Result<(), agent_client_protocol::Error> {
621        tracing::debug!(
622            source_component_index,
623            proxies_len = self.proxies.len(),
624            "send_notification_to_predecessor_of"
625        );
626        if source_component_index == 0 {
627            tracing::debug!("Sending notification directly to client");
628            client.send_notification_to(Client, notification)
629        } else {
630            tracing::debug!(
631                target_proxy = source_component_index - 1,
632                "Sending notification wrapped as SuccessorMessage to proxy"
633            );
634            self.proxies[source_component_index - 1].send_notification(SuccessorMessage {
635                message: notification,
636                meta: None,
637            })
638        }
639    }
640
641    /// Send a message (request or notification) from 'left to right'.
642    /// Left-to-right means from the client or an intermediate proxy to the component
643    /// at `target_component_index` (could be a proxy or the agent).
644    /// Ensures the component chain is initialized before forwarding the message.
645    async fn forward_client_to_agent_message(
646        &mut self,
647        target_component_index: usize,
648        message: Dispatch,
649        client: ConnectionTo<Host::Counterpart>,
650    ) -> Result<(), agent_client_protocol::Error> {
651        tracing::trace!(
652            target_component_index,
653            ?message,
654            "forward_client_to_agent_message"
655        );
656
657        // Ensure components are initialized before processing any message.
658        let Some(message) = self.ensure_initialized(client.clone(), message).await? else {
659            return Ok(());
660        };
661
662        // In proxy mode, if the target is beyond our component chain,
663        // forward to the conductor's own successor (via client connection)
664        if target_component_index < self.proxies.len() {
665            self.forward_message_from_client_to_proxy(target_component_index, message)
666                .await
667        } else {
668            assert_eq!(target_component_index, self.proxies.len());
669
670            debug!(
671                target_component_index,
672                proxies_count = self.proxies.len(),
673                "Proxy mode: forwarding successor message to conductor's successor"
674            );
675            let successor = self.successor.clone();
676            successor.send_message(message, client, self).await
677        }
678    }
679
680    /// Ensures components are initialized before processing messages.
681    ///
682    /// If components haven't been initialized yet, this expects the first message
683    /// to be an `initialize` request and uses it to spawn the component chain.
684    ///
685    /// Returns:
686    /// - `Ok(Some(message))` - Components are initialized, continue processing this message
687    /// - `Ok(None)` - An error response was sent, caller should return early
688    /// - `Err(_)` - A fatal error occurred
689    async fn ensure_initialized(
690        &mut self,
691        client: ConnectionTo<Host::Counterpart>,
692        message: Dispatch,
693    ) -> Result<Option<Dispatch>, Error> {
694        #[cfg(not(feature = "unstable_protocol_v2"))]
695        {
696            let Some(instantiator) = self.instantiator.take() else {
697                return Ok(Some(message));
698            };
699
700            let host = self.host.clone();
701            let message = host.initialize(message, client, instantiator, self).await?;
702            Ok(Some(message))
703        }
704
705        #[cfg(feature = "unstable_protocol_v2")]
706        {
707            let state =
708                std::mem::replace(&mut self.initialization, InitializationState::Initializing);
709            match state {
710                InitializationState::Pending(instantiator) => {
711                    let host = self.host.clone();
712                    match host
713                        .initialize_with_outcome(message, client, instantiator, self)
714                        .await?
715                    {
716                        InitializationOutcome::Forward(message) => {
717                            self.initialization = InitializationState::Ready;
718                            Ok(Some(message))
719                        }
720                        InitializationOutcome::Rejected(error) => {
721                            self.initialization = InitializationState::Failed(error);
722                            Ok(None)
723                        }
724                    }
725                }
726                InitializationState::Ready => {
727                    self.initialization = InitializationState::Ready;
728                    Ok(Some(message))
729                }
730                InitializationState::Failed(error) => {
731                    let result = match message {
732                        Dispatch::Request(_, responder) => {
733                            responder.respond_with_error(error.clone())
734                        }
735                        Dispatch::Notification(_) => Ok(()),
736                        Dispatch::Response(_, router) => router.route_with_error(error.clone()),
737                    };
738                    self.initialization = InitializationState::Failed(error);
739                    result?;
740                    Ok(None)
741                }
742                InitializationState::Initializing => {
743                    Err(Error::internal_error().data("conductor initialization was re-entered"))
744                }
745            }
746        }
747    }
748
749    /// Wrap a proxy component with tracing if tracing is enabled.
750    ///
751    /// Returns the component unchanged if tracing is disabled.
752    fn trace_proxy(
753        &self,
754        proxy_index: ComponentIndex,
755        successor_index: ComponentIndex,
756        component: impl ConnectTo<Conductor>,
757    ) -> DynConnectTo<Conductor> {
758        match &self.trace_handle {
759            Some(trace_handle) => {
760                trace_handle.bridge_component(proxy_index, successor_index, component)
761            }
762            None => DynConnectTo::new(component),
763        }
764    }
765
766    /// Spawn proxy components and add them to the proxies list.
767    fn spawn_proxies(
768        &mut self,
769        client: ConnectionTo<Host::Counterpart>,
770        proxy_components: Vec<DynConnectTo<Conductor>>,
771    ) -> Result<(), agent_client_protocol::Error> {
772        assert!(self.proxies.is_empty());
773
774        let num_proxies = proxy_components.len();
775        info!(proxy_count = num_proxies, "spawn_proxies");
776
777        // Special case: if there are no user-defined proxies
778        // but tracing is enabled, we make a dummy proxy that just
779        // passes through messages but which can trigger the
780        // tracing events.
781        if self.trace_handle.is_some() && num_proxies == 0 {
782            let trace_proxy = Proxy.builder();
783            #[cfg(feature = "unstable_protocol_v2")]
784            let trace_proxy = trace_proxy.without_acp_version_guard();
785
786            self.connect_to_proxy(
787                &client,
788                0,
789                ComponentIndex::Client,
790                ComponentIndex::Agent,
791                trace_proxy,
792            )?;
793        } else {
794            // Spawn each proxy component
795            for (component_index, dyn_component) in proxy_components.into_iter().enumerate() {
796                debug!(component_index, "spawning proxy");
797
798                self.connect_to_proxy(
799                    &client,
800                    component_index,
801                    ComponentIndex::Proxy(component_index),
802                    ComponentIndex::successor_of(component_index, num_proxies),
803                    dyn_component,
804                )?;
805            }
806        }
807
808        info!(proxy_count = self.proxies.len(), "Proxies spawned");
809
810        Ok(())
811    }
812
813    /// Create a connection to the proxy with index `component_index` implemented in `component`.
814    ///
815    /// If tracing is enabled, the proxy's index is `trace_proxy_index` and its successor is `trace_successor_index`.
816    fn connect_to_proxy(
817        &mut self,
818        client: &ConnectionTo<Host::Counterpart>,
819        component_index: usize,
820        trace_proxy_index: ComponentIndex,
821        trace_successor_index: ComponentIndex,
822        component: impl ConnectTo<Conductor>,
823    ) -> Result<(), Error> {
824        let connection_builder = self.connection_to_proxy(component_index);
825        let connect_component =
826            self.trace_proxy(trace_proxy_index, trace_successor_index, component);
827        let proxy_connection = client.spawn_connection(connection_builder, connect_component)?;
828        self.proxies.push(proxy_connection);
829        Ok(())
830    }
831
832    /// Create the conductor's connection to the proxy with index `component_index`.
833    ///
834    /// Outgoing messages received from the proxy are sent to `self.conductor_tx` as either
835    /// left-to-right or right-to-left messages depending on whether they are wrapped
836    /// in `SuccessorMessage`.
837    fn connection_to_proxy(
838        &mut self,
839        component_index: usize,
840    ) -> Builder<Conductor, impl HandleDispatchFrom<Proxy> + 'static> {
841        type SuccessorDispatch = Dispatch<SuccessorMessage, SuccessorMessage>;
842        let mut conductor_tx = self.conductor_tx.clone();
843        Conductor
844            .builder()
845            .name(format!("conductor-to-component({component_index})"))
846            // Intercept messages sent by the proxy.
847            .on_receive_dispatch(
848                async move |dispatch: Dispatch, _connection| {
849                    MatchDispatch::new(dispatch)
850                        .if_dispatch(async |dispatch: SuccessorDispatch| {
851                            //                         ------------------
852                            // SuccessorMessages sent by the proxy go to its successor.
853                            //
854                            // Subtle point:
855                            //
856                            // `ConductorToProxy` has only a single peer, `Agent`. This means that we see
857                            // "successor messages" in their "desugared form". So when we intercept an *outgoing*
858                            // message that matches `SuccessorMessage`, it could be one of three things
859                            //
860                            // - A request being sent by the proxy to its successor (hence going left->right)
861                            // - A notification being sent by the proxy to its successor (hence going left->right)
862                            // - A response to a request sent to the proxy *by* its successor. Here, the *request*
863                            //   was going right->left, but the *response* (the message we are processing now)
864                            //   is going left->right.
865                            //
866                            // So, in all cases, we forward as a left->right message.
867
868                            conductor_tx
869                                .send(ConductorMessage::LeftToRight {
870                                    target_component_index: component_index + 1,
871                                    message: dispatch.map(|r, cx| (r.message, cx), |n| n.message),
872                                })
873                                .await
874                                .map_err(agent_client_protocol::util::internal_error)
875                        })
876                        .await
877                        .otherwise(async |dispatch| {
878                            // Other messagrs send by the proxy go its predecessor.
879                            // As in the previous handler:
880                            //
881                            // Messages here are seen in their "desugared form", so we are seeing
882                            // one of three things
883                            //
884                            // - A request being sent by the proxy to its predecessor (hence going right->left)
885                            // - A notification being sent by the proxy to its predecessor (hence going right->left)
886                            // - A response to a request sent to the proxy *by* its predecessor. Here, the *request*
887                            //   was going left->right, but the *response* (the message we are processing now)
888                            //   is going right->left.
889                            //
890                            // So, in all cases, we forward as a right->left message.
891
892                            let message = ConductorMessage::RightToLeft {
893                                source_component_index: SourceComponentIndex::Proxy(
894                                    component_index,
895                                ),
896                                message: dispatch,
897                            };
898                            conductor_tx
899                                .send(message)
900                                .await
901                                .map_err(agent_client_protocol::util::internal_error)
902                        })
903                        .await
904                },
905                agent_client_protocol::on_receive_dispatch!(),
906            )
907    }
908
909    // The feature-off implementation awaits typed dispatch matchers; the v2
910    // implementation is intentionally raw and completes synchronously.
911    #[allow(unknown_lints, clippy::unused_async, clippy::unused_async_trait_impl)]
912    async fn forward_message_from_client_to_proxy(
913        &mut self,
914        target_component_index: usize,
915        message: Dispatch,
916    ) -> Result<(), agent_client_protocol::Error> {
917        tracing::debug!(?message, "forward_message_to_proxy");
918
919        #[cfg(not(feature = "unstable_protocol_v2"))]
920        {
921            MatchDispatch::new(message)
922                .if_request(async |_request: InitializeProxyRequest, responder| {
923                    responder.respond_with_error(
924                        agent_client_protocol::Error::invalid_request()
925                            .data("initialize/proxy requests are only sent by the conductor"),
926                    )
927                })
928                .await
929                .if_request(async |request: InitializeRequest, responder| {
930                    // The pattern for `Initialize` messages is a bit subtle.
931                    // Proxies receive incoming `Initialize` messages as if they
932                    // were a client. The conductor (us) intercepts these and
933                    // converts them to an `InitializeProxyRequest`.
934                    //
935                    // The proxy will then initialize itself and forward an `Initialize`
936                    // request to its successor.
937                    let sent = self.proxies[target_component_index]
938                        .send_request(InitializeProxyRequest::from(request));
939                    // The request is rewritten, so `forward_response_to` cannot be
940                    // used here; wire up cancellation forwarding explicitly to
941                    // keep `initialize` cancellable like every other forwarded
942                    // request.
943                    let sent = sent.forward_cancellation_from(responder.cancellation());
944                    sent.on_receiving_result(async move |result| {
945                        tracing::debug!(?result, "got initialize_proxy response from proxy");
946                        responder.respond_with_result(result)
947                    })
948                })
949                .await
950                .otherwise(async |message| {
951                    self.proxies[target_component_index].send_proxied_message(message)
952                })
953                .await
954        }
955
956        #[cfg(feature = "unstable_protocol_v2")]
957        {
958            match message {
959                Dispatch::Request(request, responder)
960                    if request.method()
961                        == agent_client_protocol::schema::METHOD_INITIALIZE_PROXY =>
962                {
963                    responder.respond_with_error(
964                        agent_client_protocol::Error::invalid_request()
965                            .data("initialize/proxy requests are only sent by the conductor"),
966                    )
967                }
968                Dispatch::Request(mut request, responder)
969                    if InitializeRequest::matches_method(request.method()) =>
970                {
971                    request.method =
972                        agent_client_protocol::schema::METHOD_INITIALIZE_PROXY.to_string();
973                    let sent = self.proxies[target_component_index].send_request(request);
974                    let sent = sent.forward_cancellation_from(responder.cancellation());
975                    sent.on_receiving_result(async move |result| {
976                        tracing::debug!(?result, "got initialize_proxy response from proxy");
977                        responder.respond_with_result(result)
978                    })
979                }
980                message => self.proxies[target_component_index].send_proxied_message(message),
981            }
982        }
983    }
984
985    /// Invoked when sending a message from the conductor to the agent that it manages.
986    /// This is called by `self.successor`'s [`ConductorSuccessor::send_message`]
987    /// method when `Link = ConductorToClient` (i.e., the conductor is not itself
988    /// running as a proxy).
989    // The feature-off implementation awaits typed dispatch matchers; the v2
990    // implementation is intentionally raw and completes synchronously.
991    #[allow(unknown_lints, clippy::unused_async, clippy::unused_async_trait_impl)]
992    async fn forward_message_to_agent(
993        &mut self,
994        _client_connection: ConnectionTo<Host::Counterpart>,
995        message: Dispatch,
996        agent_connection: ConnectionTo<Agent>,
997    ) -> Result<(), Error> {
998        #[cfg(not(feature = "unstable_protocol_v2"))]
999        {
1000            MatchDispatch::new(message)
1001                .if_request(async |_request: InitializeProxyRequest, responder| {
1002                    responder.respond_with_error(
1003                        agent_client_protocol::Error::invalid_request()
1004                            .data("initialize/proxy requests are only sent by the conductor"),
1005                    )
1006                })
1007                .await
1008                .otherwise(async |message| agent_connection.send_proxied_message_to(Agent, message))
1009                .await
1010        }
1011
1012        #[cfg(feature = "unstable_protocol_v2")]
1013        {
1014            match message {
1015                Dispatch::Request(request, responder)
1016                    if request.method()
1017                        == agent_client_protocol::schema::METHOD_INITIALIZE_PROXY =>
1018                {
1019                    responder.respond_with_error(
1020                        agent_client_protocol::Error::invalid_request()
1021                            .data("initialize/proxy requests are only sent by the conductor"),
1022                    )
1023                }
1024                message => agent_connection.send_proxied_message_to(Agent, message),
1025            }
1026        }
1027    }
1028}
1029
1030/// Identifies a component in the conductor's chain for tracing purposes.
1031///
1032/// Used to track message sources and destinations through the proxy chain.
1033#[derive(Debug, Clone, Copy)]
1034pub enum ComponentIndex {
1035    /// The client (editor) at the start of the chain.
1036    Client,
1037
1038    /// A proxy component at the given index.
1039    Proxy(usize),
1040
1041    /// The successor (agent in agent mode, outer conductor in proxy mode).
1042    Agent,
1043}
1044
1045impl ComponentIndex {
1046    /// Return the index for the predecessor of `proxy_index`, which might be `Client`.
1047    #[must_use]
1048    pub fn predecessor_of(proxy_index: usize) -> Self {
1049        match proxy_index.checked_sub(1) {
1050            Some(p_i) => ComponentIndex::Proxy(p_i),
1051            None => ComponentIndex::Client,
1052        }
1053    }
1054
1055    /// Return the index for the predecessor of `proxy_index`, which might be `Client`.
1056    #[must_use]
1057    pub fn successor_of(proxy_index: usize, num_proxies: usize) -> Self {
1058        if proxy_index == num_proxies {
1059            ComponentIndex::Agent
1060        } else {
1061            ComponentIndex::Proxy(proxy_index + 1)
1062        }
1063    }
1064}
1065
1066/// Identifies the source of an agent-to-client message.
1067///
1068/// This enum handles the fact that the conductor may receive messages from two different sources:
1069/// 1. From one of its managed components (identified by index)
1070/// 2. From the conductor's own successor in a larger proxy chain (when in proxy mode)
1071#[derive(Debug, Clone, Copy)]
1072pub enum SourceComponentIndex {
1073    /// Message from a specific component at the given index in the managed chain.
1074    Proxy(usize),
1075
1076    /// Message from the conductor's agent or successor.
1077    Successor,
1078}
1079
1080/// Trait for lazy proxy instantiation (proxy mode).
1081///
1082/// Used by conductors in proxy mode (`ConductorToConductor`) where all components
1083/// are proxies that forward to an outer conductor.
1084pub trait InstantiateProxies: Send {
1085    /// Instantiate proxy components based on the Initialize request.
1086    ///
1087    /// Returns proxy components typed as `DynConnectTo<Conductor>` since proxies
1088    /// communicate with the conductor.
1089    fn instantiate_proxies(
1090        self: Box<Self>,
1091        req: InitializeRequest,
1092    ) -> futures::future::BoxFuture<
1093        'static,
1094        Result<(InitializeRequest, Vec<DynConnectTo<Conductor>>), agent_client_protocol::Error>,
1095    >;
1096
1097    /// Instantiate proxy components for a protocol-v2 connection.
1098    ///
1099    /// Implementors that only support v1 can rely on the default rejection.
1100    /// Static component collections supplied by this crate support both
1101    /// versions and pass the request through unchanged.
1102    #[cfg(feature = "unstable_protocol_v2")]
1103    fn instantiate_v2_proxies(
1104        self: Box<Self>,
1105        req: v2::InitializeRequest,
1106    ) -> futures::future::BoxFuture<
1107        'static,
1108        Result<(v2::InitializeRequest, Vec<DynConnectTo<Conductor>>), agent_client_protocol::Error>,
1109    > {
1110        drop((self, req));
1111        Box::pin(async {
1112            Err(Error::invalid_request()
1113                .data("this conductor proxy instantiator does not support ACP protocol v2"))
1114        })
1115    }
1116}
1117
1118/// Simple implementation: provide all proxy components unconditionally.
1119///
1120/// Requires `T: ConnectTo<Conductor>`.
1121impl<T> InstantiateProxies for Vec<T>
1122where
1123    T: ConnectTo<Conductor> + 'static,
1124{
1125    fn instantiate_proxies(
1126        self: Box<Self>,
1127        req: InitializeRequest,
1128    ) -> futures::future::BoxFuture<
1129        'static,
1130        Result<(InitializeRequest, Vec<DynConnectTo<Conductor>>), agent_client_protocol::Error>,
1131    > {
1132        Box::pin(async move {
1133            let components: Vec<DynConnectTo<Conductor>> =
1134                (*self).into_iter().map(|c| DynConnectTo::new(c)).collect();
1135            Ok((req, components))
1136        })
1137    }
1138
1139    #[cfg(feature = "unstable_protocol_v2")]
1140    fn instantiate_v2_proxies(
1141        self: Box<Self>,
1142        req: v2::InitializeRequest,
1143    ) -> futures::future::BoxFuture<
1144        'static,
1145        Result<(v2::InitializeRequest, Vec<DynConnectTo<Conductor>>), agent_client_protocol::Error>,
1146    > {
1147        Box::pin(async move {
1148            let components = (*self).into_iter().map(DynConnectTo::new).collect();
1149            Ok((req, components))
1150        })
1151    }
1152}
1153
1154/// Dynamic implementation: closure receives the Initialize request and returns proxies.
1155impl<F, Fut> InstantiateProxies for F
1156where
1157    F: FnOnce(InitializeRequest) -> Fut + Send + 'static,
1158    Fut: std::future::Future<
1159            Output = Result<
1160                (InitializeRequest, Vec<DynConnectTo<Conductor>>),
1161                agent_client_protocol::Error,
1162            >,
1163        > + Send
1164        + 'static,
1165{
1166    fn instantiate_proxies(
1167        self: Box<Self>,
1168        req: InitializeRequest,
1169    ) -> futures::future::BoxFuture<
1170        'static,
1171        Result<(InitializeRequest, Vec<DynConnectTo<Conductor>>), agent_client_protocol::Error>,
1172    > {
1173        Box::pin(async move { (*self)(req).await })
1174    }
1175}
1176
1177/// Trait for lazy proxy and agent instantiation (agent mode).
1178///
1179/// Used by conductors in agent mode (`ConductorToClient`) where there are
1180/// zero or more proxies followed by an agent component.
1181pub trait InstantiateProxiesAndAgent: Send {
1182    /// Instantiate proxy and agent components based on the Initialize request.
1183    ///
1184    /// Returns the (possibly modified) request, a vector of proxy components
1185    /// (typed as `DynConnectTo<Conductor>`), and the agent component
1186    /// (typed as `DynConnectTo<Client>`).
1187    fn instantiate_proxies_and_agent(
1188        self: Box<Self>,
1189        req: InitializeRequest,
1190    ) -> futures::future::BoxFuture<
1191        'static,
1192        Result<
1193            (
1194                InitializeRequest,
1195                Vec<DynConnectTo<Conductor>>,
1196                DynConnectTo<Client>,
1197            ),
1198            agent_client_protocol::Error,
1199        >,
1200    >;
1201
1202    /// Instantiate proxy and agent components for a protocol-v2 connection.
1203    ///
1204    /// Implementors that only support v1 can rely on the default rejection.
1205    /// [`AgentOnly`] and [`ProxiesAndAgent`] support both versions.
1206    #[cfg(feature = "unstable_protocol_v2")]
1207    fn instantiate_v2_proxies_and_agent(
1208        self: Box<Self>,
1209        req: v2::InitializeRequest,
1210    ) -> futures::future::BoxFuture<
1211        'static,
1212        Result<
1213            (
1214                v2::InitializeRequest,
1215                Vec<DynConnectTo<Conductor>>,
1216                DynConnectTo<Client>,
1217            ),
1218            agent_client_protocol::Error,
1219        >,
1220    > {
1221        drop((self, req));
1222        Box::pin(async {
1223            Err(Error::invalid_request()
1224                .data("this conductor agent instantiator does not support ACP protocol v2"))
1225        })
1226    }
1227}
1228
1229/// Wrapper to convert a single agent component (no proxies) into InstantiateProxiesAndAgent.
1230#[derive(Debug)]
1231pub struct AgentOnly<A>(pub A);
1232
1233impl<A: ConnectTo<Client> + 'static> InstantiateProxiesAndAgent for AgentOnly<A> {
1234    fn instantiate_proxies_and_agent(
1235        self: Box<Self>,
1236        req: InitializeRequest,
1237    ) -> futures::future::BoxFuture<
1238        'static,
1239        Result<
1240            (
1241                InitializeRequest,
1242                Vec<DynConnectTo<Conductor>>,
1243                DynConnectTo<Client>,
1244            ),
1245            agent_client_protocol::Error,
1246        >,
1247    > {
1248        Box::pin(async move { Ok((req, Vec::new(), DynConnectTo::new(self.0))) })
1249    }
1250
1251    #[cfg(feature = "unstable_protocol_v2")]
1252    fn instantiate_v2_proxies_and_agent(
1253        self: Box<Self>,
1254        req: v2::InitializeRequest,
1255    ) -> futures::future::BoxFuture<
1256        'static,
1257        Result<
1258            (
1259                v2::InitializeRequest,
1260                Vec<DynConnectTo<Conductor>>,
1261                DynConnectTo<Client>,
1262            ),
1263            agent_client_protocol::Error,
1264        >,
1265    > {
1266        Box::pin(async move { Ok((req, Vec::new(), DynConnectTo::new(self.0))) })
1267    }
1268}
1269
1270/// Builder for creating proxies and agent components.
1271///
1272/// # Example
1273/// ```ignore
1274/// ProxiesAndAgent::new(ElizaAgent::new())
1275///     .proxy(LoggingProxy::new())
1276///     .proxy(AuthProxy::new())
1277/// ```
1278#[derive(Debug)]
1279pub struct ProxiesAndAgent {
1280    proxies: Vec<DynConnectTo<Conductor>>,
1281    agent: DynConnectTo<Client>,
1282}
1283
1284impl ProxiesAndAgent {
1285    /// Create a new builder with the given agent component.
1286    pub fn new(agent: impl ConnectTo<Client> + 'static) -> Self {
1287        Self {
1288            proxies: vec![],
1289            agent: DynConnectTo::new(agent),
1290        }
1291    }
1292
1293    /// Add a single proxy component.
1294    #[must_use]
1295    pub fn proxy(mut self, proxy: impl ConnectTo<Conductor> + 'static) -> Self {
1296        self.proxies.push(DynConnectTo::new(proxy));
1297        self
1298    }
1299
1300    /// Add multiple proxy components.
1301    #[must_use]
1302    pub fn proxies<P, I>(mut self, proxies: I) -> Self
1303    where
1304        P: ConnectTo<Conductor> + 'static,
1305        I: IntoIterator<Item = P>,
1306    {
1307        self.proxies
1308            .extend(proxies.into_iter().map(DynConnectTo::new));
1309        self
1310    }
1311}
1312
1313impl InstantiateProxiesAndAgent for ProxiesAndAgent {
1314    fn instantiate_proxies_and_agent(
1315        self: Box<Self>,
1316        req: InitializeRequest,
1317    ) -> futures::future::BoxFuture<
1318        'static,
1319        Result<
1320            (
1321                InitializeRequest,
1322                Vec<DynConnectTo<Conductor>>,
1323                DynConnectTo<Client>,
1324            ),
1325            agent_client_protocol::Error,
1326        >,
1327    > {
1328        Box::pin(async move { Ok((req, self.proxies, self.agent)) })
1329    }
1330
1331    #[cfg(feature = "unstable_protocol_v2")]
1332    fn instantiate_v2_proxies_and_agent(
1333        self: Box<Self>,
1334        req: v2::InitializeRequest,
1335    ) -> futures::future::BoxFuture<
1336        'static,
1337        Result<
1338            (
1339                v2::InitializeRequest,
1340                Vec<DynConnectTo<Conductor>>,
1341                DynConnectTo<Client>,
1342            ),
1343            agent_client_protocol::Error,
1344        >,
1345    > {
1346        Box::pin(async move { Ok((req, self.proxies, self.agent)) })
1347    }
1348}
1349
1350/// Dynamic implementation: closure receives the Initialize request and returns proxies + agent.
1351impl<F, Fut> InstantiateProxiesAndAgent for F
1352where
1353    F: FnOnce(InitializeRequest) -> Fut + Send + 'static,
1354    Fut: std::future::Future<
1355            Output = Result<
1356                (
1357                    InitializeRequest,
1358                    Vec<DynConnectTo<Conductor>>,
1359                    DynConnectTo<Client>,
1360                ),
1361                agent_client_protocol::Error,
1362            >,
1363        > + Send
1364        + 'static,
1365{
1366    fn instantiate_proxies_and_agent(
1367        self: Box<Self>,
1368        req: InitializeRequest,
1369    ) -> futures::future::BoxFuture<
1370        'static,
1371        Result<
1372            (
1373                InitializeRequest,
1374                Vec<DynConnectTo<Conductor>>,
1375                DynConnectTo<Client>,
1376            ),
1377            agent_client_protocol::Error,
1378        >,
1379    > {
1380        Box::pin(async move { (*self)(req).await })
1381    }
1382}
1383
1384/// Messages sent to the conductor's main event loop for routing.
1385///
1386/// These messages enable the conductor to route communication between:
1387/// - The editor and the first component
1388/// - Components and their successors in the chain
1389/// - Components and their clients (editor or predecessor)
1390///
1391/// All spawned tasks send messages via this enum through a shared channel,
1392/// allowing centralized routing logic in the conductor runner's event loop.
1393#[derive(Debug)]
1394pub enum ConductorMessage {
1395    /// If this message is a request or notification, then it is going "left-to-right"
1396    /// (e.g., a component making a request of its successor).
1397    ///
1398    /// If this message is a response, then it is going right-to-left
1399    /// (i.e., the successor answering a request made by its predecessor).
1400    LeftToRight {
1401        target_component_index: usize,
1402        message: Dispatch,
1403    },
1404
1405    /// If this message is a request or notification, then it is going "right-to-left"
1406    /// (e.g., a component making a request of its predecessor).
1407    ///
1408    /// If this message is a response, then it is going "left-to-right"
1409    /// (i.e., the predecessor answering a request made by its successor).
1410    RightToLeft {
1411        source_component_index: SourceComponentIndex,
1412        message: Dispatch,
1413    },
1414}
1415
1416/// Trait implemented for the two links the conductor can use:
1417///
1418/// * ConductorToClient -- conductor is acting as an agent, so when its last proxy sends to its successor, the conductor sends that message to its agent component
1419/// * ConductorToConductor -- conductor is acting as a proxy, so when its last proxy sends to its successor, the (inner) conductor sends that message to its successor, via the outer conductor
1420pub trait ConductorHostRole: Role<Counterpart: HasPeer<Client>> {
1421    /// The type used to instantiate components for this link type.
1422    type Instantiator: Send;
1423
1424    /// Handle initialization: parse the init request, instantiate components, and spawn them.
1425    ///
1426    /// Takes ownership of the instantiator and returns the (possibly modified) init request
1427    /// wrapped in a Dispatch for forwarding.
1428    fn initialize(
1429        &self,
1430        message: Dispatch,
1431        connection: ConnectionTo<Self::Counterpart>,
1432        instantiator: Self::Instantiator,
1433        runner: &mut ConductorRunner<Self>,
1434    ) -> impl Future<Output = Result<Dispatch, agent_client_protocol::Error>> + Send;
1435
1436    /// Handle initialization while distinguishing a protocol rejection from a
1437    /// fatal connection failure.
1438    ///
1439    /// The conductor uses this hook only when `unstable_protocol_v2` is
1440    /// enabled. The default preserves existing implementations by delegating
1441    /// to [`Self::initialize`].
1442    #[cfg(feature = "unstable_protocol_v2")]
1443    fn initialize_with_outcome(
1444        &self,
1445        message: Dispatch,
1446        connection: ConnectionTo<Self::Counterpart>,
1447        instantiator: Self::Instantiator,
1448        runner: &mut ConductorRunner<Self>,
1449    ) -> impl Future<Output = Result<InitializationOutcome, agent_client_protocol::Error>> + Send
1450    {
1451        async move {
1452            self.initialize(message, connection, instantiator, runner)
1453                .await
1454                .map(InitializationOutcome::Forward)
1455        }
1456    }
1457
1458    /// Handle an incoming message from the client or conductor, depending on `Self`
1459    fn handle_dispatch(
1460        &self,
1461        message: Dispatch,
1462        connection: ConnectionTo<Self::Counterpart>,
1463        conductor_tx: &mut mpsc::Sender<ConductorMessage>,
1464    ) -> impl Future<Output = Result<Handled<Dispatch>, agent_client_protocol::Error>> + Send;
1465}
1466
1467/// Result of handling the conductor's first protocol message.
1468#[cfg(feature = "unstable_protocol_v2")]
1469#[derive(Debug)]
1470pub enum InitializationOutcome {
1471    /// Initialization succeeded; forward this request into the component chain.
1472    Forward(Dispatch),
1473    /// An error response was sent and the connection must reject later traffic.
1474    Rejected(Error),
1475}
1476
1477#[cfg(feature = "unstable_protocol_v2")]
1478fn reject_initialization(
1479    responder: agent_client_protocol::Responder,
1480    error: Error,
1481) -> Result<InitializationOutcome, Error> {
1482    responder.respond_with_error(error.clone())?;
1483    Ok(InitializationOutcome::Rejected(error))
1484}
1485
1486#[cfg(feature = "unstable_protocol_v2")]
1487async fn initialize_agent_for_selected_protocol(
1488    message: Dispatch,
1489    client_connection: ConnectionTo<Client>,
1490    instantiator: Box<dyn InstantiateProxiesAndAgent>,
1491    runner: &mut ConductorRunner<Agent>,
1492) -> Result<InitializationOutcome, Error> {
1493    let invalid_request = || Error::invalid_request().data("expected `initialize` request");
1494
1495    let Dispatch::Request(raw_request, init_responder) = message else {
1496        let error = invalid_request();
1497        if let Dispatch::Response(_, router) = message {
1498            router.route_with_error(error.clone())?;
1499        }
1500        return Ok(InitializationOutcome::Rejected(error));
1501    };
1502    if !InitializeRequest::matches_method(raw_request.method()) {
1503        return reject_initialization(init_responder, invalid_request());
1504    }
1505
1506    let selection = match InitializeProtocol::from_request(&raw_request) {
1507        Ok(selection) => selection,
1508        Err(error) => return reject_initialization(init_responder, error),
1509    };
1510    let protocol = selection.protocol;
1511
1512    // Select the schema before deserializing. Parsing v2 as permissive v1
1513    // would silently discard fields such as `info` and `capabilities`.
1514    let initialization = match protocol {
1515        InitializeProtocol::V1 => {
1516            let mut init_request = match InitializeRequest::parse_message(
1517                raw_request.method(),
1518                raw_request.params(),
1519            ) {
1520                Ok(request) => request,
1521                Err(error) => return reject_initialization(init_responder, error),
1522            };
1523            init_request.protocol_version = protocol.version();
1524            let original_request = init_request.clone();
1525            match instantiator
1526                .instantiate_proxies_and_agent(init_request)
1527                .await
1528            {
1529                Ok((mut request, proxies, agent)) => {
1530                    request.protocol_version = protocol.version();
1531                    forwarded_initialize_request(
1532                        &raw_request,
1533                        selection,
1534                        &original_request,
1535                        request,
1536                    )
1537                    .map(|request| (request, proxies, agent))
1538                }
1539                Err(error) => Err(error),
1540            }
1541        }
1542        InitializeProtocol::V2 => {
1543            let mut init_request = match v2::InitializeRequest::parse_message(
1544                raw_request.method(),
1545                raw_request.params(),
1546            ) {
1547                Ok(request) => request,
1548                Err(error) => return reject_initialization(init_responder, error),
1549            };
1550            init_request.protocol_version = protocol.version();
1551            let original_request = init_request.clone();
1552            match instantiator
1553                .instantiate_v2_proxies_and_agent(init_request)
1554                .await
1555            {
1556                Ok((mut request, proxies, agent)) => {
1557                    request.protocol_version = protocol.version();
1558                    forwarded_initialize_request(
1559                        &raw_request,
1560                        selection,
1561                        &original_request,
1562                        request,
1563                    )
1564                    .map(|request| (request, proxies, agent))
1565                }
1566                Err(error) => Err(error),
1567            }
1568        }
1569    };
1570    let (modified_req, proxy_components, agent_component) = match initialization {
1571        Ok(initialization) => initialization,
1572        Err(error) => return reject_initialization(init_responder, error),
1573    };
1574
1575    debug!(?agent_component, "spawning agent");
1576
1577    let agent_builder = match protocol {
1578        InitializeProtocol::V1 => Builder::new(Client),
1579        InitializeProtocol::V2 => Builder::new(Client).with_v2_protocol_guard(),
1580    };
1581    let connection_to_agent = client_connection.spawn_connection(
1582        agent_builder
1583            .name("conductor-to-agent")
1584            .on_receive_dispatch(
1585                {
1586                    let mut conductor_tx = runner.conductor_tx.clone();
1587                    async move |dispatch: Dispatch, _cx| {
1588                        conductor_tx
1589                            .send(ConductorMessage::RightToLeft {
1590                                source_component_index: SourceComponentIndex::Successor,
1591                                message: dispatch,
1592                            })
1593                            .await
1594                            .map_err(agent_client_protocol::util::internal_error)
1595                    }
1596                },
1597                agent_client_protocol::on_receive_dispatch!(),
1598            ),
1599        agent_component,
1600    )?;
1601    runner.successor = Arc::new(connection_to_agent);
1602
1603    runner.spawn_proxies(client_connection, proxy_components)?;
1604
1605    Ok(InitializationOutcome::Forward(Dispatch::Request(
1606        modified_req,
1607        init_responder,
1608    )))
1609}
1610
1611/// Conductor acting as an agent
1612impl ConductorHostRole for Agent {
1613    type Instantiator = Box<dyn InstantiateProxiesAndAgent>;
1614
1615    async fn initialize(
1616        &self,
1617        message: Dispatch,
1618        client_connection: ConnectionTo<Client>,
1619        instantiator: Self::Instantiator,
1620        runner: &mut ConductorRunner<Self>,
1621    ) -> Result<Dispatch, agent_client_protocol::Error> {
1622        let invalid_request = || Error::invalid_request().data("expected `initialize` request");
1623
1624        let Dispatch::Request(request, init_responder) = message else {
1625            if let Dispatch::Response(_, router) = message {
1626                router.route_with_error(invalid_request())?;
1627            }
1628            return Err(invalid_request());
1629        };
1630        if !InitializeRequest::matches_method(request.method()) {
1631            init_responder.respond_with_error(invalid_request())?;
1632            return Err(invalid_request());
1633        }
1634
1635        let init_request =
1636            match InitializeRequest::parse_message(request.method(), request.params()) {
1637                Ok(request) => request,
1638                Err(error) => {
1639                    init_responder.respond_with_error(error)?;
1640                    return Err(invalid_request());
1641                }
1642            };
1643
1644        let (modified_req, proxy_components, agent_component) = instantiator
1645            .instantiate_proxies_and_agent(init_request)
1646            .await?;
1647
1648        debug!(?agent_component, "spawning agent");
1649
1650        let connection_to_agent = client_connection.spawn_connection(
1651            Client
1652                .builder()
1653                .name("conductor-to-agent")
1654                .on_receive_dispatch(
1655                    {
1656                        let mut conductor_tx = runner.conductor_tx.clone();
1657                        async move |dispatch: Dispatch, _cx| {
1658                            conductor_tx
1659                                .send(ConductorMessage::RightToLeft {
1660                                    source_component_index: SourceComponentIndex::Successor,
1661                                    message: dispatch,
1662                                })
1663                                .await
1664                                .map_err(agent_client_protocol::util::internal_error)
1665                        }
1666                    },
1667                    agent_client_protocol::on_receive_dispatch!(),
1668                ),
1669            agent_component,
1670        )?;
1671        runner.successor = Arc::new(connection_to_agent);
1672
1673        runner.spawn_proxies(client_connection.clone(), proxy_components)?;
1674
1675        Ok(Dispatch::Request(
1676            modified_req.to_untyped_message()?,
1677            init_responder,
1678        ))
1679    }
1680
1681    #[cfg(feature = "unstable_protocol_v2")]
1682    async fn initialize_with_outcome(
1683        &self,
1684        message: Dispatch,
1685        client_connection: ConnectionTo<Client>,
1686        instantiator: Self::Instantiator,
1687        runner: &mut ConductorRunner<Self>,
1688    ) -> Result<InitializationOutcome, agent_client_protocol::Error> {
1689        initialize_agent_for_selected_protocol(message, client_connection, instantiator, runner)
1690            .await
1691    }
1692
1693    async fn handle_dispatch(
1694        &self,
1695        message: Dispatch,
1696        client_connection: ConnectionTo<Client>,
1697        conductor_tx: &mut mpsc::Sender<ConductorMessage>,
1698    ) -> Result<Handled<Dispatch>, agent_client_protocol::Error> {
1699        tracing::debug!(
1700            method = ?message.method(),
1701            "ConductorToClient::handle_dispatch"
1702        );
1703        MatchDispatchFrom::new(message, &client_connection)
1704            // Any incoming messages from the client are client-to-agent messages targeting the first component.
1705            .if_dispatch_from(Client, async move |message: Dispatch| {
1706                tracing::debug!(
1707                    method = ?message.method(),
1708                    "ConductorToClient::handle_dispatch - matched Client"
1709                );
1710                ConductorImpl::<Self>::incoming_message_from_client(conductor_tx, message).await
1711            })
1712            .await
1713            .done()
1714    }
1715}
1716
1717#[cfg(feature = "unstable_protocol_v2")]
1718async fn initialize_proxy_for_selected_protocol(
1719    message: Dispatch,
1720    client_connection: ConnectionTo<Conductor>,
1721    instantiator: Box<dyn InstantiateProxies>,
1722    runner: &mut ConductorRunner<Proxy>,
1723) -> Result<InitializationOutcome, Error> {
1724    let invalid_request = || Error::invalid_request().data("expected `initialize` request");
1725
1726    let Dispatch::Request(raw_request, init_responder) = message else {
1727        let error = invalid_request();
1728        if let Dispatch::Response(_, router) = message {
1729            router.route_with_error(error.clone())?;
1730        }
1731        return Ok(InitializationOutcome::Rejected(error));
1732    };
1733    if !InitializeProxyRequest::matches_method(raw_request.method()) {
1734        return reject_initialization(init_responder, invalid_request());
1735    }
1736
1737    let selection = match InitializeProtocol::from_request(&raw_request) {
1738        Ok(selection) => selection,
1739        Err(error) => return reject_initialization(init_responder, error),
1740    };
1741    let protocol = selection.protocol;
1742
1743    tracing::debug!(?protocol, "ensure_initialized: proxy initialize");
1744
1745    let initialization = match protocol {
1746        InitializeProtocol::V1 => {
1747            let InitializeProxyRequest { mut initialize } =
1748                match InitializeProxyRequest::parse_message(
1749                    raw_request.method(),
1750                    raw_request.params(),
1751                ) {
1752                    Ok(request) => request,
1753                    Err(error) => return reject_initialization(init_responder, error),
1754                };
1755            initialize.protocol_version = protocol.version();
1756            let original_request = initialize.clone();
1757            match instantiator.instantiate_proxies(initialize).await {
1758                Ok((mut request, proxies)) => {
1759                    request.protocol_version = protocol.version();
1760                    forwarded_initialize_request(
1761                        &raw_request,
1762                        selection,
1763                        &original_request,
1764                        request,
1765                    )
1766                    .map(|request| (request, proxies))
1767                }
1768                Err(error) => Err(error),
1769            }
1770        }
1771        InitializeProtocol::V2 => {
1772            let v2::InitializeProxyRequest { mut initialize } =
1773                match v2::InitializeProxyRequest::parse_message(
1774                    raw_request.method(),
1775                    raw_request.params(),
1776                ) {
1777                    Ok(request) => request,
1778                    Err(error) => return reject_initialization(init_responder, error),
1779                };
1780            initialize.protocol_version = protocol.version();
1781            let original_request = initialize.clone();
1782            match instantiator.instantiate_v2_proxies(initialize).await {
1783                Ok((mut request, proxies)) => {
1784                    request.protocol_version = protocol.version();
1785                    forwarded_initialize_request(
1786                        &raw_request,
1787                        selection,
1788                        &original_request,
1789                        request,
1790                    )
1791                    .map(|request| (request, proxies))
1792                }
1793                Err(error) => Err(error),
1794            }
1795        }
1796    };
1797    let (modified_req, proxy_components) = match initialization {
1798        Ok(initialization) => initialization,
1799        Err(error) => return reject_initialization(init_responder, error),
1800    };
1801
1802    runner.successor = Arc::new(GrandSuccessor);
1803    runner.spawn_proxies(client_connection, proxy_components)?;
1804
1805    Ok(InitializationOutcome::Forward(Dispatch::Request(
1806        modified_req,
1807        init_responder,
1808    )))
1809}
1810
1811/// Conductor acting as a proxy
1812impl ConductorHostRole for Proxy {
1813    type Instantiator = Box<dyn InstantiateProxies>;
1814
1815    async fn initialize(
1816        &self,
1817        message: Dispatch,
1818        client_connection: ConnectionTo<Conductor>,
1819        instantiator: Self::Instantiator,
1820        runner: &mut ConductorRunner<Self>,
1821    ) -> Result<Dispatch, agent_client_protocol::Error> {
1822        let invalid_request = || Error::invalid_request().data("expected `initialize` request");
1823
1824        let Dispatch::Request(request, init_responder) = message else {
1825            if let Dispatch::Response(_, router) = message {
1826                router.route_with_error(invalid_request())?;
1827            }
1828            return Err(invalid_request());
1829        };
1830        if !InitializeProxyRequest::matches_method(request.method()) {
1831            init_responder.respond_with_error(invalid_request())?;
1832            return Err(invalid_request());
1833        }
1834
1835        let InitializeProxyRequest { initialize } =
1836            match InitializeProxyRequest::parse_message(request.method(), request.params()) {
1837                Ok(request) => request,
1838                Err(error) => {
1839                    init_responder.respond_with_error(error)?;
1840                    return Err(invalid_request());
1841                }
1842            };
1843
1844        tracing::debug!("ensure_initialized: InitializeProxyRequest (proxy mode)");
1845
1846        let (modified_req, proxy_components) = instantiator.instantiate_proxies(initialize).await?;
1847
1848        runner.successor = Arc::new(GrandSuccessor);
1849        runner.spawn_proxies(client_connection.clone(), proxy_components)?;
1850
1851        Ok(Dispatch::Request(
1852            modified_req.to_untyped_message()?,
1853            init_responder,
1854        ))
1855    }
1856
1857    #[cfg(feature = "unstable_protocol_v2")]
1858    async fn initialize_with_outcome(
1859        &self,
1860        message: Dispatch,
1861        client_connection: ConnectionTo<Conductor>,
1862        instantiator: Self::Instantiator,
1863        runner: &mut ConductorRunner<Self>,
1864    ) -> Result<InitializationOutcome, agent_client_protocol::Error> {
1865        initialize_proxy_for_selected_protocol(message, client_connection, instantiator, runner)
1866            .await
1867    }
1868
1869    async fn handle_dispatch(
1870        &self,
1871        message: Dispatch,
1872        client_connection: ConnectionTo<Conductor>,
1873        conductor_tx: &mut mpsc::Sender<ConductorMessage>,
1874    ) -> Result<Handled<Dispatch>, agent_client_protocol::Error> {
1875        tracing::debug!(
1876            method = ?message.method(),
1877            ?message,
1878            "ConductorToConductor::handle_dispatch"
1879        );
1880        MatchDispatchFrom::new(message, &client_connection)
1881            .if_dispatch_from(Agent, {
1882                // Messages from our successor arrive already unwrapped
1883                // (RemoteRoleStyle::Successor strips the SuccessorMessage envelope).
1884                async |message: Dispatch| {
1885                    tracing::debug!(
1886                        method = ?message.method(),
1887                        "ConductorToConductor::handle_dispatch - matched Agent"
1888                    );
1889                    let mut conductor_tx = conductor_tx.clone();
1890                    ConductorImpl::<Self>::incoming_message_from_agent(&mut conductor_tx, message)
1891                        .await
1892                }
1893            })
1894            .await
1895            // Any incoming messages from the client are client-to-agent messages targeting the first component.
1896            .if_dispatch_from(Client, async |message: Dispatch| {
1897                tracing::debug!(
1898                    method = ?message.method(),
1899                    "ConductorToConductor::handle_dispatch - matched Client"
1900                );
1901                let mut conductor_tx = conductor_tx.clone();
1902                ConductorImpl::<Self>::incoming_message_from_client(&mut conductor_tx, message)
1903                    .await
1904            })
1905            .await
1906            .done()
1907    }
1908}
1909
1910pub trait ConductorSuccessor<Host: ConductorHostRole>: Send + Sync + 'static {
1911    /// Send a message to the successor.
1912    fn send_message<'a>(
1913        &self,
1914        message: Dispatch,
1915        connection_to_conductor: ConnectionTo<Host::Counterpart>,
1916        runner: &'a mut ConductorRunner<Host>,
1917    ) -> BoxFuture<'a, Result<(), agent_client_protocol::Error>>;
1918}
1919
1920impl<Host: ConductorHostRole> ConductorSuccessor<Host> for agent_client_protocol::Error {
1921    fn send_message<'a>(
1922        &self,
1923        _message: Dispatch,
1924        _connection_to_conductor: ConnectionTo<Host::Counterpart>,
1925        _runner: &'a mut ConductorRunner<Host>,
1926    ) -> BoxFuture<'a, Result<(), agent_client_protocol::Error>> {
1927        let error = self.clone();
1928        Box::pin(std::future::ready(Err(error)))
1929    }
1930}
1931
1932/// A dummy type handling messages sent to the conductor's
1933/// successor when it is acting as a proxy.
1934struct GrandSuccessor;
1935
1936/// When the conductor is acting as an proxy, messages sent by
1937/// the last proxy go to the conductor's successor.
1938///
1939/// ```text
1940/// client --> Conductor -----------------------------> GrandSuccessor
1941///            |                                  |
1942///            +-> Proxy[0] -> ... -> Proxy[n-1] -+
1943/// ```
1944impl ConductorSuccessor<Proxy> for GrandSuccessor {
1945    fn send_message<'a>(
1946        &self,
1947        message: Dispatch,
1948        connection: ConnectionTo<Conductor>,
1949        _runner: &'a mut ConductorRunner<Proxy>,
1950    ) -> BoxFuture<'a, Result<(), agent_client_protocol::Error>> {
1951        Box::pin(async move {
1952            debug!("Proxy mode: forwarding successor message to conductor's successor");
1953            connection.send_proxied_message_to(Agent, message)
1954        })
1955    }
1956}
1957
1958/// When the conductor is acting as an agent, messages sent by
1959/// the last proxy to its successor go to the internal agent
1960/// (`self`).
1961impl ConductorSuccessor<Agent> for ConnectionTo<Agent> {
1962    fn send_message<'a>(
1963        &self,
1964        message: Dispatch,
1965        connection: ConnectionTo<Client>,
1966        runner: &'a mut ConductorRunner<Agent>,
1967    ) -> BoxFuture<'a, Result<(), agent_client_protocol::Error>> {
1968        let connection_to_agent = self.clone();
1969        Box::pin(async move {
1970            debug!("Proxy mode: forwarding successor message to conductor's successor");
1971            runner
1972                .forward_message_to_agent(connection, message, connection_to_agent)
1973                .await
1974        })
1975    }
1976}