Skip to main content

okbot_core/
store.rs

1//! The conversation store port.
2
3use std::time::Duration;
4
5use jiff::Timestamp;
6use serde::{Deserialize, Serialize};
7
8use crate::agent::TokenUsage;
9use crate::caps::ToolPolicy;
10use crate::id::{ChannelId, ConversationKey, MessageId, RunId, UserId};
11use crate::{BoxError, async_trait};
12
13/// A storage failure.
14#[derive(Debug, thiserror::Error)]
15#[non_exhaustive]
16pub enum StoreError {
17    /// The database was written by a newer okbot.
18    #[error("the database schema version {found} is newer than this okbot supports ({supported})")]
19    TooNew {
20        /// Version in the database.
21        found: u32,
22        /// Newest version this build knows.
23        supported: u32,
24    },
25    /// Any other failure.
26    #[error("storage error: {0}")]
27    Backend(#[source] BoxError),
28}
29
30/// Wrong pairing codes tolerated before pending codes are cancelled.
31pub const PAIRING_MAX_FAILURES: u32 = 5;
32
33/// Who wrote a stored message.
34#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
35#[serde(rename_all = "kebab-case")]
36pub enum MessageRole {
37    /// A chat user.
38    User,
39    /// The bot.
40    Bot,
41}
42
43/// A message to store.
44#[derive(Debug, Clone, PartialEq, Eq)]
45pub struct NewMessage {
46    /// Conversation.
47    pub conversation: ConversationKey,
48    /// Author role.
49    pub role: MessageRole,
50    /// Author id (the user, or the bot's own id when known).
51    pub author_id: Option<UserId>,
52    /// Author display name.
53    pub author_name: String,
54    /// Text (canonical Markdown for bot messages).
55    pub text: String,
56    /// Platform message id.
57    pub message_id: Option<MessageId>,
58    /// When it was sent.
59    pub at: Timestamp,
60}
61
62/// A stored message.
63#[derive(Debug, Clone, PartialEq, Eq)]
64pub struct StoredMessage {
65    /// Author role.
66    pub role: MessageRole,
67    /// Author display name.
68    pub author_name: String,
69    /// Text.
70    pub text: String,
71    /// When it was sent.
72    pub at: Timestamp,
73}
74
75/// One agent run, for the audit table.
76#[derive(Debug, Clone, PartialEq, Eq)]
77pub struct AgentRunRecord {
78    /// Run id.
79    pub run_id: RunId,
80    /// Conversation.
81    pub conversation: ConversationKey,
82    /// Backend that produced the result (or the last one tried).
83    pub backend: String,
84    /// Model, when known.
85    pub model: Option<String>,
86    /// Tool policy of the run.
87    pub tools: ToolPolicy,
88    /// Token usage, when reported.
89    pub usage: Option<TokenUsage>,
90    /// Duration.
91    pub duration: Duration,
92    /// Error text, when the run failed.
93    pub error: Option<String>,
94    /// When the run started.
95    pub started_at: Timestamp,
96    /// Who the run was for ([`person_key`]), counted against `[limits]`; `None` for runs nobody
97    /// asked for in chat.
98    pub person: Option<String>,
99}
100
101/// The key a person's usage and audit entries are stored under: `person:<id>` for linked
102/// accounts (`[people]`, or the owner of a personal bot), else the account
103/// `platform:instance/user-id`.
104pub fn person_key(channel: &ChannelId, user: &UserId, person: Option<&str>) -> String {
105    match person {
106        Some(id) => format!("person:{id}"),
107        None => format!("{channel}/{user}"),
108    }
109}
110
111/// What one person used over a period.
112#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
113pub struct Usage {
114    /// Agent turns.
115    pub turns: u64,
116    /// Input plus output tokens, as reported.
117    pub tokens: u64,
118}
119
120/// What happened to a tool call (ADR 0011).
121#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
122#[non_exhaustive]
123pub enum ToolDecision {
124    /// Sent to the server.
125    Allowed,
126    /// Refused by the rules (role, room, tool lists).
127    Denied,
128    /// A person declined it (or did not answer) when asked to confirm.
129    Declined,
130}
131
132impl ToolDecision {
133    /// The name stored and shown.
134    pub fn as_str(self) -> &'static str {
135        match self {
136            Self::Allowed => "allowed",
137            Self::Denied => "denied",
138            Self::Declined => "declined",
139        }
140    }
141
142    /// Parses [`ToolDecision::as_str`].
143    pub fn parse(value: &str) -> Option<Self> {
144        match value {
145            "allowed" => Some(Self::Allowed),
146            "denied" => Some(Self::Denied),
147            "declined" => Some(Self::Declined),
148            _ => None,
149        }
150    }
151}
152
153/// One call of an `[[mcp]]` tool, by an agent or a chat command. Arguments and results are not
154/// kept: they may hold secrets.
155#[derive(Debug, Clone, PartialEq, Eq)]
156pub struct ToolCallRecord {
157    /// The agent run, when an agent made the call.
158    pub run_id: Option<RunId>,
159    /// The `[[command]]`, when a person ran one.
160    pub command: Option<String>,
161    /// Who it was for ([`person_key`]).
162    pub person: String,
163    /// Where.
164    pub conversation: ConversationKey,
165    /// The `[[mcp]]` server.
166    pub server: String,
167    /// The tool.
168    pub tool: String,
169    /// What happened.
170    pub decision: ToolDecision,
171    /// The transport error or the tool's own error flag, when the call failed.
172    pub error: Option<String>,
173    /// Time from the request to the answer (including a confirmation).
174    pub duration: Duration,
175    /// When the call was made.
176    pub at: Timestamp,
177}
178
179/// A role given to an account for a limited time (`okbot role grant`, `/grant`).
180#[derive(Debug, Clone, PartialEq, Eq)]
181pub struct RoleGrant {
182    /// A role from `[roles]`.
183    pub role: String,
184    /// The account, as `platform:instance/user-id`.
185    pub account: String,
186    /// Who granted it (an account, or `operator` for the terminal).
187    pub granted_by: String,
188    /// When it was granted.
189    pub granted_at: Timestamp,
190    /// When it ends.
191    pub expires_at: Timestamp,
192}
193
194/// Where one firing of a task is (design §5.6).
195#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
196pub enum FireState {
197    /// Claimed; being delivered or waiting for a retry.
198    Sending,
199    /// Delivered.
200    Sent,
201    /// Gave up after retries.
202    Failed,
203    /// Too late to deliver (the bot was down).
204    Missed,
205}
206
207impl FireState {
208    /// The name stored and shown.
209    pub fn as_str(self) -> &'static str {
210        match self {
211            Self::Sending => "sending",
212            Self::Sent => "sent",
213            Self::Failed => "failed",
214            Self::Missed => "missed",
215        }
216    }
217
218    /// Parses [`FireState::as_str`].
219    pub fn parse(value: &str) -> Option<Self> {
220        match value {
221            "sending" => Some(Self::Sending),
222            "sent" => Some(Self::Sent),
223            "failed" => Some(Self::Failed),
224            "missed" => Some(Self::Missed),
225            _ => None,
226        }
227    }
228
229    /// No more delivery attempts follow.
230    pub fn is_final(self) -> bool {
231        !matches!(self, Self::Sending)
232    }
233}
234
235/// One firing of a task: the task id and its planned time identify it.
236#[derive(Debug, Clone, PartialEq, Eq)]
237pub struct Fire {
238    /// Task id (`tasks/people/owner/t-7k2m9q`).
239    pub task_id: String,
240    /// The planned time.
241    pub fire_at: Timestamp,
242    /// Where it is.
243    pub state: FireState,
244    /// Delivery attempts so far.
245    pub attempts: u32,
246    /// Last change.
247    pub updated_at: Timestamp,
248}
249
250/// Stores conversations, deduplication marks and the audit trail (design §6.1).
251#[async_trait]
252pub trait ConversationStore: Send + Sync {
253    /// Records an inbound event id. Returns `false` when it was already seen within `ttl`.
254    async fn mark_seen(
255        &self,
256        channel: &ChannelId,
257        event_id: &str,
258        ttl: Duration,
259    ) -> Result<bool, StoreError>;
260
261    /// Appends a message to a conversation.
262    async fn append(&self, message: NewMessage) -> Result<(), StoreError>;
263
264    /// The newest `limit` messages after the last reset, oldest first.
265    async fn recent(
266        &self,
267        conversation: &ConversationKey,
268        limit: usize,
269    ) -> Result<Vec<StoredMessage>, StoreError>;
270
271    /// Starts a fresh history for the conversation (`/reset`). Old messages stay for audit
272    /// until retention removes them.
273    async fn reset(&self, conversation: &ConversationKey) -> Result<(), StoreError>;
274
275    /// The native agent session of a conversation: `(backend, session id)`.
276    async fn session(
277        &self,
278        conversation: &ConversationKey,
279    ) -> Result<Option<(String, String)>, StoreError>;
280
281    /// Remembers the native agent session of a conversation (cleared by [`ConversationStore::reset`]).
282    async fn set_session(
283        &self,
284        conversation: &ConversationKey,
285        backend: &str,
286        session: &str,
287    ) -> Result<(), StoreError>;
288
289    /// Records an agent run.
290    async fn record_run(&self, run: AgentRunRecord) -> Result<(), StoreError>;
291
292    /// Accounts of the owner of a personal bot (paired with [`ConversationStore::redeem_pairing_code`]).
293    async fn owners(&self) -> Result<Vec<(ChannelId, UserId)>, StoreError>;
294
295    /// Removes an owner account. Returns whether it was one.
296    async fn remove_owner(&self, channel: &ChannelId, user: &UserId) -> Result<bool, StoreError>;
297
298    /// Stores a one-time pairing code, valid for `ttl`.
299    async fn create_pairing_code(&self, code: &str, ttl: Duration) -> Result<(), StoreError>;
300
301    /// Consumes `code` and makes `(channel, user)` an owner account when the code is valid.
302    ///
303    /// A wrong code counts as a failed attempt against every pending code; after
304    /// [`PAIRING_MAX_FAILURES`] failures the pending codes are cancelled, so codes cannot be
305    /// guessed.
306    async fn redeem_pairing_code(
307        &self,
308        code: &str,
309        channel: &ChannelId,
310        user: &UserId,
311    ) -> Result<bool, StoreError>;
312
313    /// Claims a task firing: records it as [`FireState::Sending`] the first time and returns its
314    /// current record, so a firing that was already sent is never sent again.
315    async fn claim_fire(&self, task_id: &str, fire_at: Timestamp) -> Result<Fire, StoreError>;
316
317    /// Updates a claimed firing.
318    async fn set_fire(
319        &self,
320        task_id: &str,
321        fire_at: Timestamp,
322        state: FireState,
323        attempts: u32,
324    ) -> Result<(), StoreError>;
325
326    /// The newest `limit` firings, newest first (for `okbot doctor` and `okbot task`).
327    async fn recent_fires(&self, limit: usize) -> Result<Vec<Fire>, StoreError>;
328
329    /// Remembers the direct conversation of `user` on `channel` (reminders sent there).
330    async fn set_direct_chat(
331        &self,
332        channel: &ChannelId,
333        user: &UserId,
334        conversation: &ConversationKey,
335    ) -> Result<(), StoreError>;
336
337    /// The latest direct conversation of `user` on `channel`.
338    async fn direct_chat(
339        &self,
340        channel: &ChannelId,
341        user: &UserId,
342    ) -> Result<Option<ConversationKey>, StoreError>;
343
344    /// Turns and tokens of `person` ([`person_key`]) in runs started at or after `since`.
345    async fn usage_since(&self, person: &str, since: Timestamp) -> Result<Usage, StoreError>;
346
347    /// Records a tool call.
348    async fn record_tool_call(&self, call: ToolCallRecord) -> Result<(), StoreError>;
349
350    /// Tool calls made at or after `since`, newest first, at most `limit`.
351    async fn tool_calls(
352        &self,
353        since: Timestamp,
354        limit: usize,
355    ) -> Result<Vec<ToolCallRecord>, StoreError>;
356
357    /// Grants a role until `grant.expires_at`, replacing an earlier grant of the same role to the
358    /// same account.
359    async fn grant_role(&self, grant: RoleGrant) -> Result<(), StoreError>;
360
361    /// Ends a grant before it expires. Returns whether there was one.
362    async fn revoke_role(&self, role: &str, account: &str) -> Result<bool, StoreError>;
363
364    /// Grants that have not expired at `now`.
365    async fn role_grants(&self, now: Timestamp) -> Result<Vec<RoleGrant>, StoreError>;
366
367    /// Deletes messages, runs, tool calls, expired grants, dedupe marks and finished firings older
368    /// than `before`. Returns rows removed.
369    async fn purge_before(&self, before: Timestamp) -> Result<u64, StoreError>;
370}