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    #[must_use]
1104    fn instantiate_v2_proxies(
1105        self: Box<Self>,
1106        req: v2::InitializeRequest,
1107    ) -> futures::future::BoxFuture<
1108        'static,
1109        Result<(v2::InitializeRequest, Vec<DynConnectTo<Conductor>>), agent_client_protocol::Error>,
1110    > {
1111        drop((self, req));
1112        Box::pin(async {
1113            Err(Error::invalid_request()
1114                .data("this conductor proxy instantiator does not support ACP protocol v2"))
1115        })
1116    }
1117}
1118
1119/// Simple implementation: provide all proxy components unconditionally.
1120///
1121/// Requires `T: ConnectTo<Conductor>`.
1122impl<T> InstantiateProxies for Vec<T>
1123where
1124    T: ConnectTo<Conductor> + 'static,
1125{
1126    fn instantiate_proxies(
1127        self: Box<Self>,
1128        req: InitializeRequest,
1129    ) -> futures::future::BoxFuture<
1130        'static,
1131        Result<(InitializeRequest, Vec<DynConnectTo<Conductor>>), agent_client_protocol::Error>,
1132    > {
1133        Box::pin(async move {
1134            let components: Vec<DynConnectTo<Conductor>> =
1135                (*self).into_iter().map(|c| DynConnectTo::new(c)).collect();
1136            Ok((req, components))
1137        })
1138    }
1139
1140    #[cfg(feature = "unstable_protocol_v2")]
1141    fn instantiate_v2_proxies(
1142        self: Box<Self>,
1143        req: v2::InitializeRequest,
1144    ) -> futures::future::BoxFuture<
1145        'static,
1146        Result<(v2::InitializeRequest, Vec<DynConnectTo<Conductor>>), agent_client_protocol::Error>,
1147    > {
1148        Box::pin(async move {
1149            let components = (*self).into_iter().map(DynConnectTo::new).collect();
1150            Ok((req, components))
1151        })
1152    }
1153}
1154
1155/// Dynamic implementation: closure receives the Initialize request and returns proxies.
1156impl<F, Fut> InstantiateProxies for F
1157where
1158    F: FnOnce(InitializeRequest) -> Fut + Send + 'static,
1159    Fut: std::future::Future<
1160            Output = Result<
1161                (InitializeRequest, Vec<DynConnectTo<Conductor>>),
1162                agent_client_protocol::Error,
1163            >,
1164        > + Send
1165        + 'static,
1166{
1167    fn instantiate_proxies(
1168        self: Box<Self>,
1169        req: InitializeRequest,
1170    ) -> futures::future::BoxFuture<
1171        'static,
1172        Result<(InitializeRequest, Vec<DynConnectTo<Conductor>>), agent_client_protocol::Error>,
1173    > {
1174        Box::pin(async move { (*self)(req).await })
1175    }
1176}
1177
1178/// Trait for lazy proxy and agent instantiation (agent mode).
1179///
1180/// Used by conductors in agent mode (`ConductorToClient`) where there are
1181/// zero or more proxies followed by an agent component.
1182pub trait InstantiateProxiesAndAgent: Send {
1183    /// Instantiate proxy and agent components based on the Initialize request.
1184    ///
1185    /// Returns the (possibly modified) request, a vector of proxy components
1186    /// (typed as `DynConnectTo<Conductor>`), and the agent component
1187    /// (typed as `DynConnectTo<Client>`).
1188    fn instantiate_proxies_and_agent(
1189        self: Box<Self>,
1190        req: InitializeRequest,
1191    ) -> futures::future::BoxFuture<
1192        'static,
1193        Result<
1194            (
1195                InitializeRequest,
1196                Vec<DynConnectTo<Conductor>>,
1197                DynConnectTo<Client>,
1198            ),
1199            agent_client_protocol::Error,
1200        >,
1201    >;
1202
1203    /// Instantiate proxy and agent components for a protocol-v2 connection.
1204    ///
1205    /// Implementors that only support v1 can rely on the default rejection.
1206    /// [`AgentOnly`] and [`ProxiesAndAgent`] support both versions.
1207    #[cfg(feature = "unstable_protocol_v2")]
1208    #[must_use]
1209    fn instantiate_v2_proxies_and_agent(
1210        self: Box<Self>,
1211        req: v2::InitializeRequest,
1212    ) -> futures::future::BoxFuture<
1213        'static,
1214        Result<
1215            (
1216                v2::InitializeRequest,
1217                Vec<DynConnectTo<Conductor>>,
1218                DynConnectTo<Client>,
1219            ),
1220            agent_client_protocol::Error,
1221        >,
1222    > {
1223        drop((self, req));
1224        Box::pin(async {
1225            Err(Error::invalid_request()
1226                .data("this conductor agent instantiator does not support ACP protocol v2"))
1227        })
1228    }
1229}
1230
1231/// Wrapper to convert a single agent component (no proxies) into InstantiateProxiesAndAgent.
1232#[derive(Debug)]
1233pub struct AgentOnly<A>(pub A);
1234
1235impl<A: ConnectTo<Client> + 'static> InstantiateProxiesAndAgent for AgentOnly<A> {
1236    fn instantiate_proxies_and_agent(
1237        self: Box<Self>,
1238        req: InitializeRequest,
1239    ) -> futures::future::BoxFuture<
1240        'static,
1241        Result<
1242            (
1243                InitializeRequest,
1244                Vec<DynConnectTo<Conductor>>,
1245                DynConnectTo<Client>,
1246            ),
1247            agent_client_protocol::Error,
1248        >,
1249    > {
1250        Box::pin(async move { Ok((req, Vec::new(), DynConnectTo::new(self.0))) })
1251    }
1252
1253    #[cfg(feature = "unstable_protocol_v2")]
1254    fn instantiate_v2_proxies_and_agent(
1255        self: Box<Self>,
1256        req: v2::InitializeRequest,
1257    ) -> futures::future::BoxFuture<
1258        'static,
1259        Result<
1260            (
1261                v2::InitializeRequest,
1262                Vec<DynConnectTo<Conductor>>,
1263                DynConnectTo<Client>,
1264            ),
1265            agent_client_protocol::Error,
1266        >,
1267    > {
1268        Box::pin(async move { Ok((req, Vec::new(), DynConnectTo::new(self.0))) })
1269    }
1270}
1271
1272/// Builder for creating proxies and agent components.
1273///
1274/// # Example
1275/// ```ignore
1276/// ProxiesAndAgent::new(ElizaAgent::new())
1277///     .proxy(LoggingProxy::new())
1278///     .proxy(AuthProxy::new())
1279/// ```
1280#[derive(Debug)]
1281pub struct ProxiesAndAgent {
1282    proxies: Vec<DynConnectTo<Conductor>>,
1283    agent: DynConnectTo<Client>,
1284}
1285
1286impl ProxiesAndAgent {
1287    /// Create a new builder with the given agent component.
1288    pub fn new(agent: impl ConnectTo<Client> + 'static) -> Self {
1289        Self {
1290            proxies: vec![],
1291            agent: DynConnectTo::new(agent),
1292        }
1293    }
1294
1295    /// Add a single proxy component.
1296    #[must_use]
1297    pub fn proxy(mut self, proxy: impl ConnectTo<Conductor> + 'static) -> Self {
1298        self.proxies.push(DynConnectTo::new(proxy));
1299        self
1300    }
1301
1302    /// Add multiple proxy components.
1303    #[must_use]
1304    pub fn proxies<P, I>(mut self, proxies: I) -> Self
1305    where
1306        P: ConnectTo<Conductor> + 'static,
1307        I: IntoIterator<Item = P>,
1308    {
1309        self.proxies
1310            .extend(proxies.into_iter().map(DynConnectTo::new));
1311        self
1312    }
1313}
1314
1315impl InstantiateProxiesAndAgent for ProxiesAndAgent {
1316    fn instantiate_proxies_and_agent(
1317        self: Box<Self>,
1318        req: InitializeRequest,
1319    ) -> futures::future::BoxFuture<
1320        'static,
1321        Result<
1322            (
1323                InitializeRequest,
1324                Vec<DynConnectTo<Conductor>>,
1325                DynConnectTo<Client>,
1326            ),
1327            agent_client_protocol::Error,
1328        >,
1329    > {
1330        Box::pin(async move { Ok((req, self.proxies, self.agent)) })
1331    }
1332
1333    #[cfg(feature = "unstable_protocol_v2")]
1334    fn instantiate_v2_proxies_and_agent(
1335        self: Box<Self>,
1336        req: v2::InitializeRequest,
1337    ) -> futures::future::BoxFuture<
1338        'static,
1339        Result<
1340            (
1341                v2::InitializeRequest,
1342                Vec<DynConnectTo<Conductor>>,
1343                DynConnectTo<Client>,
1344            ),
1345            agent_client_protocol::Error,
1346        >,
1347    > {
1348        Box::pin(async move { Ok((req, self.proxies, self.agent)) })
1349    }
1350}
1351
1352/// Dynamic implementation: closure receives the Initialize request and returns proxies + agent.
1353impl<F, Fut> InstantiateProxiesAndAgent for F
1354where
1355    F: FnOnce(InitializeRequest) -> Fut + Send + 'static,
1356    Fut: std::future::Future<
1357            Output = Result<
1358                (
1359                    InitializeRequest,
1360                    Vec<DynConnectTo<Conductor>>,
1361                    DynConnectTo<Client>,
1362                ),
1363                agent_client_protocol::Error,
1364            >,
1365        > + Send
1366        + 'static,
1367{
1368    fn instantiate_proxies_and_agent(
1369        self: Box<Self>,
1370        req: InitializeRequest,
1371    ) -> futures::future::BoxFuture<
1372        'static,
1373        Result<
1374            (
1375                InitializeRequest,
1376                Vec<DynConnectTo<Conductor>>,
1377                DynConnectTo<Client>,
1378            ),
1379            agent_client_protocol::Error,
1380        >,
1381    > {
1382        Box::pin(async move { (*self)(req).await })
1383    }
1384}
1385
1386/// Messages sent to the conductor's main event loop for routing.
1387///
1388/// These messages enable the conductor to route communication between:
1389/// - The editor and the first component
1390/// - Components and their successors in the chain
1391/// - Components and their clients (editor or predecessor)
1392///
1393/// All spawned tasks send messages via this enum through a shared channel,
1394/// allowing centralized routing logic in the conductor runner's event loop.
1395#[derive(Debug)]
1396pub enum ConductorMessage {
1397    /// If this message is a request or notification, then it is going "left-to-right"
1398    /// (e.g., a component making a request of its successor).
1399    ///
1400    /// If this message is a response, then it is going right-to-left
1401    /// (i.e., the successor answering a request made by its predecessor).
1402    LeftToRight {
1403        target_component_index: usize,
1404        message: Dispatch,
1405    },
1406
1407    /// If this message is a request or notification, then it is going "right-to-left"
1408    /// (e.g., a component making a request of its predecessor).
1409    ///
1410    /// If this message is a response, then it is going "left-to-right"
1411    /// (i.e., the predecessor answering a request made by its successor).
1412    RightToLeft {
1413        source_component_index: SourceComponentIndex,
1414        message: Dispatch,
1415    },
1416}
1417
1418/// Trait implemented for the two links the conductor can use:
1419///
1420/// * 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
1421/// * 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
1422pub trait ConductorHostRole: Role<Counterpart: HasPeer<Client>> {
1423    /// The type used to instantiate components for this link type.
1424    type Instantiator: Send;
1425
1426    /// Handle initialization: parse the init request, instantiate components, and spawn them.
1427    ///
1428    /// Takes ownership of the instantiator and returns the (possibly modified) init request
1429    /// wrapped in a Dispatch for forwarding.
1430    fn initialize(
1431        &self,
1432        message: Dispatch,
1433        connection: ConnectionTo<Self::Counterpart>,
1434        instantiator: Self::Instantiator,
1435        runner: &mut ConductorRunner<Self>,
1436    ) -> impl Future<Output = Result<Dispatch, agent_client_protocol::Error>> + Send;
1437
1438    /// Handle initialization while distinguishing a protocol rejection from a
1439    /// fatal connection failure.
1440    ///
1441    /// The conductor uses this hook only when `unstable_protocol_v2` is
1442    /// enabled. The default preserves existing implementations by delegating
1443    /// to [`Self::initialize`].
1444    #[cfg(feature = "unstable_protocol_v2")]
1445    fn initialize_with_outcome(
1446        &self,
1447        message: Dispatch,
1448        connection: ConnectionTo<Self::Counterpart>,
1449        instantiator: Self::Instantiator,
1450        runner: &mut ConductorRunner<Self>,
1451    ) -> impl Future<Output = Result<InitializationOutcome, agent_client_protocol::Error>> + Send
1452    {
1453        async move {
1454            self.initialize(message, connection, instantiator, runner)
1455                .await
1456                .map(InitializationOutcome::Forward)
1457        }
1458    }
1459
1460    /// Handle an incoming message from the client or conductor, depending on `Self`
1461    fn handle_dispatch(
1462        &self,
1463        message: Dispatch,
1464        connection: ConnectionTo<Self::Counterpart>,
1465        conductor_tx: &mut mpsc::Sender<ConductorMessage>,
1466    ) -> impl Future<Output = Result<Handled<Dispatch>, agent_client_protocol::Error>> + Send;
1467}
1468
1469/// Result of handling the conductor's first protocol message.
1470#[cfg(feature = "unstable_protocol_v2")]
1471#[derive(Debug)]
1472pub enum InitializationOutcome {
1473    /// Initialization succeeded; forward this request into the component chain.
1474    Forward(Dispatch),
1475    /// An error response was sent and the connection must reject later traffic.
1476    Rejected(Error),
1477}
1478
1479#[cfg(feature = "unstable_protocol_v2")]
1480fn reject_initialization(
1481    responder: agent_client_protocol::Responder,
1482    error: Error,
1483) -> Result<InitializationOutcome, Error> {
1484    responder.respond_with_error(error.clone())?;
1485    Ok(InitializationOutcome::Rejected(error))
1486}
1487
1488#[cfg(feature = "unstable_protocol_v2")]
1489async fn initialize_agent_for_selected_protocol(
1490    message: Dispatch,
1491    client_connection: ConnectionTo<Client>,
1492    instantiator: Box<dyn InstantiateProxiesAndAgent>,
1493    runner: &mut ConductorRunner<Agent>,
1494) -> Result<InitializationOutcome, Error> {
1495    let invalid_request = || Error::invalid_request().data("expected `initialize` request");
1496
1497    let Dispatch::Request(raw_request, init_responder) = message else {
1498        let error = invalid_request();
1499        if let Dispatch::Response(_, router) = message {
1500            router.route_with_error(error.clone())?;
1501        }
1502        return Ok(InitializationOutcome::Rejected(error));
1503    };
1504    if !InitializeRequest::matches_method(raw_request.method()) {
1505        return reject_initialization(init_responder, invalid_request());
1506    }
1507
1508    let selection = match InitializeProtocol::from_request(&raw_request) {
1509        Ok(selection) => selection,
1510        Err(error) => return reject_initialization(init_responder, error),
1511    };
1512    let protocol = selection.protocol;
1513
1514    // Select the schema before deserializing. Parsing v2 as permissive v1
1515    // would silently discard fields such as `info` and `capabilities`.
1516    let initialization = match protocol {
1517        InitializeProtocol::V1 => {
1518            let mut init_request = match InitializeRequest::parse_message(
1519                raw_request.method(),
1520                raw_request.params(),
1521            ) {
1522                Ok(request) => request,
1523                Err(error) => return reject_initialization(init_responder, error),
1524            };
1525            init_request.protocol_version = protocol.version();
1526            let original_request = init_request.clone();
1527            match instantiator
1528                .instantiate_proxies_and_agent(init_request)
1529                .await
1530            {
1531                Ok((mut request, proxies, agent)) => {
1532                    request.protocol_version = protocol.version();
1533                    forwarded_initialize_request(
1534                        &raw_request,
1535                        selection,
1536                        &original_request,
1537                        request,
1538                    )
1539                    .map(|request| (request, proxies, agent))
1540                }
1541                Err(error) => Err(error),
1542            }
1543        }
1544        InitializeProtocol::V2 => {
1545            let mut init_request = match v2::InitializeRequest::parse_message(
1546                raw_request.method(),
1547                raw_request.params(),
1548            ) {
1549                Ok(request) => request,
1550                Err(error) => return reject_initialization(init_responder, error),
1551            };
1552            init_request.protocol_version = protocol.version();
1553            let original_request = init_request.clone();
1554            match instantiator
1555                .instantiate_v2_proxies_and_agent(init_request)
1556                .await
1557            {
1558                Ok((mut request, proxies, agent)) => {
1559                    request.protocol_version = protocol.version();
1560                    forwarded_initialize_request(
1561                        &raw_request,
1562                        selection,
1563                        &original_request,
1564                        request,
1565                    )
1566                    .map(|request| (request, proxies, agent))
1567                }
1568                Err(error) => Err(error),
1569            }
1570        }
1571    };
1572    let (modified_req, proxy_components, agent_component) = match initialization {
1573        Ok(initialization) => initialization,
1574        Err(error) => return reject_initialization(init_responder, error),
1575    };
1576
1577    debug!(?agent_component, "spawning agent");
1578
1579    let agent_builder = match protocol {
1580        InitializeProtocol::V1 => Builder::new(Client),
1581        InitializeProtocol::V2 => Builder::new(Client).with_v2_protocol_guard(),
1582    };
1583    let connection_to_agent = client_connection.spawn_connection(
1584        agent_builder
1585            .name("conductor-to-agent")
1586            .on_receive_dispatch(
1587                {
1588                    let mut conductor_tx = runner.conductor_tx.clone();
1589                    async move |dispatch: Dispatch, _cx| {
1590                        conductor_tx
1591                            .send(ConductorMessage::RightToLeft {
1592                                source_component_index: SourceComponentIndex::Successor,
1593                                message: dispatch,
1594                            })
1595                            .await
1596                            .map_err(agent_client_protocol::util::internal_error)
1597                    }
1598                },
1599                agent_client_protocol::on_receive_dispatch!(),
1600            ),
1601        agent_component,
1602    )?;
1603    runner.successor = Arc::new(connection_to_agent);
1604
1605    runner.spawn_proxies(client_connection, proxy_components)?;
1606
1607    Ok(InitializationOutcome::Forward(Dispatch::Request(
1608        modified_req,
1609        init_responder,
1610    )))
1611}
1612
1613/// Conductor acting as an agent
1614impl ConductorHostRole for Agent {
1615    type Instantiator = Box<dyn InstantiateProxiesAndAgent>;
1616
1617    async fn initialize(
1618        &self,
1619        message: Dispatch,
1620        client_connection: ConnectionTo<Client>,
1621        instantiator: Self::Instantiator,
1622        runner: &mut ConductorRunner<Self>,
1623    ) -> Result<Dispatch, agent_client_protocol::Error> {
1624        let invalid_request = || Error::invalid_request().data("expected `initialize` request");
1625
1626        let Dispatch::Request(request, init_responder) = message else {
1627            if let Dispatch::Response(_, router) = message {
1628                router.route_with_error(invalid_request())?;
1629            }
1630            return Err(invalid_request());
1631        };
1632        if !InitializeRequest::matches_method(request.method()) {
1633            init_responder.respond_with_error(invalid_request())?;
1634            return Err(invalid_request());
1635        }
1636
1637        let init_request =
1638            match InitializeRequest::parse_message(request.method(), request.params()) {
1639                Ok(request) => request,
1640                Err(error) => {
1641                    init_responder.respond_with_error(error)?;
1642                    return Err(invalid_request());
1643                }
1644            };
1645
1646        let (modified_req, proxy_components, agent_component) = instantiator
1647            .instantiate_proxies_and_agent(init_request)
1648            .await?;
1649
1650        debug!(?agent_component, "spawning agent");
1651
1652        let connection_to_agent = client_connection.spawn_connection(
1653            Client
1654                .builder()
1655                .name("conductor-to-agent")
1656                .on_receive_dispatch(
1657                    {
1658                        let mut conductor_tx = runner.conductor_tx.clone();
1659                        async move |dispatch: Dispatch, _cx| {
1660                            conductor_tx
1661                                .send(ConductorMessage::RightToLeft {
1662                                    source_component_index: SourceComponentIndex::Successor,
1663                                    message: dispatch,
1664                                })
1665                                .await
1666                                .map_err(agent_client_protocol::util::internal_error)
1667                        }
1668                    },
1669                    agent_client_protocol::on_receive_dispatch!(),
1670                ),
1671            agent_component,
1672        )?;
1673        runner.successor = Arc::new(connection_to_agent);
1674
1675        runner.spawn_proxies(client_connection.clone(), proxy_components)?;
1676
1677        Ok(Dispatch::Request(
1678            modified_req.to_untyped_message()?,
1679            init_responder,
1680        ))
1681    }
1682
1683    #[cfg(feature = "unstable_protocol_v2")]
1684    async fn initialize_with_outcome(
1685        &self,
1686        message: Dispatch,
1687        client_connection: ConnectionTo<Client>,
1688        instantiator: Self::Instantiator,
1689        runner: &mut ConductorRunner<Self>,
1690    ) -> Result<InitializationOutcome, agent_client_protocol::Error> {
1691        initialize_agent_for_selected_protocol(message, client_connection, instantiator, runner)
1692            .await
1693    }
1694
1695    async fn handle_dispatch(
1696        &self,
1697        message: Dispatch,
1698        client_connection: ConnectionTo<Client>,
1699        conductor_tx: &mut mpsc::Sender<ConductorMessage>,
1700    ) -> Result<Handled<Dispatch>, agent_client_protocol::Error> {
1701        tracing::debug!(
1702            method = ?message.method(),
1703            "ConductorToClient::handle_dispatch"
1704        );
1705        MatchDispatchFrom::new(message, &client_connection)
1706            // Any incoming messages from the client are client-to-agent messages targeting the first component.
1707            .if_dispatch_from(Client, async move |message: Dispatch| {
1708                tracing::debug!(
1709                    method = ?message.method(),
1710                    "ConductorToClient::handle_dispatch - matched Client"
1711                );
1712                ConductorImpl::<Self>::incoming_message_from_client(conductor_tx, message).await
1713            })
1714            .await
1715            .done()
1716    }
1717}
1718
1719#[cfg(feature = "unstable_protocol_v2")]
1720async fn initialize_proxy_for_selected_protocol(
1721    message: Dispatch,
1722    client_connection: ConnectionTo<Conductor>,
1723    instantiator: Box<dyn InstantiateProxies>,
1724    runner: &mut ConductorRunner<Proxy>,
1725) -> Result<InitializationOutcome, Error> {
1726    let invalid_request = || Error::invalid_request().data("expected `initialize` request");
1727
1728    let Dispatch::Request(raw_request, init_responder) = message else {
1729        let error = invalid_request();
1730        if let Dispatch::Response(_, router) = message {
1731            router.route_with_error(error.clone())?;
1732        }
1733        return Ok(InitializationOutcome::Rejected(error));
1734    };
1735    if !InitializeProxyRequest::matches_method(raw_request.method()) {
1736        return reject_initialization(init_responder, invalid_request());
1737    }
1738
1739    let selection = match InitializeProtocol::from_request(&raw_request) {
1740        Ok(selection) => selection,
1741        Err(error) => return reject_initialization(init_responder, error),
1742    };
1743    let protocol = selection.protocol;
1744
1745    tracing::debug!(?protocol, "ensure_initialized: proxy initialize");
1746
1747    let initialization = match protocol {
1748        InitializeProtocol::V1 => {
1749            let InitializeProxyRequest { mut initialize } =
1750                match InitializeProxyRequest::parse_message(
1751                    raw_request.method(),
1752                    raw_request.params(),
1753                ) {
1754                    Ok(request) => request,
1755                    Err(error) => return reject_initialization(init_responder, error),
1756                };
1757            initialize.protocol_version = protocol.version();
1758            let original_request = initialize.clone();
1759            match instantiator.instantiate_proxies(initialize).await {
1760                Ok((mut request, proxies)) => {
1761                    request.protocol_version = protocol.version();
1762                    forwarded_initialize_request(
1763                        &raw_request,
1764                        selection,
1765                        &original_request,
1766                        request,
1767                    )
1768                    .map(|request| (request, proxies))
1769                }
1770                Err(error) => Err(error),
1771            }
1772        }
1773        InitializeProtocol::V2 => {
1774            let v2::InitializeProxyRequest { mut initialize } =
1775                match v2::InitializeProxyRequest::parse_message(
1776                    raw_request.method(),
1777                    raw_request.params(),
1778                ) {
1779                    Ok(request) => request,
1780                    Err(error) => return reject_initialization(init_responder, error),
1781                };
1782            initialize.protocol_version = protocol.version();
1783            let original_request = initialize.clone();
1784            match instantiator.instantiate_v2_proxies(initialize).await {
1785                Ok((mut request, proxies)) => {
1786                    request.protocol_version = protocol.version();
1787                    forwarded_initialize_request(
1788                        &raw_request,
1789                        selection,
1790                        &original_request,
1791                        request,
1792                    )
1793                    .map(|request| (request, proxies))
1794                }
1795                Err(error) => Err(error),
1796            }
1797        }
1798    };
1799    let (modified_req, proxy_components) = match initialization {
1800        Ok(initialization) => initialization,
1801        Err(error) => return reject_initialization(init_responder, error),
1802    };
1803
1804    runner.successor = Arc::new(GrandSuccessor);
1805    runner.spawn_proxies(client_connection, proxy_components)?;
1806
1807    Ok(InitializationOutcome::Forward(Dispatch::Request(
1808        modified_req,
1809        init_responder,
1810    )))
1811}
1812
1813/// Conductor acting as a proxy
1814impl ConductorHostRole for Proxy {
1815    type Instantiator = Box<dyn InstantiateProxies>;
1816
1817    async fn initialize(
1818        &self,
1819        message: Dispatch,
1820        client_connection: ConnectionTo<Conductor>,
1821        instantiator: Self::Instantiator,
1822        runner: &mut ConductorRunner<Self>,
1823    ) -> Result<Dispatch, agent_client_protocol::Error> {
1824        let invalid_request = || Error::invalid_request().data("expected `initialize` request");
1825
1826        let Dispatch::Request(request, init_responder) = message else {
1827            if let Dispatch::Response(_, router) = message {
1828                router.route_with_error(invalid_request())?;
1829            }
1830            return Err(invalid_request());
1831        };
1832        if !InitializeProxyRequest::matches_method(request.method()) {
1833            init_responder.respond_with_error(invalid_request())?;
1834            return Err(invalid_request());
1835        }
1836
1837        let InitializeProxyRequest { initialize } =
1838            match InitializeProxyRequest::parse_message(request.method(), request.params()) {
1839                Ok(request) => request,
1840                Err(error) => {
1841                    init_responder.respond_with_error(error)?;
1842                    return Err(invalid_request());
1843                }
1844            };
1845
1846        tracing::debug!("ensure_initialized: InitializeProxyRequest (proxy mode)");
1847
1848        let (modified_req, proxy_components) = instantiator.instantiate_proxies(initialize).await?;
1849
1850        runner.successor = Arc::new(GrandSuccessor);
1851        runner.spawn_proxies(client_connection.clone(), proxy_components)?;
1852
1853        Ok(Dispatch::Request(
1854            modified_req.to_untyped_message()?,
1855            init_responder,
1856        ))
1857    }
1858
1859    #[cfg(feature = "unstable_protocol_v2")]
1860    async fn initialize_with_outcome(
1861        &self,
1862        message: Dispatch,
1863        client_connection: ConnectionTo<Conductor>,
1864        instantiator: Self::Instantiator,
1865        runner: &mut ConductorRunner<Self>,
1866    ) -> Result<InitializationOutcome, agent_client_protocol::Error> {
1867        initialize_proxy_for_selected_protocol(message, client_connection, instantiator, runner)
1868            .await
1869    }
1870
1871    async fn handle_dispatch(
1872        &self,
1873        message: Dispatch,
1874        client_connection: ConnectionTo<Conductor>,
1875        conductor_tx: &mut mpsc::Sender<ConductorMessage>,
1876    ) -> Result<Handled<Dispatch>, agent_client_protocol::Error> {
1877        tracing::debug!(
1878            method = ?message.method(),
1879            ?message,
1880            "ConductorToConductor::handle_dispatch"
1881        );
1882        MatchDispatchFrom::new(message, &client_connection)
1883            .if_dispatch_from(Agent, {
1884                // Messages from our successor arrive already unwrapped
1885                // (RemoteRoleStyle::Successor strips the SuccessorMessage envelope).
1886                async |message: Dispatch| {
1887                    tracing::debug!(
1888                        method = ?message.method(),
1889                        "ConductorToConductor::handle_dispatch - matched Agent"
1890                    );
1891                    let mut conductor_tx = conductor_tx.clone();
1892                    ConductorImpl::<Self>::incoming_message_from_agent(&mut conductor_tx, message)
1893                        .await
1894                }
1895            })
1896            .await
1897            // Any incoming messages from the client are client-to-agent messages targeting the first component.
1898            .if_dispatch_from(Client, async |message: Dispatch| {
1899                tracing::debug!(
1900                    method = ?message.method(),
1901                    "ConductorToConductor::handle_dispatch - matched Client"
1902                );
1903                let mut conductor_tx = conductor_tx.clone();
1904                ConductorImpl::<Self>::incoming_message_from_client(&mut conductor_tx, message)
1905                    .await
1906            })
1907            .await
1908            .done()
1909    }
1910}
1911
1912pub trait ConductorSuccessor<Host: ConductorHostRole>: Send + Sync + 'static {
1913    /// Send a message to the successor.
1914    fn send_message<'a>(
1915        &self,
1916        message: Dispatch,
1917        connection_to_conductor: ConnectionTo<Host::Counterpart>,
1918        runner: &'a mut ConductorRunner<Host>,
1919    ) -> BoxFuture<'a, Result<(), agent_client_protocol::Error>>;
1920}
1921
1922impl<Host: ConductorHostRole> ConductorSuccessor<Host> for agent_client_protocol::Error {
1923    fn send_message<'a>(
1924        &self,
1925        _message: Dispatch,
1926        _connection_to_conductor: ConnectionTo<Host::Counterpart>,
1927        _runner: &'a mut ConductorRunner<Host>,
1928    ) -> BoxFuture<'a, Result<(), agent_client_protocol::Error>> {
1929        let error = self.clone();
1930        Box::pin(std::future::ready(Err(error)))
1931    }
1932}
1933
1934/// A dummy type handling messages sent to the conductor's
1935/// successor when it is acting as a proxy.
1936struct GrandSuccessor;
1937
1938/// When the conductor is acting as an proxy, messages sent by
1939/// the last proxy go to the conductor's successor.
1940///
1941/// ```text
1942/// client --> Conductor -----------------------------> GrandSuccessor
1943///            |                                  |
1944///            +-> Proxy[0] -> ... -> Proxy[n-1] -+
1945/// ```
1946impl ConductorSuccessor<Proxy> for GrandSuccessor {
1947    fn send_message<'a>(
1948        &self,
1949        message: Dispatch,
1950        connection: ConnectionTo<Conductor>,
1951        _runner: &'a mut ConductorRunner<Proxy>,
1952    ) -> BoxFuture<'a, Result<(), agent_client_protocol::Error>> {
1953        Box::pin(async move {
1954            debug!("Proxy mode: forwarding successor message to conductor's successor");
1955            connection.send_proxied_message_to(Agent, message)
1956        })
1957    }
1958}
1959
1960/// When the conductor is acting as an agent, messages sent by
1961/// the last proxy to its successor go to the internal agent
1962/// (`self`).
1963impl ConductorSuccessor<Agent> for ConnectionTo<Agent> {
1964    fn send_message<'a>(
1965        &self,
1966        message: Dispatch,
1967        connection: ConnectionTo<Client>,
1968        runner: &'a mut ConductorRunner<Agent>,
1969    ) -> BoxFuture<'a, Result<(), agent_client_protocol::Error>> {
1970        let connection_to_agent = self.clone();
1971        Box::pin(async move {
1972            debug!("Proxy mode: forwarding successor message to conductor's successor");
1973            runner
1974                .forward_message_to_agent(connection, message, connection_to_agent)
1975                .await
1976        })
1977    }
1978}