1use std::collections::HashMap;
13use std::fs::{self, OpenOptions};
14use std::io::{BufRead, BufReader, Write};
15use std::path::{Path, PathBuf};
16use std::sync::{Arc, LazyLock, Mutex, MutexGuard};
17
18static SESSION_INDEX_LOCK: LazyLock<Mutex<()>> = LazyLock::new(|| Mutex::new(()));
22
23use anyhow::{Context, Result};
24use chrono::Utc;
25use codewhale_paths::{CODEWHALE_APP_DIR, LEGACY_APP_DIR, codewhale_home_override};
26use rusqlite::{Connection, OptionalExtension, params};
27use serde::{Deserialize, Serialize};
28use serde_json::Value;
29
30pub use codewhale_protocol::ThreadStatus;
33
34#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
38#[serde(rename_all = "snake_case")]
39pub enum SessionSource {
40 Interactive,
42 Resume,
44 Fork,
46 Api,
48 Unknown,
50}
51
52#[derive(Debug, Clone, Serialize, Deserialize)]
57pub struct ThreadMetadata {
58 pub id: String,
60 pub rollout_path: Option<PathBuf>,
62 pub preview: String,
64 pub ephemeral: bool,
66 pub model_provider: String,
68 pub created_at: i64,
70 pub updated_at: i64,
72 pub status: ThreadStatus,
74 pub path: Option<PathBuf>,
76 pub cwd: PathBuf,
78 pub cli_version: String,
80 pub source: SessionSource,
82 pub name: Option<String>,
84 pub sandbox_policy: Option<String>,
86 pub approval_mode: Option<String>,
88 pub archived: bool,
90 pub archived_at: Option<i64>,
92 pub git_sha: Option<String>,
94 pub git_branch: Option<String>,
96 pub git_origin_url: Option<String>,
98 pub memory_mode: Option<String>,
100 pub current_leaf_id: Option<i64>,
102}
103
104#[derive(Debug, Clone, Serialize, Deserialize)]
106pub struct DynamicToolRecord {
107 pub position: i64,
109 pub name: String,
111 pub description: Option<String>,
113 pub input_schema: Value,
115}
116
117#[derive(Debug, Clone, Serialize, Deserialize)]
122pub struct MessageRecord {
123 pub id: i64,
125 pub thread_id: String,
127 pub role: String,
129 pub content: String,
131 pub item: Option<Value>,
133 pub created_at: i64,
135 pub parent_entry_id: Option<i64>,
137}
138
139#[derive(Debug, Clone, Serialize, Deserialize)]
141pub struct CheckpointRecord {
142 pub thread_id: String,
144 pub checkpoint_id: String,
146 pub state: Value,
148 pub created_at: i64,
150}
151
152#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
156#[serde(rename_all = "snake_case")]
157pub enum JobStateStatus {
158 Queued,
160 Running,
162 Paused,
164 Completed,
166 Failed,
168 Cancelled,
170}
171
172#[derive(Debug, Clone, Serialize, Deserialize)]
174pub struct JobStateRecord {
175 pub id: String,
177 pub name: String,
179 pub status: JobStateStatus,
181 pub progress: Option<u8>,
183 pub detail: Option<String>,
185 pub created_at: i64,
187 pub updated_at: i64,
189}
190
191#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
193#[serde(rename_all = "snake_case")]
194pub enum ThreadGoalStatus {
195 Active,
197 Paused,
199 Blocked,
201 UsageLimited,
203 BudgetLimited,
205 Complete,
207}
208
209#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
211pub struct ThreadGoalRecord {
212 pub thread_id: String,
214 pub goal_id: String,
216 pub objective: String,
218 pub status: ThreadGoalStatus,
220 pub token_budget: Option<i64>,
222 pub tokens_used: i64,
224 pub time_used_seconds: i64,
226 pub continuation_count: i64,
228 pub created_at: i64,
230 pub updated_at: i64,
232}
233
234#[derive(Debug, Clone)]
236pub struct ThreadListFilters {
237 pub include_archived: bool,
239 pub limit: Option<usize>,
241}
242
243impl Default for ThreadListFilters {
244 fn default() -> Self {
245 Self {
246 include_archived: false,
247 limit: Some(50),
248 }
249 }
250}
251
252#[derive(Debug, Clone, Serialize, Deserialize)]
253struct SessionIndexEntry {
254 thread_id: String,
255 thread_name: Option<String>,
256 updated_at: i64,
257 rollout_path: Option<PathBuf>,
258}
259
260fn session_index_compact_line_threshold() -> usize {
264 if cfg!(test) { 5 } else { 5_000 }
265}
266
267#[derive(Debug, Clone)]
272pub struct StateStore {
273 db_path: PathBuf,
274 session_index_path: PathBuf,
275 conn: Arc<Mutex<Connection>>,
279}
280
281impl StateStore {
282 pub fn open(path: Option<PathBuf>) -> Result<Self> {
288 let db_path = path.unwrap_or_else(default_state_db_path);
289 let session_index_path = db_path
290 .parent()
291 .unwrap_or_else(|| Path::new("."))
292 .join("session_index.jsonl");
293 if let Some(parent) = db_path.parent() {
294 fs::create_dir_all(parent).with_context(|| {
295 format!("failed to create state directory {}", parent.display())
296 })?;
297 }
298 let conn = Connection::open(&db_path)
299 .with_context(|| format!("failed to open state db {}", db_path.display()))?;
300 Self::configure_connection(&conn, &db_path)?;
301 Self::init_schema(&conn)?;
302 Ok(Self {
303 db_path,
304 session_index_path,
305 conn: Arc::new(Mutex::new(conn)),
306 })
307 }
308
309 fn configure_connection(conn: &Connection, db_path: &Path) -> Result<()> {
316 conn.busy_timeout(std::time::Duration::from_secs(5))
320 .with_context(|| format!("failed to set busy_timeout for {}", db_path.display()))?;
321 conn.pragma_update(None, "foreign_keys", "ON")
322 .with_context(|| format!("failed to enable foreign keys for {}", db_path.display()))?;
323
324 let journal_mode: String = conn
330 .pragma_query_value(None, "journal_mode", |row| row.get(0))
331 .with_context(|| format!("failed to read journal mode for {}", db_path.display()))?;
332 if !journal_mode.eq_ignore_ascii_case("wal") {
333 let configured_mode: String = conn
334 .pragma_update_and_check(None, "journal_mode", "WAL", |row| row.get(0))
335 .with_context(|| format!("failed to enable WAL for {}", db_path.display()))?;
336 if !configured_mode.eq_ignore_ascii_case("wal") {
337 anyhow::bail!(
338 "failed to enable WAL for {}: SQLite retained journal mode {configured_mode}",
339 db_path.display()
340 );
341 }
342 }
343 Ok(())
344 }
345
346 pub fn db_path(&self) -> &Path {
348 &self.db_path
349 }
350
351 fn conn(&self) -> Result<MutexGuard<'_, Connection>> {
352 self.conn
356 .lock()
357 .map_err(|_| anyhow::anyhow!("state db connection mutex poisoned"))
358 }
359
360 fn init_schema(conn: &Connection) -> Result<()> {
361 let mut user_version: u32 = conn.query_row("PRAGMA user_version;", [], |row| row.get(0))?;
362 if user_version == 0 {
363 let add_parent_entry_id = if column_exists(conn, "messages", "parent_entry_id")? {
369 ""
370 } else {
371 "ALTER TABLE messages ADD COLUMN parent_entry_id INTEGER NULL;"
372 };
373 let add_current_leaf_id = if column_exists(conn, "threads", "current_leaf_id")? {
374 ""
375 } else {
376 "ALTER TABLE threads ADD COLUMN current_leaf_id INTEGER NULL;"
377 };
378 conn.execute_batch(&format!(
379 r#"
380 BEGIN;
381 CREATE TABLE IF NOT EXISTS threads (
382 id TEXT PRIMARY KEY,
383 rollout_path TEXT,
384 preview TEXT NOT NULL,
385 ephemeral INTEGER NOT NULL,
386 model_provider TEXT NOT NULL,
387 created_at INTEGER NOT NULL,
388 updated_at INTEGER NOT NULL,
389 status TEXT NOT NULL,
390 path TEXT,
391 cwd TEXT NOT NULL,
392 cli_version TEXT NOT NULL,
393 source TEXT NOT NULL,
394 title TEXT,
395 sandbox_policy TEXT,
396 approval_mode TEXT,
397 archived INTEGER NOT NULL DEFAULT 0,
398 archived_at INTEGER,
399 git_sha TEXT,
400 git_branch TEXT,
401 git_origin_url TEXT,
402 memory_mode TEXT
403 );
404 CREATE INDEX IF NOT EXISTS idx_threads_updated_at ON threads(updated_at DESC);
405 CREATE INDEX IF NOT EXISTS idx_threads_archived_at ON threads(archived_at DESC);
406 CREATE INDEX IF NOT EXISTS idx_threads_archived_updated ON threads(archived, updated_at DESC);
407
408 CREATE TABLE IF NOT EXISTS thread_dynamic_tools (
409 thread_id TEXT NOT NULL,
410 position INTEGER NOT NULL,
411 name TEXT NOT NULL,
412 description TEXT,
413 input_schema TEXT NOT NULL,
414 PRIMARY KEY (thread_id, position),
415 FOREIGN KEY(thread_id) REFERENCES threads(id) ON DELETE CASCADE
416 );
417
418 CREATE TABLE IF NOT EXISTS messages (
419 id INTEGER PRIMARY KEY AUTOINCREMENT,
420 thread_id TEXT NOT NULL,
421 role TEXT NOT NULL,
422 content TEXT NOT NULL,
423 item_json TEXT,
424 created_at INTEGER NOT NULL,
425 FOREIGN KEY(thread_id) REFERENCES threads(id) ON DELETE CASCADE
426 );
427 CREATE INDEX IF NOT EXISTS idx_messages_thread_created_at ON messages(thread_id, created_at ASC);
428
429 CREATE TABLE IF NOT EXISTS checkpoints (
430 thread_id TEXT NOT NULL,
431 checkpoint_id TEXT NOT NULL,
432 state_json TEXT NOT NULL,
433 created_at INTEGER NOT NULL,
434 PRIMARY KEY(thread_id, checkpoint_id),
435 FOREIGN KEY(thread_id) REFERENCES threads(id) ON DELETE CASCADE
436 );
437 CREATE INDEX IF NOT EXISTS idx_checkpoints_thread_created_at ON checkpoints(thread_id, created_at DESC);
438
439 CREATE TABLE IF NOT EXISTS jobs (
440 id TEXT PRIMARY KEY,
441 name TEXT NOT NULL,
442 status TEXT NOT NULL,
443 progress INTEGER,
444 detail TEXT,
445 created_at INTEGER NOT NULL,
446 updated_at INTEGER NOT NULL
447 );
448 CREATE INDEX IF NOT EXISTS idx_jobs_updated_at ON jobs(updated_at DESC);
449
450 -- Add parent_entry_id column, and set to last message before current message
451 {add_parent_entry_id}
452 UPDATE messages
453 SET parent_entry_id = (
454 SELECT m2.id
455 FROM messages m2
456 WHERE m2.thread_id = messages.thread_id
457 AND (
458 m2.created_at < messages.created_at
459 OR (
460 m2.created_at = messages.created_at
461 AND m2.id < messages.id
462 )
463 )
464 ORDER BY m2.created_at DESC, m2.id DESC
465 LIMIT 1
466 );
467 CREATE INDEX IF NOT EXISTS idx_messages_parent_entry_id ON messages(parent_entry_id);
468
469 -- Add current_leaf_id column, and set to last message in thread
470 {add_current_leaf_id}
471 UPDATE threads
472 SET current_leaf_id = (
473 SELECT m.id
474 FROM messages m
475 WHERE m.thread_id = threads.id
476 ORDER BY m.id DESC
477 LIMIT 1
478 );
479
480 PRAGMA user_version = 1;
481 COMMIT;
482 "#
483 ))
484 .context("failed to initialize thread schema")?;
485 user_version = 1;
486 }
487 if user_version < 2 {
488 conn.execute_batch(
489 r#"
490 BEGIN;
491 CREATE TABLE IF NOT EXISTS workflow_runs (
492 id TEXT PRIMARY KEY,
493 workflow_id TEXT NOT NULL,
494 goal TEXT NOT NULL,
495 status TEXT NOT NULL,
496 input_hash TEXT,
497 started_at INTEGER NOT NULL,
498 completed_at INTEGER,
499 metadata_json TEXT NOT NULL DEFAULT '{}'
500 );
501 CREATE INDEX IF NOT EXISTS idx_workflow_runs_status_started_at
502 ON workflow_runs(status, started_at DESC);
503 CREATE INDEX IF NOT EXISTS idx_workflow_runs_workflow_started_at
504 ON workflow_runs(workflow_id, started_at DESC);
505
506 CREATE TABLE IF NOT EXISTS branch_runs (
507 id TEXT PRIMARY KEY,
508 workflow_run_id TEXT NOT NULL,
509 branch_id TEXT NOT NULL,
510 node_id TEXT NOT NULL,
511 status TEXT NOT NULL,
512 started_at INTEGER NOT NULL,
513 completed_at INTEGER,
514 result_json TEXT NOT NULL DEFAULT '{}',
515 FOREIGN KEY(workflow_run_id) REFERENCES workflow_runs(id) ON DELETE CASCADE
516 );
517 CREATE INDEX IF NOT EXISTS idx_branch_runs_workflow_run_id
518 ON branch_runs(workflow_run_id);
519 CREATE INDEX IF NOT EXISTS idx_branch_runs_branch_id
520 ON branch_runs(branch_id);
521
522 CREATE TABLE IF NOT EXISTS leaf_runs (
523 id TEXT PRIMARY KEY,
524 workflow_run_id TEXT NOT NULL,
525 branch_run_id TEXT,
526 leaf_id TEXT NOT NULL,
527 task_id TEXT NOT NULL,
528 input_hash TEXT,
529 status TEXT NOT NULL,
530 output_json TEXT NOT NULL DEFAULT '{}',
531 artifacts_json TEXT NOT NULL DEFAULT '[]',
532 started_at INTEGER NOT NULL,
533 completed_at INTEGER,
534 FOREIGN KEY(workflow_run_id) REFERENCES workflow_runs(id) ON DELETE CASCADE,
535 FOREIGN KEY(branch_run_id) REFERENCES branch_runs(id) ON DELETE SET NULL
536 );
537 CREATE INDEX IF NOT EXISTS idx_leaf_runs_workflow_run_id
538 ON leaf_runs(workflow_run_id);
539 CREATE INDEX IF NOT EXISTS idx_leaf_runs_replay_lookup
540 ON leaf_runs(workflow_run_id, leaf_id, input_hash);
541
542 CREATE TABLE IF NOT EXISTS control_node_runs (
543 id TEXT PRIMARY KEY,
544 workflow_run_id TEXT NOT NULL,
545 node_id TEXT NOT NULL,
546 kind TEXT NOT NULL,
547 status TEXT NOT NULL,
548 selected_children_json TEXT NOT NULL DEFAULT '[]',
549 result_json TEXT NOT NULL DEFAULT '{}',
550 started_at INTEGER NOT NULL,
551 completed_at INTEGER,
552 FOREIGN KEY(workflow_run_id) REFERENCES workflow_runs(id) ON DELETE CASCADE
553 );
554 CREATE INDEX IF NOT EXISTS idx_control_node_runs_workflow_run_id
555 ON control_node_runs(workflow_run_id);
556 CREATE INDEX IF NOT EXISTS idx_control_node_runs_node_id
557 ON control_node_runs(node_id);
558
559 CREATE TABLE IF NOT EXISTS teacher_candidates (
560 id TEXT PRIMARY KEY,
561 workflow_run_id TEXT NOT NULL,
562 control_node_run_id TEXT NOT NULL,
563 candidate_id TEXT NOT NULL,
564 branch_run_id TEXT,
565 score REAL,
566 passed INTEGER,
567 rationale_json TEXT NOT NULL DEFAULT '{}',
568 created_at INTEGER NOT NULL,
569 FOREIGN KEY(workflow_run_id) REFERENCES workflow_runs(id) ON DELETE CASCADE,
570 FOREIGN KEY(control_node_run_id) REFERENCES control_node_runs(id) ON DELETE CASCADE,
571 FOREIGN KEY(branch_run_id) REFERENCES branch_runs(id) ON DELETE SET NULL
572 );
573 CREATE INDEX IF NOT EXISTS idx_teacher_candidates_workflow_run_id
574 ON teacher_candidates(workflow_run_id);
575 CREATE INDEX IF NOT EXISTS idx_teacher_candidates_control_node_run_id
576 ON teacher_candidates(control_node_run_id);
577
578 PRAGMA user_version = 2;
579 COMMIT;
580 "#,
581 )
582 .context("failed to initialize workflow trace schema")?;
583 user_version = 2;
584 }
585 if user_version < 3 {
586 conn.execute_batch(
587 r#"
588 BEGIN;
589 CREATE TABLE IF NOT EXISTS thread_goals (
590 thread_id TEXT PRIMARY KEY NOT NULL,
591 goal_id TEXT NOT NULL,
592 objective TEXT NOT NULL,
593 status TEXT NOT NULL CHECK(status IN (
594 'active',
595 'paused',
596 'blocked',
597 'usage_limited',
598 'budget_limited',
599 'complete'
600 )),
601 token_budget INTEGER,
602 tokens_used INTEGER NOT NULL DEFAULT 0,
603 time_used_seconds INTEGER NOT NULL DEFAULT 0,
604 created_at INTEGER NOT NULL,
605 updated_at INTEGER NOT NULL,
606 FOREIGN KEY(thread_id) REFERENCES threads(id) ON DELETE CASCADE
607 );
608
609 PRAGMA user_version = 3;
610 COMMIT;
611 "#,
612 )
613 .context("failed to initialize thread goal schema")?;
614 user_version = 3;
615 }
616 if user_version < 4 {
617 let add_continuation_count = if column_exists(
620 conn,
621 "thread_goals",
622 "continuation_count",
623 )? {
624 ""
625 } else {
626 "ALTER TABLE thread_goals\n ADD COLUMN continuation_count INTEGER NOT NULL DEFAULT 0;"
627 };
628 conn.execute_batch(&format!(
629 r#"
630 BEGIN;
631 {add_continuation_count}
632
633 PRAGMA user_version = 4;
634 COMMIT;
635 "#
636 ))
637 .context("failed to initialize thread goal continuation schema")?;
638 }
639 Ok(())
640 }
641
642 pub fn upsert_thread(&self, thread: &ThreadMetadata) -> Result<()> {
647 let conn = self.conn()?;
648 conn.execute(
649 r#"
650 INSERT INTO threads (
651 id, rollout_path, preview, ephemeral, model_provider, created_at, updated_at, status, path, cwd,
652 cli_version, source, title, sandbox_policy, approval_mode, archived, archived_at,
653 git_sha, git_branch, git_origin_url, memory_mode
654 ) VALUES (
655 ?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10,
656 ?11, ?12, ?13, ?14, ?15, ?16, ?17,
657 ?18, ?19, ?20, ?21
658 )
659 ON CONFLICT(id) DO UPDATE SET
660 rollout_path=excluded.rollout_path,
661 preview=excluded.preview,
662 ephemeral=excluded.ephemeral,
663 model_provider=excluded.model_provider,
664 created_at=excluded.created_at,
665 updated_at=excluded.updated_at,
666 status=excluded.status,
667 path=excluded.path,
668 cwd=excluded.cwd,
669 cli_version=excluded.cli_version,
670 source=excluded.source,
671 title=excluded.title,
672 sandbox_policy=excluded.sandbox_policy,
673 approval_mode=excluded.approval_mode,
674 archived=excluded.archived,
675 archived_at=excluded.archived_at,
676 git_sha=excluded.git_sha,
677 git_branch=excluded.git_branch,
678 git_origin_url=excluded.git_origin_url,
679 memory_mode=excluded.memory_mode
680 "#,
681 params![
682 thread.id,
683 path_to_opt_string(thread.rollout_path.as_deref()),
684 thread.preview,
685 bool_to_i64(thread.ephemeral),
686 thread.model_provider,
687 thread.created_at,
688 thread.updated_at,
689 thread_status_to_str(&thread.status),
690 path_to_opt_string(thread.path.as_deref()),
691 thread.cwd.display().to_string(),
692 thread.cli_version,
693 session_source_to_str(&thread.source),
694 thread.name,
695 thread.sandbox_policy,
696 thread.approval_mode,
697 bool_to_i64(thread.archived),
698 thread.archived_at,
699 thread.git_sha,
700 thread.git_branch,
701 thread.git_origin_url,
702 thread.memory_mode,
703 ],
704 )
705 .context("failed to upsert thread metadata")?;
706
707 self.append_thread_name(
708 &thread.id,
709 thread.name.clone(),
710 thread.updated_at,
711 thread.rollout_path.clone(),
712 )?;
713 Ok(())
714 }
715
716 pub fn get_thread(&self, id: &str) -> Result<Option<ThreadMetadata>> {
720 let conn = self.conn()?;
721 conn.query_row(
722 r#"
723 SELECT id, rollout_path, preview, ephemeral, model_provider, created_at, updated_at, status, path, cwd,
724 cli_version, source, title, sandbox_policy, approval_mode, archived, archived_at,
725 git_sha, git_branch, git_origin_url, memory_mode, current_leaf_id
726 FROM threads
727 WHERE id = ?1
728 "#,
729 params![id],
730 row_to_thread,
731 )
732 .optional()
733 .context("failed to read thread")
734 }
735
736 pub fn list_threads(&self, filters: ThreadListFilters) -> Result<Vec<ThreadMetadata>> {
741 let conn = self.conn()?;
742 let sql = if filters.include_archived {
743 "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"
744 } else {
745 "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"
746 };
747
748 let mut stmt = conn.prepare(sql).context("failed to prepare list query")?;
749 let limit = i64::try_from(filters.limit.unwrap_or(50)).unwrap_or(50);
750 let mut rows = stmt
751 .query(params![limit])
752 .context("failed to query threads")?;
753 let mut out = Vec::new();
754 while let Some(row) = rows.next().context("failed to iterate thread rows")? {
755 out.push(row_to_thread(row)?);
756 }
757 Ok(out)
758 }
759
760 pub fn mark_archived(&self, id: &str) -> Result<()> {
763 let conn = self.conn()?;
764 conn.execute(
765 "UPDATE threads SET archived = 1, archived_at = ?2, status = ?3 WHERE id = ?1",
766 params![
767 id,
768 Utc::now().timestamp(),
769 thread_status_to_str(&ThreadStatus::Archived)
770 ],
771 )
772 .context("failed to archive thread")?;
773 Ok(())
774 }
775
776 pub fn mark_unarchived(&self, id: &str) -> Result<()> {
778 let conn = self.conn()?;
779 conn.execute(
780 "UPDATE threads SET archived = 0, archived_at = NULL, status = CASE WHEN status = ?2 THEN ?3 ELSE status END WHERE id = ?1",
781 params![
782 id,
783 thread_status_to_str(&ThreadStatus::Archived),
784 thread_status_to_str(&ThreadStatus::Idle),
785 ],
786 )
787 .context("failed to unarchive thread")?;
788 Ok(())
789 }
790
791 pub fn delete_thread(&self, id: &str) -> Result<()> {
794 let conn = self.conn()?;
795 conn.execute("DELETE FROM threads WHERE id = ?1", params![id])
796 .context("failed to delete thread")?;
797 Ok(())
798 }
799
800 pub fn set_thread_memory_mode(&self, id: &str, mode: Option<&str>) -> Result<()> {
804 let conn = self.conn()?;
805 conn.execute(
806 "UPDATE threads SET memory_mode = ?2 WHERE id = ?1",
807 params![id, mode],
808 )
809 .context("failed to update thread memory mode")?;
810 Ok(())
811 }
812
813 pub fn get_thread_memory_mode(&self, id: &str) -> Result<Option<String>> {
817 let conn = self.conn()?;
818 conn.query_row(
819 "SELECT memory_mode FROM threads WHERE id = ?1",
820 params![id],
821 |row| row.get::<_, Option<String>>(0),
822 )
823 .optional()
824 .context("failed to read thread memory mode")
825 .map(Option::flatten)
826 }
827
828 pub fn upsert_thread_goal(&self, goal: &ThreadGoalRecord) -> Result<()> {
830 let conn = self.conn()?;
831 let exists: Option<i64> = conn
832 .query_row(
833 "SELECT 1 FROM threads WHERE id = ?1",
834 params![goal.thread_id],
835 |row| row.get(0),
836 )
837 .optional()
838 .context("failed to verify thread before saving goal")?;
839 if exists.is_none() {
840 anyhow::bail!("thread {} not found", goal.thread_id);
841 }
842
843 conn.execute(
844 r#"
845 INSERT INTO thread_goals (
846 thread_id, goal_id, objective, status, token_budget, tokens_used,
847 time_used_seconds, continuation_count, created_at, updated_at
848 ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10)
849 ON CONFLICT(thread_id) DO UPDATE SET
850 goal_id=excluded.goal_id,
851 objective=excluded.objective,
852 status=excluded.status,
853 token_budget=excluded.token_budget,
854 tokens_used=excluded.tokens_used,
855 time_used_seconds=excluded.time_used_seconds,
856 continuation_count=excluded.continuation_count,
857 created_at=excluded.created_at,
858 updated_at=excluded.updated_at
859 "#,
860 params![
861 goal.thread_id,
862 goal.goal_id,
863 goal.objective,
864 thread_goal_status_to_str(&goal.status),
865 goal.token_budget,
866 goal.tokens_used,
867 goal.time_used_seconds,
868 goal.continuation_count,
869 goal.created_at,
870 goal.updated_at,
871 ],
872 )
873 .context("failed to upsert thread goal")?;
874 Ok(())
875 }
876
877 pub fn record_thread_goal_usage(
893 &self,
894 thread_id: &str,
895 token_delta: i64,
896 time_delta_seconds: i64,
897 now: i64,
898 ) -> Result<Option<ThreadGoalRecord>> {
899 let conn = self.conn()?;
900 let changed = conn
901 .execute(
902 r#"
903 UPDATE thread_goals
904 SET tokens_used = tokens_used + ?2,
905 time_used_seconds = time_used_seconds + ?3,
906 updated_at = MAX(updated_at, ?4)
907 WHERE thread_id = ?1
908 "#,
909 params![thread_id, token_delta, time_delta_seconds, now],
910 )
911 .context("failed to record thread goal usage")?;
912 if changed == 0 {
913 return Ok(None);
914 }
915 Self::read_thread_goal(&conn, thread_id)
916 }
917
918 pub fn record_thread_goal_continuation(
924 &self,
925 thread_id: &str,
926 now: i64,
927 ) -> Result<Option<ThreadGoalRecord>> {
928 let conn = self.conn()?;
929 let changed = conn
930 .execute(
931 r#"
932 UPDATE thread_goals
933 SET continuation_count = continuation_count + 1,
934 updated_at = MAX(updated_at, ?2)
935 WHERE thread_id = ?1
936 "#,
937 params![thread_id, now],
938 )
939 .context("failed to record thread goal continuation")?;
940 if changed == 0 {
941 return Ok(None);
942 }
943 Self::read_thread_goal(&conn, thread_id)
944 }
945
946 pub fn get_thread_goal(&self, thread_id: &str) -> Result<Option<ThreadGoalRecord>> {
948 let conn = self.conn()?;
949 Self::read_thread_goal(&conn, thread_id)
950 }
951
952 fn read_thread_goal(conn: &Connection, thread_id: &str) -> Result<Option<ThreadGoalRecord>> {
956 conn.query_row(
957 r#"
958 SELECT thread_id, goal_id, objective, status, token_budget, tokens_used,
959 time_used_seconds, continuation_count, created_at, updated_at
960 FROM thread_goals
961 WHERE thread_id = ?1
962 "#,
963 params![thread_id],
964 row_to_thread_goal,
965 )
966 .optional()
967 .context("failed to read thread goal")
968 }
969
970 pub fn delete_thread_goal(&self, thread_id: &str) -> Result<bool> {
972 let conn = self.conn()?;
973 let changed = conn
974 .execute(
975 "DELETE FROM thread_goals WHERE thread_id = ?1",
976 params![thread_id],
977 )
978 .context("failed to delete thread goal")?;
979 Ok(changed > 0)
980 }
981
982 pub fn list_leaf_messages(&self, thread_id: &str) -> Result<Vec<MessageRecord>> {
987 let conn = self.conn()?;
988 let mut stmt = conn
989 .prepare(
990 r#"
991 SELECT m1.id, m1.thread_id, m1.role, m1.content, m1.item_json, m1.created_at, m1.parent_entry_id
992 FROM messages m1
993 LEFT JOIN messages m2 ON m1.id = m2.parent_entry_id
994 WHERE m1.thread_id = ?1 AND m2.id IS NULL
995 "#,
996 )
997 .context("failed to prepare message listing query")?;
998 let mut rows = stmt
999 .query(params![thread_id])
1000 .with_context(|| format!("failed to list leaf messages for thread {thread_id}"))?;
1001 let mut out = Vec::new();
1002 while let Some(row) = rows.next().context("failed to iterate message rows")? {
1003 let item_json: Option<String> = row.get(4).context("failed to read item json")?;
1004 let item = item_json
1005 .as_deref()
1006 .map(serde_json::from_str)
1007 .transpose()
1008 .with_context(|| {
1009 format!("failed to parse message item json in thread {thread_id}")
1010 })?;
1011 out.push(MessageRecord {
1012 id: row.get(0).context("failed to read message id")?,
1013 thread_id: row.get(1).context("failed to read message thread id")?,
1014 role: row.get(2).context("failed to read message role")?,
1015 content: row.get(3).context("failed to read message content")?,
1016 item,
1017 created_at: row.get(5).context("failed to read message timestamp")?,
1018 parent_entry_id: row.get(6).context("failed to read parent entry id")?,
1019 });
1020 }
1021 Ok(out)
1022 }
1023
1024 pub fn set_current_leaf_id(&self, thread_id: &str, current_leaf_id: &str) -> Result<()> {
1029 let conn = self.conn()?;
1030 conn.execute(
1031 "UPDATE threads SET current_leaf_id = ?1 WHERE id = ?2",
1032 params![current_leaf_id, thread_id],
1033 )
1034 .context("failed to update thread current leaf id")?;
1035 Ok(())
1036 }
1037
1038 pub fn persist_dynamic_tools(
1043 &self,
1044 thread_id: &str,
1045 tools: &[DynamicToolRecord],
1046 ) -> Result<()> {
1047 let mut conn = self.conn()?;
1048 let tx = conn
1049 .transaction()
1050 .context("failed to begin dynamic tools transaction")?;
1051 tx.execute(
1052 "DELETE FROM thread_dynamic_tools WHERE thread_id = ?1",
1053 params![thread_id],
1054 )
1055 .context("failed to clear dynamic tools")?;
1056 for tool in tools {
1057 tx.execute(
1058 "INSERT INTO thread_dynamic_tools(thread_id, position, name, description, input_schema) VALUES (?1, ?2, ?3, ?4, ?5)",
1059 params![
1060 thread_id,
1061 tool.position,
1062 tool.name,
1063 tool.description,
1064 tool.input_schema.to_string()
1065 ],
1066 )
1067 .with_context(|| format!("failed to persist dynamic tool {}", tool.name))?;
1068 }
1069 tx.commit().context("failed to commit dynamic tools")?;
1070 Ok(())
1071 }
1072
1073 pub fn get_dynamic_tools(&self, thread_id: &str) -> Result<Vec<DynamicToolRecord>> {
1075 let conn = self.conn()?;
1076 let mut stmt = conn
1077 .prepare(
1078 "SELECT position, name, description, input_schema FROM thread_dynamic_tools WHERE thread_id = ?1 ORDER BY position ASC",
1079 )
1080 .context("failed to prepare get dynamic tools query")?;
1081 let mut rows = stmt
1082 .query(params![thread_id])
1083 .context("failed to query dynamic tools")?;
1084 let mut out = Vec::new();
1085 while let Some(row) = rows.next().context("failed to iterate dynamic tools")? {
1086 let input_schema_raw: String =
1087 row.get(3).context("failed to read tool input schema")?;
1088 let input_schema: Value =
1089 serde_json::from_str(&input_schema_raw).with_context(|| {
1090 format!("failed to parse input schema for dynamic tool in thread {thread_id}")
1091 })?;
1092 out.push(DynamicToolRecord {
1093 position: row.get(0).context("failed to read tool position")?,
1094 name: row.get(1).context("failed to read tool name")?,
1095 description: row.get(2).context("failed to read tool description")?,
1096 input_schema,
1097 });
1098 }
1099 Ok(out)
1100 }
1101
1102 pub fn append_message(
1108 &self,
1109 thread_id: &str,
1110 role: &str,
1111 content: &str,
1112 item: Option<Value>,
1113 ) -> Result<i64> {
1114 let mut conn = self.conn()?;
1115 let created_at = Utc::now().timestamp();
1116 let item_json = item
1117 .as_ref()
1118 .map(serde_json::to_string)
1119 .transpose()
1120 .context("failed to serialize message item payload")?;
1121
1122 let tx = conn
1123 .transaction()
1124 .context("failed to begin append message transaction")?;
1125
1126 let current_leaf_id: Option<i64> = tx
1127 .query_row(
1128 "SELECT current_leaf_id FROM threads WHERE id = ?1",
1129 params![thread_id],
1130 |row| row.get(0),
1131 )
1132 .with_context(|| {
1133 format!("failed to query thread current leaf id for thread {thread_id}")
1134 })?;
1135
1136 let next_leaf_id: i64 = tx.query_row(
1137 r#"
1138 INSERT INTO messages(thread_id, role, content, item_json, created_at, parent_entry_id)
1139 SELECT ?1, ?2, ?3, ?4, ?5, ?6
1140 RETURNING id
1141 "#, params![thread_id, role, content, item_json, created_at, current_leaf_id], |row| row.get(0)
1142 ).with_context(|| format!("failed to append message for thread {thread_id}"))?;
1143
1144 tx.execute(
1145 r#"
1146 UPDATE threads
1147 SET current_leaf_id = ?1
1148 WHERE id = ?2;
1149 "#,
1150 params![next_leaf_id, thread_id],
1151 )
1152 .with_context(|| {
1153 format!("failed to update thread current leaf id for thread {thread_id}")
1154 })?;
1155
1156 tx.commit()
1157 .context("failed to commit append message transaction")?;
1158
1159 Ok(next_leaf_id)
1160 }
1161
1162 pub fn list_messages(
1168 &self,
1169 thread_id: &str,
1170 limit: Option<usize>,
1171 ) -> Result<Vec<MessageRecord>> {
1172 let conn = self.conn()?;
1173 let limit = i64::try_from(limit.unwrap_or(500)).unwrap_or(500);
1174 let mut stmt = conn
1175 .prepare(
1176 r#"
1177 WITH RECURSIVE
1178 leaf_id AS (
1179 SELECT current_leaf_id FROM threads WHERE id = ?1
1180 ),
1181 ancestors AS (
1182 SELECT id, thread_id, role, content, item_json, created_at, parent_entry_id, 0 AS depth
1183 FROM messages
1184 WHERE id = (SELECT current_leaf_id FROM leaf_id)
1185
1186 UNION ALL
1187
1188 SELECT m.id, m.thread_id, m.role, m.content, m.item_json, m.created_at, m.parent_entry_id, a.depth + 1
1189 FROM messages m
1190 JOIN ancestors a ON m.id = a.parent_entry_id
1191 WHERE a.depth < ?2
1192 )
1193 SELECT id, thread_id, role, content, item_json, created_at, parent_entry_id FROM ancestors
1194 ORDER BY depth DESC
1195 "#
1196 )
1197 .context("failed to prepare message listing query")?;
1198 let mut rows = stmt
1199 .query(params![thread_id, limit - 1])
1200 .with_context(|| format!("failed to list messages for thread {thread_id}"))?;
1201 let mut out = Vec::new();
1202 while let Some(row) = rows.next().context("failed to iterate message rows")? {
1203 let item_json: Option<String> = row.get(4).context("failed to read item json")?;
1204 let item = item_json
1205 .as_deref()
1206 .map(serde_json::from_str)
1207 .transpose()
1208 .with_context(|| {
1209 format!("failed to parse message item json in thread {thread_id}")
1210 })?;
1211 out.push(MessageRecord {
1212 id: row.get(0).context("failed to read message id")?,
1213 thread_id: row.get(1).context("failed to read message thread id")?,
1214 role: row.get(2).context("failed to read message role")?,
1215 content: row.get(3).context("failed to read message content")?,
1216 item,
1217 created_at: row.get(5).context("failed to read message timestamp")?,
1218 parent_entry_id: row.get(6).context("failed to read parent entry id")?,
1219 });
1220 }
1221 Ok(out)
1222 }
1223
1224 pub fn fork_at_message(
1230 &self,
1231 message_id: &str,
1232 role: &str,
1233 content: &str,
1234 item: Option<Value>,
1235 ) -> Result<i64> {
1236 let mut conn = self.conn()?;
1237 let created_at = Utc::now().timestamp();
1238 let item_json = item
1239 .as_ref()
1240 .map(serde_json::to_string)
1241 .transpose()
1242 .context("failed to serialize message item payload")?;
1243
1244 let tx = conn
1245 .transaction()
1246 .context("failed to begin fork message transaction")?;
1247
1248 let thread_id: String = tx
1249 .query_row(
1250 "SELECT thread_id FROM messages WHERE id = ?1",
1251 params![message_id],
1252 |row| row.get(0),
1253 )
1254 .with_context(|| format!("failed to query thread id for message {message_id}"))?;
1255
1256 let next_leaf_id: i64 = tx.query_row(
1257 r#"
1258 INSERT INTO messages(thread_id, role, content, item_json, created_at, parent_entry_id)
1259 SELECT ?1, ?2, ?3, ?4, ?5, ?6
1260 RETURNING id
1261 "#, params![thread_id, role, content, item_json, created_at, message_id], |row| row.get(0)
1262 ).with_context(|| format!("failed to fork at message for thread {thread_id:?}"))?;
1263
1264 tx.execute(
1265 r#"
1266 UPDATE threads
1267 SET current_leaf_id = ?1
1268 WHERE id = ?2;
1269 "#,
1270 params![next_leaf_id, thread_id],
1271 )
1272 .with_context(|| {
1273 format!("failed to update thread current leaf id for thread {thread_id:?}")
1274 })?;
1275
1276 tx.commit()
1277 .context("failed to commit fork message transaction")?;
1278
1279 Ok(next_leaf_id)
1280 }
1281
1282 pub fn clear_messages(&self, thread_id: &str) -> Result<usize> {
1286 let mut conn = self.conn()?;
1287 let tx = conn
1288 .transaction()
1289 .context("failed to begin clear messages transaction")?;
1290
1291 tx.execute(
1292 r#"
1293 UPDATE threads
1294 SET current_leaf_id = NULL
1295 WHERE id = ?1;
1296 "#,
1297 params![thread_id],
1298 )
1299 .with_context(|| format!("failed to clear messages for thread {thread_id}"))?;
1300 let result = tx
1301 .execute(
1302 r#"
1303 DELETE FROM messages WHERE thread_id = ?1
1304 "#,
1305 params![thread_id],
1306 )
1307 .with_context(|| format!("failed to clear messages for thread {thread_id}"))?;
1308 tx.commit()
1309 .context("failed to commit clear messages transaction")?;
1310
1311 Ok(result)
1312 }
1313
1314 pub fn save_checkpoint(
1319 &self,
1320 thread_id: &str,
1321 checkpoint_id: &str,
1322 state: &Value,
1323 ) -> Result<()> {
1324 let conn = self.conn()?;
1325 let state_json =
1326 serde_json::to_string(state).context("failed to encode checkpoint state")?;
1327 conn.execute(
1328 r#"
1329 INSERT INTO checkpoints(thread_id, checkpoint_id, state_json, created_at)
1330 VALUES (?1, ?2, ?3, ?4)
1331 ON CONFLICT(thread_id, checkpoint_id) DO UPDATE SET
1332 state_json = excluded.state_json,
1333 created_at = excluded.created_at
1334 "#,
1335 params![thread_id, checkpoint_id, state_json, Utc::now().timestamp()],
1336 )
1337 .with_context(|| {
1338 format!("failed to save checkpoint {checkpoint_id} for thread {thread_id}")
1339 })?;
1340 Ok(())
1341 }
1342
1343 pub fn load_checkpoint(
1349 &self,
1350 thread_id: &str,
1351 checkpoint_id: Option<&str>,
1352 ) -> Result<Option<CheckpointRecord>> {
1353 let conn = self.conn()?;
1354 if let Some(checkpoint_id) = checkpoint_id {
1355 let row = conn
1356 .query_row(
1357 "SELECT thread_id, checkpoint_id, state_json, created_at FROM checkpoints WHERE thread_id = ?1 AND checkpoint_id = ?2",
1358 params![thread_id, checkpoint_id],
1359 |row| {
1360 Ok((
1361 row.get::<_, String>(0)?,
1362 row.get::<_, String>(1)?,
1363 row.get::<_, String>(2)?,
1364 row.get::<_, i64>(3)?,
1365 ))
1366 },
1367 )
1368 .optional()
1369 .with_context(|| {
1370 format!("failed to load checkpoint {checkpoint_id} for thread {thread_id}")
1371 })?;
1372 if let Some((thread_id, checkpoint_id, state_json, created_at)) = row {
1373 let state = parse_checkpoint_state(&state_json)?;
1374 return Ok(Some(CheckpointRecord {
1375 thread_id,
1376 checkpoint_id,
1377 state,
1378 created_at,
1379 }));
1380 }
1381 return Ok(None);
1382 }
1383
1384 let row = conn
1385 .query_row(
1386 "SELECT thread_id, checkpoint_id, state_json, created_at FROM checkpoints WHERE thread_id = ?1 ORDER BY created_at DESC LIMIT 1",
1387 params![thread_id],
1388 |row| {
1389 Ok((
1390 row.get::<_, String>(0)?,
1391 row.get::<_, String>(1)?,
1392 row.get::<_, String>(2)?,
1393 row.get::<_, i64>(3)?,
1394 ))
1395 },
1396 )
1397 .optional()
1398 .with_context(|| format!("failed to load latest checkpoint for thread {thread_id}"))?;
1399 if let Some((thread_id, checkpoint_id, state_json, created_at)) = row {
1400 let state = parse_checkpoint_state(&state_json)?;
1401 return Ok(Some(CheckpointRecord {
1402 thread_id,
1403 checkpoint_id,
1404 state,
1405 created_at,
1406 }));
1407 }
1408 Ok(None)
1409 }
1410
1411 pub fn list_checkpoints(
1415 &self,
1416 thread_id: &str,
1417 limit: Option<usize>,
1418 ) -> Result<Vec<CheckpointRecord>> {
1419 let conn = self.conn()?;
1420 let limit = i64::try_from(limit.unwrap_or(100)).unwrap_or(100);
1421 let mut stmt = conn
1422 .prepare(
1423 "SELECT thread_id, checkpoint_id, state_json, created_at FROM checkpoints WHERE thread_id = ?1 ORDER BY created_at DESC LIMIT ?2",
1424 )
1425 .context("failed to prepare checkpoint list query")?;
1426 let mut rows = stmt
1427 .query(params![thread_id, limit])
1428 .with_context(|| format!("failed to list checkpoints for thread {thread_id}"))?;
1429
1430 let mut out = Vec::new();
1431 while let Some(row) = rows.next().context("failed to iterate checkpoint rows")? {
1432 let state_json: String = row.get(2).context("failed to read checkpoint state json")?;
1433 let state = parse_checkpoint_state(&state_json)?;
1434 out.push(CheckpointRecord {
1435 thread_id: row.get(0).context("failed to read checkpoint thread id")?,
1436 checkpoint_id: row.get(1).context("failed to read checkpoint id")?,
1437 state,
1438 created_at: row.get(3).context("failed to read checkpoint timestamp")?,
1439 });
1440 }
1441 Ok(out)
1442 }
1443
1444 pub fn delete_checkpoint(&self, thread_id: &str, checkpoint_id: &str) -> Result<()> {
1446 let conn = self.conn()?;
1447 conn.execute(
1448 "DELETE FROM checkpoints WHERE thread_id = ?1 AND checkpoint_id = ?2",
1449 params![thread_id, checkpoint_id],
1450 )
1451 .with_context(|| {
1452 format!("failed to delete checkpoint {checkpoint_id} for thread {thread_id}")
1453 })?;
1454 Ok(())
1455 }
1456
1457 pub fn upsert_job(&self, job: &JobStateRecord) -> Result<()> {
1459 let conn = self.conn()?;
1460 conn.execute(
1461 r#"
1462 INSERT INTO jobs(id, name, status, progress, detail, created_at, updated_at)
1463 VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7)
1464 ON CONFLICT(id) DO UPDATE SET
1465 name = excluded.name,
1466 status = excluded.status,
1467 progress = excluded.progress,
1468 detail = excluded.detail,
1469 created_at = excluded.created_at,
1470 updated_at = excluded.updated_at
1471 "#,
1472 params![
1473 job.id,
1474 job.name,
1475 job_state_status_to_str(&job.status),
1476 job.progress.map(i64::from),
1477 job.detail,
1478 job.created_at,
1479 job.updated_at
1480 ],
1481 )
1482 .with_context(|| format!("failed to upsert job {}", job.id))?;
1483 Ok(())
1484 }
1485
1486 pub fn get_job(&self, id: &str) -> Result<Option<JobStateRecord>> {
1490 let conn = self.conn()?;
1491 conn.query_row(
1492 "SELECT id, name, status, progress, detail, created_at, updated_at FROM jobs WHERE id = ?1",
1493 params![id],
1494 |row| {
1495 let status_raw: String = row.get(2)?;
1496 let progress: Option<i64> = row.get(3)?;
1497 Ok(JobStateRecord {
1498 id: row.get(0)?,
1499 name: row.get(1)?,
1500 status: job_state_status_from_str(&status_raw),
1501 progress: progress.and_then(|v| u8::try_from(v).ok()),
1502 detail: row.get(4)?,
1503 created_at: row.get(5)?,
1504 updated_at: row.get(6)?,
1505 })
1506 },
1507 )
1508 .optional()
1509 .with_context(|| format!("failed to read job {id}"))
1510 }
1511
1512 pub fn list_jobs(&self, limit: Option<usize>) -> Result<Vec<JobStateRecord>> {
1516 let conn = self.conn()?;
1517 let limit = i64::try_from(limit.unwrap_or(100)).unwrap_or(100);
1518 let mut stmt = conn
1519 .prepare(
1520 "SELECT id, name, status, progress, detail, created_at, updated_at FROM jobs ORDER BY updated_at DESC LIMIT ?1",
1521 )
1522 .context("failed to prepare job list query")?;
1523 let mut rows = stmt
1524 .query(params![limit])
1525 .context("failed to query persisted jobs")?;
1526 let mut out = Vec::new();
1527 while let Some(row) = rows.next().context("failed to iterate persisted jobs")? {
1528 let status_raw: String = row.get(2).context("failed to read job status")?;
1529 let progress: Option<i64> = row.get(3).context("failed to read job progress")?;
1530 out.push(JobStateRecord {
1531 id: row.get(0).context("failed to read job id")?,
1532 name: row.get(1).context("failed to read job name")?,
1533 status: job_state_status_from_str(&status_raw),
1534 progress: progress.and_then(|v| u8::try_from(v).ok()),
1535 detail: row.get(4).context("failed to read job detail")?,
1536 created_at: row.get(5).context("failed to read job created_at")?,
1537 updated_at: row.get(6).context("failed to read job updated_at")?,
1538 });
1539 }
1540 Ok(out)
1541 }
1542
1543 pub fn delete_job(&self, id: &str) -> Result<()> {
1545 let conn = self.conn()?;
1546 conn.execute("DELETE FROM jobs WHERE id = ?1", params![id])
1547 .with_context(|| format!("failed to delete job {id}"))?;
1548 Ok(())
1549 }
1550
1551 pub fn find_rollout_path_by_id(&self, id: &str) -> Result<Option<PathBuf>> {
1553 let conn = self.conn()?;
1554 conn.query_row(
1555 "SELECT rollout_path FROM threads WHERE id = ?1",
1556 params![id],
1557 |row| row.get::<_, Option<String>>(0),
1558 )
1559 .optional()
1560 .context("failed to lookup rollout path")
1561 .map(|opt| opt.flatten().map(PathBuf::from))
1562 }
1563
1564 pub fn append_thread_name(
1570 &self,
1571 thread_id: &str,
1572 thread_name: Option<String>,
1573 updated_at: i64,
1574 rollout_path: Option<PathBuf>,
1575 ) -> Result<()> {
1576 let _guard = SESSION_INDEX_LOCK.lock().unwrap();
1579 if let Some(parent) = self.session_index_path.parent() {
1580 fs::create_dir_all(parent).with_context(|| {
1581 format!(
1582 "failed to create session index directory {}",
1583 parent.display()
1584 )
1585 })?;
1586 }
1587 let entry = SessionIndexEntry {
1588 thread_id: thread_id.to_string(),
1589 thread_name,
1590 updated_at,
1591 rollout_path,
1592 };
1593 let encoded =
1594 serde_json::to_string(&entry).context("failed to serialize session index entry")?;
1595 self.with_session_index_lock(|| {
1600 let mut file = OpenOptions::new()
1601 .create(true)
1602 .append(true)
1603 .open(&self.session_index_path)
1604 .with_context(|| {
1605 format!(
1606 "failed to open session index {}",
1607 self.session_index_path.display()
1608 )
1609 })?;
1610 writeln!(file, "{encoded}").context("failed to append session index entry")?;
1611 file.sync_data()
1615 .context("failed to flush session index entry")?;
1616 drop(file);
1617 self.compact_session_index_locked()
1618 })
1619 }
1620
1621 fn with_session_index_lock<T>(&self, operation: impl FnOnce() -> Result<T>) -> Result<T> {
1627 if let Some(parent) = self.session_index_path.parent() {
1628 fs::create_dir_all(parent).with_context(|| {
1629 format!(
1630 "failed to create session index directory {}",
1631 parent.display()
1632 )
1633 })?;
1634 }
1635 let lock_path = self.session_index_path.with_extension("jsonl.lock");
1636 let lock_file = OpenOptions::new()
1637 .create(true)
1638 .read(true)
1639 .write(true)
1640 .truncate(false)
1643 .open(&lock_path)
1644 .with_context(|| {
1645 format!("failed to open session index lock {}", lock_path.display())
1646 })?;
1647 #[cfg(unix)]
1648 {
1649 use std::os::unix::fs::PermissionsExt as _;
1650 lock_file
1651 .set_permissions(fs::Permissions::from_mode(0o600))
1652 .with_context(|| {
1653 format!(
1654 "failed to secure session index lock {}",
1655 lock_path.display()
1656 )
1657 })?;
1658 }
1659 let mut lock = fd_lock::RwLock::new(lock_file);
1660 let _guard = lock
1661 .write()
1662 .with_context(|| format!("failed to lock session index {}", lock_path.display()))?;
1663 operation()
1664 }
1665
1666 pub fn find_thread_name_by_id(&self, thread_id: &str) -> Result<Option<String>> {
1670 let map = self.session_index_map()?;
1671 Ok(map
1672 .get(thread_id)
1673 .and_then(|entry| entry.thread_name.clone()))
1674 }
1675
1676 pub fn find_thread_names_by_ids(
1680 &self,
1681 ids: &[String],
1682 ) -> Result<HashMap<String, Option<String>>> {
1683 let map = self.session_index_map()?;
1684 let mut out = HashMap::new();
1685 for id in ids {
1686 let name = map.get(id).and_then(|entry| entry.thread_name.clone());
1687 out.insert(id.clone(), name);
1688 }
1689 Ok(out)
1690 }
1691
1692 pub fn find_thread_path_by_name_str(&self, name: &str) -> Result<Option<PathBuf>> {
1697 let map = self.session_index_map()?;
1698 let matched = map
1699 .values()
1700 .filter(|entry| {
1701 entry
1702 .thread_name
1703 .as_deref()
1704 .is_some_and(|n| n.eq_ignore_ascii_case(name))
1705 })
1706 .max_by_key(|entry| entry.updated_at);
1707 Ok(matched.and_then(|entry| entry.rollout_path.clone()))
1708 }
1709
1710 fn compact_session_index_locked(&self) -> Result<()> {
1715 if !self.session_index_path.exists() {
1716 return Ok(());
1717 }
1718 let line_count = BufReader::new(
1719 OpenOptions::new()
1720 .read(true)
1721 .open(&self.session_index_path)
1722 .with_context(|| {
1723 format!(
1724 "failed to read session index {}",
1725 self.session_index_path.display()
1726 )
1727 })?,
1728 )
1729 .lines()
1730 .filter(|line| {
1731 line.as_ref()
1732 .map(|value| !value.trim().is_empty())
1733 .unwrap_or(false)
1734 })
1735 .count();
1736 if line_count <= session_index_compact_line_threshold() {
1737 return Ok(());
1738 }
1739
1740 let latest = self.session_index_map()?;
1741 let compact_path = self.session_index_path.with_extension("jsonl.compact");
1742 {
1743 let mut file = OpenOptions::new()
1744 .create(true)
1745 .write(true)
1746 .truncate(true)
1747 .open(&compact_path)
1748 .with_context(|| {
1749 format!(
1750 "failed to open compact session index {}",
1751 compact_path.display()
1752 )
1753 })?;
1754 for entry in latest.values() {
1755 let encoded = serde_json::to_string(entry)
1756 .context("failed to serialize compact session index entry")?;
1757 writeln!(file, "{encoded}")
1758 .context("failed to write compact session index entry")?;
1759 }
1760 }
1761 #[cfg(test)]
1765 tests::compaction_midpoint(&self.session_index_path);
1766 fs::rename(&compact_path, &self.session_index_path).with_context(|| {
1767 format!(
1768 "failed to replace session index {}",
1769 self.session_index_path.display()
1770 )
1771 })?;
1772 Ok(())
1773 }
1774
1775 #[cfg(test)]
1776 fn session_index_line_count(&self) -> Result<usize> {
1777 if !self.session_index_path.exists() {
1778 return Ok(0);
1779 }
1780 Ok(BufReader::new(
1781 OpenOptions::new()
1782 .read(true)
1783 .open(&self.session_index_path)
1784 .with_context(|| {
1785 format!(
1786 "failed to read session index {}",
1787 self.session_index_path.display()
1788 )
1789 })?,
1790 )
1791 .lines()
1792 .filter(|line| {
1793 line.as_ref()
1794 .map(|value| !value.trim().is_empty())
1795 .unwrap_or(false)
1796 })
1797 .count())
1798 }
1799
1800 fn session_index_map(&self) -> Result<HashMap<String, SessionIndexEntry>> {
1801 if !self.session_index_path.exists() {
1802 return Ok(HashMap::new());
1803 }
1804 let file = OpenOptions::new()
1805 .read(true)
1806 .open(&self.session_index_path)
1807 .with_context(|| {
1808 format!(
1809 "failed to read session index {}",
1810 self.session_index_path.display()
1811 )
1812 })?;
1813 let reader = BufReader::new(file);
1814 let mut latest = HashMap::<String, SessionIndexEntry>::new();
1815 for line in reader.lines() {
1816 let line = line.context("failed to read session index line")?;
1817 if line.trim().is_empty() {
1818 continue;
1819 }
1820 match serde_json::from_str::<SessionIndexEntry>(&line) {
1827 Ok(parsed) => {
1828 latest.insert(parsed.thread_id.clone(), parsed);
1829 }
1830 Err(err) => {
1831 tracing::warn!(
1832 "skipping unparseable session index entry in {}: {err}",
1833 self.session_index_path.display()
1834 );
1835 }
1836 }
1837 }
1838 Ok(latest)
1839 }
1840}
1841
1842#[must_use]
1848pub fn default_state_db_path() -> PathBuf {
1849 if let Some(overridden) = codewhale_home_override().ok().flatten() {
1856 return overridden.join("state.db");
1857 }
1858 let home = codewhale_paths::user_home().unwrap_or_else(|| PathBuf::from("."));
1859 let primary = home.join(CODEWHALE_APP_DIR).join("state.db");
1862 if primary.exists() || !home.join(LEGACY_APP_DIR).join("state.db").exists() {
1863 primary
1864 } else {
1865 home.join(LEGACY_APP_DIR).join("state.db")
1866 }
1867}
1868
1869fn bool_to_i64(value: bool) -> i64 {
1870 if value { 1 } else { 0 }
1871}
1872
1873fn column_exists(conn: &Connection, table: &str, column: &str) -> Result<bool> {
1880 let mut stmt = conn.prepare(&format!("PRAGMA table_info({table})"))?;
1881 let names = stmt.query_map([], |row| row.get::<_, String>(1))?;
1882 for name in names {
1883 if name? == column {
1884 return Ok(true);
1885 }
1886 }
1887 Ok(false)
1888}
1889
1890fn i64_to_bool(value: i64) -> bool {
1891 value != 0
1892}
1893
1894fn thread_status_to_str(status: &ThreadStatus) -> &'static str {
1895 match status {
1896 ThreadStatus::Running => "running",
1897 ThreadStatus::Idle => "idle",
1898 ThreadStatus::Completed => "completed",
1899 ThreadStatus::Failed => "failed",
1900 ThreadStatus::Paused => "paused",
1901 ThreadStatus::Archived => "archived",
1902 }
1903}
1904
1905fn thread_status_from_str(value: &str) -> ThreadStatus {
1906 match value {
1907 "running" => ThreadStatus::Running,
1908 "idle" => ThreadStatus::Idle,
1909 "completed" => ThreadStatus::Completed,
1910 "failed" => ThreadStatus::Failed,
1911 "paused" => ThreadStatus::Paused,
1912 "archived" => ThreadStatus::Archived,
1913 _ => ThreadStatus::Idle,
1914 }
1915}
1916
1917fn session_source_to_str(source: &SessionSource) -> &'static str {
1918 match source {
1919 SessionSource::Interactive => "interactive",
1920 SessionSource::Resume => "resume",
1921 SessionSource::Fork => "fork",
1922 SessionSource::Api => "api",
1923 SessionSource::Unknown => "unknown",
1924 }
1925}
1926
1927fn session_source_from_str(value: &str) -> SessionSource {
1928 match value {
1929 "interactive" => SessionSource::Interactive,
1930 "resume" => SessionSource::Resume,
1931 "fork" => SessionSource::Fork,
1932 "api" => SessionSource::Api,
1933 _ => SessionSource::Unknown,
1934 }
1935}
1936
1937fn path_to_opt_string(path: Option<&Path>) -> Option<String> {
1938 path.map(|p| p.display().to_string())
1939}
1940
1941fn parse_checkpoint_state(state_json: &str) -> Result<Value> {
1942 serde_json::from_str(state_json).context("failed to parse checkpoint state json")
1943}
1944
1945fn job_state_status_to_str(status: &JobStateStatus) -> &'static str {
1946 match status {
1947 JobStateStatus::Queued => "queued",
1948 JobStateStatus::Running => "running",
1949 JobStateStatus::Paused => "paused",
1950 JobStateStatus::Completed => "completed",
1951 JobStateStatus::Failed => "failed",
1952 JobStateStatus::Cancelled => "cancelled",
1953 }
1954}
1955
1956fn job_state_status_from_str(value: &str) -> JobStateStatus {
1957 match value {
1958 "queued" => JobStateStatus::Queued,
1959 "running" => JobStateStatus::Running,
1960 "paused" => JobStateStatus::Paused,
1961 "completed" => JobStateStatus::Completed,
1962 "failed" => JobStateStatus::Failed,
1963 "cancelled" => JobStateStatus::Cancelled,
1964 _ => JobStateStatus::Queued,
1965 }
1966}
1967
1968fn thread_goal_status_to_str(status: &ThreadGoalStatus) -> &'static str {
1969 match status {
1970 ThreadGoalStatus::Active => "active",
1971 ThreadGoalStatus::Paused => "paused",
1972 ThreadGoalStatus::Blocked => "blocked",
1973 ThreadGoalStatus::UsageLimited => "usage_limited",
1974 ThreadGoalStatus::BudgetLimited => "budget_limited",
1975 ThreadGoalStatus::Complete => "complete",
1976 }
1977}
1978
1979fn thread_goal_status_from_str(value: &str) -> ThreadGoalStatus {
1980 match value {
1981 "active" => ThreadGoalStatus::Active,
1982 "paused" => ThreadGoalStatus::Paused,
1983 "blocked" => ThreadGoalStatus::Blocked,
1984 "usage_limited" => ThreadGoalStatus::UsageLimited,
1985 "budget_limited" => ThreadGoalStatus::BudgetLimited,
1986 "complete" => ThreadGoalStatus::Complete,
1987 _ => ThreadGoalStatus::Paused,
1991 }
1992}
1993
1994fn row_to_thread(row: &rusqlite::Row<'_>) -> rusqlite::Result<ThreadMetadata> {
1995 let status_raw: String = row.get(7)?;
1996 let source_raw: String = row.get(11)?;
1997 let rollout_path: Option<String> = row.get(1)?;
1998 let path: Option<String> = row.get(8)?;
1999 Ok(ThreadMetadata {
2000 id: row.get(0)?,
2001 rollout_path: rollout_path.map(PathBuf::from),
2002 preview: row.get(2)?,
2003 ephemeral: i64_to_bool(row.get(3)?),
2004 model_provider: row.get(4)?,
2005 created_at: row.get(5)?,
2006 updated_at: row.get(6)?,
2007 status: thread_status_from_str(&status_raw),
2008 path: path.map(PathBuf::from),
2009 cwd: PathBuf::from(row.get::<_, String>(9)?),
2010 cli_version: row.get(10)?,
2011 source: session_source_from_str(&source_raw),
2012 name: row.get(12)?,
2013 sandbox_policy: row.get(13)?,
2014 approval_mode: row.get(14)?,
2015 archived: i64_to_bool(row.get(15)?),
2016 archived_at: row.get(16)?,
2017 git_sha: row.get(17)?,
2018 git_branch: row.get(18)?,
2019 git_origin_url: row.get(19)?,
2020 memory_mode: row.get(20)?,
2021 current_leaf_id: row.get(21)?,
2022 })
2023}
2024
2025fn row_to_thread_goal(row: &rusqlite::Row<'_>) -> rusqlite::Result<ThreadGoalRecord> {
2026 let status_raw: String = row.get(3)?;
2027 Ok(ThreadGoalRecord {
2028 thread_id: row.get(0)?,
2029 goal_id: row.get(1)?,
2030 objective: row.get(2)?,
2031 status: thread_goal_status_from_str(&status_raw),
2032 token_budget: row.get(4)?,
2033 tokens_used: row.get(5)?,
2034 time_used_seconds: row.get(6)?,
2035 continuation_count: row.get(7)?,
2036 created_at: row.get(8)?,
2037 updated_at: row.get(9)?,
2038 })
2039}
2040
2041#[cfg(test)]
2042mod tests {
2043 use super::*;
2044 use serde_json::json;
2045 use std::sync::{Arc, Barrier, Mutex, mpsc};
2046 use std::thread;
2047 use std::time::{Duration, SystemTime, UNIX_EPOCH};
2048
2049 fn temp_state_dir(name: &str) -> PathBuf {
2050 let suffix = SystemTime::now()
2051 .duration_since(UNIX_EPOCH)
2052 .expect("system time")
2053 .as_nanos();
2054 let dir = std::env::temp_dir().join(format!(
2055 "codewhale-state-{name}-{}-{suffix}",
2056 std::process::id()
2057 ));
2058 fs::create_dir_all(&dir).expect("create temp state dir");
2059 dir
2060 }
2061
2062 fn temp_state_store(name: &str) -> StateStore {
2063 let dir = temp_state_dir(name);
2064 StateStore::open(Some(dir.join("state.db"))).expect("open state store")
2065 }
2066
2067 fn test_thread(id: &str) -> ThreadMetadata {
2068 ThreadMetadata {
2069 id: id.to_string(),
2070 rollout_path: None,
2071 preview: "test thread".to_string(),
2072 ephemeral: false,
2073 model_provider: "deepseek".to_string(),
2074 created_at: 10,
2075 updated_at: 10,
2076 status: ThreadStatus::Running,
2077 path: None,
2078 cwd: PathBuf::from("/tmp/codewhale"),
2079 cli_version: "0.0.0-test".to_string(),
2080 source: SessionSource::Interactive,
2081 name: None,
2082 sandbox_policy: None,
2083 approval_mode: None,
2084 archived: false,
2085 archived_at: None,
2086 git_sha: None,
2087 git_branch: None,
2088 git_origin_url: None,
2089 memory_mode: None,
2090 current_leaf_id: None,
2091 }
2092 }
2093
2094 fn test_goal(thread_id: &str, objective: &str) -> ThreadGoalRecord {
2095 ThreadGoalRecord {
2096 thread_id: thread_id.to_string(),
2097 goal_id: "goal-1".to_string(),
2098 objective: objective.to_string(),
2099 status: ThreadGoalStatus::Active,
2100 token_budget: Some(123),
2101 tokens_used: 7,
2102 time_used_seconds: 11,
2103 continuation_count: 0,
2104 created_at: 100,
2105 updated_at: 101,
2106 }
2107 }
2108
2109 #[test]
2110 fn unknown_persisted_goal_status_fails_closed() {
2111 assert_eq!(
2112 thread_goal_status_from_str("future_or_corrupt_status"),
2113 ThreadGoalStatus::Paused
2114 );
2115 }
2116
2117 #[test]
2118 fn thread_goal_crud_round_trips_and_replaces() {
2119 let store = temp_state_store("thread-goal-crud");
2120 store
2121 .upsert_thread(&test_thread("thread-1"))
2122 .expect("upsert thread");
2123
2124 let goal = test_goal("thread-1", "Ship v0.8.59");
2125 store.upsert_thread_goal(&goal).expect("upsert goal");
2126 assert_eq!(
2127 store
2128 .get_thread_goal("thread-1")
2129 .expect("read goal")
2130 .as_ref(),
2131 Some(&goal)
2132 );
2133
2134 let mut replacement = test_goal("thread-1", "Ship v0.8.59 safely");
2135 replacement.goal_id = "goal-2".to_string();
2136 replacement.status = ThreadGoalStatus::BudgetLimited;
2137 replacement.token_budget = None;
2138 replacement.updated_at = 202;
2139 store
2140 .upsert_thread_goal(&replacement)
2141 .expect("replace goal");
2142 assert_eq!(
2143 store.get_thread_goal("thread-1").expect("read replacement"),
2144 Some(replacement)
2145 );
2146
2147 assert!(store.delete_thread_goal("thread-1").expect("delete goal"));
2148 assert!(
2149 store
2150 .get_thread_goal("thread-1")
2151 .expect("read empty")
2152 .is_none()
2153 );
2154 assert!(!store.delete_thread_goal("thread-1").expect("delete empty"));
2155 }
2156
2157 #[test]
2158 fn thread_goal_requires_existing_thread() {
2159 let store = temp_state_store("thread-goal-missing-thread");
2160 let err = store
2161 .upsert_thread_goal(&test_goal("missing-thread", "nope"))
2162 .expect_err("goal without a thread should fail");
2163 assert!(err.to_string().contains("thread missing-thread not found"));
2164 }
2165
2166 #[test]
2167 fn delete_thread_cascades_child_rows() {
2168 let store = temp_state_store("thread-delete-cascade");
2169 store
2170 .upsert_thread(&test_thread("thread-1"))
2171 .expect("upsert thread");
2172 store
2173 .append_message("thread-1", "user", "hello", None)
2174 .expect("append message");
2175 store
2176 .save_checkpoint("thread-1", "checkpoint-1", &serde_json::json!({"ok": true}))
2177 .expect("save checkpoint");
2178 store
2179 .persist_dynamic_tools(
2180 "thread-1",
2181 &[DynamicToolRecord {
2182 position: 0,
2183 name: "test_tool".to_string(),
2184 description: Some("test".to_string()),
2185 input_schema: serde_json::json!({"type": "object"}),
2186 }],
2187 )
2188 .expect("persist dynamic tools");
2189 store
2190 .upsert_thread_goal(&test_goal("thread-1", "Ship v0.8.67"))
2191 .expect("upsert goal");
2192
2193 store.delete_thread("thread-1").expect("delete thread");
2194
2195 let conn = store.conn().expect("conn");
2196 for table in [
2197 "messages",
2198 "checkpoints",
2199 "thread_dynamic_tools",
2200 "thread_goals",
2201 ] {
2202 let sql = format!("SELECT COUNT(*) FROM {table} WHERE thread_id = ?1");
2203 let count: i64 = conn
2204 .query_row(&sql, params!["thread-1"], |row| row.get(0))
2205 .expect("count child rows");
2206 assert_eq!(count, 0, "{table} row survived thread deletion");
2207 }
2208 }
2209
2210 #[test]
2211 fn state_store_reuses_one_connection_across_operations_and_clones() {
2212 let store = temp_state_store("conn-reuse");
2213 {
2214 let conn = store.conn().expect("conn");
2215 conn.execute_batch("CREATE TEMP TABLE conn_reuse_probe(id INTEGER);")
2216 .expect("create temp table");
2217 }
2218 let clone = store.clone();
2223 clone
2224 .upsert_thread(&test_thread("thread-conn-reuse"))
2225 .expect("upsert thread");
2226 let conn = clone.conn().expect("conn");
2227 let probe_count: i64 = conn
2228 .query_row(
2229 "SELECT COUNT(*) FROM sqlite_temp_master WHERE name = 'conn_reuse_probe'",
2230 [],
2231 |row| row.get(0),
2232 )
2233 .expect("query temp master");
2234 assert_eq!(
2235 probe_count, 1,
2236 "temp table not visible: a fresh connection was opened"
2237 );
2238 let foreign_keys: i64 = conn
2240 .query_row("PRAGMA foreign_keys;", [], |row| row.get(0))
2241 .expect("read foreign_keys pragma");
2242 assert_eq!(foreign_keys, 1);
2243 let journal_mode: String = conn
2244 .query_row("PRAGMA journal_mode;", [], |row| row.get(0))
2245 .expect("read journal_mode pragma");
2246 assert_eq!(
2247 journal_mode.to_ascii_lowercase(),
2248 "wal",
2249 "open should enable WAL for multi-process readers/writers"
2250 );
2251 }
2252
2253 #[test]
2254 fn connection_setup_waits_for_database_lock_before_enabling_wal() {
2255 let dir = temp_state_dir("locked-open");
2256 let db_path = dir.join("state.db");
2257
2258 let candidate = Connection::open(&db_path).expect("open candidate connection");
2259 candidate
2263 .busy_timeout(Duration::ZERO)
2264 .expect("disable dependency default timeout");
2265 let blocker = Connection::open(&db_path).expect("open blocking connection");
2266 let (locked_tx, locked_rx) = mpsc::sync_channel(0);
2267 let blocker_thread = thread::spawn(move || {
2268 blocker
2269 .execute_batch("BEGIN EXCLUSIVE;")
2270 .expect("acquire exclusive database lock");
2271 locked_tx.send(()).expect("announce database lock");
2272 thread::sleep(Duration::from_millis(200));
2273 blocker
2274 .execute_batch("COMMIT;")
2275 .expect("release exclusive database lock");
2276 });
2277
2278 locked_rx.recv().expect("wait for database lock");
2279 StateStore::configure_connection(&candidate, &db_path)
2280 .expect("connection setup should wait for the brief database lock");
2281 blocker_thread.join().expect("blocking thread panicked");
2282
2283 let journal_mode: String = candidate
2284 .query_row("PRAGMA journal_mode;", [], |row| row.get(0))
2285 .expect("read journal_mode");
2286 assert_eq!(journal_mode.to_ascii_lowercase(), "wal");
2287
2288 drop(candidate);
2289 let _ = fs::remove_dir_all(dir);
2290 }
2291
2292 #[test]
2299 fn second_connection_waits_for_active_writer() {
2300 let dir = temp_state_dir("concurrent-write");
2301 let db_path = dir.join("state.db");
2302
2303 let store_a = StateStore::open(Some(db_path.clone())).expect("open store a");
2304 let store_b = StateStore::open(Some(db_path.clone())).expect("open store b");
2305 let (locked_tx, locked_rx) = mpsc::sync_channel(0);
2306 let (release_tx, release_rx) = mpsc::sync_channel(0);
2307
2308 let writer_a = thread::spawn(move || {
2309 let conn = store_a.conn().expect("connection a");
2310 conn.execute_batch(
2311 r#"
2312 BEGIN IMMEDIATE;
2313 INSERT INTO jobs(id, name, status, created_at, updated_at)
2314 VALUES ('job-a', 'writer-a', 'running', 0, 0);
2315 "#,
2316 )
2317 .expect("writer a should acquire the database write lock");
2318 locked_tx.send(()).expect("announce active writer");
2319 release_rx.recv().expect("wait to release active writer");
2320 conn.execute_batch("COMMIT;")
2321 .expect("writer a should commit");
2322 });
2323
2324 locked_rx.recv().expect("wait for active writer");
2325 let (attempting_tx, attempting_rx) = mpsc::sync_channel(0);
2326 let writer_b = thread::spawn(move || {
2327 attempting_tx.send(()).expect("announce second write");
2328 store_b.upsert_job(&JobStateRecord {
2329 id: "job-b".to_string(),
2330 name: "writer-b".to_string(),
2331 status: JobStateStatus::Running,
2332 progress: None,
2333 detail: Some("waited for writer a".to_string()),
2334 created_at: 1,
2335 updated_at: 1,
2336 })
2337 });
2338
2339 attempting_rx.recv().expect("wait for second write attempt");
2340 thread::sleep(Duration::from_millis(100));
2341 assert!(
2342 !writer_b.is_finished(),
2343 "second writer should still be waiting while the first holds the lock"
2344 );
2345 release_tx.send(()).expect("release active writer");
2346 writer_a.join().expect("writer a panicked");
2347 writer_b
2348 .join()
2349 .expect("writer b panicked")
2350 .expect("writer b should succeed after the lock is released");
2351
2352 let store = StateStore::open(Some(db_path)).expect("reopen for verify");
2353 let listed = store.list_jobs(Some(2)).expect("list jobs");
2354 assert_eq!(listed.len(), 2, "both writers should persist their jobs");
2355
2356 let _ = fs::remove_dir_all(dir);
2357 }
2358
2359 #[test]
2360 fn migration_runs_cleanly_when_schema_predates_user_version_header() {
2361 let dir = temp_state_dir("migration-v0-idempotent");
2367 let db_path = dir.join("state.db");
2368 drop(StateStore::open(Some(db_path.clone())).expect("initial open"));
2369 {
2370 let conn = Connection::open(&db_path).expect("raw connection");
2371 conn.pragma_update(None, "user_version", 0)
2372 .expect("reset user_version");
2373 }
2374
2375 let store = StateStore::open(Some(db_path.clone())).expect("reopen with v0 header");
2376 store
2377 .upsert_thread(&test_thread("thread-migrated"))
2378 .expect("write after guarded migration");
2379
2380 drop(store);
2382 let store = StateStore::open(Some(db_path)).expect("third open");
2383 let persisted = store
2384 .get_thread("thread-migrated")
2385 .expect("read after reopen");
2386 assert!(persisted.is_some());
2387
2388 let _ = fs::remove_dir_all(dir);
2389 }
2390
2391 #[test]
2392 fn record_thread_goal_usage_accumulates_tokens_and_time() {
2393 let store = temp_state_store("thread-goal-usage");
2394 store
2395 .upsert_thread(&test_thread("thread-1"))
2396 .expect("upsert thread");
2397
2398 let mut goal = test_goal("thread-1", "Ship the persistent goal loop");
2400 goal.tokens_used = 0;
2401 goal.time_used_seconds = 0;
2402 goal.updated_at = 100;
2403 store.upsert_thread_goal(&goal).expect("upsert goal");
2404
2405 let after_first = store
2407 .record_thread_goal_usage("thread-1", 250, 12, 150)
2408 .expect("record usage")
2409 .expect("goal exists");
2410 assert_eq!(after_first.tokens_used, 250);
2411 assert_eq!(after_first.time_used_seconds, 12);
2412 assert_eq!(after_first.updated_at, 150);
2413 assert_eq!(after_first.goal_id, goal.goal_id);
2415 assert_eq!(after_first.objective, goal.objective);
2416 assert_eq!(after_first.status, goal.status);
2417 assert_eq!(after_first.token_budget, goal.token_budget);
2418 assert_eq!(after_first.created_at, goal.created_at);
2419 assert_eq!(after_first.continuation_count, 0);
2420
2421 let after_second = store
2423 .record_thread_goal_usage("thread-1", 75, 8, 200)
2424 .expect("record usage")
2425 .expect("goal exists");
2426 assert_eq!(after_second.tokens_used, 325);
2427 assert_eq!(after_second.time_used_seconds, 20);
2428 assert_eq!(after_second.updated_at, 200);
2429
2430 let after_stale = store
2432 .record_thread_goal_usage("thread-1", 5, 1, 1)
2433 .expect("record usage")
2434 .expect("goal exists");
2435 assert_eq!(after_stale.tokens_used, 330);
2436 assert_eq!(after_stale.time_used_seconds, 21);
2437 assert_eq!(after_stale.updated_at, 200);
2438
2439 let persisted = store
2441 .get_thread_goal("thread-1")
2442 .expect("read goal")
2443 .expect("goal exists");
2444 assert_eq!(persisted.tokens_used, 330);
2445 assert_eq!(persisted.time_used_seconds, 21);
2446 }
2447
2448 #[test]
2449 fn record_thread_goal_usage_returns_none_without_goal() {
2450 let store = temp_state_store("thread-goal-usage-missing");
2451 store
2452 .upsert_thread(&test_thread("thread-1"))
2453 .expect("upsert thread");
2454 let result = store
2457 .record_thread_goal_usage("thread-1", 100, 5, 999)
2458 .expect("record usage on goalless thread");
2459 assert!(result.is_none());
2460 assert!(
2461 store
2462 .get_thread_goal("thread-1")
2463 .expect("read goal")
2464 .is_none()
2465 );
2466 }
2467
2468 #[test]
2469 fn record_thread_goal_continuation_accumulates_durably() {
2470 let store = temp_state_store("thread-goal-continuation");
2471 store
2472 .upsert_thread(&test_thread("thread-1"))
2473 .expect("upsert thread");
2474
2475 let mut goal = test_goal("thread-1", "Keep working across turns");
2476 goal.updated_at = 100;
2477 store.upsert_thread_goal(&goal).expect("upsert goal");
2478
2479 let after_first = store
2480 .record_thread_goal_continuation("thread-1", 120)
2481 .expect("record continuation")
2482 .expect("goal exists");
2483 assert_eq!(after_first.continuation_count, 1);
2484 assert_eq!(after_first.tokens_used, goal.tokens_used);
2485 assert_eq!(after_first.time_used_seconds, goal.time_used_seconds);
2486 assert_eq!(after_first.updated_at, 120);
2487
2488 let after_second = store
2489 .record_thread_goal_continuation("thread-1", 110)
2490 .expect("record second continuation")
2491 .expect("goal exists");
2492 assert_eq!(after_second.continuation_count, 2);
2493 assert_eq!(after_second.updated_at, 120);
2494
2495 let persisted = store
2496 .get_thread_goal("thread-1")
2497 .expect("read goal")
2498 .expect("goal exists");
2499 assert_eq!(persisted.continuation_count, 2);
2500 }
2501
2502 static CODEWHALE_HOME_TEST_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
2509
2510 struct CodeWhaleHomeGuard {
2511 prior: Option<std::ffi::OsString>,
2512 }
2513 impl CodeWhaleHomeGuard {
2514 fn set(value: &str) -> Self {
2515 let prior = std::env::var_os("CODEWHALE_HOME");
2516 unsafe { std::env::set_var("CODEWHALE_HOME", value) };
2518 Self { prior }
2519 }
2520 fn remove() -> Self {
2521 let prior = std::env::var_os("CODEWHALE_HOME");
2522 unsafe { std::env::remove_var("CODEWHALE_HOME") };
2524 Self { prior }
2525 }
2526 }
2527 impl Drop for CodeWhaleHomeGuard {
2528 fn drop(&mut self) {
2529 unsafe {
2531 match &self.prior {
2532 Some(value) => std::env::set_var("CODEWHALE_HOME", value),
2533 None => std::env::remove_var("CODEWHALE_HOME"),
2534 }
2535 }
2536 }
2537 }
2538
2539 #[test]
2540 fn codewhale_home_override_returns_the_env_value_verbatim() {
2541 let _lock = CODEWHALE_HOME_TEST_LOCK.lock().unwrap();
2542 let override_path = std::env::temp_dir().join("cw-isolated-state");
2543 let _g = CodeWhaleHomeGuard::set(override_path.to_str().unwrap());
2544 assert_eq!(
2547 codewhale_home_override().unwrap().as_deref(),
2548 Some(override_path.as_path())
2549 );
2550 }
2551
2552 #[test]
2553 fn codewhale_home_override_none_when_unset() {
2554 let _lock = CODEWHALE_HOME_TEST_LOCK.lock().unwrap();
2555 let _g = CodeWhaleHomeGuard::remove();
2556 assert!(codewhale_home_override().unwrap().is_none());
2557 }
2558
2559 #[test]
2560 fn codewhale_home_override_none_when_whitespace_only() {
2561 let _lock = CODEWHALE_HOME_TEST_LOCK.lock().unwrap();
2562 let _g = CodeWhaleHomeGuard::set(" ");
2563 assert!(
2564 codewhale_home_override().unwrap().is_none(),
2565 "whitespace-only CODEWHALE_HOME must not establish isolation"
2566 );
2567 }
2568
2569 #[test]
2570 fn default_state_db_path_uses_codewhale_home_when_set() {
2571 let _lock = CODEWHALE_HOME_TEST_LOCK.lock().unwrap();
2572 let dir = std::env::temp_dir().join(format!(
2573 "cw-home-state-{}-{}",
2574 std::process::id(),
2575 std::time::SystemTime::now()
2576 .duration_since(std::time::UNIX_EPOCH)
2577 .unwrap()
2578 .as_nanos()
2579 ));
2580 let _g = CodeWhaleHomeGuard::set(dir.to_str().unwrap());
2581 assert_eq!(default_state_db_path(), dir.join("state.db"));
2585 }
2586
2587 #[test]
2588 fn load_checkpoint_propagates_invalid_state_json() {
2589 let store = temp_state_store("checkpoint-parse-error");
2590 store
2591 .upsert_thread(&test_thread("thread-1"))
2592 .expect("upsert thread");
2593 store
2594 .save_checkpoint("thread-1", "broken", &json!({"ok": true}))
2595 .expect("save checkpoint");
2596
2597 {
2598 let conn = store.conn().expect("conn");
2599 conn.execute(
2600 "UPDATE checkpoints SET state_json = ?1 WHERE thread_id = ?2 AND checkpoint_id = ?3",
2601 params!["not-json", "thread-1", "broken"],
2602 )
2603 .expect("corrupt checkpoint");
2604 }
2605
2606 let err = store
2607 .load_checkpoint("thread-1", Some("broken"))
2608 .expect_err("invalid checkpoint json should fail");
2609 assert!(
2610 err.to_string()
2611 .contains("failed to parse checkpoint state json")
2612 );
2613 }
2614
2615 #[test]
2616 fn session_index_compacts_after_threshold() {
2617 let store = temp_state_store("session-index-compact");
2618 for idx in 0..6 {
2619 store
2620 .append_thread_name("thread-1", Some(format!("name-{idx}")), idx, None)
2621 .expect("append session index entry");
2622 }
2623
2624 let line_count = store
2625 .session_index_line_count()
2626 .expect("count session index lines");
2627 assert_eq!(line_count, 1);
2628
2629 let name = store
2630 .find_thread_name_by_id("thread-1")
2631 .expect("lookup thread name");
2632 assert_eq!(name.as_deref(), Some("name-5"));
2633 }
2634
2635 #[test]
2636 fn session_index_read_skips_a_torn_line() {
2637 let store = temp_state_store("session-index-torn");
2641 store
2642 .append_thread_name("thread-1", Some("first".to_string()), 1, None)
2643 .expect("append first entry");
2644
2645 {
2646 let mut file = OpenOptions::new()
2647 .append(true)
2648 .open(&store.session_index_path)
2649 .expect("open session index");
2650 writeln!(file, "{{\"thread_id\":\"thread-2\",\"thread_na").expect("write torn line");
2652 }
2653
2654 store
2655 .append_thread_name("thread-3", Some("third".to_string()), 3, None)
2656 .expect("append third entry");
2657
2658 assert_eq!(
2659 store
2660 .find_thread_name_by_id("thread-1")
2661 .expect("lookup thread-1")
2662 .as_deref(),
2663 Some("first"),
2664 );
2665 assert_eq!(
2666 store
2667 .find_thread_name_by_id("thread-3")
2668 .expect("lookup thread-3")
2669 .as_deref(),
2670 Some("third"),
2671 );
2672 }
2673
2674 type MidpointHook = Box<dyn Fn() + Send + Sync>;
2679 static COMPACTION_MIDPOINT: Mutex<Option<(PathBuf, MidpointHook)>> = Mutex::new(None);
2680
2681 pub(super) fn compaction_midpoint(index_path: &Path) {
2684 let mut hook = COMPACTION_MIDPOINT.lock().expect("midpoint hook lock");
2685 let registered_for_this_store = match hook.as_ref() {
2686 Some((registered, _)) => registered == index_path,
2687 None => return,
2688 };
2689 if !registered_for_this_store {
2690 return;
2691 }
2692 let (_, callback) = hook.take().expect("presence checked above");
2693 drop(hook);
2694 callback();
2695 }
2696
2697 #[test]
2698 fn session_index_compaction_does_not_drop_a_concurrent_append() {
2699 let store = Arc::new(temp_state_store("session-index-race"));
2710 let threshold = session_index_compact_line_threshold();
2711
2712 for idx in 0..threshold {
2714 store
2715 .append_thread_name(
2716 &format!("thread-{idx}"),
2717 Some(format!("name-{idx}")),
2718 1,
2719 None,
2720 )
2721 .expect("append filler entry");
2722 }
2723
2724 let appender_released = Arc::new(Barrier::new(2));
2725 {
2726 let released = Arc::clone(&appender_released);
2727 *COMPACTION_MIDPOINT.lock().expect("midpoint hook lock") = Some((
2728 store.session_index_path.clone(),
2729 Box::new(move || {
2730 released.wait();
2731 thread::sleep(Duration::from_millis(300));
2734 }),
2735 ));
2736 }
2737
2738 let appender = {
2739 let store = Arc::clone(&store);
2740 let released = Arc::clone(&appender_released);
2741 thread::spawn(move || {
2742 released.wait();
2743 store
2744 .append_thread_name("racer", Some("racer-name".to_string()), 2, None)
2745 .expect("append racing entry");
2746 })
2747 };
2748
2749 store
2750 .append_thread_name("trigger", Some("trigger-name".to_string()), 1, None)
2751 .expect("append entry that triggers compaction");
2752 appender.join().expect("appender thread");
2753 *COMPACTION_MIDPOINT.lock().expect("midpoint hook lock") = None;
2754
2755 assert_eq!(
2756 store
2757 .find_thread_name_by_id("racer")
2758 .expect("lookup racer")
2759 .as_deref(),
2760 Some("racer-name"),
2761 "append was dropped by a concurrent compaction",
2762 );
2763 assert_eq!(
2764 store
2765 .find_thread_name_by_id("trigger")
2766 .expect("lookup trigger")
2767 .as_deref(),
2768 Some("trigger-name"),
2769 );
2770 }
2771}