simulator_client/managed/

parallel.rs

//! ParallelControlManager — owns the multiplexed parallel-control WebSocket.
//!
//! A single control connection drives N server-created sub-sessions. The server
//! splits the requested range, streams one `SessionCreated` per sub-session plus
//! a terminal `SessionsCreatedV2`, then multiplexes per-session control events as
//! `SessionEventV2 { session_id, seq_id, event }`. This task demultiplexes those
//! into one [`ControlEvent`] channel per sub-session and routes each driver's
//! `Continue` back as `ContinueSessionV1`. On reconnect it re-attaches with
//! `AttachParallelControlSessionV2` and the per-session sequence cursors, dropping
//! any replayed events the driver already saw.
//!
//! Transactions are NOT multiplexed: each sub-session exposes its own
//! `rpc_endpoint`, so [`ParallelSubSession`] spawns its own data-plane
//! subscription managers exactly like [`super::ManagedBacktestSession`].

use std::{
    collections::{HashMap, VecDeque},
    sync::Arc,
};

use futures::StreamExt;
use simulator_api::{
    BacktestRequest, BacktestResponse, ContinueParams, ContinueSessionRequestV1,
    CreateBacktestSessionRequest, SessionEventKind,
};
use tokio::{
    sync::{mpsc, oneshot, watch},
    task::JoinHandle,
};
use tokio_tungstenite::tungstenite::Message;
use tokio_util::sync::CancellationToken;
use tracing::{debug, info, warn};

use super::{
    ConnectionStatus, ControlConnection, ControlEvent, HANDSHAKE_RESPONSE_TIMEOUT, HandshakeError,
    InboundFrame, KEEPALIVE_INTERVAL, ManagedEvent, ManagedSessionError, MessageLoopExit,
    ReconnectCoordinator, SessionInfo, SubscriptionHandle, Ws, classify_inbound, graceful_close,
    handshake_error_for_response, is_terminal_backtest_error, resolve_rpc_url, run_control_loop,
    send_keepalive_ping, send_request,
    session::{
        DrainOutcome, drain_subscriptions_until_complete, try_next_subscription_event,
        wait_any_subscription_event, wait_connections_up,
    },
    spawn_account_diff_subscription_manager, spawn_transaction_subscription_manager,
};
use crate::{error::err_chain, urls::http_base_from_ws_url};

/// Per-read timeout while streaming the create response. Generous because the
/// server may park between `SessionCreated`s waiting for sub-session capacity.
const CREATE_RESPONSE_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(900);
/// Capacity of the shared `Continue` request channel feeding the control task.
/// One in-flight `Continue` per sub-session is the steady state.
const CONTINUE_CHANNEL_CAPACITY: usize = 256;
/// Liveness backstop for a sub-session's completion drain (see
/// [`super::ManagedBacktestSession`]'s equivalent).
const COMPLETION_DRAIN_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(60);

/// A `Continue` request tagged with the sub-session it targets.
type TaggedContinue = (String, ContinueParams);

/// Result of the parallel-create handshake handed back to the caller.
struct ParallelCreated {
    control_session_id: String,
    sessions: Vec<CreatedSubSession>,
}

struct CreatedSubSession {
    info: SessionInfo,
    events: mpsc::UnboundedReceiver<ControlEvent>,
    /// Server-authoritative sub-session range, back-filled from
    /// `SessionsCreatedV2` before the create handshake returns (a server that
    /// omits the ranges is rejected as too old, so this is never left at `0`).
    start_slot: u64,
    end_slot: u64,
}

/// Handle to a running [`ParallelControlTask`].
struct ParallelControlHandle {
    continues: mpsc::Sender<TaggedContinue>,
    status: watch::Receiver<ConnectionStatus>,
    created: Option<oneshot::Receiver<Result<ParallelCreated, String>>>,
    join: JoinHandle<()>,
}

impl ParallelControlHandle {
    /// Resolve once the parallel create handshake has produced the sub-session
    /// set (or failed). Consumes the one-shot; callable only once.
    async fn wait_created(&mut self) -> Result<ParallelCreated, String> {
        let rx = self
            .created
            .take()
            .ok_or_else(|| "parallel create already consumed".to_string())?;
        rx.await
            .map_err(|_| "control manager exited before creating sessions".to_string())?
    }

    /// Await the control task's exit; dropping the request channels nudges it to
    /// stop if the caller is giving up.
    async fn join(self) {
        drop(self.continues);
        let _ = self.join.await;
    }
}

fn spawn_parallel_control_manager(
    url: String,
    api_key: String,
    create: CreateBacktestSessionRequest,
    cancel: CancellationToken,
) -> ParallelControlHandle {
    let (continues_tx, continues_rx) = mpsc::channel::<TaggedContinue>(CONTINUE_CHANNEL_CAPACITY);
    let (status_tx, status_rx) = watch::channel(ConnectionStatus::Down);
    let (created_tx, created_rx) = oneshot::channel::<Result<ParallelCreated, String>>();

    let task = ParallelControlTask {
        url,
        api_key,
        create: Some(create),
        control_session_id: None,
        event_txs: HashMap::new(),
        last_sequences: HashMap::new(),
        completed: std::collections::HashSet::new(),
        continues_rx,
        status_tx,
        created_tx: Some(created_tx),
        cancel,
    };

    let join = tokio::spawn(run_control_loop(task));

    ParallelControlHandle {
        continues: continues_tx,
        status: status_rx,
        created: Some(created_rx),
        join,
    }
}

struct ParallelControlTask {
    url: String,
    api_key: String,
    /// Set on first connect; consumed and cleared after the create handshake.
    create: Option<CreateBacktestSessionRequest>,
    /// Populated after create; used to build `AttachParallelControlSessionV2`.
    control_session_id: Option<String>,
    /// Per-sub-session demux senders, keyed by `session_id`. Unbounded so a
    /// slow driver never head-of-line blocks the shared control connection.
    event_txs: HashMap<String, mpsc::UnboundedSender<ControlEvent>>,
    /// Highest `seq_id` delivered per sub-session — the reconnect cursor and the
    /// replay-dedup key.
    last_sequences: HashMap<String, u64>,
    /// Sub-sessions that have reported `Completed`. Once all have, the control
    /// stream is done and the task closes gracefully.
    completed: std::collections::HashSet<String>,
    continues_rx: mpsc::Receiver<TaggedContinue>,
    status_tx: watch::Sender<ConnectionStatus>,
    created_tx: Option<oneshot::Sender<Result<ParallelCreated, String>>>,
    cancel: CancellationToken,
}

impl ControlConnection for ParallelControlTask {
    fn url(&self) -> &str {
        &self.url
    }
    fn api_key(&self) -> &str {
        &self.api_key
    }
    fn cancel(&self) -> &CancellationToken {
        &self.cancel
    }
    fn label(&self) -> &'static str {
        "parallel control"
    }
    fn status_tx(&self) -> &watch::Sender<ConnectionStatus> {
        &self.status_tx
    }

    fn fail_pending(&mut self, reason: String) {
        if let Some(tx) = self.created_tx.take() {
            let _ = tx.send(Err(reason));
        }
    }

    async fn handshake(&mut self, ws: Ws) -> Result<Ws, HandshakeError> {
        if let Some(control_session_id) = self.control_session_id.clone() {
            self.attach(ws, &control_session_id).await
        } else if let Some(create) = self.create.clone() {
            // Clone rather than take: a transient drop mid-create must leave the
            // request intact so the reconnect can recreate. `create_sessions`
            // clears it (and sets `control_session_id`) only on success, after
            // which reconnects route to `attach` above.
            self.create_sessions(ws, create).await
        } else {
            Err(HandshakeError::Fatal(
                "no create request and no control_session_id".into(),
            ))
        }
    }

    async fn message_loop(&mut self, mut ws: Ws) -> MessageLoopExit {
        let mut ping_timer = tokio::time::interval(KEEPALIVE_INTERVAL);
        ping_timer.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
        let mut last_inbound = std::time::Instant::now();

        let exit = loop {
            tokio::select! {
                biased;
                _ = self.cancel.cancelled() => break MessageLoopExit::Cancelled,

                _ = ping_timer.tick() => {
                    if let Some(why) = send_keepalive_ping(&mut ws, last_inbound).await {
                        break MessageLoopExit::ConnectionLost(why);
                    }
                }

                msg = ws.next() => {
                    last_inbound = std::time::Instant::now();
                    match classify_inbound(msg) {
                        InboundFrame::Text(t) => {
                            if let Some(exit) = self.handle_text(&t) {
                                break exit;
                            }
                        }
                        InboundFrame::Ignore => {}
                        InboundFrame::Lost(why) => break MessageLoopExit::ConnectionLost(why),
                    }
                }

                req = self.continues_rx.recv() => {
                    match req {
                        Some((session_id, request)) => {
                            let msg = BacktestRequest::ContinueSessionV1(ContinueSessionRequestV1 { session_id, request });
                            if let Err(e) = send_request(&mut ws, &msg).await {
                                break MessageLoopExit::ConnectionLost(format!("continue send: {e}"));
                            }
                        }
                        None => break MessageLoopExit::SessionEnded,
                    }
                }
            }
        };

        if matches!(
            exit,
            MessageLoopExit::SessionEnded | MessageLoopExit::Cancelled
        ) {
            graceful_close(&mut ws).await;
        }
        exit
    }
}

impl ParallelControlTask {
    /// First-connect handshake: send the parallel create, collect each streamed
    /// `SessionCreated`, and finish on `SessionsCreatedV2`. Fires `created_tx`
    /// with the sub-session set so the caller can begin driving.
    async fn create_sessions(
        &mut self,
        mut ws: Ws,
        create: CreateBacktestSessionRequest,
    ) -> Result<Ws, HandshakeError> {
        send_request(&mut ws, &BacktestRequest::CreateBacktestSession(create))
            .await
            .map_err(HandshakeError::Transient)?;

        let rpc_base = http_base_from_ws_url(&self.url);
        let mut sessions: Vec<CreatedSubSession> = Vec::new();

        loop {
            let response = next_response(&mut ws, CREATE_RESPONSE_TIMEOUT)
                .await
                .map_err(HandshakeError::Transient)?;
            match response {
                BacktestResponse::SessionCreated {
                    session_id,
                    rpc_endpoint,
                    task_id,
                } => {
                    let (event_tx, event_rx) = mpsc::unbounded_channel::<ControlEvent>();
                    self.event_txs.insert(session_id.clone(), event_tx);
                    sessions.push(CreatedSubSession {
                        info: SessionInfo {
                            rpc_endpoint: resolve_rpc_url(&rpc_base, &rpc_endpoint),
                            session_id,
                            task_id,
                        },
                        events: event_rx,
                        start_slot: 0,
                        end_slot: 0,
                    });
                }
                BacktestResponse::SessionEventV2 {
                    session_id,
                    seq_id,
                    event,
                } => {
                    self.route_event(&session_id, seq_id, event);
                }
                BacktestResponse::SessionsCreatedV2 {
                    control_session_id,
                    session_ids,
                    start_slots,
                    end_slots,
                    ..
                } => {
                    info!(
                        %control_session_id,
                        sessions = sessions.len(),
                        "parallel sessions created"
                    );
                    // Bind each sub-session to its server-authoritative range
                    // (parallel arrays keyed by session_id), so the driver never
                    // re-derives the split or guesses the range from the first slot.
                    // Absent or mismatched-length arrays mean the server is too old
                    // to report ranges; fail loudly rather than silently leaving
                    // sub-sessions at `(0, 0)` (which degrades to advance-by-one).
                    if session_ids.len() != start_slots.len()
                        || session_ids.len() != end_slots.len()
                    {
                        return Err(HandshakeError::Fatal(format!(
                            "server did not report per-sub-session ranges \
                             (session_ids={}, start_slots={}, end_slots={}); \
                             server is too old for the multiplexed parallel client",
                            session_ids.len(),
                            start_slots.len(),
                            end_slots.len(),
                        )));
                    }
                    for ((id, start), end) in session_ids.iter().zip(&start_slots).zip(&end_slots) {
                        if let Some(s) = sessions.iter_mut().find(|s| s.info.session_id == *id) {
                            s.start_slot = *start;
                            s.end_slot = *end;
                        }
                    }
                    self.control_session_id = Some(control_session_id.clone());
                    // Create succeeded: drop the request so reconnects re-attach
                    // instead of recreating.
                    self.create = None;
                    if let Some(tx) = self.created_tx.take() {
                        let _ = tx.send(Ok(ParallelCreated {
                            control_session_id,
                            sessions,
                        }));
                    }
                    return Ok(ws);
                }
                BacktestResponse::Error(err) => {
                    return Err(HandshakeError::Fatal(format!(
                        "server error: {}",
                        err_chain(&err)
                    )));
                }
                _ => {}
            }
        }
    }

    /// Reconnect handshake: re-attach with the per-session sequence cursors and
    /// wait for `ParallelSessionAttachedV2`. Replayed events arrive in the
    /// message loop and are deduplicated by `last_sequences`.
    async fn attach(&mut self, mut ws: Ws, control_session_id: &str) -> Result<Ws, HandshakeError> {
        send_request(
            &mut ws,
            &BacktestRequest::AttachParallelControlSessionV2 {
                control_session_id: control_session_id.to_string(),
                last_sequences: self.last_sequences.clone().into_iter().collect(),
            },
        )
        .await
        .map_err(HandshakeError::Transient)?;

        loop {
            let response = next_response(&mut ws, HANDSHAKE_RESPONSE_TIMEOUT)
                .await
                .map_err(HandshakeError::Transient)?;
            match response {
                BacktestResponse::ParallelSessionAttachedV2 { .. } => {
                    debug!(%control_session_id, "parallel control reattached");
                    return Ok(ws);
                }
                BacktestResponse::SessionEventV2 {
                    session_id,
                    seq_id,
                    event,
                } => {
                    self.route_event(&session_id, seq_id, event);
                }
                BacktestResponse::Error(err) => {
                    return Err(handshake_error_for_response("attach", err));
                }
                _ => {}
            }
        }
    }

    /// Parse one inbound control frame. Returns `Some(exit)` when the loop should
    /// stop (all sub-sessions completed, or a fatal control-session error).
    fn handle_text(&mut self, text: &str) -> Option<MessageLoopExit> {
        let response = match serde_json::from_str::<BacktestResponse>(text) {
            Ok(r) => r,
            Err(e) => {
                warn!(error = %err_chain(&e), "discarding undeserializable parallel control message");
                return None;
            }
        };

        match response {
            BacktestResponse::SessionEventV2 {
                session_id,
                seq_id,
                event,
            } => {
                self.route_event(&session_id, seq_id, event);
                if self.completed.len() == self.event_txs.len() && !self.event_txs.is_empty() {
                    return Some(MessageLoopExit::SessionEnded);
                }
            }
            BacktestResponse::Error(err) => {
                if is_terminal_backtest_error(&err) {
                    return Some(MessageLoopExit::Terminal(format!(
                        "control session error: {}",
                        err_chain(&err)
                    )));
                }
                warn!(error = %err_chain(&err), "non-terminal parallel control error");
            }
            other => {
                debug!(?other, "ignoring unexpected parallel control response");
            }
        }
        None
    }

    /// Route one multiplexed event to its sub-session, deduplicating replays by
    /// per-session sequence and tracking completion.
    fn route_event(&mut self, session_id: &str, seq_id: u64, event: SessionEventKind) {
        // Advance the per-session sequence cursor, dropping replays. Update in
        // place when the session is already known so the steady-state path
        // doesn't re-allocate the key on every event.
        match self.last_sequences.get_mut(session_id) {
            Some(last) if seq_id <= *last => return,
            Some(last) => *last = seq_id,
            None => {
                self.last_sequences.insert(session_id.to_string(), seq_id);
            }
        }

        let is_completed = matches!(event, SessionEventKind::Completed { .. });
        let Some(control_event) = session_event_to_control(event) else {
            return;
        };
        if let Some(tx) = self.event_txs.get(session_id) {
            // Ignore send error if the driver has finished and dropped its receiver.
            let _ = tx.send(control_event);
        }
        if is_completed {
            self.completed.insert(session_id.to_string());
        }
    }
}

/// Map a multiplexed `SessionEventKind` onto the driver-facing [`ControlEvent`].
/// `Success` carries no driver signal and is dropped (`None`).
fn session_event_to_control(event: SessionEventKind) -> Option<ControlEvent> {
    Some(match event {
        SessionEventKind::ReadyForContinue => ControlEvent::ReadyForContinue,
        SessionEventKind::SlotNotification(slot) => ControlEvent::Slot(slot),
        SessionEventKind::Paused(event) => ControlEvent::Paused(event),
        SessionEventKind::DiscoveryBatch(event) => ControlEvent::DiscoveryBatch(event),
        SessionEventKind::Error(error) => ControlEvent::Error(error),
        SessionEventKind::Completed { summary } => ControlEvent::Completed {
            summary,
            agent_stats: None,
        },
        SessionEventKind::Status { status } => ControlEvent::Status(status),
        SessionEventKind::Success => return None,
    })
}

/// Read the next control frame with a bounded timeout, skipping non-text frames.
async fn next_response(
    ws: &mut Ws,
    timeout: std::time::Duration,
) -> Result<BacktestResponse, String> {
    let deadline = tokio::time::Instant::now() + timeout;
    loop {
        let msg = tokio::time::timeout_at(deadline, ws.next())
            .await
            .map_err(|_| format!("handshake timeout after {timeout:?}"))?;

        let Some(msg) = msg else {
            return Err("ws ended during handshake".into());
        };
        let msg = msg.map_err(|e| format!("ws read: {}", err_chain(&e)))?;

        let text = match msg {
            Message::Text(t) => t,
            Message::Binary(b) => match String::from_utf8(b) {
                Ok(t) => t,
                Err(_) => continue,
            },
            Message::Close(frame) => {
                return Err(format!("remote close during handshake: {frame:?}"));
            }
            _ => continue,
        };

        return serde_json::from_str::<BacktestResponse>(&text)
            .map_err(|e| format!("deserialize: {}; raw={text}", err_chain(&e)));
    }
}

/// High-level managed parallel session: owns the single multiplexed control
/// connection and hands out one [`ParallelSubSession`] per server-created
/// sub-session for the driver to run concurrently.
pub struct ManagedParallelSession {
    control_session_id: String,
    control: Option<ParallelControlHandle>,
    sub_sessions: Vec<ParallelSubSession>,
    session_cancel: CancellationToken,
}

impl ManagedParallelSession {
    /// Start a parallel session tied to a caller-owned cancellation token.
    pub async fn start_with_cancel(
        url: String,
        api_key: String,
        create: CreateBacktestSessionRequest,
        parent_cancel: CancellationToken,
    ) -> Result<Self, ManagedSessionError> {
        let session_cancel = parent_cancel.child_token();
        let mut control =
            spawn_parallel_control_manager(url, api_key, create, session_cancel.clone());

        let created = tokio::select! {
            biased;
            _ = parent_cancel.cancelled() => {
                session_cancel.cancel();
                control.join().await;
                return Err(ManagedSessionError::Cancelled);
            }
            result = control.wait_created() => {
                result.map_err(ManagedSessionError::Create)?
            }
        };

        // One coordinator shared across the sub-sessions: their data planes share
        // the client's link, so a dropped subscription parks until its streaming
        // siblings finish before reconnecting into the freed bandwidth.
        let reconnect_coordinator = Arc::new(ReconnectCoordinator::new());
        let sub_sessions = created
            .sessions
            .into_iter()
            .map(|s| ParallelSubSession {
                session_info: s.info,
                events: s.events,
                continues: control.continues.clone(),
                status: control.status.clone(),
                subscriptions: Vec::new(),
                session_cancel: session_cancel.child_token(),
                post_completion: None,
                post_completion_error: None,
                reconnect_coordinator: Some(reconnect_coordinator.clone()),
                start_slot: s.start_slot,
                end_slot: s.end_slot,
            })
            .collect();

        Ok(Self {
            control_session_id: created.control_session_id,
            control: Some(control),
            sub_sessions,
            session_cancel,
        })
    }

    /// The server-assigned id grouping these sub-sessions (used for reconnect).
    pub fn control_session_id(&self) -> &str {
        &self.control_session_id
    }

    /// Take the sub-sessions for driving. Returns them once; subsequent calls
    /// yield an empty vec.
    pub fn take_sub_sessions(&mut self) -> Vec<ParallelSubSession> {
        std::mem::take(&mut self.sub_sessions)
    }

    /// Cancel the session and join the control task.
    pub async fn shutdown(mut self) {
        self.session_cancel.cancel();
        if let Some(control) = self.control.take() {
            control.join().await;
        }
    }
}

impl Drop for ManagedParallelSession {
    fn drop(&mut self) {
        self.session_cancel.cancel();
    }
}

/// One sub-session of a parallel run. Mirrors the driver-facing surface of
/// [`super::ManagedBacktestSession`]: control events arrive over a demuxed
/// channel from the shared control connection, while transactions flow over this
/// sub-session's own data-plane subscription managers.
pub struct ParallelSubSession {
    session_info: SessionInfo,
    events: mpsc::UnboundedReceiver<ControlEvent>,
    continues: mpsc::Sender<TaggedContinue>,
    status: watch::Receiver<ConnectionStatus>,
    subscriptions: Vec<SubscriptionHandle>,
    session_cancel: CancellationToken,
    /// Trailing notifications drained on `Completed`, followed by `Completed`;
    /// served in order by `next_event`. `None` until completion.
    post_completion: Option<VecDeque<ManagedEvent>>,
    /// Surfaced after `post_completion` drains: set when the completion drain
    /// stalled, so a silently-truncated sub-session fails loudly.
    post_completion_error: Option<ManagedSessionError>,
    /// Coordinator shared across the parallel batch's sub-sessions so their
    /// subscription reconnects step aside for still-streaming siblings.
    reconnect_coordinator: Option<Arc<ReconnectCoordinator>>,
    /// Server-authoritative slot range for this sub-session (from
    /// `SessionsCreatedV2`).
    start_slot: u64,
    end_slot: u64,
}

impl ParallelSubSession {
    pub fn session_info(&self) -> &SessionInfo {
        &self.session_info
    }

    /// Server-authoritative `(start_slot, end_slot)` for this sub-session.
    pub fn range(&self) -> (u64, u64) {
        (self.start_slot, self.end_slot)
    }

    pub fn subscribe_transactions(&mut self, program_ids: Vec<String>) {
        self.subscriptions
            .push(spawn_transaction_subscription_manager(
                self.session_info.rpc_endpoint.clone(),
                program_ids,
                self.session_cancel.clone(),
                self.reconnect_coordinator.clone(),
            ));
    }

    pub fn subscribe_account_diffs(&mut self, program_ids: Vec<String>) {
        self.subscriptions
            .push(spawn_account_diff_subscription_manager(
                self.session_info.rpc_endpoint.clone(),
                program_ids,
                self.session_cancel.clone(),
                self.reconnect_coordinator.clone(),
            ));
    }

    /// Receive the next control or subscription event for this sub-session. On
    /// `Completed`, trailing subscription notifications are drained and delivered
    /// before the `Completed` event, mirroring [`super::ManagedBacktestSession`].
    pub async fn next_event(&mut self) -> Result<ManagedEvent, ManagedSessionError> {
        // Serve buffered post-completion events (trailing notifications, then
        // `Completed`); the control stream is gone once they're exhausted.
        if let Some(buffered) = self.post_completion.as_mut() {
            if let Some(event) = buffered.pop_front() {
                return Ok(event);
            }
            // Buffer drained: surface a pending stall error so an incomplete
            // run fails loudly, else the control stream is simply done.
            return Err(self
                .post_completion_error
                .take()
                .unwrap_or(ManagedSessionError::ControlClosed));
        }

        if let Some(event) = try_next_subscription_event(&mut self.subscriptions) {
            return Ok(event);
        }

        // Scope the borrows to the `select!` so the completion drain can
        // re-borrow `self`.
        let event = {
            let cancel = &self.session_cancel;
            let subscriptions = &mut self.subscriptions;
            tokio::select! {
                biased;
                _ = cancel.cancelled() => return Err(ManagedSessionError::Cancelled),
                event = self.events.recv() => {
                    event.map(ManagedEvent::from).ok_or(ManagedSessionError::ControlClosed)?
                }
                event = wait_any_subscription_event(subscriptions) => event,
            }
        };

        // Bind the payload so the re-emitted `Completed` below carries it.
        let ManagedEvent::Completed {
            summary,
            agent_stats,
        } = event
        else {
            return Ok(event);
        };

        // Flush trailing notifications up to each subscription's terminal,
        // delivering them before `Completed`, mirroring
        // [`super::ManagedBacktestSession`].
        let (mut buffered, terminal): (VecDeque<ManagedEvent>, _) = match self
            .drain_until_subscriptions_complete(COMPLETION_DRAIN_TIMEOUT)
            .await
        {
            DrainOutcome::Complete(events) => (
                events.into(),
                Ok(ManagedEvent::Completed {
                    summary,
                    agent_stats,
                }),
            ),
            // The data plane stalled before every subscription finished:
            // trailing notifications are missing, so report failure rather
            // than a silently-truncated `Completed`.
            DrainOutcome::Stalled(events) => (
                events.into(),
                Err(ManagedSessionError::SubscriptionFailed(
                    "completion drain stalled: subscriptions did not deliver their \
                     end-of-stream terminals; the captured stream is incomplete"
                        .to_string(),
                )),
            ),
        };
        match terminal {
            Ok(completed) => buffered.push_back(completed),
            Err(err) => self.post_completion_error = Some(err),
        }
        let first = buffered.pop_front();
        self.post_completion = Some(buffered);
        match first {
            Some(event) => Ok(event),
            None => Err(self
                .post_completion_error
                .take()
                .unwrap_or(ManagedSessionError::ControlClosed)),
        }
    }

    /// Wait until the shared control connection and this sub-session's
    /// subscriptions are up, then route a `Continue` as `ContinueSessionV1`.
    pub async fn send_continue(
        &mut self,
        params: ContinueParams,
    ) -> Result<(), ManagedSessionError> {
        self.wait_all_up().await?;
        self.continues
            .send((self.session_info.session_id.clone(), params))
            .await
            .map_err(|e| ManagedSessionError::ContinueSend(e.to_string()))
    }

    /// See [`drain_subscriptions_until_complete`].
    async fn drain_until_subscriptions_complete(
        &mut self,
        idle_timeout: std::time::Duration,
    ) -> DrainOutcome {
        drain_subscriptions_until_complete(
            &mut self.subscriptions,
            &self.session_cancel,
            idle_timeout,
        )
        .await
    }

    /// Cancel this sub-session's subscriptions and join them. The shared control
    /// connection is owned by [`ManagedParallelSession`] and is not touched here.
    pub async fn shutdown(mut self) {
        self.session_cancel.cancel();
        for sub in std::mem::take(&mut self.subscriptions) {
            let _ = sub.join.await;
        }
    }

    async fn wait_all_up(&self) -> Result<(), ManagedSessionError> {
        let subscriptions = self
            .subscriptions
            .iter()
            .map(|s| s.status.clone())
            .collect();
        wait_connections_up(self.status.clone(), subscriptions, &self.session_cancel).await
    }
}

impl Drop for ParallelSubSession {
    fn drop(&mut self) {
        self.session_cancel.cancel();
    }
}