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