aion_server/assistant/sessions/

registry.rs

//! The registry: every assistant session this server knows about, live or not.
//!
//! # Two halves, one answer
//!
//! The DURABLE half is the store: a record per session and a transcript per
//! session, both surviving any number of restarts. The LIVE half is a map of
//! harness processes, which does not survive one. A session's state is the join
//! of the two, projected — never a field either half stores.
//!
//! # The boot sweep writes back
//!
//! A record whose process is gone is settled at boot by an APPENDED record
//! saying so, with its cause: `dormant` when the agent advertised `loadSession`
//! (its own storage can reopen the conversation) and `ended` when it did not.
//! Written back rather than merely displayed, because a listing that showed
//! "ended" while the record said nothing would be a projection with no durable
//! answer behind it — and because the resume decision has to be readable by
//! whoever asks next, not recomputed identically in three places.

use std::sync::Arc;

use aion_core::{
    AssistantSessionEvent, AssistantSessionFrame, AssistantSessionId, AssistantSessionProjection,
    AssistantSessionState, AssistantSessionSummary, AssistantTurnContext, TITLE_CHARACTERS,
};
use aion_integration_acp::catalogue::CatalogueHarness;
use aion_store::assistant::{
    AssistantSessionRecord, AssistantSessionStore, AssistantTranscriptEvent,
};
use chrono::Utc;
use dashmap::DashMap;
use tokio::sync::broadcast;

use crate::config::ResolvedAssistantConfig;

use super::error::AssistantSessionError;
use super::launch::AssistantEndpoints;
use super::live::{LiveSession, Recorder};

/// Every assistant session this server knows about.
#[derive(Clone)]
pub struct AssistantSessions {
    inner: Arc<Inner>,
}

struct Inner {
    live: DashMap<AssistantSessionId, Arc<LiveSession>>,
    /// One lock per session, serialising SPAWNS on it.
    ///
    /// Per session rather than one lock for the registry: a spawn waits on an
    /// ACP handshake, and one session's slow agent must not hold up another
    /// operator's first turn. Entries are kept for the life of the registry —
    /// a mutex is a word, and the alternative is a removal that could race the
    /// acquisition it is meant to protect.
    spawn_locks: DashMap<AssistantSessionId, Arc<tokio::sync::Mutex<()>>>,
    /// One live-frame channel per session, INDEPENDENT of whether a process is
    /// running.
    ///
    /// 🔴 The channel cannot belong to the process. A client that opens the
    /// socket on a dormant session and then sends a turn must see that turn's
    /// frames: if the channel were minted with the harness, the subscriber
    /// taken before the spawn would be listening to a channel nothing would
    /// ever publish to, and the socket would sit silent through a whole
    /// conversation with no error to show for it.
    channels: DashMap<AssistantSessionId, broadcast::Sender<AssistantSessionFrame>>,
    /// Per-session locks serializing document-edit validate-then-append
    /// (`sessions/document_edits.rs`), minted on first use like the channels.
    document_edits: DashMap<AssistantSessionId, Arc<tokio::sync::Mutex<()>>>,
    store: Arc<dyn AssistantSessionStore>,
    config: ResolvedAssistantConfig,
    endpoint: Option<AssistantEndpoints>,
    /// The harnesses this registry can name and launch. The server passes the
    /// shipped catalogue; a test passes one whose programs it controls, so a
    /// cell about what happens AFTER a spawn is attempted never has a host
    /// launcher (`npx`, node) on its clock.
    catalogue: &'static [CatalogueHarness],
    /// The store failure that makes this surface unavailable, once one has been
    /// observed. A `RwLock` because it is read on every description and written
    /// at most twice in a server's life.
    store_fault: std::sync::RwLock<Option<String>>,
}

impl AssistantSessions {
    /// Build the registry over one durable store and the operator's
    /// configuration.
    #[must_use]
    pub fn new(
        store: Arc<dyn AssistantSessionStore>,
        config: ResolvedAssistantConfig,
        endpoint: Option<AssistantEndpoints>,
        catalogue: &'static [CatalogueHarness],
    ) -> Self {
        Self {
            inner: Arc::new(Inner {
                live: DashMap::new(),
                document_edits: DashMap::new(),
                spawn_locks: DashMap::new(),
                channels: DashMap::new(),
                store,
                config,
                endpoint,
                catalogue,
                store_fault: std::sync::RwLock::new(None),
            }),
        }
    }

    /// The harnesses this registry resolves names through.
    pub(crate) fn catalogue(&self) -> &'static [CatalogueHarness] {
        self.inner.catalogue
    }

    /// The operator's resolved `[assistant]` configuration.
    #[must_use]
    pub fn config(&self) -> &ResolvedAssistantConfig {
        &self.inner.config
    }

    /// Whether this server can open a session at all, and why not when it
    /// cannot.
    ///
    /// A stock server can: there is no `[assistant]` section to write, the
    /// harness catalogue ships with the build, and "not configured" is no longer
    /// a reason anything may give (RULED 2026-08-29). What CAN take the surface
    /// down is the durable store the sessions live in — a session is a record
    /// and a transcript before it is a process — and that is a refusal the
    /// product can name, with the store's own error in it.
    ///
    /// Read from the boot sweep, which is the one place this server has already
    /// exercised the store end to end. Nothing probes on the descriptor path: a
    /// full listing per description would make the panel's own refresh the
    /// heaviest read on the box.
    ///
    /// Whether a PARTICULAR harness can run is a different question with a
    /// different answer per entry, and it is answered on the descriptor's
    /// `harnesses[]` (`available`, with the install hint) rather than folded
    /// into one sentence here.
    #[must_use]
    pub fn availability(&self) -> Availability {
        match self.store_fault() {
            Some(reason) => Availability::Unavailable { reason },
            None => Availability::Available,
        }
    }

    /// The store failure the boot sweep observed, if it observed one.
    fn store_fault(&self) -> Option<String> {
        self.inner
            .store_fault
            .read()
            .unwrap_or_else(std::sync::PoisonError::into_inner)
            .clone()
    }

    /// Record that the durable store could not be read or written.
    ///
    /// Called by the boot sweep, which is the first thing this server does with
    /// the assistant store. It is what turns `sessions_disabled_reason` from a
    /// field nothing ever fills into the product's own sentence about a store it
    /// cannot use.
    pub(crate) fn report_store_fault(&self, error: &AssistantSessionError) {
        let reason = format!("{STORE_UNUSABLE}: {error}");
        tracing::error!(%reason, "assistant sessions are unavailable on this server");
        *self
            .inner
            .store_fault
            .write()
            .unwrap_or_else(std::sync::PoisonError::into_inner) = Some(reason);
    }

    /// Clear a previously reported store fault, the store having answered.
    pub(crate) fn clear_store_fault(&self) {
        self.inner
            .store_fault
            .write()
            .unwrap_or_else(std::sync::PoisonError::into_inner)
            .take();
    }

    /// The harness this caller last opened a session on, or [`None`].
    ///
    /// Read from the store, so it survives the restart that a remembered
    /// in-process choice would not. [`None`] is a complete answer — a caller who
    /// has picked nothing has picked nothing — and the console preselects the
    /// first available catalogue entry rather than the server inventing one.
    ///
    /// # Errors
    ///
    /// Whatever the store reports.
    pub async fn last_harness_pick(
        &self,
        subject: &str,
    ) -> Result<Option<String>, AssistantSessionError> {
        Ok(self.inner.store.assistant_default_harness(subject).await?)
    }

    /// Whether a session's agent is handed this server's own assistant tool
    /// server — the `assistant_context` route.
    ///
    /// Independent of `[mcp] enabled` and of `[assistant.tools] aion`: the only
    /// thing that can take it away is this server being unable to state an
    /// address an agent could dial back on.
    #[must_use]
    pub fn hands_over_assistant_tools(&self) -> bool {
        self.inner.endpoint.is_some()
    }

    /// Whether a session's agent is handed this server's GENERAL MCP endpoint —
    /// the workflow tools.
    ///
    /// One switch, `[mcp] enabled`, and no second one: the `[assistant.tools]
    /// aion` knob was retired with the rest of the section. Whether this server
    /// publishes workflow tools at all is a question an operator answers once,
    /// where the tools are; asking it again under the assistant would be a
    /// second thing to keep in step, and a session whose agent silently lacked
    /// the tools the server publishes is exactly the confusion that costs.
    #[must_use]
    pub fn hands_over_general_mcp(&self) -> bool {
        self.inner
            .endpoint
            .as_ref()
            .and_then(AssistantEndpoints::aion_mcp_url)
            .is_some()
    }

    /// This server's own MCP endpoint, when it has one to hand an agent.
    pub(crate) fn aion_endpoint(&self) -> Option<&AssistantEndpoints> {
        self.inner.endpoint.as_ref()
    }

    /// The lock serialising spawns for one session.
    pub(crate) fn spawn_lock(&self, session_id: AssistantSessionId) -> Arc<tokio::sync::Mutex<()>> {
        Arc::clone(
            self.inner
                .spawn_locks
                .entry(session_id)
                .or_insert_with(|| Arc::new(tokio::sync::Mutex::new(())))
                .value(),
        )
    }

    /// Every session this server currently holds a process for.
    pub(crate) fn live_ids(&self) -> Vec<AssistantSessionId> {
        self.inner.live.iter().map(|entry| *entry.key()).collect()
    }

    /// The durable store behind this registry.
    pub(crate) fn store(&self) -> &Arc<dyn AssistantSessionStore> {
        &self.inner.store
    }

    /// The per-session document-edit lock map (`sessions/document_edits.rs`).
    pub(super) fn document_edit_locks(
        &self,
    ) -> &DashMap<AssistantSessionId, Arc<tokio::sync::Mutex<()>>> {
        &self.inner.document_edits
    }

    /// The live process for a session, if this server is holding one.
    pub(crate) fn live(&self, session_id: AssistantSessionId) -> Option<Arc<LiveSession>> {
        self.inner
            .live
            .get(&session_id)
            .map(|entry| Arc::clone(entry.value()))
    }

    /// Adopt a started harness process for a session.
    pub(crate) fn adopt(&self, session: Arc<LiveSession>) {
        self.inner.live.insert(session.session_id(), session);
    }

    /// Forget a session's process, returning it so a caller can close it.
    pub(crate) fn forget(&self, session_id: AssistantSessionId) -> Option<Arc<LiveSession>> {
        self.inner
            .live
            .remove(&session_id)
            .map(|(_id, session)| session)
    }

    /// A recorder for a session, live or not.
    ///
    /// A session with no process still records: the boot sweep settles one by
    /// APPENDING to its transcript, and a settlement that could only be written
    /// while a process was running would be a settlement that could never be
    /// written at all.
    ///
    /// Every recorder for one session publishes to the SAME channel — the one
    /// this registry holds, not one the harness minted — so a subscriber taken
    /// before a spawn still receives what the spawn goes on to produce.
    pub(crate) fn recorder(&self, session_id: AssistantSessionId) -> Recorder {
        Recorder::new(
            session_id,
            Arc::clone(&self.inner.store),
            self.channel(session_id),
        )
    }

    /// The live-frame channel for a session, minted on first use.
    fn channel(&self, session_id: AssistantSessionId) -> broadcast::Sender<AssistantSessionFrame> {
        self.inner
            .channels
            .entry(session_id)
            .or_insert_with(|| {
                let (sender, _receiver) = broadcast::channel(LIVE_FRAME_BUFFER);
                sender
            })
            .value()
            .clone()
    }

    /// The record for a session the caller owns.
    ///
    /// # Errors
    ///
    /// [`AssistantSessionError::NotFound`] when no record exists, and
    /// [`AssistantSessionError::NotYours`] when one exists under another
    /// subject — reported to the caller as the same not-found, so nobody learns
    /// that another subject's session exists.
    pub(crate) async fn owned_record(
        &self,
        subject: &str,
        session_id: AssistantSessionId,
    ) -> Result<AssistantSessionRecord, AssistantSessionError> {
        let record = self
            .inner
            .store
            .get_assistant_session(&session_id)
            .await?
            .ok_or(AssistantSessionError::NotFound { session_id })?;
        if record.subject != subject {
            return Err(AssistantSessionError::NotYours {
                session_id,
                subject: subject.to_owned(),
            });
        }
        Ok(record)
    }

    /// One session's record, with NO caller narrowing.
    ///
    /// The narrowed read ([`Self::owned_record`]) is for a human caller, whose
    /// authority is a subject. The assistant MCP route's caller is a SESSION —
    /// it holds that session's own minted bearer and no subject at all — so it
    /// reads the record it is about to prove it is, and the proof is the digest
    /// comparison rather than a subject match.
    ///
    /// # Errors
    ///
    /// Whatever the store reports.
    pub async fn record(
        &self,
        session_id: AssistantSessionId,
    ) -> Result<Option<AssistantSessionRecord>, AssistantSessionError> {
        Ok(self.inner.store.get_assistant_session(&session_id).await?)
    }

    /// A session's projected state and cause.
    ///
    /// The public read of [`Self::state_of`], for the assistant MCP route: a
    /// bearer is honoured only while the session it names is not `ended`, and
    /// that decision has to read the same projection every other surface reads.
    ///
    /// # Errors
    ///
    /// Whatever the store reports.
    pub async fn state_of_session(
        &self,
        session_id: AssistantSessionId,
    ) -> Result<(AssistantSessionState, Option<String>), AssistantSessionError> {
        self.state_of(session_id).await
    }

    /// Every session this caller owns, newest first.
    ///
    /// CALLER-scoped, not namespace-scoped: a session is one operator's
    /// conversation with an agent and lives in no namespace. The store
    /// enumerates what it holds and the narrowing happens here, where the caller
    /// identity is.
    ///
    /// # Errors
    ///
    /// Whatever the store reports.
    pub async fn list(
        &self,
        subject: &str,
    ) -> Result<Vec<AssistantSessionSummary>, AssistantSessionError> {
        let listing = self.inner.store.list_assistant_sessions().await?;
        for row in &listing.undecodable {
            // Named, never silent: an operator whose session is unreadable must
            // be able to find out from the log that it EXISTS, even though it
            // cannot appear in a typed listing.
            tracing::warn!(
                session = %row.session_id,
                error = %row.error,
                "an assistant session record could not be decoded and is omitted from the listing"
            );
        }
        let mut summaries = Vec::new();
        for record in listing.sessions {
            if record.subject != subject {
                continue;
            }
            let session_id = record.session_id;
            let (state, reason) = self.state_of(session_id).await?;
            summaries.push(record.summary(state, reason));
        }
        // Newest first: the operator's current conversation is the one they
        // came for.
        summaries.sort_by_key(|summary| std::cmp::Reverse(summary.created_at));
        Ok(summaries)
    }

    /// A session's state and cause.
    ///
    /// A running process always wins. With no process, the LAST record that
    /// settled the session decides — read from the transcript's tail rather than
    /// its whole length, because a listing must not cost one full conversation
    /// read per row. The boot sweep guarantees a settling record is the last
    /// thing on a non-live session's transcript; when the tail does not carry
    /// one, the whole transcript is read rather than a state guessed.
    pub(crate) async fn state_of(
        &self,
        session_id: AssistantSessionId,
    ) -> Result<(AssistantSessionState, Option<String>), AssistantSessionError> {
        if self.is_live(session_id).await {
            return Ok((AssistantSessionState::Live, None));
        }
        let head = self
            .inner
            .store
            .assistant_transcript_head(&session_id)
            .await?;
        let tail_from = head.saturating_sub(SETTLEMENT_TAIL);
        let tail = self
            .transcript_from(session_id, tail_from.checked_sub(1))
            .await?;
        let projection = AssistantSessionProjection::of(tail.iter().map(|frame| &frame.event));
        if projection.settled.is_some() {
            return Ok(projection.state(None));
        }
        // The tail carried no settlement. Rather than guess, read the whole
        // conversation: this is bounded by one session and only happens on a
        // transcript the boot sweep has not written back to.
        let whole = self.transcript_from(session_id, None).await?;
        Ok(AssistantSessionProjection::of(whole.iter().map(|frame| &frame.event)).state(None))
    }

    /// Whether a process is running for this session.
    pub(crate) async fn is_live(&self, session_id: AssistantSessionId) -> bool {
        match self.live(session_id) {
            Some(live) => live.is_alive().await,
            None => false,
        }
    }

    /// A session's transcript as wire frames, strictly after `after`.
    ///
    /// # Errors
    ///
    /// Whatever the store reports, or
    /// [`AssistantSessionError::Internal`] when a stored frame cannot be
    /// decoded — reported rather than skipped, because a transcript with a hole
    /// in it that nobody mentions is worse than one that refuses to be read.
    pub(crate) async fn transcript_from(
        &self,
        session_id: AssistantSessionId,
        after: Option<u64>,
    ) -> Result<Vec<AssistantSessionFrame>, AssistantSessionError> {
        let stored = self
            .inner
            .store
            .assistant_transcript(&session_id, after)
            .await?;
        stored.iter().map(decode_frame).collect()
    }

    /// The projection over a session's whole transcript — every fact the resume
    /// path and the context tool read.
    ///
    /// # Errors
    ///
    /// Whatever [`Self::transcript_from`] reports.
    pub(crate) async fn projection(
        &self,
        session_id: AssistantSessionId,
    ) -> Result<AssistantSessionProjection, AssistantSessionError> {
        let frames = self.transcript_from(session_id, None).await?;
        Ok(AssistantSessionProjection::of(
            frames.iter().map(|frame| &frame.event),
        ))
    }

    /// The most recently shared on-screen context for a session.
    ///
    /// What the `assistant_context` MCP tool answers with, read off the
    /// transcript: the transcript IS the shared context, so there is no second
    /// store to keep in step and a restart loses nothing.
    ///
    /// # Errors
    ///
    /// Whatever [`Self::projection`] reports.
    pub async fn latest_context(
        &self,
        session_id: AssistantSessionId,
    ) -> Result<Option<AssistantTurnContext>, AssistantSessionError> {
        Ok(self.projection(session_id).await?.latest_context)
    }

    /// Record a state transition and its cause.
    pub(crate) async fn settle(
        &self,
        session_id: AssistantSessionId,
        state: AssistantSessionState,
        reason: impl Into<String>,
    ) -> Result<(), AssistantSessionError> {
        let reason = reason.into();
        // Ended is terminal: no settling record may move a session out of it.
        // The projection already ignores such a frame; refusing to write one
        // keeps the transcript from carrying a state it will never project.
        if state != AssistantSessionState::Ended
            && let Some((AssistantSessionState::Ended, cause)) =
                self.projection(session_id).await?.settled
        {
            tracing::warn!(
                %session_id,
                requested = ?state,
                %reason,
                "a settling record that would leave `ended` was refused: ended is terminal"
            );
            return Err(AssistantSessionError::Ended {
                session_id,
                reason: cause.unwrap_or_else(|| "ended".to_owned()),
            });
        }
        self.recorder(session_id)
            .record(AssistantSessionEvent::State {
                state,
                reason: Some(reason),
            })
            .await
            .map(drop)
    }

    /// Update a session's bookkeeping — the fields a listing reads without
    /// touching a transcript.
    ///
    /// The record is RE-READ from the store here rather than taken from the
    /// caller. A turn holds its copy of the record across the spawn, and the
    /// spawn stores the token digest of the secret it actually handed the
    /// child; writing the caller's copy back would silently restore the digest
    /// that copy remembers — a bearer no live child holds, so every MCP call
    /// the agent makes would fail verification from that write on. Every
    /// writer of this record mutates what the store holds NOW, never a copy
    /// carried across an await.
    pub(crate) async fn touch(
        &self,
        session_id: AssistantSessionId,
        first_turn_text: Option<&str>,
    ) -> Result<(), AssistantSessionError> {
        let Some(mut record) = self.inner.store.get_assistant_session(&session_id).await? else {
            return Err(AssistantSessionError::NotFound { session_id });
        };
        record.updated_at = Utc::now();
        record.turns = record.turns.saturating_add(1);
        if record.title.is_none()
            && let Some(text) = first_turn_text
        {
            record.title = Some(text.trim().chars().take(TITLE_CHARACTERS).collect());
        }
        self.inner.store.put_assistant_session(record).await?;
        Ok(())
    }
}

/// Whether this server can open assistant sessions.
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum Availability {
    /// Sessions can be opened.
    Available,
    /// They cannot, and this is why — in words an operator can act on.
    Unavailable {
        /// The reason, for the descriptor and the panel.
        reason: String,
    },
}

impl Availability {
    /// Whether sessions can be opened.
    #[must_use]
    pub const fn is_available(&self) -> bool {
        matches!(self, Self::Available)
    }

    /// The reason sessions are off, or `None` when they are on.
    #[must_use]
    pub fn reason(&self) -> Option<&str> {
        match self {
            Self::Available => None,
            Self::Unavailable { reason } => Some(reason),
        }
    }
}

/// How many trailing transcript events a state projection reads before falling
/// back to the whole conversation.
///
/// Not a tuning knob and not a cap on anything: the boot sweep guarantees the
/// LAST event of a session with no process is the record that settled it, so one
/// event would do. A few are read because a settlement immediately followed by,
/// say, a straggling frame is cheap to tolerate and expensive to be wrong about,
/// and the fallback reads everything rather than guessing.
const SETTLEMENT_TAIL: u64 = 8;

/// How many live frames a subscriber may fall behind by before it is dropped.
///
/// A BACKPRESSURE bound, not an operator knob: the round-2 amendment retired the
/// `event_buffer` setting because the number is not a policy anybody can choose
/// usefully — a socket that cannot keep up with a token stream has to be told so
/// rather than allowed to consume the server's memory. The durable transcript is
/// what a dropped subscriber reads to catch up (`?after=`), so falling behind
/// costs a reconnect and never a lost frame.
const LIVE_FRAME_BUFFER: usize = 1_024;

/// The first half of the sentence a server with an unusable store answers with.
///
/// The one reason a stock server states for sessions being off, and it names a
/// real fault rather than a missing configuration — there is no configuration to
/// miss.
pub const STORE_UNUSABLE: &str = "the durable store this server keeps assistant sessions in could not be read at start-up, so \
     a conversation could not be recorded";

/// Decode one stored transcript row into the wire frame it holds.
fn decode_frame(
    stored: &AssistantTranscriptEvent,
) -> Result<AssistantSessionFrame, AssistantSessionError> {
    let event: AssistantSessionEvent =
        serde_json::from_slice(stored.payload.bytes()).map_err(|error| {
            AssistantSessionError::Internal(format!(
                "the assistant transcript event at index {} does not decode as a session frame \
                 ({error}); the transcript is reported as unreadable rather than served with a \
                 hole in it",
                stored.index
            ))
        })?;
    Ok(AssistantSessionFrame {
        index: stored.index,
        event,
    })
}