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 self.with_session_index_lock(|| {
1549 let mut file = OpenOptions::new()
1550 .create(true)
1551 .append(true)
1552 .open(&self.session_index_path)
1553 .with_context(|| {
1554 format!(
1555 "failed to open session index {}",
1556 self.session_index_path.display()
1557 )
1558 })?;
1559 writeln!(file, "{encoded}").context("failed to append session index entry")?;
1560 file.sync_data()
1564 .context("failed to flush session index entry")?;
1565 drop(file);
1566 self.compact_session_index_locked()
1567 })
1568 }
1569
1570 fn with_session_index_lock<T>(&self, operation: impl FnOnce() -> Result<T>) -> Result<T> {
1576 if let Some(parent) = self.session_index_path.parent() {
1577 fs::create_dir_all(parent).with_context(|| {
1578 format!(
1579 "failed to create session index directory {}",
1580 parent.display()
1581 )
1582 })?;
1583 }
1584 let lock_path = self.session_index_path.with_extension("jsonl.lock");
1585 let lock_file = OpenOptions::new()
1586 .create(true)
1587 .read(true)
1588 .write(true)
1589 .truncate(false)
1592 .open(&lock_path)
1593 .with_context(|| {
1594 format!("failed to open session index lock {}", lock_path.display())
1595 })?;
1596 #[cfg(unix)]
1597 {
1598 use std::os::unix::fs::PermissionsExt as _;
1599 lock_file
1600 .set_permissions(fs::Permissions::from_mode(0o600))
1601 .with_context(|| {
1602 format!(
1603 "failed to secure session index lock {}",
1604 lock_path.display()
1605 )
1606 })?;
1607 }
1608 let mut lock = fd_lock::RwLock::new(lock_file);
1609 let _guard = lock
1610 .write()
1611 .with_context(|| format!("failed to lock session index {}", lock_path.display()))?;
1612 operation()
1613 }
1614
1615 pub fn find_thread_name_by_id(&self, thread_id: &str) -> Result<Option<String>> {
1619 let map = self.session_index_map()?;
1620 Ok(map
1621 .get(thread_id)
1622 .and_then(|entry| entry.thread_name.clone()))
1623 }
1624
1625 pub fn find_thread_names_by_ids(
1629 &self,
1630 ids: &[String],
1631 ) -> Result<HashMap<String, Option<String>>> {
1632 let map = self.session_index_map()?;
1633 let mut out = HashMap::new();
1634 for id in ids {
1635 let name = map.get(id).and_then(|entry| entry.thread_name.clone());
1636 out.insert(id.clone(), name);
1637 }
1638 Ok(out)
1639 }
1640
1641 pub fn find_thread_path_by_name_str(&self, name: &str) -> Result<Option<PathBuf>> {
1646 let map = self.session_index_map()?;
1647 let matched = map
1648 .values()
1649 .filter(|entry| {
1650 entry
1651 .thread_name
1652 .as_deref()
1653 .is_some_and(|n| n.eq_ignore_ascii_case(name))
1654 })
1655 .max_by_key(|entry| entry.updated_at);
1656 Ok(matched.and_then(|entry| entry.rollout_path.clone()))
1657 }
1658
1659 fn compact_session_index_locked(&self) -> Result<()> {
1664 if !self.session_index_path.exists() {
1665 return Ok(());
1666 }
1667 let line_count = BufReader::new(
1668 OpenOptions::new()
1669 .read(true)
1670 .open(&self.session_index_path)
1671 .with_context(|| {
1672 format!(
1673 "failed to read session index {}",
1674 self.session_index_path.display()
1675 )
1676 })?,
1677 )
1678 .lines()
1679 .filter(|line| {
1680 line.as_ref()
1681 .map(|value| !value.trim().is_empty())
1682 .unwrap_or(false)
1683 })
1684 .count();
1685 if line_count <= session_index_compact_line_threshold() {
1686 return Ok(());
1687 }
1688
1689 let latest = self.session_index_map()?;
1690 let compact_path = self.session_index_path.with_extension("jsonl.compact");
1691 {
1692 let mut file = OpenOptions::new()
1693 .create(true)
1694 .write(true)
1695 .truncate(true)
1696 .open(&compact_path)
1697 .with_context(|| {
1698 format!(
1699 "failed to open compact session index {}",
1700 compact_path.display()
1701 )
1702 })?;
1703 for entry in latest.values() {
1704 let encoded = serde_json::to_string(entry)
1705 .context("failed to serialize compact session index entry")?;
1706 writeln!(file, "{encoded}")
1707 .context("failed to write compact session index entry")?;
1708 }
1709 }
1710 #[cfg(test)]
1714 tests::compaction_midpoint(&self.session_index_path);
1715 fs::rename(&compact_path, &self.session_index_path).with_context(|| {
1716 format!(
1717 "failed to replace session index {}",
1718 self.session_index_path.display()
1719 )
1720 })?;
1721 Ok(())
1722 }
1723
1724 #[cfg(test)]
1725 fn session_index_line_count(&self) -> Result<usize> {
1726 if !self.session_index_path.exists() {
1727 return Ok(0);
1728 }
1729 Ok(BufReader::new(
1730 OpenOptions::new()
1731 .read(true)
1732 .open(&self.session_index_path)
1733 .with_context(|| {
1734 format!(
1735 "failed to read session index {}",
1736 self.session_index_path.display()
1737 )
1738 })?,
1739 )
1740 .lines()
1741 .filter(|line| {
1742 line.as_ref()
1743 .map(|value| !value.trim().is_empty())
1744 .unwrap_or(false)
1745 })
1746 .count())
1747 }
1748
1749 fn session_index_map(&self) -> Result<HashMap<String, SessionIndexEntry>> {
1750 if !self.session_index_path.exists() {
1751 return Ok(HashMap::new());
1752 }
1753 let file = OpenOptions::new()
1754 .read(true)
1755 .open(&self.session_index_path)
1756 .with_context(|| {
1757 format!(
1758 "failed to read session index {}",
1759 self.session_index_path.display()
1760 )
1761 })?;
1762 let reader = BufReader::new(file);
1763 let mut latest = HashMap::<String, SessionIndexEntry>::new();
1764 for line in reader.lines() {
1765 let line = line.context("failed to read session index line")?;
1766 if line.trim().is_empty() {
1767 continue;
1768 }
1769 match serde_json::from_str::<SessionIndexEntry>(&line) {
1776 Ok(parsed) => {
1777 latest.insert(parsed.thread_id.clone(), parsed);
1778 }
1779 Err(err) => {
1780 tracing::warn!(
1781 "skipping unparseable session index entry in {}: {err}",
1782 self.session_index_path.display()
1783 );
1784 }
1785 }
1786 }
1787 Ok(latest)
1788 }
1789}
1790
1791fn default_state_db_path() -> PathBuf {
1792 if let Some(overridden) = codewhale_home_override() {
1799 return overridden.join("state.db");
1800 }
1801 let home = dirs::home_dir().unwrap_or_else(|| PathBuf::from("."));
1802 let primary = home.join(".codewhale").join("state.db");
1805 if primary.exists() || !home.join(".deepseek").join("state.db").exists() {
1806 primary
1807 } else {
1808 home.join(".deepseek").join("state.db")
1809 }
1810}
1811
1812fn codewhale_home_override() -> Option<PathBuf> {
1822 std::env::var_os("CODEWHALE_HOME")
1823 .filter(|value| !value.is_empty())
1824 .map(PathBuf::from)
1825}
1826
1827fn bool_to_i64(value: bool) -> i64 {
1828 if value { 1 } else { 0 }
1829}
1830
1831fn i64_to_bool(value: i64) -> bool {
1832 value != 0
1833}
1834
1835fn thread_status_to_str(status: &ThreadStatus) -> &'static str {
1836 match status {
1837 ThreadStatus::Running => "running",
1838 ThreadStatus::Idle => "idle",
1839 ThreadStatus::Completed => "completed",
1840 ThreadStatus::Failed => "failed",
1841 ThreadStatus::Paused => "paused",
1842 ThreadStatus::Archived => "archived",
1843 }
1844}
1845
1846fn thread_status_from_str(value: &str) -> ThreadStatus {
1847 match value {
1848 "running" => ThreadStatus::Running,
1849 "idle" => ThreadStatus::Idle,
1850 "completed" => ThreadStatus::Completed,
1851 "failed" => ThreadStatus::Failed,
1852 "paused" => ThreadStatus::Paused,
1853 "archived" => ThreadStatus::Archived,
1854 _ => ThreadStatus::Idle,
1855 }
1856}
1857
1858fn session_source_to_str(source: &SessionSource) -> &'static str {
1859 match source {
1860 SessionSource::Interactive => "interactive",
1861 SessionSource::Resume => "resume",
1862 SessionSource::Fork => "fork",
1863 SessionSource::Api => "api",
1864 SessionSource::Unknown => "unknown",
1865 }
1866}
1867
1868fn session_source_from_str(value: &str) -> SessionSource {
1869 match value {
1870 "interactive" => SessionSource::Interactive,
1871 "resume" => SessionSource::Resume,
1872 "fork" => SessionSource::Fork,
1873 "api" => SessionSource::Api,
1874 _ => SessionSource::Unknown,
1875 }
1876}
1877
1878fn path_to_opt_string(path: Option<&Path>) -> Option<String> {
1879 path.map(|p| p.display().to_string())
1880}
1881
1882fn parse_checkpoint_state(state_json: &str) -> Result<Value> {
1883 serde_json::from_str(state_json).context("failed to parse checkpoint state json")
1884}
1885
1886fn job_state_status_to_str(status: &JobStateStatus) -> &'static str {
1887 match status {
1888 JobStateStatus::Queued => "queued",
1889 JobStateStatus::Running => "running",
1890 JobStateStatus::Paused => "paused",
1891 JobStateStatus::Completed => "completed",
1892 JobStateStatus::Failed => "failed",
1893 JobStateStatus::Cancelled => "cancelled",
1894 }
1895}
1896
1897fn job_state_status_from_str(value: &str) -> JobStateStatus {
1898 match value {
1899 "queued" => JobStateStatus::Queued,
1900 "running" => JobStateStatus::Running,
1901 "paused" => JobStateStatus::Paused,
1902 "completed" => JobStateStatus::Completed,
1903 "failed" => JobStateStatus::Failed,
1904 "cancelled" => JobStateStatus::Cancelled,
1905 _ => JobStateStatus::Queued,
1906 }
1907}
1908
1909fn thread_goal_status_to_str(status: &ThreadGoalStatus) -> &'static str {
1910 match status {
1911 ThreadGoalStatus::Active => "active",
1912 ThreadGoalStatus::Paused => "paused",
1913 ThreadGoalStatus::Blocked => "blocked",
1914 ThreadGoalStatus::UsageLimited => "usage_limited",
1915 ThreadGoalStatus::BudgetLimited => "budget_limited",
1916 ThreadGoalStatus::Complete => "complete",
1917 }
1918}
1919
1920fn thread_goal_status_from_str(value: &str) -> ThreadGoalStatus {
1921 match value {
1922 "active" => ThreadGoalStatus::Active,
1923 "paused" => ThreadGoalStatus::Paused,
1924 "blocked" => ThreadGoalStatus::Blocked,
1925 "usage_limited" => ThreadGoalStatus::UsageLimited,
1926 "budget_limited" => ThreadGoalStatus::BudgetLimited,
1927 "complete" => ThreadGoalStatus::Complete,
1928 _ => ThreadGoalStatus::Paused,
1932 }
1933}
1934
1935fn row_to_thread(row: &rusqlite::Row<'_>) -> rusqlite::Result<ThreadMetadata> {
1936 let status_raw: String = row.get(7)?;
1937 let source_raw: String = row.get(11)?;
1938 let rollout_path: Option<String> = row.get(1)?;
1939 let path: Option<String> = row.get(8)?;
1940 Ok(ThreadMetadata {
1941 id: row.get(0)?,
1942 rollout_path: rollout_path.map(PathBuf::from),
1943 preview: row.get(2)?,
1944 ephemeral: i64_to_bool(row.get(3)?),
1945 model_provider: row.get(4)?,
1946 created_at: row.get(5)?,
1947 updated_at: row.get(6)?,
1948 status: thread_status_from_str(&status_raw),
1949 path: path.map(PathBuf::from),
1950 cwd: PathBuf::from(row.get::<_, String>(9)?),
1951 cli_version: row.get(10)?,
1952 source: session_source_from_str(&source_raw),
1953 name: row.get(12)?,
1954 sandbox_policy: row.get(13)?,
1955 approval_mode: row.get(14)?,
1956 archived: i64_to_bool(row.get(15)?),
1957 archived_at: row.get(16)?,
1958 git_sha: row.get(17)?,
1959 git_branch: row.get(18)?,
1960 git_origin_url: row.get(19)?,
1961 memory_mode: row.get(20)?,
1962 current_leaf_id: row.get(21)?,
1963 })
1964}
1965
1966fn row_to_thread_goal(row: &rusqlite::Row<'_>) -> rusqlite::Result<ThreadGoalRecord> {
1967 let status_raw: String = row.get(3)?;
1968 Ok(ThreadGoalRecord {
1969 thread_id: row.get(0)?,
1970 goal_id: row.get(1)?,
1971 objective: row.get(2)?,
1972 status: thread_goal_status_from_str(&status_raw),
1973 token_budget: row.get(4)?,
1974 tokens_used: row.get(5)?,
1975 time_used_seconds: row.get(6)?,
1976 continuation_count: row.get(7)?,
1977 created_at: row.get(8)?,
1978 updated_at: row.get(9)?,
1979 })
1980}
1981
1982#[cfg(test)]
1983mod tests {
1984 use super::*;
1985 use serde_json::json;
1986 use std::sync::{Arc, Barrier, Mutex};
1987 use std::thread;
1988 use std::time::{Duration, SystemTime, UNIX_EPOCH};
1989
1990 fn temp_state_store(name: &str) -> StateStore {
1991 let suffix = SystemTime::now()
1992 .duration_since(UNIX_EPOCH)
1993 .expect("system time")
1994 .as_nanos();
1995 let dir = std::env::temp_dir().join(format!(
1996 "codewhale-state-{name}-{}-{suffix}",
1997 std::process::id()
1998 ));
1999 fs::create_dir_all(&dir).expect("create temp state dir");
2000 StateStore::open(Some(dir.join("state.db"))).expect("open state store")
2001 }
2002
2003 fn test_thread(id: &str) -> ThreadMetadata {
2004 ThreadMetadata {
2005 id: id.to_string(),
2006 rollout_path: None,
2007 preview: "test thread".to_string(),
2008 ephemeral: false,
2009 model_provider: "deepseek".to_string(),
2010 created_at: 10,
2011 updated_at: 10,
2012 status: ThreadStatus::Running,
2013 path: None,
2014 cwd: PathBuf::from("/tmp/codewhale"),
2015 cli_version: "0.0.0-test".to_string(),
2016 source: SessionSource::Interactive,
2017 name: None,
2018 sandbox_policy: None,
2019 approval_mode: None,
2020 archived: false,
2021 archived_at: None,
2022 git_sha: None,
2023 git_branch: None,
2024 git_origin_url: None,
2025 memory_mode: None,
2026 current_leaf_id: None,
2027 }
2028 }
2029
2030 fn test_goal(thread_id: &str, objective: &str) -> ThreadGoalRecord {
2031 ThreadGoalRecord {
2032 thread_id: thread_id.to_string(),
2033 goal_id: "goal-1".to_string(),
2034 objective: objective.to_string(),
2035 status: ThreadGoalStatus::Active,
2036 token_budget: Some(123),
2037 tokens_used: 7,
2038 time_used_seconds: 11,
2039 continuation_count: 0,
2040 created_at: 100,
2041 updated_at: 101,
2042 }
2043 }
2044
2045 #[test]
2046 fn unknown_persisted_goal_status_fails_closed() {
2047 assert_eq!(
2048 thread_goal_status_from_str("future_or_corrupt_status"),
2049 ThreadGoalStatus::Paused
2050 );
2051 }
2052
2053 #[test]
2054 fn thread_goal_crud_round_trips_and_replaces() {
2055 let store = temp_state_store("thread-goal-crud");
2056 store
2057 .upsert_thread(&test_thread("thread-1"))
2058 .expect("upsert thread");
2059
2060 let goal = test_goal("thread-1", "Ship v0.8.59");
2061 store.upsert_thread_goal(&goal).expect("upsert goal");
2062 assert_eq!(
2063 store
2064 .get_thread_goal("thread-1")
2065 .expect("read goal")
2066 .as_ref(),
2067 Some(&goal)
2068 );
2069
2070 let mut replacement = test_goal("thread-1", "Ship v0.8.59 safely");
2071 replacement.goal_id = "goal-2".to_string();
2072 replacement.status = ThreadGoalStatus::BudgetLimited;
2073 replacement.token_budget = None;
2074 replacement.updated_at = 202;
2075 store
2076 .upsert_thread_goal(&replacement)
2077 .expect("replace goal");
2078 assert_eq!(
2079 store.get_thread_goal("thread-1").expect("read replacement"),
2080 Some(replacement)
2081 );
2082
2083 assert!(store.delete_thread_goal("thread-1").expect("delete goal"));
2084 assert!(
2085 store
2086 .get_thread_goal("thread-1")
2087 .expect("read empty")
2088 .is_none()
2089 );
2090 assert!(!store.delete_thread_goal("thread-1").expect("delete empty"));
2091 }
2092
2093 #[test]
2094 fn thread_goal_requires_existing_thread() {
2095 let store = temp_state_store("thread-goal-missing-thread");
2096 let err = store
2097 .upsert_thread_goal(&test_goal("missing-thread", "nope"))
2098 .expect_err("goal without a thread should fail");
2099 assert!(err.to_string().contains("thread missing-thread not found"));
2100 }
2101
2102 #[test]
2103 fn delete_thread_cascades_child_rows() {
2104 let store = temp_state_store("thread-delete-cascade");
2105 store
2106 .upsert_thread(&test_thread("thread-1"))
2107 .expect("upsert thread");
2108 store
2109 .append_message("thread-1", "user", "hello", None)
2110 .expect("append message");
2111 store
2112 .save_checkpoint("thread-1", "checkpoint-1", &serde_json::json!({"ok": true}))
2113 .expect("save checkpoint");
2114 store
2115 .persist_dynamic_tools(
2116 "thread-1",
2117 &[DynamicToolRecord {
2118 position: 0,
2119 name: "test_tool".to_string(),
2120 description: Some("test".to_string()),
2121 input_schema: serde_json::json!({"type": "object"}),
2122 }],
2123 )
2124 .expect("persist dynamic tools");
2125 store
2126 .upsert_thread_goal(&test_goal("thread-1", "Ship v0.8.67"))
2127 .expect("upsert goal");
2128
2129 store.delete_thread("thread-1").expect("delete thread");
2130
2131 let conn = store.conn().expect("conn");
2132 for table in [
2133 "messages",
2134 "checkpoints",
2135 "thread_dynamic_tools",
2136 "thread_goals",
2137 ] {
2138 let sql = format!("SELECT COUNT(*) FROM {table} WHERE thread_id = ?1");
2139 let count: i64 = conn
2140 .query_row(&sql, params!["thread-1"], |row| row.get(0))
2141 .expect("count child rows");
2142 assert_eq!(count, 0, "{table} row survived thread deletion");
2143 }
2144 }
2145
2146 #[test]
2147 fn state_store_reuses_one_connection_across_operations_and_clones() {
2148 let store = temp_state_store("conn-reuse");
2149 {
2150 let conn = store.conn().expect("conn");
2151 conn.execute_batch("CREATE TEMP TABLE conn_reuse_probe(id INTEGER);")
2152 .expect("create temp table");
2153 }
2154 let clone = store.clone();
2159 clone
2160 .upsert_thread(&test_thread("thread-conn-reuse"))
2161 .expect("upsert thread");
2162 let conn = clone.conn().expect("conn");
2163 let probe_count: i64 = conn
2164 .query_row(
2165 "SELECT COUNT(*) FROM sqlite_temp_master WHERE name = 'conn_reuse_probe'",
2166 [],
2167 |row| row.get(0),
2168 )
2169 .expect("query temp master");
2170 assert_eq!(
2171 probe_count, 1,
2172 "temp table not visible: a fresh connection was opened"
2173 );
2174 let foreign_keys: i64 = conn
2176 .query_row("PRAGMA foreign_keys;", [], |row| row.get(0))
2177 .expect("read foreign_keys pragma");
2178 assert_eq!(foreign_keys, 1);
2179 let journal_mode: String = conn
2180 .query_row("PRAGMA journal_mode;", [], |row| row.get(0))
2181 .expect("read journal_mode pragma");
2182 assert_eq!(
2183 journal_mode.to_ascii_lowercase(),
2184 "wal",
2185 "open should enable WAL for multi-process readers/writers"
2186 );
2187 }
2188
2189 #[test]
2196 fn two_connections_can_write_concurrently_without_sqlite_busy() {
2197 let suffix = SystemTime::now()
2198 .duration_since(UNIX_EPOCH)
2199 .expect("system time")
2200 .as_nanos();
2201 let dir = std::env::temp_dir().join(format!(
2202 "codewhale-state-concurrent-write-{}-{suffix}",
2203 std::process::id()
2204 ));
2205 fs::create_dir_all(&dir).expect("create temp state dir");
2206 let db_path = dir.join("state.db");
2207
2208 let bootstrap = StateStore::open(Some(db_path.clone())).expect("bootstrap open");
2210 {
2211 let conn = bootstrap.conn().expect("bootstrap conn");
2212 let journal_mode: String = conn
2213 .query_row("PRAGMA journal_mode;", [], |row| row.get(0))
2214 .expect("journal_mode");
2215 assert_eq!(journal_mode.to_ascii_lowercase(), "wal");
2216 }
2217 drop(bootstrap);
2218
2219 let path_a = db_path.clone();
2220 let path_b = db_path.clone();
2221 const WRITES_PER_CONN: usize = 50;
2222
2223 let handle_a = std::thread::spawn(move || {
2224 let store = StateStore::open(Some(path_a)).expect("open store a");
2225 for i in 0..WRITES_PER_CONN {
2226 let job = JobStateRecord {
2227 id: format!("job-a-{i}"),
2228 name: format!("writer-a-{i}"),
2229 status: JobStateStatus::Running,
2230 progress: Some((i % 100) as u8),
2231 detail: Some("concurrent write a".to_string()),
2232 created_at: i as i64,
2233 updated_at: i as i64,
2234 };
2235 store.upsert_job(&job).unwrap_or_else(|err| {
2236 panic!("connection A write {i} failed (must not be SQLITE_BUSY): {err:#}");
2237 });
2238 }
2239 });
2240 let handle_b = std::thread::spawn(move || {
2241 let store = StateStore::open(Some(path_b)).expect("open store b");
2242 for i in 0..WRITES_PER_CONN {
2243 let job = JobStateRecord {
2244 id: format!("job-b-{i}"),
2245 name: format!("writer-b-{i}"),
2246 status: JobStateStatus::Running,
2247 progress: Some((i % 100) as u8),
2248 detail: Some("concurrent write b".to_string()),
2249 created_at: i as i64,
2250 updated_at: i as i64,
2251 };
2252 store.upsert_job(&job).unwrap_or_else(|err| {
2253 panic!("connection B write {i} failed (must not be SQLITE_BUSY): {err:#}");
2254 });
2255 }
2256 });
2257
2258 handle_a.join().expect("writer A panicked");
2259 handle_b.join().expect("writer B panicked");
2260
2261 let store = StateStore::open(Some(db_path)).expect("reopen for verify");
2262 let listed = store
2263 .list_jobs(Some(WRITES_PER_CONN * 2))
2264 .expect("list jobs");
2265 assert_eq!(
2266 listed.len(),
2267 WRITES_PER_CONN * 2,
2268 "both connections should have persisted all jobs"
2269 );
2270
2271 let _ = fs::remove_dir_all(dir);
2272 }
2273
2274 #[test]
2275 fn record_thread_goal_usage_accumulates_tokens_and_time() {
2276 let store = temp_state_store("thread-goal-usage");
2277 store
2278 .upsert_thread(&test_thread("thread-1"))
2279 .expect("upsert thread");
2280
2281 let mut goal = test_goal("thread-1", "Ship the persistent goal loop");
2283 goal.tokens_used = 0;
2284 goal.time_used_seconds = 0;
2285 goal.updated_at = 100;
2286 store.upsert_thread_goal(&goal).expect("upsert goal");
2287
2288 let after_first = store
2290 .record_thread_goal_usage("thread-1", 250, 12, 150)
2291 .expect("record usage")
2292 .expect("goal exists");
2293 assert_eq!(after_first.tokens_used, 250);
2294 assert_eq!(after_first.time_used_seconds, 12);
2295 assert_eq!(after_first.updated_at, 150);
2296 assert_eq!(after_first.goal_id, goal.goal_id);
2298 assert_eq!(after_first.objective, goal.objective);
2299 assert_eq!(after_first.status, goal.status);
2300 assert_eq!(after_first.token_budget, goal.token_budget);
2301 assert_eq!(after_first.created_at, goal.created_at);
2302 assert_eq!(after_first.continuation_count, 0);
2303
2304 let after_second = store
2306 .record_thread_goal_usage("thread-1", 75, 8, 200)
2307 .expect("record usage")
2308 .expect("goal exists");
2309 assert_eq!(after_second.tokens_used, 325);
2310 assert_eq!(after_second.time_used_seconds, 20);
2311 assert_eq!(after_second.updated_at, 200);
2312
2313 let after_stale = store
2315 .record_thread_goal_usage("thread-1", 5, 1, 1)
2316 .expect("record usage")
2317 .expect("goal exists");
2318 assert_eq!(after_stale.tokens_used, 330);
2319 assert_eq!(after_stale.time_used_seconds, 21);
2320 assert_eq!(after_stale.updated_at, 200);
2321
2322 let persisted = store
2324 .get_thread_goal("thread-1")
2325 .expect("read goal")
2326 .expect("goal exists");
2327 assert_eq!(persisted.tokens_used, 330);
2328 assert_eq!(persisted.time_used_seconds, 21);
2329 }
2330
2331 #[test]
2332 fn record_thread_goal_usage_returns_none_without_goal() {
2333 let store = temp_state_store("thread-goal-usage-missing");
2334 store
2335 .upsert_thread(&test_thread("thread-1"))
2336 .expect("upsert thread");
2337 let result = store
2340 .record_thread_goal_usage("thread-1", 100, 5, 999)
2341 .expect("record usage on goalless thread");
2342 assert!(result.is_none());
2343 assert!(
2344 store
2345 .get_thread_goal("thread-1")
2346 .expect("read goal")
2347 .is_none()
2348 );
2349 }
2350
2351 #[test]
2352 fn record_thread_goal_continuation_accumulates_durably() {
2353 let store = temp_state_store("thread-goal-continuation");
2354 store
2355 .upsert_thread(&test_thread("thread-1"))
2356 .expect("upsert thread");
2357
2358 let mut goal = test_goal("thread-1", "Keep working across turns");
2359 goal.updated_at = 100;
2360 store.upsert_thread_goal(&goal).expect("upsert goal");
2361
2362 let after_first = store
2363 .record_thread_goal_continuation("thread-1", 120)
2364 .expect("record continuation")
2365 .expect("goal exists");
2366 assert_eq!(after_first.continuation_count, 1);
2367 assert_eq!(after_first.tokens_used, goal.tokens_used);
2368 assert_eq!(after_first.time_used_seconds, goal.time_used_seconds);
2369 assert_eq!(after_first.updated_at, 120);
2370
2371 let after_second = store
2372 .record_thread_goal_continuation("thread-1", 110)
2373 .expect("record second continuation")
2374 .expect("goal exists");
2375 assert_eq!(after_second.continuation_count, 2);
2376 assert_eq!(after_second.updated_at, 120);
2377
2378 let persisted = store
2379 .get_thread_goal("thread-1")
2380 .expect("read goal")
2381 .expect("goal exists");
2382 assert_eq!(persisted.continuation_count, 2);
2383 }
2384
2385 static CODEWHALE_HOME_TEST_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
2392
2393 struct CodeWhaleHomeGuard {
2394 prior: Option<std::ffi::OsString>,
2395 }
2396 impl CodeWhaleHomeGuard {
2397 fn set(value: &str) -> Self {
2398 let prior = std::env::var_os("CODEWHALE_HOME");
2399 unsafe { std::env::set_var("CODEWHALE_HOME", value) };
2401 Self { prior }
2402 }
2403 fn remove() -> Self {
2404 let prior = std::env::var_os("CODEWHALE_HOME");
2405 unsafe { std::env::remove_var("CODEWHALE_HOME") };
2407 Self { prior }
2408 }
2409 }
2410 impl Drop for CodeWhaleHomeGuard {
2411 fn drop(&mut self) {
2412 unsafe {
2414 match &self.prior {
2415 Some(value) => std::env::set_var("CODEWHALE_HOME", value),
2416 None => std::env::remove_var("CODEWHALE_HOME"),
2417 }
2418 }
2419 }
2420 }
2421
2422 #[test]
2423 fn codewhale_home_override_returns_the_env_value_verbatim() {
2424 let _lock = CODEWHALE_HOME_TEST_LOCK.lock().unwrap();
2425 let _g = CodeWhaleHomeGuard::set("/tmp/cw-isolated-state");
2426 assert_eq!(
2429 codewhale_home_override().as_deref(),
2430 Some(std::path::Path::new("/tmp/cw-isolated-state"))
2431 );
2432 }
2433
2434 #[test]
2435 fn codewhale_home_override_none_when_unset() {
2436 let _lock = CODEWHALE_HOME_TEST_LOCK.lock().unwrap();
2437 let _g = CodeWhaleHomeGuard::remove();
2438 assert!(codewhale_home_override().is_none());
2439 }
2440
2441 #[test]
2442 fn codewhale_home_override_none_when_empty() {
2443 let _lock = CODEWHALE_HOME_TEST_LOCK.lock().unwrap();
2444 let _g = CodeWhaleHomeGuard::set(" ");
2445 assert!(
2451 codewhale_home_override().is_some(),
2452 "non-empty (even whitespace) counts as set; trimming is the caller's job"
2453 );
2454 }
2455
2456 #[test]
2457 fn default_state_db_path_uses_codewhale_home_when_set() {
2458 let _lock = CODEWHALE_HOME_TEST_LOCK.lock().unwrap();
2459 let dir = std::env::temp_dir().join(format!(
2460 "cw-home-state-{}-{}",
2461 std::process::id(),
2462 std::time::SystemTime::now()
2463 .duration_since(std::time::UNIX_EPOCH)
2464 .unwrap()
2465 .as_nanos()
2466 ));
2467 let _g = CodeWhaleHomeGuard::set(dir.to_str().unwrap());
2468 assert_eq!(default_state_db_path(), dir.join("state.db"));
2472 }
2473
2474 #[test]
2475 fn load_checkpoint_propagates_invalid_state_json() {
2476 let store = temp_state_store("checkpoint-parse-error");
2477 store
2478 .upsert_thread(&test_thread("thread-1"))
2479 .expect("upsert thread");
2480 store
2481 .save_checkpoint("thread-1", "broken", &json!({"ok": true}))
2482 .expect("save checkpoint");
2483
2484 {
2485 let conn = store.conn().expect("conn");
2486 conn.execute(
2487 "UPDATE checkpoints SET state_json = ?1 WHERE thread_id = ?2 AND checkpoint_id = ?3",
2488 params!["not-json", "thread-1", "broken"],
2489 )
2490 .expect("corrupt checkpoint");
2491 }
2492
2493 let err = store
2494 .load_checkpoint("thread-1", Some("broken"))
2495 .expect_err("invalid checkpoint json should fail");
2496 assert!(
2497 err.to_string()
2498 .contains("failed to parse checkpoint state json")
2499 );
2500 }
2501
2502 #[test]
2503 fn session_index_compacts_after_threshold() {
2504 let store = temp_state_store("session-index-compact");
2505 for idx in 0..6 {
2506 store
2507 .append_thread_name("thread-1", Some(format!("name-{idx}")), idx, None)
2508 .expect("append session index entry");
2509 }
2510
2511 let line_count = store
2512 .session_index_line_count()
2513 .expect("count session index lines");
2514 assert_eq!(line_count, 1);
2515
2516 let name = store
2517 .find_thread_name_by_id("thread-1")
2518 .expect("lookup thread name");
2519 assert_eq!(name.as_deref(), Some("name-5"));
2520 }
2521
2522 #[test]
2523 fn session_index_read_skips_a_torn_line() {
2524 let store = temp_state_store("session-index-torn");
2528 store
2529 .append_thread_name("thread-1", Some("first".to_string()), 1, None)
2530 .expect("append first entry");
2531
2532 {
2533 let mut file = OpenOptions::new()
2534 .append(true)
2535 .open(&store.session_index_path)
2536 .expect("open session index");
2537 writeln!(file, "{{\"thread_id\":\"thread-2\",\"thread_na").expect("write torn line");
2539 }
2540
2541 store
2542 .append_thread_name("thread-3", Some("third".to_string()), 3, None)
2543 .expect("append third entry");
2544
2545 assert_eq!(
2546 store
2547 .find_thread_name_by_id("thread-1")
2548 .expect("lookup thread-1")
2549 .as_deref(),
2550 Some("first"),
2551 );
2552 assert_eq!(
2553 store
2554 .find_thread_name_by_id("thread-3")
2555 .expect("lookup thread-3")
2556 .as_deref(),
2557 Some("third"),
2558 );
2559 }
2560
2561 type MidpointHook = Box<dyn Fn() + Send + Sync>;
2566 static COMPACTION_MIDPOINT: Mutex<Option<(PathBuf, MidpointHook)>> = Mutex::new(None);
2567
2568 pub(super) fn compaction_midpoint(index_path: &Path) {
2571 let mut hook = COMPACTION_MIDPOINT.lock().expect("midpoint hook lock");
2572 let registered_for_this_store = match hook.as_ref() {
2573 Some((registered, _)) => registered == index_path,
2574 None => return,
2575 };
2576 if !registered_for_this_store {
2577 return;
2578 }
2579 let (_, callback) = hook.take().expect("presence checked above");
2580 drop(hook);
2581 callback();
2582 }
2583
2584 #[test]
2585 fn session_index_compaction_does_not_drop_a_concurrent_append() {
2586 let store = Arc::new(temp_state_store("session-index-race"));
2597 let threshold = session_index_compact_line_threshold();
2598
2599 for idx in 0..threshold {
2601 store
2602 .append_thread_name(
2603 &format!("thread-{idx}"),
2604 Some(format!("name-{idx}")),
2605 1,
2606 None,
2607 )
2608 .expect("append filler entry");
2609 }
2610
2611 let appender_released = Arc::new(Barrier::new(2));
2612 {
2613 let released = Arc::clone(&appender_released);
2614 *COMPACTION_MIDPOINT.lock().expect("midpoint hook lock") = Some((
2615 store.session_index_path.clone(),
2616 Box::new(move || {
2617 released.wait();
2618 thread::sleep(Duration::from_millis(300));
2621 }),
2622 ));
2623 }
2624
2625 let appender = {
2626 let store = Arc::clone(&store);
2627 let released = Arc::clone(&appender_released);
2628 thread::spawn(move || {
2629 released.wait();
2630 store
2631 .append_thread_name("racer", Some("racer-name".to_string()), 2, None)
2632 .expect("append racing entry");
2633 })
2634 };
2635
2636 store
2637 .append_thread_name("trigger", Some("trigger-name".to_string()), 1, None)
2638 .expect("append entry that triggers compaction");
2639 appender.join().expect("appender thread");
2640 *COMPACTION_MIDPOINT.lock().expect("midpoint hook lock") = None;
2641
2642 assert_eq!(
2643 store
2644 .find_thread_name_by_id("racer")
2645 .expect("lookup racer")
2646 .as_deref(),
2647 Some("racer-name"),
2648 "append was dropped by a concurrent compaction",
2649 );
2650 assert_eq!(
2651 store
2652 .find_thread_name_by_id("trigger")
2653 .expect("lookup trigger")
2654 .as_deref(),
2655 Some("trigger-name"),
2656 );
2657 }
2658}