Skip to main content

kaizen/collect/tail/
cursor_state_db.rs

1// SPDX-License-Identifier: AGPL-3.0-or-later
2//! Read-only Cursor/VS Code `state.vscdb` helpers.
3
4use super::cursor_state_db_fields::{
5    file_mtime_ms, id_field, key_suffix, text, ts_any, ts_field, workspace_field, workspace_matches,
6};
7use crate::core::event::{Event, EventKind, EventSource, SessionRecord, SessionStatus};
8use anyhow::{Context, Result};
9use rusqlite::types::ValueRef;
10use rusqlite::{Connection, OpenFlags};
11use serde_json::{Value, json};
12use std::path::Path;
13use std::path::PathBuf;
14
15#[derive(Debug, Clone, PartialEq, Eq)]
16pub struct CursorStateItem {
17    pub key: String,
18    pub value: String,
19}
20
21pub fn read_items_with_prefix(db_path: &Path, prefix: &str) -> Result<Vec<CursorStateItem>> {
22    let conn = open_read_only(db_path)?;
23    if !has_item_table(&conn)? {
24        return Ok(Vec::new());
25    }
26    let like = format!("{}%", escape_like(prefix));
27    let mut stmt = conn
28        .prepare("SELECT key, value FROM ItemTable WHERE key LIKE ?1 ESCAPE '\\' ORDER BY key")?;
29    let rows = stmt.query_map([like], row_item)?;
30    rows.map(|r| r.map_err(anyhow::Error::from)).collect()
31}
32
33pub fn scan_cursor_state_db_workspace(workspace: &Path) -> Vec<(SessionRecord, Vec<Event>)> {
34    db_paths(workspace)
35        .into_iter()
36        .flat_map(|path| scan_db(&path, workspace).unwrap_or_default())
37        .collect()
38}
39
40fn open_read_only(path: &Path) -> Result<Connection> {
41    Connection::open_with_flags(
42        path,
43        OpenFlags::SQLITE_OPEN_READ_ONLY | OpenFlags::SQLITE_OPEN_NO_MUTEX,
44    )
45    .with_context(|| format!("open cursor state db read-only: {}", path.display()))
46}
47
48fn db_paths(workspace: &Path) -> Vec<PathBuf> {
49    [
50        std::env::var("CURSOR_STATE_DB").ok().map(PathBuf::from),
51        Some(workspace.join(".cursor/state.vscdb")),
52        Some(workspace.join("state.vscdb")),
53    ]
54    .into_iter()
55    .flatten()
56    .filter(|path| path.is_file())
57    .collect()
58}
59
60fn scan_db(path: &Path, workspace: &Path) -> Result<Vec<(SessionRecord, Vec<Event>)>> {
61    Ok(read_items_with_prefix(path, "composerData:")?
62        .into_iter()
63        .filter_map(|item| session_from_item(item, path, workspace))
64        .collect())
65}
66
67fn session_from_item(
68    item: CursorStateItem,
69    path: &Path,
70    workspace: &Path,
71) -> Option<(SessionRecord, Vec<Event>)> {
72    let value: Value = serde_json::from_str(&item.value).ok()?;
73    let root = workspace_field(&value).unwrap_or_else(|| workspace.to_string_lossy().into());
74    workspace_matches(&root, workspace).then(|| record_and_events(item, value, path, root))
75}
76
77fn record_and_events(
78    item: CursorStateItem,
79    value: Value,
80    path: &Path,
81    workspace: String,
82) -> (SessionRecord, Vec<Event>) {
83    let id = id_field(&value).unwrap_or_else(|| key_suffix(&item.key));
84    let started = ts_field(&value).unwrap_or_else(|| file_mtime_ms(path));
85    let events = vec![event(&id, started, item.key, value.clone())];
86    (record(id, value, path, workspace, started), events)
87}
88
89fn record(
90    id: String,
91    value: Value,
92    path: &Path,
93    workspace: String,
94    started_at_ms: u64,
95) -> SessionRecord {
96    SessionRecord {
97        id,
98        agent: "cursor".into(),
99        model: text(&value, "model"),
100        workspace,
101        started_at_ms,
102        ended_at_ms: ts_any(&value, &["ended_at_ms", "updated_at_ms"]),
103        status: SessionStatus::Done,
104        trace_path: path.to_string_lossy().into(),
105        start_commit: None,
106        end_commit: None,
107        branch: None,
108        dirty_start: None,
109        dirty_end: None,
110        repo_binding_source: None,
111        prompt_fingerprint: None,
112        parent_session_id: None,
113        agent_version: None,
114        os: None,
115        arch: None,
116        repo_file_count: None,
117        repo_total_loc: None,
118    }
119}
120
121fn event(session_id: &str, ts_ms: u64, key: String, value: Value) -> Event {
122    Event {
123        session_id: session_id.into(),
124        seq: 0,
125        ts_ms,
126        ts_exact: ts_field(&value).is_some(),
127        kind: EventKind::Message,
128        source: EventSource::Tail,
129        tool: None,
130        tool_call_id: None,
131        tokens_in: None,
132        tokens_out: None,
133        reasoning_tokens: None,
134        cost_usd_e6: None,
135        stop_reason: None,
136        latency_ms: None,
137        ttft_ms: None,
138        retry_count: None,
139        context_used_tokens: None,
140        context_max_tokens: None,
141        cache_creation_tokens: None,
142        cache_read_tokens: None,
143        system_prompt_tokens: None,
144        payload: json!({"cursor_state_key": key, "value": value}),
145    }
146}
147
148fn has_item_table(conn: &Connection) -> Result<bool> {
149    let n: i64 = conn.query_row(
150        "SELECT COUNT(*) FROM sqlite_master WHERE type='table' AND name='ItemTable'",
151        [],
152        |row| row.get(0),
153    )?;
154    Ok(n > 0)
155}
156
157fn row_item(row: &rusqlite::Row<'_>) -> rusqlite::Result<CursorStateItem> {
158    Ok(CursorStateItem {
159        key: row.get(0)?,
160        value: value_text(row.get_ref(1)?),
161    })
162}
163
164fn value_text(value: ValueRef<'_>) -> String {
165    match value {
166        ValueRef::Null => String::new(),
167        ValueRef::Integer(v) => v.to_string(),
168        ValueRef::Real(v) => v.to_string(),
169        ValueRef::Text(v) | ValueRef::Blob(v) => String::from_utf8_lossy(v).into_owned(),
170    }
171}
172
173fn escape_like(prefix: &str) -> String {
174    prefix
175        .replace('\\', "\\\\")
176        .replace('%', "\\%")
177        .replace('_', "\\_")
178}