meerkat 0.8.16

//! Runtime-state observers and cleanup primitives.
//!
//! Populated by W1-E (`PendingSessionEventStreams`,
//! `PendingSessionEventStreamDrop`), W2-E (`SessionInfo`,
//! `SessionState` enums + state observers), and W3-A (skill-identity
//! registry plumbing + `ArchiveRuntimeCleanup`).
//!
//! `ArchiveRuntimeCleanup` lives here in trait-shaped form: the MCP
//! and Mob cleanup hooks are abstracted behind
//! [`ArchiveRuntimeMcpState`] and [`ArchiveRuntimeMobState`] so the
//! struct does not need to depend on `meerkat-mob-mcp` (which the
//! `meerkat` facade does not pull in) or know about the RPC-private
//! `SessionMcpState`. Surfaces wire their own implementations.

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

use async_trait::async_trait;
use meerkat_core::EventEnvelope;
use meerkat_core::event::AgentEvent;
use meerkat_core::skills::{SkillError, SourceIdentityRegistry};
use meerkat_core::types::SessionId;
use tokio::sync::{Mutex, Notify, broadcast};

/// Observable lifecycle state of a session.
///
/// Surfaces project this onto their wire format (`session/list`,
/// `GET /sessions/<id>`, …). The serde encoding mirrors the
/// long-standing JSON-RPC contract — surfaces depending on the wire
/// strings (`"idle"`, `"running"`, `"shutting_down"`) MUST not break.
#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum SessionState {
    /// The session is idle and ready to accept a new turn.
    Idle,
    /// A turn is currently running.
    Running,
    /// The session is shutting down.
    ShuttingDown,
}

impl SessionState {
    /// Return a stable string representation matching the serde
    /// `rename_all` convention. Used by surfaces that bypass serde
    /// (e.g. tracing tags, telemetry counters) and need the canonical
    /// lowercase slug.
    #[must_use]
    pub fn as_str(self) -> &'static str {
        match self {
            Self::Idle => "idle",
            Self::Running => "running",
            Self::ShuttingDown => "shutting_down",
        }
    }
}

/// Summary information about a session: identity, lifecycle marker,
/// and durable labels. Surfaces hand this back from their list / read
/// endpoints; the wire encoding mirrors the canonical contract.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SessionInfo {
    /// The session id this summary describes.
    pub session_id: SessionId,
    /// Lifecycle marker.
    pub state: SessionState,
    /// Durable labels stored alongside the session.
    pub labels: BTreeMap<String, String>,
}

/// Per-session pending event-stream channel pair.
///
/// While a session is staged or archived but still emitting tail
/// events, the runtime keeps a broadcast sender alive so subscribers
/// drain remaining envelopes; `receiver_dropped` notifies the runtime
/// when the last subscriber drops so the channel can be reaped.
#[derive(Clone)]
pub struct PendingSessionEventStreams {
    /// Broadcast sender that fans out [`AgentEvent`] envelopes to live
    /// subscribers.
    pub events: broadcast::Sender<EventEnvelope<AgentEvent>>,
    /// Notifier that fires when the receiver-side handle is dropped, so
    /// the runtime can prune the entry from its event-stream map.
    pub receiver_dropped: Arc<Notify>,
}

/// RAII drop guard that fires `receiver_dropped` exactly once when the
/// receiver-side handle is dropped. Paired with
/// [`PendingSessionEventStreams::receiver_dropped`].
pub struct PendingSessionEventStreamDrop {
    /// Notifier shared with the corresponding
    /// [`PendingSessionEventStreams`] entry.
    pub receiver_dropped: Arc<Notify>,
}

impl Drop for PendingSessionEventStreamDrop {
    fn drop(&mut self) {
        self.receiver_dropped.notify_one();
    }
}

/// Skill-identity registry slot held by the session runtime.
///
/// `generation` is a monotone version stamp that lets the runtime guard
/// against stale writes after a config reload races a still-in-flight
/// builder. Surfaces inject the latest registry into
/// `ContextWindowBuilder` per-session via the runtime accessor.
#[derive(Clone, Default)]
pub struct SkillIdentityRegistryState {
    /// Monotonic generation counter; writes only land if their
    /// generation is `>=` the currently-stored generation.
    pub generation: u64,
    /// The active registry projected onto agents.
    pub registry: SourceIdentityRegistry,
}

/// Build a `SourceIdentityRegistry` from runtime config.
///
/// Surface-agnostic: surfaces (RPC, REST, CLI, embedded examples) call
/// this from their config-load path and feed the result into a
/// runtime accessor. Identical body to the original RPC-side helper.
#[allow(clippy::missing_errors_doc)]
pub fn build_skill_identity_registry(
    config: &meerkat_core::Config,
    context_root: Option<&std::path::Path>,
    user_root: Option<&std::path::Path>,
) -> Result<SourceIdentityRegistry, SkillError> {
    #[cfg(not(target_arch = "wasm32"))]
    {
        let _ = (context_root, user_root);
        config.skills.build_source_identity_registry()
    }
    #[cfg(target_arch = "wasm32")]
    {
        let _ = (context_root, user_root);
        config.skills.build_source_identity_registry()
    }
}

/// Per-surface MCP cleanup hook used by [`ArchiveRuntimeCleanup`].
///
/// Surfaces that own MCP session adapters (e.g. `meerkat-rpc`) implement
/// this to shutdown their adapter when a session is archived. The
/// trait is intentionally narrow: `cleanup` is the only side-effect the
/// archive flow needs.
#[async_trait]
pub trait ArchiveRuntimeMcpState: Send + Sync {
    /// Tear down the MCP adapter (and any per-session lifecycle plumbing)
    /// for `session_id`. Implementations are expected to be idempotent —
    /// a missing entry is a no-op.
    async fn cleanup(&self, session_id: &SessionId);
}

/// Per-surface Mob cleanup hook used by [`ArchiveRuntimeCleanup`].
///
/// `meerkat-rpc` (the only surface today that wires mob orchestration)
/// implements this on top of `meerkat_mob_mcp::MobMcpState`. The trait
/// keeps `meerkat` from depending on `meerkat-mob-mcp`.
#[async_trait]
pub trait ArchiveRuntimeMobState: Send + Sync {
    /// Drop session-scoped mob state. The result is propagated by
    /// [`ArchiveRuntimeCleanup::run`].
    async fn cleanup(
        &self,
        session_id: &SessionId,
    ) -> Result<(), meerkat_core::service::SessionError>;
}

/// Surface-agnostic archive cleanup orchestrator.
///
/// Composes the runtime adapter, the per-session pending event-stream
/// map, and the optional MCP / Mob cleanup hooks. Surfaces build one of
/// these per archive call and either invoke `archive_service` followed
/// by `run`, or just `run` when the service step has already happened.
#[derive(Clone)]
pub struct ArchiveRuntimeCleanup {
    /// Runtime adapter; used to unregister the session and abort comms
    /// drain.
    pub runtime_adapter: Arc<meerkat_runtime::MeerkatMachine>,
    /// Per-session pending event streams. `None` for surfaces that
    /// don't track them (e.g. mob session service flows).
    pub pending_session_event_streams:
        Option<Arc<Mutex<HashMap<SessionId, PendingSessionEventStreams>>>>,
    /// Optional MCP cleanup hook; surfaces implement
    /// [`ArchiveRuntimeMcpState`] for their adapter map.
    pub mcp_state: Option<Arc<dyn ArchiveRuntimeMcpState>>,
    /// Optional Mob cleanup hook; surfaces implement
    /// [`ArchiveRuntimeMobState`] for their `MobMcpState` wrapper.
    pub mob_state: Option<Arc<dyn ArchiveRuntimeMobState>>,
}

impl ArchiveRuntimeCleanup {
    fn runtime_cleanup_error(error: impl std::fmt::Display) -> meerkat_core::service::SessionError {
        meerkat_core::service::SessionError::Agent(meerkat_core::error::AgentError::InternalError(
            error.to_string(),
        ))
    }

    /// Run the durable-store archive step. Surfaces that already
    /// archived through the service skip this and call [`run`]
    /// directly.
    #[cfg(all(feature = "session-store", not(target_arch = "wasm32")))]
    pub async fn archive_service(
        &self,
        service: &crate::PersistentSessionService<crate::service_factory::FactoryAgentBuilder>,
        session_id: &SessionId,
    ) -> Result<(), meerkat_core::service::SessionError> {
        service
            .archive_with_machine_protocol(
                session_id,
                meerkat_session::MachineSessionArchiveProtocol::from_machine(
                    self.runtime_adapter.as_ref(),
                ),
            )
            .await
    }

    /// Run the per-surface cleanup steps that follow a successful
    /// archive: unregister from the runtime adapter, drop pending event
    /// streams, tear down MCP adapters, destroy mob state, abort comms
    /// drain.
    pub async fn run(
        &self,
        session_id: &SessionId,
    ) -> Result<(), meerkat_core::service::SessionError> {
        self.runtime_adapter
            .unregister_session(session_id)
            .await
            .map_err(Self::runtime_cleanup_error)?;
        self.run_after_runtime_unregistered(session_id).await
    }

    /// Run terminal cleanup when generated completion authority proves that
    /// the runtime was terminated.
    ///
    /// Only the machine-owned saga removes runtime registration. This entry
    /// point covers the saga-won ordering after the executor's external-only
    /// cleanup: both the typed, session-bound observation and an already-absent
    /// registration are required, so the completion relay never treats bare
    /// registry absence as terminal proof.
    pub async fn run_after_runtime_termination(
        &self,
        session_id: &SessionId,
        cleanup_observation: &meerkat_runtime::CompletionCleanupObservation,
    ) -> Result<(), meerkat_core::service::SessionError> {
        if !cleanup_observation.proves_runtime_termination_for(session_id) {
            return Err(Self::runtime_cleanup_error(format!(
                "runtime termination cleanup for session {session_id} lacks machine-owned termination proof"
            )));
        }

        if self.runtime_adapter.contains_session(session_id).await {
            return Err(Self::runtime_cleanup_error(format!(
                "runtime termination cleanup for session {session_id} expected an already-absent registration"
            )));
        }

        self.run_after_runtime_unregistered(session_id).await
    }

    async fn run_after_runtime_unregistered(
        &self,
        session_id: &SessionId,
    ) -> Result<(), meerkat_core::service::SessionError> {
        if let Some(streams) = self.pending_session_event_streams.as_ref() {
            streams.lock().await.remove(session_id);
        }
        if let Some(mcp_state) = self.mcp_state.as_ref() {
            mcp_state.cleanup(session_id).await;
        }
        if let Some(mob_state) = self.mob_state.as_ref() {
            mob_state.cleanup(session_id).await?;
        }
        #[cfg(feature = "comms")]
        match self.runtime_adapter.abort_comms_drain(session_id).await {
            Ok(()) => {}
            // The session was unregistered above, so the machine may
            // legitimately report the runtime as absent or already terminal;
            // both verdicts prove no drain remains running, which is exactly
            // the post-condition this cleanup wants.
            Err(
                meerkat_runtime::RuntimeDriverError::NotFound { .. }
                | meerkat_runtime::RuntimeDriverError::Destroyed
                | meerkat_runtime::RuntimeDriverError::NotReady { .. },
            ) => {}
            Err(error) => {
                return Err(meerkat_core::service::SessionError::Agent(
                    meerkat_core::error::AgentError::InternalError(error.to_string()),
                ));
            }
        }
        Ok(())
    }
}

/// `RuntimeStateOps` orchestrator (gated on `session-store`).
///
/// Owns the surface-agnostic session-state observers ([`discard_live_session`],
/// [`discard_stale_live_session`], [`live_session_is_stale`]). Surfaces
/// build one per call from their own SessionRuntime borrows.
///
/// RPC-facing archived-session rejection and `try_recover_persisted_session`
/// remain in `meerkat-rpc`: they compose RPC-private `TurnOverrides`,
/// `RpcError`, and surface policy around these shared state operations.
#[cfg(all(feature = "session-store", not(target_arch = "wasm32")))]
mod ops {
    use std::sync::Arc;

    use meerkat_core::service::SessionError;
    use meerkat_core::types::SessionId;
    use meerkat_runtime::MeerkatMachine;

    use crate::PersistentSessionService;
    use crate::service_factory::FactoryAgentBuilder;
    use crate::session_runtime::admission::{
        StagedCapacityAdmissions, discard_staged_capacity_admission,
    };
    use crate::{StagedSessionRegistry, session_runtime::recovery::RecoveryContext};

    /// Surface-agnostic session-state observers shared across surfaces.
    pub struct RuntimeStateOps<'a> {
        /// Persistent session service.
        pub service: &'a Arc<PersistentSessionService<FactoryAgentBuilder>>,
        /// Staged session registry.
        pub staged_sessions: &'a Arc<StagedSessionRegistry>,
        /// Staged capacity ledger; consumed when discarding a live session
        /// that is not currently staged so capacity returns to the pool.
        pub staged_capacity_admissions: &'a StagedCapacityAdmissions,
        /// Runtime adapter.
        pub runtime_adapter: &'a Arc<MeerkatMachine>,
    }

    impl RuntimeStateOps<'_> {
        /// Discard a live session. If the session is not currently
        /// staged, the staged-capacity admission is released back to
        /// the pool.
        pub async fn discard_live_session(
            &self,
            session_id: &SessionId,
        ) -> Result<(), SessionError> {
            let result = self.service.discard_live_session(session_id).await;
            if result.is_ok() && !self.staged_sessions.contains(session_id).await {
                discard_staged_capacity_admission(self.staged_capacity_admissions, session_id);
            }
            result
        }

        /// Recovery-gate-ordered service cleanup for the exact machine-owned
        /// post-stop unregister window.
        pub async fn discard_live_session_after_runtime_stop_terminalized(
            &self,
            session_id: &SessionId,
        ) -> Result<(), SessionError> {
            let result = self
                .service
                .discard_live_session_after_runtime_stop_terminalized(session_id)
                .await;
            if result.is_ok() && !self.staged_sessions.contains(session_id).await {
                discard_staged_capacity_admission(self.staged_capacity_admissions, session_id);
            }
            result
        }

        pub async fn discard_live_session_after_runtime_stop_terminalized_under_runtime_turn_boundary(
            &self,
            session_id: &SessionId,
        ) -> Result<(), SessionError> {
            let result = self
                .service
                .discard_live_session_after_runtime_stop_terminalized_under_runtime_turn_boundary(
                    session_id,
                )
                .await;
            if result.is_ok() && !self.staged_sessions.contains(session_id).await {
                discard_staged_capacity_admission(self.staged_capacity_admissions, session_id);
            }
            result
        }

        /// Discard a stale live session and unregister it from the
        /// runtime adapter.
        pub async fn discard_stale_live_session(
            &self,
            session_id: &SessionId,
        ) -> Result<(), SessionError> {
            let discard_error = match self.discard_live_session(session_id).await {
                Ok(()) | Err(SessionError::NotFound { .. }) => None,
                Err(error) => Some(error),
            };
            let unregister_error = self
                .runtime_adapter
                .unregister_session(session_id)
                .await
                .err();
            match (discard_error, unregister_error) {
                (None, None) => Ok(()),
                (Some(error), None) => Err(error),
                (None, Some(error)) => Err(SessionError::Agent(
                    meerkat_core::error::AgentError::InternalError(error.to_string()),
                )),
                (Some(primary), Some(cleanup)) => Err(SessionError::Agent(
                    meerkat_core::error::AgentError::InternalError(format!(
                        "{primary}; additionally failed to unregister stale runtime session {session_id}: {cleanup}"
                    )),
                )),
            }
        }

        /// Discard a stale live projection while the runtime loop already
        /// owns the non-reentrant turn-finalization boundary. This path keeps
        /// the runtime registration and its current executor intact so the
        /// same apply can rematerialize the actor without self-unregistering.
        pub async fn discard_stale_live_session_under_runtime_turn_boundary(
            &self,
            session_id: &SessionId,
        ) -> Result<(), SessionError> {
            let result = match self
                .service
                .discard_live_session_under_runtime_turn_boundary(session_id)
                .await
            {
                Ok(()) | Err(SessionError::NotFound { .. }) => Ok(()),
                Err(error) => Err(error),
            };
            if result.is_ok() && !self.staged_sessions.contains(session_id).await {
                discard_staged_capacity_admission(self.staged_capacity_admissions, session_id);
            }
            result
        }

        /// Determine whether the live projection for `session_id` has
        /// fallen behind the durable authoritative snapshot. Returns
        /// `Ok(false)` once the synchronization shortcut has fired or
        /// the live snapshot already mirrors the durable record.
        ///
        /// `recovery_ctx` provides the
        /// [`RecoveryContext::load_persisted_session`] flow used to
        /// cross-check the durable snapshot — surfaces pass their own
        /// recovery wiring so this observer does not embed the lookup.
        pub async fn live_session_is_stale(
            &self,
            session_id: &SessionId,
            recovery_ctx: &RecoveryContext<'_>,
        ) -> Result<bool, SessionError> {
            if self
                .service
                .synchronize_live_session_from_durable_authority_if_needed(session_id)
                .await?
            {
                return Ok(false);
            }

            let live = match self.service.export_live_session(session_id).await {
                Ok(session) => session,
                Err(SessionError::NotFound { .. }) => {
                    return Ok(recovery_ctx
                        .load_persisted_session(session_id)
                        .await?
                        .is_some());
                }
                Err(err) => return Err(err),
            };
            let Some(stored) = recovery_ctx.load_persisted_session(session_id).await? else {
                return Ok(false);
            };
            Ok(stored.messages().len() > live.messages().len())
        }
    }
}

#[cfg(all(feature = "session-store", not(target_arch = "wasm32")))]
pub use ops::RuntimeStateOps;