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
24pub use codewhale_protocol::ThreadStatus;
27
28#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
32#[serde(rename_all = "snake_case")]
33pub enum SessionSource {
34 Interactive,
36 Resume,
38 Fork,
40 Api,
42 Unknown,
44}
45
46#[derive(Debug, Clone, Serialize, Deserialize)]
51pub struct ThreadMetadata {
52 pub id: String,
54 pub rollout_path: Option<PathBuf>,
56 pub preview: String,
58 pub ephemeral: bool,
60 pub model_provider: String,
62 pub created_at: i64,
64 pub updated_at: i64,
66 pub status: ThreadStatus,
68 pub path: Option<PathBuf>,
70 pub cwd: PathBuf,
72 pub cli_version: String,
74 pub source: SessionSource,
76 pub name: Option<String>,
78 pub sandbox_policy: Option<String>,
80 pub approval_mode: Option<String>,
82 pub archived: bool,
84 pub archived_at: Option<i64>,
86 pub git_sha: Option<String>,
88 pub git_branch: Option<String>,
90 pub git_origin_url: Option<String>,
92 pub memory_mode: Option<String>,
94 pub current_leaf_id: Option<i64>,
96}
97
98#[derive(Debug, Clone, Serialize, Deserialize)]
100pub struct DynamicToolRecord {
101 pub position: i64,
103 pub name: String,
105 pub description: Option<String>,
107 pub input_schema: Value,
109}
110
111#[derive(Debug, Clone, Serialize, Deserialize)]
116pub struct MessageRecord {
117 pub id: i64,
119 pub thread_id: String,
121 pub role: String,
123 pub content: String,
125 pub item: Option<Value>,
127 pub created_at: i64,
129 pub parent_entry_id: Option<i64>,
131}
132
133#[derive(Debug, Clone, Serialize, Deserialize)]
135pub struct CheckpointRecord {
136 pub thread_id: String,
138 pub checkpoint_id: String,
140 pub state: Value,
142 pub created_at: i64,
144}
145
146#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
150#[serde(rename_all = "snake_case")]
151pub enum JobStateStatus {
152 Queued,
154 Running,
156 Paused,
158 Completed,
160 Failed,
162 Cancelled,
164}
165
166#[derive(Debug, Clone, Serialize, Deserialize)]
168pub struct JobStateRecord {
169 pub id: String,
171 pub name: String,
173 pub status: JobStateStatus,
175 pub progress: Option<u8>,
177 pub detail: Option<String>,
179 pub created_at: i64,
181 pub updated_at: i64,
183}
184
185#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
187#[serde(rename_all = "snake_case")]
188pub enum ThreadGoalStatus {
189 Active,
191 Paused,
193 Blocked,
195 UsageLimited,
197 BudgetLimited,
199 Complete,
201}
202
203#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
205pub struct ThreadGoalRecord {
206 pub thread_id: String,
208 pub goal_id: String,
210 pub objective: String,
212 pub status: ThreadGoalStatus,
214 pub token_budget: Option<i64>,
216 pub tokens_used: i64,
218 pub time_used_seconds: i64,
220 pub continuation_count: i64,
222 pub created_at: i64,
224 pub updated_at: i64,
226}
227
228#[derive(Debug, Clone)]
230pub struct ThreadListFilters {
231 pub include_archived: bool,
233 pub limit: Option<usize>,
235}
236
237impl Default for ThreadListFilters {
238 fn default() -> Self {
239 Self {
240 include_archived: false,
241 limit: Some(50),
242 }
243 }
244}
245
246#[derive(Debug, Clone, Serialize, Deserialize)]
247struct SessionIndexEntry {
248 thread_id: String,
249 thread_name: Option<String>,
250 updated_at: i64,
251 rollout_path: Option<PathBuf>,
252}
253
254fn session_index_compact_line_threshold() -> usize {
258 if cfg!(test) { 5 } else { 5_000 }
259}
260
261#[derive(Debug, Clone)]
266pub struct StateStore {
267 db_path: PathBuf,
268 session_index_path: PathBuf,
269 conn: Arc<Mutex<Connection>>,
273}
274
275impl StateStore {
276 pub fn open(path: Option<PathBuf>) -> Result<Self> {
282 let db_path = path.unwrap_or_else(default_state_db_path);
283 let session_index_path = db_path
284 .parent()
285 .unwrap_or_else(|| Path::new("."))
286 .join("session_index.jsonl");
287 if let Some(parent) = db_path.parent() {
288 fs::create_dir_all(parent).with_context(|| {
289 format!("failed to create state directory {}", parent.display())
290 })?;
291 }
292 let conn = Connection::open(&db_path)
293 .with_context(|| format!("failed to open state db {}", db_path.display()))?;
294 Self::configure_connection(&conn, &db_path)?;
295 Self::init_schema(&conn)?;
296 Ok(Self {
297 db_path,
298 session_index_path,
299 conn: Arc::new(Mutex::new(conn)),
300 })
301 }
302
303 fn configure_connection(conn: &Connection, db_path: &Path) -> Result<()> {
310 conn.pragma_update(None, "foreign_keys", "ON")
311 .with_context(|| format!("failed to enable foreign keys for {}", db_path.display()))?;
312 conn.pragma_update(None, "journal_mode", "WAL")
315 .with_context(|| format!("failed to enable WAL for {}", db_path.display()))?;
316 conn.busy_timeout(std::time::Duration::from_secs(5))
319 .with_context(|| format!("failed to set busy_timeout for {}", db_path.display()))?;
320 Ok(())
321 }
322
323 pub fn db_path(&self) -> &Path {
325 &self.db_path
326 }
327
328 fn conn(&self) -> Result<MutexGuard<'_, Connection>> {
329 self.conn
333 .lock()
334 .map_err(|_| anyhow::anyhow!("state db connection mutex poisoned"))
335 }
336
337 fn init_schema(conn: &Connection) -> Result<()> {
338 let mut user_version: u32 = conn.query_row("PRAGMA user_version;", [], |row| row.get(0))?;
339 if user_version == 0 {
340 conn.execute_batch(
341 r#"
342 BEGIN;
343 CREATE TABLE IF NOT EXISTS threads (
344 id TEXT PRIMARY KEY,
345 rollout_path TEXT,
346 preview TEXT NOT NULL,
347 ephemeral INTEGER NOT NULL,
348 model_provider TEXT NOT NULL,
349 created_at INTEGER NOT NULL,
350 updated_at INTEGER NOT NULL,
351 status TEXT NOT NULL,
352 path TEXT,
353 cwd TEXT NOT NULL,
354 cli_version TEXT NOT NULL,
355 source TEXT NOT NULL,
356 title TEXT,
357 sandbox_policy TEXT,
358 approval_mode TEXT,
359 archived INTEGER NOT NULL DEFAULT 0,
360 archived_at INTEGER,
361 git_sha TEXT,
362 git_branch TEXT,
363 git_origin_url TEXT,
364 memory_mode TEXT
365 );
366 CREATE INDEX IF NOT EXISTS idx_threads_updated_at ON threads(updated_at DESC);
367 CREATE INDEX IF NOT EXISTS idx_threads_archived_at ON threads(archived_at DESC);
368 CREATE INDEX IF NOT EXISTS idx_threads_archived_updated ON threads(archived, updated_at DESC);
369
370 CREATE TABLE IF NOT EXISTS thread_dynamic_tools (
371 thread_id TEXT NOT NULL,
372 position INTEGER NOT NULL,
373 name TEXT NOT NULL,
374 description TEXT,
375 input_schema TEXT NOT NULL,
376 PRIMARY KEY (thread_id, position),
377 FOREIGN KEY(thread_id) REFERENCES threads(id) ON DELETE CASCADE
378 );
379
380 CREATE TABLE IF NOT EXISTS messages (
381 id INTEGER PRIMARY KEY AUTOINCREMENT,
382 thread_id TEXT NOT NULL,
383 role TEXT NOT NULL,
384 content TEXT NOT NULL,
385 item_json TEXT,
386 created_at INTEGER NOT NULL,
387 FOREIGN KEY(thread_id) REFERENCES threads(id) ON DELETE CASCADE
388 );
389 CREATE INDEX IF NOT EXISTS idx_messages_thread_created_at ON messages(thread_id, created_at ASC);
390
391 CREATE TABLE IF NOT EXISTS checkpoints (
392 thread_id TEXT NOT NULL,
393 checkpoint_id TEXT NOT NULL,
394 state_json TEXT NOT NULL,
395 created_at INTEGER NOT NULL,
396 PRIMARY KEY(thread_id, checkpoint_id),
397 FOREIGN KEY(thread_id) REFERENCES threads(id) ON DELETE CASCADE
398 );
399 CREATE INDEX IF NOT EXISTS idx_checkpoints_thread_created_at ON checkpoints(thread_id, created_at DESC);
400
401 CREATE TABLE IF NOT EXISTS jobs (
402 id TEXT PRIMARY KEY,
403 name TEXT NOT NULL,
404 status TEXT NOT NULL,
405 progress INTEGER,
406 detail TEXT,
407 created_at INTEGER NOT NULL,
408 updated_at INTEGER NOT NULL
409 );
410 CREATE INDEX IF NOT EXISTS idx_jobs_updated_at ON jobs(updated_at DESC);
411
412 -- Add parent_entry_id column, and set to last message before current message
413 ALTER TABLE messages ADD COLUMN parent_entry_id INTEGER NULL;
414 UPDATE messages
415 SET parent_entry_id = (
416 SELECT m2.id
417 FROM messages m2
418 WHERE m2.thread_id = messages.thread_id
419 AND (
420 m2.created_at < messages.created_at
421 OR (
422 m2.created_at = messages.created_at
423 AND m2.id < messages.id
424 )
425 )
426 ORDER BY m2.created_at DESC, m2.id DESC
427 LIMIT 1
428 );
429 CREATE INDEX idx_messages_parent_entry_id ON messages(parent_entry_id);
430
431 -- Add current_leaf_id column, and set to last message in thread
432 ALTER TABLE threads ADD COLUMN current_leaf_id INTEGER NULL;
433 UPDATE threads
434 SET current_leaf_id = (
435 SELECT m.id
436 FROM messages m
437 WHERE m.thread_id = threads.id
438 ORDER BY m.id DESC
439 LIMIT 1
440 );
441
442 PRAGMA user_version = 1;
443 COMMIT;
444 "#,
445 )
446 .context("failed to initialize thread schema")?;
447 user_version = 1;
448 }
449 if user_version < 2 {
450 conn.execute_batch(
451 r#"
452 BEGIN;
453 CREATE TABLE IF NOT EXISTS workflow_runs (
454 id TEXT PRIMARY KEY,
455 workflow_id TEXT NOT NULL,
456 goal TEXT NOT NULL,
457 status TEXT NOT NULL,
458 input_hash TEXT,
459 started_at INTEGER NOT NULL,
460 completed_at INTEGER,
461 metadata_json TEXT NOT NULL DEFAULT '{}'
462 );
463 CREATE INDEX IF NOT EXISTS idx_workflow_runs_status_started_at
464 ON workflow_runs(status, started_at DESC);
465 CREATE INDEX IF NOT EXISTS idx_workflow_runs_workflow_started_at
466 ON workflow_runs(workflow_id, started_at DESC);
467
468 CREATE TABLE IF NOT EXISTS branch_runs (
469 id TEXT PRIMARY KEY,
470 workflow_run_id TEXT NOT NULL,
471 branch_id TEXT NOT NULL,
472 node_id TEXT NOT NULL,
473 status TEXT NOT NULL,
474 started_at INTEGER NOT NULL,
475 completed_at INTEGER,
476 result_json TEXT NOT NULL DEFAULT '{}',
477 FOREIGN KEY(workflow_run_id) REFERENCES workflow_runs(id) ON DELETE CASCADE
478 );
479 CREATE INDEX IF NOT EXISTS idx_branch_runs_workflow_run_id
480 ON branch_runs(workflow_run_id);
481 CREATE INDEX IF NOT EXISTS idx_branch_runs_branch_id
482 ON branch_runs(branch_id);
483
484 CREATE TABLE IF NOT EXISTS leaf_runs (
485 id TEXT PRIMARY KEY,
486 workflow_run_id TEXT NOT NULL,
487 branch_run_id TEXT,
488 leaf_id TEXT NOT NULL,
489 task_id TEXT NOT NULL,
490 input_hash TEXT,
491 status TEXT NOT NULL,
492 output_json TEXT NOT NULL DEFAULT '{}',
493 artifacts_json TEXT NOT NULL DEFAULT '[]',
494 started_at INTEGER NOT NULL,
495 completed_at INTEGER,
496 FOREIGN KEY(workflow_run_id) REFERENCES workflow_runs(id) ON DELETE CASCADE,
497 FOREIGN KEY(branch_run_id) REFERENCES branch_runs(id) ON DELETE SET NULL
498 );
499 CREATE INDEX IF NOT EXISTS idx_leaf_runs_workflow_run_id
500 ON leaf_runs(workflow_run_id);
501 CREATE INDEX IF NOT EXISTS idx_leaf_runs_replay_lookup
502 ON leaf_runs(workflow_run_id, leaf_id, input_hash);
503
504 CREATE TABLE IF NOT EXISTS control_node_runs (
505 id TEXT PRIMARY KEY,
506 workflow_run_id TEXT NOT NULL,
507 node_id TEXT NOT NULL,
508 kind TEXT NOT NULL,
509 status TEXT NOT NULL,
510 selected_children_json TEXT NOT NULL DEFAULT '[]',
511 result_json TEXT NOT NULL DEFAULT '{}',
512 started_at INTEGER NOT NULL,
513 completed_at INTEGER,
514 FOREIGN KEY(workflow_run_id) REFERENCES workflow_runs(id) ON DELETE CASCADE
515 );
516 CREATE INDEX IF NOT EXISTS idx_control_node_runs_workflow_run_id
517 ON control_node_runs(workflow_run_id);
518 CREATE INDEX IF NOT EXISTS idx_control_node_runs_node_id
519 ON control_node_runs(node_id);
520
521 CREATE TABLE IF NOT EXISTS teacher_candidates (
522 id TEXT PRIMARY KEY,
523 workflow_run_id TEXT NOT NULL,
524 control_node_run_id TEXT NOT NULL,
525 candidate_id TEXT NOT NULL,
526 branch_run_id TEXT,
527 score REAL,
528 passed INTEGER,
529 rationale_json TEXT NOT NULL DEFAULT '{}',
530 created_at INTEGER NOT NULL,
531 FOREIGN KEY(workflow_run_id) REFERENCES workflow_runs(id) ON DELETE CASCADE,
532 FOREIGN KEY(control_node_run_id) REFERENCES control_node_runs(id) ON DELETE CASCADE,
533 FOREIGN KEY(branch_run_id) REFERENCES branch_runs(id) ON DELETE SET NULL
534 );
535 CREATE INDEX IF NOT EXISTS idx_teacher_candidates_workflow_run_id
536 ON teacher_candidates(workflow_run_id);
537 CREATE INDEX IF NOT EXISTS idx_teacher_candidates_control_node_run_id
538 ON teacher_candidates(control_node_run_id);
539
540 PRAGMA user_version = 2;
541 COMMIT;
542 "#,
543 )
544 .context("failed to initialize workflow trace schema")?;
545 user_version = 2;
546 }
547 if user_version < 3 {
548 conn.execute_batch(
549 r#"
550 BEGIN;
551 CREATE TABLE IF NOT EXISTS thread_goals (
552 thread_id TEXT PRIMARY KEY NOT NULL,
553 goal_id TEXT NOT NULL,
554 objective TEXT NOT NULL,
555 status TEXT NOT NULL CHECK(status IN (
556 'active',
557 'paused',
558 'blocked',
559 'usage_limited',
560 'budget_limited',
561 'complete'
562 )),
563 token_budget INTEGER,
564 tokens_used INTEGER NOT NULL DEFAULT 0,
565 time_used_seconds INTEGER NOT NULL DEFAULT 0,
566 created_at INTEGER NOT NULL,
567 updated_at INTEGER NOT NULL,
568 FOREIGN KEY(thread_id) REFERENCES threads(id) ON DELETE CASCADE
569 );
570
571 PRAGMA user_version = 3;
572 COMMIT;
573 "#,
574 )
575 .context("failed to initialize thread goal schema")?;
576 user_version = 3;
577 }
578 if user_version < 4 {
579 conn.execute_batch(
580 r#"
581 BEGIN;
582 ALTER TABLE thread_goals
583 ADD COLUMN continuation_count INTEGER NOT NULL DEFAULT 0;
584
585 PRAGMA user_version = 4;
586 COMMIT;
587 "#,
588 )
589 .context("failed to initialize thread goal continuation schema")?;
590 }
591 Ok(())
592 }
593
594 pub fn upsert_thread(&self, thread: &ThreadMetadata) -> Result<()> {
599 let conn = self.conn()?;
600 conn.execute(
601 r#"
602 INSERT INTO threads (
603 id, rollout_path, preview, ephemeral, model_provider, created_at, updated_at, status, path, cwd,
604 cli_version, source, title, sandbox_policy, approval_mode, archived, archived_at,
605 git_sha, git_branch, git_origin_url, memory_mode
606 ) VALUES (
607 ?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10,
608 ?11, ?12, ?13, ?14, ?15, ?16, ?17,
609 ?18, ?19, ?20, ?21
610 )
611 ON CONFLICT(id) DO UPDATE SET
612 rollout_path=excluded.rollout_path,
613 preview=excluded.preview,
614 ephemeral=excluded.ephemeral,
615 model_provider=excluded.model_provider,
616 created_at=excluded.created_at,
617 updated_at=excluded.updated_at,
618 status=excluded.status,
619 path=excluded.path,
620 cwd=excluded.cwd,
621 cli_version=excluded.cli_version,
622 source=excluded.source,
623 title=excluded.title,
624 sandbox_policy=excluded.sandbox_policy,
625 approval_mode=excluded.approval_mode,
626 archived=excluded.archived,
627 archived_at=excluded.archived_at,
628 git_sha=excluded.git_sha,
629 git_branch=excluded.git_branch,
630 git_origin_url=excluded.git_origin_url,
631 memory_mode=excluded.memory_mode
632 "#,
633 params![
634 thread.id,
635 path_to_opt_string(thread.rollout_path.as_deref()),
636 thread.preview,
637 bool_to_i64(thread.ephemeral),
638 thread.model_provider,
639 thread.created_at,
640 thread.updated_at,
641 thread_status_to_str(&thread.status),
642 path_to_opt_string(thread.path.as_deref()),
643 thread.cwd.display().to_string(),
644 thread.cli_version,
645 session_source_to_str(&thread.source),
646 thread.name,
647 thread.sandbox_policy,
648 thread.approval_mode,
649 bool_to_i64(thread.archived),
650 thread.archived_at,
651 thread.git_sha,
652 thread.git_branch,
653 thread.git_origin_url,
654 thread.memory_mode,
655 ],
656 )
657 .context("failed to upsert thread metadata")?;
658
659 self.append_thread_name(
660 &thread.id,
661 thread.name.clone(),
662 thread.updated_at,
663 thread.rollout_path.clone(),
664 )?;
665 Ok(())
666 }
667
668 pub fn get_thread(&self, id: &str) -> Result<Option<ThreadMetadata>> {
672 let conn = self.conn()?;
673 conn.query_row(
674 r#"
675 SELECT id, rollout_path, preview, ephemeral, model_provider, created_at, updated_at, status, path, cwd,
676 cli_version, source, title, sandbox_policy, approval_mode, archived, archived_at,
677 git_sha, git_branch, git_origin_url, memory_mode, current_leaf_id
678 FROM threads
679 WHERE id = ?1
680 "#,
681 params![id],
682 row_to_thread,
683 )
684 .optional()
685 .context("failed to read thread")
686 }
687
688 pub fn list_threads(&self, filters: ThreadListFilters) -> Result<Vec<ThreadMetadata>> {
693 let conn = self.conn()?;
694 let sql = if filters.include_archived {
695 "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"
696 } else {
697 "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"
698 };
699
700 let mut stmt = conn.prepare(sql).context("failed to prepare list query")?;
701 let limit = i64::try_from(filters.limit.unwrap_or(50)).unwrap_or(50);
702 let mut rows = stmt
703 .query(params![limit])
704 .context("failed to query threads")?;
705 let mut out = Vec::new();
706 while let Some(row) = rows.next().context("failed to iterate thread rows")? {
707 out.push(row_to_thread(row)?);
708 }
709 Ok(out)
710 }
711
712 pub fn mark_archived(&self, id: &str) -> Result<()> {
715 let conn = self.conn()?;
716 conn.execute(
717 "UPDATE threads SET archived = 1, archived_at = ?2, status = ?3 WHERE id = ?1",
718 params![
719 id,
720 Utc::now().timestamp(),
721 thread_status_to_str(&ThreadStatus::Archived)
722 ],
723 )
724 .context("failed to archive thread")?;
725 Ok(())
726 }
727
728 pub fn mark_unarchived(&self, id: &str) -> Result<()> {
730 let conn = self.conn()?;
731 conn.execute(
732 "UPDATE threads SET archived = 0, archived_at = NULL, status = CASE WHEN status = ?2 THEN ?3 ELSE status END WHERE id = ?1",
733 params![
734 id,
735 thread_status_to_str(&ThreadStatus::Archived),
736 thread_status_to_str(&ThreadStatus::Idle),
737 ],
738 )
739 .context("failed to unarchive thread")?;
740 Ok(())
741 }
742
743 pub fn delete_thread(&self, id: &str) -> Result<()> {
746 let conn = self.conn()?;
747 conn.execute("DELETE FROM threads WHERE id = ?1", params![id])
748 .context("failed to delete thread")?;
749 Ok(())
750 }
751
752 pub fn set_thread_memory_mode(&self, id: &str, mode: Option<&str>) -> Result<()> {
756 let conn = self.conn()?;
757 conn.execute(
758 "UPDATE threads SET memory_mode = ?2 WHERE id = ?1",
759 params![id, mode],
760 )
761 .context("failed to update thread memory mode")?;
762 Ok(())
763 }
764
765 pub fn get_thread_memory_mode(&self, id: &str) -> Result<Option<String>> {
769 let conn = self.conn()?;
770 conn.query_row(
771 "SELECT memory_mode FROM threads WHERE id = ?1",
772 params![id],
773 |row| row.get::<_, Option<String>>(0),
774 )
775 .optional()
776 .context("failed to read thread memory mode")
777 .map(Option::flatten)
778 }
779
780 pub fn upsert_thread_goal(&self, goal: &ThreadGoalRecord) -> Result<()> {
782 let conn = self.conn()?;
783 let exists: Option<i64> = conn
784 .query_row(
785 "SELECT 1 FROM threads WHERE id = ?1",
786 params![goal.thread_id],
787 |row| row.get(0),
788 )
789 .optional()
790 .context("failed to verify thread before saving goal")?;
791 if exists.is_none() {
792 anyhow::bail!("thread {} not found", goal.thread_id);
793 }
794
795 conn.execute(
796 r#"
797 INSERT INTO thread_goals (
798 thread_id, goal_id, objective, status, token_budget, tokens_used,
799 time_used_seconds, continuation_count, created_at, updated_at
800 ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10)
801 ON CONFLICT(thread_id) DO UPDATE SET
802 goal_id=excluded.goal_id,
803 objective=excluded.objective,
804 status=excluded.status,
805 token_budget=excluded.token_budget,
806 tokens_used=excluded.tokens_used,
807 time_used_seconds=excluded.time_used_seconds,
808 continuation_count=excluded.continuation_count,
809 created_at=excluded.created_at,
810 updated_at=excluded.updated_at
811 "#,
812 params![
813 goal.thread_id,
814 goal.goal_id,
815 goal.objective,
816 thread_goal_status_to_str(&goal.status),
817 goal.token_budget,
818 goal.tokens_used,
819 goal.time_used_seconds,
820 goal.continuation_count,
821 goal.created_at,
822 goal.updated_at,
823 ],
824 )
825 .context("failed to upsert thread goal")?;
826 Ok(())
827 }
828
829 pub fn record_thread_goal_usage(
845 &self,
846 thread_id: &str,
847 token_delta: i64,
848 time_delta_seconds: i64,
849 now: i64,
850 ) -> Result<Option<ThreadGoalRecord>> {
851 let conn = self.conn()?;
852 let changed = conn
853 .execute(
854 r#"
855 UPDATE thread_goals
856 SET tokens_used = tokens_used + ?2,
857 time_used_seconds = time_used_seconds + ?3,
858 updated_at = MAX(updated_at, ?4)
859 WHERE thread_id = ?1
860 "#,
861 params![thread_id, token_delta, time_delta_seconds, now],
862 )
863 .context("failed to record thread goal usage")?;
864 if changed == 0 {
865 return Ok(None);
866 }
867 Self::read_thread_goal(&conn, thread_id)
868 }
869
870 pub fn record_thread_goal_continuation(
876 &self,
877 thread_id: &str,
878 now: i64,
879 ) -> Result<Option<ThreadGoalRecord>> {
880 let conn = self.conn()?;
881 let changed = conn
882 .execute(
883 r#"
884 UPDATE thread_goals
885 SET continuation_count = continuation_count + 1,
886 updated_at = MAX(updated_at, ?2)
887 WHERE thread_id = ?1
888 "#,
889 params![thread_id, now],
890 )
891 .context("failed to record thread goal continuation")?;
892 if changed == 0 {
893 return Ok(None);
894 }
895 Self::read_thread_goal(&conn, thread_id)
896 }
897
898 pub fn get_thread_goal(&self, thread_id: &str) -> Result<Option<ThreadGoalRecord>> {
900 let conn = self.conn()?;
901 Self::read_thread_goal(&conn, thread_id)
902 }
903
904 fn read_thread_goal(conn: &Connection, thread_id: &str) -> Result<Option<ThreadGoalRecord>> {
908 conn.query_row(
909 r#"
910 SELECT thread_id, goal_id, objective, status, token_budget, tokens_used,
911 time_used_seconds, continuation_count, created_at, updated_at
912 FROM thread_goals
913 WHERE thread_id = ?1
914 "#,
915 params![thread_id],
916 row_to_thread_goal,
917 )
918 .optional()
919 .context("failed to read thread goal")
920 }
921
922 pub fn delete_thread_goal(&self, thread_id: &str) -> Result<bool> {
924 let conn = self.conn()?;
925 let changed = conn
926 .execute(
927 "DELETE FROM thread_goals WHERE thread_id = ?1",
928 params![thread_id],
929 )
930 .context("failed to delete thread goal")?;
931 Ok(changed > 0)
932 }
933
934 pub fn list_leaf_messages(&self, thread_id: &str) -> Result<Vec<MessageRecord>> {
939 let conn = self.conn()?;
940 let mut stmt = conn
941 .prepare(
942 r#"
943 SELECT m1.id, m1.thread_id, m1.role, m1.content, m1.item_json, m1.created_at, m1.parent_entry_id
944 FROM messages m1
945 LEFT JOIN messages m2 ON m1.id = m2.parent_entry_id
946 WHERE m1.thread_id = ?1 AND m2.id IS NULL
947 "#,
948 )
949 .context("failed to prepare message listing query")?;
950 let mut rows = stmt
951 .query(params![thread_id])
952 .with_context(|| format!("failed to list leaf messages for thread {thread_id}"))?;
953 let mut out = Vec::new();
954 while let Some(row) = rows.next().context("failed to iterate message rows")? {
955 let item_json: Option<String> = row.get(4).context("failed to read item json")?;
956 let item = item_json
957 .as_deref()
958 .map(serde_json::from_str)
959 .transpose()
960 .with_context(|| {
961 format!("failed to parse message item json in thread {thread_id}")
962 })?;
963 out.push(MessageRecord {
964 id: row.get(0).context("failed to read message id")?,
965 thread_id: row.get(1).context("failed to read message thread id")?,
966 role: row.get(2).context("failed to read message role")?,
967 content: row.get(3).context("failed to read message content")?,
968 item,
969 created_at: row.get(5).context("failed to read message timestamp")?,
970 parent_entry_id: row.get(6).context("failed to read parent entry id")?,
971 });
972 }
973 Ok(out)
974 }
975
976 pub fn set_current_leaf_id(&self, thread_id: &str, current_leaf_id: &str) -> Result<()> {
981 let conn = self.conn()?;
982 conn.execute(
983 "UPDATE threads SET current_leaf_id = ?1 WHERE id = ?2",
984 params![current_leaf_id, thread_id],
985 )
986 .context("failed to update thread current leaf id")?;
987 Ok(())
988 }
989
990 pub fn persist_dynamic_tools(
995 &self,
996 thread_id: &str,
997 tools: &[DynamicToolRecord],
998 ) -> Result<()> {
999 let mut conn = self.conn()?;
1000 let tx = conn
1001 .transaction()
1002 .context("failed to begin dynamic tools transaction")?;
1003 tx.execute(
1004 "DELETE FROM thread_dynamic_tools WHERE thread_id = ?1",
1005 params![thread_id],
1006 )
1007 .context("failed to clear dynamic tools")?;
1008 for tool in tools {
1009 tx.execute(
1010 "INSERT INTO thread_dynamic_tools(thread_id, position, name, description, input_schema) VALUES (?1, ?2, ?3, ?4, ?5)",
1011 params![
1012 thread_id,
1013 tool.position,
1014 tool.name,
1015 tool.description,
1016 tool.input_schema.to_string()
1017 ],
1018 )
1019 .with_context(|| format!("failed to persist dynamic tool {}", tool.name))?;
1020 }
1021 tx.commit().context("failed to commit dynamic tools")?;
1022 Ok(())
1023 }
1024
1025 pub fn get_dynamic_tools(&self, thread_id: &str) -> Result<Vec<DynamicToolRecord>> {
1027 let conn = self.conn()?;
1028 let mut stmt = conn
1029 .prepare(
1030 "SELECT position, name, description, input_schema FROM thread_dynamic_tools WHERE thread_id = ?1 ORDER BY position ASC",
1031 )
1032 .context("failed to prepare get dynamic tools query")?;
1033 let mut rows = stmt
1034 .query(params![thread_id])
1035 .context("failed to query dynamic tools")?;
1036 let mut out = Vec::new();
1037 while let Some(row) = rows.next().context("failed to iterate dynamic tools")? {
1038 let input_schema_raw: String =
1039 row.get(3).context("failed to read tool input schema")?;
1040 let input_schema: Value =
1041 serde_json::from_str(&input_schema_raw).with_context(|| {
1042 format!("failed to parse input schema for dynamic tool in thread {thread_id}")
1043 })?;
1044 out.push(DynamicToolRecord {
1045 position: row.get(0).context("failed to read tool position")?,
1046 name: row.get(1).context("failed to read tool name")?,
1047 description: row.get(2).context("failed to read tool description")?,
1048 input_schema,
1049 });
1050 }
1051 Ok(out)
1052 }
1053
1054 pub fn append_message(
1060 &self,
1061 thread_id: &str,
1062 role: &str,
1063 content: &str,
1064 item: Option<Value>,
1065 ) -> Result<i64> {
1066 let mut conn = self.conn()?;
1067 let created_at = Utc::now().timestamp();
1068 let item_json = item
1069 .as_ref()
1070 .map(serde_json::to_string)
1071 .transpose()
1072 .context("failed to serialize message item payload")?;
1073
1074 let tx = conn
1075 .transaction()
1076 .context("failed to begin append message transaction")?;
1077
1078 let current_leaf_id: Option<i64> = tx
1079 .query_row(
1080 "SELECT current_leaf_id FROM threads WHERE id = ?1",
1081 params![thread_id],
1082 |row| row.get(0),
1083 )
1084 .with_context(|| {
1085 format!("failed to query thread current leaf id for thread {thread_id}")
1086 })?;
1087
1088 let next_leaf_id: i64 = tx.query_row(
1089 r#"
1090 INSERT INTO messages(thread_id, role, content, item_json, created_at, parent_entry_id)
1091 SELECT ?1, ?2, ?3, ?4, ?5, ?6
1092 RETURNING id
1093 "#, params![thread_id, role, content, item_json, created_at, current_leaf_id], |row| row.get(0)
1094 ).with_context(|| format!("failed to append message for thread {thread_id}"))?;
1095
1096 tx.execute(
1097 r#"
1098 UPDATE threads
1099 SET current_leaf_id = ?1
1100 WHERE id = ?2;
1101 "#,
1102 params![next_leaf_id, thread_id],
1103 )
1104 .with_context(|| {
1105 format!("failed to update thread current leaf id for thread {thread_id}")
1106 })?;
1107
1108 tx.commit()
1109 .context("failed to commit append message transaction")?;
1110
1111 Ok(next_leaf_id)
1112 }
1113
1114 pub fn list_messages(
1120 &self,
1121 thread_id: &str,
1122 limit: Option<usize>,
1123 ) -> Result<Vec<MessageRecord>> {
1124 let conn = self.conn()?;
1125 let limit = i64::try_from(limit.unwrap_or(500)).unwrap_or(500);
1126 let mut stmt = conn
1127 .prepare(
1128 r#"
1129 WITH RECURSIVE
1130 leaf_id AS (
1131 SELECT current_leaf_id FROM threads WHERE id = ?1
1132 ),
1133 ancestors AS (
1134 SELECT id, thread_id, role, content, item_json, created_at, parent_entry_id, 0 AS depth
1135 FROM messages
1136 WHERE id = (SELECT current_leaf_id FROM leaf_id)
1137
1138 UNION ALL
1139
1140 SELECT m.id, m.thread_id, m.role, m.content, m.item_json, m.created_at, m.parent_entry_id, a.depth + 1
1141 FROM messages m
1142 JOIN ancestors a ON m.id = a.parent_entry_id
1143 WHERE a.depth < ?2
1144 )
1145 SELECT id, thread_id, role, content, item_json, created_at, parent_entry_id FROM ancestors
1146 ORDER BY depth DESC
1147 "#
1148 )
1149 .context("failed to prepare message listing query")?;
1150 let mut rows = stmt
1151 .query(params![thread_id, limit - 1])
1152 .with_context(|| format!("failed to list messages for thread {thread_id}"))?;
1153 let mut out = Vec::new();
1154 while let Some(row) = rows.next().context("failed to iterate message rows")? {
1155 let item_json: Option<String> = row.get(4).context("failed to read item json")?;
1156 let item = item_json
1157 .as_deref()
1158 .map(serde_json::from_str)
1159 .transpose()
1160 .with_context(|| {
1161 format!("failed to parse message item json in thread {thread_id}")
1162 })?;
1163 out.push(MessageRecord {
1164 id: row.get(0).context("failed to read message id")?,
1165 thread_id: row.get(1).context("failed to read message thread id")?,
1166 role: row.get(2).context("failed to read message role")?,
1167 content: row.get(3).context("failed to read message content")?,
1168 item,
1169 created_at: row.get(5).context("failed to read message timestamp")?,
1170 parent_entry_id: row.get(6).context("failed to read parent entry id")?,
1171 });
1172 }
1173 Ok(out)
1174 }
1175
1176 pub fn fork_at_message(
1182 &self,
1183 message_id: &str,
1184 role: &str,
1185 content: &str,
1186 item: Option<Value>,
1187 ) -> Result<i64> {
1188 let mut conn = self.conn()?;
1189 let created_at = Utc::now().timestamp();
1190 let item_json = item
1191 .as_ref()
1192 .map(serde_json::to_string)
1193 .transpose()
1194 .context("failed to serialize message item payload")?;
1195
1196 let tx = conn
1197 .transaction()
1198 .context("failed to begin fork message transaction")?;
1199
1200 let thread_id: String = tx
1201 .query_row(
1202 "SELECT thread_id FROM messages WHERE id = ?1",
1203 params![message_id],
1204 |row| row.get(0),
1205 )
1206 .with_context(|| format!("failed to query thread id for message {message_id}"))?;
1207
1208 let next_leaf_id: i64 = tx.query_row(
1209 r#"
1210 INSERT INTO messages(thread_id, role, content, item_json, created_at, parent_entry_id)
1211 SELECT ?1, ?2, ?3, ?4, ?5, ?6
1212 RETURNING id
1213 "#, params![thread_id, role, content, item_json, created_at, message_id], |row| row.get(0)
1214 ).with_context(|| format!("failed to fork at message for thread {thread_id:?}"))?;
1215
1216 tx.execute(
1217 r#"
1218 UPDATE threads
1219 SET current_leaf_id = ?1
1220 WHERE id = ?2;
1221 "#,
1222 params![next_leaf_id, thread_id],
1223 )
1224 .with_context(|| {
1225 format!("failed to update thread current leaf id for thread {thread_id:?}")
1226 })?;
1227
1228 tx.commit()
1229 .context("failed to commit fork message transaction")?;
1230
1231 Ok(next_leaf_id)
1232 }
1233
1234 pub fn clear_messages(&self, thread_id: &str) -> Result<usize> {
1238 let mut conn = self.conn()?;
1239 let tx = conn
1240 .transaction()
1241 .context("failed to begin clear messages transaction")?;
1242
1243 tx.execute(
1244 r#"
1245 UPDATE threads
1246 SET current_leaf_id = NULL
1247 WHERE id = ?1;
1248 "#,
1249 params![thread_id],
1250 )
1251 .with_context(|| format!("failed to clear messages for thread {thread_id}"))?;
1252 let result = tx
1253 .execute(
1254 r#"
1255 DELETE FROM messages WHERE thread_id = ?1
1256 "#,
1257 params![thread_id],
1258 )
1259 .with_context(|| format!("failed to clear messages for thread {thread_id}"))?;
1260 tx.commit()
1261 .context("failed to commit clear messages transaction")?;
1262
1263 Ok(result)
1264 }
1265
1266 pub fn save_checkpoint(
1271 &self,
1272 thread_id: &str,
1273 checkpoint_id: &str,
1274 state: &Value,
1275 ) -> Result<()> {
1276 let conn = self.conn()?;
1277 let state_json =
1278 serde_json::to_string(state).context("failed to encode checkpoint state")?;
1279 conn.execute(
1280 r#"
1281 INSERT INTO checkpoints(thread_id, checkpoint_id, state_json, created_at)
1282 VALUES (?1, ?2, ?3, ?4)
1283 ON CONFLICT(thread_id, checkpoint_id) DO UPDATE SET
1284 state_json = excluded.state_json,
1285 created_at = excluded.created_at
1286 "#,
1287 params![thread_id, checkpoint_id, state_json, Utc::now().timestamp()],
1288 )
1289 .with_context(|| {
1290 format!("failed to save checkpoint {checkpoint_id} for thread {thread_id}")
1291 })?;
1292 Ok(())
1293 }
1294
1295 pub fn load_checkpoint(
1301 &self,
1302 thread_id: &str,
1303 checkpoint_id: Option<&str>,
1304 ) -> Result<Option<CheckpointRecord>> {
1305 let conn = self.conn()?;
1306 if let Some(checkpoint_id) = checkpoint_id {
1307 let row = conn
1308 .query_row(
1309 "SELECT thread_id, checkpoint_id, state_json, created_at FROM checkpoints WHERE thread_id = ?1 AND checkpoint_id = ?2",
1310 params![thread_id, checkpoint_id],
1311 |row| {
1312 Ok((
1313 row.get::<_, String>(0)?,
1314 row.get::<_, String>(1)?,
1315 row.get::<_, String>(2)?,
1316 row.get::<_, i64>(3)?,
1317 ))
1318 },
1319 )
1320 .optional()
1321 .with_context(|| {
1322 format!("failed to load checkpoint {checkpoint_id} for thread {thread_id}")
1323 })?;
1324 if let Some((thread_id, checkpoint_id, state_json, created_at)) = row {
1325 let state = parse_checkpoint_state(&state_json)?;
1326 return Ok(Some(CheckpointRecord {
1327 thread_id,
1328 checkpoint_id,
1329 state,
1330 created_at,
1331 }));
1332 }
1333 return Ok(None);
1334 }
1335
1336 let row = conn
1337 .query_row(
1338 "SELECT thread_id, checkpoint_id, state_json, created_at FROM checkpoints WHERE thread_id = ?1 ORDER BY created_at DESC LIMIT 1",
1339 params![thread_id],
1340 |row| {
1341 Ok((
1342 row.get::<_, String>(0)?,
1343 row.get::<_, String>(1)?,
1344 row.get::<_, String>(2)?,
1345 row.get::<_, i64>(3)?,
1346 ))
1347 },
1348 )
1349 .optional()
1350 .with_context(|| format!("failed to load latest checkpoint for thread {thread_id}"))?;
1351 if let Some((thread_id, checkpoint_id, state_json, created_at)) = row {
1352 let state = parse_checkpoint_state(&state_json)?;
1353 return Ok(Some(CheckpointRecord {
1354 thread_id,
1355 checkpoint_id,
1356 state,
1357 created_at,
1358 }));
1359 }
1360 Ok(None)
1361 }
1362
1363 pub fn list_checkpoints(
1367 &self,
1368 thread_id: &str,
1369 limit: Option<usize>,
1370 ) -> Result<Vec<CheckpointRecord>> {
1371 let conn = self.conn()?;
1372 let limit = i64::try_from(limit.unwrap_or(100)).unwrap_or(100);
1373 let mut stmt = conn
1374 .prepare(
1375 "SELECT thread_id, checkpoint_id, state_json, created_at FROM checkpoints WHERE thread_id = ?1 ORDER BY created_at DESC LIMIT ?2",
1376 )
1377 .context("failed to prepare checkpoint list query")?;
1378 let mut rows = stmt
1379 .query(params![thread_id, limit])
1380 .with_context(|| format!("failed to list checkpoints for thread {thread_id}"))?;
1381
1382 let mut out = Vec::new();
1383 while let Some(row) = rows.next().context("failed to iterate checkpoint rows")? {
1384 let state_json: String = row.get(2).context("failed to read checkpoint state json")?;
1385 let state = parse_checkpoint_state(&state_json)?;
1386 out.push(CheckpointRecord {
1387 thread_id: row.get(0).context("failed to read checkpoint thread id")?,
1388 checkpoint_id: row.get(1).context("failed to read checkpoint id")?,
1389 state,
1390 created_at: row.get(3).context("failed to read checkpoint timestamp")?,
1391 });
1392 }
1393 Ok(out)
1394 }
1395
1396 pub fn delete_checkpoint(&self, thread_id: &str, checkpoint_id: &str) -> Result<()> {
1398 let conn = self.conn()?;
1399 conn.execute(
1400 "DELETE FROM checkpoints WHERE thread_id = ?1 AND checkpoint_id = ?2",
1401 params![thread_id, checkpoint_id],
1402 )
1403 .with_context(|| {
1404 format!("failed to delete checkpoint {checkpoint_id} for thread {thread_id}")
1405 })?;
1406 Ok(())
1407 }
1408
1409 pub fn upsert_job(&self, job: &JobStateRecord) -> Result<()> {
1411 let conn = self.conn()?;
1412 conn.execute(
1413 r#"
1414 INSERT INTO jobs(id, name, status, progress, detail, created_at, updated_at)
1415 VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7)
1416 ON CONFLICT(id) DO UPDATE SET
1417 name = excluded.name,
1418 status = excluded.status,
1419 progress = excluded.progress,
1420 detail = excluded.detail,
1421 created_at = excluded.created_at,
1422 updated_at = excluded.updated_at
1423 "#,
1424 params![
1425 job.id,
1426 job.name,
1427 job_state_status_to_str(&job.status),
1428 job.progress.map(i64::from),
1429 job.detail,
1430 job.created_at,
1431 job.updated_at
1432 ],
1433 )
1434 .with_context(|| format!("failed to upsert job {}", job.id))?;
1435 Ok(())
1436 }
1437
1438 pub fn get_job(&self, id: &str) -> Result<Option<JobStateRecord>> {
1442 let conn = self.conn()?;
1443 conn.query_row(
1444 "SELECT id, name, status, progress, detail, created_at, updated_at FROM jobs WHERE id = ?1",
1445 params![id],
1446 |row| {
1447 let status_raw: String = row.get(2)?;
1448 let progress: Option<i64> = row.get(3)?;
1449 Ok(JobStateRecord {
1450 id: row.get(0)?,
1451 name: row.get(1)?,
1452 status: job_state_status_from_str(&status_raw),
1453 progress: progress.and_then(|v| u8::try_from(v).ok()),
1454 detail: row.get(4)?,
1455 created_at: row.get(5)?,
1456 updated_at: row.get(6)?,
1457 })
1458 },
1459 )
1460 .optional()
1461 .with_context(|| format!("failed to read job {id}"))
1462 }
1463
1464 pub fn list_jobs(&self, limit: Option<usize>) -> Result<Vec<JobStateRecord>> {
1468 let conn = self.conn()?;
1469 let limit = i64::try_from(limit.unwrap_or(100)).unwrap_or(100);
1470 let mut stmt = conn
1471 .prepare(
1472 "SELECT id, name, status, progress, detail, created_at, updated_at FROM jobs ORDER BY updated_at DESC LIMIT ?1",
1473 )
1474 .context("failed to prepare job list query")?;
1475 let mut rows = stmt
1476 .query(params![limit])
1477 .context("failed to query persisted jobs")?;
1478 let mut out = Vec::new();
1479 while let Some(row) = rows.next().context("failed to iterate persisted jobs")? {
1480 let status_raw: String = row.get(2).context("failed to read job status")?;
1481 let progress: Option<i64> = row.get(3).context("failed to read job progress")?;
1482 out.push(JobStateRecord {
1483 id: row.get(0).context("failed to read job id")?,
1484 name: row.get(1).context("failed to read job name")?,
1485 status: job_state_status_from_str(&status_raw),
1486 progress: progress.and_then(|v| u8::try_from(v).ok()),
1487 detail: row.get(4).context("failed to read job detail")?,
1488 created_at: row.get(5).context("failed to read job created_at")?,
1489 updated_at: row.get(6).context("failed to read job updated_at")?,
1490 });
1491 }
1492 Ok(out)
1493 }
1494
1495 pub fn delete_job(&self, id: &str) -> Result<()> {
1497 let conn = self.conn()?;
1498 conn.execute("DELETE FROM jobs WHERE id = ?1", params![id])
1499 .with_context(|| format!("failed to delete job {id}"))?;
1500 Ok(())
1501 }
1502
1503 pub fn find_rollout_path_by_id(&self, id: &str) -> Result<Option<PathBuf>> {
1505 let conn = self.conn()?;
1506 conn.query_row(
1507 "SELECT rollout_path FROM threads WHERE id = ?1",
1508 params![id],
1509 |row| row.get::<_, Option<String>>(0),
1510 )
1511 .optional()
1512 .context("failed to lookup rollout path")
1513 .map(|opt| opt.flatten().map(PathBuf::from))
1514 }
1515
1516 pub fn append_thread_name(
1522 &self,
1523 thread_id: &str,
1524 thread_name: Option<String>,
1525 updated_at: i64,
1526 rollout_path: Option<PathBuf>,
1527 ) -> Result<()> {
1528 if let Some(parent) = self.session_index_path.parent() {
1529 fs::create_dir_all(parent).with_context(|| {
1530 format!(
1531 "failed to create session index directory {}",
1532 parent.display()
1533 )
1534 })?;
1535 }
1536 let entry = SessionIndexEntry {
1537 thread_id: thread_id.to_string(),
1538 thread_name,
1539 updated_at,
1540 rollout_path,
1541 };
1542 let encoded =
1543 serde_json::to_string(&entry).context("failed to serialize session index entry")?;
1544 let mut file = OpenOptions::new()
1545 .create(true)
1546 .append(true)
1547 .open(&self.session_index_path)
1548 .with_context(|| {
1549 format!(
1550 "failed to open session index {}",
1551 self.session_index_path.display()
1552 )
1553 })?;
1554 writeln!(file, "{encoded}").context("failed to append session index entry")?;
1555 self.maybe_compact_session_index()?;
1556 Ok(())
1557 }
1558
1559 pub fn find_thread_name_by_id(&self, thread_id: &str) -> Result<Option<String>> {
1563 let map = self.session_index_map()?;
1564 Ok(map
1565 .get(thread_id)
1566 .and_then(|entry| entry.thread_name.clone()))
1567 }
1568
1569 pub fn find_thread_names_by_ids(
1573 &self,
1574 ids: &[String],
1575 ) -> Result<HashMap<String, Option<String>>> {
1576 let map = self.session_index_map()?;
1577 let mut out = HashMap::new();
1578 for id in ids {
1579 let name = map.get(id).and_then(|entry| entry.thread_name.clone());
1580 out.insert(id.clone(), name);
1581 }
1582 Ok(out)
1583 }
1584
1585 pub fn find_thread_path_by_name_str(&self, name: &str) -> Result<Option<PathBuf>> {
1590 let map = self.session_index_map()?;
1591 let matched = map
1592 .values()
1593 .filter(|entry| {
1594 entry
1595 .thread_name
1596 .as_deref()
1597 .is_some_and(|n| n.eq_ignore_ascii_case(name))
1598 })
1599 .max_by_key(|entry| entry.updated_at);
1600 Ok(matched.and_then(|entry| entry.rollout_path.clone()))
1601 }
1602
1603 fn maybe_compact_session_index(&self) -> Result<()> {
1604 if !self.session_index_path.exists() {
1605 return Ok(());
1606 }
1607 let line_count = BufReader::new(
1608 OpenOptions::new()
1609 .read(true)
1610 .open(&self.session_index_path)
1611 .with_context(|| {
1612 format!(
1613 "failed to read session index {}",
1614 self.session_index_path.display()
1615 )
1616 })?,
1617 )
1618 .lines()
1619 .filter(|line| {
1620 line.as_ref()
1621 .map(|value| !value.trim().is_empty())
1622 .unwrap_or(false)
1623 })
1624 .count();
1625 if line_count <= session_index_compact_line_threshold() {
1626 return Ok(());
1627 }
1628
1629 let latest = self.session_index_map()?;
1630 let compact_path = self.session_index_path.with_extension("jsonl.compact");
1631 {
1632 let mut file = OpenOptions::new()
1633 .create(true)
1634 .write(true)
1635 .truncate(true)
1636 .open(&compact_path)
1637 .with_context(|| {
1638 format!(
1639 "failed to open compact session index {}",
1640 compact_path.display()
1641 )
1642 })?;
1643 for entry in latest.values() {
1644 let encoded = serde_json::to_string(entry)
1645 .context("failed to serialize compact session index entry")?;
1646 writeln!(file, "{encoded}")
1647 .context("failed to write compact session index entry")?;
1648 }
1649 }
1650 fs::rename(&compact_path, &self.session_index_path).with_context(|| {
1651 format!(
1652 "failed to replace session index {}",
1653 self.session_index_path.display()
1654 )
1655 })?;
1656 Ok(())
1657 }
1658
1659 #[cfg(test)]
1660 fn session_index_line_count(&self) -> Result<usize> {
1661 if !self.session_index_path.exists() {
1662 return Ok(0);
1663 }
1664 Ok(BufReader::new(
1665 OpenOptions::new()
1666 .read(true)
1667 .open(&self.session_index_path)
1668 .with_context(|| {
1669 format!(
1670 "failed to read session index {}",
1671 self.session_index_path.display()
1672 )
1673 })?,
1674 )
1675 .lines()
1676 .filter(|line| {
1677 line.as_ref()
1678 .map(|value| !value.trim().is_empty())
1679 .unwrap_or(false)
1680 })
1681 .count())
1682 }
1683
1684 fn session_index_map(&self) -> Result<HashMap<String, SessionIndexEntry>> {
1685 if !self.session_index_path.exists() {
1686 return Ok(HashMap::new());
1687 }
1688 let file = OpenOptions::new()
1689 .read(true)
1690 .open(&self.session_index_path)
1691 .with_context(|| {
1692 format!(
1693 "failed to read session index {}",
1694 self.session_index_path.display()
1695 )
1696 })?;
1697 let reader = BufReader::new(file);
1698 let mut latest = HashMap::<String, SessionIndexEntry>::new();
1699 for line in reader.lines() {
1700 let line = line.context("failed to read session index line")?;
1701 if line.trim().is_empty() {
1702 continue;
1703 }
1704 let parsed: SessionIndexEntry =
1705 serde_json::from_str(&line).context("failed to parse session index entry")?;
1706 latest.insert(parsed.thread_id.clone(), parsed);
1707 }
1708 Ok(latest)
1709 }
1710}
1711
1712fn default_state_db_path() -> PathBuf {
1713 if let Some(overridden) = codewhale_home_override() {
1720 return overridden.join("state.db");
1721 }
1722 let home = dirs::home_dir().unwrap_or_else(|| PathBuf::from("."));
1723 let primary = home.join(".codewhale").join("state.db");
1726 if primary.exists() || !home.join(".deepseek").join("state.db").exists() {
1727 primary
1728 } else {
1729 home.join(".deepseek").join("state.db")
1730 }
1731}
1732
1733fn codewhale_home_override() -> Option<PathBuf> {
1743 std::env::var_os("CODEWHALE_HOME")
1744 .filter(|value| !value.is_empty())
1745 .map(PathBuf::from)
1746}
1747
1748fn bool_to_i64(value: bool) -> i64 {
1749 if value { 1 } else { 0 }
1750}
1751
1752fn i64_to_bool(value: i64) -> bool {
1753 value != 0
1754}
1755
1756fn thread_status_to_str(status: &ThreadStatus) -> &'static str {
1757 match status {
1758 ThreadStatus::Running => "running",
1759 ThreadStatus::Idle => "idle",
1760 ThreadStatus::Completed => "completed",
1761 ThreadStatus::Failed => "failed",
1762 ThreadStatus::Paused => "paused",
1763 ThreadStatus::Archived => "archived",
1764 }
1765}
1766
1767fn thread_status_from_str(value: &str) -> ThreadStatus {
1768 match value {
1769 "running" => ThreadStatus::Running,
1770 "idle" => ThreadStatus::Idle,
1771 "completed" => ThreadStatus::Completed,
1772 "failed" => ThreadStatus::Failed,
1773 "paused" => ThreadStatus::Paused,
1774 "archived" => ThreadStatus::Archived,
1775 _ => ThreadStatus::Idle,
1776 }
1777}
1778
1779fn session_source_to_str(source: &SessionSource) -> &'static str {
1780 match source {
1781 SessionSource::Interactive => "interactive",
1782 SessionSource::Resume => "resume",
1783 SessionSource::Fork => "fork",
1784 SessionSource::Api => "api",
1785 SessionSource::Unknown => "unknown",
1786 }
1787}
1788
1789fn session_source_from_str(value: &str) -> SessionSource {
1790 match value {
1791 "interactive" => SessionSource::Interactive,
1792 "resume" => SessionSource::Resume,
1793 "fork" => SessionSource::Fork,
1794 "api" => SessionSource::Api,
1795 _ => SessionSource::Unknown,
1796 }
1797}
1798
1799fn path_to_opt_string(path: Option<&Path>) -> Option<String> {
1800 path.map(|p| p.display().to_string())
1801}
1802
1803fn parse_checkpoint_state(state_json: &str) -> Result<Value> {
1804 serde_json::from_str(state_json).context("failed to parse checkpoint state json")
1805}
1806
1807fn job_state_status_to_str(status: &JobStateStatus) -> &'static str {
1808 match status {
1809 JobStateStatus::Queued => "queued",
1810 JobStateStatus::Running => "running",
1811 JobStateStatus::Paused => "paused",
1812 JobStateStatus::Completed => "completed",
1813 JobStateStatus::Failed => "failed",
1814 JobStateStatus::Cancelled => "cancelled",
1815 }
1816}
1817
1818fn job_state_status_from_str(value: &str) -> JobStateStatus {
1819 match value {
1820 "queued" => JobStateStatus::Queued,
1821 "running" => JobStateStatus::Running,
1822 "paused" => JobStateStatus::Paused,
1823 "completed" => JobStateStatus::Completed,
1824 "failed" => JobStateStatus::Failed,
1825 "cancelled" => JobStateStatus::Cancelled,
1826 _ => JobStateStatus::Queued,
1827 }
1828}
1829
1830fn thread_goal_status_to_str(status: &ThreadGoalStatus) -> &'static str {
1831 match status {
1832 ThreadGoalStatus::Active => "active",
1833 ThreadGoalStatus::Paused => "paused",
1834 ThreadGoalStatus::Blocked => "blocked",
1835 ThreadGoalStatus::UsageLimited => "usage_limited",
1836 ThreadGoalStatus::BudgetLimited => "budget_limited",
1837 ThreadGoalStatus::Complete => "complete",
1838 }
1839}
1840
1841fn thread_goal_status_from_str(value: &str) -> ThreadGoalStatus {
1842 match value {
1843 "active" => ThreadGoalStatus::Active,
1844 "paused" => ThreadGoalStatus::Paused,
1845 "blocked" => ThreadGoalStatus::Blocked,
1846 "usage_limited" => ThreadGoalStatus::UsageLimited,
1847 "budget_limited" => ThreadGoalStatus::BudgetLimited,
1848 "complete" => ThreadGoalStatus::Complete,
1849 _ => ThreadGoalStatus::Paused,
1853 }
1854}
1855
1856fn row_to_thread(row: &rusqlite::Row<'_>) -> rusqlite::Result<ThreadMetadata> {
1857 let status_raw: String = row.get(7)?;
1858 let source_raw: String = row.get(11)?;
1859 let rollout_path: Option<String> = row.get(1)?;
1860 let path: Option<String> = row.get(8)?;
1861 Ok(ThreadMetadata {
1862 id: row.get(0)?,
1863 rollout_path: rollout_path.map(PathBuf::from),
1864 preview: row.get(2)?,
1865 ephemeral: i64_to_bool(row.get(3)?),
1866 model_provider: row.get(4)?,
1867 created_at: row.get(5)?,
1868 updated_at: row.get(6)?,
1869 status: thread_status_from_str(&status_raw),
1870 path: path.map(PathBuf::from),
1871 cwd: PathBuf::from(row.get::<_, String>(9)?),
1872 cli_version: row.get(10)?,
1873 source: session_source_from_str(&source_raw),
1874 name: row.get(12)?,
1875 sandbox_policy: row.get(13)?,
1876 approval_mode: row.get(14)?,
1877 archived: i64_to_bool(row.get(15)?),
1878 archived_at: row.get(16)?,
1879 git_sha: row.get(17)?,
1880 git_branch: row.get(18)?,
1881 git_origin_url: row.get(19)?,
1882 memory_mode: row.get(20)?,
1883 current_leaf_id: row.get(21)?,
1884 })
1885}
1886
1887fn row_to_thread_goal(row: &rusqlite::Row<'_>) -> rusqlite::Result<ThreadGoalRecord> {
1888 let status_raw: String = row.get(3)?;
1889 Ok(ThreadGoalRecord {
1890 thread_id: row.get(0)?,
1891 goal_id: row.get(1)?,
1892 objective: row.get(2)?,
1893 status: thread_goal_status_from_str(&status_raw),
1894 token_budget: row.get(4)?,
1895 tokens_used: row.get(5)?,
1896 time_used_seconds: row.get(6)?,
1897 continuation_count: row.get(7)?,
1898 created_at: row.get(8)?,
1899 updated_at: row.get(9)?,
1900 })
1901}
1902
1903#[cfg(test)]
1904mod tests {
1905 use super::*;
1906 use serde_json::json;
1907 use std::time::{SystemTime, UNIX_EPOCH};
1908
1909 fn temp_state_store(name: &str) -> StateStore {
1910 let suffix = SystemTime::now()
1911 .duration_since(UNIX_EPOCH)
1912 .expect("system time")
1913 .as_nanos();
1914 let dir = std::env::temp_dir().join(format!(
1915 "codewhale-state-{name}-{}-{suffix}",
1916 std::process::id()
1917 ));
1918 fs::create_dir_all(&dir).expect("create temp state dir");
1919 StateStore::open(Some(dir.join("state.db"))).expect("open state store")
1920 }
1921
1922 fn test_thread(id: &str) -> ThreadMetadata {
1923 ThreadMetadata {
1924 id: id.to_string(),
1925 rollout_path: None,
1926 preview: "test thread".to_string(),
1927 ephemeral: false,
1928 model_provider: "deepseek".to_string(),
1929 created_at: 10,
1930 updated_at: 10,
1931 status: ThreadStatus::Running,
1932 path: None,
1933 cwd: PathBuf::from("/tmp/codewhale"),
1934 cli_version: "0.0.0-test".to_string(),
1935 source: SessionSource::Interactive,
1936 name: None,
1937 sandbox_policy: None,
1938 approval_mode: None,
1939 archived: false,
1940 archived_at: None,
1941 git_sha: None,
1942 git_branch: None,
1943 git_origin_url: None,
1944 memory_mode: None,
1945 current_leaf_id: None,
1946 }
1947 }
1948
1949 fn test_goal(thread_id: &str, objective: &str) -> ThreadGoalRecord {
1950 ThreadGoalRecord {
1951 thread_id: thread_id.to_string(),
1952 goal_id: "goal-1".to_string(),
1953 objective: objective.to_string(),
1954 status: ThreadGoalStatus::Active,
1955 token_budget: Some(123),
1956 tokens_used: 7,
1957 time_used_seconds: 11,
1958 continuation_count: 0,
1959 created_at: 100,
1960 updated_at: 101,
1961 }
1962 }
1963
1964 #[test]
1965 fn unknown_persisted_goal_status_fails_closed() {
1966 assert_eq!(
1967 thread_goal_status_from_str("future_or_corrupt_status"),
1968 ThreadGoalStatus::Paused
1969 );
1970 }
1971
1972 #[test]
1973 fn thread_goal_crud_round_trips_and_replaces() {
1974 let store = temp_state_store("thread-goal-crud");
1975 store
1976 .upsert_thread(&test_thread("thread-1"))
1977 .expect("upsert thread");
1978
1979 let goal = test_goal("thread-1", "Ship v0.8.59");
1980 store.upsert_thread_goal(&goal).expect("upsert goal");
1981 assert_eq!(
1982 store
1983 .get_thread_goal("thread-1")
1984 .expect("read goal")
1985 .as_ref(),
1986 Some(&goal)
1987 );
1988
1989 let mut replacement = test_goal("thread-1", "Ship v0.8.59 safely");
1990 replacement.goal_id = "goal-2".to_string();
1991 replacement.status = ThreadGoalStatus::BudgetLimited;
1992 replacement.token_budget = None;
1993 replacement.updated_at = 202;
1994 store
1995 .upsert_thread_goal(&replacement)
1996 .expect("replace goal");
1997 assert_eq!(
1998 store.get_thread_goal("thread-1").expect("read replacement"),
1999 Some(replacement)
2000 );
2001
2002 assert!(store.delete_thread_goal("thread-1").expect("delete goal"));
2003 assert!(
2004 store
2005 .get_thread_goal("thread-1")
2006 .expect("read empty")
2007 .is_none()
2008 );
2009 assert!(!store.delete_thread_goal("thread-1").expect("delete empty"));
2010 }
2011
2012 #[test]
2013 fn thread_goal_requires_existing_thread() {
2014 let store = temp_state_store("thread-goal-missing-thread");
2015 let err = store
2016 .upsert_thread_goal(&test_goal("missing-thread", "nope"))
2017 .expect_err("goal without a thread should fail");
2018 assert!(err.to_string().contains("thread missing-thread not found"));
2019 }
2020
2021 #[test]
2022 fn delete_thread_cascades_child_rows() {
2023 let store = temp_state_store("thread-delete-cascade");
2024 store
2025 .upsert_thread(&test_thread("thread-1"))
2026 .expect("upsert thread");
2027 store
2028 .append_message("thread-1", "user", "hello", None)
2029 .expect("append message");
2030 store
2031 .save_checkpoint("thread-1", "checkpoint-1", &serde_json::json!({"ok": true}))
2032 .expect("save checkpoint");
2033 store
2034 .persist_dynamic_tools(
2035 "thread-1",
2036 &[DynamicToolRecord {
2037 position: 0,
2038 name: "test_tool".to_string(),
2039 description: Some("test".to_string()),
2040 input_schema: serde_json::json!({"type": "object"}),
2041 }],
2042 )
2043 .expect("persist dynamic tools");
2044 store
2045 .upsert_thread_goal(&test_goal("thread-1", "Ship v0.8.67"))
2046 .expect("upsert goal");
2047
2048 store.delete_thread("thread-1").expect("delete thread");
2049
2050 let conn = store.conn().expect("conn");
2051 for table in [
2052 "messages",
2053 "checkpoints",
2054 "thread_dynamic_tools",
2055 "thread_goals",
2056 ] {
2057 let sql = format!("SELECT COUNT(*) FROM {table} WHERE thread_id = ?1");
2058 let count: i64 = conn
2059 .query_row(&sql, params!["thread-1"], |row| row.get(0))
2060 .expect("count child rows");
2061 assert_eq!(count, 0, "{table} row survived thread deletion");
2062 }
2063 }
2064
2065 #[test]
2066 fn state_store_reuses_one_connection_across_operations_and_clones() {
2067 let store = temp_state_store("conn-reuse");
2068 {
2069 let conn = store.conn().expect("conn");
2070 conn.execute_batch("CREATE TEMP TABLE conn_reuse_probe(id INTEGER);")
2071 .expect("create temp table");
2072 }
2073 let clone = store.clone();
2078 clone
2079 .upsert_thread(&test_thread("thread-conn-reuse"))
2080 .expect("upsert thread");
2081 let conn = clone.conn().expect("conn");
2082 let probe_count: i64 = conn
2083 .query_row(
2084 "SELECT COUNT(*) FROM sqlite_temp_master WHERE name = 'conn_reuse_probe'",
2085 [],
2086 |row| row.get(0),
2087 )
2088 .expect("query temp master");
2089 assert_eq!(
2090 probe_count, 1,
2091 "temp table not visible: a fresh connection was opened"
2092 );
2093 let foreign_keys: i64 = conn
2095 .query_row("PRAGMA foreign_keys;", [], |row| row.get(0))
2096 .expect("read foreign_keys pragma");
2097 assert_eq!(foreign_keys, 1);
2098 let journal_mode: String = conn
2099 .query_row("PRAGMA journal_mode;", [], |row| row.get(0))
2100 .expect("read journal_mode pragma");
2101 assert_eq!(
2102 journal_mode.to_ascii_lowercase(),
2103 "wal",
2104 "open should enable WAL for multi-process readers/writers"
2105 );
2106 }
2107
2108 #[test]
2115 fn two_connections_can_write_concurrently_without_sqlite_busy() {
2116 let suffix = SystemTime::now()
2117 .duration_since(UNIX_EPOCH)
2118 .expect("system time")
2119 .as_nanos();
2120 let dir = std::env::temp_dir().join(format!(
2121 "codewhale-state-concurrent-write-{}-{suffix}",
2122 std::process::id()
2123 ));
2124 fs::create_dir_all(&dir).expect("create temp state dir");
2125 let db_path = dir.join("state.db");
2126
2127 let bootstrap = StateStore::open(Some(db_path.clone())).expect("bootstrap open");
2129 {
2130 let conn = bootstrap.conn().expect("bootstrap conn");
2131 let journal_mode: String = conn
2132 .query_row("PRAGMA journal_mode;", [], |row| row.get(0))
2133 .expect("journal_mode");
2134 assert_eq!(journal_mode.to_ascii_lowercase(), "wal");
2135 }
2136 drop(bootstrap);
2137
2138 let path_a = db_path.clone();
2139 let path_b = db_path.clone();
2140 const WRITES_PER_CONN: usize = 50;
2141
2142 let handle_a = std::thread::spawn(move || {
2143 let store = StateStore::open(Some(path_a)).expect("open store a");
2144 for i in 0..WRITES_PER_CONN {
2145 let job = JobStateRecord {
2146 id: format!("job-a-{i}"),
2147 name: format!("writer-a-{i}"),
2148 status: JobStateStatus::Running,
2149 progress: Some((i % 100) as u8),
2150 detail: Some("concurrent write a".to_string()),
2151 created_at: i as i64,
2152 updated_at: i as i64,
2153 };
2154 store.upsert_job(&job).unwrap_or_else(|err| {
2155 panic!("connection A write {i} failed (must not be SQLITE_BUSY): {err:#}");
2156 });
2157 }
2158 });
2159 let handle_b = std::thread::spawn(move || {
2160 let store = StateStore::open(Some(path_b)).expect("open store b");
2161 for i in 0..WRITES_PER_CONN {
2162 let job = JobStateRecord {
2163 id: format!("job-b-{i}"),
2164 name: format!("writer-b-{i}"),
2165 status: JobStateStatus::Running,
2166 progress: Some((i % 100) as u8),
2167 detail: Some("concurrent write b".to_string()),
2168 created_at: i as i64,
2169 updated_at: i as i64,
2170 };
2171 store.upsert_job(&job).unwrap_or_else(|err| {
2172 panic!("connection B write {i} failed (must not be SQLITE_BUSY): {err:#}");
2173 });
2174 }
2175 });
2176
2177 handle_a.join().expect("writer A panicked");
2178 handle_b.join().expect("writer B panicked");
2179
2180 let store = StateStore::open(Some(db_path)).expect("reopen for verify");
2181 let listed = store
2182 .list_jobs(Some(WRITES_PER_CONN * 2))
2183 .expect("list jobs");
2184 assert_eq!(
2185 listed.len(),
2186 WRITES_PER_CONN * 2,
2187 "both connections should have persisted all jobs"
2188 );
2189
2190 let _ = fs::remove_dir_all(dir);
2191 }
2192
2193 #[test]
2194 fn record_thread_goal_usage_accumulates_tokens_and_time() {
2195 let store = temp_state_store("thread-goal-usage");
2196 store
2197 .upsert_thread(&test_thread("thread-1"))
2198 .expect("upsert thread");
2199
2200 let mut goal = test_goal("thread-1", "Ship the persistent goal loop");
2202 goal.tokens_used = 0;
2203 goal.time_used_seconds = 0;
2204 goal.updated_at = 100;
2205 store.upsert_thread_goal(&goal).expect("upsert goal");
2206
2207 let after_first = store
2209 .record_thread_goal_usage("thread-1", 250, 12, 150)
2210 .expect("record usage")
2211 .expect("goal exists");
2212 assert_eq!(after_first.tokens_used, 250);
2213 assert_eq!(after_first.time_used_seconds, 12);
2214 assert_eq!(after_first.updated_at, 150);
2215 assert_eq!(after_first.goal_id, goal.goal_id);
2217 assert_eq!(after_first.objective, goal.objective);
2218 assert_eq!(after_first.status, goal.status);
2219 assert_eq!(after_first.token_budget, goal.token_budget);
2220 assert_eq!(after_first.created_at, goal.created_at);
2221 assert_eq!(after_first.continuation_count, 0);
2222
2223 let after_second = store
2225 .record_thread_goal_usage("thread-1", 75, 8, 200)
2226 .expect("record usage")
2227 .expect("goal exists");
2228 assert_eq!(after_second.tokens_used, 325);
2229 assert_eq!(after_second.time_used_seconds, 20);
2230 assert_eq!(after_second.updated_at, 200);
2231
2232 let after_stale = store
2234 .record_thread_goal_usage("thread-1", 5, 1, 1)
2235 .expect("record usage")
2236 .expect("goal exists");
2237 assert_eq!(after_stale.tokens_used, 330);
2238 assert_eq!(after_stale.time_used_seconds, 21);
2239 assert_eq!(after_stale.updated_at, 200);
2240
2241 let persisted = store
2243 .get_thread_goal("thread-1")
2244 .expect("read goal")
2245 .expect("goal exists");
2246 assert_eq!(persisted.tokens_used, 330);
2247 assert_eq!(persisted.time_used_seconds, 21);
2248 }
2249
2250 #[test]
2251 fn record_thread_goal_usage_returns_none_without_goal() {
2252 let store = temp_state_store("thread-goal-usage-missing");
2253 store
2254 .upsert_thread(&test_thread("thread-1"))
2255 .expect("upsert thread");
2256 let result = store
2259 .record_thread_goal_usage("thread-1", 100, 5, 999)
2260 .expect("record usage on goalless thread");
2261 assert!(result.is_none());
2262 assert!(
2263 store
2264 .get_thread_goal("thread-1")
2265 .expect("read goal")
2266 .is_none()
2267 );
2268 }
2269
2270 #[test]
2271 fn record_thread_goal_continuation_accumulates_durably() {
2272 let store = temp_state_store("thread-goal-continuation");
2273 store
2274 .upsert_thread(&test_thread("thread-1"))
2275 .expect("upsert thread");
2276
2277 let mut goal = test_goal("thread-1", "Keep working across turns");
2278 goal.updated_at = 100;
2279 store.upsert_thread_goal(&goal).expect("upsert goal");
2280
2281 let after_first = store
2282 .record_thread_goal_continuation("thread-1", 120)
2283 .expect("record continuation")
2284 .expect("goal exists");
2285 assert_eq!(after_first.continuation_count, 1);
2286 assert_eq!(after_first.tokens_used, goal.tokens_used);
2287 assert_eq!(after_first.time_used_seconds, goal.time_used_seconds);
2288 assert_eq!(after_first.updated_at, 120);
2289
2290 let after_second = store
2291 .record_thread_goal_continuation("thread-1", 110)
2292 .expect("record second continuation")
2293 .expect("goal exists");
2294 assert_eq!(after_second.continuation_count, 2);
2295 assert_eq!(after_second.updated_at, 120);
2296
2297 let persisted = store
2298 .get_thread_goal("thread-1")
2299 .expect("read goal")
2300 .expect("goal exists");
2301 assert_eq!(persisted.continuation_count, 2);
2302 }
2303
2304 static CODEWHALE_HOME_TEST_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
2311
2312 struct CodeWhaleHomeGuard {
2313 prior: Option<std::ffi::OsString>,
2314 }
2315 impl CodeWhaleHomeGuard {
2316 fn set(value: &str) -> Self {
2317 let prior = std::env::var_os("CODEWHALE_HOME");
2318 unsafe { std::env::set_var("CODEWHALE_HOME", value) };
2320 Self { prior }
2321 }
2322 fn remove() -> Self {
2323 let prior = std::env::var_os("CODEWHALE_HOME");
2324 unsafe { std::env::remove_var("CODEWHALE_HOME") };
2326 Self { prior }
2327 }
2328 }
2329 impl Drop for CodeWhaleHomeGuard {
2330 fn drop(&mut self) {
2331 unsafe {
2333 match &self.prior {
2334 Some(value) => std::env::set_var("CODEWHALE_HOME", value),
2335 None => std::env::remove_var("CODEWHALE_HOME"),
2336 }
2337 }
2338 }
2339 }
2340
2341 #[test]
2342 fn codewhale_home_override_returns_the_env_value_verbatim() {
2343 let _lock = CODEWHALE_HOME_TEST_LOCK.lock().unwrap();
2344 let _g = CodeWhaleHomeGuard::set("/tmp/cw-isolated-state");
2345 assert_eq!(
2348 codewhale_home_override().as_deref(),
2349 Some(std::path::Path::new("/tmp/cw-isolated-state"))
2350 );
2351 }
2352
2353 #[test]
2354 fn codewhale_home_override_none_when_unset() {
2355 let _lock = CODEWHALE_HOME_TEST_LOCK.lock().unwrap();
2356 let _g = CodeWhaleHomeGuard::remove();
2357 assert!(codewhale_home_override().is_none());
2358 }
2359
2360 #[test]
2361 fn codewhale_home_override_none_when_empty() {
2362 let _lock = CODEWHALE_HOME_TEST_LOCK.lock().unwrap();
2363 let _g = CodeWhaleHomeGuard::set(" ");
2364 assert!(
2370 codewhale_home_override().is_some(),
2371 "non-empty (even whitespace) counts as set; trimming is the caller's job"
2372 );
2373 }
2374
2375 #[test]
2376 fn default_state_db_path_uses_codewhale_home_when_set() {
2377 let _lock = CODEWHALE_HOME_TEST_LOCK.lock().unwrap();
2378 let dir = std::env::temp_dir().join(format!(
2379 "cw-home-state-{}-{}",
2380 std::process::id(),
2381 std::time::SystemTime::now()
2382 .duration_since(std::time::UNIX_EPOCH)
2383 .unwrap()
2384 .as_nanos()
2385 ));
2386 let _g = CodeWhaleHomeGuard::set(dir.to_str().unwrap());
2387 assert_eq!(default_state_db_path(), dir.join("state.db"));
2391 }
2392
2393 #[test]
2394 fn load_checkpoint_propagates_invalid_state_json() {
2395 let store = temp_state_store("checkpoint-parse-error");
2396 store
2397 .upsert_thread(&test_thread("thread-1"))
2398 .expect("upsert thread");
2399 store
2400 .save_checkpoint("thread-1", "broken", &json!({"ok": true}))
2401 .expect("save checkpoint");
2402
2403 {
2404 let conn = store.conn().expect("conn");
2405 conn.execute(
2406 "UPDATE checkpoints SET state_json = ?1 WHERE thread_id = ?2 AND checkpoint_id = ?3",
2407 params!["not-json", "thread-1", "broken"],
2408 )
2409 .expect("corrupt checkpoint");
2410 }
2411
2412 let err = store
2413 .load_checkpoint("thread-1", Some("broken"))
2414 .expect_err("invalid checkpoint json should fail");
2415 assert!(
2416 err.to_string()
2417 .contains("failed to parse checkpoint state json")
2418 );
2419 }
2420
2421 #[test]
2422 fn session_index_compacts_after_threshold() {
2423 let store = temp_state_store("session-index-compact");
2424 for idx in 0..6 {
2425 store
2426 .append_thread_name("thread-1", Some(format!("name-{idx}")), idx, None)
2427 .expect("append session index entry");
2428 }
2429
2430 let line_count = store
2431 .session_index_line_count()
2432 .expect("count session index lines");
2433 assert_eq!(line_count, 1);
2434
2435 let name = store
2436 .find_thread_name_by_id("thread-1")
2437 .expect("lookup thread name");
2438 assert_eq!(name.as_deref(), Some("name-5"));
2439 }
2440}