rlmesh_runtime/

driver.rs

use std::future::Future;
use std::sync::Arc;
use std::time::{Duration, Instant};

use async_trait::async_trait;
use prost::Message;
use rlmesh_proto::common::v1::MessageBytes;
use rlmesh_proto::core::v1::OperationTelemetry;
use rlmesh_proto::env::v1::{
    EpisodeMetadata, ResetRequest, ResetResponse, StepRequest, StepResponse,
};
use rlmesh_proto::model::v1::{CloseRouteRequest, PredictRequest, PredictResponse};
use rlmesh_proto::spaces::v1::SpaceValue;
use tokio_util::sync::CancellationToken;

use crate::episodes::EpisodeRecord;
use crate::hooks::{
    ActionReceivedEvent, EpisodeCompletedEvent, EpisodeStartedEvent, HookError, LogEvent, LogLevel,
    NoopRuntimeHooks, ObservationEmittedEvent, RuntimeHooks, SessionEndedEvent, SessionFailedEvent,
    SessionStartedEvent, StepCompletedEvent, TelemetrySummaryEvent, TelemetryWindowEvent,
};
use crate::route::requests::RequestPhase;
use crate::spec::{RuntimeReport, RuntimeSessionSpec};
use crate::state::{RouteSnapshot, RouteState, StartedEpisode};
use crate::timing::{RuntimeTiming, StepTimingSample};

#[derive(Debug, thiserror::Error)]
pub enum RuntimeError {
    #[error("invalid runtime session spec: {0}")]
    InvalidSpec(String),

    #[error(
        "{operation} timed out on route {route_id} component {component_id} at runtime step {step} after {timeout:?}"
    )]
    OperationTimeout {
        route_id: String,
        component_id: String,
        operation: &'static str,
        step: i64,
        timeout: Duration,
    },

    #[error("route {route_id} cancelled at runtime step {step}: {reason}")]
    RouteCancelled {
        route_id: String,
        step: i64,
        reason: String,
    },

    #[error(
        "environment {operation} failed at runtime step {step}: {message}. If the source is 'transport error: connection closed', the environment server exited, crashed, or received SIGTERM before replying; inspect the environment container logs immediately before the runtime error timestamp"
    )]
    EnvRpc {
        operation: &'static str,
        step: i64,
        message: String,
    },

    #[error("model endpoint {component_id} request failed: {message}")]
    ModelRpc {
        component_id: String,
        message: String,
    },

    #[error(
        "model endpoint {component_id} returned mismatched route identity for request {request_id}"
    )]
    ModelRouteMismatch {
        component_id: String,
        request_id: String,
    },

    #[error("protocol error: {0}")]
    Protocol(String),

    #[error("runtime hook failed: {0}")]
    Hook(HookError),
}

impl RuntimeError {
    pub fn operation_timeout(
        route_id: impl Into<String>,
        component_id: impl Into<String>,
        operation: &'static str,
        step: i64,
        timeout: Duration,
    ) -> Self {
        Self::OperationTimeout {
            route_id: route_id.into(),
            component_id: component_id.into(),
            operation,
            step,
            timeout,
        }
    }

    pub fn route_cancelled(
        route_id: impl Into<String>,
        step: i64,
        reason: impl Into<String>,
    ) -> Self {
        Self::RouteCancelled {
            route_id: route_id.into(),
            step,
            reason: reason.into(),
        }
    }
}

pub struct RuntimeEnvReset {
    pub response: ResetResponse,
    pub telemetry: Option<OperationTelemetry>,
}

pub struct RuntimeEnvStep {
    pub response: StepResponse,
    pub telemetry: Option<OperationTelemetry>,
}

pub struct RuntimeModelPrediction {
    pub response: PredictResponse,
    pub telemetry: Option<OperationTelemetry>,
}

#[async_trait]
pub trait RuntimeEnv: Send {
    async fn reset(&mut self, request: ResetRequest) -> Result<RuntimeEnvReset, RuntimeError>;

    async fn step(&mut self, request: StepRequest) -> Result<RuntimeEnvStep, RuntimeError>;

    async fn close(&mut self, _timeout: Duration) -> Result<(), String> {
        Ok(())
    }
}

#[async_trait]
pub trait RuntimeModel: Send {
    async fn predict(
        &mut self,
        request: PredictRequest,
    ) -> Result<RuntimeModelPrediction, RuntimeError>;

    async fn close_route(
        &mut self,
        _request: CloseRouteRequest,
        _timeout: Duration,
    ) -> Result<(), String> {
        Ok(())
    }
}

pub struct RuntimeDriver<E, M> {
    spec: RuntimeSessionSpec,
    env: E,
    model: M,
    hooks: Arc<dyn RuntimeHooks>,
}

impl<E, M> RuntimeDriver<E, M>
where
    E: RuntimeEnv,
    M: RuntimeModel,
{
    pub fn new(spec: RuntimeSessionSpec, env: E, model: M, hooks: Arc<dyn RuntimeHooks>) -> Self {
        Self {
            spec,
            env,
            model,
            hooks,
        }
    }

    pub fn without_hooks(spec: RuntimeSessionSpec, env: E, model: M) -> Self {
        Self::new(spec, env, model, Arc::new(NoopRuntimeHooks))
    }

    pub async fn run(self) -> Result<RuntimeReport, RuntimeError> {
        self.run_with_cancellation(CancellationToken::new()).await
    }

    pub async fn run_with_cancellation(
        mut self,
        cancellation: CancellationToken,
    ) -> Result<RuntimeReport, RuntimeError> {
        self.spec.validate().map_err(RuntimeError::InvalidSpec)?;
        let mut state = RouteState::new(&self.spec);
        let result = self.run_loop(&mut state, &cancellation).await;
        if let Err(error) = &result {
            self.shutdown_after_failure(&mut state, error).await;
        }
        result
    }

    async fn run_loop(
        &mut self,
        state: &mut RouteState,
        cancellation: &CancellationToken,
    ) -> Result<RuntimeReport, RuntimeError> {
        let mut timings = RuntimeTiming::default();
        self.invoke_session_started(state, &self.spec.env_id).await;

        let reset_started = Instant::now();
        let reset_timeout = self.spec.limits.env_reset_timeout;
        let reset_timeout_ms = self.spec.limits.env_reset_timeout_ms();
        let reset_ok = await_runtime_operation(
            cancellation,
            reset_timeout,
            RuntimeError::operation_timeout(
                state.route_id(),
                state.env_component_id(),
                "env.reset",
                0,
                reset_timeout,
            ),
            self.cancelled_error(state, 0),
            self.env.reset(ResetRequest {
                seeds: vec![],
                options: None,
                timeout_ms: reset_timeout_ms,
            }),
        )
        .await?;
        let reset_latency = reset_started.elapsed();
        timings.reset.record(reset_latency);
        timings
            .window
            .record_operation_telemetry(state.env_component_id(), reset_ok.telemetry.as_ref());
        self.invoke_log(
            state,
            LogLevel::Info,
            format!(
                "env reset complete in {:.0}ms ({} episode(s) ready)",
                reset_latency.as_secs_f64() * 1000.0,
                reset_ok
                    .response
                    .episode_ids
                    .iter()
                    .filter(|value| !value.is_empty())
                    .count()
            ),
        )
        .await;

        let reset_observation = value_bytes(reset_ok.response.observation.as_ref())?;
        let started_episodes = state.start_episodes(reset_ok.response.episode_ids, false);
        self.invoke_started_episodes(state, started_episodes).await;

        let mut reset_msg =
            state.predict_request(reset_observation.clone(), RequestPhase::ResetObservation);
        let mut reset_event =
            self.observation_event(state, state.snapshot(), true, reset_observation.clone());
        let transformed_reset_observation = self
            .invoke_transform_observation(reset_event.clone())
            .await?;
        reset_event.observation = transformed_reset_observation.clone();
        reset_msg.observation = transformed_reset_observation.map(bytes_value);
        self.invoke_observation_emitted(reset_event).await;

        let mut pending_observation_msg = reset_msg;

        loop {
            if cancellation.is_cancelled() {
                return Err(self.cancelled_error(state, state.snapshot().step));
            }

            let predict_snapshot = state.snapshot();
            let model_wait_started = Instant::now();
            let predict_timeout = self.spec.limits.model_predict_timeout;
            let expected_context = pending_observation_msg.context.clone();
            let action_msg = await_runtime_operation(
                cancellation,
                predict_timeout,
                RuntimeError::operation_timeout(
                    state.route_id(),
                    state.model_component_id(),
                    "model.predict",
                    predict_snapshot.step,
                    predict_timeout,
                ),
                self.cancelled_error(state, predict_snapshot.step),
                self.model.predict(pending_observation_msg),
            )
            .await?;
            if action_msg.response.context != expected_context {
                let request_id = expected_context
                    .as_ref()
                    .map(|context| context.request_id.clone())
                    .unwrap_or_default();
                return Err(RuntimeError::ModelRouteMismatch {
                    component_id: state.model_component_id().to_string(),
                    request_id,
                });
            }
            let model_action = value_bytes(action_msg.response.action.as_ref())?;
            let model_wait_latency = model_wait_started.elapsed();
            timings.model_wait.record(model_wait_latency);
            timings.window.record_operation_telemetry(
                state.model_component_id(),
                action_msg.telemetry.as_ref(),
            );

            let action_step = predict_snapshot.step + 1;
            let mut action_event = ActionReceivedEvent {
                session_id: state.session_id().to_string(),
                route: state.route_context(),
                episode_id: predict_snapshot.episode_id.clone(),
                episode_record_id: predict_snapshot.episode_record_id.clone(),
                episode_ids: predict_snapshot.episode_ids.clone(),
                episode_record_ids: predict_snapshot.episode_record_ids.clone(),
                step: action_step,
                env_index: predict_snapshot.env_index,
                action_space: self.spec.action_space().clone(),
                action: model_action,
            };
            action_event.action = self.invoke_transform_action(action_event.clone()).await?;
            let request_bytes = action_event
                .action
                .as_ref()
                .map(|action| action.data.len())
                .unwrap_or_default();
            self.invoke_action_received(action_event.clone()).await;

            let env_step_started = Instant::now();
            let step_timeout = self.spec.limits.env_step_timeout;
            let step_timeout_ms = self.spec.limits.env_step_timeout_ms();
            let step_ok = await_runtime_operation(
                cancellation,
                step_timeout,
                RuntimeError::operation_timeout(
                    state.route_id(),
                    state.env_component_id(),
                    "env.step",
                    action_step,
                    step_timeout,
                ),
                self.cancelled_error(state, action_step),
                self.env.step(StepRequest {
                    action: action_event.action.map(bytes_value),
                    timeout_ms: step_timeout_ms,
                }),
            )
            .await?;
            let env_step_latency = env_step_started.elapsed();
            let step_observation = value_bytes(step_ok.response.observation.as_ref())?;
            timings.env_step.record(env_step_latency);
            timings
                .window
                .record_operation_telemetry(state.env_component_id(), step_ok.telemetry.as_ref());
            let response_bytes = step_observation
                .as_ref()
                .map(|obs| obs.data.len())
                .unwrap_or_default()
                + step_ok
                    .response
                    .infos
                    .as_ref()
                    .map(Message::encoded_len)
                    .unwrap_or_default();
            timings.window.record_step(StepTimingSample {
                model_wait: model_wait_latency,
                env_step: env_step_latency,
                request_bytes,
                response_bytes,
                env_component_id: state.env_component_id(),
                model_component_id: state.model_component_id(),
            });

            state.record_step();
            let step_snapshot = state.snapshot();
            self.invoke_step_completed(StepCompletedEvent {
                session_id: state.session_id().to_string(),
                route: state.route_context(),
                episode_id: step_snapshot.episode_id.clone(),
                episode_record_id: step_snapshot.episode_record_id.clone(),
                step: step_snapshot.step,
                env_index: step_snapshot.env_index,
                rewards: step_ok.response.rewards.clone(),
            })
            .await;

            if let Some(event) = timings.maybe_emit_window(
                state.session_id(),
                state.route_context(),
                self.spec.limits.telemetry_window,
            ) {
                self.invoke_telemetry_window(event).await;
            }

            let observation_snapshot = if step_ok.response.episode_ids.is_empty() {
                state.snapshot()
            } else {
                let started_episodes = state.observe_episode_ids(step_ok.response.episode_ids);
                self.invoke_started_episodes(state, started_episodes).await;
                state.snapshot()
            };
            self.invoke_observation_emitted(self.observation_event(
                state,
                observation_snapshot,
                false,
                step_observation.clone(),
            ))
            .await;

            self.emit_completed_episodes(state, &step_ok.response.completed_episodes)
                .await;

            if self
                .spec
                .max_episodes
                .is_some_and(|limit| state.total_episodes() >= limit as i64)
            {
                if let Some(event) = timings.flush_window(state.session_id(), state.route_context())
                {
                    self.invoke_telemetry_window(event).await;
                }
                if let Some(event) =
                    timings.telemetry_summary(state.session_id(), state.route_context())
                {
                    self.invoke_telemetry_summary(event).await;
                }
                let close_request = state.close_route_request("completed requested episodes");
                self.shutdown_terminal_route(state, "completed requested episodes", close_request)
                    .await;
                self.invoke_session_ended(
                    state,
                    "completed requested episodes",
                    state.total_steps(),
                    state.total_episodes(),
                )
                .await;
                timings.log_summary(state.total_steps(), state.total_episodes());
                return Ok(RuntimeReport {
                    session_id: state.session_id().to_string(),
                    route_id: self.spec.route_id.clone(),
                    total_steps: state.total_steps(),
                    total_episodes: state.total_episodes(),
                });
            }

            let need_reset = !step_ok.response.completed_episodes.is_empty();
            let (next_obs, phase, is_reset_msg) = if need_reset {
                let reset_started = Instant::now();
                let step = state.snapshot().step;
                let reset_timeout = self.spec.limits.env_reset_timeout;
                let reset_timeout_ms = self.spec.limits.env_reset_timeout_ms();
                let reset_ok = await_runtime_operation(
                    cancellation,
                    reset_timeout,
                    RuntimeError::operation_timeout(
                        state.route_id(),
                        state.env_component_id(),
                        "env.reset",
                        step,
                        reset_timeout,
                    ),
                    self.cancelled_error(state, step),
                    self.env.reset(ResetRequest {
                        seeds: vec![],
                        options: None,
                        timeout_ms: reset_timeout_ms,
                    }),
                )
                .await?;
                timings.reset.record(reset_started.elapsed());
                timings.window.record_operation_telemetry(
                    state.env_component_id(),
                    reset_ok.telemetry.as_ref(),
                );
                let next_obs = value_bytes(reset_ok.response.observation.as_ref())?;
                let started_episodes = state.start_episodes(reset_ok.response.episode_ids, true);
                self.invoke_started_episodes(state, started_episodes).await;
                (next_obs, RequestPhase::ResetObservation, true)
            } else {
                (
                    step_observation.clone(),
                    RequestPhase::StepObservation,
                    false,
                )
            };

            let mut obs_msg = state.predict_request(next_obs.clone(), phase);
            let mut outgoing_observation_event =
                self.observation_event(state, state.snapshot(), is_reset_msg, next_obs);
            let transformed_observation = self
                .invoke_transform_observation(outgoing_observation_event.clone())
                .await?;
            outgoing_observation_event.observation = transformed_observation.clone();
            obs_msg.observation = transformed_observation.map(bytes_value);
            if is_reset_msg {
                self.invoke_observation_emitted(outgoing_observation_event)
                    .await;
            }

            pending_observation_msg = obs_msg;
        }
    }

    async fn shutdown_after_failure(&mut self, state: &mut RouteState, error: &RuntimeError) {
        let reason = error.to_string();
        let request = state.close_route_request(reason.clone());
        self.shutdown_terminal_route(state, &reason, request).await;

        if let Err(err) = self
            .hooks
            .session_failed(SessionFailedEvent {
                session_id: state.session_id().to_string(),
                route: state.route_context(),
                reason,
            })
            .await
        {
            tracing::warn!("runtime hook session_failed failed: {err}");
        }
    }

    async fn shutdown_terminal_route(
        &mut self,
        state: &RouteState,
        reason: &str,
        request: CloseRouteRequest,
    ) {
        let timeout = self.spec.limits.service_close_timeout;
        let model_close = self.model.close_route(request, timeout);
        if self.spec.close_env_on_end {
            let env_close = self.env.close(timeout);
            let (env_result, model_result) = tokio::join!(env_close, model_close);
            if let Err(err) = env_result {
                tracing::warn!(error = %err, "environment close failed during route shutdown");
            }
            if let Err(err) = model_result {
                tracing::warn!(
                    error = %err,
                    reason,
                    "model route close failed during route shutdown; relying on owner shutdown"
                );
            }
            return;
        }

        tracing::debug!(
            route_id = %state.route_id(),
            reason,
            "skipping environment close for route; endpoint remains owned by the run"
        );
        if let Err(err) = model_close.await {
            tracing::warn!(
                error = %err,
                reason,
                "model route close failed during route shutdown; relying on owner shutdown"
            );
        }
    }

    fn cancelled_error(&self, state: &RouteState, step: i64) -> RuntimeError {
        RuntimeError::route_cancelled(
            state.route_id(),
            step,
            "cancelled after sibling route failure",
        )
    }

    async fn invoke_session_started(&self, state: &RouteState, env_id: &str) {
        if let Err(err) = self
            .hooks
            .session_started(SessionStartedEvent {
                session_id: state.session_id().to_string(),
                route: state.route_context(),
                env_id: env_id.to_string(),
            })
            .await
        {
            tracing::warn!("runtime hook session_started failed: {err}");
        }
    }

    async fn invoke_started_episodes(&self, state: &RouteState, episodes: Vec<StartedEpisode>) {
        for episode in episodes {
            self.invoke_episode_started(state, &episode.episode_id, &episode.record)
                .await;
        }
    }

    async fn invoke_episode_started(
        &self,
        state: &RouteState,
        episode_id: &str,
        record: &EpisodeRecord,
    ) {
        if let Err(err) = self
            .hooks
            .episode_started(EpisodeStartedEvent {
                session_id: state.session_id().to_string(),
                route: state.route_context(),
                episode_id: episode_id.to_string(),
                episode_record_id: record.record_id.clone(),
                episode_index: record.index,
                env_index: record.env_index,
                started_from_auto_reset: record.started_from_auto_reset,
            })
            .await
        {
            tracing::warn!("runtime hook episode_started failed: {err}");
        }
    }

    async fn emit_completed_episodes(&self, state: &mut RouteState, episodes: &[EpisodeMetadata]) {
        for completed in episodes {
            let record = state.complete_episode(&completed.episode_id);
            self.invoke_episode_completed(EpisodeCompletedEvent {
                session_id: state.session_id().to_string(),
                route: state.route_context(),
                episode_id: completed.episode_id.clone(),
                episode_record_id: record
                    .as_ref()
                    .map(|record| record.record_id.clone())
                    .unwrap_or_default(),
                episode_index: record.as_ref().map_or(0, |record| record.index),
                env_index: completed.env_index,
                step_count: completed.step_count,
                cumulative_reward: completed.cumulative_reward,
                terminated: completed.terminated,
                truncated: completed.truncated,
                duration_ms: completed.duration_ms,
                final_info: completed.final_info.clone(),
            })
            .await;
        }
    }

    async fn invoke_episode_completed(&self, event: EpisodeCompletedEvent) {
        if let Err(err) = self.hooks.episode_completed(event).await {
            tracing::warn!("runtime hook episode_completed failed: {err}");
        }
    }

    async fn invoke_action_received(&self, event: ActionReceivedEvent) {
        if let Err(err) = self.hooks.action_received(event).await {
            tracing::warn!("runtime hook action_received failed: {err}");
        }
    }

    async fn invoke_transform_action(
        &self,
        event: ActionReceivedEvent,
    ) -> Result<Option<MessageBytes>, RuntimeError> {
        match self.hooks.transform_action(event).await {
            Ok(action) => Ok(action),
            Err(err) => {
                tracing::warn!("runtime hook transform_action failed: {err}");
                Err(RuntimeError::Hook(err))
            }
        }
    }

    async fn invoke_step_completed(&self, event: StepCompletedEvent) {
        if let Err(err) = self.hooks.step_completed(event).await {
            tracing::warn!("runtime hook step_completed failed: {err}");
        }
    }

    async fn invoke_observation_emitted(&self, event: ObservationEmittedEvent) {
        if let Err(err) = self.hooks.observation_emitted(event).await {
            tracing::warn!("runtime hook observation_emitted failed: {err}");
        }
    }

    async fn invoke_transform_observation(
        &self,
        event: ObservationEmittedEvent,
    ) -> Result<Option<MessageBytes>, RuntimeError> {
        match self.hooks.transform_observation(event).await {
            Ok(observation) => Ok(observation),
            Err(err) => {
                tracing::warn!("runtime hook transform_observation failed: {err}");
                Err(RuntimeError::Hook(err))
            }
        }
    }

    async fn invoke_telemetry_window(&self, event: TelemetryWindowEvent) {
        if let Err(err) = self.hooks.telemetry_window(event).await {
            tracing::warn!("runtime hook telemetry_window failed: {err}");
        }
    }

    async fn invoke_telemetry_summary(&self, event: TelemetrySummaryEvent) {
        if let Err(err) = self.hooks.telemetry_summary(event).await {
            tracing::warn!("runtime hook telemetry_summary failed: {err}");
        }
    }

    async fn invoke_log(&self, state: &RouteState, level: LogLevel, message: impl Into<String>) {
        if let Err(err) = self
            .hooks
            .log(LogEvent {
                session_id: state.session_id().to_string(),
                route: state.route_context(),
                level,
                message: message.into(),
                source: Some("runtime".to_string()),
            })
            .await
        {
            tracing::warn!("runtime hook log failed: {err}");
        }
    }

    async fn invoke_session_ended(
        &self,
        state: &RouteState,
        reason: &str,
        total_steps: i64,
        total_episodes: i64,
    ) {
        if let Err(err) = self
            .hooks
            .session_ended(SessionEndedEvent {
                session_id: state.session_id().to_string(),
                route: state.route_context(),
                reason: reason.to_string(),
                total_steps,
                total_episodes,
            })
            .await
        {
            tracing::warn!("runtime hook session_ended failed: {err}");
        }
    }

    fn observation_event(
        &self,
        state: &RouteState,
        snapshot: RouteSnapshot,
        is_reset: bool,
        observation: Option<MessageBytes>,
    ) -> ObservationEmittedEvent {
        ObservationEmittedEvent {
            session_id: state.session_id().to_string(),
            route: state.route_context(),
            episode_id: snapshot.episode_id,
            episode_record_id: snapshot.episode_record_id,
            episode_ids: snapshot.episode_ids,
            episode_record_ids: snapshot.episode_record_ids,
            step: snapshot.step,
            env_index: snapshot.env_index,
            is_reset,
            num_envs: self.spec.num_envs as u32,
            observation_space: self.spec.observation_space().clone(),
            observation,
        }
    }
}

async fn await_runtime_operation<T, F>(
    cancellation: &CancellationToken,
    timeout: Duration,
    timeout_error: RuntimeError,
    cancelled_error: RuntimeError,
    operation: F,
) -> Result<T, RuntimeError>
where
    F: Future<Output = Result<T, RuntimeError>>,
{
    tokio::select! {
        _ = cancellation.cancelled() => Err(cancelled_error),
        result = tokio::time::timeout(timeout, operation) => match result {
            Ok(result) => result,
            Err(_) => Err(timeout_error),
        },
    }
}

fn bytes_value(value: MessageBytes) -> SpaceValue {
    SpaceValue { bytes: Some(value) }
}

fn value_bytes(payload: Option<&SpaceValue>) -> Result<Option<MessageBytes>, RuntimeError> {
    let Some(payload) = payload else {
        return Ok(None);
    };
    Ok(payload.bytes.clone())
}