okbot_testkit/

memory_store.rs

//! An in-memory [`ConversationStore`].

use std::collections::HashMap;
use std::sync::{Arc, Mutex};
use std::time::Duration;

use jiff::{SignedDuration, Timestamp};
use okbot_core::{
    AgentRunRecord, AutomationRun, ChannelId, Clock, ConversationKey, ConversationStore, Fire,
    FireState, NewMessage, PAIRING_MAX_FAILURES, PeerTaskRecord, PeerTaskState, RoleGrant,
    SealedEntry, StoreError, StoredMessage, ToolCallRecord, Usage, UserId, async_trait,
};

#[derive(Default)]
struct State {
    seen: HashMap<(ChannelId, String), Timestamp>,
    messages: Vec<(u64, NewMessage)>,
    reset_after: HashMap<ConversationKey, u64>,
    sessions: HashMap<ConversationKey, (String, String)>,
    runs: Vec<AgentRunRecord>,
    next: u64,
    owners: Vec<(ChannelId, UserId)>,
    /// code → (expires, failures)
    codes: HashMap<String, (Timestamp, u32)>,
    fires: HashMap<(String, Timestamp), Fire>,
    direct: HashMap<(ChannelId, UserId), ConversationKey>,
    tool_calls: Vec<ToolCallRecord>,
    grants: HashMap<(String, String), RoleGrant>,
    /// (kind, server, subject) → (sealed, updated)
    sealed: HashMap<(String, String, String), (Vec<u8>, Timestamp)>,
    /// (automation, item) → seen
    items: HashMap<(String, String), Timestamp>,
    automation_runs: Vec<AutomationRun>,
    /// (inbound, id) → task
    peer_tasks: HashMap<(bool, String), PeerTaskRecord>,
}

/// A [`ConversationStore`] kept in memory, for runtime tests.
pub struct MemoryStore {
    state: Mutex<State>,
    clock: Arc<dyn Clock>,
}

impl std::fmt::Debug for MemoryStore {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("MemoryStore").finish_non_exhaustive()
    }
}

impl MemoryStore {
    /// An empty store reading time from `clock`.
    pub fn new(clock: Arc<dyn Clock>) -> Self {
        Self {
            state: Mutex::new(State::default()),
            clock,
        }
    }

    fn lock(&self) -> std::sync::MutexGuard<'_, State> {
        self.state
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner)
    }

    /// Recorded agent runs, oldest first.
    pub fn runs(&self) -> Vec<AgentRunRecord> {
        self.lock().runs.clone()
    }

    /// Recorded tool calls, oldest first.
    pub fn recorded_tool_calls(&self) -> Vec<ToolCallRecord> {
        self.lock().tool_calls.clone()
    }
}

fn signed(d: Duration) -> SignedDuration {
    SignedDuration::try_from(d).unwrap_or(SignedDuration::MAX)
}

#[async_trait]
impl ConversationStore for MemoryStore {
    async fn mark_seen(
        &self,
        channel: &ChannelId,
        event_id: &str,
        ttl: Duration,
    ) -> Result<bool, StoreError> {
        let now = self.clock.now();
        let mut state = self.lock();
        let key = (channel.clone(), event_id.to_string());
        if state.seen.get(&key).is_some_and(|expires| *expires > now) {
            return Ok(false);
        }
        state
            .seen
            .insert(key, now.saturating_add(signed(ttl)).unwrap_or(now));
        Ok(true)
    }

    async fn append(&self, message: NewMessage) -> Result<(), StoreError> {
        let mut state = self.lock();
        state.next += 1;
        let id = state.next;
        state.messages.push((id, message));
        Ok(())
    }

    async fn recent(
        &self,
        conversation: &ConversationKey,
        limit: usize,
    ) -> Result<Vec<StoredMessage>, StoreError> {
        let state = self.lock();
        let after = state.reset_after.get(conversation).copied().unwrap_or(0);
        let mut out: Vec<StoredMessage> = state
            .messages
            .iter()
            .rev()
            .filter(|(id, m)| *id > after && m.conversation == *conversation)
            .take(limit)
            .map(|(_, m)| StoredMessage {
                role: m.role,
                author_name: m.author_name.clone(),
                text: m.text.clone(),
                at: m.at,
            })
            .collect();
        out.reverse();
        Ok(out)
    }

    async fn reset(&self, conversation: &ConversationKey) -> Result<(), StoreError> {
        let mut state = self.lock();
        let last = state.next;
        state.reset_after.insert(conversation.clone(), last);
        state.sessions.remove(conversation);
        Ok(())
    }

    async fn session(
        &self,
        conversation: &ConversationKey,
    ) -> Result<Option<(String, String)>, StoreError> {
        Ok(self.lock().sessions.get(conversation).cloned())
    }

    async fn set_session(
        &self,
        conversation: &ConversationKey,
        backend: &str,
        session: &str,
    ) -> Result<(), StoreError> {
        self.lock().sessions.insert(
            conversation.clone(),
            (backend.to_string(), session.to_string()),
        );
        Ok(())
    }

    async fn owners(&self) -> Result<Vec<(ChannelId, UserId)>, StoreError> {
        let mut owners = self.lock().owners.clone();
        owners.sort();
        Ok(owners)
    }

    async fn remove_owner(&self, channel: &ChannelId, user: &UserId) -> Result<bool, StoreError> {
        let mut state = self.lock();
        let before = state.owners.len();
        state.owners.retain(|(c, u)| c != channel || u != user);
        Ok(state.owners.len() != before)
    }

    async fn create_pairing_code(&self, code: &str, ttl: Duration) -> Result<(), StoreError> {
        let expires = self.clock.now() + signed(ttl);
        self.lock().codes.insert(code.to_string(), (expires, 0));
        Ok(())
    }

    async fn redeem_pairing_code(
        &self,
        code: &str,
        channel: &ChannelId,
        user: &UserId,
    ) -> Result<bool, StoreError> {
        let now = self.clock.now();
        let mut state = self.lock();
        state.codes.retain(|_, (expires, _)| *expires > now);
        if state.codes.remove(code).is_some() {
            let account = (channel.clone(), user.clone());
            if !state.owners.contains(&account) {
                state.owners.push(account);
            }
            return Ok(true);
        }
        for (_, failures) in state.codes.values_mut() {
            *failures += 1;
        }
        state
            .codes
            .retain(|_, (_, failures)| *failures < PAIRING_MAX_FAILURES);
        Ok(false)
    }

    async fn record_run(&self, run: AgentRunRecord) -> Result<(), StoreError> {
        self.lock().runs.push(run);
        Ok(())
    }

    async fn claim_fire(&self, task_id: &str, fire_at: Timestamp) -> Result<Fire, StoreError> {
        let now = self.clock.now();
        Ok(self
            .lock()
            .fires
            .entry((task_id.to_string(), fire_at))
            .or_insert_with(|| Fire {
                task_id: task_id.to_string(),
                fire_at,
                state: FireState::Sending,
                attempts: 0,
                updated_at: now,
            })
            .clone())
    }

    async fn set_fire(
        &self,
        task_id: &str,
        fire_at: Timestamp,
        state: FireState,
        attempts: u32,
    ) -> Result<(), StoreError> {
        let fire = Fire {
            task_id: task_id.to_string(),
            fire_at,
            state,
            attempts,
            updated_at: self.clock.now(),
        };
        self.lock()
            .fires
            .insert((task_id.to_string(), fire_at), fire);
        Ok(())
    }

    async fn recent_fires(&self, limit: usize) -> Result<Vec<Fire>, StoreError> {
        let mut fires: Vec<Fire> = self.lock().fires.values().cloned().collect();
        fires.sort_by_key(|f| std::cmp::Reverse((f.updated_at, f.fire_at)));
        fires.truncate(limit);
        Ok(fires)
    }

    async fn set_direct_chat(
        &self,
        channel: &ChannelId,
        user: &UserId,
        conversation: &ConversationKey,
    ) -> Result<(), StoreError> {
        self.lock()
            .direct
            .insert((channel.clone(), user.clone()), conversation.clone());
        Ok(())
    }

    async fn direct_chat(
        &self,
        channel: &ChannelId,
        user: &UserId,
    ) -> Result<Option<ConversationKey>, StoreError> {
        Ok(self
            .lock()
            .direct
            .get(&(channel.clone(), user.clone()))
            .cloned())
    }

    async fn usage_since(&self, person: &str, since: Timestamp) -> Result<Usage, StoreError> {
        let state = self.lock();
        let mut usage = Usage::default();
        for run in state
            .runs
            .iter()
            .filter(|r| r.person.as_deref() == Some(person) && r.started_at >= since)
        {
            usage.turns += 1;
            usage.tokens += run.usage.map_or(0, |u| u.input + u.output);
        }
        Ok(usage)
    }

    async fn record_tool_call(&self, call: ToolCallRecord) -> Result<(), StoreError> {
        self.lock().tool_calls.push(call);
        Ok(())
    }

    async fn tool_calls(
        &self,
        since: Timestamp,
        limit: usize,
    ) -> Result<Vec<ToolCallRecord>, StoreError> {
        Ok(self
            .lock()
            .tool_calls
            .iter()
            .rev()
            .filter(|c| c.at >= since)
            .take(limit)
            .cloned()
            .collect())
    }

    async fn grant_role(&self, grant: RoleGrant) -> Result<(), StoreError> {
        self.lock()
            .grants
            .insert((grant.role.clone(), grant.account.clone()), grant);
        Ok(())
    }

    async fn revoke_role(&self, role: &str, account: &str) -> Result<bool, StoreError> {
        let now = self.clock.now();
        Ok(self
            .lock()
            .grants
            .remove(&(role.to_string(), account.to_string()))
            .is_some_and(|g| g.expires_at > now))
    }

    async fn role_grants(&self, now: Timestamp) -> Result<Vec<RoleGrant>, StoreError> {
        let mut grants: Vec<RoleGrant> = self
            .lock()
            .grants
            .values()
            .filter(|g| g.expires_at > now)
            .cloned()
            .collect();
        grants.sort_by(|a, b| (&a.role, &a.account).cmp(&(&b.role, &b.account)));
        Ok(grants)
    }

    async fn purge_before(&self, before: Timestamp) -> Result<u64, StoreError> {
        let now = self.clock.now();
        let mut state = self.lock();
        let count = |s: &State| {
            s.messages.len()
                + s.runs.len()
                + s.seen.len()
                + s.fires.len()
                + s.tool_calls.len()
                + s.grants.len()
                + s.items.len()
                + s.automation_runs.len()
                + s.peer_tasks.len()
        };
        let start = count(&state);
        state.messages.retain(|(_, m)| m.at >= before);
        state.runs.retain(|r| r.started_at >= before);
        state.tool_calls.retain(|c| c.at >= before);
        state.grants.retain(|_, g| g.expires_at > now);
        state.items.retain(|_, seen| *seen >= before);
        state.automation_runs.retain(|r| r.at >= before);
        state.peer_tasks.retain(|_, t| t.updated_at >= before);
        state.seen.retain(|_, expires| *expires > now);
        state
            .fires
            .retain(|_, f| f.updated_at >= before || f.state == FireState::Sending);
        let end = count(&state);
        Ok((start - end) as u64)
    }

    async fn sealed(
        &self,
        kind: &str,
        server: &str,
        subject: &str,
    ) -> Result<Option<Vec<u8>>, StoreError> {
        let key = (kind.to_string(), server.to_string(), subject.to_string());
        Ok(self.lock().sealed.get(&key).map(|(v, _)| v.clone()))
    }

    async fn set_sealed(
        &self,
        kind: &str,
        server: &str,
        subject: &str,
        sealed: Vec<u8>,
    ) -> Result<(), StoreError> {
        let key = (kind.to_string(), server.to_string(), subject.to_string());
        let now = self.clock.now();
        self.lock().sealed.insert(key, (sealed, now));
        Ok(())
    }

    async fn delete_sealed(
        &self,
        kind: &str,
        server: &str,
        subject: &str,
    ) -> Result<bool, StoreError> {
        let key = (kind.to_string(), server.to_string(), subject.to_string());
        Ok(self.lock().sealed.remove(&key).is_some())
    }

    async fn sealed_entries(&self, kind: &str) -> Result<Vec<SealedEntry>, StoreError> {
        let mut out: Vec<SealedEntry> = self
            .lock()
            .sealed
            .iter()
            .filter(|((k, _, _), _)| k == kind)
            .map(|((k, server, subject), (_, at))| SealedEntry {
                kind: k.clone(),
                server: server.clone(),
                subject: subject.clone(),
                updated_at: *at,
            })
            .collect();
        out.sort_by(|a, b| (&a.server, &a.subject).cmp(&(&b.server, &b.subject)));
        Ok(out)
    }

    async fn claim_items(
        &self,
        automation: &str,
        items: Vec<String>,
    ) -> Result<Vec<String>, StoreError> {
        let now = self.clock.now();
        let mut state = self.lock();
        let mut new = Vec::new();
        for item in items {
            let fresh = state
                .items
                .insert((automation.to_string(), item.clone()), now)
                .is_none();
            if fresh {
                new.push(item);
            }
        }
        Ok(new)
    }

    async fn release_items(&self, automation: &str, items: Vec<String>) -> Result<(), StoreError> {
        let mut state = self.lock();
        for item in items {
            state.items.remove(&(automation.to_string(), item));
        }
        Ok(())
    }

    async fn record_automation_run(&self, run: AutomationRun) -> Result<(), StoreError> {
        self.lock().automation_runs.push(run);
        Ok(())
    }

    async fn automation_runs(
        &self,
        automation: Option<&str>,
        since: Timestamp,
        limit: usize,
    ) -> Result<Vec<AutomationRun>, StoreError> {
        Ok(self
            .lock()
            .automation_runs
            .iter()
            .rev()
            .filter(|r| r.at >= since && automation.is_none_or(|a| r.automation == a))
            .take(limit)
            .cloned()
            .collect())
    }

    async fn put_peer_task(&self, task: PeerTaskRecord) -> Result<(), StoreError> {
        let mut state = self.lock();
        let key = (task.inbound, task.id.clone());
        let duplicate = state.peer_tasks.values().any(|t| {
            t.inbound == task.inbound
                && t.team == task.team
                && t.peer == task.peer
                && t.message_id == task.message_id
                && t.id != task.id
        });
        if duplicate {
            return Err(StoreError::Backend("duplicate peer message".into()));
        }
        match state.peer_tasks.get_mut(&key) {
            Some(old) => {
                old.state = task.state;
                old.conversation = task.conversation;
                old.text = task.text;
                old.tokens = task.tokens;
                old.updated_at = task.updated_at;
            }
            None => {
                state.peer_tasks.insert(key, task);
            }
        }
        Ok(())
    }

    async fn peer_task(
        &self,
        inbound: bool,
        id: &str,
    ) -> Result<Option<PeerTaskRecord>, StoreError> {
        Ok(self
            .lock()
            .peer_tasks
            .get(&(inbound, id.to_string()))
            .cloned())
    }

    async fn peer_task_by_message(
        &self,
        inbound: bool,
        team: &str,
        peer: &str,
        message_id: &str,
    ) -> Result<Option<PeerTaskRecord>, StoreError> {
        Ok(self
            .lock()
            .peer_tasks
            .values()
            .find(|t| {
                t.inbound == inbound
                    && t.team == team
                    && t.peer == peer
                    && t.message_id == message_id
            })
            .cloned())
    }

    async fn peer_tasks(&self, limit: usize) -> Result<Vec<PeerTaskRecord>, StoreError> {
        let mut tasks: Vec<PeerTaskRecord> = self.lock().peer_tasks.values().cloned().collect();
        tasks.sort_by_key(|t| std::cmp::Reverse((t.updated_at, t.created_at)));
        tasks.truncate(limit);
        Ok(tasks)
    }

    async fn interrupt_peer_tasks(&self, now: Timestamp) -> Result<u64, StoreError> {
        let mut changed = 0;
        for task in self.lock().peer_tasks.values_mut() {
            if task.inbound && !task.state.is_final() {
                task.state = PeerTaskState::Failed;
                task.text = Some("interrupted: the bot stopped".into());
                task.updated_at = now;
                changed += 1;
            }
        }
        Ok(changed)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::ManualClock;

    #[tokio::test]
    async fn conformance() {
        let clock = Arc::new(ManualClock::default());
        let store = MemoryStore::new(clock.clone());
        crate::conformance::conversation_store(&store, &clock).await;
    }
}