Skip to main content

koda_core/
persistence.rs

1//! Persistence trait — the storage contract for koda.
2//!
3//! Types and trait definition for the storage layer. The engine
4//! depends on this trait, not the concrete SQLite implementation.
5//!
6//! The default implementation is `Database` in `db.rs`.
7
8use anyhow::Result;
9use std::path::Path;
10
11/// Message roles in the conversation.
12#[derive(Debug, Clone, PartialEq)]
13#[allow(dead_code)]
14pub enum Role {
15    /// System prompt.
16    System,
17    /// User message.
18    User,
19    /// Assistant (LLM) response.
20    Assistant,
21    /// Tool result.
22    Tool,
23}
24
25impl Role {
26    /// String representation for database storage.
27    pub fn as_str(&self) -> &'static str {
28        match self {
29            Self::System => "system",
30            Self::User => "user",
31            Self::Assistant => "assistant",
32            Self::Tool => "tool",
33        }
34    }
35}
36
37impl std::fmt::Display for Role {
38    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
39        write!(f, "{}", self.as_str())
40    }
41}
42
43impl std::str::FromStr for Role {
44    type Err = String;
45    fn from_str(s: &str) -> Result<Self, Self::Err> {
46        match s {
47            "system" => Ok(Self::System),
48            "user" => Ok(Self::User),
49            "assistant" => Ok(Self::Assistant),
50            "tool" => Ok(Self::Tool),
51            other => Err(format!("unknown role: {other}")),
52        }
53    }
54}
55
56/// A stored message row.
57#[derive(Debug, Clone)]
58#[allow(dead_code)]
59pub struct Message {
60    /// Database row ID.
61    pub id: i64,
62    /// Session this message belongs to.
63    pub session_id: String,
64    /// Message role (system, user, assistant, tool).
65    pub role: Role,
66    /// Text content.
67    pub content: Option<String>,
68    /// Serialized tool calls JSON.
69    pub tool_calls: Option<String>,
70    /// ID of the tool call this responds to.
71    pub tool_call_id: Option<String>,
72    /// Input tokens for this message.
73    pub prompt_tokens: Option<i64>,
74    /// Output tokens for this message.
75    pub completion_tokens: Option<i64>,
76    /// Cached input tokens.
77    pub cache_read_tokens: Option<i64>,
78    /// Tokens written to cache.
79    pub cache_creation_tokens: Option<i64>,
80    /// Reasoning/thinking tokens.
81    pub thinking_tokens: Option<i64>,
82}
83
84/// Token usage totals for a session.
85#[derive(Debug, Clone, Default)]
86pub struct SessionUsage {
87    /// Total input tokens.
88    pub prompt_tokens: i64,
89    /// Total output tokens.
90    pub completion_tokens: i64,
91    /// Total cached input tokens.
92    pub cache_read_tokens: i64,
93    /// Total tokens written to cache.
94    pub cache_creation_tokens: i64,
95    /// Total reasoning/thinking tokens.
96    pub thinking_tokens: i64,
97    /// Number of API calls made.
98    pub api_calls: i64,
99}
100
101/// Summary info for a stored session.
102#[derive(Debug, Clone)]
103pub struct SessionInfo {
104    /// Session identifier.
105    pub id: String,
106    /// Agent name for this session.
107    pub agent_name: String,
108    /// ISO 8601 creation timestamp.
109    pub created_at: String,
110    /// Total messages in the session.
111    pub message_count: i64,
112    /// Cumulative token count.
113    pub total_tokens: i64,
114}
115
116/// Stats about compacted (archived) messages in the database.
117#[derive(Debug, Clone, Default)]
118pub struct CompactedStats {
119    /// Number of compacted messages.
120    pub message_count: i64,
121    /// Number of sessions with compacted messages.
122    pub session_count: i64,
123    /// Approximate size in bytes of compacted message content.
124    pub size_bytes: i64,
125    /// ISO 8601 timestamp of the oldest compacted message.
126    pub oldest: Option<String>,
127}
128
129/// Core storage contract for sessions, messages, and metadata.
130#[async_trait::async_trait]
131pub trait Persistence: Send + Sync {
132    // ── Sessions ──
133
134    /// Create a new session, returning its unique ID.
135    async fn create_session(&self, agent_name: &str, project_root: &Path) -> Result<String>;
136    /// List recent sessions for the given project root.
137    async fn list_sessions(&self, limit: i64, project_root: &Path) -> Result<Vec<SessionInfo>>;
138    /// Delete a session by ID. Returns `true` if it existed.
139    async fn delete_session(&self, session_id: &str) -> Result<bool>;
140
141    // ── Messages ──
142
143    /// Insert a message into a session.
144    async fn insert_message(
145        &self,
146        session_id: &str,
147        role: &Role,
148        content: Option<&str>,
149        tool_calls: Option<&str>,
150        tool_call_id: Option<&str>,
151        usage: Option<&crate::providers::TokenUsage>,
152    ) -> Result<i64>;
153
154    /// Insert a message with an explicit agent name (for sub-agent tracking).
155    #[allow(clippy::too_many_arguments)]
156    async fn insert_message_with_agent(
157        &self,
158        session_id: &str,
159        role: &Role,
160        content: Option<&str>,
161        tool_calls: Option<&str>,
162        tool_call_id: Option<&str>,
163        usage: Option<&crate::providers::TokenUsage>,
164        agent_name: Option<&str>,
165    ) -> Result<i64>;
166
167    /// Load active (non-compacted) conversation context for a session.
168    async fn load_context(&self, session_id: &str) -> Result<Vec<Message>>;
169    /// Load all messages in a session (no token limit).
170    async fn load_all_messages(&self, session_id: &str) -> Result<Vec<Message>>;
171    /// Recent user messages across all sessions (for startup hints).
172    async fn recent_user_messages(&self, limit: i64) -> Result<Vec<String>>;
173    /// Last assistant message in a session.
174    async fn last_assistant_message(&self, session_id: &str) -> Result<String>;
175    /// Last user message in a session.
176    async fn last_user_message(&self, session_id: &str) -> Result<String>;
177    /// Check if the session has unresolved tool calls.
178    async fn has_pending_tool_calls(&self, session_id: &str) -> Result<bool>;
179
180    // ── Token usage ──
181
182    /// Token usage totals for a session.
183    async fn session_token_usage(&self, session_id: &str) -> Result<SessionUsage>;
184    /// Token usage broken down by agent name.
185    async fn session_usage_by_agent(&self, session_id: &str)
186    -> Result<Vec<(String, SessionUsage)>>;
187
188    // ── Compaction ──
189
190    /// Compact old messages into a summary, preserving the last N messages.
191    async fn compact_session(
192        &self,
193        session_id: &str,
194        summary: &str,
195        preserve_count: usize,
196    ) -> Result<usize>;
197
198    // ── Purge ──
199
200    /// Stats about compacted (archived) messages across all sessions.
201    async fn compacted_stats(&self) -> Result<CompactedStats>;
202    /// Permanently delete compacted messages older than `min_age_days`.
203    /// Returns the number of messages deleted.
204    async fn purge_compacted(&self, min_age_days: u32) -> Result<usize>;
205
206    // ── Metadata ──
207
208    /// Get a session metadata value by key.
209    async fn get_metadata(&self, session_id: &str, key: &str) -> Result<Option<String>>;
210    /// Set a session metadata value.
211    async fn set_metadata(&self, session_id: &str, key: &str, value: &str) -> Result<()>;
212    /// Get the TODO list for a session.
213    async fn get_todo(&self, session_id: &str) -> Result<Option<String>>;
214    /// Set the TODO list for a session.
215    async fn set_todo(&self, session_id: &str, content: &str) -> Result<()>;
216}