use std::time::Duration;
use jiff::Timestamp;
use serde::{Deserialize, Serialize};
use crate::agent::TokenUsage;
use crate::caps::ToolPolicy;
use crate::id::{ChannelId, ConversationKey, MessageId, RunId, UserId};
use crate::{BoxError, async_trait};
#[derive(Debug, thiserror::Error)]
#[non_exhaustive]
pub enum StoreError {
#[error("the database schema version {found} is newer than this okbot supports ({supported})")]
TooNew {
found: u32,
supported: u32,
},
#[error("storage error: {0}")]
Backend(#[source] BoxError),
}
pub const PAIRING_MAX_FAILURES: u32 = 5;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "kebab-case")]
pub enum MessageRole {
User,
Bot,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct NewMessage {
pub conversation: ConversationKey,
pub role: MessageRole,
pub author_id: Option<UserId>,
pub author_name: String,
pub text: String,
pub message_id: Option<MessageId>,
pub at: Timestamp,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct StoredMessage {
pub role: MessageRole,
pub author_name: String,
pub text: String,
pub at: Timestamp,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AgentRunRecord {
pub run_id: RunId,
pub conversation: ConversationKey,
pub backend: String,
pub model: Option<String>,
pub tools: ToolPolicy,
pub usage: Option<TokenUsage>,
pub duration: Duration,
pub error: Option<String>,
pub started_at: Timestamp,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum FireState {
Sending,
Sent,
Failed,
Missed,
}
impl FireState {
pub fn as_str(self) -> &'static str {
match self {
Self::Sending => "sending",
Self::Sent => "sent",
Self::Failed => "failed",
Self::Missed => "missed",
}
}
pub fn parse(value: &str) -> Option<Self> {
match value {
"sending" => Some(Self::Sending),
"sent" => Some(Self::Sent),
"failed" => Some(Self::Failed),
"missed" => Some(Self::Missed),
_ => None,
}
}
pub fn is_final(self) -> bool {
!matches!(self, Self::Sending)
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Fire {
pub task_id: String,
pub fire_at: Timestamp,
pub state: FireState,
pub attempts: u32,
pub updated_at: Timestamp,
}
#[async_trait]
pub trait ConversationStore: Send + Sync {
async fn mark_seen(
&self,
channel: &ChannelId,
event_id: &str,
ttl: Duration,
) -> Result<bool, StoreError>;
async fn append(&self, message: NewMessage) -> Result<(), StoreError>;
async fn recent(
&self,
conversation: &ConversationKey,
limit: usize,
) -> Result<Vec<StoredMessage>, StoreError>;
async fn reset(&self, conversation: &ConversationKey) -> Result<(), StoreError>;
async fn session(
&self,
conversation: &ConversationKey,
) -> Result<Option<(String, String)>, StoreError>;
async fn set_session(
&self,
conversation: &ConversationKey,
backend: &str,
session: &str,
) -> Result<(), StoreError>;
async fn record_run(&self, run: AgentRunRecord) -> Result<(), StoreError>;
async fn owners(&self) -> Result<Vec<(ChannelId, UserId)>, StoreError>;
async fn remove_owner(&self, channel: &ChannelId, user: &UserId) -> Result<bool, StoreError>;
async fn create_pairing_code(&self, code: &str, ttl: Duration) -> Result<(), StoreError>;
async fn redeem_pairing_code(
&self,
code: &str,
channel: &ChannelId,
user: &UserId,
) -> Result<bool, StoreError>;
async fn claim_fire(&self, task_id: &str, fire_at: Timestamp) -> Result<Fire, StoreError>;
async fn set_fire(
&self,
task_id: &str,
fire_at: Timestamp,
state: FireState,
attempts: u32,
) -> Result<(), StoreError>;
async fn recent_fires(&self, limit: usize) -> Result<Vec<Fire>, StoreError>;
async fn set_direct_chat(
&self,
channel: &ChannelId,
user: &UserId,
conversation: &ConversationKey,
) -> Result<(), StoreError>;
async fn direct_chat(
&self,
channel: &ChannelId,
user: &UserId,
) -> Result<Option<ConversationKey>, StoreError>;
async fn purge_before(&self, before: Timestamp) -> Result<u64, StoreError>;
}