sim-cli 0.7.0

use std::{
    collections::{BTreeMap, HashMap},
    path::{Path, PathBuf},
    sync::{Arc, Mutex},
    time::Duration,
};

use clap::Parser;
use console::style;
use eyre::{Context, Result};
use indicatif::{MultiProgress, ProgressBar};
use sim_cli::{
    engine::{
        self, DrivableSession, ModificationsProvider, Outcome, SessionFailure, SessionObserver,
        SlotRange, format_slot,
    },
    signals::spawn_ctrlc_cursor_fix,
};
use simulator_client::{
    BacktestClient, CreateSession, backtest_ws_url,
    managed::{
        ManagedBacktestSession, ManagedEvent, ManagedParallelSession, ManagedSessionError,
        ParallelSubSession, ReconnectCoordinator,
    },
    modify_program_via_rpc, split_range,
};
use solana_client::nonblocking::rpc_client::RpcClient;
use solana_rpc_client::{http_sender::HttpSender, rpc_client::RpcClientConfig};
use tokio_util::sync::CancellationToken;
use tracing::info;

use crate::{
    output::{OutputEvent, SessionStartedRecord, SimulationMetadata, SimulationSummary},
    subscription::{RecordsStore, on_account_diff_notification, on_transaction_notification},
};

const MAX_PARALLEL_SESSIONS: usize = 50;
const CREATION_TIMEOUT_SECS: u16 = 900;
/// How long the server holds a session's runtime after the control connection
/// drops, giving the managed session time to reconnect and resume. Applied to
/// both the single-session and parallel creates so the reconnect window matches.
const DISCONNECT_TIMEOUT_SECS: u16 = 180;
const ESTIMATED_MS_PER_BLOCK: u64 = 250;
const HTTP_REQUEST_TIMEOUT: Duration = Duration::from_secs(30);
const HTTP_POOL_IDLE_TIMEOUT: Duration = Duration::from_secs(30);

/// `(program_id, elf_bytes)` pairs preloaded from disk before sessions start.
/// Read once and shared across parallel sessions to keep the file-descriptor
/// count bounded — opening the .so per-session blew through the per-process
/// fd limit when many sessions ran concurrently.
type PreloadedPrograms = Vec<(String, Vec<u8>)>;

/// Shared infrastructure threaded into every session — built once in `run()`
/// and cloned (cheaply) into each parallel session task.
#[derive(Clone)]
struct SessionInfra {
    preloaded_programs: Arc<PreloadedPrograms>,
    http_client: reqwest::Client,
}

/// Build an `RpcClient` that reuses the given shared `reqwest::Client`.
///
/// `RpcClient::new(url)` builds its own internal `reqwest::Client` per call,
/// so per-session clients each have their own connection pool and trigger an
/// independent DNS lookup on first use. Under high parallelism this stresses
/// the macOS system resolver enough to produce sporadic
/// `getaddrinfo`/`task cancelled` failures. Sharing one HTTP client folds all
/// sessions onto a single pool so DNS is hit once per host.
fn make_rpc_client(rpc_endpoint: &str, http_client: &reqwest::Client) -> RpcClient {
    let sender = HttpSender::new_with_client(rpc_endpoint, http_client.clone());
    RpcClient::new_sender(sender, RpcClientConfig::default())
}

impl ModificationsProvider for SessionInfra {
    fn is_empty(&self) -> bool {
        self.preloaded_programs.is_empty()
    }

    async fn build(
        &self,
        rpc_endpoint: &str,
    ) -> Result<BTreeMap<solana_address::Address, simulator_api::AccountData>> {
        if self.preloaded_programs.is_empty() {
            return Ok(BTreeMap::new());
        }
        let rpc = make_rpc_client(rpc_endpoint, &self.http_client);
        let mut mods = BTreeMap::new();
        for (id, elf) in self.preloaded_programs.iter() {
            let m = modify_program_via_rpc(&rpc, id, elf)
                .await
                .map_err(|e| eyre::eyre!("building program injection for {id}: {e}"))?;
            mods.extend(m);
        }
        Ok(mods)
    }
}

#[derive(Clone, Debug, PartialEq, Eq, clap::ValueEnum)]
pub enum SubscriptionType {
    /// Alias for `transaction`. The server-pushed transaction payload already
    /// includes logs, so logs mode is just transaction mode under a different
    /// name.
    Logs,
    /// Per-account state diffs. Pairs with a concurrent transactionSubscribe
    /// stream that populates balance/token enrichment without any RPC calls.
    AccountDiff,
    /// Full per-tx records pushed by the server, matching the shape of a
    /// `getTransaction` response. The default — no RPC fetch needed.
    Transaction,
}

#[derive(Clone)]
struct OutputSink {
    /// Per-tx records used for the trailing `summary` totals and (in
    /// account-diff mode) for matching each diff back to its tx's
    /// balance/token data. See [`RecordsStore`] for retention behavior.
    records: RecordsStore,
    /// Primary output channel — every event (metadata / session_started /
    /// tx / diff / summary) flows through here to the NDJSON writer task.
    output_tx: tokio::sync::mpsc::UnboundedSender<OutputEvent>,
}

impl SessionObserver for OutputSink {
    fn on_data_event(&self, event: ManagedEvent) {
        match event {
            ManagedEvent::Transaction(notification) => on_transaction_notification(
                self.records.clone(),
                self.output_tx.clone(),
                *notification,
            ),
            ManagedEvent::AccountDiff(notification) => on_account_diff_notification(
                self.records.clone(),
                self.output_tx.clone(),
                notification,
            ),
            // Other variants are handled inline in the engine's match arms.
            _ => {}
        }
    }
}

#[derive(Parser, Debug, Clone)]
pub struct RunArgs {
    /// First slot (inclusive) to simulate
    #[arg(long, env = "SIMULATOR_START_SLOT")]
    pub start_slot: u64,

    /// Last slot (inclusive) to simulate
    #[arg(long, env = "SIMULATOR_END_SLOT")]
    pub end_slot: u64,

    /// Number of slots to advance per continue request (default: full range)
    #[arg(long, env = "SIMULATOR_ADVANCE_COUNT")]
    pub advance_count: Option<u64>,

    /// Split the range across multiple parallel sessions (up to 50)
    #[arg(long, env = "SIMULATOR_PARALLEL", default_value_t = false)]
    pub parallel: bool,

    /// In `--parallel` mode, abort the remaining sub-ranges as soon as one
    /// fails instead of running them all and reporting failures at the end.
    #[arg(long, env = "SIMULATOR_FAIL_FAST", default_value_t = false)]
    pub fail_fast: bool,

    /// Program ID(s) to collect transactions for (repeatable)
    #[arg(long, env = "SIMULATOR_PROGRAM_ID")]
    pub program_id: Vec<String>,

    /// Path(s) to compiled .so file(s) to inject before simulation (one per --program-id, same order)
    #[arg(long, env = "SIMULATOR_PROGRAM_SO")]
    pub program_so: Vec<PathBuf>,

    /// File path to write JSON output to (default: <start_slot>_<end_slot>_baseline.json or _override.json)
    #[arg(long, env = "SIMULATOR_OUTPUT_FILE")]
    pub output_file: Option<PathBuf>,

    /// Extra compute units to add to each transaction's SetComputeUnitLimit budget.
    #[arg(long, env = "SIMULATOR_EXTRA_COMPUTE_UNITS")]
    pub extra_compute_units: Option<u32>,

    /// Subscription type to collect per transaction (requires --program-id).
    /// One of: `logs`, `account-diff`, `transaction`. Defaults to `transaction`
    /// — the only mode that produces a full per-tx record with balances/tokens
    /// without any follow-up RPC calls.
    #[arg(long, env = "SIMULATOR_SUBSCRIPTION", requires = "program_id")]
    pub subscription: Option<SubscriptionType>,

    /// Show tracing logs instead of a progress bar
    #[arg(long, default_value_t = false)]
    pub verbose: bool,

    /// Show session plan.
    #[arg(long, env = "SIMULATOR_PLAN", default_value_t = false)]
    pub plan: bool,
}

fn validate_args(args: &RunArgs) -> Result<()> {
    if !args.program_so.is_empty() {
        eyre::ensure!(
            args.program_so.len() == args.program_id.len(),
            "--program-so count ({}) must match --program-id count ({})",
            args.program_so.len(),
            args.program_id.len(),
        );
    }
    Ok(())
}

async fn generate_plan(client: &BacktestClient, args: &RunArgs) -> Result<()> {
    let ranges = client
        .available_ranges()
        .await
        .map_err(|e| eyre::eyre!("failed to fetch available ranges: {e}"))?;

    let sub_ranges = split_range(&ranges, args.start_slot, args.end_slot)
        .map_err(|e| eyre::eyre!("cannot split range: {e}"))?;

    if sub_ranges.len() > MAX_PARALLEL_SESSIONS {
        eyre::bail!(
            "unable to generate plan: requires {} sessions but max is {MAX_PARALLEL_SESSIONS}",
            sub_ranges.len(),
        );
    }
    for window in sub_ranges.windows(2) {
        let (_, end) = window[0];
        let (next_start, _) = window[1];
        if end + 1 != next_start {
            eyre::bail!(
                "unable to generate plan: gap in data between slots {} and {}",
                format_slot(end),
                format_slot(next_start),
            );
        }
    }

    let max_slots = sub_ranges
        .iter()
        .map(|(start, end)| end.saturating_sub(*start) + 1)
        .max()
        .unwrap_or(0);
    let estimated_ms = max_slots * ESTIMATED_MS_PER_BLOCK;
    let estimated = if estimated_ms < 60 * 60 * 1000 {
        format!("~{}m", estimated_ms / 60_000)
    } else {
        format!("~{:.1}h", estimated_ms as f64 / 3_600_000.0)
    };

    println!(
        "\n{} {estimated} runtime · {} session(s)\n{} [{} → {}] · {} slots",
        style(">").cyan(),
        sub_ranges.len(),
        style(">").cyan(),
        format_slot(args.start_slot),
        format_slot(args.end_slot),
        format_slot(args.end_slot - args.start_slot + 1),
    );
    println!();
    println!(
        "{:<5} {:<20} {:<20} {:<15}",
        "#", "Start Slot", "End Slot", "Slots"
    );
    println!("{}", "-".repeat(62));
    for (i, (start, end)) in sub_ranges.iter().enumerate() {
        let slots = end.saturating_sub(*start) + 1;
        println!(
            "{:<5} {:<20} {:<20} {:<15}",
            i + 1,
            format_slot(*start),
            format_slot(*end),
            format_slot(slots),
        );
    }

    Ok(())
}

fn resolve_output_file(args: &RunArgs) -> PathBuf {
    args.output_file.clone().unwrap_or_else(|| {
        let suffix = if !args.program_so.is_empty() {
            "override"
        } else {
            "baseline"
        };
        PathBuf::from(format!(
            "{}_{}_{}.ndjson",
            args.start_slot, args.end_slot, suffix
        ))
    })
}

/// Emit the `metadata` envelope first on the channel, ahead of any session output.
fn setup_output(
    output_file: &Path,
    args: &RunArgs,
) -> Result<(OutputSink, tokio::task::JoinHandle<Result<()>>)> {
    let output_writer = open_output_writer(output_file)?;
    let (output_tx, output_rx) = tokio::sync::mpsc::unbounded_channel::<OutputEvent>();
    let output_writer_handle =
        spawn_output_writer(output_file.to_path_buf(), output_writer, output_rx);

    let metadata = SimulationMetadata {
        start_slot: args.start_slot,
        end_slot: args.end_slot,
        program_ids: args.program_id.clone(),
        program_so: args
            .program_so
            .iter()
            .map(|p| p.display().to_string())
            .collect(),
        ran_at_unix_secs: std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .map(|d| d.as_secs())
            .unwrap_or(0),
        // Session IDs aren't known yet — they appear on the trailing `summary`
        // line once sessions have been created.
        session_ids: Vec::new(),
    };
    let _ = output_tx.send(OutputEvent::Metadata(metadata));

    let output = OutputSink {
        records: RecordsStore::new(),
        output_tx,
    };
    Ok((output, output_writer_handle))
}

fn build_infra(args: &RunArgs) -> Result<SessionInfra> {
    let preloaded_programs: Arc<PreloadedPrograms> = Arc::new(
        args.program_id
            .iter()
            .zip(&args.program_so)
            .map(|(id, path)| {
                let elf = std::fs::read(path)
                    .with_context(|| format!("failed to read {}", path.display()))?;
                Ok((id.clone(), elf))
            })
            .collect::<Result<_>>()?,
    );

    // `hickory_dns(true)` swaps reqwest's default resolver (tokio's wrapper
    // around the OS `getaddrinfo`) for an in-process async resolver with TTL
    // caching. Without it, every fresh TCP connection triggers a system-call
    // DNS lookup; on macOS those serialize through `mDNSResponder` and fail
    // sporadically under high concurrency.
    let http_client = reqwest::Client::builder()
        .timeout(HTTP_REQUEST_TIMEOUT)
        .pool_idle_timeout(HTTP_POOL_IDLE_TIMEOUT)
        .hickory_dns(true)
        .build()
        .context("failed to build shared HTTP client")?;

    Ok(SessionInfra {
        preloaded_programs,
        http_client,
    })
}

pub async fn run(
    args: RunArgs,
    url: String,
    api_key: String,
    cancellation: CancellationToken,
) -> Result<()> {
    validate_args(&args)?;

    let base_url = backtest_ws_url(&url);
    let client = BacktestClient::builder()
        .url(base_url.clone())
        .api_key(api_key.clone())
        .build();

    if args.plan {
        return generate_plan(&client, &args).await;
    }

    let output_file = resolve_output_file(&args);
    let (output, output_writer_handle) = setup_output(&output_file, &args)?;
    let session_ids: Arc<Mutex<Vec<String>>> = Arc::new(Mutex::new(Vec::new()));
    let infra = build_infra(&args)?;

    let args = Arc::new(args);

    spawn_ctrlc_cursor_fix();

    if args.verbose {
        tracing_subscriber::fmt()
            .with_env_filter(
                tracing_subscriber::EnvFilter::try_from_default_env()
                    .unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("info")),
            )
            .init();
    }

    // Sub-ranges that did not complete, collected across the run so the tail can
    // render one de-duplicated summary rather than aborting on the first failure.
    let mut session_failures: Vec<SessionFailure> = Vec::new();
    let mut session_count = 1;

    if args.parallel {
        let progress = MultiProgress::new();

        // Compute the split locally so each sub-session gets a determinate
        // `[start, end]` bar. The server splits the range with the same shared
        // `split_range`, so these ranges match the created sub-sessions exactly.
        let ranges = client
            .available_ranges()
            .await
            .map_err(|e| eyre::eyre!("failed to fetch available ranges: {e}"))?;

        let sub_ranges = split_range(&ranges, args.start_slot, args.end_slot)
            .map_err(|e| eyre::eyre!("cannot split range: {e}"))?;

        session_count = sub_ranges.len();
        info!("splitting into {session_count} parallel session(s)");
        let _ = progress.println(format!(
            "{} Splitting into {session_count} parallel session(s)",
            style(">").cyan(),
        ));

        // One display bar per sub-range, keyed by start slot. Each sub-session
        // binds to its bar by the server-authoritative range it reports, so this
        // local split only seeds the bars (a mismatch falls back to a fresh bar).
        let mut bars_by_start: HashMap<u64, ProgressBar> = HashMap::new();
        for &(start, end) in &sub_ranges {
            let bar = if args.verbose {
                ProgressBar::hidden()
            } else {
                let bar = ProgressBar::new(end.saturating_sub(start) + 1);
                bar.set_style(engine::spinner_style(start, end));
                let bar = progress.add(bar);
                bar.enable_steady_tick(Duration::from_millis(100));
                bar
            };
            bars_by_start.insert(start, bar);
        }

        let create = CreateSession::builder()
            .start_slot(args.start_slot)
            .end_slot(args.end_slot)
            .parallel(true)
            .disconnect_timeout_secs(DISCONNECT_TIMEOUT_SECS)
            .capacity_wait_timeout_secs(CREATION_TIMEOUT_SECS)
            .maybe_extra_compute_units(args.extra_compute_units)
            .build()
            .into_request()
            .map_err(|e| eyre::eyre!("building parallel create request: {e}"))?;

        let parallel = match ManagedParallelSession::start_with_cancel(
            base_url.clone(),
            api_key.clone(),
            create,
            cancellation.clone(),
        )
        .await
        {
            Ok(p) => Some(p),
            Err(ManagedSessionError::Cancelled) => None,
            Err(e) => return Err(eyre::eyre!("parallel session create failed: {e}")),
        };

        if let Some(parallel) = parallel {
            let bars_by_start = Arc::new(bars_by_start);
            let drive = {
                let args = args.clone();
                let output = output.clone();
                let session_ids = session_ids.clone();
                let infra = infra.clone();
                let bars = bars_by_start.clone();
                move |session: ParallelSubSession| {
                    let args = args.clone();
                    let output = output.clone();
                    let session_ids = session_ids.clone();
                    let infra = infra.clone();
                    let bars = bars.clone();
                    async move {
                        drive_parallel_sub(session, &args, &bars, &output, &session_ids, &infra)
                            .await
                    }
                }
            };
            session_failures = engine::run_parallel_sessions(
                parallel,
                bars_by_start,
                &progress,
                args.fail_fast,
                None,
                &cancellation,
                drive,
            )
            .await;
        } else {
            for bar in bars_by_start.values() {
                if !bar.is_finished() {
                    bar.abandon_with_message("cancelled");
                }
            }
        }
    } else {
        let range = SlotRange {
            start: args.start_slot,
            end: args.end_slot,
        };
        let bar = if args.verbose {
            ProgressBar::hidden()
        } else {
            let bar = ProgressBar::new(range.end.saturating_sub(range.start) + 1);
            bar.set_style(engine::spinner_style(range.start, range.end));
            bar.set_message("starting session...");
            bar.enable_steady_tick(Duration::from_millis(100));
            bar
        };

        drive_single(
            base_url,
            api_key,
            &args,
            &bar,
            range,
            &output,
            &session_ids,
            &infra,
            &cancellation,
            None,
        )
        .await
        .inspect_err(|e| bar.abandon_with_message(format!("failed: {e}")))?;
    }

    if cancellation.is_cancelled() {
        eprintln!("\n{} Cancelled by user", style("!").yellow());
        // Still close the writer cleanly so partial output is flushed to disk.
        let OutputSink { output_tx, .. } = output;
        drop(output_tx);
        let _ = output_writer_handle.await;
        return Ok(());
    }

    let OutputSink {
        records, output_tx, ..
    } = output;

    let session_ids = session_ids.lock().unwrap().clone();
    let total = records.total();
    let successes = records.successes();
    let summary = SimulationSummary {
        total_transactions: total,
        successes,
        failures: total.saturating_sub(successes),
        session_ids,
    };

    let _ = output_tx.send(OutputEvent::Summary(summary.clone()));
    drop(output_tx);
    output_writer_handle
        .await
        .context("output writer task panicked")?
        .context("output writer failed")?;

    render_outcome(&session_failures, session_count, &summary, &output_file)
}

/// Render the run tail. Runs after the output writer flushed, so the printed
/// file path points at a complete file.
fn render_outcome(
    session_failures: &[SessionFailure],
    session_count: usize,
    summary: &SimulationSummary,
    output_file: &Path,
) -> Result<()> {
    if !session_failures.is_empty() {
        // De-duplicate by failure reason: identical errors across many
        // sub-ranges collapse to one line listing the affected slot ranges,
        // instead of one noisy stanza per session.
        let mut by_reason: BTreeMap<&str, Vec<(u64, u64)>> = BTreeMap::new();
        for failure in session_failures {
            by_reason
                .entry(failure.reason.as_str())
                .or_default()
                .extend(failure.range);
        }

        println!(
            "\n{} {} of {} sub-range(s) did not complete",
            style("✖").red(),
            session_failures.len(),
            session_count,
        );
        for (reason, ranges) in &mut by_reason {
            ranges.sort_unstable();
            let slots = if ranges.is_empty() {
                String::new()
            } else {
                let list = ranges
                    .iter()
                    .map(|(s, e)| format!("{s}–{e}"))
                    .collect::<Vec<_>>()
                    .join(", ");
                format!(" {}", style(format!("[slots {list}]")).dim())
            };
            println!("  {}{}", style(reason).yellow(), slots);
        }
        println!(
            "\n  Captured {} transaction(s) from the sub-ranges that finished",
            summary.total_transactions,
        );
        println!("  Output file  : {}", output_file.display());

        // Phrased as a summary so the wrapping red mark in `main.rs` reads as a
        // total, not another per-sub-range failure line (which looked like a
        // second failure when only one sub-range failed).
        return Err(eyre::eyre!(
            "summary: {} of {} parallel sub-ranges failed",
            session_failures.len(),
            session_count,
        ));
    }

    if summary.total_transactions == 0 {
        println!(
            "\n{} Simulation finished with no transactions",
            style("⚠").yellow()
        );
        println!("  Output file  : {}", output_file.display());
        return Err(eyre::eyre!(
            "no matching transactions found in the requested slot range"
        ));
    }

    println!("\n{} Simulation complete", style("✔").green());
    println!("  Transactions : {}", summary.total_transactions);
    println!("  Successes    : {}", summary.successes);
    println!("  Failures     : {}", summary.failures);
    println!("  Output file  : {}", output_file.display());

    Ok(())
}

/// Push the session id and emit the opening `session_started` output row.
fn record_session_started<S: DrivableSession>(
    session: &S,
    range: SlotRange,
    output: &OutputSink,
    session_ids: &Arc<Mutex<Vec<String>>>,
) {
    let info = session.session_info();
    session_ids.lock().unwrap().push(info.session_id.clone());
    info!(
        session_id = %info.session_id,
        task_id = ?info.task_id,
        start = range.start,
        end = range.end,
        "session created"
    );
    let _ = output
        .output_tx
        .send(OutputEvent::SessionStarted(SessionStartedRecord {
            session_id: info.session_id.clone(),
            task_id: info.task_id.clone(),
            start_slot: range.start,
            end_slot: range.end,
        }));
}

/// Subscribe (if a program filter is set) and drive a created session to
/// completion. Reconnect, keepalive, and Up-gating on `send_continue` are
/// handled by the wrapper; this just maps events to progress and output rows.
async fn subscribe_and_drive<S: DrivableSession>(
    mut session: S,
    args: &RunArgs,
    pb: &ProgressBar,
    range: SlotRange,
    advance_count: u64,
    output: &OutputSink,
    infra: &SessionInfra,
) -> Result<Outcome> {
    // Stay in spinner style until the server signals `ReadyForContinue`.
    pb.set_message("waiting for runtime");

    // AccountDiff mode runs two subscriptions in parallel: transactionSubscribe
    // (for tx records used to enrich diffs) plus accountDiffSubscribe. All
    // other modes get just transactionSubscribe.
    if !args.program_id.is_empty() {
        session.subscribe_transactions(args.program_id.clone());
        if matches!(args.subscription, Some(SubscriptionType::AccountDiff)) {
            session.subscribe_account_diffs(args.program_id.clone());
        }
    }

    let outcome =
        engine::run_session_loop(&mut session, pb, range, advance_count, infra, output).await;
    session.shutdown().await;
    outcome
}

/// Create and drive a single client-side session to completion.
#[allow(clippy::too_many_arguments)]
async fn drive_single(
    base_url: String,
    api_key: String,
    args: &RunArgs,
    pb: &ProgressBar,
    range: SlotRange,
    output: &OutputSink,
    session_ids: &Arc<Mutex<Vec<String>>>,
    infra: &SessionInfra,
    parent_cancel: &CancellationToken,
    reconnect_coordinator: Option<Arc<ReconnectCoordinator>>,
) -> Result<()> {
    let advance_count = args
        .advance_count
        .unwrap_or(range.end.saturating_sub(range.start) + 1);

    let create = CreateSession::builder()
        .start_slot(range.start)
        .end_slot(range.end)
        .disconnect_timeout_secs(DISCONNECT_TIMEOUT_SECS)
        .capacity_wait_timeout_secs(CREATION_TIMEOUT_SECS)
        .maybe_extra_compute_units(args.extra_compute_units)
        .build()
        .into_request()
        .map_err(|e| eyre::eyre!("building create request: {e}"))?;

    let session = match ManagedBacktestSession::start_with_cancel(
        base_url,
        api_key,
        create,
        parent_cancel.clone(),
        reconnect_coordinator,
    )
    .await
    {
        Ok(s) => s,
        Err(ManagedSessionError::Cancelled) => {
            pb.abandon_with_message("cancelled");
            return Ok(());
        }
        Err(e) => return Err(eyre::eyre!("session create failed: {e}")),
    };

    record_session_started(&session, range, output, session_ids);
    let outcome = subscribe_and_drive(session, args, pb, range, advance_count, output, infra).await;
    engine::finish_bar(pb, outcome)
}

/// Drive one sub-session of a server-multiplexed parallel run. The sub-session's
/// range is server-authoritative (from `SessionsCreatedV2`), so we bind to its
/// pre-created bar by that range's start and run the shared driver loop — no
/// guessing the range from the first slot, and a client/server split mismatch
/// just falls back to a fresh bar rather than failing the sub-session.
#[allow(clippy::too_many_arguments)]
async fn drive_parallel_sub(
    session: ParallelSubSession,
    args: &RunArgs,
    bars: &HashMap<u64, ProgressBar>,
    output: &OutputSink,
    session_ids: &Arc<Mutex<Vec<String>>>,
    infra: &SessionInfra,
) -> std::result::Result<(), SessionFailure> {
    let (start, end) = session.range();
    let range = SlotRange { start, end };
    let pb = bars
        .get(&start)
        .cloned()
        .unwrap_or_else(ProgressBar::hidden);
    let advance_count = args.advance_count.unwrap_or(end.saturating_sub(start) + 1);

    record_session_started(&session, range, output, session_ids);

    let outcome =
        subscribe_and_drive(session, args, &pb, range, advance_count, output, infra).await;
    engine::finish_bar(&pb, outcome).map_err(|e| SessionFailure {
        range: Some((start, end)),
        reason: format!("{e:#}"),
    })
}

fn open_output_writer(path: &std::path::Path) -> Result<std::fs::File> {
    if let Some(parent) = path.parent()
        && !parent.as_os_str().is_empty()
        && !parent.exists()
    {
        std::fs::create_dir_all(parent)
            .with_context(|| format!("could not create output directory '{}'", parent.display()))?;
    }
    std::fs::File::create(path).with_context(|| format!("could not create '{}'", path.display()))
}

fn spawn_output_writer(
    path: PathBuf,
    file: std::fs::File,
    mut rx: tokio::sync::mpsc::UnboundedReceiver<OutputEvent>,
) -> tokio::task::JoinHandle<Result<()>> {
    tokio::spawn(async move {
        use std::io::Write;
        let write_err = || format!("could not write to '{}'", path.display());
        let mut writer = std::io::BufWriter::new(file);

        while let Some(event) = rx.recv().await {
            serde_json::to_writer(&mut writer, &event).with_context(write_err)?;
            writer.write_all(b"\n").with_context(write_err)?;
            while let Ok(event) = rx.try_recv() {
                serde_json::to_writer(&mut writer, &event).with_context(write_err)?;
                writer.write_all(b"\n").with_context(write_err)?;
            }
            writer.flush().with_context(write_err)?;
        }

        writer
            .flush()
            .with_context(|| format!("could not finalize '{}'", path.display()))?;
        Ok(())
    })
}