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//!
8//! ## Design (DESIGN.md)
9//!
10//! - **Database Backend: SQLite + Persistence Trait (P3)**: This trait
11//!   exists because P3 says the storage backend will change. But there's
12//!   only one real implementation because P1 says we don't need a second
13//!   one yet. When P1 and P3 conflict: P1 wins on timing, P3 wins on
14//!   architecture.
15
16use anyhow::Result;
17use std::path::Path;
18
19/// Message roles in the conversation.
20#[derive(Debug, Clone, PartialEq)]
21#[allow(dead_code)]
22pub enum Role {
23    /// System prompt.
24    System,
25    /// User message.
26    User,
27    /// Assistant (LLM) response.
28    Assistant,
29    /// Tool result.
30    Tool,
31}
32
33impl Role {
34    /// String representation for database storage.
35    pub fn as_str(&self) -> &'static str {
36        match self {
37            Self::System => "system",
38            Self::User => "user",
39            Self::Assistant => "assistant",
40            Self::Tool => "tool",
41        }
42    }
43}
44
45impl std::fmt::Display for Role {
46    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
47        write!(f, "{}", self.as_str())
48    }
49}
50
51impl std::str::FromStr for Role {
52    type Err = String;
53    fn from_str(s: &str) -> Result<Self, Self::Err> {
54        match s {
55            "system" => Ok(Self::System),
56            "user" => Ok(Self::User),
57            "assistant" => Ok(Self::Assistant),
58            "tool" => Ok(Self::Tool),
59            other => Err(format!("unknown role: {other}")),
60        }
61    }
62}
63
64/// A stored message row.
65#[derive(Debug, Clone)]
66#[allow(dead_code)]
67pub struct Message {
68    /// Database row ID.
69    pub id: i64,
70    /// Session this message belongs to.
71    pub session_id: String,
72    /// Message role (system, user, assistant, tool).
73    pub role: Role,
74    /// Text content (may be a summary for Bash results).
75    pub content: Option<String>,
76    /// Full untruncated output (only set for Bash tool results).
77    pub full_content: Option<String>,
78    /// Serialized tool calls JSON.
79    pub tool_calls: Option<String>,
80    /// ID of the tool call this responds to.
81    pub tool_call_id: Option<String>,
82    /// Input tokens for this message.
83    pub prompt_tokens: Option<i64>,
84    /// Output tokens for this message.
85    pub completion_tokens: Option<i64>,
86    /// Cached input tokens.
87    pub cache_read_tokens: Option<i64>,
88    /// Tokens written to cache.
89    pub cache_creation_tokens: Option<i64>,
90    /// Reasoning/thinking tokens.
91    pub thinking_tokens: Option<i64>,
92    /// ISO 8601 creation timestamp.
93    pub created_at: Option<String>,
94}
95
96/// Detected interruption state for a resumed session.
97///
98/// Returned by [`detect_interruption`](crate::db::queries::detect_interruption)
99/// after inspecting the tail of the message history.
100///
101/// ## Design decision: banner, not auto-resume
102///
103/// Claude Code auto-continues interrupted turns (re-sends the prompt or
104/// injects "Continue from where you left off"). Koda deliberately shows a
105/// banner and lets the user decide, for three reasons:
106///
107/// 1. **Safety** — auto-resuming a destructive tool call (e.g. `rm -rf`)
108///    after a VPN drop is surprising. The user should see the state first.
109/// 2. **Stale context** — the user may have fixed the issue manually while
110///    Koda was disconnected. Auto-resume wastes tokens re-doing work.
111/// 3. **Cost** — resuming to *check history* shouldn't burn an API call.
112///
113/// A single "type `continue` or rephrase" banner is near-zero friction
114/// (one word + Enter) and handles all three cases.
115#[derive(Debug, Clone, PartialEq)]
116pub enum InterruptionKind {
117    /// The user's prompt was never answered (last message is `Role::User`).
118    /// Contains a preview of the unanswered prompt.
119    Prompt(String),
120    /// A tool finished but the assistant never processed the result
121    /// (last message is `Role::Tool`).
122    Tool,
123}
124
125/// Token usage totals for a session.
126#[derive(Debug, Clone, Default)]
127pub struct SessionUsage {
128    /// Total input tokens.
129    pub prompt_tokens: i64,
130    /// Total output tokens.
131    pub completion_tokens: i64,
132    /// Total cached input tokens.
133    pub cache_read_tokens: i64,
134    /// Total tokens written to cache.
135    pub cache_creation_tokens: i64,
136    /// Total reasoning/thinking tokens.
137    pub thinking_tokens: i64,
138    /// Number of API calls made.
139    pub api_calls: i64,
140}
141
142/// Summary info for a stored session.
143#[derive(Debug, Clone)]
144pub struct SessionInfo {
145    /// Session identifier.
146    pub id: String,
147    /// Agent name for this session.
148    pub agent_name: String,
149    /// ISO 8601 creation timestamp.
150    pub created_at: String,
151    /// Total messages in the session.
152    pub message_count: i64,
153    /// Cumulative token count.
154    pub total_tokens: i64,
155    /// Auto-generated title from first user message.
156    pub title: Option<String>,
157    /// Last active approval mode (for restore on resume).
158    pub mode: Option<String>,
159}
160
161/// Stats about compacted (archived) messages in the database.
162#[derive(Debug, Clone, Default)]
163pub struct CompactedStats {
164    /// Number of compacted messages.
165    pub message_count: i64,
166    /// Number of sessions with compacted messages.
167    pub session_count: i64,
168    /// Approximate size in bytes of compacted message content.
169    pub size_bytes: i64,
170    /// ISO 8601 timestamp of the oldest compacted message.
171    pub oldest: Option<String>,
172}
173
174/// Core storage contract for sessions, messages, and metadata.
175#[async_trait::async_trait]
176pub trait Persistence: Send + Sync {
177    // ── Sessions ──
178
179    /// Create a new session, returning its unique ID.
180    async fn create_session(&self, agent_name: &str, project_root: &Path) -> Result<String>;
181    /// List recent sessions for the given project root.
182    async fn list_sessions(&self, limit: i64, project_root: &Path) -> Result<Vec<SessionInfo>>;
183    /// Delete a session by ID. Returns `true` if it existed.
184    async fn delete_session(&self, session_id: &str) -> Result<bool>;
185    /// Set the auto-generated title for a session.
186    async fn set_session_title(&self, session_id: &str, title: &str) -> Result<()>;
187    /// Persist the current approval mode for a session (restored on resume).
188    async fn set_session_mode(&self, session_id: &str, mode: &str) -> Result<()>;
189    /// Get the stored approval mode for a session.
190    async fn get_session_mode(&self, session_id: &str) -> Result<Option<String>>;
191    /// Seconds elapsed since the session was last accessed (`last_accessed_at`).
192    /// Returns `None` if the column is NULL (session never had a context load).
193    async fn get_session_idle_secs(&self, session_id: &str) -> Result<Option<i64>>;
194
195    // ── Messages ──
196
197    /// Insert a message into a session.
198    async fn insert_message(
199        &self,
200        session_id: &str,
201        role: &Role,
202        content: Option<&str>,
203        tool_calls: Option<&str>,
204        tool_call_id: Option<&str>,
205        usage: Option<&crate::providers::TokenUsage>,
206    ) -> Result<i64>;
207
208    /// Insert a message with an explicit agent name (for sub-agent tracking).
209    #[allow(clippy::too_many_arguments)]
210    async fn insert_message_with_agent(
211        &self,
212        session_id: &str,
213        role: &Role,
214        content: Option<&str>,
215        tool_calls: Option<&str>,
216        tool_call_id: Option<&str>,
217        usage: Option<&crate::providers::TokenUsage>,
218        agent_name: Option<&str>,
219    ) -> Result<i64>;
220
221    /// Insert a tool message with full (untruncated) output stored separately.
222    ///
223    /// `content` holds the model-facing summary; `full_content` holds the
224    /// complete output for later retrieval via RecallContext.
225    #[allow(clippy::too_many_arguments)]
226    async fn insert_tool_message_with_full(
227        &self,
228        session_id: &str,
229        content: &str,
230        tool_call_id: &str,
231        full_content: &str,
232    ) -> Result<i64>;
233
234    /// Load active (non-compacted) conversation context for a session.
235    async fn load_context(&self, session_id: &str) -> Result<Vec<Message>>;
236    /// Load all messages in a session (no token limit).
237    async fn load_all_messages(&self, session_id: &str) -> Result<Vec<Message>>;
238    /// Recent user messages across all sessions (for startup hints).
239    async fn recent_user_messages(&self, limit: i64) -> Result<Vec<String>>;
240    /// Last assistant message in a session.
241    async fn last_assistant_message(&self, session_id: &str) -> Result<String>;
242    /// Last user message in a session.
243    async fn last_user_message(&self, session_id: &str) -> Result<String>;
244    /// Check if the session has unresolved tool calls.
245    async fn has_pending_tool_calls(&self, session_id: &str) -> Result<bool>;
246
247    /// Mark an assistant message as fully delivered.
248    ///
249    /// Sets `completed_at = CURRENT_TIMESTAMP`. Only called after a legitimate
250    /// `StreamChunk::Done` — not after user cancellation or a network error.
251    /// A `NULL` `completed_at` means the message is in-progress or was interrupted.
252    async fn mark_message_complete(&self, message_id: i64) -> Result<()>;
253
254    // ── Token usage ──
255
256    /// Token usage totals for a session.
257    async fn session_token_usage(&self, session_id: &str) -> Result<SessionUsage>;
258    /// Token usage broken down by agent name.
259    async fn session_usage_by_agent(&self, session_id: &str)
260    -> Result<Vec<(String, SessionUsage)>>;
261
262    // ── Compaction ──
263
264    /// Compact old messages into a summary, preserving the last N messages.
265    async fn compact_session(
266        &self,
267        session_id: &str,
268        summary: &str,
269        preserve_count: usize,
270    ) -> Result<usize>;
271
272    // ── Microcompact ──
273
274    /// Replace message content for the given IDs with a stub string.
275    /// Used by microcompact to clear old tool results without full compaction.
276    async fn clear_message_content(&self, message_ids: &[i64], stub: &str) -> Result<()>;
277
278    // ── Purge ──
279
280    /// Stats about compacted (archived) messages across all sessions.
281    async fn compacted_stats(&self) -> Result<CompactedStats>;
282    /// Permanently delete compacted messages older than `min_age_days`.
283    /// Returns the number of messages deleted.
284    async fn purge_compacted(&self, min_age_days: u32) -> Result<usize>;
285
286    // ── Metadata ──
287
288    /// Get a session metadata value by key.
289    async fn get_metadata(&self, session_id: &str, key: &str) -> Result<Option<String>>;
290    /// Set a session metadata value.
291    async fn set_metadata(&self, session_id: &str, key: &str, value: &str) -> Result<()>;
292    /// Get the TODO list for a session.
293    async fn get_todo(&self, session_id: &str) -> Result<Option<String>>;
294    /// Set the TODO list for a session.
295    async fn set_todo(&self, session_id: &str, content: &str) -> Result<()>;
296}