1use std::collections::HashMap;
13use std::fs::{self, OpenOptions};
14use std::io::{BufRead, BufReader, Write};
15use std::path::{Path, PathBuf};
16use std::sync::{Arc, Mutex, MutexGuard};
17
18use anyhow::{Context, Result};
19use chrono::Utc;
20use rusqlite::{Connection, OptionalExtension, params};
21use serde::{Deserialize, Serialize};
22use serde_json::Value;
23
24#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
28#[serde(rename_all = "snake_case")]
29pub enum ThreadStatus {
30 Running,
32 Idle,
34 Completed,
36 Failed,
38 Paused,
40 Archived,
42}
43
44#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
48#[serde(rename_all = "snake_case")]
49pub enum SessionSource {
50 Interactive,
52 Resume,
54 Fork,
56 Api,
58 Unknown,
60}
61
62#[derive(Debug, Clone, Serialize, Deserialize)]
67pub struct ThreadMetadata {
68 pub id: String,
70 pub rollout_path: Option<PathBuf>,
72 pub preview: String,
74 pub ephemeral: bool,
76 pub model_provider: String,
78 pub created_at: i64,
80 pub updated_at: i64,
82 pub status: ThreadStatus,
84 pub path: Option<PathBuf>,
86 pub cwd: PathBuf,
88 pub cli_version: String,
90 pub source: SessionSource,
92 pub name: Option<String>,
94 pub sandbox_policy: Option<String>,
96 pub approval_mode: Option<String>,
98 pub archived: bool,
100 pub archived_at: Option<i64>,
102 pub git_sha: Option<String>,
104 pub git_branch: Option<String>,
106 pub git_origin_url: Option<String>,
108 pub memory_mode: Option<String>,
110 pub current_leaf_id: Option<i64>,
112}
113
114#[derive(Debug, Clone, Serialize, Deserialize)]
116pub struct DynamicToolRecord {
117 pub position: i64,
119 pub name: String,
121 pub description: Option<String>,
123 pub input_schema: Value,
125}
126
127#[derive(Debug, Clone, Serialize, Deserialize)]
132pub struct MessageRecord {
133 pub id: i64,
135 pub thread_id: String,
137 pub role: String,
139 pub content: String,
141 pub item: Option<Value>,
143 pub created_at: i64,
145 pub parent_entry_id: Option<i64>,
147}
148
149#[derive(Debug, Clone, Serialize, Deserialize)]
151pub struct CheckpointRecord {
152 pub thread_id: String,
154 pub checkpoint_id: String,
156 pub state: Value,
158 pub created_at: i64,
160}
161
162#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
166#[serde(rename_all = "snake_case")]
167pub enum JobStateStatus {
168 Queued,
170 Running,
172 Completed,
174 Failed,
176 Cancelled,
178}
179
180#[derive(Debug, Clone, Serialize, Deserialize)]
182pub struct JobStateRecord {
183 pub id: String,
185 pub name: String,
187 pub status: JobStateStatus,
189 pub progress: Option<u8>,
191 pub detail: Option<String>,
193 pub created_at: i64,
195 pub updated_at: i64,
197}
198
199#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
201#[serde(rename_all = "snake_case")]
202pub enum ThreadGoalStatus {
203 Active,
205 Paused,
207 Blocked,
209 UsageLimited,
211 BudgetLimited,
213 Complete,
215}
216
217#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
219pub struct ThreadGoalRecord {
220 pub thread_id: String,
222 pub goal_id: String,
224 pub objective: String,
226 pub status: ThreadGoalStatus,
228 pub token_budget: Option<i64>,
230 pub tokens_used: i64,
232 pub time_used_seconds: i64,
234 pub continuation_count: i64,
236 pub created_at: i64,
238 pub updated_at: i64,
240}
241
242#[derive(Debug, Clone)]
244pub struct ThreadListFilters {
245 pub include_archived: bool,
247 pub limit: Option<usize>,
249}
250
251impl Default for ThreadListFilters {
252 fn default() -> Self {
253 Self {
254 include_archived: false,
255 limit: Some(50),
256 }
257 }
258}
259
260#[derive(Debug, Clone, Serialize, Deserialize)]
261struct SessionIndexEntry {
262 thread_id: String,
263 thread_name: Option<String>,
264 updated_at: i64,
265 rollout_path: Option<PathBuf>,
266}
267
268#[derive(Debug, Clone)]
273pub struct StateStore {
274 db_path: PathBuf,
275 session_index_path: PathBuf,
276 conn: Arc<Mutex<Connection>>,
280}
281
282impl StateStore {
283 pub fn open(path: Option<PathBuf>) -> Result<Self> {
289 let db_path = path.unwrap_or_else(default_state_db_path);
290 let session_index_path = db_path
291 .parent()
292 .unwrap_or_else(|| Path::new("."))
293 .join("session_index.jsonl");
294 if let Some(parent) = db_path.parent() {
295 fs::create_dir_all(parent).with_context(|| {
296 format!("failed to create state directory {}", parent.display())
297 })?;
298 }
299 let conn = Connection::open(&db_path)
300 .with_context(|| format!("failed to open state db {}", db_path.display()))?;
301 conn.pragma_update(None, "foreign_keys", "ON")
302 .with_context(|| format!("failed to enable foreign keys for {}", db_path.display()))?;
303 Self::init_schema(&conn)?;
304 Ok(Self {
305 db_path,
306 session_index_path,
307 conn: Arc::new(Mutex::new(conn)),
308 })
309 }
310
311 pub fn db_path(&self) -> &Path {
313 &self.db_path
314 }
315
316 fn conn(&self) -> Result<MutexGuard<'_, Connection>> {
317 self.conn
321 .lock()
322 .map_err(|_| anyhow::anyhow!("state db connection mutex poisoned"))
323 }
324
325 fn init_schema(conn: &Connection) -> Result<()> {
326 let mut user_version: u32 = conn.query_row("PRAGMA user_version;", [], |row| row.get(0))?;
327 if user_version == 0 {
328 conn.execute_batch(
329 r#"
330 BEGIN;
331 CREATE TABLE IF NOT EXISTS threads (
332 id TEXT PRIMARY KEY,
333 rollout_path TEXT,
334 preview TEXT NOT NULL,
335 ephemeral INTEGER NOT NULL,
336 model_provider TEXT NOT NULL,
337 created_at INTEGER NOT NULL,
338 updated_at INTEGER NOT NULL,
339 status TEXT NOT NULL,
340 path TEXT,
341 cwd TEXT NOT NULL,
342 cli_version TEXT NOT NULL,
343 source TEXT NOT NULL,
344 title TEXT,
345 sandbox_policy TEXT,
346 approval_mode TEXT,
347 archived INTEGER NOT NULL DEFAULT 0,
348 archived_at INTEGER,
349 git_sha TEXT,
350 git_branch TEXT,
351 git_origin_url TEXT,
352 memory_mode TEXT
353 );
354 CREATE INDEX IF NOT EXISTS idx_threads_updated_at ON threads(updated_at DESC);
355 CREATE INDEX IF NOT EXISTS idx_threads_archived_at ON threads(archived_at DESC);
356 CREATE INDEX IF NOT EXISTS idx_threads_archived_updated ON threads(archived, updated_at DESC);
357
358 CREATE TABLE IF NOT EXISTS thread_dynamic_tools (
359 thread_id TEXT NOT NULL,
360 position INTEGER NOT NULL,
361 name TEXT NOT NULL,
362 description TEXT,
363 input_schema TEXT NOT NULL,
364 PRIMARY KEY (thread_id, position),
365 FOREIGN KEY(thread_id) REFERENCES threads(id) ON DELETE CASCADE
366 );
367
368 CREATE TABLE IF NOT EXISTS messages (
369 id INTEGER PRIMARY KEY AUTOINCREMENT,
370 thread_id TEXT NOT NULL,
371 role TEXT NOT NULL,
372 content TEXT NOT NULL,
373 item_json TEXT,
374 created_at INTEGER NOT NULL,
375 FOREIGN KEY(thread_id) REFERENCES threads(id) ON DELETE CASCADE
376 );
377 CREATE INDEX IF NOT EXISTS idx_messages_thread_created_at ON messages(thread_id, created_at ASC);
378
379 CREATE TABLE IF NOT EXISTS checkpoints (
380 thread_id TEXT NOT NULL,
381 checkpoint_id TEXT NOT NULL,
382 state_json TEXT NOT NULL,
383 created_at INTEGER NOT NULL,
384 PRIMARY KEY(thread_id, checkpoint_id),
385 FOREIGN KEY(thread_id) REFERENCES threads(id) ON DELETE CASCADE
386 );
387 CREATE INDEX IF NOT EXISTS idx_checkpoints_thread_created_at ON checkpoints(thread_id, created_at DESC);
388
389 CREATE TABLE IF NOT EXISTS jobs (
390 id TEXT PRIMARY KEY,
391 name TEXT NOT NULL,
392 status TEXT NOT NULL,
393 progress INTEGER,
394 detail TEXT,
395 created_at INTEGER NOT NULL,
396 updated_at INTEGER NOT NULL
397 );
398 CREATE INDEX IF NOT EXISTS idx_jobs_updated_at ON jobs(updated_at DESC);
399
400 -- Add parent_entry_id column, and set to last message before current message
401 ALTER TABLE messages ADD COLUMN parent_entry_id INTEGER NULL;
402 UPDATE messages
403 SET parent_entry_id = (
404 SELECT m2.id
405 FROM messages m2
406 WHERE m2.thread_id = messages.thread_id
407 AND (
408 m2.created_at < messages.created_at
409 OR (
410 m2.created_at = messages.created_at
411 AND m2.id < messages.id
412 )
413 )
414 ORDER BY m2.created_at DESC, m2.id DESC
415 LIMIT 1
416 );
417 CREATE INDEX idx_messages_parent_entry_id ON messages(parent_entry_id);
418
419 -- Add current_leaf_id column, and set to last message in thread
420 ALTER TABLE threads ADD COLUMN current_leaf_id INTEGER NULL;
421 UPDATE threads
422 SET current_leaf_id = (
423 SELECT m.id
424 FROM messages m
425 WHERE m.thread_id = threads.id
426 ORDER BY m.id DESC
427 LIMIT 1
428 );
429
430 PRAGMA user_version = 1;
431 COMMIT;
432 "#,
433 )
434 .context("failed to initialize thread schema")?;
435 user_version = 1;
436 }
437 if user_version < 2 {
438 conn.execute_batch(
439 r#"
440 BEGIN;
441 CREATE TABLE IF NOT EXISTS workflow_runs (
442 id TEXT PRIMARY KEY,
443 workflow_id TEXT NOT NULL,
444 goal TEXT NOT NULL,
445 status TEXT NOT NULL,
446 input_hash TEXT,
447 started_at INTEGER NOT NULL,
448 completed_at INTEGER,
449 metadata_json TEXT NOT NULL DEFAULT '{}'
450 );
451 CREATE INDEX IF NOT EXISTS idx_workflow_runs_status_started_at
452 ON workflow_runs(status, started_at DESC);
453 CREATE INDEX IF NOT EXISTS idx_workflow_runs_workflow_started_at
454 ON workflow_runs(workflow_id, started_at DESC);
455
456 CREATE TABLE IF NOT EXISTS branch_runs (
457 id TEXT PRIMARY KEY,
458 workflow_run_id TEXT NOT NULL,
459 branch_id TEXT NOT NULL,
460 node_id TEXT NOT NULL,
461 status TEXT NOT NULL,
462 started_at INTEGER NOT NULL,
463 completed_at INTEGER,
464 result_json TEXT NOT NULL DEFAULT '{}',
465 FOREIGN KEY(workflow_run_id) REFERENCES workflow_runs(id) ON DELETE CASCADE
466 );
467 CREATE INDEX IF NOT EXISTS idx_branch_runs_workflow_run_id
468 ON branch_runs(workflow_run_id);
469 CREATE INDEX IF NOT EXISTS idx_branch_runs_branch_id
470 ON branch_runs(branch_id);
471
472 CREATE TABLE IF NOT EXISTS leaf_runs (
473 id TEXT PRIMARY KEY,
474 workflow_run_id TEXT NOT NULL,
475 branch_run_id TEXT,
476 leaf_id TEXT NOT NULL,
477 task_id TEXT NOT NULL,
478 input_hash TEXT,
479 status TEXT NOT NULL,
480 output_json TEXT NOT NULL DEFAULT '{}',
481 artifacts_json TEXT NOT NULL DEFAULT '[]',
482 started_at INTEGER NOT NULL,
483 completed_at INTEGER,
484 FOREIGN KEY(workflow_run_id) REFERENCES workflow_runs(id) ON DELETE CASCADE,
485 FOREIGN KEY(branch_run_id) REFERENCES branch_runs(id) ON DELETE SET NULL
486 );
487 CREATE INDEX IF NOT EXISTS idx_leaf_runs_workflow_run_id
488 ON leaf_runs(workflow_run_id);
489 CREATE INDEX IF NOT EXISTS idx_leaf_runs_replay_lookup
490 ON leaf_runs(workflow_run_id, leaf_id, input_hash);
491
492 CREATE TABLE IF NOT EXISTS control_node_runs (
493 id TEXT PRIMARY KEY,
494 workflow_run_id TEXT NOT NULL,
495 node_id TEXT NOT NULL,
496 kind TEXT NOT NULL,
497 status TEXT NOT NULL,
498 selected_children_json TEXT NOT NULL DEFAULT '[]',
499 result_json TEXT NOT NULL DEFAULT '{}',
500 started_at INTEGER NOT NULL,
501 completed_at INTEGER,
502 FOREIGN KEY(workflow_run_id) REFERENCES workflow_runs(id) ON DELETE CASCADE
503 );
504 CREATE INDEX IF NOT EXISTS idx_control_node_runs_workflow_run_id
505 ON control_node_runs(workflow_run_id);
506 CREATE INDEX IF NOT EXISTS idx_control_node_runs_node_id
507 ON control_node_runs(node_id);
508
509 CREATE TABLE IF NOT EXISTS teacher_candidates (
510 id TEXT PRIMARY KEY,
511 workflow_run_id TEXT NOT NULL,
512 control_node_run_id TEXT NOT NULL,
513 candidate_id TEXT NOT NULL,
514 branch_run_id TEXT,
515 score REAL,
516 passed INTEGER,
517 rationale_json TEXT NOT NULL DEFAULT '{}',
518 created_at INTEGER NOT NULL,
519 FOREIGN KEY(workflow_run_id) REFERENCES workflow_runs(id) ON DELETE CASCADE,
520 FOREIGN KEY(control_node_run_id) REFERENCES control_node_runs(id) ON DELETE CASCADE,
521 FOREIGN KEY(branch_run_id) REFERENCES branch_runs(id) ON DELETE SET NULL
522 );
523 CREATE INDEX IF NOT EXISTS idx_teacher_candidates_workflow_run_id
524 ON teacher_candidates(workflow_run_id);
525 CREATE INDEX IF NOT EXISTS idx_teacher_candidates_control_node_run_id
526 ON teacher_candidates(control_node_run_id);
527
528 PRAGMA user_version = 2;
529 COMMIT;
530 "#,
531 )
532 .context("failed to initialize workflow trace schema")?;
533 user_version = 2;
534 }
535 if user_version < 3 {
536 conn.execute_batch(
537 r#"
538 BEGIN;
539 CREATE TABLE IF NOT EXISTS thread_goals (
540 thread_id TEXT PRIMARY KEY NOT NULL,
541 goal_id TEXT NOT NULL,
542 objective TEXT NOT NULL,
543 status TEXT NOT NULL CHECK(status IN (
544 'active',
545 'paused',
546 'blocked',
547 'usage_limited',
548 'budget_limited',
549 'complete'
550 )),
551 token_budget INTEGER,
552 tokens_used INTEGER NOT NULL DEFAULT 0,
553 time_used_seconds INTEGER NOT NULL DEFAULT 0,
554 created_at INTEGER NOT NULL,
555 updated_at INTEGER NOT NULL,
556 FOREIGN KEY(thread_id) REFERENCES threads(id) ON DELETE CASCADE
557 );
558
559 PRAGMA user_version = 3;
560 COMMIT;
561 "#,
562 )
563 .context("failed to initialize thread goal schema")?;
564 user_version = 3;
565 }
566 if user_version < 4 {
567 conn.execute_batch(
568 r#"
569 BEGIN;
570 ALTER TABLE thread_goals
571 ADD COLUMN continuation_count INTEGER NOT NULL DEFAULT 0;
572
573 PRAGMA user_version = 4;
574 COMMIT;
575 "#,
576 )
577 .context("failed to initialize thread goal continuation schema")?;
578 }
579 Ok(())
580 }
581
582 pub fn upsert_thread(&self, thread: &ThreadMetadata) -> Result<()> {
587 let conn = self.conn()?;
588 conn.execute(
589 r#"
590 INSERT INTO threads (
591 id, rollout_path, preview, ephemeral, model_provider, created_at, updated_at, status, path, cwd,
592 cli_version, source, title, sandbox_policy, approval_mode, archived, archived_at,
593 git_sha, git_branch, git_origin_url, memory_mode
594 ) VALUES (
595 ?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10,
596 ?11, ?12, ?13, ?14, ?15, ?16, ?17,
597 ?18, ?19, ?20, ?21
598 )
599 ON CONFLICT(id) DO UPDATE SET
600 rollout_path=excluded.rollout_path,
601 preview=excluded.preview,
602 ephemeral=excluded.ephemeral,
603 model_provider=excluded.model_provider,
604 created_at=excluded.created_at,
605 updated_at=excluded.updated_at,
606 status=excluded.status,
607 path=excluded.path,
608 cwd=excluded.cwd,
609 cli_version=excluded.cli_version,
610 source=excluded.source,
611 title=excluded.title,
612 sandbox_policy=excluded.sandbox_policy,
613 approval_mode=excluded.approval_mode,
614 archived=excluded.archived,
615 archived_at=excluded.archived_at,
616 git_sha=excluded.git_sha,
617 git_branch=excluded.git_branch,
618 git_origin_url=excluded.git_origin_url,
619 memory_mode=excluded.memory_mode
620 "#,
621 params![
622 thread.id,
623 path_to_opt_string(thread.rollout_path.as_deref()),
624 thread.preview,
625 bool_to_i64(thread.ephemeral),
626 thread.model_provider,
627 thread.created_at,
628 thread.updated_at,
629 thread_status_to_str(&thread.status),
630 path_to_opt_string(thread.path.as_deref()),
631 thread.cwd.display().to_string(),
632 thread.cli_version,
633 session_source_to_str(&thread.source),
634 thread.name,
635 thread.sandbox_policy,
636 thread.approval_mode,
637 bool_to_i64(thread.archived),
638 thread.archived_at,
639 thread.git_sha,
640 thread.git_branch,
641 thread.git_origin_url,
642 thread.memory_mode,
643 ],
644 )
645 .context("failed to upsert thread metadata")?;
646
647 self.append_thread_name(
648 &thread.id,
649 thread.name.clone(),
650 thread.updated_at,
651 thread.rollout_path.clone(),
652 )?;
653 Ok(())
654 }
655
656 pub fn get_thread(&self, id: &str) -> Result<Option<ThreadMetadata>> {
660 let conn = self.conn()?;
661 conn.query_row(
662 r#"
663 SELECT id, rollout_path, preview, ephemeral, model_provider, created_at, updated_at, status, path, cwd,
664 cli_version, source, title, sandbox_policy, approval_mode, archived, archived_at,
665 git_sha, git_branch, git_origin_url, memory_mode, current_leaf_id
666 FROM threads
667 WHERE id = ?1
668 "#,
669 params![id],
670 row_to_thread,
671 )
672 .optional()
673 .context("failed to read thread")
674 }
675
676 pub fn list_threads(&self, filters: ThreadListFilters) -> Result<Vec<ThreadMetadata>> {
681 let conn = self.conn()?;
682 let sql = if filters.include_archived {
683 "SELECT id, rollout_path, preview, ephemeral, model_provider, created_at, updated_at, status, path, cwd, cli_version, source, title, sandbox_policy, approval_mode, archived, archived_at, git_sha, git_branch, git_origin_url, memory_mode, current_leaf_id FROM threads ORDER BY updated_at DESC LIMIT ?1"
684 } else {
685 "SELECT id, rollout_path, preview, ephemeral, model_provider, created_at, updated_at, status, path, cwd, cli_version, source, title, sandbox_policy, approval_mode, archived, archived_at, git_sha, git_branch, git_origin_url, memory_mode, current_leaf_id FROM threads WHERE archived = 0 ORDER BY updated_at DESC LIMIT ?1"
686 };
687
688 let mut stmt = conn.prepare(sql).context("failed to prepare list query")?;
689 let limit = i64::try_from(filters.limit.unwrap_or(50)).unwrap_or(50);
690 let mut rows = stmt
691 .query(params![limit])
692 .context("failed to query threads")?;
693 let mut out = Vec::new();
694 while let Some(row) = rows.next().context("failed to iterate thread rows")? {
695 out.push(row_to_thread(row)?);
696 }
697 Ok(out)
698 }
699
700 pub fn mark_archived(&self, id: &str) -> Result<()> {
703 let conn = self.conn()?;
704 conn.execute(
705 "UPDATE threads SET archived = 1, archived_at = ?2, status = ?3 WHERE id = ?1",
706 params![
707 id,
708 Utc::now().timestamp(),
709 thread_status_to_str(&ThreadStatus::Archived)
710 ],
711 )
712 .context("failed to archive thread")?;
713 Ok(())
714 }
715
716 pub fn mark_unarchived(&self, id: &str) -> Result<()> {
718 let conn = self.conn()?;
719 conn.execute(
720 "UPDATE threads SET archived = 0, archived_at = NULL WHERE id = ?1",
721 params![id],
722 )
723 .context("failed to unarchive thread")?;
724 Ok(())
725 }
726
727 pub fn delete_thread(&self, id: &str) -> Result<()> {
730 let conn = self.conn()?;
731 conn.execute("DELETE FROM threads WHERE id = ?1", params![id])
732 .context("failed to delete thread")?;
733 Ok(())
734 }
735
736 pub fn set_thread_memory_mode(&self, id: &str, mode: Option<&str>) -> Result<()> {
740 let conn = self.conn()?;
741 conn.execute(
742 "UPDATE threads SET memory_mode = ?2 WHERE id = ?1",
743 params![id, mode],
744 )
745 .context("failed to update thread memory mode")?;
746 Ok(())
747 }
748
749 pub fn get_thread_memory_mode(&self, id: &str) -> Result<Option<String>> {
753 let conn = self.conn()?;
754 conn.query_row(
755 "SELECT memory_mode FROM threads WHERE id = ?1",
756 params![id],
757 |row| row.get::<_, Option<String>>(0),
758 )
759 .optional()
760 .context("failed to read thread memory mode")
761 .map(Option::flatten)
762 }
763
764 pub fn upsert_thread_goal(&self, goal: &ThreadGoalRecord) -> Result<()> {
766 let conn = self.conn()?;
767 let exists: Option<i64> = conn
768 .query_row(
769 "SELECT 1 FROM threads WHERE id = ?1",
770 params![goal.thread_id],
771 |row| row.get(0),
772 )
773 .optional()
774 .context("failed to verify thread before saving goal")?;
775 if exists.is_none() {
776 anyhow::bail!("thread {} not found", goal.thread_id);
777 }
778
779 conn.execute(
780 r#"
781 INSERT INTO thread_goals (
782 thread_id, goal_id, objective, status, token_budget, tokens_used,
783 time_used_seconds, continuation_count, created_at, updated_at
784 ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10)
785 ON CONFLICT(thread_id) DO UPDATE SET
786 goal_id=excluded.goal_id,
787 objective=excluded.objective,
788 status=excluded.status,
789 token_budget=excluded.token_budget,
790 tokens_used=excluded.tokens_used,
791 time_used_seconds=excluded.time_used_seconds,
792 continuation_count=excluded.continuation_count,
793 created_at=excluded.created_at,
794 updated_at=excluded.updated_at
795 "#,
796 params![
797 goal.thread_id,
798 goal.goal_id,
799 goal.objective,
800 thread_goal_status_to_str(&goal.status),
801 goal.token_budget,
802 goal.tokens_used,
803 goal.time_used_seconds,
804 goal.continuation_count,
805 goal.created_at,
806 goal.updated_at,
807 ],
808 )
809 .context("failed to upsert thread goal")?;
810 Ok(())
811 }
812
813 pub fn record_thread_goal_usage(
829 &self,
830 thread_id: &str,
831 token_delta: i64,
832 time_delta_seconds: i64,
833 now: i64,
834 ) -> Result<Option<ThreadGoalRecord>> {
835 let changed = {
836 let conn = self.conn()?;
837 conn.execute(
838 r#"
839 UPDATE thread_goals
840 SET tokens_used = tokens_used + ?2,
841 time_used_seconds = time_used_seconds + ?3,
842 updated_at = MAX(updated_at, ?4)
843 WHERE thread_id = ?1
844 "#,
845 params![thread_id, token_delta, time_delta_seconds, now],
846 )
847 .context("failed to record thread goal usage")?
848 };
849 if changed == 0 {
850 return Ok(None);
851 }
852 self.get_thread_goal(thread_id)
853 }
854
855 pub fn record_thread_goal_continuation(
861 &self,
862 thread_id: &str,
863 now: i64,
864 ) -> Result<Option<ThreadGoalRecord>> {
865 let changed = {
866 let conn = self.conn()?;
867 conn.execute(
868 r#"
869 UPDATE thread_goals
870 SET continuation_count = continuation_count + 1,
871 updated_at = MAX(updated_at, ?2)
872 WHERE thread_id = ?1
873 "#,
874 params![thread_id, now],
875 )
876 .context("failed to record thread goal continuation")?
877 };
878 if changed == 0 {
879 return Ok(None);
880 }
881 self.get_thread_goal(thread_id)
882 }
883
884 pub fn get_thread_goal(&self, thread_id: &str) -> Result<Option<ThreadGoalRecord>> {
886 let conn = self.conn()?;
887 conn.query_row(
888 r#"
889 SELECT thread_id, goal_id, objective, status, token_budget, tokens_used,
890 time_used_seconds, continuation_count, created_at, updated_at
891 FROM thread_goals
892 WHERE thread_id = ?1
893 "#,
894 params![thread_id],
895 row_to_thread_goal,
896 )
897 .optional()
898 .context("failed to read thread goal")
899 }
900
901 pub fn delete_thread_goal(&self, thread_id: &str) -> Result<bool> {
903 let conn = self.conn()?;
904 let changed = conn
905 .execute(
906 "DELETE FROM thread_goals WHERE thread_id = ?1",
907 params![thread_id],
908 )
909 .context("failed to delete thread goal")?;
910 Ok(changed > 0)
911 }
912
913 pub fn list_leaf_messages(&self, thread_id: &str) -> Result<Vec<MessageRecord>> {
918 let conn = self.conn()?;
919 let mut stmt = conn
920 .prepare(
921 r#"
922 SELECT m1.id, m1.thread_id, m1.role, m1.content, m1.item_json, m1.created_at, m1.parent_entry_id
923 FROM messages m1
924 LEFT JOIN messages m2 ON m1.id = m2.parent_entry_id
925 WHERE m1.thread_id = ?1 AND m2.id IS NULL
926 "#,
927 )
928 .context("failed to prepare message listing query")?;
929 let mut rows = stmt
930 .query(params![thread_id])
931 .with_context(|| format!("failed to list leaf messages for thread {thread_id}"))?;
932 let mut out = Vec::new();
933 while let Some(row) = rows.next().context("failed to iterate message rows")? {
934 let item_json: Option<String> = row.get(4).context("failed to read item json")?;
935 let item = item_json
936 .as_deref()
937 .map(serde_json::from_str)
938 .transpose()
939 .with_context(|| {
940 format!("failed to parse message item json in thread {thread_id}")
941 })?;
942 out.push(MessageRecord {
943 id: row.get(0).context("failed to read message id")?,
944 thread_id: row.get(1).context("failed to read message thread id")?,
945 role: row.get(2).context("failed to read message role")?,
946 content: row.get(3).context("failed to read message content")?,
947 item,
948 created_at: row.get(5).context("failed to read message timestamp")?,
949 parent_entry_id: row.get(6).context("failed to read parent entry id")?,
950 });
951 }
952 Ok(out)
953 }
954
955 pub fn set_current_leaf_id(&self, thread_id: &str, current_leaf_id: &str) -> Result<()> {
960 let conn = self.conn()?;
961 conn.execute(
962 "UPDATE threads SET current_leaf_id = ?1 WHERE id = ?2",
963 params![current_leaf_id, thread_id],
964 )
965 .context("failed to update thread current leaf id")?;
966 Ok(())
967 }
968
969 pub fn persist_dynamic_tools(
974 &self,
975 thread_id: &str,
976 tools: &[DynamicToolRecord],
977 ) -> Result<()> {
978 let mut conn = self.conn()?;
979 let tx = conn
980 .transaction()
981 .context("failed to begin dynamic tools transaction")?;
982 tx.execute(
983 "DELETE FROM thread_dynamic_tools WHERE thread_id = ?1",
984 params![thread_id],
985 )
986 .context("failed to clear dynamic tools")?;
987 for tool in tools {
988 tx.execute(
989 "INSERT INTO thread_dynamic_tools(thread_id, position, name, description, input_schema) VALUES (?1, ?2, ?3, ?4, ?5)",
990 params![
991 thread_id,
992 tool.position,
993 tool.name,
994 tool.description,
995 tool.input_schema.to_string()
996 ],
997 )
998 .with_context(|| format!("failed to persist dynamic tool {}", tool.name))?;
999 }
1000 tx.commit().context("failed to commit dynamic tools")?;
1001 Ok(())
1002 }
1003
1004 pub fn get_dynamic_tools(&self, thread_id: &str) -> Result<Vec<DynamicToolRecord>> {
1006 let conn = self.conn()?;
1007 let mut stmt = conn
1008 .prepare(
1009 "SELECT position, name, description, input_schema FROM thread_dynamic_tools WHERE thread_id = ?1 ORDER BY position ASC",
1010 )
1011 .context("failed to prepare get dynamic tools query")?;
1012 let mut rows = stmt
1013 .query(params![thread_id])
1014 .context("failed to query dynamic tools")?;
1015 let mut out = Vec::new();
1016 while let Some(row) = rows.next().context("failed to iterate dynamic tools")? {
1017 let input_schema_raw: String =
1018 row.get(3).context("failed to read tool input schema")?;
1019 let input_schema: Value =
1020 serde_json::from_str(&input_schema_raw).with_context(|| {
1021 format!("failed to parse input schema for dynamic tool in thread {thread_id}")
1022 })?;
1023 out.push(DynamicToolRecord {
1024 position: row.get(0).context("failed to read tool position")?,
1025 name: row.get(1).context("failed to read tool name")?,
1026 description: row.get(2).context("failed to read tool description")?,
1027 input_schema,
1028 });
1029 }
1030 Ok(out)
1031 }
1032
1033 pub fn append_message(
1039 &self,
1040 thread_id: &str,
1041 role: &str,
1042 content: &str,
1043 item: Option<Value>,
1044 ) -> Result<i64> {
1045 let mut conn = self.conn()?;
1046 let created_at = Utc::now().timestamp();
1047 let item_json = item
1048 .as_ref()
1049 .map(serde_json::to_string)
1050 .transpose()
1051 .context("failed to serialize message item payload")?;
1052
1053 let tx = conn
1054 .transaction()
1055 .context("failed to begin append message transaction")?;
1056
1057 let current_leaf_id: Option<i64> = tx
1058 .query_row(
1059 "SELECT current_leaf_id FROM threads WHERE id = ?1",
1060 params![thread_id],
1061 |row| row.get(0),
1062 )
1063 .with_context(|| {
1064 format!("failed to query thread current leaf id for thread {thread_id}")
1065 })?;
1066
1067 let next_leaf_id: i64 = tx.query_row(
1068 r#"
1069 INSERT INTO messages(thread_id, role, content, item_json, created_at, parent_entry_id)
1070 SELECT ?1, ?2, ?3, ?4, ?5, ?6
1071 RETURNING id
1072 "#, params![thread_id, role, content, item_json, created_at, current_leaf_id], |row| row.get(0)
1073 ).with_context(|| format!("failed to append message for thread {thread_id}"))?;
1074
1075 tx.execute(
1076 r#"
1077 UPDATE threads
1078 SET current_leaf_id = ?1
1079 WHERE id = ?2;
1080 "#,
1081 params![next_leaf_id, thread_id],
1082 )
1083 .with_context(|| {
1084 format!("failed to update thread current leaf id for thread {thread_id}")
1085 })?;
1086
1087 tx.commit()
1088 .context("failed to commit append message transaction")?;
1089
1090 Ok(next_leaf_id)
1091 }
1092
1093 pub fn list_messages(
1099 &self,
1100 thread_id: &str,
1101 limit: Option<usize>,
1102 ) -> Result<Vec<MessageRecord>> {
1103 let conn = self.conn()?;
1104 let limit = i64::try_from(limit.unwrap_or(500)).unwrap_or(500);
1105 let mut stmt = conn
1106 .prepare(
1107 r#"
1108 WITH RECURSIVE
1109 leaf_id AS (
1110 SELECT current_leaf_id FROM threads WHERE id = ?1
1111 ),
1112 ancestors AS (
1113 SELECT id, thread_id, role, content, item_json, created_at, parent_entry_id, 0 AS depth
1114 FROM messages
1115 WHERE id = (SELECT current_leaf_id FROM leaf_id)
1116
1117 UNION ALL
1118
1119 SELECT m.id, m.thread_id, m.role, m.content, m.item_json, m.created_at, m.parent_entry_id, a.depth + 1
1120 FROM messages m
1121 JOIN ancestors a ON m.id = a.parent_entry_id
1122 WHERE a.depth < ?2
1123 )
1124 SELECT id, thread_id, role, content, item_json, created_at, parent_entry_id FROM ancestors
1125 ORDER BY depth DESC
1126 "#
1127 )
1128 .context("failed to prepare message listing query")?;
1129 let mut rows = stmt
1130 .query(params![thread_id, limit - 1])
1131 .with_context(|| format!("failed to list messages for thread {thread_id}"))?;
1132 let mut out = Vec::new();
1133 while let Some(row) = rows.next().context("failed to iterate message rows")? {
1134 let item_json: Option<String> = row.get(4).context("failed to read item json")?;
1135 let item = item_json
1136 .as_deref()
1137 .map(serde_json::from_str)
1138 .transpose()
1139 .with_context(|| {
1140 format!("failed to parse message item json in thread {thread_id}")
1141 })?;
1142 out.push(MessageRecord {
1143 id: row.get(0).context("failed to read message id")?,
1144 thread_id: row.get(1).context("failed to read message thread id")?,
1145 role: row.get(2).context("failed to read message role")?,
1146 content: row.get(3).context("failed to read message content")?,
1147 item,
1148 created_at: row.get(5).context("failed to read message timestamp")?,
1149 parent_entry_id: row.get(6).context("failed to read parent entry id")?,
1150 });
1151 }
1152 Ok(out)
1153 }
1154
1155 pub fn fork_at_message(
1161 &self,
1162 message_id: &str,
1163 role: &str,
1164 content: &str,
1165 item: Option<Value>,
1166 ) -> Result<i64> {
1167 let mut conn = self.conn()?;
1168 let created_at = Utc::now().timestamp();
1169 let item_json = item
1170 .as_ref()
1171 .map(serde_json::to_string)
1172 .transpose()
1173 .context("failed to serialize message item payload")?;
1174
1175 let tx = conn
1176 .transaction()
1177 .context("failed to begin fork message transaction")?;
1178
1179 let thread_id: String = tx
1180 .query_row(
1181 "SELECT thread_id FROM messages WHERE id = ?1",
1182 params![message_id],
1183 |row| row.get(0),
1184 )
1185 .with_context(|| format!("failed to query thread id for message {message_id}"))?;
1186
1187 let next_leaf_id: i64 = tx.query_row(
1188 r#"
1189 INSERT INTO messages(thread_id, role, content, item_json, created_at, parent_entry_id)
1190 SELECT ?1, ?2, ?3, ?4, ?5, ?6
1191 RETURNING id
1192 "#, params![thread_id, role, content, item_json, created_at, message_id], |row| row.get(0)
1193 ).with_context(|| format!("failed to fork at message for thread {thread_id:?}"))?;
1194
1195 tx.execute(
1196 r#"
1197 UPDATE threads
1198 SET current_leaf_id = ?1
1199 WHERE id = ?2;
1200 "#,
1201 params![next_leaf_id, thread_id],
1202 )
1203 .with_context(|| {
1204 format!("failed to update thread current leaf id for thread {thread_id:?}")
1205 })?;
1206
1207 tx.commit()
1208 .context("failed to commit fork message transaction")?;
1209
1210 Ok(next_leaf_id)
1211 }
1212
1213 pub fn clear_messages(&self, thread_id: &str) -> Result<usize> {
1217 let mut conn = self.conn()?;
1218 let tx = conn
1219 .transaction()
1220 .context("failed to begin clear messages transaction")?;
1221
1222 tx.execute(
1223 r#"
1224 UPDATE threads
1225 SET current_leaf_id = NULL
1226 WHERE id = ?1;
1227 "#,
1228 params![thread_id],
1229 )
1230 .with_context(|| format!("failed to clear messages for thread {thread_id}"))?;
1231 let result = tx
1232 .execute(
1233 r#"
1234 DELETE FROM messages WHERE thread_id = ?1
1235 "#,
1236 params![thread_id],
1237 )
1238 .with_context(|| format!("failed to clear messages for thread {thread_id}"))?;
1239 tx.commit()
1240 .context("failed to commit clear messages transaction")?;
1241
1242 Ok(result)
1243 }
1244
1245 pub fn save_checkpoint(
1250 &self,
1251 thread_id: &str,
1252 checkpoint_id: &str,
1253 state: &Value,
1254 ) -> Result<()> {
1255 let conn = self.conn()?;
1256 let state_json =
1257 serde_json::to_string(state).context("failed to encode checkpoint state")?;
1258 conn.execute(
1259 r#"
1260 INSERT INTO checkpoints(thread_id, checkpoint_id, state_json, created_at)
1261 VALUES (?1, ?2, ?3, ?4)
1262 ON CONFLICT(thread_id, checkpoint_id) DO UPDATE SET
1263 state_json = excluded.state_json,
1264 created_at = excluded.created_at
1265 "#,
1266 params![thread_id, checkpoint_id, state_json, Utc::now().timestamp()],
1267 )
1268 .with_context(|| {
1269 format!("failed to save checkpoint {checkpoint_id} for thread {thread_id}")
1270 })?;
1271 Ok(())
1272 }
1273
1274 pub fn load_checkpoint(
1280 &self,
1281 thread_id: &str,
1282 checkpoint_id: Option<&str>,
1283 ) -> Result<Option<CheckpointRecord>> {
1284 let conn = self.conn()?;
1285 if let Some(checkpoint_id) = checkpoint_id {
1286 let row = conn
1287 .query_row(
1288 "SELECT thread_id, checkpoint_id, state_json, created_at FROM checkpoints WHERE thread_id = ?1 AND checkpoint_id = ?2",
1289 params![thread_id, checkpoint_id],
1290 |row| {
1291 let state_json: String = row.get(2)?;
1292 let state = serde_json::from_str(&state_json).unwrap_or(Value::Null);
1293 Ok(CheckpointRecord {
1294 thread_id: row.get(0)?,
1295 checkpoint_id: row.get(1)?,
1296 state,
1297 created_at: row.get(3)?,
1298 })
1299 },
1300 )
1301 .optional()
1302 .with_context(|| {
1303 format!("failed to load checkpoint {checkpoint_id} for thread {thread_id}")
1304 })?;
1305 return Ok(row);
1306 }
1307
1308 conn.query_row(
1309 "SELECT thread_id, checkpoint_id, state_json, created_at FROM checkpoints WHERE thread_id = ?1 ORDER BY created_at DESC LIMIT 1",
1310 params![thread_id],
1311 |row| {
1312 let state_json: String = row.get(2)?;
1313 let state = serde_json::from_str(&state_json).unwrap_or(Value::Null);
1314 Ok(CheckpointRecord {
1315 thread_id: row.get(0)?,
1316 checkpoint_id: row.get(1)?,
1317 state,
1318 created_at: row.get(3)?,
1319 })
1320 },
1321 )
1322 .optional()
1323 .with_context(|| format!("failed to load latest checkpoint for thread {thread_id}"))
1324 }
1325
1326 pub fn list_checkpoints(
1330 &self,
1331 thread_id: &str,
1332 limit: Option<usize>,
1333 ) -> Result<Vec<CheckpointRecord>> {
1334 let conn = self.conn()?;
1335 let limit = i64::try_from(limit.unwrap_or(100)).unwrap_or(100);
1336 let mut stmt = conn
1337 .prepare(
1338 "SELECT thread_id, checkpoint_id, state_json, created_at FROM checkpoints WHERE thread_id = ?1 ORDER BY created_at DESC LIMIT ?2",
1339 )
1340 .context("failed to prepare checkpoint list query")?;
1341 let mut rows = stmt
1342 .query(params![thread_id, limit])
1343 .with_context(|| format!("failed to list checkpoints for thread {thread_id}"))?;
1344
1345 let mut out = Vec::new();
1346 while let Some(row) = rows.next().context("failed to iterate checkpoint rows")? {
1347 let state_json: String = row.get(2).context("failed to read checkpoint state json")?;
1348 let state = serde_json::from_str(&state_json).unwrap_or(Value::Null);
1349 out.push(CheckpointRecord {
1350 thread_id: row.get(0).context("failed to read checkpoint thread id")?,
1351 checkpoint_id: row.get(1).context("failed to read checkpoint id")?,
1352 state,
1353 created_at: row.get(3).context("failed to read checkpoint timestamp")?,
1354 });
1355 }
1356 Ok(out)
1357 }
1358
1359 pub fn delete_checkpoint(&self, thread_id: &str, checkpoint_id: &str) -> Result<()> {
1361 let conn = self.conn()?;
1362 conn.execute(
1363 "DELETE FROM checkpoints WHERE thread_id = ?1 AND checkpoint_id = ?2",
1364 params![thread_id, checkpoint_id],
1365 )
1366 .with_context(|| {
1367 format!("failed to delete checkpoint {checkpoint_id} for thread {thread_id}")
1368 })?;
1369 Ok(())
1370 }
1371
1372 pub fn upsert_job(&self, job: &JobStateRecord) -> Result<()> {
1374 let conn = self.conn()?;
1375 conn.execute(
1376 r#"
1377 INSERT INTO jobs(id, name, status, progress, detail, created_at, updated_at)
1378 VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7)
1379 ON CONFLICT(id) DO UPDATE SET
1380 name = excluded.name,
1381 status = excluded.status,
1382 progress = excluded.progress,
1383 detail = excluded.detail,
1384 created_at = excluded.created_at,
1385 updated_at = excluded.updated_at
1386 "#,
1387 params![
1388 job.id,
1389 job.name,
1390 job_state_status_to_str(&job.status),
1391 job.progress.map(i64::from),
1392 job.detail,
1393 job.created_at,
1394 job.updated_at
1395 ],
1396 )
1397 .with_context(|| format!("failed to upsert job {}", job.id))?;
1398 Ok(())
1399 }
1400
1401 pub fn get_job(&self, id: &str) -> Result<Option<JobStateRecord>> {
1405 let conn = self.conn()?;
1406 conn.query_row(
1407 "SELECT id, name, status, progress, detail, created_at, updated_at FROM jobs WHERE id = ?1",
1408 params![id],
1409 |row| {
1410 let status_raw: String = row.get(2)?;
1411 let progress: Option<i64> = row.get(3)?;
1412 Ok(JobStateRecord {
1413 id: row.get(0)?,
1414 name: row.get(1)?,
1415 status: job_state_status_from_str(&status_raw),
1416 progress: progress.and_then(|v| u8::try_from(v).ok()),
1417 detail: row.get(4)?,
1418 created_at: row.get(5)?,
1419 updated_at: row.get(6)?,
1420 })
1421 },
1422 )
1423 .optional()
1424 .with_context(|| format!("failed to read job {id}"))
1425 }
1426
1427 pub fn list_jobs(&self, limit: Option<usize>) -> Result<Vec<JobStateRecord>> {
1431 let conn = self.conn()?;
1432 let limit = i64::try_from(limit.unwrap_or(100)).unwrap_or(100);
1433 let mut stmt = conn
1434 .prepare(
1435 "SELECT id, name, status, progress, detail, created_at, updated_at FROM jobs ORDER BY updated_at DESC LIMIT ?1",
1436 )
1437 .context("failed to prepare job list query")?;
1438 let mut rows = stmt
1439 .query(params![limit])
1440 .context("failed to query persisted jobs")?;
1441 let mut out = Vec::new();
1442 while let Some(row) = rows.next().context("failed to iterate persisted jobs")? {
1443 let status_raw: String = row.get(2).context("failed to read job status")?;
1444 let progress: Option<i64> = row.get(3).context("failed to read job progress")?;
1445 out.push(JobStateRecord {
1446 id: row.get(0).context("failed to read job id")?,
1447 name: row.get(1).context("failed to read job name")?,
1448 status: job_state_status_from_str(&status_raw),
1449 progress: progress.and_then(|v| u8::try_from(v).ok()),
1450 detail: row.get(4).context("failed to read job detail")?,
1451 created_at: row.get(5).context("failed to read job created_at")?,
1452 updated_at: row.get(6).context("failed to read job updated_at")?,
1453 });
1454 }
1455 Ok(out)
1456 }
1457
1458 pub fn delete_job(&self, id: &str) -> Result<()> {
1460 let conn = self.conn()?;
1461 conn.execute("DELETE FROM jobs WHERE id = ?1", params![id])
1462 .with_context(|| format!("failed to delete job {id}"))?;
1463 Ok(())
1464 }
1465
1466 pub fn find_rollout_path_by_id(&self, id: &str) -> Result<Option<PathBuf>> {
1468 let conn = self.conn()?;
1469 conn.query_row(
1470 "SELECT rollout_path FROM threads WHERE id = ?1",
1471 params![id],
1472 |row| row.get::<_, Option<String>>(0),
1473 )
1474 .optional()
1475 .context("failed to lookup rollout path")
1476 .map(|opt| opt.flatten().map(PathBuf::from))
1477 }
1478
1479 pub fn append_thread_name(
1485 &self,
1486 thread_id: &str,
1487 thread_name: Option<String>,
1488 updated_at: i64,
1489 rollout_path: Option<PathBuf>,
1490 ) -> Result<()> {
1491 if let Some(parent) = self.session_index_path.parent() {
1492 fs::create_dir_all(parent).with_context(|| {
1493 format!(
1494 "failed to create session index directory {}",
1495 parent.display()
1496 )
1497 })?;
1498 }
1499 let entry = SessionIndexEntry {
1500 thread_id: thread_id.to_string(),
1501 thread_name,
1502 updated_at,
1503 rollout_path,
1504 };
1505 let encoded =
1506 serde_json::to_string(&entry).context("failed to serialize session index entry")?;
1507 let mut file = OpenOptions::new()
1508 .create(true)
1509 .append(true)
1510 .open(&self.session_index_path)
1511 .with_context(|| {
1512 format!(
1513 "failed to open session index {}",
1514 self.session_index_path.display()
1515 )
1516 })?;
1517 writeln!(file, "{encoded}").context("failed to append session index entry")?;
1518 Ok(())
1519 }
1520
1521 pub fn find_thread_name_by_id(&self, thread_id: &str) -> Result<Option<String>> {
1525 let map = self.session_index_map()?;
1526 Ok(map
1527 .get(thread_id)
1528 .and_then(|entry| entry.thread_name.clone()))
1529 }
1530
1531 pub fn find_thread_names_by_ids(
1535 &self,
1536 ids: &[String],
1537 ) -> Result<HashMap<String, Option<String>>> {
1538 let map = self.session_index_map()?;
1539 let mut out = HashMap::new();
1540 for id in ids {
1541 let name = map.get(id).and_then(|entry| entry.thread_name.clone());
1542 out.insert(id.clone(), name);
1543 }
1544 Ok(out)
1545 }
1546
1547 pub fn find_thread_path_by_name_str(&self, name: &str) -> Result<Option<PathBuf>> {
1552 let map = self.session_index_map()?;
1553 let matched = map
1554 .values()
1555 .filter(|entry| {
1556 entry
1557 .thread_name
1558 .as_deref()
1559 .is_some_and(|n| n.eq_ignore_ascii_case(name))
1560 })
1561 .max_by_key(|entry| entry.updated_at);
1562 Ok(matched.and_then(|entry| entry.rollout_path.clone()))
1563 }
1564
1565 fn session_index_map(&self) -> Result<HashMap<String, SessionIndexEntry>> {
1566 if !self.session_index_path.exists() {
1567 return Ok(HashMap::new());
1568 }
1569 let file = OpenOptions::new()
1570 .read(true)
1571 .open(&self.session_index_path)
1572 .with_context(|| {
1573 format!(
1574 "failed to read session index {}",
1575 self.session_index_path.display()
1576 )
1577 })?;
1578 let reader = BufReader::new(file);
1579 let mut latest = HashMap::<String, SessionIndexEntry>::new();
1580 for line in reader.lines() {
1581 let line = line.context("failed to read session index line")?;
1582 if line.trim().is_empty() {
1583 continue;
1584 }
1585 let parsed: SessionIndexEntry =
1586 serde_json::from_str(&line).context("failed to parse session index entry")?;
1587 latest.insert(parsed.thread_id.clone(), parsed);
1588 }
1589 Ok(latest)
1590 }
1591}
1592
1593fn default_state_db_path() -> PathBuf {
1594 if let Some(overridden) = codewhale_home_override() {
1601 return overridden.join("state.db");
1602 }
1603 let home = dirs::home_dir().unwrap_or_else(|| PathBuf::from("."));
1604 let primary = home.join(".codewhale").join("state.db");
1607 if primary.exists() || !home.join(".deepseek").join("state.db").exists() {
1608 primary
1609 } else {
1610 home.join(".deepseek").join("state.db")
1611 }
1612}
1613
1614fn codewhale_home_override() -> Option<PathBuf> {
1624 std::env::var_os("CODEWHALE_HOME")
1625 .filter(|value| !value.is_empty())
1626 .map(PathBuf::from)
1627}
1628
1629fn bool_to_i64(value: bool) -> i64 {
1630 if value { 1 } else { 0 }
1631}
1632
1633fn i64_to_bool(value: i64) -> bool {
1634 value != 0
1635}
1636
1637fn thread_status_to_str(status: &ThreadStatus) -> &'static str {
1638 match status {
1639 ThreadStatus::Running => "running",
1640 ThreadStatus::Idle => "idle",
1641 ThreadStatus::Completed => "completed",
1642 ThreadStatus::Failed => "failed",
1643 ThreadStatus::Paused => "paused",
1644 ThreadStatus::Archived => "archived",
1645 }
1646}
1647
1648fn thread_status_from_str(value: &str) -> ThreadStatus {
1649 match value {
1650 "running" => ThreadStatus::Running,
1651 "idle" => ThreadStatus::Idle,
1652 "completed" => ThreadStatus::Completed,
1653 "failed" => ThreadStatus::Failed,
1654 "paused" => ThreadStatus::Paused,
1655 "archived" => ThreadStatus::Archived,
1656 _ => ThreadStatus::Idle,
1657 }
1658}
1659
1660fn session_source_to_str(source: &SessionSource) -> &'static str {
1661 match source {
1662 SessionSource::Interactive => "interactive",
1663 SessionSource::Resume => "resume",
1664 SessionSource::Fork => "fork",
1665 SessionSource::Api => "api",
1666 SessionSource::Unknown => "unknown",
1667 }
1668}
1669
1670fn session_source_from_str(value: &str) -> SessionSource {
1671 match value {
1672 "interactive" => SessionSource::Interactive,
1673 "resume" => SessionSource::Resume,
1674 "fork" => SessionSource::Fork,
1675 "api" => SessionSource::Api,
1676 _ => SessionSource::Unknown,
1677 }
1678}
1679
1680fn path_to_opt_string(path: Option<&Path>) -> Option<String> {
1681 path.map(|p| p.display().to_string())
1682}
1683
1684fn job_state_status_to_str(status: &JobStateStatus) -> &'static str {
1685 match status {
1686 JobStateStatus::Queued => "queued",
1687 JobStateStatus::Running => "running",
1688 JobStateStatus::Completed => "completed",
1689 JobStateStatus::Failed => "failed",
1690 JobStateStatus::Cancelled => "cancelled",
1691 }
1692}
1693
1694fn job_state_status_from_str(value: &str) -> JobStateStatus {
1695 match value {
1696 "queued" => JobStateStatus::Queued,
1697 "running" => JobStateStatus::Running,
1698 "completed" => JobStateStatus::Completed,
1699 "failed" => JobStateStatus::Failed,
1700 "cancelled" => JobStateStatus::Cancelled,
1701 _ => JobStateStatus::Queued,
1702 }
1703}
1704
1705fn thread_goal_status_to_str(status: &ThreadGoalStatus) -> &'static str {
1706 match status {
1707 ThreadGoalStatus::Active => "active",
1708 ThreadGoalStatus::Paused => "paused",
1709 ThreadGoalStatus::Blocked => "blocked",
1710 ThreadGoalStatus::UsageLimited => "usage_limited",
1711 ThreadGoalStatus::BudgetLimited => "budget_limited",
1712 ThreadGoalStatus::Complete => "complete",
1713 }
1714}
1715
1716fn thread_goal_status_from_str(value: &str) -> ThreadGoalStatus {
1717 match value {
1718 "active" => ThreadGoalStatus::Active,
1719 "paused" => ThreadGoalStatus::Paused,
1720 "blocked" => ThreadGoalStatus::Blocked,
1721 "usage_limited" => ThreadGoalStatus::UsageLimited,
1722 "budget_limited" => ThreadGoalStatus::BudgetLimited,
1723 "complete" => ThreadGoalStatus::Complete,
1724 _ => ThreadGoalStatus::Active,
1725 }
1726}
1727
1728fn row_to_thread(row: &rusqlite::Row<'_>) -> rusqlite::Result<ThreadMetadata> {
1729 let status_raw: String = row.get(7)?;
1730 let source_raw: String = row.get(11)?;
1731 let rollout_path: Option<String> = row.get(1)?;
1732 let path: Option<String> = row.get(8)?;
1733 Ok(ThreadMetadata {
1734 id: row.get(0)?,
1735 rollout_path: rollout_path.map(PathBuf::from),
1736 preview: row.get(2)?,
1737 ephemeral: i64_to_bool(row.get(3)?),
1738 model_provider: row.get(4)?,
1739 created_at: row.get(5)?,
1740 updated_at: row.get(6)?,
1741 status: thread_status_from_str(&status_raw),
1742 path: path.map(PathBuf::from),
1743 cwd: PathBuf::from(row.get::<_, String>(9)?),
1744 cli_version: row.get(10)?,
1745 source: session_source_from_str(&source_raw),
1746 name: row.get(12)?,
1747 sandbox_policy: row.get(13)?,
1748 approval_mode: row.get(14)?,
1749 archived: i64_to_bool(row.get(15)?),
1750 archived_at: row.get(16)?,
1751 git_sha: row.get(17)?,
1752 git_branch: row.get(18)?,
1753 git_origin_url: row.get(19)?,
1754 memory_mode: row.get(20)?,
1755 current_leaf_id: row.get(21)?,
1756 })
1757}
1758
1759fn row_to_thread_goal(row: &rusqlite::Row<'_>) -> rusqlite::Result<ThreadGoalRecord> {
1760 let status_raw: String = row.get(3)?;
1761 Ok(ThreadGoalRecord {
1762 thread_id: row.get(0)?,
1763 goal_id: row.get(1)?,
1764 objective: row.get(2)?,
1765 status: thread_goal_status_from_str(&status_raw),
1766 token_budget: row.get(4)?,
1767 tokens_used: row.get(5)?,
1768 time_used_seconds: row.get(6)?,
1769 continuation_count: row.get(7)?,
1770 created_at: row.get(8)?,
1771 updated_at: row.get(9)?,
1772 })
1773}
1774
1775#[cfg(test)]
1776mod tests {
1777 use super::*;
1778 use std::time::{SystemTime, UNIX_EPOCH};
1779
1780 fn temp_state_store(name: &str) -> StateStore {
1781 let suffix = SystemTime::now()
1782 .duration_since(UNIX_EPOCH)
1783 .expect("system time")
1784 .as_nanos();
1785 let dir = std::env::temp_dir().join(format!(
1786 "codewhale-state-{name}-{}-{suffix}",
1787 std::process::id()
1788 ));
1789 fs::create_dir_all(&dir).expect("create temp state dir");
1790 StateStore::open(Some(dir.join("state.db"))).expect("open state store")
1791 }
1792
1793 fn test_thread(id: &str) -> ThreadMetadata {
1794 ThreadMetadata {
1795 id: id.to_string(),
1796 rollout_path: None,
1797 preview: "test thread".to_string(),
1798 ephemeral: false,
1799 model_provider: "deepseek".to_string(),
1800 created_at: 10,
1801 updated_at: 10,
1802 status: ThreadStatus::Running,
1803 path: None,
1804 cwd: PathBuf::from("/tmp/codewhale"),
1805 cli_version: "0.0.0-test".to_string(),
1806 source: SessionSource::Interactive,
1807 name: None,
1808 sandbox_policy: None,
1809 approval_mode: None,
1810 archived: false,
1811 archived_at: None,
1812 git_sha: None,
1813 git_branch: None,
1814 git_origin_url: None,
1815 memory_mode: None,
1816 current_leaf_id: None,
1817 }
1818 }
1819
1820 fn test_goal(thread_id: &str, objective: &str) -> ThreadGoalRecord {
1821 ThreadGoalRecord {
1822 thread_id: thread_id.to_string(),
1823 goal_id: "goal-1".to_string(),
1824 objective: objective.to_string(),
1825 status: ThreadGoalStatus::Active,
1826 token_budget: Some(123),
1827 tokens_used: 7,
1828 time_used_seconds: 11,
1829 continuation_count: 0,
1830 created_at: 100,
1831 updated_at: 101,
1832 }
1833 }
1834
1835 #[test]
1836 fn thread_goal_crud_round_trips_and_replaces() {
1837 let store = temp_state_store("thread-goal-crud");
1838 store
1839 .upsert_thread(&test_thread("thread-1"))
1840 .expect("upsert thread");
1841
1842 let goal = test_goal("thread-1", "Ship v0.8.59");
1843 store.upsert_thread_goal(&goal).expect("upsert goal");
1844 assert_eq!(
1845 store
1846 .get_thread_goal("thread-1")
1847 .expect("read goal")
1848 .as_ref(),
1849 Some(&goal)
1850 );
1851
1852 let mut replacement = test_goal("thread-1", "Ship v0.8.59 safely");
1853 replacement.goal_id = "goal-2".to_string();
1854 replacement.status = ThreadGoalStatus::BudgetLimited;
1855 replacement.token_budget = None;
1856 replacement.updated_at = 202;
1857 store
1858 .upsert_thread_goal(&replacement)
1859 .expect("replace goal");
1860 assert_eq!(
1861 store.get_thread_goal("thread-1").expect("read replacement"),
1862 Some(replacement)
1863 );
1864
1865 assert!(store.delete_thread_goal("thread-1").expect("delete goal"));
1866 assert!(
1867 store
1868 .get_thread_goal("thread-1")
1869 .expect("read empty")
1870 .is_none()
1871 );
1872 assert!(!store.delete_thread_goal("thread-1").expect("delete empty"));
1873 }
1874
1875 #[test]
1876 fn thread_goal_requires_existing_thread() {
1877 let store = temp_state_store("thread-goal-missing-thread");
1878 let err = store
1879 .upsert_thread_goal(&test_goal("missing-thread", "nope"))
1880 .expect_err("goal without a thread should fail");
1881 assert!(err.to_string().contains("thread missing-thread not found"));
1882 }
1883
1884 #[test]
1885 fn delete_thread_cascades_child_rows() {
1886 let store = temp_state_store("thread-delete-cascade");
1887 store
1888 .upsert_thread(&test_thread("thread-1"))
1889 .expect("upsert thread");
1890 store
1891 .append_message("thread-1", "user", "hello", None)
1892 .expect("append message");
1893 store
1894 .save_checkpoint("thread-1", "checkpoint-1", &serde_json::json!({"ok": true}))
1895 .expect("save checkpoint");
1896 store
1897 .persist_dynamic_tools(
1898 "thread-1",
1899 &[DynamicToolRecord {
1900 position: 0,
1901 name: "test_tool".to_string(),
1902 description: Some("test".to_string()),
1903 input_schema: serde_json::json!({"type": "object"}),
1904 }],
1905 )
1906 .expect("persist dynamic tools");
1907 store
1908 .upsert_thread_goal(&test_goal("thread-1", "Ship v0.8.67"))
1909 .expect("upsert goal");
1910
1911 store.delete_thread("thread-1").expect("delete thread");
1912
1913 let conn = store.conn().expect("conn");
1914 for table in [
1915 "messages",
1916 "checkpoints",
1917 "thread_dynamic_tools",
1918 "thread_goals",
1919 ] {
1920 let sql = format!("SELECT COUNT(*) FROM {table} WHERE thread_id = ?1");
1921 let count: i64 = conn
1922 .query_row(&sql, params!["thread-1"], |row| row.get(0))
1923 .expect("count child rows");
1924 assert_eq!(count, 0, "{table} row survived thread deletion");
1925 }
1926 }
1927
1928 #[test]
1929 fn state_store_reuses_one_connection_across_operations_and_clones() {
1930 let store = temp_state_store("conn-reuse");
1931 {
1932 let conn = store.conn().expect("conn");
1933 conn.execute_batch("CREATE TEMP TABLE conn_reuse_probe(id INTEGER);")
1934 .expect("create temp table");
1935 }
1936 let clone = store.clone();
1941 clone
1942 .upsert_thread(&test_thread("thread-conn-reuse"))
1943 .expect("upsert thread");
1944 let conn = clone.conn().expect("conn");
1945 let probe_count: i64 = conn
1946 .query_row(
1947 "SELECT COUNT(*) FROM sqlite_temp_master WHERE name = 'conn_reuse_probe'",
1948 [],
1949 |row| row.get(0),
1950 )
1951 .expect("query temp master");
1952 assert_eq!(
1953 probe_count, 1,
1954 "temp table not visible: a fresh connection was opened"
1955 );
1956 let foreign_keys: i64 = conn
1958 .query_row("PRAGMA foreign_keys;", [], |row| row.get(0))
1959 .expect("read foreign_keys pragma");
1960 assert_eq!(foreign_keys, 1);
1961 }
1962
1963 #[test]
1964 fn record_thread_goal_usage_accumulates_tokens_and_time() {
1965 let store = temp_state_store("thread-goal-usage");
1966 store
1967 .upsert_thread(&test_thread("thread-1"))
1968 .expect("upsert thread");
1969
1970 let mut goal = test_goal("thread-1", "Ship the persistent goal loop");
1972 goal.tokens_used = 0;
1973 goal.time_used_seconds = 0;
1974 goal.updated_at = 100;
1975 store.upsert_thread_goal(&goal).expect("upsert goal");
1976
1977 let after_first = store
1979 .record_thread_goal_usage("thread-1", 250, 12, 150)
1980 .expect("record usage")
1981 .expect("goal exists");
1982 assert_eq!(after_first.tokens_used, 250);
1983 assert_eq!(after_first.time_used_seconds, 12);
1984 assert_eq!(after_first.updated_at, 150);
1985 assert_eq!(after_first.goal_id, goal.goal_id);
1987 assert_eq!(after_first.objective, goal.objective);
1988 assert_eq!(after_first.status, goal.status);
1989 assert_eq!(after_first.token_budget, goal.token_budget);
1990 assert_eq!(after_first.created_at, goal.created_at);
1991 assert_eq!(after_first.continuation_count, 0);
1992
1993 let after_second = store
1995 .record_thread_goal_usage("thread-1", 75, 8, 200)
1996 .expect("record usage")
1997 .expect("goal exists");
1998 assert_eq!(after_second.tokens_used, 325);
1999 assert_eq!(after_second.time_used_seconds, 20);
2000 assert_eq!(after_second.updated_at, 200);
2001
2002 let after_stale = store
2004 .record_thread_goal_usage("thread-1", 5, 1, 1)
2005 .expect("record usage")
2006 .expect("goal exists");
2007 assert_eq!(after_stale.tokens_used, 330);
2008 assert_eq!(after_stale.time_used_seconds, 21);
2009 assert_eq!(after_stale.updated_at, 200);
2010
2011 let persisted = store
2013 .get_thread_goal("thread-1")
2014 .expect("read goal")
2015 .expect("goal exists");
2016 assert_eq!(persisted.tokens_used, 330);
2017 assert_eq!(persisted.time_used_seconds, 21);
2018 }
2019
2020 #[test]
2021 fn record_thread_goal_usage_returns_none_without_goal() {
2022 let store = temp_state_store("thread-goal-usage-missing");
2023 store
2024 .upsert_thread(&test_thread("thread-1"))
2025 .expect("upsert thread");
2026 let result = store
2029 .record_thread_goal_usage("thread-1", 100, 5, 999)
2030 .expect("record usage on goalless thread");
2031 assert!(result.is_none());
2032 assert!(
2033 store
2034 .get_thread_goal("thread-1")
2035 .expect("read goal")
2036 .is_none()
2037 );
2038 }
2039
2040 #[test]
2041 fn record_thread_goal_continuation_accumulates_durably() {
2042 let store = temp_state_store("thread-goal-continuation");
2043 store
2044 .upsert_thread(&test_thread("thread-1"))
2045 .expect("upsert thread");
2046
2047 let mut goal = test_goal("thread-1", "Keep working across turns");
2048 goal.updated_at = 100;
2049 store.upsert_thread_goal(&goal).expect("upsert goal");
2050
2051 let after_first = store
2052 .record_thread_goal_continuation("thread-1", 120)
2053 .expect("record continuation")
2054 .expect("goal exists");
2055 assert_eq!(after_first.continuation_count, 1);
2056 assert_eq!(after_first.tokens_used, goal.tokens_used);
2057 assert_eq!(after_first.time_used_seconds, goal.time_used_seconds);
2058 assert_eq!(after_first.updated_at, 120);
2059
2060 let after_second = store
2061 .record_thread_goal_continuation("thread-1", 110)
2062 .expect("record second continuation")
2063 .expect("goal exists");
2064 assert_eq!(after_second.continuation_count, 2);
2065 assert_eq!(after_second.updated_at, 120);
2066
2067 let persisted = store
2068 .get_thread_goal("thread-1")
2069 .expect("read goal")
2070 .expect("goal exists");
2071 assert_eq!(persisted.continuation_count, 2);
2072 }
2073
2074 static CODEWHALE_HOME_TEST_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
2081
2082 struct CodeWhaleHomeGuard {
2083 prior: Option<std::ffi::OsString>,
2084 }
2085 impl CodeWhaleHomeGuard {
2086 fn set(value: &str) -> Self {
2087 let prior = std::env::var_os("CODEWHALE_HOME");
2088 unsafe { std::env::set_var("CODEWHALE_HOME", value) };
2090 Self { prior }
2091 }
2092 fn remove() -> Self {
2093 let prior = std::env::var_os("CODEWHALE_HOME");
2094 unsafe { std::env::remove_var("CODEWHALE_HOME") };
2096 Self { prior }
2097 }
2098 }
2099 impl Drop for CodeWhaleHomeGuard {
2100 fn drop(&mut self) {
2101 unsafe {
2103 match &self.prior {
2104 Some(value) => std::env::set_var("CODEWHALE_HOME", value),
2105 None => std::env::remove_var("CODEWHALE_HOME"),
2106 }
2107 }
2108 }
2109 }
2110
2111 #[test]
2112 fn codewhale_home_override_returns_the_env_value_verbatim() {
2113 let _lock = CODEWHALE_HOME_TEST_LOCK.lock().unwrap();
2114 let _g = CodeWhaleHomeGuard::set("/tmp/cw-isolated-state");
2115 assert_eq!(
2118 codewhale_home_override().as_deref(),
2119 Some(std::path::Path::new("/tmp/cw-isolated-state"))
2120 );
2121 }
2122
2123 #[test]
2124 fn codewhale_home_override_none_when_unset() {
2125 let _lock = CODEWHALE_HOME_TEST_LOCK.lock().unwrap();
2126 let _g = CodeWhaleHomeGuard::remove();
2127 assert!(codewhale_home_override().is_none());
2128 }
2129
2130 #[test]
2131 fn codewhale_home_override_none_when_empty() {
2132 let _lock = CODEWHALE_HOME_TEST_LOCK.lock().unwrap();
2133 let _g = CodeWhaleHomeGuard::set(" ");
2134 assert!(
2140 codewhale_home_override().is_some(),
2141 "non-empty (even whitespace) counts as set; trimming is the caller's job"
2142 );
2143 }
2144
2145 #[test]
2146 fn default_state_db_path_uses_codewhale_home_when_set() {
2147 let _lock = CODEWHALE_HOME_TEST_LOCK.lock().unwrap();
2148 let dir = std::env::temp_dir().join(format!(
2149 "cw-home-state-{}-{}",
2150 std::process::id(),
2151 std::time::SystemTime::now()
2152 .duration_since(std::time::UNIX_EPOCH)
2153 .unwrap()
2154 .as_nanos()
2155 ));
2156 let _g = CodeWhaleHomeGuard::set(dir.to_str().unwrap());
2157 assert_eq!(default_state_db_path(), dir.join("state.db"));
2161 }
2162}