simulator_client/managed/

control.rs

//! ControlManager — owns the backtest control WebSocket.
//!
//! Responsibilities:
//! - Establish the WS with a bounded connect timeout
//! - Perform the correct handshake (`Create` the first time, `Attach`+`Resume`
//!   on reconnect) with a bounded per-response timeout
//! - Bridge inbound `BacktestResponse` ↔ outbound `Continue` via channels
//! - Keep the connection alive with WebSocket ping/pong
//! - Reconnect with bounded backoff; publish `ConnectionStatus` transitions
//! - On total-budget exhaustion, publish `Failed(reason)` and exit

use std::time::Instant;

use futures::StreamExt;
use simulator_api::{
    AgentStatsReport, BacktestError, BacktestRequest, BacktestResponse, BacktestStatus,
    ContinueParams, ContinueToParams, CreateBacktestSessionRequest, DiscoveryBatchEvent,
    PausedEvent, SequencedResponse, SessionSummary,
};
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, HANDSHAKE_RESPONSE_TIMEOUT, HandshakeError, InboundFrame,
    KEEPALIVE_INTERVAL, MessageLoopExit, SessionInfo, Ws, classify_inbound, graceful_close,
    handshake_error_for_response, is_terminal_backtest_error, resolve_rpc_url, run_control_loop,
    send_keepalive_ping, send_request,
};
use crate::{error::err_chain, urls::http_base_from_ws_url};

/// Events the driver observes from the control connection.
///
/// Session-lifecycle responses (`SessionCreated`, `SessionAttached`,
/// `ResumeSuccess`) are handled internally and not forwarded.
#[derive(Debug)]
pub enum ControlEvent {
    ReadyForContinue,
    /// Server paused at a `ContinueTo` target. The session is ready for
    /// another `Continue` or `ContinueTo` from this point.
    Paused(PausedEvent),
    /// Server discovered an upcoming batch matching a registered
    /// `DiscoveryFilter`. Send `ContinueTo(slot, batch_index)` to pause
    /// immediately before it executes.
    DiscoveryBatch(DiscoveryBatchEvent),
    Slot(u64),
    /// High-level progress phase during session startup (e.g. `StartingRuntime`).
    /// Useful for showing what the server is doing while waiting for the first
    /// `ReadyForContinue`.
    Status(BacktestStatus),
    Completed {
        summary: Option<SessionSummary>,
        agent_stats: Option<Vec<AgentStatsReport>>,
    },
    Error(BacktestError),
}

/// Handle to a running `ControlManager` task.
pub struct ControlHandle {
    continues: mpsc::Sender<ContinueParams>,
    continue_tos: mpsc::Sender<ContinueToParams>,
    pub events: mpsc::Receiver<ControlEvent>,
    pub status: watch::Receiver<ConnectionStatus>,
    session_info: Option<oneshot::Receiver<Result<SessionInfo, String>>>,
    join: JoinHandle<()>,
}

impl ControlHandle {
    /// Resolve once the session has been created (or the manager has failed
    /// before reaching that point). Consumes the one-shot; callable only once.
    pub async fn wait_for_session(&mut self) -> Result<SessionInfo, String> {
        let rx = self
            .session_info
            .take()
            .ok_or_else(|| "session_info already consumed".to_string())?;
        rx.await
            .map_err(|_| "control manager exited before creating session".to_string())?
    }

    /// Send a `Continue` request to the control task. Errors if the manager has
    /// exited.
    pub async fn send_continue(
        &self,
        params: ContinueParams,
    ) -> Result<(), mpsc::error::SendError<ContinueParams>> {
        self.continues.send(params).await
    }

    /// Send a `ContinueTo` request to step to a specific slot/batch boundary.
    /// Pair with `ControlEvent::DiscoveryBatch` to pause before each
    /// discovered batch.
    pub async fn send_continue_to(
        &self,
        params: ContinueToParams,
    ) -> Result<(), mpsc::error::SendError<ContinueToParams>> {
        self.continue_tos.send(params).await
    }

    /// Await the control task's exit. The task exits on its own when the
    /// server reports `Completed`, the cancel token fires, or it hits a
    /// terminal error; dropping the request channels here nudges it in the
    /// case where the driver is giving up without having seen `Completed`.
    pub async fn join(self) {
        drop(self.continues);
        drop(self.continue_tos);
        let _ = self.join.await;
    }
}

/// Spawn a `ControlManager` task and return a handle.
///
/// The `continues` channel has a bounded capacity of 1: we only ever have one
/// Continue in flight, and backpressuring the driver is the correct behavior
/// if the connection is temporarily down.
pub fn spawn_control_manager(
    url: String,
    api_key: String,
    create: CreateBacktestSessionRequest,
    cancel: CancellationToken,
) -> ControlHandle {
    let (continues_tx, continues_rx) = mpsc::channel::<ContinueParams>(1);
    let (continue_tos_tx, continue_tos_rx) = mpsc::channel::<ContinueToParams>(1);
    let (events_tx, events_rx) = mpsc::channel::<ControlEvent>(256);
    let (status_tx, status_rx) = watch::channel(ConnectionStatus::Down);
    let (session_tx, session_rx) = oneshot::channel::<Result<SessionInfo, String>>();

    let manager = ControlTask {
        url,
        api_key,
        create: Some(create),
        session_info: None,
        session_tx: Some(session_tx),
        last_sequence: None,
        continues_rx,
        continue_tos_rx,
        events_tx,
        status_tx,
        cancel,
    };

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

    ControlHandle {
        continues: continues_tx,
        continue_tos: continue_tos_tx,
        events: events_rx,
        status: status_rx,
        session_info: Some(session_rx),
        join,
    }
}

struct ControlTask {
    url: String,
    api_key: String,
    /// Set on first connect; consumed and cleared after `Create` succeeds.
    create: Option<CreateBacktestSessionRequest>,
    /// Populated after `Create` succeeds. On reconnect, used to build `Attach`.
    session_info: Option<SessionInfo>,
    /// One-shot result to the handle; fired exactly once.
    session_tx: Option<oneshot::Sender<Result<SessionInfo, String>>>,
    /// Highest sequence number observed from the server.
    last_sequence: Option<u64>,
    continues_rx: mpsc::Receiver<ContinueParams>,
    continue_tos_rx: mpsc::Receiver<ContinueToParams>,
    events_tx: mpsc::Sender<ControlEvent>,
    status_tx: watch::Sender<ConnectionStatus>,
    cancel: CancellationToken,
}

impl ControlConnection for ControlTask {
    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 {
        "control"
    }
    fn status_tx(&self) -> &watch::Sender<ConnectionStatus> {
        &self.status_tx
    }

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

    async fn handshake(&mut self, mut ws: Ws) -> Result<Ws, HandshakeError> {
        if let Some(info) = &self.session_info {
            let info = info.clone();
            attach(
                &mut ws,
                &info.session_id,
                self.last_sequence,
                &mut self.events_tx,
                &mut self.last_sequence,
            )
            .await?;
            resume(&mut ws, &mut self.events_tx, &mut self.last_sequence).await?;
            debug!(session_id = info.session_id, "control reattached");
        } else if let Some(create) = self.create.take() {
            let info = create_session(
                &mut ws,
                create,
                &self.url,
                &mut self.events_tx,
                &mut self.last_sequence,
            )
            .await?;
            info!(session_id = info.session_id, "control session created");
            self.session_info = Some(info.clone());
            if let Some(tx) = self.session_tx.take() {
                let _ = tx.send(Ok(info));
            }
        } else {
            return Err(HandshakeError::Fatal(
                "no create request and no session_id".into(),
            ));
        }

        Ok(ws)
    }

    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 = 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 = Instant::now();
                    match classify_inbound(msg) {
                        InboundFrame::Text(t) => {
                            if let Err(exit) = self.handle_text(&t).await {
                                break exit;
                            }
                        }
                        InboundFrame::Ignore => {}
                        InboundFrame::Lost(why) => break MessageLoopExit::ConnectionLost(why),
                    }
                }

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

                req = self.continue_tos_rx.recv() => {
                    match req {
                        Some(params) => {
                            if let Err(e) = send_request(&mut ws, &BacktestRequest::ContinueTo(params)).await {
                                break MessageLoopExit::ConnectionLost(format!("continue_to send: {e}"));
                            }
                        }
                        None => break MessageLoopExit::SessionEnded,
                    }
                }
            }
        };

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

impl ControlTask {
    /// Returns `Err(MessageLoopExit)` if the message signals we should exit the loop.
    async fn handle_text(&mut self, text: &str) -> Result<(), MessageLoopExit> {
        let (seq, response) = match serde_json::from_str::<SequencedResponse>(text) {
            Ok(s) => (Some(s.seq_id), s.response),
            Err(_) => match serde_json::from_str::<BacktestResponse>(text) {
                Ok(r) => (None, r),
                Err(e) => {
                    warn!(error = %err_chain(&e), "discarding undeserializable control message");
                    return Ok(());
                }
            },
        };

        if let Some(s) = seq {
            self.last_sequence = Some(s);
        }

        match response {
            BacktestResponse::ReadyForContinue => {
                let _ = self.events_tx.send(ControlEvent::ReadyForContinue).await;
            }
            BacktestResponse::Paused(event) => {
                let _ = self.events_tx.send(ControlEvent::Paused(event)).await;
            }
            BacktestResponse::DiscoveryBatch(event) => {
                let _ = self
                    .events_tx
                    .send(ControlEvent::DiscoveryBatch(event))
                    .await;
            }
            BacktestResponse::SlotNotification(slot) => {
                let _ = self.events_tx.send(ControlEvent::Slot(slot)).await;
            }
            BacktestResponse::Completed {
                summary,
                agent_stats,
            } => {
                let _ = self
                    .events_tx
                    .send(ControlEvent::Completed {
                        summary,
                        agent_stats,
                    })
                    .await;
                return Err(MessageLoopExit::SessionEnded);
            }
            BacktestResponse::Error(err) => {
                // Per-slot simulation errors are non-fatal: log and keep going.
                if matches!(&err, BacktestError::SimulationError { .. }) {
                    warn!(error = %err_chain(&err), "simulation error");
                    return Ok(());
                }
                let terminal = is_terminal_backtest_error(&err);
                let _ = self.events_tx.send(ControlEvent::Error(err)).await;
                if terminal {
                    return Err(MessageLoopExit::Terminal(
                        "server reported terminal error".into(),
                    ));
                }
            }
            BacktestResponse::Status { status } => {
                let _ = self.events_tx.send(ControlEvent::Status(status)).await;
            }
            BacktestResponse::Success => {
                // Ack for Close or similar; nothing to forward.
            }
            other => {
                // SessionCreated/Attached/etc. during the message loop are unexpected.
                debug!(?other, "ignoring unexpected control response");
            }
        }

        Ok(())
    }
}

async fn create_session(
    ws: &mut Ws,
    request: CreateBacktestSessionRequest,
    url: &str,
    events: &mut mpsc::Sender<ControlEvent>,
    last_sequence: &mut Option<u64>,
) -> Result<SessionInfo, HandshakeError> {
    send_request(ws, &BacktestRequest::CreateBacktestSession(request))
        .await
        .map_err(HandshakeError::Transient)?;

    let rpc_base = http_base_from_ws_url(url);

    loop {
        let response = next_response_with_timeout(ws, events, last_sequence)
            .await
            .map_err(HandshakeError::Transient)?;
        match response {
            BacktestResponse::SessionCreated {
                session_id,
                rpc_endpoint,
                task_id,
            } => {
                let rpc_endpoint = resolve_rpc_url(&rpc_base, &rpc_endpoint);
                return Ok(SessionInfo {
                    session_id,
                    rpc_endpoint,
                    task_id,
                });
            }
            BacktestResponse::Error(err) => {
                return Err(HandshakeError::Fatal(format!(
                    "server error: {}",
                    err_chain(&err)
                )));
            }
            _ => {
                // Any unexpected response before SessionCreated — ignore and
                // keep waiting. (e.g. statuses, early events.)
            }
        }
    }
}

async fn attach(
    ws: &mut Ws,
    session_id: &str,
    last_sequence: Option<u64>,
    events: &mut mpsc::Sender<ControlEvent>,
    last_seq_state: &mut Option<u64>,
) -> Result<(), HandshakeError> {
    send_request(
        ws,
        &BacktestRequest::AttachBacktestSession {
            session_id: session_id.to_string(),
            last_sequence,
        },
    )
    .await
    .map_err(HandshakeError::Transient)?;

    loop {
        let response = next_response_with_timeout(ws, events, last_seq_state)
            .await
            .map_err(HandshakeError::Transient)?;
        match response {
            BacktestResponse::SessionAttached { .. } => return Ok(()),
            BacktestResponse::Error(err) => {
                return Err(handshake_error_for_response("attach", err));
            }
            _ => {}
        }
    }
}

async fn resume(
    ws: &mut Ws,
    events: &mut mpsc::Sender<ControlEvent>,
    last_seq_state: &mut Option<u64>,
) -> Result<(), HandshakeError> {
    send_request(ws, &BacktestRequest::ResumeAttachedSession)
        .await
        .map_err(HandshakeError::Transient)?;

    loop {
        let response = next_response_with_timeout(ws, events, last_seq_state)
            .await
            .map_err(HandshakeError::Transient)?;
        match response {
            BacktestResponse::Success => return Ok(()),
            BacktestResponse::Error(err) => {
                return Err(handshake_error_for_response("resume", err));
            }
            _ => {}
        }
    }
}

/// Read the next response during a handshake, with a bounded timeout.
///
/// Any non-handshake responses observed during the wait are forwarded to the
/// driver (slot notifications, errors) so we don't lose them.
async fn next_response_with_timeout(
    ws: &mut Ws,
    events: &mut mpsc::Sender<ControlEvent>,
    last_sequence: &mut Option<u64>,
) -> Result<BacktestResponse, String> {
    let deadline = tokio::time::Instant::now() + HANDSHAKE_RESPONSE_TIMEOUT;
    loop {
        let msg = tokio::time::timeout_at(deadline, ws.next())
            .await
            .map_err(|_| format!("handshake timeout after {HANDSHAKE_RESPONSE_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 std::str::from_utf8(&b) {
                Ok(t) => t.to_string(),
                Err(_) => continue,
            },
            Message::Close(frame) => {
                return Err(format!("remote close during handshake: {frame:?}"));
            }
            _ => continue,
        };

        let (seq, response) = match serde_json::from_str::<SequencedResponse>(&text) {
            Ok(s) => (Some(s.seq_id), s.response),
            Err(_) => (
                None,
                serde_json::from_str::<BacktestResponse>(&text)
                    .map_err(|e| format!("deserialize: {}; raw={text}", err_chain(&e)))?,
            ),
        };
        if let Some(s) = seq {
            *last_sequence = Some(s);
        }

        // Forward noisy event kinds to the driver so nothing is lost while we
        // wait for the handshake response.
        match response {
            BacktestResponse::SlotNotification(slot) => {
                let _ = events.send(ControlEvent::Slot(slot)).await;
            }
            BacktestResponse::ReadyForContinue => {
                let _ = events.send(ControlEvent::ReadyForContinue).await;
            }
            BacktestResponse::Paused(event) => {
                let _ = events.send(ControlEvent::Paused(event)).await;
            }
            BacktestResponse::DiscoveryBatch(event) => {
                let _ = events.send(ControlEvent::DiscoveryBatch(event)).await;
            }
            BacktestResponse::Completed {
                summary,
                agent_stats,
            } => {
                let _ = events
                    .send(ControlEvent::Completed {
                        summary,
                        agent_stats,
                    })
                    .await;
            }
            other => return Ok(other),
        }
    }
}