rlmesh-runtime 0.1.0-rc.9

//! [`RouteState`]: the per-route bookkeeping the driver advances each step, and
//! the request messages it builds from that state.
//!
//! Slots are per lane (`env_index`). The driver works on *groups* of lanes
//! (one group per lane for a lane endpoint, one group for the whole vector
//! otherwise), so every accessor takes the group's slot `positions` and reads
//! or advances only those lanes. Episode indices and seeds come from the
//! route-global slot counter ([`claim_slots`](RouteState::claim_slots)), so
//! they depend on start order alone, never on which group ran the episode.

use prost::bytes::Bytes;
use rlmesh_proto::model::v1::{
    AdapterContext, EpisodeInfo, PredictRequest, ReleaseAdapterRequest, ResetAdapterRequest,
};
use rlmesh_proto::spaces::v1::SpaceValue;

use std::collections::HashMap;

use crate::episodes::{EpisodeRecord, EpisodeRecordRegistry};
use crate::hooks::RuntimeEnvContext;
use crate::spec::{EpisodeSummary, RuntimeSessionSpec};

use super::{EpisodeState, RouteSnapshot, SlotState, StartedEpisode};

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum RequestPhase {
    ResetObservation,
    StepObservation,
}

impl RequestPhase {
    pub(crate) fn as_str(self) -> &'static str {
        match self {
            Self::ResetObservation => "reset_observation",
            Self::StepObservation => "step_observation",
        }
    }
}

fn leaves_value(leaves: Vec<Bytes>) -> SpaceValue {
    SpaceValue { leaves }
}

#[derive(Debug)]
pub(crate) struct RouteState {
    session_id: String,
    env_id: String,
    env_component_id: String,
    model_component_id: String,
    slots: Vec<SlotState>,
    request_seq: u64,
    total_steps: i64,
    total_episodes: i64,
    records: EpisodeRecordRegistry,
    episode_summaries: Vec<EpisodeSummary>,
    /// The explicit seed each live episode was reset with, keyed by episode id;
    /// drained into that episode's summary at completion.
    seed_by_episode: HashMap<String, i64>,
    /// The next route-global episode slot. A slot is claimed when an episode
    /// starts; it is the episode's index and picks its seed.
    next_slot: u64,
    /// The episode budget the bounded claim honors.
    max_episodes: Option<u64>,
    /// The trial ordinal each live episode was reset with, keyed by episode id;
    /// drained into that episode's summary at completion.
    trial_by_episode: HashMap<String, u64>,
    /// How many trial ordinals have been minted off `trial_index_base`
    /// (episode-start order). Monotone for the life of the route.
    trial_cursor: u64,
}

impl RouteState {
    pub(crate) fn new(spec: &RuntimeSessionSpec) -> Self {
        let slots = (0..spec.num_envs.max(1))
            .map(|index| SlotState {
                env_index: index.try_into().unwrap_or(i32::MAX),
                episode: None,
                step: 0,
                reset: true,
                cumulative_reward: 0.0,
                started_at_ns: now_unix_ns(),
            })
            .collect();

        Self {
            session_id: spec.session_id.clone(),
            env_id: spec.env_id.clone(),
            env_component_id: spec.env_component_id.clone(),
            model_component_id: spec.model_component_id.clone(),
            slots,
            request_seq: 0,
            total_steps: 0,
            total_episodes: 0,
            records: EpisodeRecordRegistry::default(),
            episode_summaries: Vec::new(),
            seed_by_episode: HashMap::new(),
            next_slot: 0,
            max_episodes: spec.max_episodes,
            trial_by_episode: HashMap::new(),
            trial_cursor: 0,
        }
    }

    /// Claim the next `lanes` trial ordinals off `trial_index_base`, in
    /// episode-start order and positionally aligned to the lanes being reset.
    ///
    /// The window rule: a route's trial window is its `max_episodes` budget M --
    /// the run ends once `trials_completed_in_window >= M`, so a shard walks
    /// exactly the ordinals `[base, base + M)` and the next shard's base is
    /// `base + M`. The cursor only walks forward, and every reset claims one
    /// ordinal per restarted lane, so no two episodes on a route ever share one.
    pub(crate) fn claim_trial_indices(&mut self, base: u64, lanes: usize) -> Vec<u64> {
        let start = base.saturating_add(self.trial_cursor);
        self.trial_cursor += lanes as u64;
        (0..lanes as u64).map(|offset| start + offset).collect()
    }

    /// Remember which trial ordinal each episode in a reset batch received, so
    /// its completion summary can report it. No-op for an unsequenced batch.
    pub(crate) fn note_episode_trials(&mut self, episode_ids: &[String], trials: &[u64]) {
        for (episode_id, trial) in episode_ids.iter().zip(trials) {
            self.trial_by_episode.insert(episode_id.clone(), *trial);
        }
    }

    pub(crate) fn trial_for_episode(&self, episode_id: &str) -> Option<u64> {
        self.trial_by_episode.get(episode_id).copied()
    }

    /// The position of lane `env_index` in the slot vector, or `None` when the
    /// lane is not part of this route.
    pub(crate) fn slot_position(&self, env_index: u32) -> Option<usize> {
        let env_index = i32::try_from(env_index).ok()?;
        self.slots
            .iter()
            .position(|slot| slot.env_index == env_index)
    }

    /// Claim the next `count` consecutive episode slots. `bounded` refuses the
    /// claim (returning `None`, claiming nothing) once any of them would fall
    /// past `max_episodes`; unbounded always claims.
    pub(crate) fn claim_slots(&mut self, count: usize, bounded: bool) -> Option<Vec<u64>> {
        let first = self.next_slot;
        let last = first + count as u64;
        if bounded && self.max_episodes.is_some_and(|max| last > max) {
            return None;
        }
        self.next_slot = last;
        Some((first..last).collect())
    }

    /// Remember which explicit seed each episode in a reset batch received, so
    /// its completion summary can report it. No-op for an unseeded batch.
    pub(crate) fn note_episode_seeds(&mut self, episode_ids: &[String], seeds: &[i64]) {
        for (episode_id, seed) in episode_ids.iter().zip(seeds) {
            self.seed_by_episode.insert(episode_id.clone(), *seed);
        }
    }

    /// Record one completed episode's summary (completion order) for the
    /// session report.
    pub(crate) fn record_episode_summary(&mut self, summary: EpisodeSummary) {
        self.episode_summaries.push(summary);
    }

    /// Drain the recorded episode summaries into the returned report.
    pub(crate) fn take_episode_summaries(&mut self) -> Vec<EpisodeSummary> {
        std::mem::take(&mut self.episode_summaries)
    }

    pub(crate) fn session_id(&self) -> &str {
        &self.session_id
    }

    pub(crate) fn env_id(&self) -> &str {
        &self.env_id
    }

    pub(crate) fn env_component_id(&self) -> &str {
        &self.env_component_id
    }

    pub(crate) fn model_component_id(&self) -> &str {
        &self.model_component_id
    }

    pub(crate) fn env_context(&self) -> RuntimeEnvContext {
        RuntimeEnvContext {
            env_id: self.env_id.clone(),
            env_component_id: self.env_component_id.clone(),
            model_component_id: self.model_component_id.clone(),
            lane: None,
        }
    }

    /// The route context for a group: the lane is stamped when the group is
    /// one lane of a lane endpoint, so events and telemetry slice per lane.
    pub(crate) fn group_context(&self, positions: &[usize], lane_group: bool) -> RuntimeEnvContext {
        RuntimeEnvContext {
            lane: if lane_group {
                positions
                    .first()
                    .and_then(|&position| self.slots.get(position))
                    .and_then(|slot| u32::try_from(slot.env_index).ok())
            } else {
                None
            },
            ..self.env_context()
        }
    }

    pub(crate) fn total_steps(&self) -> i64 {
        self.total_steps
    }

    pub(crate) fn total_episodes(&self) -> i64 {
        self.total_episodes
    }

    pub(crate) fn next_request_id(&mut self, phase: &str) -> String {
        self.request_seq += 1;
        // env_id is globally unique (UUIDv7), so it alone disambiguates request
        // ids across every adapter; request_seq restarts at 0 per RouteState.
        format!("{}:{}:{:06}", self.env_id, phase, self.request_seq)
    }

    /// The slots at `positions`, in that order.
    pub(crate) fn slots_at(&self, positions: &[usize]) -> Vec<&SlotState> {
        positions
            .iter()
            .filter_map(|&position| self.slots.get(position))
            .collect()
    }

    /// Ordered per-row episode ids for a group — the self-describing batch.
    /// Row `i` belongs to `positions[i]`. Empty string for a lane with no
    /// active episode.
    pub(crate) fn episode_ids_at(&self, positions: &[usize]) -> Vec<String> {
        self.slots_at(positions)
            .into_iter()
            .map(|slot| {
                slot.episode
                    .as_ref()
                    .map(|episode| episode.episode_id.clone())
                    .unwrap_or_default()
            })
            .collect()
    }

    pub(crate) fn snapshot_at(&self, positions: &[usize]) -> RouteSnapshot {
        let slots = self.slots_at(positions);
        let episode_ids = slots
            .iter()
            .map(|slot| {
                slot.episode
                    .as_ref()
                    .map(|episode| episode.episode_id.clone())
                    .unwrap_or_default()
            })
            .collect::<Vec<_>>();
        let episode_record_ids = slots
            .iter()
            .map(|slot| {
                slot.episode
                    .as_ref()
                    .map(|episode| episode.episode_record_id.clone())
                    .unwrap_or_default()
            })
            .collect::<Vec<_>>();
        let primary = slots.first().copied();
        RouteSnapshot {
            episode_id: episode_ids.first().cloned().unwrap_or_default(),
            episode_record_id: episode_record_ids.first().cloned().unwrap_or_default(),
            episode_ids,
            episode_record_ids,
            step: primary.map_or(0, |slot| slot.step),
            env_index: primary.map_or(0, |slot| slot.env_index),
            reset: primary.is_some_and(|slot| slot.reset),
        }
    }

    /// Start one episode per position (a driver-owned reset). `slots` are the
    /// claimed route-global slots, aligned to `positions`; each becomes its
    /// episode's index.
    pub(crate) fn start_episodes_at(
        &mut self,
        positions: &[usize],
        episode_ids: Vec<String>,
        started_from_auto_reset: bool,
        slots: &[u64],
    ) -> Vec<StartedEpisode> {
        let indices: Vec<Option<i64>> = positions
            .iter()
            .enumerate()
            .map(|(i, _)| slots.get(i).map(|slot| *slot as i64 + 1))
            .collect();
        let (record_ids, started) =
            self.records
                .ensure_for_slots(&episode_ids, started_from_auto_reset, &indices);
        self.sync_slots(
            positions,
            episode_ids,
            record_ids,
            true,
            started_from_auto_reset,
        );
        started
            .into_iter()
            .map(|(episode_id, record)| StartedEpisode { episode_id, record })
            .collect()
    }

    /// Observe the ids the env rolled itself to (NEXT_STEP autoreset): lanes
    /// whose id changed start a fresh episode at the given slot; the others keep
    /// their episode. `slots` aligns to `positions` (`None` = not rolling).
    pub(crate) fn observe_episode_ids_at(
        &mut self,
        positions: &[usize],
        episode_ids: Vec<String>,
        slots: &[Option<u64>],
    ) -> Vec<StartedEpisode> {
        let indices: Vec<Option<i64>> = positions
            .iter()
            .enumerate()
            .map(|(i, _)| slots.get(i).copied().flatten().map(|slot| slot as i64 + 1))
            .collect();
        let (record_ids, started) = self.records.ensure_for_slots(&episode_ids, true, &indices);
        self.sync_slots(positions, episode_ids, record_ids, false, true);
        started
            .into_iter()
            .map(|(episode_id, record)| StartedEpisode { episode_id, record })
            .collect()
    }

    /// Advance the group's lanes by one step; `rewards` aligns to `positions`.
    pub(crate) fn record_step_at(&mut self, positions: &[usize], rewards: &[f64]) {
        self.total_steps += 1;
        for (i, &position) in positions.iter().enumerate() {
            if let Some(slot) = self.slots.get_mut(position) {
                slot.step += 1;
                slot.reset = false;
                slot.cumulative_reward += rewards.get(i).copied().unwrap_or(0.0);
            }
        }
    }

    pub(crate) fn complete_episode(&mut self, episode_id: &str) -> Option<EpisodeRecord> {
        self.total_episodes += 1;
        self.records.record_for(episode_id).cloned()
    }

    pub(crate) fn seed_for_episode(&self, episode_id: &str) -> Option<i64> {
        self.seed_by_episode.get(episode_id).copied()
    }

    pub(crate) fn predict_request_at(
        &mut self,
        positions: &[usize],
        observation: Option<Vec<Bytes>>,
        phase: RequestPhase,
    ) -> PredictRequest {
        let episode_info = self
            .episode_ids_at(positions)
            .into_iter()
            .map(|episode_id| {
                let seed = self.seed_for_episode(&episode_id);
                EpisodeInfo { episode_id, seed }
            })
            .collect();
        PredictRequest {
            context: Some(AdapterContext {
                session_id: self.session_id().to_string(),
                env_id: self.env_id().to_string(),
                request_id: self.next_request_id(phase.as_str()),
            }),
            observation: observation.map(leaves_value),
            episode_info,
        }
    }

    pub(crate) fn reset_adapter_request(
        &mut self,
        episode_ids: Vec<String>,
    ) -> ResetAdapterRequest {
        ResetAdapterRequest {
            context: Some(AdapterContext {
                session_id: self.session_id().to_string(),
                env_id: self.env_id().to_string(),
                request_id: self.next_request_id("reset_adapter"),
            }),
            episode_ids,
        }
    }

    pub(crate) fn release_adapter_request(
        &mut self,
        reason: impl Into<String>,
    ) -> ReleaseAdapterRequest {
        ReleaseAdapterRequest {
            context: Some(AdapterContext {
                session_id: self.session_id().to_string(),
                env_id: self.env_id().to_string(),
                request_id: self.next_request_id("release_adapter"),
            }),
            reason: reason.into(),
        }
    }

    fn sync_slots(
        &mut self,
        positions: &[usize],
        episode_ids: Vec<String>,
        record_ids: Vec<String>,
        reset_steps: bool,
        started_from_auto_reset: bool,
    ) {
        for (index, &position) in positions.iter().enumerate() {
            let Some(slot) = self.slots.get_mut(position) else {
                continue;
            };
            let episode_id = episode_ids.get(index).cloned().unwrap_or_default();
            let episode_record_id = record_ids.get(index).cloned().unwrap_or_default();
            // Did this lane's episode id flip? A NEXT_STEP autoreset rolls the id
            // on a single lane at t+1; only that lane's step counter must reset.
            let previous_id = slot
                .episode
                .as_ref()
                .map(|episode| episode.episode_id.clone())
                .unwrap_or_default();
            let rolled = !episode_id.is_empty() && episode_id != previous_id;
            if rolled && !previous_id.is_empty() {
                // The outgoing episode's seed is dropped only when its lane
                // rolls (never at completion emit), so the completion
                // iteration's final predict still reports it.
                self.seed_by_episode.remove(&previous_id);
                self.trial_by_episode.remove(&previous_id);
            }
            slot.episode = if episode_id.is_empty() {
                None
            } else {
                let record = self.records.record_for(&episode_id);
                Some(EpisodeState {
                    episode_id,
                    episode_record_id,
                    episode_index: record.map_or(0, |record| record.index),
                    started_from_auto_reset,
                })
            };
            // `reset_steps` force-resets every lane of the group (a driver-owned
            // reset); `rolled` resets only the lane whose id flipped (autoreset).
            if reset_steps || rolled {
                slot.step = 0;
                slot.reset = true;
                slot.cumulative_reward = 0.0;
                slot.started_at_ns = now_unix_ns();
            }
        }
    }
}

/// Unix time in nanoseconds; saturates at i64::MAX.
pub(crate) fn now_unix_ns() -> i64 {
    std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .map(|d| i64::try_from(d.as_nanos()).unwrap_or(i64::MAX))
        .unwrap_or(0)
}