aion-store 0.29.0

//! Shared assistant-session persistence scenarios.
//!
//! Every backend answers these, so the two properties the server relies on are
//! properties of the CONTRACT rather than of whichever store a deployment
//! happens to run:
//!
//! - **the store assigns a dense, unique index inside the append.** The
//!   interleaved-appenders cell drives two appends at once on one session and
//!   requires the two indices to be `{n, n+1}`. A backend that let a caller's
//!   read-then-write race decide the index would hand out `{n, n}` and fail
//!   here — which is the whole point of taking the index parameter away.
//! - **a poisoned row is listed, never skipped.** A session an operator can see
//!   and cannot read is a different fact from a session that is not there, and
//!   the listing has to be able to say which.

use std::sync::Arc;

use aion_core::{AssistantSessionId, ContentType, Payload};
use chrono::{DateTime, TimeZone, Utc};

use super::{AssistantSessionRawWrite, AssistantSessionReopen};
use crate::StoreError;
use crate::assistant::{AssistantSessionRecord, AssistantSessionStore};

pub(crate) async fn run(
    store: Arc<dyn AssistantSessionStore>,
    reopen: Option<AssistantSessionReopen>,
    write_raw: AssistantSessionRawWrite,
) -> Result<(), StoreError> {
    round_trip(&store).await?;
    listing_orders_by_created_at_then_id(&store).await?;
    indices_are_dense_from_zero(&store).await?;
    two_interleaved_appenders_never_share_an_index(&store).await?;
    transcript_after_excludes_the_bound(&store).await?;
    append_to_an_unknown_session_refuses(&store).await?;
    poisoned_row_is_listed_with_its_error(&store, &write_raw).await?;
    drop(write_raw);
    the_last_pick_is_per_caller_and_replaces_itself(&store).await?;
    restart_survival(store, reopen).await
}

/// The harness memory is one row per CALLER, and the newest pick is the answer.
///
/// Both halves matter. Per-caller, because a shared row would mean one
/// operator's choice preselecting another's console. Replaced rather than
/// accumulated, because the question is "what did you last open", and a store
/// that kept the first answer would preselect a harness the operator moved off
/// weeks ago.
async fn the_last_pick_is_per_caller_and_replaces_itself(
    store: &Arc<dyn AssistantSessionStore>,
) -> Result<(), StoreError> {
    require(
        store.assistant_default_harness("nobody").await?.is_none(),
        "a caller who has opened nothing has picked nothing; an invented default here would \
         report a choice its operator never made",
    )?;
    store
        .put_assistant_default_harness("ada", "claude-code")
        .await?;
    store.put_assistant_default_harness("bob", "codex").await?;
    require(
        store.assistant_default_harness("ada").await?.as_deref() == Some("claude-code"),
        "the pick is read back for the caller that made it",
    )?;
    require(
        store.assistant_default_harness("bob").await?.as_deref() == Some("codex"),
        "one caller's pick is not another's: a shared row would preselect somebody else's harness",
    )?;
    store
        .put_assistant_default_harness("ada", "opencode")
        .await?;
    require(
        store.assistant_default_harness("ada").await?.as_deref() == Some("opencode"),
        "the LAST pick is the answer",
    )?;
    require(
        store.assistant_default_harness("bob").await?.as_deref() == Some("codex"),
        "replacing one caller's pick leaves every other caller's alone",
    )
}

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("conformance instant is invalid".to_owned()))
}

fn session(nth: u128) -> AssistantSessionId {
    AssistantSessionId::new(uuid::Uuid::from_u128(nth))
}

fn record(
    session_id: AssistantSessionId,
    created_offset: i64,
) -> Result<AssistantSessionRecord, StoreError> {
    Ok(AssistantSessionRecord {
        session_id,
        subject: String::from("operator"),
        harness: String::from("claude"),
        account: Some(String::from("work")),
        title: None,
        created_at: instant(created_offset)?,
        updated_at: instant(created_offset)?,
        turns: 0,
        mcp_token_digest: Some("a".repeat(64)),
        commands: Vec::new(),
    })
}

fn frame(text: &str) -> Payload {
    Payload::new(
        ContentType::Json,
        format!("{{\"type\":\"delta\",\"turn_id\":\"t-1\",\"text\":\"{text}\"}}").into_bytes(),
    )
}

fn require(condition: bool, message: impl Into<String>) -> Result<(), StoreError> {
    if condition {
        Ok(())
    } else {
        Err(StoreError::Backend(message.into()))
    }
}

async fn round_trip(store: &Arc<dyn AssistantSessionStore>) -> Result<(), StoreError> {
    let expected = record(session(1), 0)?;
    store.put_assistant_session(expected.clone()).await?;
    let read = store.get_assistant_session(&session(1)).await?;
    require(
        read.as_ref() == Some(&expected),
        format!("a stored session must read back verbatim, got {read:?}"),
    )?;
    require(
        store.get_assistant_session(&session(99)).await?.is_none(),
        "an absent session reads as absent, never as a fabricated record",
    )?;
    // Replacement, not a second row: the record is a single-key slot.
    let mut replaced = expected;
    replaced.turns = 3;
    store.put_assistant_session(replaced.clone()).await?;
    require(
        store.get_assistant_session(&session(1)).await? == Some(replaced),
        "a second put replaces the record rather than adding one",
    )
}

async fn listing_orders_by_created_at_then_id(
    store: &Arc<dyn AssistantSessionStore>,
) -> Result<(), StoreError> {
    // Written out of order on purpose: the listing's order must come from the
    // contract, not from insertion.
    store
        .put_assistant_session(record(session(21), 300)?)
        .await?;
    store
        .put_assistant_session(record(session(23), 100)?)
        .await?;
    store
        .put_assistant_session(record(session(22), 100)?)
        .await?;
    let listing = store.list_assistant_sessions().await?;
    let ordered: Vec<AssistantSessionId> = listing
        .sessions
        .iter()
        .filter(|record| [session(21), session(22), session(23)].contains(&record.session_id))
        .map(|record| record.session_id)
        .collect();
    require(
        ordered == vec![session(22), session(23), session(21)],
        format!(
            "sessions list by created_at then id — the tie at one instant breaks by id — got \
             {ordered:?}"
        ),
    )
}

async fn indices_are_dense_from_zero(
    store: &Arc<dyn AssistantSessionStore>,
) -> Result<(), StoreError> {
    let id = session(31);
    store.put_assistant_session(record(id, 0)?).await?;
    require(
        store.assistant_transcript_head(&id).await? == 0,
        "an unwritten transcript's head is zero",
    )?;
    for expected in 0_u64..5 {
        let assigned = store
            .append_assistant_transcript_event(&id, instant(0)?, frame("chunk"))
            .await?;
        require(
            assigned == expected,
            format!("the store assigns index {expected}, got {assigned}"),
        )?;
    }
    require(
        store.assistant_transcript_head(&id).await? == 5,
        "the head is the next index the store would assign",
    )?;
    let events = store.assistant_transcript(&id, None).await?;
    let indices: Vec<u64> = events.iter().map(|event| event.index).collect();
    require(
        indices == vec![0, 1, 2, 3, 4],
        format!("a transcript reads back dense and in order, got {indices:?}"),
    )
}

async fn two_interleaved_appenders_never_share_an_index(
    store: &Arc<dyn AssistantSessionStore>,
) -> Result<(), StoreError> {
    let id = session(41);
    store.put_assistant_session(record(id, 0)?).await?;
    let base = store
        .append_assistant_transcript_event(&id, instant(0)?, frame("first"))
        .await?;

    // BOTH appends are in flight at once. Whichever wins, the loser must be
    // handed the NEXT index rather than the same one: the store's own
    // read-modify-write is what decides, and a backend that let the caller
    // decide would produce {n, n} here.
    let left = Arc::clone(store);
    let right = Arc::clone(store);
    let (first, second) = futures::future::join(
        async move {
            left.append_assistant_transcript_event(&id, instant(1)?, frame("left"))
                .await
        },
        async move {
            right
                .append_assistant_transcript_event(&id, instant(2)?, frame("right"))
                .await
        },
    )
    .await;
    let mut assigned = vec![first?, second?];
    assigned.sort_unstable();
    require(
        assigned == vec![base + 1, base + 2],
        format!(
            "two concurrent appends must be handed {{{}, {}}}, got {assigned:?}",
            base + 1,
            base + 2
        ),
    )?;
    let events = store.assistant_transcript(&id, None).await?;
    let indices: Vec<u64> = events.iter().map(|event| event.index).collect();
    require(
        indices == vec![base, base + 1, base + 2],
        format!("the transcript stays dense after a race, got {indices:?}"),
    )
}

async fn transcript_after_excludes_the_bound(
    store: &Arc<dyn AssistantSessionStore>,
) -> Result<(), StoreError> {
    let id = session(51);
    store.put_assistant_session(record(id, 0)?).await?;
    for _ in 0..4 {
        store
            .append_assistant_transcript_event(&id, instant(0)?, frame("chunk"))
            .await?;
    }
    let after_one = store.assistant_transcript(&id, Some(1)).await?;
    let indices: Vec<u64> = after_one.iter().map(|event| event.index).collect();
    require(
        indices == vec![2, 3],
        format!("`after` is EXCLUSIVE: after=1 yields 2 and 3, got {indices:?}"),
    )?;
    let after_last = store.assistant_transcript(&id, Some(3)).await?;
    require(
        after_last.is_empty(),
        "after the last index there is nothing left, and that is an empty read, not an error",
    )?;
    let unknown = store.assistant_transcript(&session(52), None).await?;
    require(
        unknown.is_empty(),
        "reading an unknown session's transcript is an absence, not a refusal",
    )
}

async fn append_to_an_unknown_session_refuses(
    store: &Arc<dyn AssistantSessionStore>,
) -> Result<(), StoreError> {
    let outcome = store
        .append_assistant_transcript_event(&session(61), instant(0)?, frame("orphan"))
        .await;
    match outcome {
        Err(StoreError::AssistantSessionNotFound { session_id }) => require(
            session_id == session(61).to_string(),
            format!("the refusal names the session it was asked about, got {session_id}"),
        ),
        Err(other) => Err(StoreError::Backend(format!(
            "an append to an unknown session must refuse as not-found, got {other}"
        ))),
        Ok(index) => Err(StoreError::Backend(format!(
            "an append to an unknown session must refuse, not silently create one at index {index}"
        ))),
    }
}

async fn poisoned_row_is_listed_with_its_error(
    store: &Arc<dyn AssistantSessionStore>,
    write_raw: &AssistantSessionRawWrite,
) -> Result<(), StoreError> {
    let poisoned = session(71).to_string();
    write_raw(poisoned.clone(), b"{ this is not a session record".to_vec()).await?;
    let listing = store.list_assistant_sessions().await?;
    let row = listing
        .undecodable
        .iter()
        .find(|row| row.session_id == poisoned);
    match row {
        Some(row) => require(
            !row.error.is_empty(),
            "a poisoned row carries the decode failure that explains it",
        )?,
        None => {
            return Err(StoreError::Backend(format!(
                "the poisoned session {poisoned} must be LISTED as undecodable, never skipped"
            )));
        }
    }
    require(
        !listing
            .sessions
            .iter()
            .any(|record| record.session_id.to_string() == poisoned),
        "a poisoned row must not also appear as a decoded record",
    )?;
    // It is still a session for append purposes: refusing its transcript would
    // destroy the evidence an operator needs to work out what happened to it.
    store
        .append_assistant_transcript_event(&session(71), instant(0)?, frame("after poisoning"))
        .await
        .map(drop)
}

async fn restart_survival(
    store: Arc<dyn AssistantSessionStore>,
    reopen: Option<AssistantSessionReopen>,
) -> Result<(), StoreError> {
    let id = session(81);
    store.put_assistant_session(record(id, 0)?).await?;
    store
        .append_assistant_transcript_event(&id, instant(0)?, frame("before the restart"))
        .await?;
    // Written before the restart, read after it: the last pick is what the
    // console preselects on a fresh morning, so a memory that died with the
    // process would ask the operator to choose a harness they already chose.
    store
        .put_assistant_default_harness("ada", "claude-code")
        .await?;
    let Some(reopen) = reopen else {
        // The in-memory backend has no durable storage to reopen; saying so is
        // the honest exemption, not a skipped assertion.
        return Ok(());
    };
    drop(store);
    let reopened = reopen().await?;
    let read = reopened.get_assistant_session(&id).await?;
    require(
        read.is_some(),
        "a session record survives a restart of the store that holds it",
    )?;
    let events = reopened.assistant_transcript(&id, None).await?;
    require(
        events.len() == 1 && events[0].index == 0,
        format!("the transcript survives the restart with its indices, got {events:?}"),
    )?;
    require(
        reopened.assistant_default_harness("ada").await?.as_deref() == Some("claude-code"),
        "the caller's last harness pick survives a restart of the store that holds it",
    )?;
    // And the NEXT index continues from where it stopped: a transcript that
    // restarted at zero after a reopen would silently overwrite its own past.
    let next = reopened
        .append_assistant_transcript_event(&id, instant(1)?, frame("after the restart"))
        .await?;
    require(
        next == 1,
        format!("the index continues dense across a restart, got {next}"),
    )
}