aion_store/memory/

assistant.rs

//! The reference [`AssistantSessionStore`] over [`InMemoryStore`].
//!
//! The whole point of the reference implementation is to define the answer the
//! durable backend must agree with, so the two properties that matter are
//! written here as plainly as they can be:
//!
//! - **the store assigns the index, under one lock.** A transcript is a `Vec`
//!   and the index is its length at the moment of the push, taken while the map
//!   is held. Two appenders cannot both read the same length, so they cannot
//!   both be handed the same index.
//! - **an append to an unknown session refuses.** The session map is consulted
//!   before the transcript map is touched, and a missing record is
//!   [`StoreError::AssistantSessionNotFound`] rather than a transcript with no
//!   owner.

use std::collections::BTreeMap;
use std::sync::MutexGuard;

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

use super::InMemoryStore;
use crate::StoreError;
use crate::assistant::{
    AssistantSessionListing, AssistantSessionRecord, AssistantSessionStore,
    AssistantTranscriptEvent, UndecodableAssistantSession,
};

impl InMemoryStore {
    /// Write assistant-session backend bytes directly for persistence
    /// conformance testing (poisoned-row listing behaviour), bypassing the
    /// entity codec.
    ///
    /// # Errors
    ///
    /// [`StoreError::Serialization`] when `session_id` is empty or whitespace.
    pub fn write_raw_assistant_session(
        &self,
        session_id: &str,
        bytes: Vec<u8>,
    ) -> Result<(), StoreError> {
        if session_id.trim().is_empty() {
            return Err(StoreError::Serialization(
                "assistant session id must not be empty".to_owned(),
            ));
        }
        self.lock_assistant_sessions()
            .insert(session_id.to_owned(), bytes);
        Ok(())
    }

    /// Whether a session record exists, without decoding it.
    ///
    /// A poisoned row still counts as existing: the session is there, it is
    /// merely unreadable, and refusing its appends would destroy the very
    /// transcript an operator would use to work out what happened to it.
    fn assistant_session_exists(&self, session_id: &AssistantSessionId) -> bool {
        self.lock_assistant_sessions()
            .contains_key(&session_id.to_string())
    }
}

#[async_trait]
impl AssistantSessionStore for InMemoryStore {
    async fn put_assistant_session(
        &self,
        record: AssistantSessionRecord,
    ) -> Result<(), StoreError> {
        let bytes = record.encode()?;
        self.lock_assistant_sessions()
            .insert(record.session_id.to_string(), bytes);
        Ok(())
    }

    async fn get_assistant_session(
        &self,
        session_id: &AssistantSessionId,
    ) -> Result<Option<AssistantSessionRecord>, StoreError> {
        self.lock_assistant_sessions()
            .get(&session_id.to_string())
            .map(|bytes| AssistantSessionRecord::decode(bytes))
            .transpose()
    }

    async fn list_assistant_sessions(&self) -> Result<AssistantSessionListing, StoreError> {
        let mut listing = AssistantSessionListing::default();
        for (session_id, bytes) in self.lock_assistant_sessions().iter() {
            match AssistantSessionRecord::decode(bytes) {
                Ok(record) => listing.sessions.push(record),
                Err(error) => listing.undecodable.push(UndecodableAssistantSession {
                    session_id: session_id.clone(),
                    error: error.to_string(),
                }),
            }
        }
        listing.sort();
        Ok(listing)
    }

    async fn append_assistant_transcript_event(
        &self,
        session_id: &AssistantSessionId,
        recorded_at: DateTime<Utc>,
        payload: Payload,
    ) -> Result<u64, StoreError> {
        if !self.assistant_session_exists(session_id) {
            return Err(StoreError::AssistantSessionNotFound {
                session_id: session_id.to_string(),
            });
        }
        // ONE lock spans reading the length and pushing at it, which is what
        // makes the index the store's to give rather than a value two appenders
        // could both observe.
        let mut transcripts = self.lock_assistant_transcripts();
        let events = transcripts.entry(session_id.to_string()).or_default();
        let index = u64::try_from(events.len()).map_err(|error| {
            StoreError::Backend(format!(
                "assistant transcript for {session_id} is longer than a u64 index can name: {error}"
            ))
        })?;
        events.push(AssistantTranscriptEvent {
            index,
            recorded_at,
            payload,
        });
        Ok(index)
    }

    async fn assistant_transcript_head(
        &self,
        session_id: &AssistantSessionId,
    ) -> Result<u64, StoreError> {
        let head = self
            .lock_assistant_transcripts()
            .get(&session_id.to_string())
            .map_or(0, Vec::len);
        u64::try_from(head).map_err(|error| {
            StoreError::Backend(format!(
                "assistant transcript for {session_id} is longer than a u64 head can name: {error}"
            ))
        })
    }

    async fn put_assistant_default_harness(
        &self,
        subject: &str,
        harness: &str,
    ) -> Result<(), StoreError> {
        self.lock_assistant_default_harnesses()
            .insert(subject.to_owned(), harness.to_owned());
        Ok(())
    }

    async fn assistant_default_harness(&self, subject: &str) -> Result<Option<String>, StoreError> {
        Ok(self
            .lock_assistant_default_harnesses()
            .get(subject)
            .cloned())
    }

    async fn assistant_transcript(
        &self,
        session_id: &AssistantSessionId,
        after: Option<u64>,
    ) -> Result<Vec<AssistantTranscriptEvent>, StoreError> {
        Ok(self
            .lock_assistant_transcripts()
            .get(&session_id.to_string())
            .map(|events| {
                events
                    .iter()
                    .filter(|event| after.is_none_or(|bound| event.index > bound))
                    .cloned()
                    .collect()
            })
            .unwrap_or_default())
    }
}

impl InMemoryStore {
    /// Poison-tolerant lock over the assistant-session record map.
    fn lock_assistant_sessions(&self) -> MutexGuard<'_, BTreeMap<String, Vec<u8>>> {
        self.assistant_sessions
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner)
    }

    /// Poison-tolerant lock over the per-caller last-pick memory.
    fn lock_assistant_default_harnesses(&self) -> MutexGuard<'_, BTreeMap<String, String>> {
        self.assistant_default_harnesses
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner)
    }

    /// Poison-tolerant lock over the assistant transcripts.
    fn lock_assistant_transcripts(
        &self,
    ) -> MutexGuard<'_, BTreeMap<String, Vec<AssistantTranscriptEvent>>> {
        self.assistant_transcripts
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner)
    }
}