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