atuinsh/atuin
1
0
Fork 0
mirror of https://github.com/atuinsh/atuin.git synced 2026-10-12 02:11:09 +00:00
Code Issues Packages Projects Releases Wiki Activity

refactor(ai-session): squash sidecar migrations, intern repeated message text, drop session transcript (atuinsh/atuin#4371)

Browse Source
This commit is contained in:
Ellie Huxtable
2026-10-07 06:11:43 +00:00
committed by Ellie Huxtable
parent 1fccfd438c
commit 6079ac2357
20 changed files with 383 additions and 1199 deletions
Download Patch File Download Diff File Expand all files Collapse all files
AGENTS.md
+1
View File
@@ -34,6 +34,7 @@ atuin-server-sqlite SQLite implementation (sqlx)
- **Client**: SQLite everywhere. Separate DBs for history, record store, KV, scripts. All use sqlx + WAL mode.
- **Server**: Postgres (primary) or SQLite. Auto-detected from URI prefix.
- Migrations live alongside each crate. Never modify existing migrations, only add new ones.
- Exception: the AI session sidecar's migrations (`atuin-client/src/ai_session/migrations`) were squashed into one `0001_init.sql` before release, and the sidecar file was renamed (`ai_harness_sessions.db` → `ai_session_sidecar.db`) so no old sidecar meets the new history. Don't add back compatibility with the old migrations; from here on, add new ones as usual.
## Hot paths
crates/atuin-ai/src/commands/session.rs
-45
View File
@@ -47,13 +47,6 @@ enum SubCmd {
session: String,
},
/// Print a session's rendered transcript. Accepts a session id or `latest`.
Transcript {
/// Session id (a unique prefix is enough), or `latest` for the most recently active session.
#[arg(value_name = "ID|latest")]
session: String,
},
#[command(about = "Full-text search across captured sessions, most relevant first.")]
Search {
#[arg(
@@ -167,7 +160,6 @@ pub async fn run(cmd: Cmd, settings: &Settings) -> Result<()> {
let result = match cmd.cmd {
SubCmd::List => list(&mut client, style).await,
SubCmd::Show { session } => show(&mut client, &session, style).await,
SubCmd::Transcript { session } => transcript(&mut client, &session, style).await,
SubCmd::Search {
query,
harness,
@@ -285,36 +277,6 @@ async fn show(client: &mut AiClient, selector: &str, style: Style) -> Result<()>
Ok(())
}
async fn transcript(client: &mut AiClient, selector: &str, style: Style) -> Result<()> {
let session = resolve(client, selector).await?;
let mut stream = client.get_transcript(session.clone()).await?;
let mut text = String::new();
while let Some(chunk) = stream.next().await {
text.push_str(&chunk?.chunk);
}
let stdout = io::stdout();
let mut out = stdout.lock();
if style.is_json() {
let record = TranscriptJson {
harness: harness_name(session.harness).to_owned(),
session_id: session.session.into(),
transcript: text,
};
serde_json::to_writer(&mut out, &record)?;
writeln!(out)?;
} else {
write!(out, "{}", sanitize(&text))?;
if !text.ends_with('\n') {
writeln!(out)?;
}
}
Ok(())
}
async fn search(
client: &mut AiClient,
query: &str,
@@ -1146,13 +1108,6 @@ enum TailEventJson {
},
}
#[derive(Serialize)]
struct TranscriptJson {
harness: String,
session_id: String,
transcript: String,
}
#[derive(Serialize)]
struct HighlightJson {
text: String,
crates/atuin-client/src/ai_session/database.rs
+78 -679
View File
@@ -14,7 +14,7 @@ use atuin_common::harnesstools::session::{
use atuin_common::string::TruncateCharsExt;
use atuin_common::string::highlighted::HighlightedString;
use atuin_domain::record::{HostId, RecordId};
use futures::{Stream, StreamExt, TryStreamExt};
use futures::{Stream, TryStreamExt};
use sqlx::migrate::MigrateError;
use time::OffsetDateTime;
use tracing::warn;
@@ -31,16 +31,15 @@ mod watermark;
pub use heads::{Analysis, FastForward, Head};
pub use watermark::{Generation, Watermark};
pub use write::PreparedMessage;
use write::Writer;
use write::{Interned, Writer};
const COMPRESS_THRESHOLD: usize = 256;
const ZSTD_LEVEL: i32 = 3;
const REINDEX_CHUNK: i64 = 512;
const SNIPPET_TOKENS: usize = 32;
/// The newest migration this build knows: [`AiSessionDatabase::open_read_only`] reads only a
/// sidecar at exactly this version.
const SCHEMA_VERSION: i64 = 8;
const SCHEMA_VERSION: i64 = 1;
/// How much more a title match weighs than a body match in bm25.
const TITLE_WEIGHT: f64 = 5.0;
@@ -83,10 +82,13 @@ macro_rules! message_columns {
() => {
"m.id, m.harness, (SELECT session_id FROM sessions WHERE id = m.session) AS session_id, \
(SELECT parent_kind FROM sessions WHERE id = m.session) AS parent_kind, m.source_id, \
m.parent_harness, m.parent_session_id, m.parent_source_id, m.timestamp, m.role, \
m.content, m.content_z, m.cwd, m.git_branch, m.model, m.usage_input, m.usage_output, \
m.usage_cache_read, m.usage_cache_write, m.usage_reasoning, m.stop_reason, \
m.usage_present, m.turn_id, m.title_change, m.host_id"
m.parent_harness, m.parent_session_id, m.parent_source_id, m.timestamp, (SELECT value \
FROM interned WHERE id = m.role) AS role, m.content, m.content_z, (SELECT value FROM \
interned WHERE id = m.cwd) AS cwd, (SELECT value FROM interned WHERE id = m.git_branch) \
AS git_branch, (SELECT value FROM interned WHERE id = m.model) AS model, m.usage_input, \
m.usage_output, m.usage_cache_read, m.usage_cache_write, m.usage_reasoning, \
m.stop_reason, m.usage_present, m.turn_id, m.title_change, (SELECT value FROM interned \
WHERE id = m.host) AS host_id"
};
}
@@ -278,7 +280,7 @@ struct SessionRow {
root_harness: Option<i64>,
root_session_id: Option<String>,
copy_of_session_id: Option<String>,
atuin_id: Option<Vec<u8>>,
atuin_id: Vec<u8>,
parent_atuin_id: Option<Vec<u8>>,
root_atuin_id: Option<Vec<u8>>,
child_atuin_ids: Option<String>,
@@ -403,14 +405,6 @@ struct SearchRow {
score: f64,
}
#[derive(sqlx::FromRow)]
struct ReindexRow {
rowid: i64,
#[sqlx(flatten)]
message: MessageRow,
session_title: Option<String>,
}
impl AiSessionDatabase {
/// Open the sidecar at `path` to write, creating and migrating it as needed.
///
@@ -433,12 +427,11 @@ impl AiSessionDatabase {
},
err => err,
})?;
db.reindex().await?;
Ok(db)
}
/// Open the sidecar to read beside the daemon that writes it (WAL keeps the two apart): no
/// migration or reindex, and no writes. Fails when the file is missing, and with
/// migration, and no writes. Fails when the file is missing, and with
/// [`DbError::Uninitialized`], [`DbError::OutdatedSchema`] or [`DbError::UnknownSchema`]
/// unless it is at exactly the schema this build reads.
pub async fn open_read_only(path: impl AsRef<Path>) -> Result<Self, DbError> {
@@ -518,33 +511,6 @@ impl AiSessionDatabase {
async fn migrate(&self) -> Result<(), DbError> {
let pool = self.db.pool();
db::migrate!(pool, "./src/ai_session/migrations").await?;
self.write(Self::group_migrated).await?;
self.write(Self::recount_migrated).await?;
self.atuin_id_migrated().await
}
/// Derive ids for sessions stored before the `atuin_id` migration.
async fn atuin_id_migrated(&self) -> Result<(), DbError> {
let mut tx = self.db.pool().begin_with("BEGIN IMMEDIATE").await?;
let unset: Vec<(i64, i64, String, i64)> = db::query_as(
"SELECT id, harness, session_id, started_at FROM sessions WHERE atuin_id IS NULL",
)
.fetch_all(&mut *tx)
.await?;
for (id, harness, session_id, started_at) in unset {
let handle = HarnessSession {
harness: Self::harness_from_repr(harness)?,
session: NativeSessionId::from(session_id),
};
let atuin_id = handle.atuin_id(Self::time_from_millis(started_at)?);
db::query("UPDATE sessions SET atuin_id = ? WHERE id = ?")
.bind(atuin_id.as_bytes().as_slice())
.bind(id)
.execute(&mut *tx)
.await?;
}
tx.commit().await?;
Ok(())
}
@@ -557,9 +523,18 @@ impl AiSessionDatabase {
self.writer.as_ref().ok_or(DbError::ReadOnly)?.run(write).await
}
/// Run `write` on the writer with its interned ids (see [`Writer::run_interning`]).
async fn write_interning<T, F>(&self, write: F) -> Result<T, DbError>
where
T: Send + 'static,
F: FnOnce(&mut rusqlite::Connection, &mut Interned) -> Result<T, DbError> + Send + 'static,
{
self.writer.as_ref().ok_or(DbError::ReadOnly)?.run_interning(write).await
}
pub async fn append(&self, msg: &Message) -> Result<Appended, DbError> {
let msg = PreparedMessage::new(msg.clone())?;
self.write(move |conn| Self::append_one(conn, &msg)).await
self.write_interning(move |conn, interned| Self::append_one(conn, interned, &msg)).await
}
/// [`Self::append`] each of `msgs` in order, in one transaction: all of them or, on an error,
@@ -577,7 +552,10 @@ impl AiSessionDatabase {
if msgs.is_empty() {
return Ok(true);
}
self.write(move |conn| Self::append_page(conn, &msgs, generation)).await
self.write_interning(move |conn, interned| {
Self::append_page(conn, interned, &msgs, generation)
})
.await
}
pub async fn contains_message(
@@ -634,13 +612,12 @@ impl AiSessionDatabase {
Ok(ids.into_iter().map(SourceId::from).collect())
}
/// Where capture resumes a session, if it checkpointed one with a digest; one without a
/// digest reads as `None`, so the session is read again from its start.
/// Where capture resumes a session, if it checkpointed one.
pub async fn checkpoint(
&self,
session: &HarnessSession,
) -> Result<Option<Checkpoint>, DbError> {
let row: Option<(i64, Option<i64>)> = db::query_as(
let row: Option<(i64, i64)> = db::query_as(
"SELECT \"offset\", digest FROM checkpoints WHERE harness = ? AND session_id = ?",
)
.bind(session.harness as i64)
@@ -648,12 +625,10 @@ impl AiSessionDatabase {
.fetch_optional(self.db.pool())
.await?;
Ok(row.and_then(|(at, digest)| {
Some(Checkpoint {
at: u64::try_from(at).unwrap_or(0),
digest: digest?.cast_unsigned(),
generation: 0,
})
Ok(row.map(|(at, digest)| Checkpoint {
at: u64::try_from(at).unwrap_or(0),
digest: digest.cast_unsigned(),
generation: 0,
}))
}
@@ -798,15 +773,15 @@ impl AiSessionDatabase {
let Some(id) = id else {
return Ok(parts);
};
// Roles are stored as their JSON, so a role is matched as that text: SQLite walks the
// session's messages in time order (`messages_session_timestamp`) and stops at the first
// match, without the roles of the rest ever reaching here.
// Roles are stored as their JSON's `interned` id, so a role is matched by that: SQLite walks
// the session's messages in time order (`messages_session_timestamp`) and stops at the
// first match, without the roles of the rest ever reaching here.
let user = serde_json::to_string(&Role::User)?;
let assistant = serde_json::to_string(&Role::Assistant)?;
let first: Option<(String, Option<Vec<u8>>)> = db::query_as(
"SELECT content, content_z FROM messages WHERE session = ? AND role = ? ORDER BY \
timestamp, id LIMIT 1",
"SELECT content, content_z FROM messages WHERE session = ? AND role = (SELECT id FROM \
interned WHERE value = ?) ORDER BY timestamp, id LIMIT 1",
)
.bind(id)
.bind(&user)
@@ -826,7 +801,8 @@ impl AiSessionDatabase {
None => {
db::query_as(
"SELECT timestamp, id, content, content_z FROM messages WHERE session = ? \
AND role = ? ORDER BY timestamp DESC, id DESC LIMIT 1",
AND role = (SELECT id FROM interned WHERE value = ?) ORDER BY timestamp \
DESC, id DESC LIMIT 1",
)
.bind(id)
.bind(&assistant)
@@ -836,8 +812,8 @@ impl AiSessionDatabase {
Some((timestamp, message)) => {
db::query_as(
"SELECT timestamp, id, content, content_z FROM messages WHERE session = ? \
AND role = ? AND (timestamp < ? OR (timestamp = ? AND id < ?)) ORDER BY \
timestamp DESC, id DESC LIMIT 1",
AND role = (SELECT id FROM interned WHERE value = ?) AND (timestamp < ? \
OR (timestamp = ? AND id < ?)) ORDER BY timestamp DESC, id DESC LIMIT 1",
)
.bind(id)
.bind(&assistant)
@@ -981,13 +957,6 @@ impl AiSessionDatabase {
}
}
pub fn transcript(
&self,
session: &HarnessSession,
) -> impl Stream<Item = Result<String, DbError>> + Send + 'static {
self.messages(session).map(|result| result.map(|msg| Self::render_transcript_chunk(&msg)))
}
/// Sessions matching `query` (its terms matching as `terms` says) and passing `filter`, most
/// relevant first, at most `limit` (0 is unbounded).
///
@@ -1143,70 +1112,6 @@ impl AiSessionDatabase {
}
}
pub async fn reindex(&self) -> Result<(), DbError> {
let pool = self.db.pool();
// messages_fts rowids are always a contiguous prefix of messages rowids: append writes the
// message and its FTS row in one tx, and this backfill (the only other writer, running
// under open() before the daemon serves) walks rowids ascending. So the highest indexed
// rowid alone determines coverage — gate on max(rowid) (O(log N)) rather than counting
// both tables. This also repopulates the index from scratch after a migration rebuilds
// messages_fts.
let mut watermark: i64 =
db::query_scalar("SELECT coalesce(max(rowid), 0) FROM messages_fts")
.fetch_one(pool)
.await?;
let last_message: i64 = db::query_scalar("SELECT coalesce(max(rowid), 0) FROM messages")
.fetch_one(pool)
.await?;
if watermark >= last_message {
return Ok(());
}
loop {
let mut tx = pool.begin_with("BEGIN IMMEDIATE").await?;
let rows: Vec<ReindexRow> = db::query_as::<_, ReindexRow>(
"SELECT m.rowid AS rowid, m.id, m.harness, s.session_id, m.source_id, \
m.parent_harness, m.parent_session_id, s.parent_kind, m.parent_source_id, \
m.timestamp, m.role, m.content, m.content_z, m.cwd, m.git_branch, m.model, \
m.usage_input, m.usage_output, m.usage_cache_read, m.usage_cache_write, \
m.usage_reasoning, m.stop_reason, m.usage_present, m.turn_id, m.title_change, \
m.host_id, s.title AS session_title FROM messages m JOIN sessions s ON s.id = \
m.session WHERE m.rowid > ? ORDER BY m.rowid LIMIT ?",
)
.bind(watermark)
.bind(REINDEX_CHUNK)
.fetch_all(&mut *tx)
.await?;
let Some(last) = rows.last() else {
break;
};
watermark = last.rowid;
for row in rows {
let rowid = row.rowid;
let title = row.session_title.unwrap_or_default();
let body = match Self::message_from_row(row.message) {
Ok(msg) => Self::searchable_body(&msg),
Err(err) => {
warn!(?err, rowid, "failed to decode ai-session message; indexing empty");
String::new()
}
};
db::query(
"INSERT OR REPLACE INTO messages_fts(rowid, title, body) VALUES (?, ?, ?)",
)
.bind(rowid)
.bind(&title)
.bind(&body)
.execute(&mut *tx)
.await?;
}
tx.commit().await?;
}
Ok(())
}
/// The index of the first whitespace word where a query term matches the way the FTS index
/// matched it (see [`QueryTerms`]).
fn preview_hit(words: &[&str], terms: &QueryTerms) -> Option<usize> {
@@ -1576,35 +1481,6 @@ impl AiSessionDatabase {
.build())
}
fn render_transcript_chunk(message: &Message) -> String {
let role = match &message.role {
Role::User => "user",
Role::Assistant => "assistant",
Role::System => "system",
Role::Tool => "tool",
Role::Other(other) => other.as_str(),
};
let body = message
.content
.iter()
.filter_map(|content| match content {
Content::Text(text) | Content::Reasoning(text) => Some(text.clone()),
Content::ReasoningSummary { tokens } => {
Some(atuin_common::harnesstools::session::model::reasoning_label(
tokens.or(message.usage.and_then(|u| u.reasoning)),
))
}
_ => None,
})
.collect::<Vec<_>>()
.join("\n");
// Trailing newline: chunks are concatenated verbatim by consumers, so the separator has
// to live in the chunk or every message would run together on one line.
format!("{role}: {body}\n")
}
const fn parent_kind_repr(kind: ParentKind) -> i64 {
match kind {
ParentKind::Subagent => 0,
@@ -1656,11 +1532,7 @@ impl AiSessionDatabase {
session: NativeSessionId::from(row.session_id),
};
let started_at = Self::time_from_millis(row.started_at)?;
let atuin_id = match row.atuin_id.as_deref() {
Some(bytes) => Uuid::from_slice(bytes).map_err(|_| DbError::InvalidAtuinId)?.into(),
// A reader can open the sidecar between the migration and its backfill.
None => handle.atuin_id(started_at),
};
let atuin_id = Uuid::from_slice(&row.atuin_id).map_err(|_| DbError::InvalidAtuinId)?.into();
// A link whose session is not stored yet (or has no stored id) resolves to nothing.
let linked = |bytes: Option<&[u8]>| {
bytes.and_then(|bytes| Uuid::from_slice(bytes).ok()).map(AtuinSessionId::from)
@@ -1957,52 +1829,6 @@ mod tests {
assert_eq!(got[0].turn_id.as_deref(), Some("msg_01"));
}
/// Upgrading keeps each message's rowid, which its contentless messages_fts row is keyed by.
#[rstest]
#[tokio::test]
async fn interning_sessions_keeps_the_search_index_aligned() {
let sqlite = atuin_common::db::sqlite::Sqlite::builder_in_memory().open().await.unwrap();
let pool = sqlite.pool();
sqlx::raw_sql(include_str!("migrations/0001_init.sql")).execute(pool).await.unwrap();
sqlx::raw_sql(
"INSERT INTO sessions (harness, session_id, started_at, updated_at) VALUES (1, 's', \
0, 0);
INSERT INTO messages (rowid, id, harness, session_id, source_id, timestamp, role, \
content)
VALUES (7, x'01', 1, 's', 'a', 0, '\"User\"', '[{\"Text\":\"hello\"}]');
INSERT INTO messages_fts (rowid, title, body) VALUES (7, '', 'hello');",
)
.execute(pool)
.await
.unwrap();
sqlx::raw_sql(include_str!("migrations/0002_intern_sessions.sql"))
.execute(pool)
.await
.unwrap();
for migration in [
include_str!("migrations/0003_last_reply.sql"),
include_str!("migrations/0004_parent_kind.sql"),
include_str!("migrations/0005_sessions_updated_at.sql"),
include_str!("migrations/0006_incremental_sidecar.sql"),
include_str!("migrations/0007_session_scoped_calls.sql"),
include_str!("migrations/0008_atuin_id.sql"),
] {
sqlx::raw_sql(migration).execute(pool).await.unwrap();
}
// Read-only but for the migration it runs, which goes through the pool.
let db = AiSessionDatabase::from_sqlite(sqlite, None);
db.atuin_id_migrated().await.unwrap();
let hits: Vec<SessionMatch> = db
.search("hello", SearchTerms::Typed, &SessionFilter::default(), 0)
.try_collect()
.await
.unwrap();
assert_eq!(hits.len(), 1);
assert_eq!(hits[0].session.handle.session.as_ref(), "s");
}
/// Rows that carry only usage, a title or session context are stored, but are no messages.
#[rstest]
#[tokio::test]
@@ -2053,20 +1879,6 @@ mod tests {
assert_eq!(db.checkpoint(&session).await.unwrap(), Some(latest));
}
#[rstest]
#[tokio::test]
async fn a_checkpoint_without_a_digest_reads_as_none() {
let db = AiSessionDatabase::in_memory().await.unwrap();
let session = sample_handle();
db::query("INSERT INTO checkpoints (harness, session_id, \"offset\") VALUES (?, ?, 42)")
.bind(session.harness as i64)
.bind(session.session.as_ref())
.execute(db.db.pool())
.await
.unwrap();
assert_eq!(db.checkpoint(&session).await.unwrap(), None);
}
#[rstest]
#[tokio::test]
async fn last_message_is_the_newest_by_timestamp() {
@@ -2420,20 +2232,6 @@ mod tests {
]);
}
#[rstest]
#[tokio::test]
async fn transcript_separates_messages() {
let db = AiSessionDatabase::in_memory().await.unwrap();
let session = sample_handle();
for m in ordered_messages(&session) {
db.append(&m).await.unwrap();
}
let chunks: Vec<String> = db.transcript(&session).try_collect().await.unwrap();
assert!(chunks.iter().all(|c| c.ends_with('\n')), "each chunk must be newline-terminated");
assert_eq!(chunks.concat().lines().count(), 3, "messages must not run together");
}
fn handle(harness: HarnessKind, id: &str) -> HarnessSession {
HarnessSession {
harness,
@@ -3014,48 +2812,33 @@ mod tests {
assert!(db.contains_message(&first.session, &first.source_id).await.unwrap());
}
/// Text interned by a page that failed is gone with the rest of it: a later row with the
/// same text stores it afresh rather than pointing at the id the failed page gave it.
#[rstest]
#[tokio::test]
async fn reindex_backfills_messages_indexed_before_the_fts_row_existed() {
async fn text_interned_by_a_failed_page_is_stored_again() {
let db = AiSessionDatabase::in_memory().await.unwrap();
let session = sample_handle();
let long: String = (0..64).map(|i| format!("chunk-{i:03}-buried ")).collect();
assert!(long.len() >= COMPRESS_THRESHOLD);
db.append(&message_in(&session, 0, &long)).await.unwrap();
db.append(&message_in(&session, 1, "plain buried token")).await.unwrap();
// Simulate a sidecar upgraded to the FTS migration with no index rows yet.
atuin_common::db::query("DELETE FROM messages_fts").execute(db.db.pool()).await.unwrap();
assert!(search(&db, "buried").await.is_empty());
db.reindex().await.unwrap();
assert_eq!(search(&db, "buried").await.len(), 1);
// A second pass must not duplicate the index or change results.
db.reindex().await.unwrap();
assert_eq!(search(&db, "buried").await.len(), 1);
}
#[rstest]
#[tokio::test]
async fn reindex_resumes_from_an_interrupted_backfill() {
let db = AiSessionDatabase::in_memory().await.unwrap();
let session = sample_handle();
db.append(&message_in(&session, 0, "alpha unique-a")).await.unwrap();
db.append(&message_in(&session, 1, "beta unique-b")).await.unwrap();
db.append(&message_in(&session, 2, "gamma unique-c")).await.unwrap();
atuin_common::db::query(
"DELETE FROM messages_fts WHERE rowid = (SELECT max(rowid) FROM messages_fts)",
let in_dir = |index| {
let mut msg = message_in(&session, index, "text");
msg.cwd = Some(std::path::PathBuf::from("/only/here"));
msg
};
db::query(
"CREATE TRIGGER fail_second BEFORE INSERT ON messages WHEN NEW.source_id = 'source-1' \
BEGIN SELECT RAISE(ABORT, 'fail'); END",
)
.execute(db.db.pool())
.await
.unwrap();
assert!(search(&db, "unique-c").await.is_empty(), "the highest-rowid row is unindexed");
assert_eq!(search(&db, "unique-a").await.len(), 1, "the prefix stays indexed");
let page = [0, 1].map(|i| super::PreparedMessage::new(in_dir(i)).unwrap());
let generation = db.projection_generation().await.unwrap();
db.append_all(page.into(), generation).await.expect_err("the trigger fails the page");
db.reindex().await.unwrap();
assert_eq!(search(&db, "unique-c").await.len(), 1, "reindex fills the suffix gap");
db.append(&in_dir(2)).await.unwrap();
let stored: Vec<Message> = db.messages(&session).try_collect().await.unwrap();
assert_eq!(stored.len(), 1);
assert_eq!(stored[0].cwd, Some(std::path::PathBuf::from("/only/here")));
}
#[rstest]
@@ -3982,60 +3765,6 @@ mod tests {
assert_eq!(output_in(&db, HarnessKind::ClaudeCode, "resumed").await, 3);
}
/// Migrating links the copies stored before by the same rule as rows arriving later.
#[rstest]
#[tokio::test]
async fn migrating_links_only_by_harness_given_turn_ids() {
let forms = turn_forms();
let db = sidecar_at(BEFORE_INCREMENTAL, &[]).await;
for (i, (harness, turn, _)) in forms.iter().enumerate() {
for (n, copy) in ["first", "second"].into_iter().enumerate() {
let id = format!("{copy}-{i}");
let at = i64::try_from(n).unwrap();
db::query(
"INSERT INTO sessions (harness, session_id, started_at, updated_at) VALUES \
(?, ?, ?, ?)",
)
.bind(*harness as i64)
.bind(&id)
.bind(at)
.bind(at)
.execute(db.db.pool())
.await
.unwrap();
db::query(
"INSERT INTO messages (id, harness, session, source_id, timestamp, role, \
content, turn_id) VALUES (randomblob(16), ?1, (SELECT id FROM sessions WHERE \
harness = ?1 AND session_id = ?2), 'a1', ?3, '\"Assistant\"', '[]', ?4)",
)
.bind(*harness as i64)
.bind(&id)
.bind(at)
.bind(*turn)
.execute(db.db.pool())
.await
.unwrap();
}
}
db.migrate().await.unwrap();
for (i, (harness, turn, links)) in forms.into_iter().enumerate() {
let second = handle(harness, &format!("second-{i}"));
let second = db.get_session(&second).await.unwrap().unwrap();
let expected = links.then(|| handle(harness, &format!("first-{i}")));
assert_eq!(second.copy_of, expected, "{harness:?} {turn}");
assert_eq!(
second.group().session.as_ref(),
if links {
format!("first-{i}")
} else {
format!("second-{i}")
}
);
}
}
/// One roots-only row for the lot, with the copies counted as its children and shown as
/// forks (the copy's subagent is grouped under it too, but is no child a person carried on).
#[rstest]
@@ -4562,47 +4291,6 @@ mod tests {
);
}
#[rstest]
#[tokio::test]
async fn migrating_to_atuin_ids_derives_them_and_replays_every_record() {
let db = sidecar_at(7, &[("s", None)]).await;
db::query("INSERT INTO reproject_watermark VALUES ('h', 't', 3, 'r')")
.execute(db.db.pool())
.await
.unwrap();
let before = db.projection_generation().await.unwrap();
db.migrate().await.unwrap();
let session = handle(HarnessKind::ClaudeCode, "s");
let stored: Option<Vec<u8>> = db::query_scalar("SELECT atuin_id FROM sessions")
.fetch_one(db.db.pool())
.await
.unwrap();
assert_eq!(
stored.as_deref(),
Some(session.atuin_id(OffsetDateTime::UNIX_EPOCH).as_bytes().as_slice())
);
let marks: i64 = db::query_scalar("SELECT count(*) FROM reproject_watermark")
.fetch_one(db.db.pool())
.await
.unwrap();
assert_eq!(marks, 0);
assert_ne!(db.projection_generation().await.unwrap(), before);
}
#[rstest]
#[tokio::test]
async fn a_reader_before_the_backfill_derives_the_id() {
let db = sidecar_at(8, &[("s", None)]).await;
let session = handle(HarnessKind::ClaudeCode, "s");
assert_eq!(
db.get_session(&session).await.unwrap().unwrap().atuin_id,
session.atuin_id(OffsetDateTime::UNIX_EPOCH)
);
}
/// Links stay (harness, native id): their atuin ids are looked up when read, so a child
/// stored before its parent names the parent's id once the parent arrives.
#[rstest]
@@ -4656,22 +4344,6 @@ mod tests {
assert_eq!(listed_child, &c);
}
/// Sessions stored before the `atuin_id` migration resolve their links once it backfills.
#[rstest]
#[tokio::test]
async fn links_of_sessions_from_before_atuin_ids_resolve() {
let db = sidecar_at(7, &[("p", None), ("c", Some("p"))]).await;
db.migrate().await.unwrap();
let (p, c) = (handle(HarnessKind::ClaudeCode, "p"), handle(HarnessKind::ClaudeCode, "c"));
let (p, c) = (
db.get_session(&p).await.unwrap().unwrap(),
db.get_session(&c).await.unwrap().unwrap(),
);
assert_eq!(c.parent_atuin_id, Some(p.atuin_id));
assert_eq!(p.child_atuin_ids, vec![c.atuin_id]);
}
#[rstest]
#[case::ascii("ab", Some("ac"))]
#[case::empty("", None)]
@@ -4762,10 +4434,14 @@ mod tests {
for index in [0, 2, 3, 6] {
corrupt(&db, &s, index).await;
}
db::query("INSERT INTO interned (value) VALUES ('not json')")
.execute(db.db.pool())
.await
.unwrap();
for index in [0, 6] {
db::query(
"UPDATE messages SET role = 'not json' WHERE session = (SELECT id FROM sessions \
WHERE session_id = ?) AND source_id = ?",
"UPDATE messages SET role = (SELECT id FROM interned WHERE value = 'not json') \
WHERE session = (SELECT id FROM sessions WHERE session_id = ?) AND source_id = ?",
)
.bind(s.session.as_ref())
.bind(format!("source-{index}"))
@@ -4816,11 +4492,11 @@ mod tests {
async fn preview_parts_use_the_session_time_index() {
let db = AiSessionDatabase::in_memory().await.unwrap();
for sql in [
"SELECT content, content_z FROM messages WHERE session = 1 AND role = '\"User\"' \
ORDER BY timestamp, id LIMIT 1",
"SELECT content, content_z FROM messages WHERE session = 1 AND role = (SELECT id FROM \
interned WHERE value = '\"User\"') ORDER BY timestamp, id LIMIT 1",
"SELECT timestamp, id, content, content_z FROM messages WHERE session = 1 AND role = \
'\"Assistant\"' AND (timestamp < 5 OR (timestamp = 5 AND id < x'00')) ORDER BY \
timestamp DESC, id DESC LIMIT 1",
(SELECT id FROM interned WHERE value = '\"Assistant\"') AND (timestamp < 5 OR \
(timestamp = 5 AND id < x'00')) ORDER BY timestamp DESC, id DESC LIMIT 1",
] {
let plan: Vec<(i64, i64, i64, String)> =
db::query_as(sqlx::AssertSqlSafe(format!("EXPLAIN QUERY PLAN {sql}")))
@@ -4945,237 +4621,6 @@ mod tests {
assert_eq!(newest, Some(SCHEMA_VERSION));
}
/// A sidecar at `version`, holding `sessions`, (harness, id, parent) rows written as that
/// version stored them.
async fn sidecar_at(version: i64, sessions: &[(&str, Option<&str>)]) -> AiSessionDatabase {
let (db, writer) = AiSessionDatabase::memory().await.unwrap();
#[allow(clippy::disallowed_macros)]
let mut migrator = sqlx::migrate!("./src/ai_session/migrations");
migrator.migrations =
migrator.migrations.iter().filter(|m| m.version <= version).cloned().collect();
migrator.run(db.pool()).await.unwrap();
for (i, (id, parent)) in sessions.iter().enumerate() {
db::query(
"INSERT INTO sessions (harness, session_id, parent_harness, parent_session_id, \
started_at, updated_at) VALUES (1, ?, ?, ?, ?, ?)",
)
.bind(*id)
.bind(parent.map(|_| 1))
.bind(*parent)
.bind(i64::try_from(i).unwrap())
.bind(i64::try_from(i).unwrap())
.execute(db.pool())
.await
.unwrap();
}
AiSessionDatabase::from_sqlite(db, Some(writer))
}
/// The version before the `incremental_sidecar` migration: a released sidecar.
const BEFORE_INCREMENTAL: i64 = 5;
/// Store a row of `session` holding model call `turn` in a sidecar at
/// [`BEFORE_INCREMENTAL`], as that version stored them.
async fn old_row(db: &AiSessionDatabase, i: usize, session: &str, turn: &str) {
db::query(
"INSERT INTO messages (id, harness, session, source_id, timestamp, role, content, \
turn_id) VALUES (?, 1, (SELECT id FROM sessions WHERE harness = 1 AND session_id = \
?), 'a1', ?, '\"Assistant\"', '[]', ?)",
)
.bind(vec![u8::try_from(i).unwrap(); 16])
.bind(session)
.bind(i64::try_from(i).unwrap())
.bind(turn)
.execute(db.db.pool())
.await
.unwrap();
}
/// Migrating groups the sessions already stored, and starts without reproject watermarks so
/// the next daemon start replays every record (filling in hosts).
#[rstest]
#[tokio::test]
async fn migrating_groups_stored_sessions_and_forces_a_reproject() {
let db = sidecar_at(BEFORE_INCREMENTAL, &[
("agent-b", Some("fork")),
("fork", Some("root")),
("root", None),
("orphan", Some("ghost")),
("a", Some("b")),
("b", Some("a")),
])
.await;
db.migrate().await.unwrap();
let roots = roots_of(&db).await;
let root_of = |id: &str| roots.iter().find(|(s, _)| s == id).unwrap().1.clone();
assert_eq!(root_of("agent-b"), "root");
assert_eq!(root_of("fork"), "root");
assert_eq!(root_of("root"), "root");
assert_eq!(root_of("orphan"), "orphan");
// A parent cycle has no top: its least member heads it.
assert_eq!(root_of("a"), "a");
assert_eq!(root_of("b"), "a");
let watermarks: i64 = db::query_scalar("SELECT count(*) FROM reproject_watermark")
.fetch_one(db.db.pool())
.await
.unwrap();
assert_eq!(watermarks, 0);
// Grouping carries on incrementally from the migrated state.
db.append(&tree_row("agent-c", Some("agent-b"), 50, "x")).await.unwrap();
assert_eq!(roots_of(&db).await.iter().find(|(s, _)| s == "agent-c").unwrap().1, "root");
}
/// Migrating groups a chain of any length under its top.
#[rstest]
#[tokio::test]
async fn migrating_groups_a_deep_chain_under_its_top() {
let ids = chain_ids(100);
let sessions: Vec<_> = ids.iter().map(|(id, p)| (id.as_str(), p.as_deref())).collect();
let db = sidecar_at(BEFORE_INCREMENTAL, &sessions).await;
db.migrate().await.unwrap();
let roots = roots_of(&db).await;
assert!(roots.iter().all(|(_, root)| root == "s000"), "{roots:?}");
}
/// Migrating links the copies already stored to their originals from the calls they share,
/// and regroups them, from the rows alone.
#[rstest]
#[tokio::test]
async fn migrating_links_stored_copies() {
let db = sidecar_at(BEFORE_INCREMENTAL, &[
("original", None),
("resumed", None),
("agent-x", Some("resumed")),
("zeta", None),
])
.await;
for (i, (session, turn)) in
[("original", "msg_A"), ("resumed", "msg_A"), ("zeta", "msg_Z")].into_iter().enumerate()
{
old_row(&db, i, session, turn).await;
}
db.migrate().await.unwrap();
let links = copy_links(&db).await;
let link = |id: &str| links.iter().find(|(s, ..)| s == id).unwrap().clone();
assert_eq!(link("resumed").1, "original");
assert_eq!(link("resumed").2.as_deref(), Some("original"));
assert_eq!(link("agent-x").1, "original");
assert_eq!(link("zeta").1, "zeta");
// Linking carries on incrementally from the migrated state.
db.append(&claude_row("twice", None, "a1", 5, Some("msg_A"), 1)).await.unwrap();
let links = copy_links(&db).await;
assert_eq!(links.iter().find(|(s, ..)| s == "twice").unwrap().1, "original");
}
/// The version before the `session_scoped_calls` migration.
const BEFORE_SCOPED_CALLS: i64 = 6;
/// Migrating attributes the usage already stored afresh from the stored rows: unrelated
/// Codex sessions whose content-derived call the old rule merged into one get theirs back,
/// a call shared under a harness-given id still counts once, and rows outside any call
/// keep counting.
#[rstest]
#[tokio::test]
async fn migrating_recounts_stored_usage() {
let db = sidecar_at(BEFORE_SCOPED_CALLS, &[]).await;
let codex = HarnessKind::Codex as i64;
let claude = HarnessKind::ClaudeCode as i64;
let token_count = "token_count:12000.0.40.0.12040";
// (harness, session, started, output as the old rule attributed it)
let sessions = [
(codex, "first", 0, 10 + 4),
(codex, "second", 5, 0),
(claude, "original", 0, 20),
(claude, "resumed", 5, 0),
];
for (harness, id, at, output) in sessions {
db::query(
"INSERT INTO sessions (harness, session_id, started_at, updated_at, usage_output, \
root_harness, root_session_id) VALUES (?1, ?2, ?3, ?3, ?4, ?1, ?2)",
)
.bind(harness)
.bind(id)
.bind(at)
.bind(output)
.execute(db.db.pool())
.await
.unwrap();
}
// (harness, session, source, turn, output)
let rows = [
(codex, "first", "c1", Some(token_count), 10),
(codex, "first", "c2", None, 4),
(codex, "second", "c1", Some(token_count), 7),
(claude, "original", "a1", Some("msg_A"), 20),
(claude, "resumed", "a1", Some("msg_A"), 20),
];
for (i, (harness, id, source, turn, output)) in rows.into_iter().enumerate() {
db::query(
"INSERT INTO messages (id, harness, session, source_id, timestamp, role, content, \
turn_id, usage_output, usage_present) VALUES (?1, ?2, (SELECT id FROM sessions \
WHERE harness = ?2 AND session_id = ?3), ?4, ?5, '\"Assistant\"', '[]', ?6, ?7, \
1)",
)
.bind(vec![u8::try_from(i).unwrap(); 16])
.bind(harness)
.bind(id)
.bind(source)
.bind(i64::try_from(i).unwrap())
.bind(turn)
.bind(output)
.execute(db.db.pool())
.await
.unwrap();
}
for (harness, turn, owner, output) in
[(codex, token_count, "first", 10), (claude, "msg_A", "original", 20)]
{
db::query(
"INSERT INTO calls (harness, turn_id, session_id, usage_input, usage_output, \
usage_cache_read, usage_cache_write, usage_reasoning) VALUES (?, ?, ?, 0, ?, 0, \
0, 0)",
)
.bind(harness)
.bind(turn)
.bind(owner)
.bind(output)
.execute(db.db.pool())
.await
.unwrap();
}
db.migrate().await.unwrap();
let sessions = [
(HarnessKind::Codex, "first"),
(HarnessKind::Codex, "second"),
(HarnessKind::ClaudeCode, "original"),
(HarnessKind::ClaudeCode, "resumed"),
];
let mut got = Vec::new();
for (harness, id) in sessions {
got.push(output_in(&db, harness, id).await);
}
assert_eq!(got, vec![14, 7, 20, 0]);
// Once: migrating again leaves the recount alone, and attribution carries on from it.
db.migrate().await.unwrap();
db.append(&usage_row(HarnessKind::Codex, "second", None, 6, token_count, 8)).await.unwrap();
let mut got = Vec::new();
for (harness, id) in sessions {
got.push(output_in(&db, harness, id).await);
}
assert_eq!(got, vec![14, 8, 20, 0]);
}
// --- read-only open -------------------------------------------------------------------------
#[fixture]
@@ -5213,8 +4658,7 @@ mod tests {
#[rstest]
#[case::uninitialized(None)]
#[case::outdated(Some(1))]
#[case::one_behind(Some(SCHEMA_VERSION - 1))]
#[case::outdated(Some(SCHEMA_VERSION - 1))]
#[case::unknown(Some(99))]
#[tokio::test]
async fn a_read_only_open_refuses_another_schema(
@@ -5226,20 +4670,12 @@ mod tests {
let db = Sqlite::builder(path.as_os_str()).open().await.unwrap();
if let Some(version) = version {
#[allow(clippy::disallowed_macros)]
let mut migrator = sqlx::migrate!("./src/ai_session/migrations");
migrator.migrations =
migrator.migrations.iter().filter(|m| m.version <= version).cloned().collect();
migrator.run(db.pool()).await.unwrap();
if version > SCHEMA_VERSION {
db::query(
"INSERT INTO _sqlx_migrations (version, description, success, checksum, \
execution_time) VALUES (?, 'future', 1, x'00', 0)",
)
sqlx::migrate!("./src/ai_session/migrations").run(db.pool()).await.unwrap();
db::query("UPDATE _sqlx_migrations SET version = ?")
.bind(version)
.execute(db.pool())
.await
.unwrap();
}
}
db.close().await;
}
@@ -5260,15 +4696,12 @@ mod tests {
// --- sidecars from other builds -------------------------------------------------------------
/// A sidecar a development build migrated with its own numbering: this build's first
/// migration, then others at versions this build uses for different ones.
/// A sidecar another build migrated past this one: this build's migration, then others it
/// does not know.
async fn dev_build_sidecar(path: &Path) {
let db = Sqlite::builder(path.as_os_str()).open().await.unwrap();
#[allow(clippy::disallowed_macros)]
let mut migrator = sqlx::migrate!("./src/ai_session/migrations");
migrator.migrations =
migrator.migrations.iter().filter(|m| m.version == 1).cloned().collect();
migrator.run(db.pool()).await.unwrap();
sqlx::migrate!("./src/ai_session/migrations").run(db.pool()).await.unwrap();
for (version, description) in [(2, "reproject watermark"), (3, "hosts and roots")] {
db::query(
"INSERT INTO _sqlx_migrations (version, description, success, checksum, \
@@ -5281,8 +4714,8 @@ mod tests {
.unwrap();
}
db::query(
"INSERT INTO sessions (harness, session_id, started_at, updated_at) VALUES (1, 'old', \
0, 0)",
"INSERT INTO sessions (harness, session_id, atuin_id, started_at, updated_at) VALUES \
(1, 'old', x'00', 0, 0)",
)
.execute(db.pool())
.await
@@ -5312,40 +4745,6 @@ mod tests {
assert_eq!(sessions, 1, "the sidecar was touched");
}
/// A sidecar a released daemon left (every migration before `incremental_sidecar`) is
/// migrated in place, what it holds kept.
#[rstest]
#[tokio::test]
async fn a_sidecar_from_main_is_migrated_in_place(dir: tempfile::TempDir) {
let path = dir.path().join("sidecar.db");
{
let db = Sqlite::builder(path.as_os_str()).open().await.unwrap();
#[allow(clippy::disallowed_macros)]
let mut migrator = sqlx::migrate!("./src/ai_session/migrations");
migrator.migrations = migrator
.migrations
.iter()
.filter(|m| m.version <= BEFORE_INCREMENTAL)
.cloned()
.collect();
migrator.run(db.pool()).await.unwrap();
db::query(
"INSERT INTO sessions (harness, session_id, parent_harness, parent_session_id, \
parent_kind, started_at, updated_at) VALUES (1, 'kept', 1, 'root', 1, 0, 0)",
)
.execute(db.pool())
.await
.unwrap();
db.close().await;
}
let db = AiSessionDatabase::open(&path).await.unwrap();
db.check_schema().await.unwrap();
let kept = db.get_session(&handle(HarnessKind::ClaudeCode, "kept")).await.unwrap().unwrap();
assert_eq!(kept.parent_kind, Some(ParentKind::Fork));
assert_eq!(kept.inferred_parent_kind(), Some(ParentKind::Fork));
}
/// A newer build's sidecar is left alone: that build owns it, and this one reports it.
#[rstest]
#[tokio::test]
crates/atuin-client/src/ai_session/database/watermark.rs
+7 -18
View File
@@ -8,7 +8,7 @@
//! Anything that clears watermarks or deletes projected rows is an invalidation, and bumps the
//! [`Generation`] in the same transaction. A reprojection moves a watermark only while the
//! generation is the one it read before replaying, so it never records as projected what an
//! invalidation removed or asked to be replayed meanwhile (see the `incremental_sidecar` migration).
//! invalidation removed or asked to be replayed meanwhile (see `projection_state` in the schema).
use std::collections::HashMap;
use std::path::Path;
@@ -179,8 +179,8 @@ impl AiSessionDatabase {
}
/// Make sure the watermarks were made with the key whose PASERK id is `key_id`, clearing
/// them all when they were not (or when that is unknown, as right after the `incremental_sidecar`
/// migration). Returns whether it cleared them.
/// them all when they were not (or when that is unknown, as in a new sidecar). Returns whether
/// it cleared them.
///
/// A record the key cannot decrypt holds its series' watermark below it, so a stale key
/// never moves a watermark past what the right one could project. This covers the rest: once
@@ -214,8 +214,7 @@ impl AiSessionDatabase {
/// Make the next reprojection into the sidecar at `path` a full one, without migrating or
/// otherwise opening it for use. For the maintenance commands that rewrite the record store
/// while the daemon (which owns the sidecar) may be running: a reprojection the daemon has in
/// flight notices, and starts over. A missing sidecar, or one from before watermarks, is
/// replayed in full anyway.
/// flight notices, and starts over. A missing sidecar is replayed in full anyway.
pub async fn invalidate_projection(path: impl AsRef<Path>) -> Result<(), DbError> {
let path = path.as_ref();
if !path.exists() {
@@ -260,21 +259,11 @@ impl AiSessionDatabase {
Ok(())
}
/// [`Self::invalidate_projection`] over an open sidecar, at whatever schema version it is.
/// [`Self::invalidate_projection`] over an open sidecar.
async fn invalidate_tables(sqlite: &Sqlite) -> Result<(), DbError> {
let mut tx = sqlite.pool().begin_with("BEGIN IMMEDIATE").await?;
let tables: Vec<String> = db::query_scalar(
"SELECT name FROM sqlite_master WHERE type = 'table' AND name IN \
('reproject_watermark', 'projection_state')",
)
.fetch_all(&mut *tx)
.await?;
if tables.iter().any(|t| t == "reproject_watermark") {
db::query("DELETE FROM reproject_watermark").execute(&mut *tx).await?;
}
if tables.iter().any(|t| t == "projection_state") {
Self::bump_generation(&mut tx).await?;
}
db::query("DELETE FROM reproject_watermark").execute(&mut *tx).await?;
Self::bump_generation(&mut tx).await?;
tx.commit().await?;
Ok(())
}
crates/atuin-client/src/ai_session/database/write.rs
+122 -117
View File
@@ -5,7 +5,7 @@
//! more than the statements themselves; here the whole write runs on a blocking thread, against
//! a rusqlite connection opened on the same database as the sqlx pool that serves the reads.
use std::collections::{BTreeMap, BTreeSet, HashSet};
use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet};
use std::sync::Arc;
use std::time::Duration;
@@ -29,7 +29,57 @@ const STATEMENT_CACHE: usize = 64;
/// The sidecar's writing connection, shared by clones. Writes take turns on it.
#[derive(Debug, Clone)]
pub(super) struct Writer(Arc<Mutex<Connection>>);
pub(super) struct Writer {
conn: Arc<Mutex<Connection>>,
interned: Arc<Mutex<Interned>>,
}
/// The `interned` ids this writer has looked up or stored, so the few values a session repeats
/// on every row cost no query after the first. Only what a committed transaction stored is
/// kept: an id from one rolled back would point at nothing.
#[derive(Debug, Default)]
pub(super) struct Interned {
committed: HashMap<String, i64>,
pending: HashMap<String, i64>,
}
impl Interned {
/// The `interned` id of `value`, storing it the first time it is seen.
fn id(&mut self, conn: &Connection, value: Option<&str>) -> Result<Option<i64>, DbError> {
let Some(value) = value else {
return Ok(None);
};
if let Some(id) = self.committed.get(value).or_else(|| self.pending.get(value)) {
return Ok(Some(*id));
}
let id = match query_scalar("SELECT id FROM interned WHERE value = ?")
.bind(value)
.fetch_optional(conn)?
{
Some(id) => id,
None => query("INSERT INTO interned (value) VALUES (?)")
.bind(value)
.execute(conn)?
.last_insert_rowid(),
};
self.pending.insert(value.to_owned(), id);
Ok(Some(id))
}
/// Run `write` in a transaction it commits, keeping the ids it stored only if it did.
fn transaction<T>(
&mut self,
write: impl FnOnce(&mut Self) -> Result<(T, bool), DbError>,
) -> Result<T, DbError> {
self.pending.clear();
let (result, committed) = write(self)?;
if committed {
let pending = std::mem::take(&mut self.pending);
self.committed.extend(pending);
}
Ok(result)
}
}
impl Writer {
/// Open `location`, a file the sqlx pool has created or a shared in-memory database's URI,
@@ -40,7 +90,20 @@ impl Writer {
conn.pragma_update(None, "synchronous", "NORMAL")?;
conn.pragma_update(None, "foreign_keys", "ON")?;
conn.set_prepared_statement_cache_capacity(STATEMENT_CACHE);
Ok(Self(Arc::new(Mutex::new(conn))))
Ok(Self {
conn: Arc::new(Mutex::new(conn)),
interned: Arc::default(),
})
}
/// [`Self::run`], with this writer's [`Interned`] ids.
pub(super) async fn run_interning<T, F>(&self, write: F) -> Result<T, DbError>
where
T: Send + 'static,
F: FnOnce(&mut Connection, &mut Interned) -> Result<T, DbError> + Send + 'static,
{
let interned = Arc::clone(&self.interned);
self.run(move |conn| write(conn, &mut interned.lock())).await
}
/// Run `write` on the connection, on a blocking thread. A panic in it is raised here, as if it
@@ -50,7 +113,7 @@ impl Writer {
T: Send + 'static,
F: FnOnce(&mut Connection) -> Result<T, DbError> + Send + 'static,
{
let conn = Arc::clone(&self.0);
let conn = Arc::clone(&self.conn);
let task = tokio::task::spawn_blocking(move || {
// A write that panicked rolled its transaction back as it unwound: the connection is
// as good as before it.
@@ -148,34 +211,44 @@ impl AiSessionDatabase {
/// `generation`: see [`Self::append_all`].
pub(super) fn append_page(
conn: &mut Connection,
interned: &mut Interned,
msgs: &[PreparedMessage],
generation: Generation,
) -> Result<bool, DbError> {
let tx = conn.transaction_with_behavior(TransactionBehavior::Immediate)?;
let current: i64 =
query_scalar("SELECT generation FROM projection_state WHERE id = 0").fetch_one(&tx)?;
if current != generation.0 {
return Ok(false);
}
for msg in msgs {
Self::append_in(&tx, msg)?;
}
tx.commit()?;
Ok(true)
interned.transaction(|interned| {
let tx = conn.transaction_with_behavior(TransactionBehavior::Immediate)?;
let current: i64 = query_scalar("SELECT generation FROM projection_state WHERE id = 0")
.fetch_one(&tx)?;
if current != generation.0 {
return Ok((false, false));
}
for msg in msgs {
Self::append_in(&tx, interned, msg)?;
}
tx.commit()?;
Ok((true, true))
})
}
/// Append `msg` in its own transaction: see [`Self::append`].
pub(super) fn append_one(
conn: &mut Connection,
interned: &mut Interned,
msg: &PreparedMessage,
) -> Result<Appended, DbError> {
let tx = conn.transaction_with_behavior(TransactionBehavior::Immediate)?;
let appended = Self::append_in(&tx, msg)?;
tx.commit()?;
Ok(appended)
interned.transaction(|interned| {
let tx = conn.transaction_with_behavior(TransactionBehavior::Immediate)?;
let appended = Self::append_in(&tx, interned, msg)?;
tx.commit()?;
Ok((appended, true))
})
}
fn append_in(conn: &Connection, prepared: &PreparedMessage) -> Result<Appended, DbError> {
fn append_in(
conn: &Connection,
interned: &mut Interned,
prepared: &PreparedMessage,
) -> Result<Appended, DbError> {
let msg = &prepared.msg;
let harness = msg.session.harness as i64;
let session_id = msg.session.session.as_ref();
@@ -226,7 +299,7 @@ impl AiSessionDatabase {
id, harness, session, source_id, parent_harness, parent_session_id,
parent_source_id, timestamp, role, content, content_z, cwd, git_branch, model,
usage_input, usage_output, usage_cache_read, usage_cache_write,
usage_reasoning, stop_reason, usage_present, turn_id, title_change, host_id
usage_reasoning, stop_reason, usage_present, turn_id, title_change, host
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
ON CONFLICT(session, source_id) DO NOTHING",
)
@@ -238,12 +311,12 @@ impl AiSessionDatabase {
.bind(parent_session_id.clone())
.bind(msg.parent_source_id.as_ref().map(|s| s.as_ref()))
.bind(timestamp)
.bind(&prepared.role)
.bind(interned.id(conn, Some(&prepared.role))?)
.bind(&prepared.content)
.bind(prepared.content_z.as_deref())
.bind(cwd.clone())
.bind(msg.git_branch.as_deref())
.bind(msg.model.as_deref())
.bind(interned.id(conn, cwd.as_deref())?)
.bind(interned.id(conn, msg.git_branch.as_deref())?)
.bind(interned.id(conn, msg.model.as_deref())?)
.bind(usage_input)
.bind(usage_output)
.bind(usage_cache_read)
@@ -253,13 +326,13 @@ impl AiSessionDatabase {
.bind(i64::from(msg.usage.is_some()))
.bind(msg.turn_id.as_deref())
.bind(prepared.title_change.as_deref())
.bind(host.as_deref())
.bind(interned.id(conn, host.as_deref())?)
.execute(conn)?;
if inserted.rows_affected() == 0 {
// A row stored before hosts were tracked learns its host when the reproject replays
// its record (see the `incremental_sidecar` migration).
// A row stored without a host learns it when the reproject replays its record.
if let Some(host) = &host {
let host = interned.id(conn, Some(host))?;
Self::backfill_host(conn, session, source_id, host)?;
}
// Another host may have captured this row under another id.
@@ -356,8 +429,8 @@ impl AiSessionDatabase {
Self::regroup(conn, harness, session_id)?;
}
// A copy names no parent, so it is linked to its original through the calls they share
// (see the `incremental_sidecar` migration). A link that moves other than by being made regroups
// everything.
// (see `sessions.copy_of_session_id` in the schema). A link that moves other than by being
// made regroups everything.
let mut relinked = false;
if let Some(before) = &before {
if gained_parent {
@@ -435,8 +508,9 @@ impl AiSessionDatabase {
type BodyRow =
(i64, String, Option<Vec<u8>>, Option<String>, Option<String>, Option<String>);
let rows: Vec<BodyRow> = query_as(
"SELECT rowid, content, content_z, cwd, git_branch, model FROM messages WHERE \
session = ?",
"SELECT rowid, content, content_z, (SELECT value FROM interned WHERE id = cwd), \
(SELECT value FROM interned WHERE id = git_branch), (SELECT value FROM interned \
WHERE id = model) FROM messages WHERE session = ?",
)
.bind(session)
.fetch_all(conn)?;
@@ -555,8 +629,9 @@ impl AiSessionDatabase {
/// For a row of `session` holding call `turn`, just stored: offer each other session holding
/// the call as the other's original. Each keeps the lowest-ranked original it is offered (see
/// the `incremental_sidecar` migration), so this only ever lowers a link, and while no session's start or
/// parent moves, the links end up the same whatever order the rows arrive in.
/// `sessions.copy_of_session_id` in the schema), so this only ever lowers a link, and while no
/// session's start or parent moves, the links end up the same whatever order the rows arrive
/// in.
///
/// A session linked for the first time was a root, and is placed like one learning its
/// parent. Returns whether an existing link moved instead, which needs
@@ -614,8 +689,8 @@ impl AiSessionDatabase {
Ok(moved)
}
/// Link `session` to its original afresh from every call it shares, as
/// [`Self::group_migrated`] does. Returns whether its link changed.
/// Link `session` to its original afresh from every call it shares. Returns whether its link
/// changed.
fn relink(conn: &Connection, harness: i64, session_id: &str) -> Result<bool, DbError> {
let (previous, parent): (Option<String>, Option<String>) = query_as(
"SELECT copy_of_session_id, parent_session_id FROM sessions WHERE harness = ? AND \
@@ -692,16 +767,16 @@ impl AiSessionDatabase {
Ok(changed)
}
/// Record `host` on a stored row that has none, and work its session's host out again.
/// Record `host` (its `interned` id) on a stored row that has none, and work its session's
/// host out again.
fn backfill_host(
conn: &Connection,
session: i64,
source_id: &str,
host: &str,
host: Option<i64>,
) -> Result<(), DbError> {
let filled = query(
"UPDATE messages SET host_id = ? WHERE session = ? AND source_id = ? AND host_id IS \
NULL",
"UPDATE messages SET host = ? WHERE session = ? AND source_id = ? AND host IS NULL",
)
.bind(host)
.bind(session)
@@ -720,8 +795,9 @@ impl AiSessionDatabase {
/// whatever it has.
fn refresh_session_host(conn: &Connection, session: i64) -> Result<(), DbError> {
query(
"UPDATE sessions SET host_id = COALESCE((SELECT host_id FROM messages WHERE session = \
?1 AND host_id IS NOT NULL ORDER BY timestamp, id LIMIT 1), host_id) WHERE id = ?1",
"UPDATE sessions SET host_id = COALESCE((SELECT (SELECT value FROM interned WHERE id \
= host) FROM messages WHERE session = ?1 AND host IS NOT NULL ORDER BY timestamp, id \
LIMIT 1), host_id) WHERE id = ?1",
)
.bind(session)
.execute(conn)?;
@@ -995,78 +1071,6 @@ impl AiSessionDatabase {
Ok(false)
}
/// Attribute the usage stored before the `session_scoped_calls` migration afresh from the
/// stored rows: that migration empties `calls`, whose rows were attributed by a rule that
/// merged unrelated sessions' calls sharing a content-derived turn id. Each session's usage
/// starts over from its rows outside any model call, and every call is attributed as rows
/// arriving later are ([`Self::attribute_call`]), which depends only on the rows stored and
/// how they are grouped (so this runs after [`Self::group_migrated`]: a content-derived call
/// is counted once per group of its holders). Runs until it has committed once, and then
/// never again.
pub(super) fn recount_migrated(conn: &mut Connection) -> Result<(), DbError> {
let tx = conn.transaction_with_behavior(TransactionBehavior::Immediate)?;
let recount: i64 = query_scalar("SELECT recount_calls FROM projection_state WHERE id = 0")
.fetch_one(&tx)?;
if recount == 0 {
return Ok(());
}
query("DELETE FROM calls").execute(&tx)?;
query(
"UPDATE sessions SET (usage_input, usage_output, usage_cache_read, usage_cache_write, \
usage_reasoning) = (SELECT COALESCE(SUM(m.usage_input), 0), \
COALESCE(SUM(m.usage_output), 0), COALESCE(SUM(m.usage_cache_read), 0), \
COALESCE(SUM(m.usage_cache_write), 0), COALESCE(SUM(m.usage_reasoning), 0) FROM \
messages m WHERE m.session = sessions.id AND m.usage_present = 1 AND m.turn_id IS \
NULL)",
)
.execute(&tx)?;
let turns: Vec<(i64, String)> = query_as(
"SELECT DISTINCT harness, turn_id FROM messages WHERE usage_present = 1 AND turn_id \
IS NOT NULL",
)
.fetch_all(&tx)?;
for (harness, turn) in &turns {
Self::attribute_call(&tx, *harness, turn)?;
}
query("UPDATE projection_state SET recount_calls = 0 WHERE id = 0").execute(&tx)?;
tx.commit()?;
Ok(())
}
/// Link and group the sessions stored before the `incremental_sidecar` migration, which
/// leaves them without a root: in Rust rather than in the migration, so it links copies by
/// the same [`linkable_turn`] rule and groups them with the same query ([`Self::regroup_all`])
/// as rows arriving later are. Every session stored since has a root, so this runs until it
/// has committed once, and then never again.
pub(super) fn group_migrated(conn: &mut Connection) -> Result<(), DbError> {
let tx = conn.transaction_with_behavior(TransactionBehavior::Immediate)?;
let ungrouped: Option<i64> =
query_scalar("SELECT 1 FROM sessions WHERE root_harness IS NULL LIMIT 1")
.fetch_optional(&tx)?;
if ungrouped.is_none() {
return Ok(());
}
// Each parentless session links to the lowest-ranked other parentless session sharing a
// call, when that ranks below it (see the migration), as `relink` does one at a time.
query(concat!(
"UPDATE sessions SET copy_of_session_id = (SELECT t.session_id FROM messages m CROSS \
JOIN messages o ON o.harness = m.harness AND o.turn_id = m.turn_id AND o.session <> \
m.session CROSS JOIN sessions t ON t.id = o.session WHERE m.session = sessions.id \
AND ",
linkable_turn!(),
" AND t.parent_session_id IS NULL AND (t.started_at, t.session_id) < \
(sessions.started_at, sessions.session_id) ORDER BY t.started_at, t.session_id LIMIT \
1) WHERE parent_session_id IS NULL"
))
.execute(&tx)?;
Self::regroup_all(&tx)?;
for (harness, turn) in Self::take_regrouped_calls(&tx)? {
Self::attribute_call(&tx, harness, &turn)?;
}
tx.commit()?;
Ok(())
}
/// [`Self::forget_host_sparing`], on the writer.
pub(super) fn forget_host_in(
conn: &mut Connection,
@@ -1081,7 +1085,7 @@ impl AiSessionDatabase {
// The sessions going, as a JSON array of their row ids for `json_each`.
let going: String = query_scalar(
"SELECT json_group_array(id) FROM sessions WHERE host_id = ?1 OR id IN (SELECT \
session FROM messages WHERE host_id = ?1)",
session FROM messages WHERE host = (SELECT id FROM interned WHERE value = ?1))",
)
.bind(&host)
.fetch_one(&tx)?;
@@ -1096,8 +1100,9 @@ impl AiSessionDatabase {
.collect();
// Every other host with a row in the sessions going: NULL for a row of unknown host.
let contributors: Vec<Option<String>> = query_scalar(
"SELECT DISTINCT m.host_id FROM messages m WHERE m.session IN (SELECT value FROM \
json_each(?1)) AND m.host_id IS NOT ?2",
"SELECT DISTINCT (SELECT value FROM interned WHERE id = m.host) FROM messages m WHERE \
m.session IN (SELECT value FROM json_each(?1)) AND m.host IS NOT (SELECT id FROM \
interned WHERE value = ?2)",
)
.bind(&going)
.bind(&host)
crates/atuin-client/src/ai_session/migrations/0001_init.sql
+150 -39
View File
@@ -1,18 +1,101 @@
CREATE TABLE messages (
id BLOB PRIMARY KEY,
-- Sessions get an INTEGER PRIMARY KEY, and messages point at theirs by it rather than repeating
-- (harness, session_id) in the table and every index. Both tables have one so VACUUM cannot
-- renumber rowids: messages_fts is keyed by messages.rowid, and messages.session by sessions.id.
CREATE TABLE sessions (
id INTEGER PRIMARY KEY,
harness INTEGER NOT NULL,
session_id TEXT NOT NULL,
-- The session's id across hosts and harnesses (`AtuinSessionId`): the smallest any of its
-- rows carries.
atuin_id BLOB NOT NULL,
parent_harness INTEGER,
parent_session_id TEXT,
-- How the session relates to its parent (`ParentKind`: 0 subagent, 1 fork, 2 continuation),
-- so readers can tell a subagent's fragment from a conversation a person carried on. NULL
-- when the harness did not say.
parent_kind INTEGER,
-- The top-most stored ancestor, following parent links else the copy link: forks, subagents
-- and copies group under it. A session whose parent is not stored (yet) is its own root, and
-- is regrouped once it arrives.
root_harness INTEGER,
root_session_id TEXT,
-- The session a parentless one was copied from, inferred from the model calls they share:
-- Claude Code's `--fork-session` (and a `--resume` it turns into a fork) copies the history
-- into a new session with only `sessionId` rewritten, naming the original nowhere, so capture
-- has no parent to record. A call's `turn_id` is the harness's own id for it, so one held by
-- two sessions was copied from one into the other. Only an id the harness gave the call
-- counts, not one capture derived from a line's content, which unrelated sessions can share
-- (see `linkable_turn!`).
--
-- It points at the lowest-ranked (earliest start, then id) other parentless session sharing a
-- call, when that ranks below this one: so links only ever go down in rank and cannot cycle,
-- the group's root is the session that owns the shared calls (see `attribute_call`), and the
-- result depends only on the rows stored. A session with a parent (a subagent, a `/branch`
-- fork) groups by it instead, and is never a link's target.
copy_of_session_id TEXT,
cwd TEXT,
git_branch TEXT,
model TEXT,
-- The host the session was captured on: its earliest row's (see `messages.host`).
host_id TEXT,
started_at INTEGER NOT NULL,
updated_at INTEGER NOT NULL,
message_count INTEGER NOT NULL DEFAULT 0,
-- Usage attributed to this session: its rows without a model call, plus every call in
-- `calls` it owns.
usage_input INTEGER NOT NULL DEFAULT 0,
usage_output INTEGER NOT NULL DEFAULT 0,
usage_cache_read INTEGER NOT NULL DEFAULT 0,
usage_cache_write INTEGER NOT NULL DEFAULT 0,
usage_reasoning INTEGER NOT NULL DEFAULT 0,
-- The newest row's title and where it came from (TitleSource: 0 summary, 1 generated,
-- 2 named, 3 agent), so a resumed capture keeps ranking titles.
title TEXT,
title_source INTEGER,
preview TEXT,
-- The session's most recent assistant reply, so a listing can say how each session ended
-- without the reader opening every transcript. `last_reply_at` orders replies that arrive out
-- of order (a reprojection replays records in record order, not message order).
last_reply TEXT,
last_reply_at INTEGER,
UNIQUE (harness, session_id)
);
-- Sessions are listed newest first, often for one harness and only those active since a given
-- time; without these every listing reads and sorts the whole table.
CREATE INDEX sessions_harness_updated_at ON sessions (harness, updated_at);
CREATE INDEX sessions_updated_at ON sessions (updated_at);
CREATE INDEX sessions_root ON sessions (root_harness, root_session_id);
CREATE INDEX sessions_parent ON sessions (parent_harness, parent_session_id);
CREATE INDEX sessions_copy_of ON sessions (harness, copy_of_session_id)
WHERE copy_of_session_id IS NOT NULL;
-- Text that many message rows repeat (roles, directories, branches, models, hosts), stored once:
-- a session's rows mostly share one of each, so a row refers to its text by id. Rows here are
-- never deleted; there are only ever a few hundred.
CREATE TABLE interned (
id INTEGER PRIMARY KEY,
value TEXT NOT NULL UNIQUE
);
CREATE TABLE messages (
seq INTEGER PRIMARY KEY,
id BLOB NOT NULL UNIQUE,
harness INTEGER NOT NULL,
session INTEGER NOT NULL,
source_id TEXT NOT NULL,
parent_harness INTEGER,
parent_session_id TEXT,
parent_source_id TEXT,
timestamp INTEGER NOT NULL,
role TEXT NOT NULL,
-- The role as its JSON (`interned`).
role INTEGER NOT NULL,
content TEXT NOT NULL,
content_z BLOB,
cwd TEXT,
git_branch TEXT,
model TEXT,
-- These and `host` are `interned` ids.
cwd INTEGER,
git_branch INTEGER,
model INTEGER,
-- The usage this row reported, as the harness reported it. Every row of one model call may
-- repeat (or grow) the same figures, so these are never summed directly: see `calls`.
usage_input INTEGER NOT NULL DEFAULT 0,
@@ -29,61 +112,59 @@ CREATE TABLE messages (
turn_id TEXT,
-- The title the row's own line set or cleared (JSON TitleChange), replayed on resume.
title_change TEXT,
UNIQUE (harness, session_id, source_id)
-- The host the row was captured on (the record's host id, hyphenated), for the host filter
-- and to tell a session recorded elsewhere from a local one.
host INTEGER,
UNIQUE (session, source_id)
);
CREATE INDEX messages_harness_session_timestamp ON messages (harness, session_id, timestamp);
CREATE INDEX messages_session_timestamp ON messages (session, timestamp);
-- Rows of one model call, across sessions (usage accounting) and within one (reasoning counts).
CREATE INDEX messages_harness_turn_session ON messages (harness, turn_id, session_id);
CREATE TABLE sessions (
harness INTEGER NOT NULL,
session_id TEXT NOT NULL,
parent_harness INTEGER,
parent_session_id TEXT,
cwd TEXT,
git_branch TEXT,
model TEXT,
started_at INTEGER NOT NULL,
updated_at INTEGER NOT NULL,
message_count INTEGER NOT NULL DEFAULT 0,
-- Usage attributed to this session: its rows without a model call, plus every call in
-- `calls` it owns.
usage_input INTEGER NOT NULL DEFAULT 0,
usage_output INTEGER NOT NULL DEFAULT 0,
usage_cache_read INTEGER NOT NULL DEFAULT 0,
usage_cache_write INTEGER NOT NULL DEFAULT 0,
usage_reasoning INTEGER NOT NULL DEFAULT 0,
-- The newest row's title and where it came from (TitleSource: 0 summary, 1 generated,
-- 2 named, 3 agent), so a resumed capture keeps ranking titles.
title TEXT,
title_source INTEGER,
preview TEXT,
PRIMARY KEY (harness, session_id)
);
CREATE INDEX messages_harness_turn_session ON messages (harness, turn_id, session);
-- One model call, counted once whichever sessions its rows were copied into: the field-wise max
-- of every row reporting it, attributed to one owning session.
--
-- A call is counted once across every session holding it only when its `turn_id` is one the
-- harness gave it (see `linkable_turn!`). An id capture derived from a line's content (Codex
-- `token_count:`/`thread:`, Pi `line:`, opencode `<millis>:`) can be shared by unrelated sessions,
-- so such a call is counted once per group instead: once among the sessions holding it that share
-- a root (a fork copying its parent's history is grouped under it), and separately in each
-- unrelated group.
--
-- `scope` tells those apart: empty for a harness-given id, else the group's root as
-- `<root_harness>:<root_session_id>` (see `attribute_call`).
CREATE TABLE calls (
harness INTEGER NOT NULL,
turn_id TEXT NOT NULL,
scope TEXT NOT NULL,
session_id TEXT NOT NULL,
usage_input INTEGER NOT NULL,
usage_output INTEGER NOT NULL,
usage_cache_read INTEGER NOT NULL,
usage_cache_write INTEGER NOT NULL,
usage_reasoning INTEGER NOT NULL,
PRIMARY KEY (harness, turn_id)
PRIMARY KEY (harness, turn_id, scope)
);
-- The sessions whose root changed since `AiSessionDatabase` last attributed afresh the
-- content-derived calls they hold, which are scoped to that root: filled by the trigger, wherever
-- sessions are regrouped, and emptied in the same transaction.
CREATE TABLE regrouped (session INTEGER PRIMARY KEY);
CREATE TRIGGER sessions_regrouped AFTER UPDATE OF root_harness, root_session_id ON sessions
WHEN OLD.root_harness IS NOT NEW.root_harness OR OLD.root_session_id IS NOT NEW.root_session_id
BEGIN
INSERT OR IGNORE INTO regrouped (session) VALUES (NEW.id);
END;
CREATE TABLE checkpoints (
harness INTEGER NOT NULL,
session_id TEXT NOT NULL,
"offset" INTEGER NOT NULL,
-- The digest of the item the checkpoint was taken after, so a transcript rewritten under it
-- or an event log reset since is noticed on resume. A checkpoint without one is ignored: its
-- session is read again from the start.
digest INTEGER,
-- or an event log reset since is noticed on resume.
digest INTEGER NOT NULL,
PRIMARY KEY (harness, session_id)
);
@@ -100,3 +181,33 @@ CREATE VIRTUAL TABLE messages_fts USING fts5(
content = '',
contentless_delete = 1
);
-- How far the sidecar has been reprojected from each record series (one host's records of one
-- tag): every record at or below `idx` is projected. `record_id` is the id of the record at `idx`,
-- so a series rewritten under the watermark (a reset, deleted or re-keyed store) is noticed and
-- replayed from the start. No row means replay that series from the start.
CREATE TABLE reproject_watermark (
host TEXT NOT NULL,
tag TEXT NOT NULL,
idx INTEGER NOT NULL,
record_id TEXT NOT NULL,
PRIMARY KEY (host, tag)
);
-- What the reproject watermarks were made against, in a single row.
--
-- `generation` counts invalidations: whatever clears watermarks or deletes projected rows bumps it
-- in the same transaction, and a reprojection moves a watermark only while the generation is still
-- the one it started from. So an invalidation landing mid-reprojection (a maintenance command, or
-- a replay forgetting a host) is noticed even for a series that had no watermark to lose.
--
-- `key_id` is the PASERK id (a public key identifier, not the key) of the encryption key the
-- watermarks were made with. A reprojection under a different key clears every watermark first:
-- what the old key projected, or held back, is no guide to what the new one can read.
CREATE TABLE projection_state (
id INTEGER PRIMARY KEY CHECK (id = 0),
generation INTEGER NOT NULL,
key_id TEXT
);
INSERT INTO projection_state (id, generation, key_id) VALUES (0, 0, NULL);
crates/atuin-client/src/ai_session/migrations/0002_intern_sessions.sql
-97
View File
@@ -1,97 +0,0 @@
-- Messages point at their session by integer id rather than repeating (harness, session_id) in the
-- table and every index: a third of the sidecar was session ids.
--
-- Both tables get an INTEGER PRIMARY KEY so VACUUM cannot renumber rowids: messages_fts is keyed by
-- messages.rowid, and messages.session by sessions.id.
CREATE TABLE sessions_new (
id INTEGER PRIMARY KEY,
harness INTEGER NOT NULL,
session_id TEXT NOT NULL,
parent_harness INTEGER,
parent_session_id TEXT,
cwd TEXT,
git_branch TEXT,
model TEXT,
started_at INTEGER NOT NULL,
updated_at INTEGER NOT NULL,
message_count INTEGER NOT NULL DEFAULT 0,
usage_input INTEGER NOT NULL DEFAULT 0,
usage_output INTEGER NOT NULL DEFAULT 0,
usage_cache_read INTEGER NOT NULL DEFAULT 0,
usage_cache_write INTEGER NOT NULL DEFAULT 0,
usage_reasoning INTEGER NOT NULL DEFAULT 0,
title TEXT,
title_source INTEGER,
preview TEXT,
UNIQUE (harness, session_id)
);
INSERT INTO sessions_new (
harness, session_id, parent_harness, parent_session_id, cwd, git_branch, model, started_at,
updated_at, message_count, usage_input, usage_output, usage_cache_read, usage_cache_write,
usage_reasoning, title, title_source, preview
)
SELECT
harness, session_id, parent_harness, parent_session_id, cwd, git_branch, model, started_at,
updated_at, message_count, usage_input, usage_output, usage_cache_read, usage_cache_write,
usage_reasoning, title, title_source, preview
FROM sessions;
CREATE TABLE messages_new (
seq INTEGER PRIMARY KEY,
id BLOB NOT NULL UNIQUE,
harness INTEGER NOT NULL,
session INTEGER NOT NULL,
source_id TEXT NOT NULL,
parent_harness INTEGER,
parent_session_id TEXT,
parent_source_id TEXT,
timestamp INTEGER NOT NULL,
role TEXT NOT NULL,
content TEXT NOT NULL,
content_z BLOB,
cwd TEXT,
git_branch TEXT,
model TEXT,
-- The usage this row reported, as the harness reported it. Every row of one model call may
-- repeat (or grow) the same figures, so these are never summed directly: see `calls`.
usage_input INTEGER NOT NULL DEFAULT 0,
usage_output INTEGER NOT NULL DEFAULT 0,
usage_cache_read INTEGER NOT NULL DEFAULT 0,
usage_cache_write INTEGER NOT NULL DEFAULT 0,
-- Reasoning tokens, part of output; NULL when the row did not break them out.
usage_reasoning INTEGER,
-- Distinguishes "no usage reported" from "zero tokens".
usage_present INTEGER NOT NULL DEFAULT 0,
stop_reason TEXT,
-- The model call this row came from; groups the rows one response is split into, across
-- every session a harness copied them into.
turn_id TEXT,
-- The title the row's own line set or cleared (JSON TitleChange), replayed on resume.
title_change TEXT,
UNIQUE (session, source_id)
);
-- Keeps each message's rowid, which its messages_fts row is keyed by.
INSERT INTO messages_new (
seq, id, harness, session, source_id, parent_harness, parent_session_id, parent_source_id,
timestamp, role, content, content_z, cwd, git_branch, model, usage_input, usage_output,
usage_cache_read, usage_cache_write, usage_reasoning, usage_present, stop_reason, turn_id,
title_change
)
SELECT
m.rowid, m.id, m.harness, s.id, m.source_id, m.parent_harness, m.parent_session_id,
m.parent_source_id, m.timestamp, m.role, m.content, m.content_z, m.cwd, m.git_branch, m.model,
m.usage_input, m.usage_output, m.usage_cache_read, m.usage_cache_write, m.usage_reasoning,
m.usage_present, m.stop_reason, m.turn_id, m.title_change
FROM messages m
JOIN sessions_new s ON s.harness = m.harness AND s.session_id = m.session_id;
DROP TABLE messages;
DROP TABLE sessions;
ALTER TABLE sessions_new RENAME TO sessions;
ALTER TABLE messages_new RENAME TO messages;
CREATE INDEX messages_session_timestamp ON messages (session, timestamp);
CREATE INDEX messages_harness_turn_session ON messages (harness, turn_id, session);
crates/atuin-client/src/ai_session/migrations/0003_last_reply.sql
-5
View File
@@ -1,5 +0,0 @@
-- The session's most recent assistant reply, so a listing can say how each session ended
-- without the reader opening every transcript. `last_reply_at` orders replies that arrive out of
-- order (a reprojection replays records in record order, not message order).
ALTER TABLE sessions ADD COLUMN last_reply TEXT;
ALTER TABLE sessions ADD COLUMN last_reply_at INTEGER;
crates/atuin-client/src/ai_session/migrations/0004_parent_kind.sql
-4
View File
@@ -1,4 +0,0 @@
-- How a session relates to its parent (`ParentKind`: 0 subagent, 1 fork, 2 continuation), so
-- readers can tell a subagent's fragment from a conversation a person carried on. NULL when the
-- records predate it or the harness did not say.
ALTER TABLE sessions ADD COLUMN parent_kind INTEGER;
crates/atuin-client/src/ai_session/migrations/0005_sessions_updated_at.sql
-3
View File
@@ -1,3 +0,0 @@
-- Sessions are listed newest first, often for one harness and only those active since a given
-- time; without this every listing reads and sorts the whole table.
CREATE INDEX sessions_harness_updated_at ON sessions (harness, updated_at);
crates/atuin-client/src/ai_session/migrations/0006_incremental_sidecar.sql
-66
View File
@@ -1,66 +0,0 @@
-- Incremental reprojection, the host each row was captured on, and session groups (forks,
-- subagents and copies under their root).
-- How far the sidecar has been reprojected from each record series (one host's records of one
-- tag): every record at or below `idx` is projected. `record_id` is the id of the record at `idx`,
-- so a series rewritten under the watermark (a reset, deleted or re-keyed store) is noticed and
-- replayed from the start. No row means replay that series from the start: starting with none,
-- the next daemon start replays every record, which fills in `host_id` below.
CREATE TABLE reproject_watermark (
host TEXT NOT NULL,
tag TEXT NOT NULL,
idx INTEGER NOT NULL,
record_id TEXT NOT NULL,
PRIMARY KEY (host, tag)
);
-- What the reproject watermarks were made against, in a single row.
--
-- `generation` counts invalidations: whatever clears watermarks or deletes projected rows bumps it
-- in the same transaction, and a reprojection moves a watermark only while the generation is still
-- the one it started from. So an invalidation landing mid-reprojection (a maintenance command, or
-- a replay forgetting a host) is noticed even for a series that had no watermark to lose.
--
-- `key_id` is the PASERK id (a public key identifier, not the key) of the encryption key the
-- watermarks were made with. A reprojection under a different key clears every watermark first:
-- what the old key projected, or held back, is no guide to what the new one can read.
CREATE TABLE projection_state (
id INTEGER PRIMARY KEY CHECK (id = 0),
generation INTEGER NOT NULL,
key_id TEXT
);
INSERT INTO projection_state (id, generation, key_id) VALUES (0, 0, NULL);
-- The host each row was captured on (the record's host id, hyphenated), for the host filter and
-- to tell a session recorded elsewhere from a local one. A session's is its earliest row's. NULL
-- until a replay fills it in.
ALTER TABLE messages ADD COLUMN host_id TEXT;
ALTER TABLE sessions ADD COLUMN host_id TEXT;
-- The top-most stored ancestor, following parent links else the copy link: forks, subagents and
-- copies group under it. A session whose parent is not stored (yet) is its own root, and is
-- regrouped once it arrives. NULL only until `AiSessionDatabase::migrate` links and groups the
-- sessions stored before this migration, in Rust, with the queries grouping uses as rows arrive.
ALTER TABLE sessions ADD COLUMN root_harness INTEGER;
ALTER TABLE sessions ADD COLUMN root_session_id TEXT;
-- The session a parentless one was copied from, inferred from the model calls they share: Claude
-- Code's `--fork-session` (and a `--resume` it turns into a fork) copies the history into a new
-- session with only `sessionId` rewritten, naming the original nowhere, so capture has no parent
-- to record. A call's `turn_id` is the harness's own id for it, so one held by two sessions was
-- copied from one into the other. Only an id the harness gave the call counts, not one capture
-- derived from a line's content, which unrelated sessions can share (see `linkable_turn!`).
--
-- It points at the lowest-ranked (earliest start, then id) other parentless session sharing a
-- call, when that ranks below this one: so links only ever go down in rank and cannot cycle, the
-- group's root is the session that owns the shared calls (see `attribute_call`), and the result
-- depends only on the rows stored. A session with a parent (a subagent, a `/branch` fork) groups
-- by it instead, and is never a link's target.
ALTER TABLE sessions ADD COLUMN copy_of_session_id TEXT;
CREATE INDEX sessions_root ON sessions (root_harness, root_session_id);
CREATE INDEX sessions_parent ON sessions (parent_harness, parent_session_id);
CREATE INDEX sessions_updated_at ON sessions (updated_at);
CREATE INDEX sessions_copy_of ON sessions (harness, copy_of_session_id)
WHERE copy_of_session_id IS NOT NULL;
crates/atuin-client/src/ai_session/migrations/0007_session_scoped_calls.sql
-40
View File
@@ -1,40 +0,0 @@
-- A model call is counted once across every session holding it only when its `turn_id` is one the
-- harness gave it (see `linkable_turn!`). An id capture derived from a line's content (Codex
-- `token_count:`/`thread:`, Pi `line:`, opencode `<millis>:`) can be shared by unrelated sessions,
-- so such a call is counted once per group instead: once among the sessions holding it that share
-- a root (`sessions.root_harness`/`root_session_id`: a fork copying its parent's history is grouped
-- under it), and separately in each unrelated group.
--
-- `scope` tells those apart: empty for a harness-given id, else the group's root as
-- `<root_harness>:<root_session_id>` (see `attribute_call`).
DROP TABLE calls;
CREATE TABLE calls (
harness INTEGER NOT NULL,
turn_id TEXT NOT NULL,
scope TEXT NOT NULL,
session_id TEXT NOT NULL,
usage_input INTEGER NOT NULL,
usage_output INTEGER NOT NULL,
usage_cache_read INTEGER NOT NULL,
usage_cache_write INTEGER NOT NULL,
usage_reasoning INTEGER NOT NULL,
PRIMARY KEY (harness, turn_id, scope)
);
-- Usage stored before was attributed by the old rule. `AiSessionDatabase::migrate` attributes it
-- afresh from the stored rows, in Rust with the same code rows arriving later go through, and
-- clears this once it has.
ALTER TABLE projection_state ADD COLUMN recount_calls INTEGER NOT NULL DEFAULT 0;
UPDATE projection_state SET recount_calls = 1;
-- The sessions whose root changed since `AiSessionDatabase` last attributed afresh the
-- content-derived calls they hold, which are scoped to that root: filled by the trigger, wherever
-- sessions are regrouped, and emptied in the same transaction.
CREATE TABLE regrouped (session INTEGER PRIMARY KEY);
CREATE TRIGGER sessions_regrouped AFTER UPDATE OF root_harness, root_session_id ON sessions
WHEN OLD.root_harness IS NOT NEW.root_harness OR OLD.root_session_id IS NOT NEW.root_session_id
BEGIN
INSERT OR IGNORE INTO regrouped (session) VALUES (NEW.id);
END;
crates/atuin-client/src/ai_session/migrations/0008_atuin_id.sql
-5
View File
@@ -1,5 +0,0 @@
ALTER TABLE sessions ADD COLUMN atuin_id BLOB;
-- Builds before record versions skip v1 records as undecodable: replay everything once.
DELETE FROM reproject_watermark;
UPDATE projection_state SET generation = generation + 1;
crates/atuin-client/src/settings.rs
+1 -1
View File
@@ -1251,7 +1251,7 @@ impl Settings {
/// read-only (`AiSessionDatabase::open_read_only`).
#[must_use]
pub fn ai_session_sidecar_path() -> PathBuf {
Self::effective_data_dir().join("ai_harness_sessions.db")
Self::effective_data_dir().join("ai_session_sidecar.db")
}
#[must_use]
crates/atuin-daemon/proto/ai/session.proto
-13
View File
@@ -49,16 +49,6 @@ message GetSessionEvent {
}
}
message GetTranscriptRequest {
ai.agent.HarnessSession session = 1;
}
// One chunk of a rendered transcript. Concatenate `chunk`s in arrival order to
// reconstruct the whole transcript.
message GetTranscriptChunk {
string chunk = 1;
}
message SearchSessionsRequest {
string query = 1;
// Maximum number of distinct sessions to return, most-relevant first. 0 means
@@ -163,9 +153,6 @@ service AiSession {
// response.
rpc GetSession(GetSessionRequest) returns (stream GetSessionEvent);
// Render one session as a plain-text transcript, streamed in chunks.
rpc GetTranscript(GetTranscriptRequest) returns (stream GetTranscriptChunk);
// Full-text search sessions (every term as a whole word, or any term as a
// prefix with any_term), streamed most-relevant-first (bm25 blended with
// recency). An empty query streams sessions newest first.
crates/atuin-daemon/src/client.rs
+3 -14
View File
@@ -26,9 +26,9 @@ use tracing::{Level, instrument, span};
use crate::grpc::ai::agent::pb::Session as AiSession;
use crate::grpc::ai::session::pb::ai_session_client::AiSessionClient as AiSessionServiceClient;
use crate::grpc::ai::session::pb::{
GetSessionEvent, GetSessionRequest, GetTranscriptChunk, GetTranscriptRequest,
ImportSessionsEvent, ImportSessionsRequest, ListSessionsRequest, RebuildSessionsRequest,
SearchSessionsMatch, SearchSessionsRequest, TailSessionsEvent, TailSessionsRequest,
GetSessionEvent, GetSessionRequest, ImportSessionsEvent, ImportSessionsRequest,
ListSessionsRequest, RebuildSessionsRequest, SearchSessionsMatch, SearchSessionsRequest,
TailSessionsEvent, TailSessionsRequest,
};
use crate::grpc::history::pb::history_client::HistoryClient as HistoryServiceClient;
use crate::grpc::history::pb::{
@@ -594,17 +594,6 @@ impl AiClient {
Ok(self.client.get_session(request).await?.into_inner())
}
/// Stream a rendered plain-text transcript in chunks; concatenate them in arrival order.
pub async fn get_transcript(
&mut self,
session: HarnessSession,
) -> Result<tonic::Streaming<GetTranscriptChunk>> {
let request = GetTranscriptRequest {
session: Some(session.into()),
};
Ok(self.client.get_transcript(request).await?.into_inner())
}
/// Follow sessions and messages as they are recorded. `harness` filters to one harness when set.
pub async fn tail_sessions(
&mut self,
crates/atuin-daemon/src/grpc/ai/session/mod.rs
+5 -24
View File
@@ -13,12 +13,11 @@ use tonic::{Code, Request, Response, Status};
use crate::grpc::ai::agent::pb as agent;
use crate::grpc::ai::session::pb::ai_session_server::AiSession as GrpcService;
use crate::grpc::ai::session::pb::{
GetSessionEvent, GetSessionRequest, GetTranscriptChunk, GetTranscriptRequest,
HarnessFilterRequest, ImportSessionsEvent, ImportSessionsProgress, ImportSessionsRequest,
ImportSessionsSummary, ListSessionsRequest, RebuildSessionsReply, RebuildSessionsRequest,
SearchSessionsMatch, SearchSessionsRequest, SessionFilterRequest, SessionRefRequest,
TailSessionsEvent, TailSessionsRequest, get_session_event, import_sessions_event,
tail_sessions_event,
GetSessionEvent, GetSessionRequest, HarnessFilterRequest, ImportSessionsEvent,
ImportSessionsProgress, ImportSessionsRequest, ImportSessionsSummary, ListSessionsRequest,
RebuildSessionsReply, RebuildSessionsRequest, SearchSessionsMatch, SearchSessionsRequest,
SessionFilterRequest, SessionRefRequest, TailSessionsEvent, TailSessionsRequest,
get_session_event, import_sessions_event, tail_sessions_event,
};
use crate::grpc::common::pb as common;
use crate::grpc::common::pb::Lagged;
@@ -91,8 +90,6 @@ pub fn is_rebuilding(status: &Status) -> bool {
impl GrpcService for Service {
type ListSessionsStream = Pin<Box<dyn Stream<Item = Result<agent::Session, Status>> + Send>>;
type GetSessionStream = Pin<Box<dyn Stream<Item = Result<GetSessionEvent, Status>> + Send>>;
type GetTranscriptStream =
Pin<Box<dyn Stream<Item = Result<GetTranscriptChunk, Status>> + Send>>;
type SearchSessionsStream =
Pin<Box<dyn Stream<Item = Result<SearchSessionsMatch, Status>> + Send>>;
type TailSessionsStream = Pin<Box<dyn Stream<Item = Result<TailSessionsEvent, Status>> + Send>>;
@@ -151,22 +148,6 @@ impl GrpcService for Service {
Ok(Response::new(Box::pin(stream)))
}
async fn get_transcript(
&self,
request: Request<GetTranscriptRequest>,
) -> Result<Response<Self::GetTranscriptStream>, Status> {
self.ensure_recovered()?;
let handle = request.into_inner().session()?;
let chunks = self.capture.transcript(&handle).map(|chunk| {
chunk
.map(|chunk| GetTranscriptChunk { chunk })
.map_err(|e| Status::internal(e.to_string()))
});
Ok(Response::new(Box::pin(chunks)))
}
async fn search_sessions(
&self,
request: Request<SearchSessionsRequest>,
crates/atuin-daemon/src/grpc/ai/session/pb.rs
-6
View File
@@ -221,12 +221,6 @@ impl SessionRefRequest for GetSessionRequest {
}
}
impl SessionRefRequest for GetTranscriptRequest {
fn session_ref(self) -> Option<agent::HarnessSession> {
self.session
}
}
#[cfg(test)]
mod tests {
use atuin_client::ai_session::NativeSessionId;
crates/atuin-daemon/src/session_capture/mod.rs
+15 -22
View File
@@ -842,13 +842,6 @@ impl AiHarnessSessionCapture {
self.sink.sidecar.messages(session)
}
pub fn transcript(
&self,
session: &HarnessSession,
) -> impl Stream<Item = Result<String, DbError>> + Send + 'static {
self.sink.sidecar.transcript(session)
}
pub fn search(
&self,
query: &str,
@@ -1443,14 +1436,11 @@ mod tests {
}
#[rstest]
#[case(None, "Reasoned")]
#[case(Some(185), "Reasoning · 185 tokens")]
#[case(Some(0), "Reasoning · 0 tokens")]
#[case(None)]
#[case(Some(185))]
#[case(Some(0))]
#[tokio::test]
async fn reasoning_metadata_survives_storage_and_rendering(
#[case] tokens: Option<u64>,
#[case] label: &str,
) {
async fn reasoning_metadata_survives_storage(#[case] tokens: Option<u64>) {
let sink = Sink::new(mem_store().await, AiSessionDatabase::in_memory().await.unwrap());
let mut msg = sample_message();
msg.role = Role::Assistant;
@@ -1468,8 +1458,6 @@ mod tests {
crate::grpc::ai::agent::pb::ReasoningSummary { tokens }
))
);
let mut transcript = Box::pin(rebuilt.transcript(&msg.session));
assert_eq!(transcript.next().await.unwrap().unwrap(), format!("assistant: {label}\n"));
}
/// A row of model call `turn` reporting `output` tokens, `reasoning` of them thinking.
@@ -2415,6 +2403,7 @@ mod tests {
#[rstest]
#[tokio::test]
async fn split_and_interleaved_calls_count_reasoning_once() {
use futures::TryStreamExt;
let sink = Sink::new(mem_store().await, AiSessionDatabase::in_memory().await.unwrap());
for (index, (turn, marker, output, reasoning)) in [
("a", true, 10, None),
@@ -2434,12 +2423,16 @@ mod tests {
sink.append(msg).await.unwrap();
}
assert_eq!(charged(&sink.sidecar).await, (1049, 227));
let transcript: Vec<String> =
sink.sidecar.transcript(&sample_handle()).map(Result::unwrap).collect().await;
assert_eq!(transcript[..2], [
"assistant: Reasoned\n".to_owned(),
"assistant: Reasoning · 42 tokens\n".to_owned()
]);
let markers: Vec<(Vec<Content>, Option<u64>)> = sink
.sidecar
.messages(&sample_handle())
.map(|msg| msg.map(|msg| (msg.content, msg.usage.and_then(|usage| usage.reasoning))))
.take(2)
.try_collect()
.await
.unwrap();
let marker = || vec![Content::ReasoningSummary { tokens: None }];
assert_eq!(markers, [(marker(), None), (marker(), Some(42))]);
}
#[rstest]
crates/atuin/tests/e2e_daemon.rs
+1 -1
View File
@@ -206,7 +206,7 @@ async fn stale_ai_session_index(env: &FreshEnv) -> PathBuf {
use atuin_common::harnesstools::session::{Content, Role};
std::fs::create_dir_all(env.data_dir()).unwrap();
let path = env.data_dir().join("ai_harness_sessions.db");
let path = env.data_dir().join("ai_session_sidecar.db");
let db = AiSessionDatabase::open(&path).await.unwrap();
let msg = Message::builder()
.id(atuin_domain::record::RecordId(atuin_common::utils::uuid_v7()))