aion_store/

assistant.rs

//! Durable assistant-session records and their transcript contract.
//!
//! An assistant session is one live agent-harness process the server owns on a
//! caller's behalf. The PROCESS dies with the server; the RECORD and the
//! TRANSCRIPT do not. That split is the whole shape of this module:
//!
//! - **[`AssistantSessionRecord`]** — one per session: who owns it, which
//!   harness and account it runs, when it was created, and the bookkeeping a
//!   listing needs. It carries NO STATUS. Whether a session is live, dormant or
//!   ended is a projection over its transcript records plus the live-process
//!   registry ([`aion_core::AssistantSessionProjection`]), for the reason
//!   `WorkflowStatus` is a projection: a stored status is a second answer that
//!   drifts from the records that justify it. A session whose process is gone is
//!   settled by an APPENDED record at the next boot — written back, not merely
//!   displayed.
//! - **[`AssistantTranscriptEvent`]** — one per event, at a dense
//!   store-assigned index. Every frame the WebSocket streams is appended here
//!   FIRST and broadcast second, so nothing a client saw is absent from the
//!   record, and the socket's `?after=` replay is a read of the same rows.
//!
//! # The store assigns the index, inside the append
//!
//! [`AssistantSessionStore::append_assistant_transcript_event`] takes no index.
//! The store reads the current head and commits at it under the backend's own
//! optimistic-concurrency discipline, retrying on a conflict — so two appenders
//! racing on one session receive `n` and `n+1`, never `n` twice. A caller
//! cannot supply an index, so a caller cannot mint a duplicate one.
//!
//! # A poisoned row is listed, never skipped
//!
//! [`AssistantSessionListing`] carries both the decoded records and every row
//! that was present and could not be decoded, with its decode error. The same
//! discipline the workloop registry uses: an operator must be able to see that a
//! session exists and is unreadable, rather than watch it silently vanish from
//! a list.

use aion_core::{
    AssistantCommand, AssistantConfigOption, AssistantSessionId, AssistantSessionState,
    AssistantSessionSummary, Payload,
};
use async_trait::async_trait;
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};

use crate::StoreError;

/// One durable assistant session.
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct AssistantSessionRecord {
    /// The session's identity (primary key).
    pub session_id: AssistantSessionId,
    /// The caller subject that owns it. The list surface is per caller, so this
    /// is the field it filters on — never a namespace, because a session is not
    /// a workflow and lives in no namespace.
    pub subject: String,
    /// The configured harness name it runs.
    pub harness: String,
    /// The configured account name, when the harness declares any.
    pub account: Option<String>,
    /// The first 80 characters of the first prompt; `None` before one.
    pub title: Option<String>,
    /// When the session was created.
    pub created_at: DateTime<Utc>,
    /// The last time anything about it changed.
    pub updated_at: DateTime<Utc>,
    /// How many turns have been submitted on it — bookkeeping the server keeps
    /// in step with the transcript projection, so a listing does not have to
    /// read every session's whole transcript to count them.
    pub turns: u64,
    /// The SHA-256 digest, lowercase hex, of the session-scoped bearer this
    /// session's harness carries to the server's own MCP endpoint. `None` when
    /// the session was given no aion MCP server.
    ///
    /// The digest, never the token: a store that held the token would be a
    /// store that could hand it out. Verification hashes what the caller
    /// presented and compares. Its LIFETIME is the session's — a call is
    /// admitted only while the session's projected lifecycle is not `ended`, so
    /// revocation is a durable fact rather than an in-memory set that dies with
    /// the process that minted it.
    pub mcp_token_digest: Option<String>,
    /// The commands the harness advertised most recently.
    ///
    /// A DERIVED CACHE of the transcript's latest
    /// [`aion_core::AssistantSessionEvent::AvailableCommands`] record, kept for
    /// exactly the reason `turns` and `title` are: a listing must not read every
    /// session's whole conversation to answer one field, and an advertisement
    /// can arrive at the first turn of a conversation that runs for hours, so a
    /// bounded tail read cannot find it.
    ///
    /// The TRANSCRIPT is the authority. The frame is appended first and this is
    /// written from it inside the same append path, so the two cannot be written
    /// out of order; `the_records_command_cache_is_what_the_transcript_projects`
    /// pins that a rebuild from the transcript reproduces this field exactly.
    /// Unlike a status, this is not a second ANSWER to a question the transcript
    /// answers differently — it is the same answer, kept where a listing can
    /// afford to read it.
    #[serde(default)]
    pub commands: Vec<AssistantCommand>,

    /// The configuration options the harness advertised most recently — the
    /// model picker among them. The same derived cache as `commands`, kept for
    /// the same listing-cost reason, with the transcript's
    /// [`aion_core::AssistantSessionEvent::ConfigOptions`] records as the
    /// authority. `#[serde(default)]` because records written before options
    /// existed decode as having none advertised, which is exactly true.
    #[serde(default)]
    pub config_options: Vec<AssistantConfigOption>,
}

impl AssistantSessionRecord {
    /// Encode the stable backend-neutral representation.
    ///
    /// # Errors
    ///
    /// [`StoreError::Serialization`] when serialization fails.
    pub fn encode(&self) -> Result<Vec<u8>, StoreError> {
        serde_json::to_vec(self).map_err(|error| StoreError::Serialization(error.to_string()))
    }

    /// Decode and re-validate a stored representation.
    ///
    /// # Errors
    ///
    /// [`StoreError::Serialization`] for malformed bytes or an unknown field.
    pub fn decode(bytes: &[u8]) -> Result<Self, StoreError> {
        serde_json::from_slice(bytes).map_err(|error| StoreError::Serialization(error.to_string()))
    }

    /// The listing row for this record under a state the caller has projected.
    ///
    /// The state, its lifecycle and its cause are PARAMETERS because none of
    /// them is stored: they are read off the session's transcript records plus
    /// whether a process is running (see
    /// [`aion_core::AssistantSessionProjection`]). A record that carried its own
    /// status would be a second answer to a question the transcript already
    /// answers, and the two would drift.
    #[must_use]
    pub fn summary(
        &self,
        state: AssistantSessionState,
        reason: Option<String>,
    ) -> AssistantSessionSummary {
        AssistantSessionSummary {
            session_id: self.session_id,
            harness: self.harness.clone(),
            account: self.account.clone(),
            state,
            reason,
            created_at: self.created_at,
            updated_at: self.updated_at,
            turns: self.turns,
            title: self.title.clone(),
            commands: self.commands.clone(),
            config_options: self.config_options.clone(),
        }
    }
}

/// One durable transcript event at its store-assigned index.
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct AssistantTranscriptEvent {
    /// The event's dense position in the session's transcript, assigned by the
    /// store inside the append. Zero-based and contiguous.
    pub index: u64,
    /// When the server recorded it.
    pub recorded_at: DateTime<Utc>,
    /// The event itself, type-erased — the serialized
    /// [`aion_core::AssistantSessionEvent`] the socket streams. Carried as a
    /// [`Payload`] because the store is type-erased by construction: it
    /// persists bytes and a content-type tag, never a domain type it would have
    /// to keep in step.
    pub payload: Payload,
}

impl AssistantTranscriptEvent {
    /// Encode the stable backend-neutral representation.
    ///
    /// # Errors
    ///
    /// [`StoreError::Serialization`] when serialization fails.
    pub fn encode(&self) -> Result<Vec<u8>, StoreError> {
        serde_json::to_vec(self).map_err(|error| StoreError::Serialization(error.to_string()))
    }

    /// Decode a stored representation.
    ///
    /// # Errors
    ///
    /// [`StoreError::Serialization`] for malformed bytes or an unknown field.
    pub fn decode(bytes: &[u8]) -> Result<Self, StoreError> {
        serde_json::from_slice(bytes).map_err(|error| StoreError::Serialization(error.to_string()))
    }
}

/// An assistant-session row that was present but could not be decoded.
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct UndecodableAssistantSession {
    /// The key the poisoned row is stored under (the session id's text form).
    pub session_id: String,
    /// The decode failure, rendered for operator diagnosis.
    pub error: String,
}

/// Complete assistant-session listing, including poisoned-row visibility.
#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct AssistantSessionListing {
    /// Successfully decoded records, ordered by `created_at` then session id.
    pub sessions: Vec<AssistantSessionRecord>,
    /// Present rows that could not be decoded, ordered by session id text.
    pub undecodable: Vec<UndecodableAssistantSession>,
}

impl AssistantSessionListing {
    /// Put the listing in its contract order: sessions by `created_at` then
    /// session id, poisoned rows by session id.
    ///
    /// One helper rather than a sort at each backend, so "ordered by `created_at`
    /// then id" is one rule with one implementation instead of two that could
    /// tie-break differently on the same instant.
    pub fn sort(&mut self) {
        self.sessions
            .sort_by(|left, right| match left.created_at.cmp(&right.created_at) {
                std::cmp::Ordering::Equal => left.session_id.cmp(&right.session_id),
                ordering => ordering,
            });
        self.undecodable
            .sort_by(|left, right| left.session_id.cmp(&right.session_id));
    }
}

/// Durable assistant-session persistence contract.
///
/// Implemented by every backend and exercised by
/// [`crate::conformance::run_assistant_session_suite`] against all of them, so
/// the index-assignment and poisoned-row guarantees are properties of the
/// CONTRACT rather than of whichever backend a deployment happens to run.
#[async_trait]
pub trait AssistantSessionStore: Send + Sync + 'static {
    /// Create or replace a session's record.
    async fn put_assistant_session(&self, record: AssistantSessionRecord)
    -> Result<(), StoreError>;

    /// Look up one session by id.
    async fn get_assistant_session(
        &self,
        session_id: &AssistantSessionId,
    ) -> Result<Option<AssistantSessionRecord>, StoreError>;

    /// List decodable sessions and report every undecodable row, ordered by
    /// `created_at` then session id.
    ///
    /// Unfiltered: a store enumerates what it holds, and the per-caller
    /// narrowing is the server's authorization decision, made where the caller
    /// identity is.
    async fn list_assistant_sessions(&self) -> Result<AssistantSessionListing, StoreError>;

    /// Append one event to a session's transcript, returning the index the
    /// store assigned it.
    ///
    /// The index is the store's to give: it is the next index after the last
    /// stored one, committed under the backend's optimistic-concurrency
    /// discipline so two appenders racing on one session get `n` and `n+1` and
    /// never `n` twice. A caller cannot pass an index and so cannot mint a
    /// duplicate.
    ///
    /// # Errors
    ///
    /// [`StoreError::AssistantSessionNotFound`] when no record exists for
    /// `session_id` — an append never silently creates a session, because a
    /// transcript with no record is a conversation with no owner. Otherwise a
    /// backend or serialization error.
    async fn append_assistant_transcript_event(
        &self,
        session_id: &AssistantSessionId,
        recorded_at: DateTime<Utc>,
        payload: Payload,
    ) -> Result<u64, StoreError>;

    /// The next index the store would assign this session — equivalently, how
    /// many events its transcript holds.
    ///
    /// An unwritten transcript reads `0`. Cheap on every backend (it reads
    /// stream metadata, never the events), which is what lets a listing find
    /// each session's last record without reading its whole conversation.
    async fn assistant_transcript_head(
        &self,
        session_id: &AssistantSessionId,
    ) -> Result<u64, StoreError>;

    /// A session's transcript in index order, excluding every event at or below
    /// `after`.
    ///
    /// `None` reads the whole transcript. An unknown session reads empty rather
    /// than refusing: a read of a session that is not there is an absence, and
    /// the refusal that matters is on the WRITE.
    async fn assistant_transcript(
        &self,
        session_id: &AssistantSessionId,
        after: Option<u64>,
    ) -> Result<Vec<AssistantTranscriptEvent>, StoreError>;

    /// Remember the harness `subject` last opened a session on.
    ///
    /// Written by `createSession` and by nothing else: the memory is a record of
    /// what the operator DID, not a preference they set, so there is no second
    /// door through which it could say something a session never said.
    ///
    /// It lives in the store rather than in the process because it must survive
    /// a restart — an operator who picked Claude Code yesterday must not be
    /// asked again this morning — and CALLER-scoped because it is one person's
    /// last choice, not a deployment default.
    async fn put_assistant_default_harness(
        &self,
        subject: &str,
        harness: &str,
    ) -> Result<(), StoreError>;

    /// The harness `subject` last opened a session on, or [`None`] before any.
    ///
    /// [`None`] is a complete answer and never a fallback: a caller who has
    /// picked nothing has picked nothing, and a surface that invented a default
    /// here would be reporting a choice its operator never made.
    async fn assistant_default_harness(&self, subject: &str) -> Result<Option<String>, StoreError>;
}

#[cfg(test)]
mod tests {
    use aion_core::ContentType;
    use chrono::TimeZone;

    use super::*;

    fn instant(offset: i64) -> Result<DateTime<Utc>, StoreError> {
        Utc.with_ymd_and_hms(2026, 8, 29, 6, 0, 0)
            .single()
            .map(|base| base + chrono::Duration::seconds(offset))
            .ok_or_else(|| StoreError::Backend("test instant must be valid".to_owned()))
    }

    fn record(offset: i64) -> Result<AssistantSessionRecord, StoreError> {
        Ok(AssistantSessionRecord {
            session_id: AssistantSessionId::new(uuid::Uuid::from_u128(7)),
            subject: String::from("operator"),
            harness: String::from("claude"),
            account: Some(String::from("work")),
            title: Some(String::from("fix the check")),
            created_at: instant(offset)?,
            updated_at: instant(offset + 10)?,
            turns: 2,
            mcp_token_digest: Some("0".repeat(64)),
            commands: vec![AssistantCommand {
                name: String::from("compact"),
                description: String::from("compact the conversation"),
                input_hint: None,
            }],
            config_options: Vec::new(),
        })
    }

    #[test]
    fn a_session_record_round_trips() -> Result<(), StoreError> {
        let expected = record(0)?;
        assert_eq!(
            AssistantSessionRecord::decode(&expected.encode()?)?,
            expected
        );
        Ok(())
    }

    #[test]
    fn a_record_carries_no_status_of_its_own() -> Result<(), StoreError> {
        // The summary's state, lifecycle and cause all arrive as arguments:
        // there is no field on the record a caller could read instead, which is
        // what makes the projection the only answer.
        let summary = record(0)?.summary(
            AssistantSessionState::Dormant,
            Some(String::from("process_exited")),
        );
        assert_eq!(summary.state, AssistantSessionState::Dormant);
        assert_eq!(summary.reason.as_deref(), Some("process_exited"));
        Ok(())
    }

    #[test]
    fn a_transcript_event_round_trips() -> Result<(), StoreError> {
        let event = AssistantTranscriptEvent {
            index: 4,
            recorded_at: instant(0)?,
            payload: Payload::new(ContentType::Json, b"{\"type\":\"delta\"}".to_vec()),
        };
        assert_eq!(AssistantTranscriptEvent::decode(&event.encode()?)?, event);
        Ok(())
    }

    #[test]
    fn the_listing_orders_by_created_at_then_id() -> Result<(), StoreError> {
        let mut early = record(0)?;
        early.session_id = AssistantSessionId::new(uuid::Uuid::from_u128(2));
        let mut tied = record(0)?;
        tied.session_id = AssistantSessionId::new(uuid::Uuid::from_u128(1));
        let late = record(100)?;
        let mut listing = AssistantSessionListing {
            sessions: vec![late.clone(), early.clone(), tied.clone()],
            undecodable: vec![
                UndecodableAssistantSession {
                    session_id: String::from("b"),
                    error: String::from("bad"),
                },
                UndecodableAssistantSession {
                    session_id: String::from("a"),
                    error: String::from("bad"),
                },
            ],
        };
        listing.sort();
        assert_eq!(listing.sessions, vec![tied, early, late]);
        assert_eq!(
            listing
                .undecodable
                .iter()
                .map(|row| row.session_id.as_str())
                .collect::<Vec<_>>(),
            vec!["a", "b"]
        );
        Ok(())
    }
}