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

fix(ai-session): harden capture, sync and resume across agents (atuinsh/atuin#4387)

Browse Source
This commit is contained in:
Ellie Huxtable
2026-10-11 23:31:56 +00:00
committed by Ellie Huxtable
parent 3e1c10b832
commit bbc2ab4a31
77 changed files with 8005 additions and 667 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Cargo.lock
Generated
+2
View File
@@ -520,6 +520,7 @@ dependencies = [
"serde",
"serde_json",
"serde_with",
"sha2 0.10.9",
"shlex",
"sqlx",
"strum 0.27.2",
@@ -531,6 +532,7 @@ dependencies = [
"tiny-bip39",
"tokio",
"tokio-util",
"toml_edit",
"tracing",
"tracing-subscriber",
"typed-builder",
Cargo.toml
+1
View File
@@ -138,6 +138,7 @@ serde = { version = "1", features = ["derive"] }
serde_json = { version = "1", features = ["preserve_order"] }
serde_regex = "1"
serde_with = "3"
sha2 = "0.10"
shellexpand = "3"
shlex = "2"
# held at 0.3: crossterm 0.29 pins signal-hook 0.3; bumping would add a 2nd copy
crates/atuin-ai/src/commands/resume.rs
+241 -50
View File
@@ -12,7 +12,9 @@
//! shell widget, so the command the widget puts on the command line works as it stands.
//! - `--in <harness>` continues the session the id names in another harness instead: it is
//! written out there as a new session (its tool calls flattened into notes; see
//! [`atuin_common::harnesstools::continuation`]) and that is resumed, the same way. In the
//! [`atuin_common::harnesstools::continuation`]) and that is resumed, the same way. A session
//! that went on separately on several machines is continued along one branch, which `--branch`
//! must name: without it, it fails listing them (the picker takes the newest, saying so). In the
//! picker, accepting a session asks where to resume it: its own harness first, then the others
//! installed here (see [`crate::resume_tui::chooser`]; `[ai.sessions] resume_chooser = false`
//! skips that).
@@ -29,8 +31,9 @@
//! same history, linked to it as a fork (see [`atuin_common::harnesstools::fork`]), and that is
//! resumed. The original is left as it is. A session that went on separately on several
//! machines forks from its newest branch.
//! - `--branch` picks the branch to catch up to, to switch to, or to fork from: `this`,
//! `@<host id>`, or the start of its head's id.
//! - `--branch` picks the branch to catch up to, to switch to, to fork from, or to continue in
//! another harness: `this`, `@<host id>`, or the start of its head's id. `--switch` with several
//! branches to switch to, and `--in` with several branches, need it.
//! - `--as-is`, `--switch` and `--branch` need an id naming a single session that can resume here:
//! anything else is an error (listing the sessions an id names), never the picker, which would
//! drop them.
@@ -45,7 +48,7 @@
use std::io::{self, IsTerminal, Write};
use std::sync::Arc;
use atuin_client::ai_session::{Analysis, AtuinSessionId, HarnessKind, SourceId};
use atuin_client::ai_session::{Analysis, AtuinSessionId, HarnessKind, HarnessSession, SourceId};
use atuin_client::settings::{AiSessionFilterMode, Settings};
use atuin_client::theme::ThemeManager;
use atuin_common::string::EscapeNonPrintablePosixExt as _;
@@ -87,7 +90,8 @@ pub struct Cmd {
print: bool,
/// Continue the session QUERY names (by id, or a unique id prefix) in another agent: it is
/// written out as a new session there, tool calls flattened into notes, and resumed.
/// written out as a new session there, tool calls flattened into notes, and resumed. One that
/// went on separately on several machines needs `--branch`.
#[arg(long = "in", value_enum, value_name = "HARNESS")]
continue_in: Option<ContinueIn>,
@@ -103,7 +107,7 @@ pub struct Cmd {
as_is: bool,
/// Switch this machine's copy of the session QUERY names to another branch (the one
/// `--branch` names, else the newest it can be switched to), in place, and resume it: the
/// `--branch` names, else the only one it can be switched to), in place, and resume it: the
/// history it shares with that branch is kept as it is, and the branch's messages since are
/// appended from sync. The branch it was on stays in atuin, for `--fork` to bring back, and
/// the copy as it was is kept as a backup. Only when sync holds all of it and no agent here
@@ -111,10 +115,11 @@ pub struct Cmd {
#[arg(long, conflicts_with_all = ["fork", "continue_in", "as_is"])]
switch: bool,
/// The branch of the session QUERY names to catch up to, to switch to with `--switch`, or to
/// fork from with `--fork`: `this`, `@<host id>`, or the start of its head's id. The newest
/// by default.
#[arg(long, value_name = "BRANCH", conflicts_with_all = ["continue_in", "as_is"])]
/// The branch of the session QUERY names to catch up to, to switch to with `--switch`, to
/// fork from with `--fork`, or to continue with `--in`: `this`, `@<host id>`, or the start of
/// its head's id. The newest by default; `--in`, and `--switch` with several branches to
/// switch to, need it named.
#[arg(long, value_name = "BRANCH", conflicts_with = "as_is")]
branch: Option<String>,
/// The filter the picker opens in (default: workspace, widening to global).
@@ -219,6 +224,7 @@ async fn direct_match(source: &dyn SessionSource, query: &str) -> Result<Option<
if let [exact] = exact.as_slice() {
return Ok(Some(found.swap_remove(*exact)));
}
let mut found = resumable_first(found);
Ok(if found.len() == 1 {
found.pop()
} else {
@@ -226,6 +232,15 @@ async fn direct_match(source: &dyn SessionSource, query: &str) -> Result<Option<
})
}
/// The sessions an id names, `found`, but for subagents named beside the session they work for:
/// one stands for that session (see [`direct_target`]), so it is not another to pick.
fn resumable_first(mut found: Vec<SessionRow>) -> Vec<SessionRow> {
let named: Vec<HarnessSession> = found.iter().map(|r| r.handle.clone()).collect();
found
.retain(|r| r.relation.is_listed() || r.parent.as_ref().is_none_or(|p| !named.contains(p)));
found
}
/// Whether the picker has a terminal to run on: stdout, or `/dev/tty` when stdout is captured.
fn picker_has_a_terminal() -> bool {
io::stdout().is_terminal()
@@ -282,19 +297,29 @@ fn be_more_specific(
/// The session named, `row`, can't be resumed here, for `why`: the picker opens on it instead, so
/// the reason shows (and another can be picked). With `flag` (`--switch`, `--branch`, `--as-is`)
/// given, which the picker would drop, an error saying why.
/// given, which the picker would drop, or without a `terminal` to open it on, an error saying
/// why.
fn picker_on(
row: SessionRow,
why: NotResumable,
flag: Option<&str>,
terminal: bool,
) -> Result<(SessionRow, NotResumable)> {
if let Some(flag) = flag {
let message = format!("can't resume {} with `{flag}`: {why}", row.handle.session);
bail!("{}", message.escape_non_printable());
let with = flag.map(|flag| format!(" with `{flag}`")).unwrap_or_default();
if flag.is_some() || !terminal {
bail!("{}", escaped(&format!("can't resume {}{with}: {why}", row.handle.session)));
}
Ok((row, why))
}
/// `message` without the control characters an agent-supplied part of it (an id, an agent's
/// stderr) may carry, line by line: its lines stay lines.
fn escaped(message: &str) -> String {
let lines: Vec<String> =
message.lines().map(|line| line.escape_non_printable().to_string()).collect();
lines.join("\n")
}
/// With `flag` (`--switch`, `--branch`, `--as-is`) given and `query` naming no single session, the
/// error saying to be more specific ([`be_more_specific`]).
async fn no_single_session_for(
@@ -306,7 +331,7 @@ async fn no_single_session_for(
return Ok(());
};
let found = if looks_like_id(query) {
source.find_by_id(query).await?
resumable_first(source.find_by_id(query).await?)
} else {
Vec::new()
};
@@ -325,7 +350,7 @@ async fn ambiguous_matches(
if !looks_like_id(query) {
return Ok(Vec::new());
}
let found = source.find_by_id(query).await?;
let found = resumable_first(source.find_by_id(query).await?);
if found.len() < 2 {
return Ok(Vec::new());
}
@@ -396,12 +421,10 @@ pub async fn run(cmd: Cmd, settings: &Settings) -> Result<()> {
if cmd.fork || cmd.continue_in.is_some() {
let (source, resumer, query) = (source.as_ref(), resumer.as_ref(), query.trim());
let (branch, here) = (cmd.branch.as_deref(), context.host_id.as_str());
let (plan, status) = match cmd.continue_in {
Some(into) => continue_plan(source, resumer, query, into).await?,
None => {
let branch = cmd.branch.as_deref();
fork_plan(source, resumer, query, branch, &context.host_id).await?
}
Some(into) => continue_plan(source, resumer, query, into, branch, here).await?,
None => fork_plan(source, resumer, query, branch, here).await?,
};
// The widget reads stderr for the command: nothing else may go there.
if output != Output::Widget {
@@ -450,7 +473,7 @@ pub async fn run(cmd: Cmd, settings: &Settings) -> Result<()> {
}
return finish(direct_outcome(plan, output, settings.enter_accept), output);
}
Err(why) => preselect = Some(picker_on(row, why, flag)?),
Err(why) => preselect = Some(picker_on(row, why, flag, picker_has_a_terminal())?),
}
}
@@ -542,9 +565,7 @@ async fn pick_head(
Ok(head) => Ok(head.source_id.clone()),
Err(why) => {
// Line by line: the listing of branches keeps its lines.
let lines: Vec<String> =
why.lines().map(|l| l.escape_non_printable().to_string()).collect();
bail!("--branch: {}", lines.join("\n"))
bail!("--branch: {}", escaped(&why))
}
}
}
@@ -568,9 +589,12 @@ fn choice_note(row: &SessionRow, held: &Held, here: &str) -> String {
branch.host
));
}
// Without `--branch`, `--switch` goes to the newest branch it can.
if let Some((n, branch)) = held.branches.iter().enumerate().find(|(_, b)| b.switch) {
let pick = if n == 0 {
// The branch `--branch` named, when the copy can be switched to it; else the first it can.
// Without `--branch`, `--switch` goes to the only branch it can, and refuses to choose
// between several.
let named = held.branches.get(held.chosen).filter(|b| b.switch && held.chosen != 0);
if let Some(branch) = named.or_else(|| held.switches().next()) {
let pick = if named.is_none() && held.switches().count() == 1 {
String::new()
} else {
format!(" --branch {}", quote(&branch.selector).escape_non_printable())
@@ -592,19 +616,55 @@ fn choice_note(row: &SessionRow, held: &Held, here: &str) -> String {
}
/// `atuin ai resume <id> --in <harness>`: write the session `query` names out as a new session
/// of `into`, and the plan that resumes it, with the status line saying what was flattened.
/// of `into`, along the branch `branch` names, and the plan that resumes it, with the status line
/// saying what was flattened. A session that went on separately on several machines needs one
/// named: it is an error listing them. A session of `into` already is resumed, as without
/// `--in`, saying so.
async fn continue_plan(
source: &dyn SessionSource,
resumer: &dyn Resumer,
query: &str,
into: ContinueIn,
branch: Option<&str>,
here: &str,
) -> Result<(ResumePlan, String)> {
let Some((row, redirected)) = direct_target(source, query).await? else {
bail!("`--in` continues the session an id names: no single session has the id {query:?}");
};
let continued = resumer.continue_in(source, &row, into.kind()).await.map_err(|why| {
let message = format!("can't continue {}: {why}", row.handle.session);
eyre::eyre!("{}", message.escape_non_printable())
let kind = row.handle.harness;
if into.kind() == kind {
let resumed = direct_plan(source, resumer, &row, Direct::CatchUp, branch, here).await?;
let (plan, said) = resumed.map_err(|why| {
eyre::eyre!("{}", escaped(&format!("can't resume {}: {why}", row.handle.session)))
})?;
let mut status = format!("it's a {} session already: resuming it", harness_label(kind));
for said in redirected.into_iter().chain(said) {
status.push_str(&format!("; {said}"));
}
return Ok((plan, status));
}
let head = match branch {
Some(selector) => Some(pick_head(source, &row, selector, here).await?),
None => {
// Which of several branches to carry on is the user's to say, as for `--switch`.
let analysis = source.analyse(&row.handle).await?;
if let Some(analysis) = analysis.filter(Analysis::diverged) {
let (heads, now) = (analysis.heads(), time::OffsetDateTime::now_utc());
let message = format!(
"can't continue {}: it went {} ways: name the branch to continue with \
`--branch`:{}",
row.handle.session,
heads.len(),
catchup::listing(heads, heads, now, here)
);
bail!("{}", escaped(&message));
}
None
}
};
let continued = resumer.continue_in(source, &row, into.kind(), head.as_ref()).await;
let continued = continued.map_err(|why| {
eyre::eyre!("{}", escaped(&format!("can't continue {}: {why}", row.handle.session)))
})?;
let mut status = continued.status();
if let Some(redirected) = redirected {
@@ -635,8 +695,7 @@ async fn fork_plan(
let from = catchup::fork_from(source, &row.handle, head.as_ref()).await;
let forked = async { resumer.fork(source, &row, from?).await };
let forked = forked.await.map_err(|why| {
let message = format!("can't fork {}: {why}", row.handle.session);
eyre::eyre!("{}", message.escape_non_printable())
eyre::eyre!("{}", escaped(&format!("can't fork {}: {why}", row.handle.session)))
})?;
let mut status = forked.status();
if let Some(redirected) = redirected {
@@ -815,6 +874,43 @@ mod tests {
);
}
/// A subagent an id names beside the session it works for is not counted apart from it;
/// named alone, it stands for that session.
#[rstest]
#[tokio::test]
async fn a_subagent_named_beside_sessions_is_left_out() {
use atuin_client::ai_session::HarnessKind;
use crate::resume_tui::source::Relation;
let root = fake::row(HarnessKind::Pi, "01a11e4a-24bd", "the work");
let other = fake::row(HarnessKind::Pi, "01a11e47-485a", "other work");
let mut subagent = fake::row(HarnessKind::Pi, "01a11e4e-af40", "a part of it");
subagent.relation = Relation::Subagent;
subagent.parent = Some(root.handle.clone());
let source = FakeSource::from_rows(vec![root.clone(), other, subagent.clone()]);
let err = ambiguous_matches(&source, "01a11e4", false).await.unwrap_err().to_string();
assert!(err.starts_with("01a11e4 names 2 sessions"), "{err}");
assert!(!err.contains("01a11e4e"), "{err}");
let err = no_single_session_for(&source, "01a11e4", Some("--as-is")).await.unwrap_err();
assert!(!err.to_string().contains("01a11e4e"), "{err}");
let source = FakeSource::from_rows(vec![root.clone(), subagent.clone()]);
let named = direct_match(&source, "01a11e4").await.unwrap().unwrap();
assert_eq!(named.handle, root.handle);
let source = FakeSource::from_rows(vec![subagent.clone()]);
let named = direct_match(&source, "01a11e4").await.unwrap().unwrap();
assert_eq!(named.handle, subagent.handle);
// Its parent not named, the subagent is another session the id names, never dropped
// for an unrelated one.
let unrelated = fake::row(HarnessKind::Pi, "01a11e4f-0000", "unrelated");
let source = FakeSource::from_rows(vec![unrelated, subagent]);
assert_eq!(direct_match(&source, "01a11e4").await.unwrap(), None);
assert_eq!(ambiguous_matches(&source, "01a11e4", true).await.unwrap().len(), 2);
}
/// `--branch` and `--as-is` only resume the session an id names: with an id naming several,
/// or none, it is an error listing them, never the picker, which would drop the flag.
#[rstest]
@@ -857,20 +953,28 @@ mod tests {
}
/// A session named that can't be resumed opens the picker on it, but not with `--branch` or
/// `--as-is`, which the picker would drop: an error saying why.
/// `--as-is`, which the picker would drop, nor without a terminal to open it on: an error
/// saying why.
#[rstest]
#[case::plain(None, None)]
#[case::branch(Some("--branch"), Some("can't resume s with `--branch`: catching it up"))]
#[case::as_is(Some("--as-is"), Some("can't resume s with `--as-is`: catching it up"))]
#[case::plain(None, true, None)]
#[case::branch(Some("--branch"), true, Some("can't resume s with `--branch`: catching it up"))]
#[case::as_is(Some("--as-is"), true, Some("can't resume s with `--as-is`: catching it up"))]
#[case::no_terminal(None, false, Some("can't resume s: catching it up"))]
#[case::no_terminal_with_a_flag(
Some("--as-is"),
false,
Some("can't resume s with `--as-is`: catching it up")
)]
fn a_session_that_cant_resume_keeps_its_flags(
#[case] flag: Option<&str>,
#[case] terminal: bool,
#[case] err: Option<&str>,
) {
use atuin_client::ai_session::HarnessKind;
let row = fake::row(HarnessKind::ClaudeCode, "s", "t");
let why = NotResumable::CatchUp("no longer a head".to_owned());
match (picker_on(row, why, flag), err) {
match (picker_on(row, why, flag, terminal), err) {
(Ok((row, _)), None) => assert_eq!(row.handle.session.as_ref(), "s"),
(Err(e), Some(want)) => assert!(e.to_string().starts_with(want), "{e}"),
(got, _) => panic!("{got:?}"),
@@ -925,7 +1029,9 @@ mod tests {
let resumer = FakeResumer::default();
let (plan, status) =
continue_plan(&source, &resumer, query, ContinueIn::Codex).await.unwrap();
continue_plan(&source, &resumer, query, ContinueIn::Codex, None, fake::THIS_HOST_ID)
.await
.unwrap();
assert_eq!(plan.args[1], format!("continued-{target}"));
assert_eq!(status.contains("is a subagent"), redirected, "{status}");
}
@@ -1012,13 +1118,24 @@ mod tests {
async fn continuing_by_id_plans_the_new_session() {
let source = FakeSource::new();
let resumer = FakeResumer::default();
let (plan, status) =
continue_plan(&source, &resumer, "7f3c9a12", ContinueIn::Codex).await.unwrap();
let (plan, status) = continue_plan(
&source,
&resumer,
"7f3c9a12",
ContinueIn::Codex,
None,
fake::THIS_HOST_ID,
)
.await
.unwrap();
assert_eq!(plan.program, "codex");
assert_eq!(plan.args, ["resume", "continued-7f3c9a12-5be0-4d7e-9c41-0a8e2b6f4d10"]);
assert_eq!(status, "continuing in Codex: 42 tool calls become notes, reasoning dropped");
let err = continue_plan(&source, &resumer, "fix flaky", ContinueIn::Pi).await.unwrap_err();
let err =
continue_plan(&source, &resumer, "fix flaky", ContinueIn::Pi, None, fake::THIS_HOST_ID)
.await
.unwrap_err();
assert!(err.to_string().contains("no single session has the id"), "{err}");
let cli = Cli::try_parse_from(["resume", "7f3c9a12", "--in", "opencode"]).unwrap();
@@ -1026,6 +1143,70 @@ mod tests {
assert!(Cli::try_parse_from(["resume", "7f3c9a12", "--in", "cursor"]).is_err());
}
/// An agent's error keeps its lines, each without control characters.
#[rstest]
fn errors_keep_their_lines() {
let row = fake::row(HarnessKind::Opencode, "ses_1", "t");
let why =
NotResumable::Restore("opencode failed: ERROR\n Unrecognized \x1b[2Jflag".into());
let err = picker_on(row, why, None, false).unwrap_err().to_string();
assert_eq!(
err,
"can't resume ses_1: restoring it from sync failed: opencode failed: ERROR\n \
Unrecognized ^[[2Jflag"
);
}
/// `--in` continues a session that went on separately on several machines along the branch
/// `--branch` names; without one, it continues nothing, listing the branches to name.
#[rstest]
#[case::by_host(Some("@00000002"), Ok("d"))]
#[case::by_head(Some("y"), Ok("y"))]
#[case::none_named(None, Err("it went 2 ways: name the branch to continue with `--branch`:"))]
#[tokio::test]
async fn continuing_by_id_follows_a_branch(
#[case] branch: Option<&str>,
#[case] want: Result<&str, &str>,
) {
let (source, _) = branched().await;
let resumer = FakeResumer::default();
let here = fake::THIS_HOST_ID;
let continued =
continue_plan(&source, &resumer, "7f3c9a12", ContinueIn::Pi, branch, here).await;
match want {
Ok(head) => {
continued.unwrap();
assert_eq!(*resumer.continues.lock(), [Some(SourceId::from(head.to_owned()))]);
}
Err(why) => {
let err = continued.unwrap_err().to_string();
assert!(err.contains(why), "{err}");
assert!(err.contains("\n y this machine · ") && err.contains("\n d @00000002"));
assert!(resumer.continues.lock().is_empty());
}
}
let cli = Cli::try_parse_from(["resume", "7f3c9a12", "--in", "pi", "--branch", "d"]);
assert_eq!(cli.unwrap().cmd.branch.as_deref(), Some("d"));
}
/// `--in` the session's own agent resumes it, as without `--in`, saying so.
#[rstest]
#[tokio::test]
async fn continuing_in_its_own_agent_resumes_it() {
let source = FakeSource::new();
let resumer = FakeResumer::default();
let here = fake::THIS_HOST_ID;
let (plan, status) =
continue_plan(&source, &resumer, "7f3c9a12", ContinueIn::Claude, None, here)
.await
.unwrap();
let row = direct_match(&source, "7f3c9a12").await.unwrap().unwrap();
assert_eq!(plan, resumer.plan(&row).await.unwrap().plan);
assert_eq!(status, "it's a Claude Code session already: resuming it");
assert!(resumer.continues.lock().is_empty());
}
/// `atuin ai resume <id> --fork` forks the session the id names and plans resuming the fork,
/// with `--print` too; it can't go with `--in`.
#[rstest]
@@ -1143,7 +1324,7 @@ mod tests {
#[case::fork_from_a_branch(&["--fork", "--branch", "d", "--print"], true)]
#[case::as_is_and_fork(&["--as-is", "--fork"], false)]
#[case::as_is_and_branch(&["--as-is", "--branch", "d"], false)]
#[case::branch_and_in(&["--branch", "d", "--in", "pi"], false)]
#[case::branch_and_in(&["--branch", "d", "--in", "pi"], true)]
fn catch_up_flags(#[case] flags: &[&str], #[case] ok: bool) {
let args = ["resume", "7f3c9a12"].iter().chain(flags);
let cli = Cli::try_parse_from(args);
@@ -1157,7 +1338,7 @@ mod tests {
}
/// `--switch` switches the copy here of the session an id names to the branch `--branch`
/// names (else the newest it can), saying so; one that can't be switched is an error saying
/// names (else the only one it can), saying so; one that can't be switched is an error saying
/// why. A choice catching up needs names `--switch` (with the branch it would take) where
/// the copy can be switched.
#[rstest]
@@ -1183,7 +1364,7 @@ mod tests {
let refused = direct_plan(&source, &fresh, &row, Direct::Switch, None, here);
let why = refused.await.unwrap().unwrap_err();
assert_eq!(why, NotResumable::Switch("no head named".to_owned()));
let err = picker_on(row.clone(), why, Some("--switch")).unwrap_err().to_string();
let err = picker_on(row.clone(), why, Some("--switch"), true).unwrap_err().to_string();
assert!(err.starts_with("can't resume 7f3c9a12"), "{err}");
assert!(err.contains("with `--switch`: switching it to another branch failed"), "{err}");
@@ -1191,7 +1372,7 @@ mod tests {
let copy = fake::local_tip(&["a", "b", "x", "y"], Some("y"));
let mut branches = branches(&analysis, Some(&copy), here);
branches[1].switch = true;
let held = Held {
let mut held = Held {
why: Why::Diverged,
harness: HarnessKind::ClaudeCode,
plan: FakeResumer::default().plan(&row).await.unwrap().plan,
@@ -1200,10 +1381,20 @@ mod tests {
};
let note = choice_note(&row, &held, here);
let id = row.handle.session.as_ref();
let want = format!(
", `atuin ai resume {id} --switch --branch d` switches this copy to @00000002's branch"
);
let want =
format!(", `atuin ai resume {id} --switch` switches this copy to @00000002's branch");
assert!(note.contains(&want), "{note}");
// Of several, `--switch` takes none: the newest is named.
held.branches[0].switch = true;
let note = choice_note(&row, &held, here);
let want = format!(", `atuin ai resume {id} --switch --branch y` switches this copy to");
assert!(note.contains(&want), "{note}");
// With a branch named (`--branch`), the switch it offers is to that one, as the fork.
held.chosen = 1;
let note = choice_note(&row, &held, here);
let want = format!(", `atuin ai resume {id} --switch --branch d` switches this copy to");
assert!(note.contains(&want), "{note}");
assert!(note.contains(&format!("`atuin ai resume {id} --fork --branch d`")), "{note}");
}
#[rstest]
crates/atuin-ai/src/commands/session.rs
+127 -16
View File
@@ -9,7 +9,8 @@ use std::collections::HashMap;
use std::io::{self, IsTerminal, Write};
use atuin_client::ai_session::{
HarnessKind, HarnessSession, Message, SearchTerms, Session, SessionFilter, SessionMatch,
AtuinSessionId, HarnessKind, HarnessSession, Message, SearchTerms, Session, SessionFilter,
SessionMatch,
};
use atuin_client::settings::Settings;
use atuin_common::harnesstools::session::model::reasoning_label;
@@ -51,6 +52,7 @@ enum SubCmd {
Search {
#[arg(
value_name = "QUERY",
allow_hyphen_values = true,
help = "Words to look for in session titles, message text, reasoning, and tool calls"
)]
query: String,
@@ -68,7 +70,10 @@ enum SubCmd {
/// Follow sessions and messages as they are recorded (until interrupted).
Tail,
/// Import every session the agents have on disk, such as ones from before capture was on.
/// What is already captured is not stored again.
Import {
/// Only import sessions from this harness.
#[arg(long, value_parser = parse_harness)]
harness: Option<HarnessKind>,
},
@@ -551,6 +556,12 @@ pub enum SelectError {
NotFound(String),
#[error("id `{0}` matches more than one session; use a longer or full id")]
Ambiguous(String),
/// Several agents' sessions have the id (a session continued in another agent can keep it).
#[error("id `{id}` names a session in more than one agent; use its atuin id:{choices}")]
Shared {
id: String,
choices: String,
},
}
/// Pure selector logic, split out from the RPC so it can be tested directly.
@@ -564,14 +575,26 @@ pub fn select_session(
}
// `list` prints ids truncated to 12 chars, so accept a unique id prefix as well as a full id.
let mut matches = sessions.into_iter().filter(|s| {
s.handle.session.as_ref().starts_with(selector) || s.atuin_id.has_prefix(selector)
});
let first = matches.next().ok_or_else(|| SelectError::NotFound(selector.to_owned()))?;
if matches.next().is_some() {
return Err(SelectError::Ambiguous(selector.to_owned()));
// An id given whole names its session, whatever longer ids it is a prefix of.
let atuin_id = selector.parse::<AtuinSessionId>().ok();
let (exact, prefixed): (Vec<Session>, Vec<Session>) = sessions
.into_iter()
.filter(|s| {
s.handle.session.as_ref().starts_with(selector) || s.atuin_id.has_prefix(selector)
})
.partition(|s| s.handle.session.as_ref() == selector || atuin_id == Some(s.atuin_id));
match (exact.as_slice(), prefixed.as_slice()) {
([one], _) | ([], [one]) => Ok(one.handle.clone()),
([], []) => Err(SelectError::NotFound(selector.to_owned())),
([], _) => Err(SelectError::Ambiguous(selector.to_owned())),
(several, _) => Err(SelectError::Shared {
id: selector.to_owned(),
choices: several
.iter()
.map(|s| format!("\n {:<12} {}", harness_name(s.handle.harness), s.atuin_id))
.collect(),
}),
}
Ok(first.handle)
}
// --- human rendering ----------------------------------------------------------------------------
@@ -617,7 +640,12 @@ fn write_session_header(out: &mut dyn Write, s: &Session) -> io::Result<()> {
writeln!(out)
}
/// `m` as `show` prints it: nothing for a row with nothing to read (an attachment, session
/// context, usage alone), which JSON output still holds.
fn write_message_text(out: &mut dyn Write, m: &Message) -> io::Result<()> {
if message_summary(m).is_none() {
return Ok(());
}
writeln!(out, "── {} · {} ──", role_name(&m.role), age(m.timestamp))?;
for block in &m.content {
match block {
@@ -635,7 +663,7 @@ fn write_message_text(out: &mut dyn Write, m: &Message) -> io::Result<()> {
Content::ToolUse(tu) => {
write!(out, "[tool-call {}]", sanitize(&tu.name))?;
if !tu.input.is_null() {
write!(out, " {}", sanitize(&tu.input.to_string()))?;
write!(out, " {}", sanitize(&input_text(&tu.input)))?;
}
writeln!(out)?;
}
@@ -662,6 +690,12 @@ fn write_message_text(out: &mut dyn Write, m: &Message) -> io::Result<()> {
writeln!(out)
}
/// A call's input as text: a string (a Codex call's arguments, a freeform tool's input) as it
/// is, any other JSON as encoded.
fn input_text(input: &serde_json::Value) -> Cow<'_, str> {
input.as_str().map_or_else(|| Cow::Owned(input.to_string()), Cow::Borrowed)
}
fn short_id(session_id: &str) -> &str {
session_id.get(..12).unwrap_or(session_id)
}
@@ -922,8 +956,11 @@ fn parse_harness(value: &str) -> Result<HarnessKind, String> {
match value {
"claude-code" => Ok(HarnessKind::ClaudeCode),
"codex" => Ok(HarnessKind::Codex),
"opencode" => Ok(HarnessKind::Opencode),
"pi" => Ok(HarnessKind::Pi),
other => Err(format!("unknown harness `{other}` (expected claude-code, codex, or pi)")),
other => {
Err(format!("unknown harness `{other}` (expected claude-code, codex, opencode, or pi)"))
}
}
}
@@ -1200,7 +1237,7 @@ fn content_json(block: &Content, reasoning: Option<u64>) -> Option<ContentJson>
input: if tu.input.is_null() {
String::new()
} else {
tu.input.to_string()
input_text(&tu.input).into_owned()
},
},
Content::ToolResult(tr) => ContentJson::ToolResult {
@@ -1277,6 +1314,7 @@ mod tests {
#[rstest]
#[case("claude-code", HarnessKind::ClaudeCode)]
#[case("codex", HarnessKind::Codex)]
#[case("opencode", HarnessKind::Opencode)]
#[case("pi", HarnessKind::Pi)]
fn parse_harness_maps_names(#[case] input: &str, #[case] want: HarnessKind) {
assert_eq!(parse_harness(input).unwrap(), want);
@@ -1284,7 +1322,7 @@ mod tests {
#[rstest]
fn parse_harness_rejects_unknown_harnesses() {
assert!(parse_harness("opencode").is_err());
assert!(parse_harness("copilot").is_err());
}
#[rstest]
@@ -1353,11 +1391,53 @@ mod tests {
assert!(select_session(sessions, "zzz").is_err());
}
/// A query starting with `-` is a query, after `--` or not; flags still parse around it.
#[rstest]
fn ambiguous_id_across_harnesses_is_an_error() {
let sessions =
vec![session(HarnessKind::Codex, "dup"), session(HarnessKind::ClaudeCode, "dup")];
assert!(select_session(sessions, "dup").is_err());
#[case::bare(&["search", "-v"])]
#[case::after_dashes(&["search", "--", "-v"])]
#[case::before_a_flag(&["search", "-v", "--limit", "3"])]
#[case::after_a_flag(&["search", "--limit", "3", "-v"])]
fn a_search_query_may_start_with_a_hyphen(#[case] args: &[&str]) {
#[derive(clap::Parser)]
struct Cli {
#[command(flatten)]
cmd: Cmd,
}
let cli = <Cli as clap::Parser>::try_parse_from(
std::iter::once("session").chain(args.iter().copied()),
)
.unwrap();
let SubCmd::Search { query, .. } = cli.cmd.cmd else {
panic!("not a search");
};
assert_eq!(query, "-v");
}
#[rstest]
fn an_exact_id_wins_over_the_ids_it_prefixes() {
let sessions = vec![
session(HarnessKind::Codex, "abc"),
session(HarnessKind::ClaudeCode, "abcdef"),
session(HarnessKind::Pi, "abcxyz"),
];
assert_eq!(select_session(sessions, "abc").unwrap(), handle(HarnessKind::Codex, "abc"));
}
/// An id several agents' sessions have names none of them, and says how to name each.
#[rstest]
fn an_id_shared_across_harnesses_lists_their_atuin_ids() {
let sessions = vec![
session(HarnessKind::Codex, "dup"),
session(HarnessKind::ClaudeCode, "dup"),
session(HarnessKind::Pi, "dupe"),
];
let ids: Vec<String> = sessions.iter().map(|s| s.atuin_id.to_string()).collect();
let err = select_session(sessions.clone(), "dup").unwrap_err().to_string();
assert!(err.contains(&ids[0]) && err.contains(&ids[1]), "{err}");
assert!(!err.contains(&ids[2]), "{err}");
let picked = select_session(sessions, &ids[1]).unwrap();
assert_eq!(picked, handle(HarnessKind::ClaudeCode, "dup"));
}
#[rstest]
@@ -1526,6 +1606,37 @@ mod tests {
assert!(message_summary(&blank).is_none());
}
#[rstest]
#[case::empty(vec![], false)]
#[case::blank(vec![Content::Text(" \n".to_owned())], false)]
#[case::raw(vec![Content::Other(json!({"type": "attachment"}))], false)]
#[case::text(vec![Content::Text("hi".to_owned())], true)]
#[case::uncaptured_call(vec![tool_call("Bash")], true)]
fn show_prints_only_rows_with_something_to_read(
#[case] content: Vec<Content>,
#[case] printed: bool,
) {
let mut out = Vec::new();
write_message_text(&mut out, &msg(Role::Other(String::new()), content)).unwrap();
assert_eq!(!out.is_empty(), printed, "{}", String::from_utf8_lossy(&out));
}
/// A call's text input and a result's text blocks read as text, not as escaped JSON.
#[rstest]
fn show_prints_tool_text_as_text() {
let call = Content::ToolUse(ToolUse {
id: ToolCallId::from("c1".to_owned()),
name: "exec".to_owned(),
input: json!("const out = \"hi\";"),
});
let result = tool_result(json!([{"type": "text", "text": "line 1\nline 2"}]), false);
let mut out = Vec::new();
write_message_text(&mut out, &msg(Role::Assistant, vec![call, result])).unwrap();
let out = String::from_utf8(out).unwrap();
assert!(out.contains("[tool-call exec] const out = \"hi\";\n"), "{out}");
assert!(out.contains("[tool-result] line 1\nline 2\n"), "{out}");
}
#[rstest]
fn summary_marks_thinking() {
let m = msg(Role::Assistant, vec![Content::ReasoningSummary { tokens: Some(7) }]);
crates/atuin-ai/src/resume_tui/catchup.rs
+55 -22
View File
@@ -222,12 +222,16 @@ const OPENCODE_SWITCH: &str = "opencode keeps its sessions in a database, which
const UNSYNCED_SWITCH: &str = "this copy has messages sync hasn't got: switching would lose them";
/// The head of `analysis` to switch the copy of a session of `harness` read as `tip` to: `named`,
/// else the newest it can be switched to ([`can_switch`]); or why it can't be.
/// else the only one it can be switched to ([`can_switch`]); or why it can't be. With several to
/// choose from and none named, which to take is the user's to say: an error listing them, as
/// `now` and on host `here` ([`describe`]).
pub fn switch_to<'a>(
analysis: &'a Analysis,
harness: HarnessKind,
named: Option<&SourceId>,
tip: &LocalTip,
now: OffsetDateTime,
here: &str,
) -> Result<&'a Head, String> {
if harness == HarnessKind::Opencode {
return Err(OPENCODE_SWITCH.to_owned());
@@ -244,10 +248,24 @@ pub fn switch_to<'a>(
.iter()
.find(|h| h.source_id == *named)
.ok_or_else(|| format!("{named} is no longer one of its branches' heads"))?,
None => heads
.iter()
.find(|h| can_switch(analysis, harness, &h.source_id, tip))
.ok_or("this copy is on the session's only line: resuming it catches it up")?,
None => {
let can: Vec<&Head> =
heads.iter().filter(|h| can_switch(analysis, harness, &h.source_id, tip)).collect();
match can.as_slice() {
[] => {
let why = "this copy is on the session's only line: resuming it catches it up";
return Err(why.to_owned());
}
[head] => *head,
several => {
return Err(format!(
"it went {} other ways: name the branch to switch to with `--branch`:{}",
several.len(),
listing(heads, several.iter().copied(), now, here)
));
}
}
}
};
match classify(analysis, harness, &head.source_id, tip) {
Step::Choice(Why::Diverged) => Ok(head),
@@ -391,28 +409,39 @@ pub fn pick_branch<'a>(
} else {
heads.iter().filter(|h| h.source_id.as_ref().starts_with(selector)).collect()
};
let listing = || {
let lines = heads.iter().enumerate().map(|(n, h)| {
let latest = if n == 0 {
" (latest)"
} else {
""
};
format!("\n {} {}{latest}", branch_selector(heads, h), describe(h, now, here))
});
lines.collect::<String>()
};
match found.as_slice() {
[head] => Ok(head),
[] => Err(format!("no branch is {selector:?}; its branches:{}", listing())),
[] => Err(format!(
"no branch is {selector:?}; its branches:{}",
listing(heads, heads, now, here)
)),
several => Err(format!(
"{selector:?} names {} branches; name one by id:{}",
several.len(),
listing()
listing(heads, heads, now, here)
)),
}
}
/// `shown`, some of `heads`, a line each: what `--branch` names it by, and [`describe`]'s line,
/// the newest of `heads` marked so.
pub fn listing<'a>(
heads: &[Head],
shown: impl IntoIterator<Item = &'a Head>,
now: OffsetDateTime,
here: &str,
) -> String {
let lines = shown.into_iter().map(|h| {
let latest = if heads.first() == Some(h) {
" (latest)"
} else {
""
};
format!("\n {} {}{latest}", branch_selector(heads, h), describe(h, now, here))
});
lines.collect()
}
/// What forking `session` from `head` starts from: the rows of that head's branch
/// ([`branch_rows`]). Without one, a session that went on separately on several machines forks
/// from its newest head, and any other from every row synced of it ([`ForkFrom::default`]).
@@ -424,7 +453,10 @@ pub async fn fork_from(
let rows = branch_rows(source, session, head)
.await
.map_err(|e| NotResumable::Fork(format!("{e:#}")))?;
Ok(ForkFrom { rows, tip: None })
Ok(ForkFrom {
rows: rows.map(|(rows, _)| rows),
tip: None,
})
}
/// The rows a transcript of `session` written out along `head` holds (see
@@ -432,12 +464,12 @@ pub async fn fork_from(
/// one, a session that went on separately on several machines follows its newest head, so no
/// branch's messages land in another's. `None` for any other session (or one whose branches
/// aren't known): it is written out from every row synced of it, as it always was, rows off every
/// branch included.
/// branch included. With the rows, the head followed, when it is one of the session's.
pub async fn branch_rows(
source: &dyn SessionSource,
session: &HarnessSession,
head: Option<&SourceId>,
) -> eyre::Result<Option<Vec<RehydrateMessage>>> {
) -> eyre::Result<Option<(Vec<RehydrateMessage>, Option<Head>)>> {
let Some(analysis) = source.analyse(session).await? else {
return Ok(None);
};
@@ -445,7 +477,8 @@ pub async fn branch_rows(
let Some(head) = head.or(newest.map(|h| &h.source_id)) else {
return Ok(None);
};
Ok(Some(analysis.rows_for(head).into_iter().map(|m| m.clone().into()).collect()))
let rows = analysis.rows_for(head).into_iter().map(|m| m.clone().into()).collect();
Ok(Some((rows, analysis.heads().iter().find(|h| h.source_id == *head).cloned())))
}
#[cfg(test)]
crates/atuin-ai/src/resume_tui/catchup/tests.rs
+39 -1
View File
@@ -114,7 +114,14 @@ async fn a_copy_switches_to_a_head_on_another_line(
let analysis = analysis(rows(true)).await;
let named = named.map(|n| SourceId::from(n.to_owned()));
let copy = tip(known, Some(at));
let to = switch_to(&analysis, HarnessKind::ClaudeCode, named.as_ref(), &copy);
let to = switch_to(
&analysis,
HarnessKind::ClaudeCode,
named.as_ref(),
&copy,
fake::now(),
fake::THIS_HOST_ID,
);
match want {
Ok(head) => assert_eq!(to.unwrap().source_id.as_ref(), head),
Err(why) => assert!(to.unwrap_err().contains(why)),
@@ -140,6 +147,37 @@ async fn a_copy_switches_to_a_head_on_another_line(
);
}
/// With several heads on other lines than the copy's and none named, the copy switches to none
/// of them: which is the user's to say, and they are listed to say it by. A third machine's
/// branch, `p`-`q`, goes on from `b` too.
#[rstest]
#[case::none_named(None, Err(()))]
#[case::named(Some("d"), Ok("d"))]
#[tokio::test]
async fn a_copy_switches_to_no_head_of_several_unless_named(
#[case] named: Option<&str>,
#[case] want: Result<&str, ()>,
) {
let mut rows = rows(true);
rows.extend(
[("p", Some("b"), 3, 6, false), ("q", Some("p"), 3, 7, true)]
.map(|spec| synced_row(&handle(), spec)),
);
let analysis = analysis(rows).await;
let named = named.map(|n| SourceId::from(n.to_owned()));
let copy = tip(&["a", "b", "x", "y"], Some("y"));
let (now, here) = (fake::now(), fake::THIS_HOST_ID);
let to = switch_to(&analysis, HarnessKind::ClaudeCode, named.as_ref(), &copy, now, here);
match want {
Ok(head) => assert_eq!(to.unwrap().source_id.as_ref(), head),
Err(()) => assert_eq!(
to.unwrap_err(),
"it went 2 other ways: name the branch to switch to with `--branch`:\n q @00000003 \
· 53m · 4 msgs (latest)\n d @00000002 · 57m · 4 msgs"
),
}
}
#[rstest]
#[case::live(Why::Live, "Claude Code is running this session here")]
#[case::unsynced(Why::Unsynced, "this copy has messages sync hasn't got")]
crates/atuin-ai/src/resume_tui/fake.rs
+6
View File
@@ -1005,6 +1005,8 @@ pub struct FakeResumer {
pub forks: Arc<Mutex<Vec<ForkFrom>>>,
/// The head each switch went to.
pub switches: Arc<Mutex<Vec<SourceId>>>,
/// The head each continuation followed, if one was named.
pub continues: Arc<Mutex<Vec<Option<SourceId>>>>,
}
impl Default for FakeResumer {
@@ -1014,6 +1016,7 @@ impl Default for FakeResumer {
catch_up: None,
forks: Arc::default(),
switches: Arc::default(),
continues: Arc::default(),
}
}
}
@@ -1110,7 +1113,9 @@ impl Resumer for FakeResumer {
_source: &dyn SessionSource,
session: &SessionRow,
target: HarnessKind,
head: Option<&SourceId>,
) -> Result<Continued, NotResumable> {
self.continues.lock().push(head.cloned());
let mut row = session.clone();
row.handle = HarnessSession {
harness: target,
@@ -1125,6 +1130,7 @@ impl Resumer for FakeResumer {
tool_results: 42,
reasoning: 7,
},
along: None,
note: None,
})
}
crates/atuin-ai/src/resume_tui/resumer.rs
+46 -25
View File
@@ -18,7 +18,7 @@ use std::collections::HashSet;
use std::path::{Path, PathBuf};
use async_trait::async_trait;
use atuin_client::ai_session::{HarnessKind, SourceId};
use atuin_client::ai_session::{HarnessKind, Head, SourceId};
use atuin_client::settings::AiSessionResume;
use atuin_common::harnesstools::continuation::{self, Flattened};
use atuin_common::harnesstools::rehydrate::{
@@ -112,14 +112,22 @@ pub struct Continued {
pub plan: ResumePlan,
/// What the target couldn't take as it was.
pub flattened: Flattened,
/// The host whose branch it carries, as [`catchup::host_label`] names it, when it carries one
/// branch of several (or one named).
pub along: Option<String>,
/// Why it resumes somewhere other than the directory the session ran in, when it does.
pub note: Option<String>,
}
impl Continued {
/// The status line: `continuing in Codex: 42 tool calls become notes, reasoning dropped`.
/// The status line: `continuing in Codex: 42 tool calls become notes, reasoning dropped`, and
/// whose branch it carries (`; along @3f9a12bc's branch`) when it carries one of several.
pub fn status(&self) -> String {
continuing(self.target, Some(&self.flattened))
let status = continuing(self.target, Some(&self.flattened));
match &self.along {
Some(host) => format!("{status}; along {host}'s branch"),
None => status,
}
}
}
@@ -218,13 +226,15 @@ pub trait Resumer: Send + Sync {
) -> Result<CatchUp, NotResumable> {
let resume = self.plan(session).await?;
let (plan, status) = match &resume.restore {
Some(restore) => (self.restore(source, session, restore).await?, restored(restore)),
Some(restore) => {
(self.restore(source, session, restore).await?, Some(restored(restore)))
}
None => (resume.plan, None),
};
Ok(CatchUp::Ready { plan, status })
}
/// Switch this machine's copy of `session` to `head`'s branch (the newest it can be switched
/// Switch this machine's copy of `session` to `head`'s branch (the only one it can be switched
/// to when `None`), from what `source` holds of it, and plan resuming it ([`super::catchup`]):
/// in place, keeping the history it shares with the branch as it is, and the copy as it was
/// as a backup. Refused, with nothing written, unless sync holds every row of the copy, no
@@ -244,13 +254,16 @@ pub trait Resumer: Send + Sync {
Vec::new()
}
/// Write `session`, from what `source` holds of it, out as a new session of `target`
/// ([`continuation`]), and plan resuming that. Only once the user has chosen to.
/// Write `session`, from what `source` holds of it along `head` (one of its heads; for a
/// session that diverged, the newest when `None`: see [`restored_session`]), out as a new
/// session of `target` ([`continuation`]), and plan resuming that. Only once the user has
/// chosen to.
async fn continue_in(
&self,
_source: &dyn SessionSource,
session: &SessionRow,
_target: HarnessKind,
_head: Option<&SourceId>,
) -> Result<Continued, NotResumable> {
Err(NotResumable::Unsupported(harness_label(session.handle.harness)))
}
@@ -350,9 +363,13 @@ pub trait Machine: Send + Sync {
}
}
/// What the status line says of a session restored as `restore`, when anything.
fn restored(restore: &Restore) -> Option<String> {
restore.note.as_ref().map(|note| format!("restored the session from sync; {note}"))
/// What the status line says of a session restored as `restore`: that it was, and why it resumes
/// where it does, when that isn't where it ran.
fn restored(restore: &Restore) -> String {
match &restore.note {
Some(note) => format!("restored the session from sync; {note}"),
None => "restored the session from sync".to_owned(),
}
}
/// This machine, as the harnesses themselves see it.
@@ -519,7 +536,7 @@ impl HarnessResumer {
restore: &Restore,
head: Option<&SourceId>,
) -> Result<ResumePlan, NotResumable> {
let data = restored_session(source, session, &restore.cwd, head)
let (data, _) = restored_session(source, session, &restore.cwd, head)
.await
.map_err(|e| NotResumable::Restore(format!("{e:#}")))?;
self.write_out(session.handle.harness, &data, NotResumable::Restore).await
@@ -528,18 +545,20 @@ impl HarnessResumer {
/// `session` as `source` holds it, to be written out and resumed in `cwd`: along `head` (or, for
/// a session that diverged, its newest head) when there is one ([`catchup::branch_rows`]), so no
/// other branch's messages come with it; else every row synced of it.
/// other branch's messages come with it; else every row synced of it. With it, the head it was
/// written along, if any.
async fn restored_session(
source: &dyn SessionSource,
session: &SessionRow,
cwd: &Path,
head: Option<&SourceId>,
) -> eyre::Result<RehydrateSession> {
) -> eyre::Result<(RehydrateSession, Option<Head>)> {
let mut data = source.rehydrate(&session.handle, cwd).await?;
if let Some(rows) = catchup::branch_rows(source, &session.handle, head).await? {
data.messages = rows;
}
Ok(data)
let Some((rows, along)) = catchup::branch_rows(source, &session.handle, head).await? else {
return Ok((data, None));
};
data.messages = rows;
Ok((data, along))
}
/// This machine's copy of a session, as [`HarnessResumer::plan_here`] read it.
@@ -635,7 +654,7 @@ impl Resumer for HarnessResumer {
let ready = |plan, status| Ok(CatchUp::Ready { plan, status });
if let Some(restore) = &resume.restore {
let plan = self.restore_along(source, session, restore, head).await?;
return ready(plan, restored(restore));
return ready(plan, Some(restored(restore)));
}
let Some(Copy { tip, dir }) = copy else {
return ready(resume.plan, None);
@@ -646,7 +665,7 @@ impl Resumer for HarnessResumer {
// Gone since it was located: restored, as a copy that isn't here is.
let restore = resolve_cwd(session.cwd.as_deref(), &self.context);
let plan = self.restore_along(source, session, &restore, head).await?;
return ready(plan, restored(&restore));
return ready(plan, Some(restored(&restore)));
}
Err(e) => Err(e),
};
@@ -670,8 +689,7 @@ impl Resumer for HarnessResumer {
let here = &self.context.host_id;
let tip = match tip {
Ok(tip) => tip,
// A copy the agent can't read, or can't catch up (opencode 2.0's): never resumed
// as if it were caught up.
// A copy the agent can't read: never resumed as if it were caught up.
Err(e) => {
let branches = catchup::branches(&analysis, None, here);
return Ok(held(Why::Refused(e.to_string()), kind, resume.plan, branches, chosen));
@@ -761,7 +779,8 @@ impl Resumer for HarnessResumer {
let analysis = analysis.ok_or_else(|| {
fail("its branches aren't known yet: the daemon hasn't synced its messages".to_owned())
})?;
let head = catchup::switch_to(&analysis, kind, head, &tip).map_err(fail)?;
let (now, here) = (time::OffsetDateTime::now_utc(), &self.context.host_id);
let head = catchup::switch_to(&analysis, kind, head, &tip, now, here).map_err(fail)?;
let running = || fail(format!("{} is running this session here", harness_label(kind)));
// Never under an agent that may be writing the session: checked again as it is written.
if resume.live {
@@ -840,7 +859,7 @@ impl Resumer for HarnessResumer {
let restore = resolve_cwd(session.cwd.as_deref(), &self.context);
// Restored as a resume would restore it (along its newest head, when it diverged), should
// pi need it written out below.
let mut original = restored_session(source, session, &restore.cwd, None)
let (mut original, _) = restored_session(source, session, &restore.cwd, None)
.await
.map_err(|e| fail(format!("{e:#}")))?;
// Checked before anything is written. A fork of no conversation is refused, as a
@@ -893,6 +912,7 @@ impl Resumer for HarnessResumer {
source: &dyn SessionSource,
session: &SessionRow,
target: HarnessKind,
head: Option<&SourceId>,
) -> Result<Continued, NotResumable> {
let from_kind = session.handle.harness;
let from =
@@ -903,8 +923,8 @@ impl Resumer for HarnessResumer {
return Err(fail("it's that agent's own session".to_owned()));
}
let restore = resolve_cwd(session.cwd.as_deref(), &self.context);
let original = source
.rehydrate(&session.handle, &restore.cwd)
// Along one branch, as a restore or a fork is: never every branch's messages at once.
let (original, along) = restored_session(source, session, &restore.cwd, head)
.await
.map_err(|e| fail(format!("{e:#}")))?;
let atuin_id = session.atuin_id.to_string();
@@ -917,6 +937,7 @@ impl Resumer for HarnessResumer {
target,
plan,
flattened: continued.flattened,
along: along.map(|head| catchup::host_label(&head, &self.context.host_id)),
note: restore.note,
})
}
crates/atuin-ai/src/resume_tui/resumer/tests.rs
+55 -7
View File
@@ -581,7 +581,7 @@ async fn continuing_writes_a_new_session_of_the_target(dirs: Dirs) {
let machine = FakeMachine::default();
let resumer = resumer(&dirs, &machine);
let session = row(HarnessKind::ClaudeCode, "abc-123", &dirs.elsewhere, true);
let continued = resumer.continue_in(&Worked, &session, HarnessKind::Pi).await.unwrap();
let continued = resumer.continue_in(&Worked, &session, HarnessKind::Pi, None).await.unwrap();
let written = machine.written.lock().clone();
let [new] = written.as_slice() else {
@@ -615,7 +615,7 @@ async fn an_empty_session_is_not_continued(
let machine = FakeMachine::default();
let resumer = resumer(&dirs, &machine);
let session = row(HarnessKind::ClaudeCode, "abc-123", &dirs.elsewhere, true);
let err = resumer.continue_in(&Synced, &session, target).await.unwrap_err();
let err = resumer.continue_in(&Synced, &session, target, None).await.unwrap_err();
assert_eq!(err, NotResumable::Empty(continuation::NothingToContinue));
assert_eq!(err.to_string(), "nothing to continue: the session has no messages");
assert!(machine.written.lock().is_empty());
@@ -654,7 +654,7 @@ async fn a_path_template_target_is_offered_and_resumes_the_written_session(dirs:
let session = row(HarnessKind::ClaudeCode, "abc-123", &dirs.elsewhere, true);
assert!(resumer.continue_targets(&session).contains(&HarnessKind::Pi));
let continued = resumer.continue_in(&Worked, &session, HarnessKind::Pi).await.unwrap();
let continued = resumer.continue_in(&Worked, &session, HarnessKind::Pi, None).await.unwrap();
let written = machine.written.lock().clone();
let path = format!("/restored/{}.jsonl", written[0].id);
assert_eq!(continued.plan.args, ["--session", path.as_str()]);
@@ -673,9 +673,11 @@ async fn only_other_installed_harnesses_are_offered(dirs: Dirs) {
let session = row(HarnessKind::ClaudeCode, "abc-123", &dirs.elsewhere, false);
assert_eq!(resumer.continue_targets(&session), [HarnessKind::Codex, HarnessKind::Pi]);
let err = resumer.continue_in(&Worked, &session, HarnessKind::Opencode).await.unwrap_err();
let err =
resumer.continue_in(&Worked, &session, HarnessKind::Opencode, None).await.unwrap_err();
assert_eq!(err, NotResumable::NotInstalled("opencode".to_owned()));
let err = resumer.continue_in(&Worked, &session, HarnessKind::ClaudeCode).await.unwrap_err();
let err =
resumer.continue_in(&Worked, &session, HarnessKind::ClaudeCode, None).await.unwrap_err();
assert!(matches!(err, NotResumable::Continue("Claude Code", _)), "{err}");
assert!(machine.written.lock().is_empty());
@@ -924,7 +926,7 @@ async fn catch_up_in(
/// it isn't here, resumed as it is when it is at the head, the rows it lacks appended when it is
/// behind; anything else is a choice, with nothing written.
#[rstest]
#[case::not_here(None, None, false, "/restored/s.jsonl", &[])]
#[case::not_here(None, None, false, "/restored/s.jsonl restored the session from sync", &[])]
#[case::at_the_head(Some((&["a", "b", "c", "d"][..], "d")), None, false, "/t/s.jsonl", &[])]
#[case::behind(
Some((&["a", "b"][..], "b")),
@@ -1153,6 +1155,52 @@ async fn a_session_not_here_is_restored_along_one_branch(
assert_eq!(ids, want);
}
/// A session continued in another harness carries one branch, as a restore does: the head named,
/// else, when it diverged, the newest (as the picker continues it), saying whose; never every
/// branch's messages at once.
#[rstest]
#[case::one_way(false, None, &["a", "b", "c", "d"], None)]
#[case::diverged(true, None, &["a", "b", "x", "y"], Some("this machine"))]
#[case::diverged_named(true, Some("d"), &["a", "b", "c", "d"], Some("@00000002"))]
#[tokio::test]
async fn a_session_is_continued_along_one_branch(
dirs: Dirs,
#[case] diverged: bool,
#[case] head: Option<&str>,
#[case] want: &[&str],
#[case] along: Option<&str>,
) {
use atuin_common::harnesstools::session::Content;
let machine = FakeMachine::default();
let resumer = resumer(&dirs, &machine);
let mut session = row(HarnessKind::ClaudeCode, "s", &dirs.elsewhere, true);
session.handle = fake::synced_handle();
let source = fake::FakeSource::from_rows(vec![session.clone()])
.with_synced(&session.handle, fake::synced_rows(diverged));
let head = head.map(|h| SourceId::from(h.to_owned()));
let continued = resumer.continue_in(&source, &session, HarnessKind::Codex, head.as_ref());
let status = continued.await.unwrap().status();
// Which branch it carries, when it carries one of several, the picker says.
let says = along.map(|host| format!("; along {host}'s branch")).unwrap_or_default();
assert_eq!(status, format!("continuing in Codex{says}"));
let written = machine.written.lock().clone();
let [continued] = written.as_slice() else {
panic!("one session written: {written:?}");
};
let said: Vec<String> = continued
.messages
.iter()
.flat_map(|m| &m.content)
.filter_map(|c| match c {
Content::Text(text) => text.strip_prefix("message ").map(str::to_owned),
_ => None,
})
.collect();
assert_eq!(said, want);
}
/// The plan, the catch-up and the append all look for an agent where this machine's copy
/// records working, when it records a directory that is there: not where the session row says
/// (another host's newer directory, maybe). Else where the session works here; else anywhere.
@@ -1229,7 +1277,7 @@ async fn a_copy_gone_since_it_was_located_is_restored(dirs: Dirs) {
machine.found.insert("s".to_owned(), PathBuf::from("/t/s.jsonl"));
let written = machine.written.clone();
let caught = catch_up(&dirs, machine, false).await;
assert_eq!(caught, ("/restored/s.jsonl".to_owned(), vec![]));
assert_eq!(caught, ("/restored/s.jsonl restored the session from sync".to_owned(), vec![]));
assert_eq!(written.lock().len(), 1);
}
crates/atuin-ai/src/resume_tui/sidecar.rs
+27 -2
View File
@@ -6,6 +6,7 @@ use std::path::{Path, PathBuf};
use async_trait::async_trait;
use atuin_client::ai_session::{AiSessionDatabase, Analysis, HarnessSession, SearchTerms, Session};
use atuin_common::harnesstools::continuation;
use atuin_common::harnesstools::rehydrate::RehydrateSession;
use atuin_common::harnesstools::session::{Content, ParentKind};
use atuin_common::string::highlighted::HighlightedString;
@@ -108,9 +109,11 @@ fn conversation_text(c: &Content) -> Option<&str> {
}
}
/// The conversation text of a message: text and summaries, never tool calls or reasoning.
/// The conversation text of a message: text and summaries, never tool calls or reasoning, nor
/// the marker opening a pi continuation's first prompt.
fn text_of(content: &[Content]) -> Option<String> {
let text: Vec<&str> = content.iter().filter_map(conversation_text).collect();
let marker = usize::from(continuation::opens_with_marker(content));
let text: Vec<&str> = content[marker..].iter().filter_map(conversation_text).collect();
(!text.is_empty()).then(|| text.join("\n"))
}
@@ -360,6 +363,28 @@ mod tests {
}
}
/// A pi continuation's first prompt opens with its marker: it is titled and previewed by
/// what the user said, never by the marker.
#[rstest]
#[tokio::test]
async fn a_continuation_is_shown_by_its_prompt_not_its_marker(#[future] source: SidecarSource) {
use atuin_common::harnesstools::AnyHarness;
use atuin_common::harnesstools::ccode::Ccode;
use atuin_common::harnesstools::continuation::marker_text;
let source = source.await;
let marker = marker_text(AnyHarness::ClaudeCode(Ccode), "root", None);
let mut prompt = message("continued", None, Role::User, "fix the flaky clock", 4);
prompt.content.insert(0, Content::Text(marker));
source.db.append(&prompt).await.unwrap();
let rows = source.search(&roots("")).await.unwrap();
let row = rows.iter().find(|r| r.handle == handle("continued")).unwrap();
assert_eq!(row.title.text, "fix the flaky clock");
let preview = source.preview(&handle("continued")).await.unwrap();
assert_eq!(preview.first_prompt.as_deref(), Some("fix the flaky clock"));
}
#[rstest]
#[tokio::test]
async fn previews_and_ids(#[future] source: SidecarSource) {
crates/atuin-ai/src/resume_tui/tests.rs
+7 -7
View File
@@ -819,7 +819,7 @@ async fn a_picked_fork_is_written_then_resumed() {
assert_eq!(pick(&mut state, picked, resumer.as_ref(), &requests), None, "waits");
let id = state.continued;
assert!(state.continuing.as_ref().is_some_and(|c| c.kind == super::state::Writing::Fork));
let continued = resumer.continue_in(&source, &row, HarnessKind::Codex).await;
let continued = resumer.continue_in(&source, &row, HarnessKind::Codex, None).await;
assert!(finish_continuation(&mut state, id, continued).is_none(), "not a continuation");
let forked = resumer.fork(&source, &row, ForkFrom::default()).await;
@@ -1091,7 +1091,7 @@ async fn changing_the_search_drops_a_continuation() {
assert!(state.continuing.is_none());
assert_eq!(state.status, None);
let continued = resumer.continue_in(&FakeSource::new(), &row, HarnessKind::Codex).await;
let continued = resumer.continue_in(&FakeSource::new(), &row, HarnessKind::Codex, None).await;
assert_eq!(finish_continuation(&mut state, id, continued), None, "ignored");
}
@@ -1311,7 +1311,7 @@ async fn a_written_continuation_resumes_or_edits_its_new_session() {
for (action, run) in [(Pending::Resume, true), (Pending::Edit, false)] {
start_continuation(&mut state, row.clone(), HarnessKind::Codex, action, &requests);
let id = state.continued;
let continued = resumer.continue_in(&source, &row, HarnessKind::Codex).await;
let continued = resumer.continue_in(&source, &row, HarnessKind::Codex, None).await;
let (outcome, status) = finish_continuation(&mut state, id, continued).unwrap();
assert_eq!(status, "continuing in Codex: 42 tool calls become notes, reasoning dropped");
let plan = match (outcome, run) {
@@ -1362,11 +1362,11 @@ async fn a_superseded_continuation_is_not_taken_for_the_newer_one() {
let second = state.continued;
assert_ne!(first, second);
let codex = resumer.continue_in(&source, &row, HarnessKind::Codex).await;
let codex = resumer.continue_in(&source, &row, HarnessKind::Codex, None).await;
assert!(finish_continuation(&mut state, first, codex).is_none(), "the first answer");
assert!(state.continuing.is_some(), "still waiting on the newer pick");
let pi = resumer.continue_in(&source, &row, HarnessKind::Pi).await;
let pi = resumer.continue_in(&source, &row, HarnessKind::Pi, None).await;
let (outcome, _) = finish_continuation(&mut state, second, pi).unwrap();
let Outcome::Edit(plan) = outcome else {
panic!("{outcome:?}");
@@ -1419,7 +1419,7 @@ async fn a_continuation_picked_supersedes_a_waiting_original(#[values(false, tru
let forked = resumer.fork(&source, &row, ForkFrom::default()).await;
finish_fork(&mut state, id, forked).unwrap()
} else {
let continued = resumer.continue_in(&source, &row, HarnessKind::Codex).await;
let continued = resumer.continue_in(&source, &row, HarnessKind::Codex, None).await;
finish_continuation(&mut state, id, continued).unwrap()
};
let Outcome::Resume(plan) = outcome else {
@@ -1485,7 +1485,7 @@ async fn an_original_chosen_supersedes_a_continuation(
let forked = resumer.fork(&source, &row, ForkFrom::default()).await;
assert_eq!(finish_fork(&mut state, id, forked), None, "ignored");
} else {
let continued = resumer.continue_in(&source, &row, HarnessKind::Codex).await;
let continued = resumer.continue_in(&source, &row, HarnessKind::Codex, None).await;
assert_eq!(finish_continuation(&mut state, id, continued), None, "ignored");
}
let outcome =
crates/atuin-ai/src/resume_tui/worker.rs
+6 -4
View File
@@ -287,9 +287,10 @@ async fn details(
);
Response::Flattened(session, flattened)
}
Request::Continue(row, target, id) => {
Response::Continued(id, resumer.continue_in(source.as_ref(), &row, target).await)
}
Request::Continue(row, target, id) => Response::Continued(
id,
resumer.continue_in(source.as_ref(), &row, target, None).await,
),
Request::Fork(row, id, head) => {
let source = source.as_ref();
let from = catchup::fork_from(source, &row.handle, head.as_ref()).await;
@@ -478,10 +479,11 @@ mod tests {
source: &dyn SessionSource,
session: &SessionRow,
target: HarnessKind,
head: Option<&SourceId>,
) -> Result<Continued, NotResumable> {
self.started.lock().push(session.handle.session.to_string());
self.gate.acquire().await.unwrap().forget();
self.inner.continue_in(source, session, target).await
self.inner.continue_in(source, session, target, head).await
}
async fn fork(
crates/atuin-client/src/ai_session/database.rs
+103 -11
View File
@@ -7,6 +7,7 @@ use atuin_common::db::sqlite::fts::{
};
use atuin_common::db::sqlite::{Sqlite, SqliteOpenOrCreateError};
use atuin_common::db::{self};
use atuin_common::harnesstools::continuation;
use atuin_common::harnesstools::rehydrate::RehydrateSession;
use atuin_common::harnesstools::session::{
Checkpoint, Content, ParentKind, Role, TitleChange, TitleSource, Usage,
@@ -51,6 +52,14 @@ const RECENCY_DAYS: f64 = 30.0;
const DAY_MILLIS: f64 = 86_400_000.0;
/// The most a row stores of each usage count: far past any model call, and small enough that a
/// session's totals, adding up millions of calls, never overflow.
const MAX_ROW_TOKENS: u64 = 1 << 40;
/// [`Message::EARLIEST`] and [`Message::LATEST`] in milliseconds: the range of a stored timestamp.
const EARLIEST_MILLIS: i64 = Message::EARLIEST.unix_timestamp() * 1000;
const LATEST_MILLIS: i64 = Message::LATEST.unix_timestamp() * 1000;
/// The shortest last term that matches as a prefix. A one-character prefix matches most of the
/// index (every row with a word starting `e`), which costs hundreds of milliseconds to rank for
/// nothing useful, so a lone character matches only as a whole token until the next keystroke.
@@ -195,8 +204,6 @@ pub enum DbError {
Json(#[from] serde_json::Error),
#[error("failed to compress ai-session message content: {0}")]
Compress(std::io::Error),
#[error("stored ai-session timestamp is out of range: {0}")]
Time(#[from] time::error::ComponentRange),
#[error("unknown ai-session harness discriminant {0}")]
UnknownHarness(i64),
#[error("failed to decompress ai-session message content: {0}")]
@@ -481,6 +488,12 @@ impl AiSessionDatabase {
self.reprojection.clone().lock_owned().await
}
/// [`Self::lock_reprojection`], unless a reprojection holds it.
#[must_use]
pub fn try_lock_reprojection(&self) -> Option<tokio::sync::OwnedMutexGuard<()>> {
self.reprojection.clone().try_lock_owned().ok()
}
/// Whether the sidecar is at exactly the schema this build reads, by the newest migration
/// applied to it.
async fn check_schema(&self) -> Result<(), DbError> {
@@ -1291,7 +1304,9 @@ impl AiSessionDatabase {
return None;
}
msg.content.iter().find_map(|content| match content {
// pi's continuation marker opens the prompt it rides on: the prompt is what was said.
let marker = usize::from(continuation::opens_with_marker(&msg.content));
msg.content[marker..].iter().find_map(|content| match content {
Content::Text(text) => Some(Self::clip_summary(text)),
_ => None,
})
@@ -1395,7 +1410,12 @@ impl AiSessionDatabase {
};
[usage.input, usage.output, usage.cache_read, usage.cache_write, usage.reasoning]
.map(|n| i64::try_from(n.unwrap_or(0)).unwrap_or(i64::MAX))
.map(|n| Self::row_tokens(n.unwrap_or(0)))
}
/// A count of a row's usage as stored: at most [`MAX_ROW_TOKENS`].
fn row_tokens(n: u64) -> i64 {
i64::try_from(n.min(MAX_ROW_TOKENS)).expect("MAX_ROW_TOKENS fits in i64")
}
/// Delete every session holding a row `host` captured, whole, for a reprojection that replays
@@ -1445,12 +1465,18 @@ impl AiSessionDatabase {
host.and_then(|h| uuid::Uuid::parse_str(&h).ok()).map(HostId)
}
/// `ts` in milliseconds, held to [`Message::EARLIEST`] and [`Message::LATEST`].
fn millis(ts: OffsetDateTime) -> i64 {
i64::try_from(ts.unix_timestamp_nanos() / 1_000_000).unwrap_or(i64::MAX)
i64::try_from(ts.unix_timestamp_nanos() / 1_000_000)
.expect("every OffsetDateTime fits in i64 milliseconds")
.clamp(EARLIEST_MILLIS, LATEST_MILLIS)
}
fn time_from_millis(ms: i64) -> Result<OffsetDateTime, DbError> {
Ok(OffsetDateTime::from_unix_timestamp_nanos(i128::from(ms) * 1_000_000)?)
/// The inverse of [`Self::millis`], held to the same range, so it never fails.
fn time_from_millis(ms: i64) -> OffsetDateTime {
let ms = ms.clamp(EARLIEST_MILLIS, LATEST_MILLIS);
OffsetDateTime::from_unix_timestamp_nanos(i128::from(ms) * 1_000_000)
.expect("the range is representable")
}
fn harness_from_repr(n: i64) -> Result<HarnessKind, DbError> {
@@ -1511,7 +1537,7 @@ impl AiSessionDatabase {
.parent(parent)
.parent_kind(parent_kind)
.parent_source_id(row.parent_source_id.map(SourceId::from))
.timestamp(Self::time_from_millis(row.timestamp)?)
.timestamp(Self::time_from_millis(row.timestamp))
.role(role)
.content(content)
.cwd(row.cwd.map(PathBuf::from))
@@ -1574,7 +1600,7 @@ impl AiSessionDatabase {
let parent = Self::optional_session(row.parent_harness, row.parent_session_id)?;
let root = Self::optional_session(row.root_harness, row.root_session_id)?
.filter(|root| root.harness != harness || root.session.as_ref() != row.session_id);
let group_updated_at = row.group_updated_at.map(Self::time_from_millis).transpose()?;
let group_updated_at = row.group_updated_at.map(Self::time_from_millis);
let copy_of = row.copy_of_session_id.map(|session| HarnessSession {
harness,
session: NativeSessionId::from(session),
@@ -1583,7 +1609,7 @@ impl AiSessionDatabase {
harness,
session: NativeSessionId::from(row.session_id),
};
let started_at = Self::time_from_millis(row.started_at)?;
let started_at = Self::time_from_millis(row.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]>| {
@@ -1605,7 +1631,7 @@ impl AiSessionDatabase {
.git_branch(row.git_branch)
.model(row.model)
.started_at(started_at)
.updated_at(Self::time_from_millis(row.updated_at)?)
.updated_at(Self::time_from_millis(row.updated_at))
.message_count(u64::try_from(row.message_count).unwrap_or(0))
.usage(Usage {
input: Some(u64::try_from(row.usage_input).unwrap_or(0)),
@@ -5090,4 +5116,70 @@ mod tests {
.unwrap();
assert_eq!(sessions, 1);
}
/// A timestamp past the years a date can name, appended or stored, is held to the range
/// rather than failing every read that meets it.
#[rstest]
#[case::appended(time::macros::datetime!(9999-12-31 23:59:59 -23:59), None)]
#[case::stored_late(OffsetDateTime::UNIX_EPOCH, Some(i64::MAX))]
#[case::stored_early(OffsetDateTime::UNIX_EPOCH, Some(i64::MIN))]
#[tokio::test]
async fn an_out_of_range_timestamp_reads_as_the_nearest_in_range(
#[case] timestamp: OffsetDateTime,
#[case] stored: Option<i64>,
) {
let db = AiSessionDatabase::in_memory().await.unwrap();
let mut odd = message_in(&sample_handle(), 0, "odd");
odd.timestamp = timestamp;
db.append(&odd).await.unwrap();
db.append(&three_sessions_oldest_first()[0]).await.unwrap();
if let Some(stored) = stored {
db::query("UPDATE messages SET timestamp = ?1 WHERE source_id = 'source-0'")
.bind(stored)
.execute(db.db.pool())
.await
.unwrap();
db::query("UPDATE sessions SET started_at = ?1, updated_at = ?1")
.bind(stored)
.execute(db.db.pool())
.await
.unwrap();
}
let sessions = db.list_sessions(&SessionFilter::default()).await.unwrap();
assert_eq!(sessions.len(), 2);
let messages: Vec<_> = db.messages(&sample_handle()).try_collect().await.unwrap();
let read = messages[0].timestamp;
assert!((Message::EARLIEST..=Message::LATEST).contains(&read), "{read}");
let found: Vec<_> = db
.search("odd", SearchTerms::All, &SessionFilter::default(), 0)
.try_collect()
.await
.unwrap();
assert_eq!(found.len(), 1);
}
/// Usage counts past any model call's are stored held to a bound, so a session's totals never
/// overflow.
#[rstest]
#[tokio::test]
async fn huge_usage_never_breaks_a_sessions_totals() {
let db = AiSessionDatabase::in_memory().await.unwrap();
let s = sample_handle();
for i in 0..3 {
let mut m = message_in(&s, i, "reply");
m.turn_id = Some(format!("msg_{i}"));
m.usage = Some(Usage {
input: Some(u64::MAX),
output: Some(u64::MAX),
cache_read: Some(u64::MAX),
cache_write: Some(u64::MAX),
reasoning: Some(u64::MAX),
});
db.append(&m).await.unwrap();
}
let session = db.get_session(&s).await.unwrap().unwrap();
assert_eq!(session.usage.input, Some(3 << 40));
assert_eq!(db.list_sessions(&SessionFilter::default()).await.unwrap().len(), 1);
}
}
crates/atuin-client/src/ai_session/database/write.rs
+1 -2
View File
@@ -267,8 +267,7 @@ impl AiSessionDatabase {
Self::fold_usage(msg.usage.as_ref());
// Unlike the others, a row's reasoning stays NULL when unreported: most harnesses never
// break it out, and zero would claim the call did not reason.
let usage_reasoning =
msg.usage.and_then(|u| u.reasoning).map(|n| i64::try_from(n).unwrap_or(i64::MAX));
let usage_reasoning = msg.usage.and_then(|u| u.reasoning).map(Self::row_tokens);
let cwd = msg.cwd.as_ref().map(|p| p.to_string_lossy().into_owned());
let host = msg.host.map(Self::host_repr);
crates/atuin-client/src/ai_session/model.rs
+9
View File
@@ -196,6 +196,15 @@ pub struct Message {
pub host: Option<HostId>,
}
impl Message {
/// The earliest and latest [`Self::timestamp`] stored: a day inside the years
/// [`OffsetDateTime`] can name, so either is still a date in any offset. A harness line can
/// name an instant past them (Claude Code's `9999-12-31T23:59:59-23:59` is in year 10000 in
/// UTC), which capture and the sidecar hold to these.
pub const EARLIEST: OffsetDateTime = time::macros::datetime!(-9999-01-02 0:00 UTC);
pub const LATEST: OffsetDateTime = time::macros::datetime!(9999-12-30 0:00 UTC);
}
/// [`Message::parent_kind`] as a record holds it: `None` for a kind this build does not know,
/// so a newer build can add kinds without its records failing to decode here.
fn known_parent_kind<'de, D>(deserializer: D) -> Result<Option<ParentKind>, D::Error>
crates/atuin-client/src/ai_session/store.rs
+90 -9
View File
@@ -10,7 +10,9 @@ use tracing::{debug, warn};
use typed_builder::TypedBuilder;
use crate::ai_session::Message;
use crate::ai_session::database::{AiSessionDatabase, DbError, PreparedMessage, Watermark};
use crate::ai_session::database::{
AiSessionDatabase, DbError, Generation, PreparedMessage, Watermark,
};
use crate::record::decode::decode_parallel;
use crate::record::sqlite_store::SqliteStore;
@@ -35,6 +37,20 @@ pub struct AiSessionStore {
key: Key,
}
/// What capture knows of this host's watermark before it pushes a record, for
/// [`AiSessionStore::advance_past_capture`] to move it past the record.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CaptureMark {
/// Nothing but the invalidation generation, read before the push.
Unknown(Generation),
/// Where capture last moved it, and the generation it moved it under: no invalidation since
/// means it still stands there.
Moved {
generation: Generation,
mark: Watermark,
},
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum AiSessionRecord {
Message(Message),
@@ -135,7 +151,8 @@ impl AiSessionStore {
Ok(self.store.contains(id).await?)
}
pub async fn push(&self, msg: &Message) -> Result<RecordId, PushError> {
/// Push `msg` to this host's series, returning where it was stored.
pub async fn push(&self, msg: &Message) -> Result<Watermark, PushError> {
let id = msg.id;
let bytes = AiSessionRecord::Message(msg.clone()).serialize();
@@ -154,11 +171,63 @@ impl AiSessionStore {
.build();
if self.store.push_unique(&record.encrypt(&self.key)).await? {
return Ok(id);
return Ok(Watermark { idx, record_id: id });
}
}
}
/// Move this host's watermark onto `pushed`, a record [`Self::push`] stored and the caller
/// has since projected, so the next reprojection (each startup's) doesn't replay what capture
/// projected as it went. Only when the watermark stands on the record just before it (still
/// the record it was, as a reprojection checks), no invalidation happened since `known`'s
/// generation, and no reprojection is running, which moves it itself.
///
/// Returns where it moved it, for the next capture's `known`: one following on from it moves
/// it with a single write. `None` if it didn't move it: the next capture reads the generation
/// before its push again.
pub async fn advance_past_capture(
&self,
db: &AiSessionDatabase,
known: CaptureMark,
pushed: Watermark,
) -> Result<Option<CaptureMark>, BuildError> {
let Some(_running) = db.try_lock_reprojection() else {
return Ok(None);
};
let series = RecordSeriesKey::new(self.host_id, RecordTag::AiSession);
let (generation, from) = match known {
CaptureMark::Moved { generation, mark } => {
if mark.idx + 1 != pushed.idx {
return Ok(None);
}
(generation, Some(mark))
}
CaptureMark::Unknown(generation) => {
let from = db.reproject_watermarks().await?.remove(&series);
let follows = match from {
None => pushed.idx == 0,
Some(mark) => {
mark.idx + 1 == pushed.idx
&& self
.store
.idx(&series, mark.idx)
.await?
.is_some_and(|record| record.id == mark.record_id)
}
};
if !follows {
return Ok(None);
}
(generation, from)
}
};
let moved = db.advance_reproject_watermark(&series, generation, from, pushed).await?;
Ok(moved.then_some(CaptureMark::Moved {
generation,
mark: pushed,
}))
}
/// Decrypt and decode `record` with `key`, and prepare it for the sidecar, which the caller
/// appends it to.
fn decode(key: &Key, record: &Record<EncryptedData>) -> Result<Projected, DbError> {
@@ -179,9 +248,16 @@ impl AiSessionStore {
// - One that decrypts but whose kind or version this build does not know is held until
// an upgrade.
// - One of a version this build no longer reads (from before records stored the
// session's id) is skipped: no later build reads it either.
// session's id) is skipped: no later build reads it either. Before it is decrypted,
// which is most of what a replay of a store full of them costs. The version is only
// authenticated by decrypting, so a record's version changed to v0 is skipped rather
// than held: no more than a server can do by withholding the record.
// - One that decrypts to a known kind but fails to decode is skipped: decryption
// authenticates it, so it is exactly what its writer wrote, and no retry reads it.
if record.version == RecordVersion::V0 {
debug!(id = %id.0, "skipping retired ai-session record");
return Ok(Projected::Skipped);
}
let decrypted = match record.decrypt(key) {
Ok(decrypted) => decrypted,
Err(err) => {
@@ -701,9 +777,13 @@ mod tests {
let store = SqliteStore::in_memory(test_local_timeout()).await.unwrap();
let s = AiSessionStore::builder().store(store.clone()).host_id(hid()).key(key()).build();
let msg = sample_message();
let id = s.push(&msg).await.unwrap();
let pushed = s.push(&msg).await.unwrap();
assert_eq!(id, msg.id, "push must return the message's own id, not a re-minted one");
assert_eq!(
pushed.record_id, msg.id,
"push must return the message's own id, not a re-minted one"
);
assert_eq!(pushed.idx, 0);
let recs = store.all_tagged(&RecordTag::AiSession).await.unwrap();
assert_eq!(recs.len(), 1);
@@ -1521,10 +1601,11 @@ mod tests {
}
/// A v0 record (from before records stored the session's id) is passed over, not held: it
/// would hold the watermark forever.
/// would hold the watermark forever. It is passed over unread, so even one this key can't
/// decrypt (which would otherwise be held) is.
#[rstest]
#[tokio::test]
async fn a_v0_record_is_skipped() {
async fn a_v0_record_is_skipped(#[values(key(), Key::from([1u8; 32]))] written_with: Key) {
let store = SqliteStore::in_memory(test_local_timeout()).await.unwrap();
let [s] = <[_; 1]>::try_from(writers(&store, 1)).ok().unwrap();
let handle = sample_handle();
@@ -1538,7 +1619,7 @@ mod tests {
.idx(1)
.data(DecryptedData(AiSessionRecord::Message(msg).serialize()))
.build();
store.push(&record.encrypt(&key())).await.unwrap();
store.push(&record.encrypt(&written_with)).await.unwrap();
push_range(&s, &handle, 1..3).await;
let db = AiSessionDatabase::in_memory().await.unwrap();
crates/atuin-client/src/meta.rs
+41 -1
View File
@@ -30,6 +30,9 @@ const KEY_FILES_MIGRATED: &str = "files_migrated";
const HUB_TOKEN_PREFIX: &str = "atapi_";
/// The longest to wait for another process's commit to the store before failing.
const LOCK_WAIT: Duration = Duration::from_mins(1);
/// Whether `token` is a Hub API token rather than a legacy CLI session token.
#[must_use]
pub fn is_hub_token(token: &SecretString) -> bool {
@@ -63,8 +66,15 @@ impl MetaStore {
}
async fn from_builder(builder: SqliteBuilder<'_>, timeout: Duration) -> Result<Self> {
// A rollback journal, as every atuin before has forced it: one that opened it in WAL mode
// would lock an older one out for as long as it ran. So a commit holds out readers
// while it syncs, which on a busy disk can take many seconds (the daemon commits on every
// sync): wait that out rather than fail. The wait outlasts the user's `local_timeout`, and
// holds up interactive paths too (Ctrl-R's host id, nushell's `history end` without the
// daemon): a hang on a stalled disk, rather than an error.
let sqlite = builder
.timeout(timeout)
.timeout(timeout.max(LOCK_WAIT))
.busy_timeout(LOCK_WAIT)
.journal(Some(Journaling::Delete))
.foreign_keys(false)
.restrict_permissions()
@@ -442,4 +452,34 @@ mod tests {
assert_eq!(store.get(KEY_FILES_MIGRATED).await.unwrap(), None);
}
/// A commit by another process held up past sqlite's usual 5s wait (the daemon saving its
/// last sync on a busy disk) is waited out: a command opening and reading the store meanwhile
/// gets there, and sees the store as committed.
#[rstest]
#[tokio::test(flavor = "multi_thread")]
async fn waits_out_another_process_committing_slowly() {
use sqlx::{ConnectOptions as _, Connection as _};
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("meta.db");
MetaStore::new(&path, Duration::from_secs(2)).await.unwrap().set("k", "v").await.unwrap();
let mut writer =
sqlx::sqlite::SqliteConnectOptions::new().filename(&path).connect().await.unwrap();
db::query("BEGIN EXCLUSIVE").execute(&mut writer).await.unwrap();
db::query("UPDATE meta SET value = 'w' WHERE key = 'k'")
.execute(&mut writer)
.await
.unwrap();
let commit = tokio::spawn(async move {
tokio::time::sleep(Duration::from_secs(6)).await;
db::query("COMMIT").execute(&mut writer).await.unwrap();
writer.close().await.unwrap();
});
let store = MetaStore::new(&path, Duration::from_secs(2)).await.unwrap();
assert_eq!(store.get("k").await.unwrap().as_deref(), Some("w"));
commit.await.unwrap();
}
}
crates/atuin-client/src/settings.rs
+7 -1
View File
@@ -1382,6 +1382,9 @@ impl Settings {
let meta = match Self::meta_store().await {
Ok(m) => m,
Err(e) => {
// Unlike not being logged in, a failure: callers tracing the result at debug
// would hide it.
tracing::warn!("failed to open the meta store for sync auth: {e}");
return SyncAuth::NotLoggedIn {
reason: format!("Failed to open meta store: {e}"),
};
@@ -1436,8 +1439,11 @@ impl Settings {
/// Delegates to [`Self::resolve_sync_auth`] and converts the result to an
/// `AuthToken`. Callers that need to distinguish between auth states
/// (e.g. to show different UI) should call `resolve_sync_auth` directly.
///
/// Not being logged in is an expected state (an offline daemon asks on every start), so the
/// error is traced at debug, and callers that fail on it report it themselves.
#[cfg(feature = "sync")]
#[instrument(level = "trace", skip_all, err)]
#[instrument(level = "trace", skip_all, err(level = "debug"))]
pub async fn sync_auth_token(&self) -> Result<crate::api_client::AuthToken> {
self.resolve_sync_auth().await.into_auth_token()
}
crates/atuin-common/Cargo.toml
+4 -1
View File
@@ -13,7 +13,7 @@ version.workspace = true
[features]
# Enables `atuin_common::harnesstools`.
ai = ["fs-watcher", "jsonl", "dep:xxhash-rust", "sqlite"]
ai = ["fs-watcher", "jsonl", "dep:sha2", "dep:toml_edit", "dep:xxhash-rust", "sqlite"]
# Enables `atuin_common::ansi`.
ansi = ["dep:atuin-vt100"]
# Enables `atuin_common::db` (but not `db::sqlite`).
@@ -71,6 +71,8 @@ semver = { workspace = true }
serde = { workspace = true }
serde_json = { workspace = true }
serde_with = { workspace = true }
# Codex keys the hooks it trusts on their SHA-256.
sha2 = { workspace = true, optional = true }
shlex = { workspace = true }
sqlx = { workspace = true, optional = true }
strum = { workspace = true }
@@ -82,6 +84,7 @@ time = { workspace = true, features = ["formatting", "parsing"] }
tiny-bip39 = { workspace = true }
tokio = { workspace = true }
tokio-util = { workspace = true, optional = true }
toml_edit = { workspace = true, optional = true }
tracing = { workspace = true }
tracing-subscriber = { workspace = true, optional = true }
typed-builder = { workspace = true }
crates/atuin-common/res/harnesstools/opencode/atuin.v2.ts
+228
View File
@@ -0,0 +1,228 @@
/**
* Atuin plugin for opencode 2.
*
* Tracks shell commands executed by opencode in Atuin history with author
* `opencode`. (opencode 1.x takes a plugin of another shape: `atuin hook
* install opencode` installs the one for the opencode it finds.)
*
* Install with:
* atuin hook install opencode
*
* Then restart opencode.
*/
import { spawn } from "node:child_process";
const ATUIN_AUTHOR = "opencode";
const ATUIN_TIMEOUT_MS = 10_000;
const SHELL_TOOL = "shell";
// A command killed for running too long reports no exit code; record the
// conventional one.
const EXIT_TIMED_OUT = 124;
// A denied command never completes, so its proposal is never claimed and would
// sit in the map for the life of the session. Bound the map to cap that.
// Evicting the oldest can in principle drop a command still waiting at its
// permission prompt, which just goes unrecorded, but only a handful of calls
// are ever genuinely in flight at once.
const MAX_PROPOSED = 128;
interface Proposal {
command: string;
intent?: string;
started?: Started;
}
interface Started {
shell: ShellCreate;
at: bigint;
}
interface AtuinResult {
code: number | null;
stdout: string;
}
// What opencode hands its tool hooks: the call (`id`) and its input, and after
// it ran, how it ended.
interface ToolCall {
tool: string;
id: string;
input: unknown;
status?: "completed" | "error";
result?: { metadata?: { exit?: unknown; status?: unknown; timeout?: unknown } };
}
// What opencode hands `create.before`: the shell it is about to start. Its
// `cwd` is made absolute once the command is cleared to run.
interface ShellCreate {
command: string;
cwd: string;
}
interface Hooks {
tool: {
hook(name: string, fn: (call: ToolCall) => void | Promise<void>): unknown;
};
shell: {
hook(name: string, fn: (shell: ShellCreate) => void | Promise<void>): unknown;
};
}
/**
* Run Atuin, resolving rather than rejecting on every failure.
*
* Spawned without a shell: a recorded command is arbitrary text that has to
* reach Atuin as a single argv entry unmangled.
*/
function atuin(args: string[], cwd: string): Promise<AtuinResult> {
return new Promise((resolve) => {
let child: ReturnType<typeof spawn>;
try {
child = spawn("atuin", args, { cwd, stdio: ["ignore", "pipe", "ignore"] });
} catch {
resolve({ code: null, stdout: "" });
return;
}
let stdout = "";
let settled = false;
const timer = setTimeout(() => child.kill("SIGKILL"), ATUIN_TIMEOUT_MS);
const settle = (result: AtuinResult) => {
if (settled) return;
settled = true;
clearTimeout(timer);
resolve(result);
};
child.stdout?.setEncoding("utf8");
child.stdout?.on("data", (chunk: string) => {
stdout += chunk;
});
child.on("error", () => settle({ code: null, stdout: "" }));
child.on("close", (code) => settle({ code, stdout }));
});
}
async function startHistory(
cwd: string,
proposal: Proposal,
): Promise<string | undefined> {
const args = [
"history",
"start",
"--author",
ATUIN_AUTHOR,
"--author-kind",
"agent",
];
if (proposal.intent) args.push("--intent", proposal.intent);
args.push("--", proposal.command);
const result = await atuin(args, cwd);
if (result.code !== 0) return undefined;
const historyId = result.stdout.trim();
return historyId.length > 0 ? historyId : undefined;
}
// A rejected hook aborts opencode's tool call and discards the command's
// output. A missing history entry is always the better failure.
async function swallowFailures(work: () => void | Promise<void>): Promise<void> {
try {
await work();
} catch {
// Deliberately ignored.
}
}
// The exit code a command that completed ended with: the shell tool completes
// a command whatever its exit code, and one it killed for running too long has
// none.
function exitCodeOf(call: ToolCall): number {
const metadata = call.result?.metadata;
if (typeof metadata?.exit === "number") return metadata.exit;
if (metadata?.timeout === true) return EXIT_TIMED_OUT;
return 1;
}
// opencode's hooks before a shell runs (`execute.before`, `create.before`) both
// come before the permission check, and none comes between that check and the
// command starting, so a command is recorded once it has ended: started then,
// and ended with the time since its shell was created. That time includes any
// wait at the permission prompt, and the entry's timestamp is when the command
// ended. A command that was denied ends the call in an error and is not
// recorded; so does one the user cancels, which is not recorded either. One
// left running in the background (`status: "running"`) has not ended, and is
// not recorded.
function setup(ctx: Hooks) {
// Commands proposed, keyed by tool call ID.
const proposed = new Map<string, Proposal>();
ctx.tool.hook("execute.before", (call) =>
swallowFailures(() => {
if (call.tool !== SHELL_TOOL) return;
const { command, description } = (call.input ?? {}) as {
command?: unknown;
description?: unknown;
};
if (typeof command !== "string" || command.length === 0) return;
if (proposed.size >= MAX_PROPOSED) {
const oldest = proposed.keys().next().value;
if (oldest !== undefined) proposed.delete(oldest);
}
proposed.set(call.id, {
command,
intent:
typeof description === "string" && description.length > 0
? description
: undefined,
});
}),
);
// Carries no call ID: it goes to the oldest proposal of the same command
// whose shell has not been created yet.
ctx.shell.hook("create.before", (shell) =>
swallowFailures(() => {
for (const proposal of proposed.values()) {
if (proposal.started || proposal.command !== shell.command) continue;
proposal.started = { shell, at: process.hrtime.bigint() };
return;
}
}),
);
ctx.tool.hook("execute.after", (call) =>
swallowFailures(async () => {
if (call.tool !== SHELL_TOOL) return;
const proposal = proposed.get(call.id);
proposed.delete(call.id);
if (!proposal?.started || call.status !== "completed") return;
if (call.result?.metadata?.status === "running") return;
const { shell, at } = proposal.started;
const duration = process.hrtime.bigint() - at;
const historyId = await startHistory(shell.cwd, proposal);
if (!historyId) return;
await atuin(
[
"history",
"end",
historyId,
"--exit",
String(exitCodeOf(call)),
"--duration",
String(duration),
],
shell.cwd,
);
}),
);
}
// opencode 2 loads a plugin's default export: its `id` and `setup`.
export default { id: "atuin", setup };
crates/atuin-common/src/db/sqlite/builder.rs
+13
View File
@@ -27,6 +27,7 @@ pub enum Journaling {
pub struct SqliteBuilder<'a> {
input: &'a OsStr,
timeout: Duration,
busy_timeout: Option<Duration>,
journal: Option<Journaling>,
synchronous: SqliteSynchronous,
foreign_keys: bool,
@@ -75,6 +76,7 @@ impl<'a> SqliteBuilder<'a> {
Self {
input,
timeout: Duration::from_secs(5),
busy_timeout: None,
journal: Some(Journaling::Wal {
max_size_hint: Self::DEFAULT_MAX_WAL_SIZE,
}),
@@ -92,6 +94,14 @@ impl<'a> SqliteBuilder<'a> {
self
}
/// How long a statement waits for another connection's lock before failing with "database is
/// locked" (sqlite's busy timeout): sqlx's 5s when unset.
#[must_use]
pub fn busy_timeout(mut self, timeout: Duration) -> Self {
self.busy_timeout = Some(timeout);
self
}
#[must_use]
pub fn journal(mut self, journal: Option<Journaling>) -> Self {
self.journal = journal;
@@ -148,6 +158,9 @@ impl<'a> SqliteBuilder<'a> {
.foreign_keys(self.foreign_keys)
.read_only(self.read_only)
.create_if_missing(!self.read_only);
if let Some(busy) = self.busy_timeout {
opts = opts.busy_timeout(busy);
}
match self.journal {
// The journal mode is the owner's to set: changing it needs a write.
crates/atuin-common/src/harnesstools/access.rs
+98
View File
@@ -6,6 +6,8 @@ use serde_json::{Map, Value};
use crate::harnesstools::note::command_line;
mod script;
/// What a tool call touches.
#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub struct Access {
@@ -24,6 +26,23 @@ const PATH_KEYS: &[&str] = &["file_path", "filePath", "path", "notebook_path", "
const PATHS_KEYS: &[&str] = &["paths", "file_paths", "filePaths", "files"];
impl Access {
/// What the call `name`, with `input`, touches: what [`Self::of`] finds, or for a Codex `exec`
/// script (`input` the JavaScript it runs, see the `script` module), what each tool call it
/// makes does; and whether that is all it touches. It isn't for a script not read whole (a
/// call's arguments not all literals, `tools` used other than in such calls, a string, regex or
/// comment left open): what that runs or writes can't be told, nor judged in pieces against
/// filters written for whole commands (`^psql`), and only the calls that were read are given.
#[must_use]
pub fn all(name: &str, input: &Value) -> (Vec<Self>, bool) {
let script = match input {
Value::String(text) if name == "exec" && object(input).is_none() => text,
_ => return (vec![Self::of(name, input)], true),
};
let calls = script::calls(script);
let accesses = calls.calls.iter().map(|call| Self::of(&call.name, &call.input)).collect();
(accesses, calls.whole)
}
/// What the call `name`, with `input`, touches. Tools are known by name, case aside and
/// without an MCP server's prefix (`mcp__fs__write_file`): shells by the command they run, and
/// a writer by the paths it's given; any other tool given a path reads it (a read, a search,
@@ -235,4 +254,83 @@ mod tests {
fn what_a_call_touches(#[case] name: &str, #[case] input: Value, #[case] want: Access) {
assert_eq!(Access::of(name, &input), want);
}
/// Codex 0.160 runs every tool through a script (`exec`); each call it makes is an access.
/// The scripts are as real runs wrote them.
#[rstest]
#[case::one_command(
"const r = await tools.exec_command({cmd:\"echo dropme-kilopear\",workdir:\"/w\",yield_time_ms:10000,max_output_tokens:100});\ntext(r.output);\n",
vec![shell("echo dropme-kilopear", Some("/w"))],
)]
#[case::quoted_keys_and_escapes(
r#"text(await tools.exec_command({cmd:"rg -n 'a|b' x \\(y\\) \"q\"","max_output_tokens":5000}));"#,
vec![shell(r#"rg -n 'a|b' x \(y\) "q""#, None)],
)]
#[case::parallel(
"const results = await Promise.allSettled([\ntools.exec_command({cmd:\"jj st\",\"max_output_tokens\":4000}),\ntools.exec_command({cmd:'cat .env', max_output_tokens:5000})]);\nresults.forEach(text);",
vec![shell("jj st", None), shell("cat .env", None)],
)]
#[case::template(
"await tools.exec_command({cmd: `printf 'a\\nb' > .env`})",
vec![Access { writes: vec![".env".into()], ..shell("printf 'a\nb' > .env", None) }],
)]
#[case::patch_by_name(
"const patch = \"*** Begin Patch\\n*** Add File: fox.txt\\n+foxtrotpear\\n*** End Patch\";\nconst r = await tools.apply_patch(patch);\ntext(r);\n",
vec![writes(&["fox.txt"])],
)]
#[case::regex_holding_a_quote(
"text(ALL_TOOLS.filter(x=>/\"|search/i.test(x.name)));\ntext(await tools.exec_command({cmd:\"ls -a .codex\"}));",
vec![shell("ls -a .codex", None)],
)]
#[case::no_tools("text(1 + 1)", vec![])]
fn what_a_script_touches(#[case] script: &str, #[case] want: Vec<Access>) {
assert_eq!(Access::all("exec", &json!(script)).0, want);
}
/// What a script runs isn't known where it builds a command or a patch from pieces, or uses
/// `tools` some way not read (from review and validation's adversarial scripts); it is where
/// every call is of literals.
#[rstest]
#[case::joined_list(
"await tools.exec_command({cmd: [\"echo\", \"dropme-e\"].join(\" \")});",
false
)]
#[case::split_joined(
"await tools.exec_command({cmd: \"echo,dropme-v\".split(\",\").join(\" \")});",
false
)]
#[case::template_of_pieces(
"await tools.exec_command({cmd: `echo ${\"drop\" + \"me\"}-h`});",
false
)]
#[case::appended(
"let c = \"echo\"; c += \" dropme-u\"; await tools.exec_command({cmd: c});",
false
)]
#[case::computed_workdir(
"await tools.exec_command({cmd: \"ls\", workdir: base + \"/secret\"});",
false
)]
#[case::patch_interpolated(
"await tools.apply_patch(`*** Begin Patch\\n*** Update File: ${name}\\n+X=3\\n*** End \
Patch`);",
false
)]
#[case::indexed("await tools[\"apply_patch\"](patch);", false)]
#[case::literal(
"const p = \"*** Begin Patch\\n*** End Patch\";\nawait tools.apply_patch(p);\ntext(`done \
${1}`);",
true
)]
fn what_a_script_runs_is_known_only_read_whole(#[case] script: &str, #[case] known: bool) {
assert_eq!(Access::all("exec", &json!(script)).1, known);
}
/// Only Codex's `exec` script is read as one; any other call is one access.
#[rstest]
#[case::other_tool("Bash", json!({"command": "ls"}))]
#[case::exec_with_arguments("exec", json!(r#"{"cmd":"ls"}"#))]
fn a_call_not_a_script_is_one_access(#[case] name: &str, #[case] input: Value) {
assert_eq!(Access::all(name, &input), (vec![Access::of(name, &input)], true));
}
}
crates/atuin-common/src/harnesstools/access/script.rs
+623
View File
@@ -0,0 +1,623 @@
//! The tool calls a Codex `exec` script makes. Since 0.160 Codex runs every tool the model asks
//! for through one (codex-rs `tools/code_mode`): the model writes JavaScript calling them as
//! `tools.<name>(<arguments>)`, as in `await tools.exec_command({cmd: "ls", workdir: "/w"})` or
//! `const patch = "*** Begin Patch ..."; await tools.apply_patch(patch)`.
//!
//! No JavaScript is run or parsed whole: the script is split into tokens, and the arguments of
//! each `tools.<name>(...)` read where they are literals (strings in any quotes, numbers, objects
//! and lists of them) or a name an earlier `const` gave a string, number or `true`/`false`/`null`.
//! Where that may have missed something (a string, regex or comment left open, a name escaped,
//! more uses of `tools` in the text than calls read, or a name read that the script mentions
//! anywhere else, as it may then hold something else), the script is not read
//! [`whole`](Script::whole).
use std::collections::HashMap;
use std::sync::LazyLock;
use regex::Regex;
use serde_json::{Map, Value};
/// One tool call a script makes, and its arguments as JSON.
#[derive(Debug, PartialEq)]
pub(super) struct Call {
pub name: String,
pub input: Value,
}
/// What [`calls`] makes of a script.
#[derive(Debug, PartialEq)]
pub(super) struct Script {
/// The calls it makes, with their arguments as far as they are literals.
pub calls: Vec<Call>,
/// Whether that is all of it: every call's arguments were literals, every use of `tools`
/// such a call, and nothing in the script was left open.
pub whole: bool,
}
/// Where a script uses `tools`, however it reads (`tools.x`, `tools [`), in strings and
/// comments too.
static TOOLS: LazyLock<Regex> =
LazyLock::new(|| Regex::new(r"\btools\s*[.\[]").expect("a valid regex"));
/// The tool calls `script` makes.
pub(super) fn calls(script: &str) -> Script {
let (tokens, closed) = tokens(script);
let mut parser = Parser {
tokens: &tokens,
at: 0,
bindings: Map::new(),
reads: Vec::new(),
whole: closed,
};
let mut calls = Vec::new();
while let Some(token) = tokens.get(parser.at) {
let after_dot = parser.at > 0 && tokens[parser.at - 1] == Token::Punct('.');
match token {
Token::Ident(word) if word == "const" => {
parser.at += 1;
parser.binding();
}
Token::Ident(word) if word == "tools" && !after_dot => {
parser.at += 1;
match parser.call() {
Some(call) => calls.push(call),
None => parser.whole = false,
}
}
_ => parser.at += 1,
}
}
// A misread `/` (division taken for a regex, or the other way about) can hide calls in what
// it takes for a string: then the text has more of them than were read.
let whole = parser.whole
&& TOOLS.find_iter(script).count() <= calls.len()
&& parser.reads_are_all_mentions();
Script { calls, whole }
}
#[derive(Debug, Clone, PartialEq)]
enum Token {
/// A string, and whether it is one (`false` for a template interpolating `${...}`, kept as
/// written).
Str(String, bool),
Ident(String),
Num(String),
Punct(char),
}
/// `script`'s tokens, and whether every string, regex and comment in it was closed and no name
/// in it escaped.
fn tokens(script: &str) -> (Vec<Token>, bool) {
let chars: Vec<char> = script.chars().collect();
let mut tokens = Vec::new();
let mut closed = true;
let mut i = 0;
while let Some(&c) = chars.get(i) {
let next = chars.get(i + 1).copied();
if c.is_whitespace() {
i += 1;
} else if c == '/' && next == Some('/') {
while chars.get(i).is_some_and(|&c| c != '\n') {
i += 1;
}
} else if c == '/' && next == Some('*') {
i += 2;
while i < chars.len() && !(chars[i] == '*' && chars.get(i + 1) == Some(&'/')) {
i += 1;
}
closed &= i < chars.len();
i += 2;
} else if c == '/' && regex_may_start(tokens.last()) {
let (end, ended) = past_regex(&chars, i + 1);
closed &= ended;
i = end;
} else if matches!(c, '"' | '\'' | '`') {
let (token, ended, end) = string(&chars, i + 1, c);
closed &= ended;
tokens.push(token);
i = end;
} else if c.is_alphabetic() || c == '_' || c == '$' {
let start = i;
while chars.get(i).is_some_and(|&c| c.is_alphanumeric() || c == '_' || c == '$') {
i += 1;
}
tokens.push(Token::Ident(chars[start..i].iter().collect()));
} else if c.is_ascii_digit() {
let start = i;
while chars.get(i).is_some_and(|&c| c.is_ascii_alphanumeric() || c == '.' || c == '_') {
i += 1;
}
tokens.push(Token::Num(chars[start..i].iter().collect()));
} else {
// Outside a string a `\` escapes a name (`\u0074ools` is `tools`): not read.
closed &= c != '\\';
tokens.push(Token::Punct(c));
i += 1;
}
}
(tokens, closed)
}
/// Whether a `/` after `previous` starts a regex rather than dividing what came before it.
fn regex_may_start(previous: Option<&Token>) -> bool {
match previous {
None => true,
Some(Token::Punct(c)) => !matches!(c, ')' | ']' | '}'),
Some(Token::Ident(word)) => matches!(word.as_str(), "return" | "typeof" | "in" | "of"),
Some(Token::Str(..) | Token::Num(_)) => false,
}
}
/// Where the regex whose body starts at `i` ends, its flags and all, and whether it was closed.
fn past_regex(chars: &[char], mut i: usize) -> (usize, bool) {
let mut class = false;
let mut ended = false;
while let Some(&c) = chars.get(i) {
i += 1;
match c {
'\\' => i += 1,
'[' => class = true,
']' => class = false,
'/' if !class => {
ended = true;
break;
}
'\n' => break,
_ => {}
}
}
while chars.get(i).is_some_and(|c| c.is_alphanumeric()) {
i += 1;
}
(i, ended)
}
/// The string whose text starts at `i`, closed by `quote`: it (see [`Token::Str`]), whether it
/// was closed (a quoted string can't hold a line break), and where it ends.
fn string(chars: &[char], mut i: usize, quote: char) -> (Token, bool, usize) {
let mut text = String::new();
let mut literal = true;
let mut ended = false;
while let Some(&c) = chars.get(i) {
i += 1;
match c {
c if c == quote => {
ended = true;
break;
}
'\\' => i = escape(chars, i, &mut text),
'$' if quote == '`' && chars.get(i) == Some(&'{') => {
literal = false;
text.push('$');
let mut depth = 0usize;
while let Some(&c) = chars.get(i) {
text.push(c);
i += 1;
match c {
'{' => depth += 1,
'}' if depth <= 1 => break,
'}' => depth -= 1,
_ => {}
}
}
}
'\n' if quote != '`' => break,
c => text.push(c),
}
}
(Token::Str(text, literal), ended, i)
}
/// Decode the escape whose character is at `i` into `text`; where it ends.
fn escape(chars: &[char], i: usize, text: &mut String) -> usize {
let Some(&c) = chars.get(i) else {
return i;
};
let hex = |from: usize, len: usize| -> Option<char> {
let digits: String = chars.get(from..from + len)?.iter().collect();
char::from_u32(u32::from_str_radix(&digits, 16).ok()?)
};
match c {
'n' => text.push('\n'),
't' => text.push('\t'),
'r' => text.push('\r'),
'b' => text.push('\u{8}'),
'f' => text.push('\u{c}'),
'v' => text.push('\u{b}'),
'0' => text.push('\0'),
// A line continuation.
'\n' => {}
'x' => {
if let Some(c) = hex(i + 1, 2) {
text.push(c);
return i + 3;
}
text.push('x');
}
'u' if chars.get(i + 1) == Some(&'{') => {
let close = chars[i..].iter().position(|&c| c == '}').map(|at| i + at);
if let Some(c) = close.and_then(|close| hex(i + 2, close - i - 2)) {
text.push(c);
return close.unwrap_or(i) + 1;
}
text.push('u');
}
'u' => {
if let Some(c) = hex(i + 1, 4) {
text.push(c);
return i + 5;
}
text.push('u');
}
c => text.push(c),
}
i + 1
}
struct Parser<'a> {
tokens: &'a [Token],
at: usize,
/// The literals the script named so far (`const patch = "..."`).
bindings: Map<String, Value>,
/// The names read for their values, once each time.
reads: Vec<String>,
/// Whether every use of `tools` so far was a call whose arguments were literals.
whole: bool,
}
impl Parser<'_> {
fn peek(&self) -> Option<&Token> {
self.tokens.get(self.at)
}
fn eat(&mut self, punct: char) -> bool {
let found = self.peek() == Some(&Token::Punct(punct));
if found {
self.at += 1;
}
found
}
/// Whether every name read is mentioned only where it is given its value and where it was
/// read: anywhere else (declared again in a block, a parameter, assigned to) it may hold
/// something else by the time it is read. A key (`{cmd: cmd}`) is not a mention.
fn reads_are_all_mentions(&self) -> bool {
let mut reads: HashMap<&str, usize> = HashMap::new();
for name in &self.reads {
*reads.entry(name).or_default() += 1;
}
let mut mentions: HashMap<&str, usize> = HashMap::new();
let next = self.tokens.iter().skip(1).map(Some).chain([None]);
for (token, next) in self.tokens.iter().zip(next) {
if let Token::Ident(word) = token
&& reads.contains_key(word.as_str())
&& next != Some(&Token::Punct(':'))
{
*mentions.entry(word).or_default() += 1;
}
}
reads.iter().all(|(name, read)| mentions.get(name) == Some(&(1 + read)))
}
/// `<name> = <literal>`, after a `const`. Objects and lists are not kept, as what they hold
/// can be changed without naming them again (in a `${...}`).
fn binding(&mut self) {
let Some(Token::Ident(name)) = self.peek().cloned() else {
return;
};
self.at += 1;
if !self.eat('=') {
return;
}
// Whether the value is all literal says nothing of the calls, unless one is given it.
let (at, reads) = (self.at, self.reads.len());
let whole = std::mem::replace(&mut self.whole, true);
match self.value() {
Some(value @ (Value::String(_) | Value::Number(_) | Value::Bool(_) | Value::Null))
if self.whole && self.ends_value() =>
{
self.bindings.insert(name, value);
}
// Not a literal: the next token is looked at again (it may be a call of `tools`).
_ => {
self.at = at;
self.reads.truncate(reads);
}
}
self.whole = whole;
}
/// `.<name>(<arguments>)`, after `tools`; `None` when that isn't what follows, or the
/// arguments are not all literals.
fn call(&mut self) -> Option<Call> {
if !self.eat('.') {
return None;
}
let Some(Token::Ident(name)) = self.peek().cloned() else {
return None;
};
self.at += 1;
if !self.eat('(') {
return None;
}
let input = if self.peek() == Some(&Token::Punct(')')) {
Value::Null
} else {
match self.value() {
Some(value) if self.ends_value() => value,
_ => {
self.whole = false;
Value::Null
}
}
};
Some(Call { name, input })
}
/// Whether the value just read is all of the expression it is in.
fn ends_value(&self) -> bool {
matches!(
self.peek(),
None | Some(Token::Punct(',' | ')' | '}' | ']' | ';') | Token::Ident(_))
)
}
/// A literal value, or a name bound to one, strings joined by `+` joined. An object or list
/// holding something else is read as far as it is literal, the rest `null`, and the script no
/// longer `whole`.
fn value(&mut self) -> Option<Value> {
let mut value = self.operand()?;
while let Value::String(text) = &mut value
&& self.peek() == Some(&Token::Punct('+'))
{
self.at += 1;
match self.operand()? {
Value::String(more) => text.push_str(&more),
_ => return None,
}
}
Some(value)
}
fn operand(&mut self) -> Option<Value> {
let token = self.peek()?.clone();
self.at += 1;
Some(match token {
Token::Str(text, literal) => {
self.whole &= literal;
Value::String(text)
}
Token::Num(digits) => serde_json::from_str(&digits).unwrap_or(Value::String(digits)),
Token::Ident(word) => match word.as_str() {
"true" => Value::Bool(true),
"false" => Value::Bool(false),
"null" | "undefined" => Value::Null,
_ => {
let value = self.bindings.get(&word).cloned()?;
self.reads.push(word);
value
}
},
Token::Punct('{') => self.object(),
Token::Punct('[') => self.list(),
Token::Punct(_) => return None,
})
}
/// The rest of an object, after its `{`.
fn object(&mut self) -> Value {
let mut object = Map::new();
loop {
if self.eat('}') {
return Value::Object(object);
}
let key = match self.peek().cloned() {
Some(Token::Ident(key) | Token::Str(key, true) | Token::Num(key)) => key,
// A bracket closing something else ends it: there is nothing more to skip.
Some(_) => {
let at = self.at;
self.skip_member();
if self.at == at {
return Value::Object(object);
}
continue;
}
None => return Value::Object(object),
};
self.at += 1;
let value = if self.eat(':') {
self.member_value()
} else {
// `{cmd}`: the value of the name.
if let Some(value) = self.bindings.get(&key).cloned() {
self.reads.push(key.clone());
value
} else {
self.whole = false;
Value::Null
}
};
object.insert(key, value);
self.eat(',');
}
}
/// The rest of a list, after its `[`.
fn list(&mut self) -> Value {
let mut items = Vec::new();
loop {
if self.eat(']') || self.peek().is_none() {
return Value::Array(items);
}
let at = self.at;
items.push(self.member_value());
self.eat(',');
// A bracket closing something else ends it (and `whole`, skipped as not a literal).
if self.at == at {
return Value::Array(items);
}
}
}
/// A value inside an object or list: `null` where it is not a literal, with the rest of the
/// expression skipped.
fn member_value(&mut self) -> Value {
let (at, reads) = (self.at, self.reads.len());
match self.value() {
Some(value) if self.ends_value() && !matches!(self.peek(), Some(Token::Ident(_))) => {
value
}
_ => {
self.at = at;
self.reads.truncate(reads);
self.skip_member();
Value::Null
}
}
}
/// Skip what is left of a member that is not a literal, to the `,` after it or the bracket
/// closing its object or list; the script is no longer `whole`.
fn skip_member(&mut self) {
self.whole = false;
let mut depth = 0usize;
while let Some(token) = self.peek() {
match token {
Token::Punct('(' | '[' | '{') => depth += 1,
Token::Punct(')' | ']' | '}') if depth == 0 => return,
Token::Punct(')' | ']' | '}') => depth -= 1,
Token::Punct(',') if depth == 0 => {
self.at += 1;
return;
}
_ => {}
}
self.at += 1;
}
}
}
#[cfg(test)]
mod tests {
use proptest::prelude::*;
use rstest::rstest;
use serde_json::json;
use super::*;
fn call(name: &str, input: Value) -> Call {
Call {
name: name.to_owned(),
input,
}
}
#[rstest]
#[case::single_quotes(r"tools.exec_command({cmd: 'it\'s'})", json!({"cmd": "it's"}))]
#[case::hex_and_unicode(r#"tools.exec_command({cmd: "\x41B\u{43}"})"#, json!({"cmd": "ABC"}))]
#[case::continuation("tools.exec_command({cmd: \"a\\\nb\"})", json!({"cmd": "ab"}))]
#[case::nested("tools.exec_command({cmd: \"ls\", env: {A: [1, true, null]}})", json!({"cmd": "ls", "env": {"A": [1, true, null]}}))]
#[case::shorthand("const cmd = 'ls'; tools.exec_command({cmd})", json!({"cmd": "ls"}))]
#[case::named("const cmd = 'ls'; tools.exec_command({cmd: cmd})", json!({"cmd": "ls"}))]
#[case::joined("const dir = '/w'; tools.exec_command({cmd: 'cd ' + dir + \" && ls\"})", json!({"cmd": "cd /w && ls"}))]
#[case::comments("/* read */ tools.exec_command({cmd: 'ls' /* , cmd: 'rm' */}) // done", json!({"cmd": "ls"}))]
fn arguments_are_read_as_literals(#[case] script: &str, #[case] input: Value) {
let found = calls(script);
assert!(found.whole);
assert_eq!(found.calls, vec![call("exec_command", input)]);
}
/// What isn't a literal is `null`, and the script no longer read whole.
#[rstest]
#[case::computed_member("tools.exec_command({cmd: \"echo \" + x, workdir: \"/w\"})", vec![call("exec_command", json!({"cmd": null, "workdir": "/w"}))])]
#[case::unknown_name("tools.apply_patch(patch)", vec![call("apply_patch", Value::Null)])]
#[case::interpolation("tools.exec_command({cmd: `cat ${f}`})", vec![call("exec_command", json!({"cmd": "cat ${f}"}))])]
#[case::reassigned("let c = 'ls'; c = 'psql'; tools.exec_command({cmd: c})", vec![call("exec_command", json!({"cmd": null}))])]
#[case::joined_with_other("tools.exec_command({cmd: 'sleep ' + 1})", vec![call("exec_command", json!({"cmd": null}))])]
#[case::indexed("tools['exec_command']({cmd: 'ls'})", vec![])]
#[case::passed_on("const run = tools.exec_command; run({cmd: 'ls'})", vec![])]
fn what_is_not_literal_is_not_whole(#[case] script: &str, #[case] want: Vec<Call>) {
let found = calls(script);
assert!(!found.whole);
assert_eq!(found.calls, want);
}
/// What may have been misread is not read whole: a `/` taken for a regex or a division,
/// swallowing calls into a string or a regex left open; a string, regex or comment left open;
/// `tools` used more often than calls were read. The first four from review.
#[rstest]
#[case::regex_after_if(
"if (ok) /\"/.test(x); await tools.exec_command({cmd:\"psql -c 1\"}); const y = \"z\";"
)]
#[case::regex_after_block(
"const f = () => {}\n/\"/.test(s); await tools.exec_command({cmd:\"psql\"}); \"\""
)]
#[case::division_after_increment(
"let i = 0; i++ / 2; await tools.exec_command({cmd:\"psql\"});"
)]
#[case::regex_after_empty_block("if(1){} /\"/.test(\"x\"); tools.exec_command({cmd:\"psql\"})")]
#[case::open_string("tools.exec_command({cmd: 'ls'}); 'psql")]
#[case::open_comment("tools.exec_command({cmd: 'ls'}); /* psql")]
#[case::tools_in_a_string("tools.exec_command({cmd: 'ls'}); text('tools.exec_command')")]
#[case::escaped_tools(r"await \u0074ools.exec_command({cmd:'cat .env'})")]
#[case::escaped_tools_braced(r"await tool\u{73}.exec_command({cmd:'cat .env'})")]
#[case::escaped_name(r"tools.\u0065xec_command({cmd:'cat .env'})")]
#[case::escaped_shadow(
r"const a='ls'; { let \u0061 = 'cat .env'; tools.exec_command({cmd:a}) }"
)]
fn what_may_be_misread_is_not_whole(#[case] script: &str) {
assert!(!calls(script).whole);
}
/// A script with a stray closing bracket is read to its end, not forever.
#[rstest]
#[case::paren_in_object("tools.exec_command({cmd:'ls')})")]
#[case::bracket_in_object("tools.exec_command({cmd:'ls']})")]
#[case::paren_in_list("tools.exec_command([1)])")]
#[case::brace_in_list("tools.exec_command({a: [1}]})")]
fn a_stray_closer_ends_the_read(#[case] script: &'static str) {
let (send, receive) = std::sync::mpsc::channel();
std::thread::spawn(move || send.send(calls(script).whole));
let whole = receive.recv_timeout(std::time::Duration::from_secs(5)).expect("the read ends");
assert!(!whole);
}
/// A name is only trusted where nothing else can change what it holds at the call: the
/// first two from review.
#[rstest]
#[case::member_assigned("const a={cmd:'ls'}; a.cmd='cat .env'; await tools.exec_command(a);")]
#[case::shadowed_in_block("const a='ls'; { let a = 'cat .env'; } tools.exec_command({cmd:a})")]
#[case::shadowed_by_parameter(
"const a='ls'; const f = a => tools.exec_command({cmd:a}); f('cat .env')"
)]
#[case::passed_elsewhere("const a={cmd:'ls'}; change(a); tools.exec_command(a)")]
#[case::pushed("const a=['ls']; a.push('cat .env'); tools.exec_command({cmd:a})")]
#[case::incremented("let n=1; n++; tools.exec_command({cmd:'ls', n})")]
#[case::changed_in_template("const a={cmd:'ls'}; `${a.cmd='cat .env'}`; tools.exec_command(a)")]
#[case::reassigned_in_template("let a='ls'; `${a='cat .env'}`; tools.exec_command({cmd:a})")]
fn a_name_that_may_change_is_not_whole(#[case] script: &str) {
assert!(!calls(script).whole);
}
/// A name read many times is checked once, not once per read: a script of thousands of
/// calls reads in linear time.
#[rstest]
fn a_name_read_many_times_reads_whole() {
let script =
format!("const cmd = 'ls';{}", "tools.exec_command({cmd: cmd});".repeat(20_000));
let found = calls(&script);
assert!(found.whole);
assert_eq!(found.calls.len(), 20_000);
}
proptest! {
/// Any command, written as a JSON string (which is a JavaScript one), reads back as it is.
#[test]
fn a_quoted_command_reads_back(command in any::<String>()) {
let script = format!("await tools.exec_command({{cmd:{}}});", json!(command));
let found = calls(&script);
prop_assert!(found.whole);
prop_assert_eq!(found.calls, vec![call("exec_command", json!({"cmd": command}))]);
}
}
}
crates/atuin-common/src/harnesstools/codex/hook_trust.rs
+245
View File
@@ -0,0 +1,245 @@
//! Trusting the hooks atuin installs for Codex. Since 0.160 Codex runs a hook from `hooks.json`
//! only once the user has trusted it: in its `config.toml`, under `[hooks.state."<key>"]`, a
//! `trusted_hash` of the hook as it was trusted, so a hook changed since runs no more. The user
//! asked for atuin's hooks by installing them, so atuin records them as trusted, each as it
//! installed it.
//!
//! The key is where the hook is: `<Codex home>/hooks.json:<event>:<group>:<handler>`, the home's
//! path with its symlinks resolved but not `hooks.json`'s own (as codex 0.160.1 keys them).
use std::path::{Path, PathBuf};
use serde_json::{Value, json};
use sha2::{Digest, Sha256};
use toml_edit::{DocumentMut, Item, Table, TableLike};
use crate::harnesstools::InstallHookError;
use crate::harnesstools::json_hooks::invokes_atuin_hook;
/// The events Codex runs atuin's hooks for, as `hooks.json` and the trust key name them.
const EVENTS: &[(&str, &str)] = &[("PreToolUse", "pre_tool_use"), ("PostToolUse", "post_tool_use")];
/// Codex's `config.toml`, read before anything is installed, so a config atuin can't record
/// trust in leaves nothing half done.
pub(super) struct Config {
path: PathBuf,
doc: DocumentMut,
}
const MALFORMED: InstallHookError =
InstallHookError::Malformed("`hooks.state` in Codex's config.toml is not a table");
/// The `config.toml` of Codex home `home` (empty where there is none yet), refused where it
/// isn't TOML or holds a `hooks.state` that isn't a table.
pub(super) async fn config(home: &Path) -> Result<Config, InstallHookError> {
let path = home.join("config.toml");
let text = match tokio::fs::read_to_string(&path).await {
Ok(text) => text,
Err(err) if err.kind() == std::io::ErrorKind::NotFound => String::new(),
Err(err) => return Err(err.into()),
};
let mut doc: DocumentMut =
text.parse().map_err(|_| InstallHookError::Malformed("Codex's config.toml is not TOML"))?;
state(&mut doc)?;
Ok(Config { path, doc })
}
/// The `hooks.state` table of `doc`, made where there is none (inline or not, as it is).
fn state(doc: &mut DocumentMut) -> Result<&mut dyn TableLike, InstallHookError> {
doc.entry("hooks")
.or_insert_with(implicit)
.as_table_like_mut()
.and_then(|hooks| hooks.entry("state").or_insert_with(implicit).as_table_like_mut())
.ok_or(MALFORMED)
}
/// Record the hooks running `atuin hook <harness>` in `home/hooks.json` as trusted in `config`,
/// returning whether that changed it.
pub(super) async fn trust(
home: &Path,
harness: &str,
mut config: Config,
) -> Result<bool, InstallHookError> {
let hooks_path = tokio::fs::canonicalize(home).await?.join("hooks.json");
let hooks: Value = serde_json::from_str(&tokio::fs::read_to_string(&hooks_path).await?)?;
let mut changed = false;
let state = state(&mut config.doc)?;
for (event, group, handler, hash) in atuin_hooks(&hooks, harness) {
let key = format!("{}:{event}:{group}:{handler}", hooks_path.display());
let hook = state
.entry(&key)
.or_insert_with(|| Item::Table(Table::new()))
.as_table_like_mut()
.ok_or(InstallHookError::Malformed("a hook's state in Codex's config.toml"))?;
if hook.get("trusted_hash").and_then(Item::as_str) != Some(hash.as_str()) {
hook.insert("trusted_hash", toml_edit::value(hash));
changed = true;
}
}
if changed {
tokio::fs::write(&config.path, config.doc.to_string()).await?;
}
Ok(changed)
}
fn implicit() -> Item {
let mut table = Table::new();
table.set_implicit(true);
Item::Table(table)
}
/// The hooks running `atuin hook <harness>` in a `hooks.json`, as atuin installs them (a
/// command alone): each one's event, as the trust key names it, where it is (its group and its
/// place in it), and the hash Codex trusts it by.
fn atuin_hooks(hooks: &Value, harness: &str) -> Vec<(&'static str, usize, usize, String)> {
let mut found = Vec::new();
for (event, key) in EVENTS {
let Some(groups) = hooks["hooks"][event].as_array() else {
continue;
};
for (g, group) in groups.iter().enumerate() {
let Some(handlers) = group["hooks"].as_array() else {
continue;
};
for (h, handler) in handlers.iter().enumerate() {
let Some(command) = handler["command"].as_str() else {
continue;
};
let installed = *handler == json!({"type": "command", "command": command});
if installed && invokes_atuin_hook(command, harness) {
found.push((*key, g, h, hash(key, &group["matcher"], command)));
}
}
}
}
found
}
/// What Codex trusts a hook by: the SHA-256 of the JSON of its event, matcher and handler (its
/// defaults filled in), keys sorted.
fn hash(event: &str, matcher: &Value, command: &str) -> String {
let hook = json!({
"event_name": event,
"hooks": [{"async": false, "command": command, "timeout": 600, "type": "command"}],
"matcher": matcher,
});
format!("sha256:{:x}", Sha256::digest(hook.to_string()))
}
#[cfg(test)]
mod tests {
use rstest::rstest;
use super::*;
const COMMAND: &str = "/opt/atuin/bin/atuin hook codex";
/// The hashes codex 0.160.1 itself gave these hooks, as its app server's `hooks/list`
/// reported them for a `hooks.json` holding them (`currentHash`): of the command a dev build
/// installed, and of [`COMMAND`].
#[rstest]
#[case::pre_dev(
"pre_tool_use",
"/home/ellie/src/mono/.claude/worktrees/agent-a6cdedcfcffbcf53a/atuin/target/debug/atuin \
hook codex",
"sha256:5c43193105a163fcd3fed7f209bb1b1169cece43716091bc6dac24d2a7a059ab"
)]
#[case::post_dev(
"post_tool_use",
"/home/ellie/src/mono/.claude/worktrees/agent-a6cdedcfcffbcf53a/atuin/target/debug/atuin \
hook codex",
"sha256:3621fd71768b67a9c87d2732c34f77bd5403afb7f27226388dced73e52c18c78"
)]
#[case::pre(
"pre_tool_use",
COMMAND,
"sha256:3b981b43c4878ddcca45ed5c03b71b76508a73e7c8492cdb3266c0867adade82"
)]
#[case::post(
"post_tool_use",
COMMAND,
"sha256:274cc2128367ee69da364e7677ebe7489cc8f29de94635225f2e40735e09e642"
)]
fn hashes_are_codexs(#[case] event: &str, #[case] command: &str, #[case] expected: &str) {
assert_eq!(hash(event, &json!("^Bash$"), command), expected);
}
fn installed() -> Value {
json!({"matcher": "^Bash$", "hooks": [{"type": "command", "command": COMMAND}]})
}
/// Atuin's hooks are trusted where they are, beside the user's, which are left untrusted;
/// done again, nothing changes. A `hooks` table written inline is written to as it is.
#[rstest]
#[case::tables("model = \"o3\"\n")]
#[case::inline("model = \"o3\"\nhooks = { state = {} }\n")]
#[tokio::test]
async fn atuins_hooks_are_trusted_where_they_are(#[case] before: &str) {
let home = tempfile::tempdir().unwrap();
let theirs = json!({"matcher": "", "hooks": [{"type": "command", "command": "echo hi"}]});
let hooks = json!({"hooks": {
"PreToolUse": [theirs, installed()],
"PostToolUse": [installed()],
"PostToolUseFailure": [installed()],
}});
std::fs::write(home.path().join("hooks.json"), hooks.to_string()).unwrap();
std::fs::write(home.path().join("config.toml"), before).unwrap();
let config_of = || config(home.path());
assert!(trust(home.path(), "codex", config_of().await.unwrap()).await.unwrap());
assert!(!trust(home.path(), "codex", config_of().await.unwrap()).await.unwrap());
let path = std::fs::canonicalize(home.path()).unwrap().join("hooks.json");
let text = std::fs::read_to_string(home.path().join("config.toml")).unwrap();
let config: toml_edit::DocumentMut = text.parse().unwrap();
let state = config["hooks"]["state"].as_table_like().unwrap();
let trusted: Vec<(String, String)> = state
.iter()
.map(|(key, hook)| (key.to_owned(), hook["trusted_hash"].as_str().unwrap().to_owned()))
.collect();
pretty_assertions::assert_eq!(trusted, [
(
format!("{}:pre_tool_use:1:0", path.display()),
hash("pre_tool_use", &json!("^Bash$"), COMMAND)
),
(
format!("{}:post_tool_use:0:0", path.display()),
hash("post_tool_use", &json!("^Bash$"), COMMAND)
),
]);
assert!(text.starts_with(before.lines().next().unwrap()), "{text}");
}
/// Codex keys a hook by its home's real path, `hooks.json` itself as it is named there, even
/// where that is a link (a dotfiles checkout's).
#[cfg(unix)]
#[rstest]
#[tokio::test]
async fn the_key_names_hooks_json_in_codexs_real_home() {
let dir = tempfile::tempdir().unwrap();
let real = dir.path().join("real");
std::fs::create_dir(&real).unwrap();
std::os::unix::fs::symlink(&real, dir.path().join("home")).unwrap();
let home = dir.path().join("home");
let hooks = json!({"hooks": {"PreToolUse": [installed()]}});
std::fs::write(dir.path().join("dotfiles.json"), hooks.to_string()).unwrap();
std::os::unix::fs::symlink(dir.path().join("dotfiles.json"), real.join("hooks.json"))
.unwrap();
trust(&home, "codex", config(&home).await.unwrap()).await.unwrap();
let text = std::fs::read_to_string(real.join("config.toml")).unwrap();
let real = std::fs::canonicalize(&real).unwrap();
assert!(text.contains(&format!("{}:pre_tool_use:0:0", real.join("hooks.json").display())));
}
/// A config whose `hooks.state` isn't a table is refused before anything is written.
#[rstest]
#[case::not_toml("model = ")]
#[case::state_not_a_table("hooks = { state = 1 }\n")]
#[tokio::test]
async fn a_config_trust_cant_be_recorded_in_is_refused(#[case] text: &str) {
let home = tempfile::tempdir().unwrap();
std::fs::write(home.path().join("config.toml"), text).unwrap();
assert!(matches!(config(home.path()).await, Err(InstallHookError::Malformed(_))));
}
}
crates/atuin-common/src/harnesstools/codex/mod.rs
+22 -5
View File
@@ -3,12 +3,18 @@ use std::path::PathBuf;
use super::rehydrate::{RehydrateError, RehydrateSession};
use super::resume::{self, CwdRequirement, ResumeError, ResumePlan, ResumeTarget};
use super::{Harness, InstallHookError, json_hooks};
use crate::utils::home_dir;
use crate::utils::{env_nonempty, home_dir};
mod hook_trust;
pub mod rehydrate;
pub mod session;
pub(crate) mod state_db;
/// Codex's home: `$CODEX_HOME`, else `~/.codex`.
pub(crate) fn home() -> PathBuf {
env_nonempty("CODEX_HOME").map_or_else(|| home_dir().join(".codex"), PathBuf::from)
}
#[derive(Debug, Clone, Copy, Default)]
pub struct Codex;
@@ -24,11 +30,22 @@ impl Harness for Codex {
"codex"
}
/// Install the hooks on the current active user's machine
/// Install the hooks on the current active user's machine, trusted (see `hook_trust`).
async fn install_hooks(&self) -> Result<PathBuf, InstallHookError> {
// Codex reads Claude-Code-style hooks from ~/.codex/hooks.json.
let config_path = home_dir().join(".codex").join("hooks.json");
json_hooks::install(&config_path, "^Bash$", self.name()).await?;
// Codex reads Claude-Code-style hooks from $CODEX_HOME/hooks.json.
let home = home();
let config_path = home.join("hooks.json");
let config = hook_trust::config(&home).await?;
let installed = match json_hooks::install(&config_path, "^Bash$", self.name()).await {
Ok(()) => true,
Err(InstallHookError::AlreadyInstalled) => false,
Err(err) => return Err(err),
};
// Hooks installed before atuin trusted them are trusted on installing them again.
let trusted = hook_trust::trust(&home, self.name(), config).await?;
if !installed && !trusted {
return Err(InstallHookError::AlreadyInstalled);
}
Ok(config_path)
}
crates/atuin-common/src/harnesstools/codex/rehydrate.rs
+102 -18
View File
@@ -12,7 +12,8 @@
//! their `id`, or with none when capture keyed them on their `call_id` (`<call_id>#out` for an
//! output);
//! - the files an applied patch changed, as the `item_completed` event of a `FileChange` item
//! (under the patch's call, in the turn it was applied in), which Codex shows the diff from on
//! (under the item it was read from, else the patch's call, in the turn it was applied in, after
//! the call it was made in), which Codex shows the diff from on
//! resume; capture reads it back as the same row, whether it first read a legacy rollout's
//! `patch_apply_end` or a paginated one's `item_completed`;
//! - usage, as the line it was captured from: a `token_usage_record` under its response (or the
@@ -63,6 +64,9 @@
//! - **Tool output**: written back as captured. A call kept with its input but not its output
//! (withheld by a policy, or over the size limit) gets [`UNCAPTURED_OUTPUT`] as its output (an
//! empty `tools` list for a tool search).
//! - **Whether a script's command failed** (Codex 0.160 runs tools through a script, `exec`):
//! Codex says it only in the command's own event, which is not written back; re-captured, the
//! output reads as a success.
//! - Rows whose line carries nothing Codex needs back: thread names (Codex keeps those in its
//! `session_index.jsonl` now) and other events.
//! - Content kinds a line of the kind cannot carry (text inside a tool call, and so on).
@@ -81,7 +85,7 @@ use crate::harnesstools::rehydrate::{
};
use crate::harnesstools::resume::is_plain_name;
use crate::harnesstools::session::{
Change, Content, Patch, Role, StopReason, ToolResult, ToolUse, Usage,
Change, Content, Patch, Role, StopReason, ToolResult, ToolUse, UNREADABLE_SUMMARY, Usage,
};
/// What Codex puts ahead of a compaction summary (codex-rs
@@ -132,7 +136,13 @@ pub(crate) async fn rehydrate_into(
// naming another rollout that exists means Codex has moved the thread on: leave it.
Err(RehydrateError::AlreadyExists(existing)) if existing.starts_with(root) => {
return match super::state_db::point_at(home, &thread, None, &existing).await {
Ok(_) | Err(super::state_db::StateDbError::Elsewhere(_)) => {
Ok(pointed) => {
if pointed == super::state_db::Pointed::Repointed {
forget_history(home, &thread).await;
}
Err(RehydrateError::AlreadyExists(existing))
}
Err(super::state_db::StateDbError::Elsewhere(_)) => {
Err(RehydrateError::AlreadyExists(existing))
}
Err(err) => Err(RehydrateError::Other(format!(
@@ -149,9 +159,18 @@ pub(crate) async fn rehydrate_into(
path.display()
)));
}
forget_history(home, &thread).await;
Ok(path)
}
/// Forget Codex's history of thread `thread`, if it holds one: it was projected from a rollout
/// now gone.
async fn forget_history(home: &Path, thread: &str) {
if let Err(err) = super::state_db::forget_history(home, thread).await {
tracing::warn!(%err, thread, "Codex may go on showing the rollout restored over");
}
}
/// [`rehydrate`] under the sessions directory `root`.
pub(crate) fn write(root: &Path, session: &RehydrateSession) -> Result<PathBuf, RehydrateError> {
if !is_plain_name(&session.id) {
@@ -521,12 +540,15 @@ fn message_lines(
};
}
match message.content.as_slice() {
// One whose summary was unreadable is empty, as Codex writes one the API compacted: Codex
// resumes it from the prompts before it, not the whole conversation.
[Content::Summary(summary)] => {
vec![line(
&at,
"compacted",
json!({"message": format!("{COMPACTION_PREAMBLE}{summary}")}),
)]
let message = if summary == UNREADABLE_SUMMARY {
String::new()
} else {
format!("{COMPACTION_PREAMBLE}{summary}")
};
vec![line(&at, "compacted", json!({ "message": message }))]
}
[Content::Error(why)] => {
let payload = match message.stop_reason {
@@ -552,7 +574,9 @@ fn message_lines(
/// The `item_completed` event of a `FileChange` item, which Codex shows an applied patch's diff
/// from on resume: each file as Codex records it (`{"type": "add", "content"}`, `delete` the
/// same, `{"type": "update", "unified_diff", "move_path"}`). Its turn is filled in by [`turns`].
/// same, `{"type": "update", "unified_diff", "move_path"}`), under the item capture read it
/// from (a Codex 0.160 script's names only itself, `exec-...`), else its call. Its turn is filled
/// in by [`turns`].
fn file_change(
session: &RehydrateSession,
message: &RehydrateMessage,
@@ -574,12 +598,13 @@ fn file_change(
.collect();
let millis =
i64::try_from(message.timestamp.unix_timestamp_nanos() / 1_000_000).unwrap_or_default();
let item = message.source_id.strip_suffix("#patch").unwrap_or(patch.call.as_ref());
line(
at,
"event_msg",
json!({"type": "item_completed", "thread_id": session::resume_id(&session.id),
"turn_id": "",
"item": {"type": "FileChange", "id": patch.call, "changes": changes,
"item": {"type": "FileChange", "id": item, "changes": changes,
"status": "completed"},
"started_at_ms": millis, "completed_at_ms": millis}),
)
@@ -1090,6 +1115,41 @@ pub(crate) mod tests {
})
}
/// A compaction comes back as a `compacted` line: with its summary, or, where that couldn't
/// be read, empty, as Codex writes one the API compacted, which it resumes from the prompts
/// before it rather than the whole conversation.
#[rstest]
#[case::readable("We fixed the parser.", true)]
#[case::unreadable(UNREADABLE_SUMMARY, false)]
#[tokio::test]
async fn a_compaction_comes_back_as_codex_writes_it(
#[case] summary: &str,
#[case] preamble: bool,
) {
let text = |t: &str| vec![Content::Text(t.to_owned())];
let messages = vec![
row(ID, 0, Role::Other("session_meta".into()), vec![]),
row("msg_u1", 1, Role::User, text("fix it")),
row("msg_a1", 2, Role::Assistant, text("fixed")),
row("syn-1", 3, Role::System, vec![Content::Summary(summary.to_owned())]),
row("msg_u2", 4, Role::User, text("thanks")),
];
let lines = turn_lines(&messages);
let compacted = lines.iter().find(|l| l["type"] == "compacted").unwrap();
let message = compacted["payload"]["message"].as_str().unwrap();
assert_eq!(message.starts_with(COMPACTION_PREAMBLE), preamble);
assert_eq!(message.ends_with(summary), preamble);
let root = tempfile::tempdir().unwrap();
let path = write(root.path(), &session(ID, messages)).unwrap();
let summaries: Vec<Content> = captured(ID, &read(ID, path).await)
.into_iter()
.flat_map(|row| row.content)
.filter(|c| matches!(c, Content::Summary(_)))
.collect();
assert_eq!(summaries, [Content::Summary(summary.to_owned())]);
}
/// Rows from older rollouts: items keyed on their `call_id`, a `token_count` keyed on the
/// running total, and one keyed on its timestamp.
#[rstest]
@@ -1138,10 +1198,21 @@ pub(crate) mod tests {
}
/// An applied patch comes back as the `FileChange` item Codex shows its diff from, in the
/// turn it was applied in, and recaptures as the same row.
/// turn it was applied in, and recaptures as the same row: the call's own item, or one of a
/// Codex 0.160 script's, which names only itself.
#[rstest]
#[case::call("call_2#patch", "apply_patch", "*** Begin Patch\n*** End Patch")]
#[case::script(
"exec-2#patch",
"exec",
"await tools.apply_patch(\"*** Begin Patch\\n*** Update File: a.rs\\n*** End Patch\")"
)]
#[tokio::test]
async fn a_patch_comes_back_as_its_file_changes() {
async fn a_patch_comes_back_as_its_file_changes(
#[case] patch_id: &str,
#[case] tool: &str,
#[case] input: &str,
) {
use crate::harnesstools::session::{FilePatch, Hunk};
let update = FilePatch {
@@ -1174,12 +1245,8 @@ pub(crate) mod tests {
let messages = vec![
row(ID, 0, Role::Other("session_meta".into()), vec![]),
row("msg_u1", 1, Role::User, text("fix it")),
row("ctc_2", 2, Role::Assistant, vec![tool_use(
"call_2",
"apply_patch",
json!("*** Begin Patch\n*** End Patch"),
)]),
row("call_2#patch", 3, Role::Tool, vec![Content::Patch(patch)]),
row("ctc_2", 2, Role::Assistant, vec![tool_use("call_2", tool, json!(input))]),
row(patch_id, 3, Role::Tool, vec![Content::Patch(patch)]),
row("ctco_2", 3, Role::Tool, vec![tool_result("call_2", "Done", false)]),
row("msg_a1", 4, Role::Assistant, text("fixed")),
];
@@ -1343,6 +1410,20 @@ pub(crate) mod tests {
assert_eq!(std::fs::read(&first).unwrap(), before);
}
/// Restoring a thread whose rollout is gone forgets the history Codex projected of it from
/// that rollout, which it projects again from the restored one.
#[rstest]
#[tokio::test]
async fn a_restore_forgets_codexs_history_of_the_rollout_gone() {
let home = tempfile::tempdir().unwrap();
let root = home.path().join("sessions");
let thread = session::resume_id(ID);
let history = super::super::state_db::tests::history(home.path(), &[thread]).await;
rehydrate_into(&root, home.path(), &session(ID, vec![])).await.unwrap();
assert_eq!(super::super::state_db::tests::projected(&history, thread).await, 0);
}
/// A rollout whose index row can't be pointed at it is left where it is: Codex may be
/// resuming the thread from it already.
#[rstest]
@@ -1404,6 +1485,7 @@ pub(crate) mod tests {
}
// The rollout an earlier restore wrote, before its index update failed.
let left = write(&root, &session(ID, vec![])).unwrap();
let history = super::super::state_db::tests::history(home.path(), &[thread]).await;
let err = rehydrate_into(&root, home.path(), &session(ID, vec![])).await.unwrap_err();
assert!(matches!(&err, RehydrateError::AlreadyExists(path) if *path == left), "{err}");
@@ -1416,6 +1498,8 @@ pub(crate) mod tests {
.unwrap();
assert_eq!(std::path::Path::new(&named), left);
<sqlx::SqliteConnection as sqlx::Connection>::close(conn).await.unwrap();
// Codex's history of the gone rollout is forgotten too.
assert_eq!(super::super::state_db::tests::projected(&history, thread).await, 0);
}
#[rstest]
crates/atuin-common/src/harnesstools/codex/session.rs
+339 -22
View File
@@ -11,10 +11,11 @@ use tokio::sync::watch;
use typed_builder::TypedBuilder;
use crate::fs::tree_watcher::{FileStat, TreeWatcher};
use crate::harnesstools::access::Access;
use crate::harnesstools::codex::Codex;
use crate::harnesstools::session::model::{
Content, MessageId, ParentKind, Role, StopReason, TitleChange, TitleSource, ToolCallId,
ToolResult, ToolUse, Usage,
ToolResult, ToolUse, UNREADABLE_SUMMARY, Usage,
};
use crate::harnesstools::session::synthetic::{content_hash, synthetic_id};
use crate::harnesstools::session::{
@@ -24,7 +25,6 @@ use crate::harnesstools::session::{
use crate::io::{FollowLines, Line, PathLineReader, PooledReadLines};
use crate::json::jsonl::JsonlExt;
use crate::sync::BlockingPool;
use crate::utils::{env_nonempty, home_dir};
#[derive(Debug, Clone, TypedBuilder)]
pub struct CodexSessions {
@@ -42,10 +42,7 @@ impl CodexSessions {
/// Codex's live rollouts: `$CODEX_HOME/sessions`, else `~/.codex/sessions`.
pub(crate) fn default_root() -> PathBuf {
env_nonempty("CODEX_HOME")
.map(PathBuf::from)
.unwrap_or_else(|| home_dir().join(".codex"))
.join("sessions")
super::home().join("sessions")
}
/// The rollout of session `id` under `root` (`<yyyy>/<mm>/<dd>/rollout-<timestamp>-<id>.jsonl`).
@@ -307,6 +304,7 @@ impl CodexSession {
usage_recorded: (start == 0).then_some(false),
history_start: (start == 0).then_some(None),
follows: None,
open: Vec::new(),
}
}
}
@@ -389,16 +387,127 @@ struct Stamper {
history_start: Option<Option<u64>>,
/// What the next line [follows](Message::follows), when reading did not start at the start.
follows: Option<MessageId>,
/// The calls made and not yet answered, oldest first.
open: Vec<OpenCall>,
}
/// A tool call made and not yet answered (see [`Stamper::track_calls`]).
struct OpenCall {
id: String,
/// What it touches: what its commands and patches are known by.
accesses: Vec<Access>,
/// What it runs is all known (see [`Access::all`]).
known: bool,
/// A command it ran failed.
failed: bool,
}
impl Stamper {
async fn stamp(&mut self, line: &Line, message: &mut CodexMessage) {
self.stamp_history(line, message).await;
self.track_calls(message);
}
/// What Codex 0.160 tells of a call only in the events of what it ran (its tools run
/// through a script, `exec`, see [`Access::all`]): which call a `FileChange` was made in (the
/// item names no call), and whether a command failed (the script's output says nothing of
/// it). A command is the open call's whose script runs it, else the only open call's, where
/// what that runs isn't all known: a command started in the background ends during a later
/// call that didn't run it (a poll). A file change is the call's whose script writes the
/// file, else the latest's. Calls left open when a turn ends were interrupted. Where reading
/// began after the call was made, it isn't known.
fn track_calls(&mut self, message: &mut CodexMessage) {
let Some(payload) = message.payload.as_ref() else {
return;
};
let call_id = payload["call_id"].as_str();
match (message.kind.as_str(), payload["type"].as_str()) {
("response_item", Some("function_call" | "custom_tool_call")) => {
let input = if payload["type"] == "function_call" {
&payload["arguments"]
} else {
&payload["input"]
};
if let Some(id) = call_id {
let name = payload["name"].as_str().unwrap_or_default();
let (accesses, known) = Access::all(name, input);
self.open.push(OpenCall {
id: id.to_owned(),
accesses,
known,
failed: false,
});
}
}
("response_item", Some("function_call_output" | "custom_tool_call_output")) => {
if let Some(at) =
self.open.iter().position(|call| Some(call.id.as_str()) == call_id)
{
message.context.failed = self.open.remove(at).failed;
}
}
("event_msg", Some("task_complete" | "turn_complete" | "turn_aborted")) => {
self.open.clear();
}
("event_msg", Some("item_completed")) => {
let item = &payload["item"];
match item["type"].as_str() {
Some("CommandExecution") => {
let failed = item["exit_code"].as_i64().is_some_and(|code| code != 0)
|| matches!(item["status"].as_str(), Some("failed" | "declined"));
let ran =
Access::of("shell", &serde_json::json!({"command": item["command"]}))
.command;
let ran_it = |call: &OpenCall| {
call.accesses.iter().any(|a| ran.is_some() && a.command == ran)
};
let at = self.open.iter().rposition(ran_it).or_else(|| {
matches!(self.open.as_slice(), [only] if !only.known).then_some(0)
});
if let Some(call) = at.and_then(|at| self.open.get_mut(at)) {
call.failed |= failed;
}
}
// Before 0.160 the item was the `apply_patch` call itself, under its id.
Some("FileChange")
if !self.open.iter().any(|call| item["id"] == call.id.as_str()) =>
{
let changed = item["changes"].as_object();
let names = |written: &String| {
changed.is_some_and(|changes| {
changes.keys().any(|path| {
path == written
|| path
.strip_suffix(written.as_str())
.is_some_and(|dir| dir.ends_with('/'))
})
})
};
let writes = |call: &&OpenCall| {
call.accesses.iter().any(|a| a.writes.iter().any(names))
};
message.context.call = self
.open
.iter()
.rfind(writes)
.or_else(|| self.open.last())
.map(|call| call.id.clone());
}
_ => {}
}
}
_ => {}
}
}
async fn stamp_history(&mut self, line: &Line, message: &mut CodexMessage) {
// The rollout's first line: reading began at its start, or restarted there after the
// file was truncated or replaced, so nothing was skipped and nothing before it counts.
if line.end == line.bytes.len() as u64 + 1 {
self.before = 0;
self.usage_recorded = Some(false);
self.history_start = Some(None);
self.open.clear();
}
if message.kind == TOKEN_USAGE_RECORD {
self.usage_recorded = Some(true);
@@ -420,6 +529,7 @@ impl Stamper {
session: Some(self.session.clone()),
usage_recorded: self.usage_recorded == Some(true),
follows: self.follows.take(),
..LineContext::default()
};
if let Some(meta) = message.own_meta() {
self.history_start = Some(history_start_of(meta));
@@ -518,6 +628,57 @@ fn records_usage(meta: &serde_json::Value) -> bool {
})
}
/// The exit code of the command Codex ran as item `item` (a `CommandExecution`, which its hooks
/// name the command by), from the end of the rollout at `path`: Codex records the item before
/// it runs the command's `PostToolUse` hooks, which say only its output.
///
/// The item's line holds the command's output three times over, so it isn't parsed: from where
/// it names the item (`"type":"CommandExecution","id":"<item>"`, as Codex writes it, the last
/// time the rollout does), its `"exit_code":` is looked for before the line ends. Only so much
/// of the rollout is read either way.
#[must_use]
pub fn command_exit_code(path: &Path, item: &str) -> Option<i64> {
const SCANNED: u64 = 16 * 1024 * 1024;
const EXIT: &[u8] = b"\"exit_code\":";
let mut file = File::open(path).ok()?;
let len = file.metadata().ok()?.len();
let named = format!("\"type\":\"CommandExecution\",\"id\":{}", serde_json::json!(item));
let at = rfind(&mut file, len, named.as_bytes(), SCANNED).ok()??;
io::Seek::seek(&mut file, io::SeekFrom::Start(at)).ok()?;
let mut line = io::BufRead::split(io::BufReader::new(file.take(SCANNED)), b'\n');
// Only the first line read, the item's: up to where its exit code is.
let line = line.next()?.ok()?;
let from = memchr::memmem::find(&line, EXIT)? + EXIT.len();
let digits: Vec<u8> =
line[from..].iter().copied().take_while(|b| *b == b'-' || b.is_ascii_digit()).collect();
std::str::from_utf8(&digits).ok()?.parse().ok()
}
/// Where `needle` last starts in the first `end` bytes of `file`, looking back no further than
/// `limit` bytes, read in chunks.
fn rfind(file: &mut File, end: u64, needle: &[u8], limit: u64) -> io::Result<Option<u64>> {
const CHUNK: u64 = 64 * 1024;
let finder = memchr::memmem::FinderRev::new(needle);
let floor = end.saturating_sub(limit);
let mut pos = end;
// The start of the chunk read last, which a match may run on into.
let mut carry = Vec::new();
while pos > floor {
let start = pos.saturating_sub(CHUNK).max(floor);
let mut window = vec![0; usize::try_from(pos - start).unwrap_or_default()];
io::Seek::seek(file, io::SeekFrom::Start(start))?;
file.read_exact(&mut window)?;
window.extend_from_slice(&carry);
if let Some(at) = finder.rfind(&window) {
return Ok(Some(start + at as u64));
}
carry = window[..needle.len().saturating_sub(1).min(window.len())].to_vec();
pos = start;
}
Ok(None)
}
/// The own `session_meta` payload of the rollout of session `id` at `path`.
pub(crate) fn own_meta_at(path: &Path, id: &str) -> Option<serde_json::Value> {
first_own_meta(path, &SessionId::from(id.to_owned()))
@@ -623,6 +784,12 @@ struct LineContext {
usage_recorded: bool,
/// What the line [follows](Message::follows), when the reader knows and capture does not.
follows: Option<MessageId>,
/// For a `FileChange` event, the call it was made in, when the reader knows (see
/// [`Stamper::track_calls`]).
call: Option<String>,
/// For a call's output, whether a command the call ran failed (see
/// [`Stamper::track_calls`]).
failed: bool,
}
#[derive(Debug, Clone, Deserialize)]
@@ -1013,9 +1180,11 @@ impl CodexMessage {
Some((call.as_str().filter(|call| !call.is_empty())?, changes))
}
/// The patch [`Self::file_changes`] names.
/// The patch [`Self::file_changes`] names, of the call it was made in: a `FileChange` item
/// Codex 0.160 made in a script names only itself (`exec-...`).
fn patch(&self) -> Option<Patch> {
let (call, changes) = self.file_changes()?;
let (item, changes) = self.file_changes()?;
let call = self.context.call.as_deref().unwrap_or(item);
let files = changes
.iter()
.map(|(path, change)| {
@@ -1159,13 +1328,16 @@ impl Message for CodexMessage {
let Some(payload) = self.payload.as_ref() else {
return Vec::new();
};
// Compacted by the API, the summary is encrypted in the line's `replacement_history`:
// the conversation before it is still no longer what the model sees.
if self.kind == "compacted" {
return payload["message"]
.as_str()
.map(compaction_summary)
.filter(|summary| !summary.is_empty())
.map(|summary| vec![Content::Summary(summary.to_owned())])
.unwrap_or_default();
let summary = payload["message"].as_str().map_or("", compaction_summary);
let summary = if summary.is_empty() {
UNREADABLE_SUMMARY
} else {
summary
};
return vec![Content::Summary(summary.to_owned())];
}
if let Some((_, why)) = self.failure() {
return vec![Content::Error(why)];
@@ -1235,7 +1407,7 @@ impl Message for CodexMessage {
vec![Content::ToolResult(ToolResult {
call: call_id("call_id"),
output: payload["output"].clone(),
error: codex_output_failed(&payload["output"]),
error: self.context.failed || codex_output_failed(&payload["output"]),
})]
}
Some("tool_search_output") => vec![Content::ToolResult(ToolResult {
@@ -2241,16 +2413,25 @@ mod tests {
}
/// A `compacted` line carries the compaction summary in `payload.message` (codex-rs
/// `CompactedItemWire.message`).
/// `CompactedItemWire.message`). Compacted by the API (as Codex 0.160 does on OpenAI), the
/// summary is encrypted in its `replacement_history` and the message empty: still a
/// compaction, its summary unreadable.
#[rstest]
fn compacted_summary_is_kept() {
#[case::readable(
serde_json::json!({"message": "Summary of the conversation so far", "replacement_history": []}),
"Summary of the conversation so far"
)]
#[case::encrypted(
serde_json::json!({"message": "", "replacement_history": [
{"type": "message", "role": "user", "content": [{"type": "input_text", "text": "hi"}]},
{"type": "compaction", "id": "cmp_1", "encrypted_content": "gAAAAB..."}]}),
UNREADABLE_SUMMARY
)]
fn compacted_summary_is_kept(#[case] payload: serde_json::Value, #[case] summary: &str) {
let m = line(&serde_json::json!({
"timestamp": "2026-09-18T10:00:00.000Z", "type": "compacted",
"payload": {"message": "Summary of the conversation so far", "replacement_history": []},
"timestamp": "2026-09-18T10:00:00.000Z", "type": "compacted", "payload": payload,
}));
assert_eq!(m.content(), vec![Content::Summary(
"Summary of the conversation so far".to_owned()
)]);
assert_eq!(m.content(), vec![Content::Summary(summary.to_owned())]);
assert_eq!(m.role(), Role::System);
}
@@ -2484,6 +2665,142 @@ mod tests {
assert_eq!(m.parent_kind(), Some(ParentKind::Continuation));
}
/// Codex 0.160 runs each tool through a script (`exec`), whose output says nothing of how
/// its commands went, and whose patches' `FileChange` items name only themselves: a failed
/// command's call fails, and a patch belongs to the call that made it. From a codex 0.160.1
/// run.
#[rstest]
#[tokio::test]
async fn a_scripts_failures_and_patches_are_its_calls() {
let id = "01a11e4a-7de0-77a3-9030-5535c9877485";
let lines = read_session(id, fixture("code-mode.jsonl")).await;
let content: Vec<Content> = lines.iter().flat_map(Message::content).collect();
let failed: Vec<&str> = content
.iter()
.filter_map(|c| match c {
Content::ToolResult(r) if r.error => Some(r.call.as_ref()),
_ => None,
})
.collect();
assert_eq!(failed, ["call_3243009da42b4e1b916288e04872b89c"]);
// Each keeps its item's id: one script can apply several.
let patches: Vec<(String, String)> = lines
.iter()
.filter_map(|m| match m.content().as_slice() {
[Content::Patch(p)] => Some((m.id()?.to_string(), p.call.to_string())),
_ => None,
})
.collect();
let patch = |item: &str, call: &str| (format!("{item}#patch"), call.to_owned());
assert_eq!(patches, [
patch(
"exec-dbb70cd4-3cec-41f3-acb0-cce249b2aada",
"call_b9854e37b255482a8e0760061b40aa72"
),
patch(
"exec-df2a946a-e49c-453a-a5a1-9af8d26eeae8",
"call_e5f92d9061964574bde820232646737f"
),
]);
}
/// A hook names the command it ran by its item, whose exit code the rollout records.
#[rstest]
#[case::failed("exec-4ba1d5ca-afce-484c-9151-6ffbca09fb0d", Some(2))]
#[case::succeeded("exec-0cc28dd9-b1a1-4b58-8230-d2b26119bd6f", Some(0))]
#[case::a_patch("exec-dbb70cd4-3cec-41f3-acb0-cce249b2aada", None)]
#[case::unknown("exec-nope", None)]
fn a_commands_exit_code_is_in_its_rollout(#[case] item: &str, #[case] exit: Option<i64>) {
assert_eq!(command_exit_code(&fixture("code-mode.jsonl"), item), exit);
}
/// A command whose output is larger than any chunk read (`seq 1 400000; exit 3` made a
/// 2.47 MB line), with lines after it, still says its exit code.
#[rstest]
fn a_large_commands_exit_code_is_found() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("rollout.jsonl");
let output = "1\n".repeat(400_000);
let item = serde_json::json!({"type": "event_msg", "payload": {"type": "item_completed",
"item": {"type": "CommandExecution", "id": "exec-big",
"command": ["/bin/zsh", "-lc", "seq 1 400000; exit 3"], "status": "failed",
"stdout": output, "aggregated_output": output, "exit_code": 3,
"formatted_output": output}}});
let after = serde_json::json!({"type": "event_msg", "payload": {"type": "token_count"}});
std::fs::write(&path, format!("{item}\n{after}\n")).unwrap();
assert_eq!(command_exit_code(&path, "exec-big"), Some(3));
}
/// Of the calls open when a command ran, the one whose script runs it is the one that ran it;
/// read past the call (resumed after it), the patch names only itself.
#[rstest]
#[tokio::test]
async fn an_event_goes_to_the_call_that_ran_it() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("rollout-2026-10-09T00-00-00-th1.jsonl");
let script = |id: &str, script: &str| {
serde_json::json!({"type": "response_item", "payload": {"type": "custom_tool_call",
"call_id": id, "name": "exec", "input": script}})
};
let call = |id: &str, cmd: &str| {
script(id, &format!("await tools.exec_command({{cmd:{}}})", serde_json::json!(cmd)))
};
let computed = "const c = await f(); await tools.exec_command({cmd: c})";
let ran = |cmd: &str, exit: i64| {
serde_json::json!({"type": "event_msg", "payload": {"type": "item_completed",
"item": {"type": "CommandExecution", "id": "exec-1",
"command": ["/bin/zsh", "-lc", cmd], "exit_code": exit,
"status": if exit == 0 { "completed" } else { "failed" }}}})
};
let output = |id: &str| {
serde_json::json!({"type": "response_item", "payload": {
"type": "custom_tool_call_output", "call_id": id,
"output": [{"type": "input_text", "text": "Script completed\nOutput:\n"}]}})
};
let lines = [
call("c1", "false"),
call("c2", "true"),
ran("false", 1),
ran("true", 0),
output("c2"),
output("c1"),
// A command started in the background ends during a poll, which didn't run it.
script("c3", "await tools.write_stdin({session_id: 7, chars: \"\"})"),
ran("sleep 1; false", 1),
output("c3"),
// The only call open ran it, as far as can be told.
script("c4", computed),
ran("make", 2),
output("c4"),
// Interrupted, never answered: no longer open once its turn ends.
script("c5", computed),
serde_json::json!({"type": "event_msg", "payload": {"type": "turn_aborted",
"turn_id": "t1", "reason": "interrupted"}}),
script("c6", computed),
ran("make", 2),
output("c6"),
];
let text: String = lines.iter().map(|l| format!("{l}\n")).collect();
std::fs::write(&path, text).unwrap();
let failed: Vec<(String, bool)> = read_session("th1", path)
.await
.iter()
.flat_map(Message::content)
.filter_map(|c| match c {
Content::ToolResult(r) => Some((r.call.to_string(), r.error)),
_ => None,
})
.collect();
let failed: Vec<(&str, bool)> = failed.iter().map(|(c, e)| (c.as_str(), *e)).collect();
assert_eq!(failed, [
("c2", false),
("c1", true),
("c3", false),
("c4", true),
("c6", true)
]);
}
/// Before 0.32.0 (codex-rs 43809a454e) a rollout had no `type`/`payload` envelope: its first
/// line is the session meta itself, items are written bare, `{"record_type": "state"}` lines
/// sit between them. The fixture follows the writer at codex-rs 2437a8d17a.
crates/atuin-common/src/harnesstools/codex/state_db.rs
+113 -13
View File
@@ -42,19 +42,70 @@ pub enum StateDbError {
/// The current thread index under Codex home `home`: the highest-numbered `state_<n>.sqlite`.
pub fn index_path(home: &Path) -> Option<PathBuf> {
newest(home, "state_")
}
/// The highest-numbered `<prefix><n>.sqlite` under Codex home `home`.
fn newest(home: &Path, prefix: &str) -> Option<PathBuf> {
std::fs::read_dir(home)
.ok()?
.filter_map(Result::ok)
.filter_map(|entry| {
let name = entry.file_name();
let version: u32 =
name.to_str()?.strip_prefix("state_")?.strip_suffix(".sqlite")?.parse().ok()?;
name.to_str()?.strip_prefix(prefix)?.strip_suffix(".sqlite")?.parse().ok()?;
Some((version, entry.path()))
})
.max()
.map(|(_, path)| path)
}
async fn connect(db: &Path) -> Result<SqliteConnection, sqlx::Error> {
let opts = SqliteConnectOptions::new()
.filename(db)
.create_if_missing(false)
.journal_mode(SqliteJournalMode::Wal)
.busy_timeout(Duration::from_secs(5));
SqliteConnection::connect_with(&opts).await
}
async fn has_table(conn: &mut SqliteConnection, table: &str) -> Result<bool, sqlx::Error> {
let found = query_scalar::<Sqlite, i64>(
"SELECT count(*) FROM sqlite_master WHERE type = 'table' AND name = ?1",
)
.bind(table)
.fetch_one(conn)
.await?;
Ok(found > 0)
}
/// Forget what Codex projected of thread `thread`'s history from its rollout
/// (`thread_history_<n>.sqlite`: the thread's turns and items, which it shows the thread from,
/// and the byte offset into the rollout it goes on projecting from), which Codex projects again
/// from the whole rollout when it next loads the thread. A rollout written anew (switched to
/// another branch, or restored) no longer holds the lines it was projected from: Codex would go
/// on showing those, reading on from an offset into lines that are gone.
pub async fn forget_history(home: &Path, thread: &str) -> Result<(), StateDbError> {
let Some(db) = newest(home, "thread_history_") else {
return Ok(());
};
let mut conn = connect(&db).await?;
if !has_table(&mut conn, "thread_history_projection_state").await? {
return Ok(());
}
let mut tx = conn.begin().await?;
for delete in [
"DELETE FROM thread_history_projection_state WHERE thread_id = ?1",
"DELETE FROM thread_items WHERE thread_id = ?1",
"DELETE FROM thread_turns WHERE thread_id = ?1",
] {
query::<Sqlite>(delete).bind(thread).execute(&mut *tx).await?;
}
tx.commit().await?;
conn.close().await?;
Ok(())
}
/// Make the index row of thread `thread` (if there is one) name `rollout`, where it names
/// `replacing` (the rollout `rollout` continues) or a file that is gone. Refuses
/// ([`StateDbError::Elsewhere`]) where it names another rollout that exists: Codex has moved the
@@ -68,18 +119,8 @@ pub async fn point_at(
let Some(db) = index_path(home) else {
return Ok(Pointed::Unindexed);
};
let opts = SqliteConnectOptions::new()
.filename(&db)
.create_if_missing(false)
.journal_mode(SqliteJournalMode::Wal)
.busy_timeout(Duration::from_secs(5));
let mut conn = SqliteConnection::connect_with(&opts).await?;
let has_threads = query_scalar::<Sqlite, i64>(
"SELECT count(*) FROM sqlite_master WHERE type = 'table' AND name = 'threads'",
)
.fetch_one(&mut conn)
.await?;
if has_threads == 0 {
let mut conn = connect(&db).await?;
if !has_table(&mut conn, "threads").await? {
return Ok(Pointed::Unindexed);
}
let stored = query_scalar::<Sqlite, String>("SELECT rollout_path FROM threads WHERE id = ?1")
@@ -147,6 +188,65 @@ pub mod tests {
db
}
/// Codex's history of threads as it projects it (codex 0.160.1's `thread_history_1.sqlite`,
/// its columns cut down), holding some of each of `threads`.
pub async fn history(home: &Path, threads: &[&str]) -> PathBuf {
let db = home.join("thread_history_1.sqlite");
let opts = SqliteConnectOptions::new().filename(&db).create_if_missing(true);
let mut conn = SqliteConnection::connect_with(&opts).await.unwrap();
for sql in [
"CREATE TABLE thread_turns (thread_id TEXT NOT NULL, turn_id TEXT NOT NULL, \
rollout_ordinal INTEGER NOT NULL, PRIMARY KEY (thread_id, turn_id))",
"CREATE TABLE thread_items (thread_id TEXT NOT NULL, turn_id TEXT NOT NULL, item_id \
TEXT NOT NULL, rollout_ordinal INTEGER NOT NULL, PRIMARY KEY (thread_id, turn_id, \
item_id))",
"CREATE TABLE thread_history_projection_state (thread_id TEXT PRIMARY KEY, \
next_rollout_byte_offset INTEGER NOT NULL, next_rollout_ordinal INTEGER NOT NULL)",
] {
query::<Sqlite>(sql).execute(&mut conn).await.unwrap();
}
for thread in threads {
for sql in [
"INSERT INTO thread_turns VALUES (?1, 't1', 1)",
"INSERT INTO thread_items VALUES (?1, 't1', 'i1', 2)",
"INSERT INTO thread_history_projection_state VALUES (?1, 83674, 36)",
] {
query::<Sqlite>(sql).bind(*thread).execute(&mut conn).await.unwrap();
}
}
db
}
/// How many rows of its history Codex holds of `thread`.
pub async fn projected(db: &Path, thread: &str) -> i64 {
let mut conn = SqliteConnection::connect_with(&SqliteConnectOptions::new().filename(db))
.await
.unwrap();
query_scalar::<Sqlite, i64>(
"SELECT (SELECT count(*) FROM thread_turns WHERE thread_id = ?1) + (SELECT count(*) \
FROM thread_items WHERE thread_id = ?1) + (SELECT count(*) FROM \
thread_history_projection_state WHERE thread_id = ?1)",
)
.bind(thread)
.fetch_one(&mut conn)
.await
.unwrap()
}
/// The thread's history goes, to be projected again from its rollout; another thread's
/// stays. With no history at all there is nothing to forget.
#[rstest]
#[tokio::test]
async fn a_threads_history_is_forgotten() {
let home = tempfile::tempdir().unwrap();
forget_history(home.path(), THREAD).await.unwrap();
let other = "01a0ea6e-0000-7000-8000-000000000000";
let db = history(home.path(), &[THREAD, other]).await;
forget_history(home.path(), THREAD).await.unwrap();
assert_eq!(projected(&db, THREAD).await, 0);
assert_eq!(projected(&db, other).await, 3);
}
async fn stored(db: &Path) -> String {
let opts = SqliteConnectOptions::new().filename(db);
let mut conn = SqliteConnection::connect_with(&opts).await.unwrap();
crates/atuin-common/src/harnesstools/continuation.rs
+113 -11
View File
@@ -21,7 +21,10 @@
//! call captured without its input (`ai.capture_tools` off, or a policy withheld it), one with no
//! result (interrupted, or a Codex web search), and one the model's provider ran itself (a
//! Claude Code web search), whose result only that provider reads.
//! - Summaries (compactions, abandoned branches), as text of the user's turn.
//! - Summaries (compactions, abandoned branches), as text of the user's turn. A compaction
//! whose summary couldn't be read is where the conversation carried over starts, and so is
//! the prompt the last ~100k tokens of a longer conversation start at: with a note, and the
//! prompts before it, as Codex goes on after a compaction.
//! - Failed model calls, as a short note of the error in the assistant's turn.
//!
//! Dropped: reasoning (only the model that wrote it can read it back), text the source harness
@@ -60,6 +63,7 @@ use super::resume::is_plain_name;
use super::{AnyHarness, Harness as _, opencode};
use crate::harnesstools::session::{
Content, Message, ParentKind, Role, StopReason, ToolCallId, ToolResult, ToolUse,
UNREADABLE_SUMMARY,
};
mod tools;
@@ -227,15 +231,10 @@ pub fn linked_from_message<M: Message + ?Sized>(m: &M) -> Option<(AnyHarness, St
let content = m.content();
let marker = match (m.role(), content.as_slice()) {
(Role::Assistant | Role::Tool, _) => return None,
(Role::User, [Content::Text(first), _, ..]) => {
let (kind, harness, id, atuin_id) = marker_parts(first)?;
let exact = match kind {
ParentKind::Fork => fork_marker_text(harness, id, atuin_id),
_ => marker_text(harness, id, atuin_id),
};
return (*first == exact).then(|| (harness, id.to_owned(), kind));
(Role::User, content) => {
return prompt_marker(content)
.map(|(kind, harness, id)| (harness, id.to_owned(), kind));
}
(Role::User, _) => return None,
(_, content) => content.iter().find_map(|c| match c {
Content::Text(text) if marker_parts(text).is_some() => Some(text),
_ => None,
@@ -244,6 +243,28 @@ pub fn linked_from_message<M: Message + ?Sized>(m: &M) -> Option<(AnyHarness, St
marker_parts(marker).map(|(kind, harness, id, _)| (harness, id.to_owned(), kind))
}
/// The link a prompt's first block makes when it is pi's [marker](marker_text): the whole marker,
/// exactly as atuin writes it, with more after it.
fn prompt_marker(content: &[Content]) -> Option<(ParentKind, AnyHarness, &str)> {
let [Content::Text(first), _, ..] = content else {
return None;
};
let (kind, harness, id, atuin_id) = marker_parts(first)?;
let exact = match kind {
ParentKind::Fork => fork_marker_text(harness, id, atuin_id),
_ => marker_text(harness, id, atuin_id),
};
(*first == exact).then_some((kind, harness, id))
}
/// Whether a prompt of `content` opens with pi's [marker](marker_text): its first block, which
/// is no part of what the user said, so neither what a session is shown by nor carried on when
/// it is continued again.
#[must_use]
pub fn opens_with_marker(content: &[Content]) -> bool {
prompt_marker(content).is_some()
}
/// That `messages` hold something of the conversation, so a continuation or fork of them isn't
/// empty.
pub fn conversational(messages: &[RehydrateMessage]) -> Result<(), NothingToContinue> {
@@ -547,9 +568,83 @@ impl<'m> Turn<'m> {
}
}
/// How much of a conversation's rows (their content, as JSON) a continuation carries, about
/// 100k tokens: more is past what a model takes, with room to go on.
const CARRIED: usize = 400 * 1024;
/// How much of the prompts before where a continuation starts it carries, as Codex does after a
/// compaction (codex-rs `compact.rs`, 20k tokens of them).
const PROMPTS: usize = 80 * 1024;
/// Where a continuation of `messages` starts, and why there, when not at the start: after the
/// last compaction whose summary couldn't be read ([`UNREADABLE_SUMMARY`]), the model saw only
/// what came after it; and a conversation longer than [`CARRIED`] (one captured before such
/// compactions were) starts at the prompt it fits from.
fn start(messages: &[RehydrateMessage]) -> (usize, Option<&'static str>) {
let compacted = messages.iter().rposition(|m| {
m.content.iter().any(|c| matches!(c, Content::Summary(s) if s == UNREADABLE_SUMMARY))
});
let (from, why) = match compacted {
Some(at) => (at, Some(COMPACTED)),
None => (0, None),
};
let mut size = 0;
let mut fits = None;
for (at, m) in messages.iter().enumerate().skip(from).rev() {
size += serde_json::to_string(&m.content).map_or(0, |json| json.len());
// The last prompt, at least, whatever its turn holds.
if m.role == Role::User && said(m) && (size <= CARRIED || fits.is_none()) {
fits = Some(at);
}
if size > CARRIED && fits.is_some() {
break;
}
}
match fits {
Some(at) if size > CARRIED && at > from => (at, Some(TOO_LONG)),
_ => (from, why),
}
}
/// Whether `m` says something in text.
fn said(m: &RehydrateMessage) -> bool {
m.content.iter().any(|c| matches!(c, Content::Text(t) if !t.trim().is_empty()))
}
const COMPACTED: &str = "[the conversation before this was compacted, and its summary can't be \
read: only the prompts of it are carried over]";
const TOO_LONG: &str =
"[the conversation before this is too long to carry over: only the prompts of it are]";
/// The conversation of `messages` as alternating turns, user first and assistant last, and
/// what was left out of it.
/// what was left out of it: from where [`start`] says, after a note saying why and the prompts
/// before it (the latest [`PROMPTS`] of them).
fn turns(messages: &[RehydrateMessage]) -> (Vec<Turn<'_>>, Flattened) {
let (from, why) = start(messages);
let mut opening = Vec::new();
if let Some(why) = why {
let mut budget = PROMPTS;
let mut prompts: Vec<Part> = messages[..from]
.iter()
.rev()
.filter(|m| m.role == Role::User)
// without an earlier continuation's marker, as in the turns after
.flat_map(|m| m.content[usize::from(opens_with_marker(&m.content))..].iter().rev())
.filter_map(|c| match c {
Content::Text(text) if !text.trim().is_empty() => Some(text),
_ => None,
})
.take_while(|text| {
budget = budget.saturating_sub(text.len());
budget > 0
})
.map(|text| Part::Text(text.clone()))
.collect();
prompts.reverse();
opening.push(Part::Note(why.to_owned()));
opening.extend(prompts);
}
let messages = &messages[from..];
// A call is carried over with its result. One of a tool the model's provider ran itself,
// answered in the assistant's own row (opencode keeps every result there), is only for that
// provider to read: its results are opaque, encrypted.
@@ -571,6 +666,9 @@ fn turns(messages: &[RehydrateMessage]) -> (Vec<Turn<'_>>, Flattened) {
let mut flattened = Flattened::default();
let mut carried: HashSet<&ToolCallId> = HashSet::new();
let mut turns: Vec<Turn> = Vec::new();
if let Some(first) = messages.first().filter(|_| !opening.is_empty()) {
turns.push(Turn::parts(Kind::User, first.timestamp, opening));
}
// Results came back since the model's last row, or its row is of another model call (its
// turn id up to a step's usage key: Claude Code's message id, opencode's call): its next is
// another call of the model. (opencode answers calls in the assistant's own rows.)
@@ -580,9 +678,13 @@ fn turns(messages: &[RehydrateMessage]) -> (Vec<Turn<'_>>, Flattened) {
let mut user = Vec::new();
let mut assistant = Vec::new();
let mut reasoned = false;
for content in &m.content {
// A continuation's own marker: the new one names this session instead.
let marker = usize::from(m.role == Role::User && opens_with_marker(&m.content));
for content in &m.content[marker..] {
match content {
Content::Reasoning(_) | Content::ReasoningSummary { .. } => reasoned = true,
// Said by the note the continuation opens with.
Content::Summary(summary) if summary == UNREADABLE_SUMMARY => {}
Content::Summary(summary) if !summary.trim().is_empty() => user.push(Part::Text(
format!("[Summary of the earlier conversation]\n{}", summary.trim()),
)),
crates/atuin-common/src/harnesstools/continuation/tests.rs
+108 -1
View File
@@ -432,7 +432,7 @@ fn conversation(rows: &[RehydrateMessage]) -> Vec<(Role, String)> {
out.push((Role::Assistant, format!("{} {}({input})", call.id, call.name)));
}
Content::ToolResult(result) => {
let output = tools::output_text(&result.output);
let output = tools::output_text(result);
out.push((Role::Tool, format!("{} -> {output}", result.call)));
}
other => panic!("only text and calls are written: {other:?}"),
@@ -586,6 +586,81 @@ fn calls_become_notes_and_turns_merge() {
assert!(!everything.contains("environment_context"), "nor the source's own context");
}
/// The continuation's turns past its marker, each one text.
fn texts(c: &Continuation) -> Vec<(Role, String)> {
c.session.messages[1..]
.iter()
.map(|m| match m.content.as_slice() {
[Content::Text(t)] => (m.role.clone(), t.clone()),
other => panic!("{other:?}"),
})
.collect()
}
/// A compaction whose summary couldn't be read (Codex's, compacted by the API) left the model
/// only what came after it: the continuation starts there, saying so, with the prompts before
/// it, as Codex goes on. An empty summary of another kind (opencode's, of a compaction that
/// stopped) is no such cut.
#[rstest]
fn an_unreadable_compaction_starts_the_conversation() {
let original = session_of("orig", vec![
msg(Role::User, vec![text("first, a long task")], 1),
msg(Role::Assistant, vec![call("exec_command", Value::Null)], 2),
msg(Role::System, vec![Content::Summary(UNREADABLE_SUMMARY.to_owned())], 3),
msg(Role::User, vec![text("then this")], 4),
msg(Role::System, vec![Content::Summary(UNREADABLE_SUMMARY.to_owned())], 5),
msg(Role::User, vec![text("and now this")], 6),
msg(Role::System, vec![Content::Summary(String::new())], 7),
msg(Role::Assistant, vec![text("Done.")], 8),
]);
let c = continue_at(CODEX, &original, None, CLAUDE, OffsetDateTime::UNIX_EPOCH).unwrap();
assert_eq!(c.flattened, Flattened::default());
pretty_assertions::assert_eq!(texts(&c), vec![
(Role::User, format!("{COMPACTED}\n\nfirst, a long task\n\nthen this\n\nand now this")),
(Role::Assistant, "Done.".to_owned()),
]);
}
/// The prompts carried over a cut leave out an earlier pi continuation's marker too.
#[rstest]
fn the_prompts_before_a_cut_leave_out_a_marker() {
let old = marker_text(CLAUDE, "4bd6b5bd-7d34-4a39-a2d6-2e6f6a6c0b1d", Some(ATUIN_ID));
let original = session_of("orig", vec![
msg(Role::User, vec![text(&old), text("first, a long task")], 1),
msg(Role::System, vec![Content::Summary(UNREADABLE_SUMMARY.to_owned())], 2),
msg(Role::User, vec![text("and now this")], 3),
msg(Role::Assistant, vec![text("Done.")], 4),
]);
let c = continue_at(PI, &original, None, CLAUDE, OffsetDateTime::UNIX_EPOCH).unwrap();
pretty_assertions::assert_eq!(texts(&c), vec![
(Role::User, format!("{COMPACTED}\n\nfirst, a long task\n\nand now this")),
(Role::Assistant, "Done.".to_owned()),
]);
}
/// A conversation longer than a continuation carries (one captured before Codex's compactions
/// were kept) starts at the prompt its last part does, saying so, with the latest prompts
/// before it.
#[rstest]
fn a_long_conversation_starts_where_it_fits() {
let big = "x".repeat(CARRIED / 3);
let mut rows = Vec::new();
for n in 0..6 {
rows.push(msg(Role::User, vec![text(&format!("prompt {n}"))], 2 * n));
rows.push(msg(Role::Assistant, vec![text(&big)], 2 * n + 1));
}
let original = session_of("orig", rows);
let c = continue_at(CODEX, &original, None, CLAUDE, OffsetDateTime::UNIX_EPOCH).unwrap();
let turns = texts(&c);
// The last two turns fit; the third from last would not.
assert_eq!(turns.len(), 4, "{turns:.200?}");
assert_eq!(
turns[0].1,
format!("{TOO_LONG}\n\nprompt 0\n\nprompt 1\n\nprompt 2\n\nprompt 3\n\nprompt 4")
);
assert_eq!(turns[2].1, "prompt 5");
}
/// Calls captured without their input (all of them, now) say only what they did, and the same
/// note in a row is written once, counted.
#[rstest]
@@ -676,6 +751,38 @@ fn a_pasted_marker_is_no_link(#[case] paste: fn(&str) -> String) {
assert!(continued_from_row(&written).is_some());
}
/// Continuing a pi continuation again carries its conversation but not its marker, which named
/// the session it continued: the new session names only the one it continues. A prompt that
/// merely quotes a marker keeps it.
#[rstest]
fn a_continuations_marker_is_not_carried_on(
#[values(CLAUDE, CODEX, OPENCODE, PI)] target: AnyHarness,
) {
let old = marker_text(CLAUDE, "4bd6b5bd-7d34-4a39-a2d6-2e6f6a6c0b1d", Some(ATUIN_ID));
let original = session_of("orig", vec![
msg(Role::User, vec![text(&old), text("fix it")], 1),
msg(Role::Assistant, vec![text("fixed")], 2),
msg(Role::User, vec![text(&old)], 3),
msg(Role::Assistant, vec![text("that's a marker")], 4),
]);
let c = continue_in(PI, &original, None, target).unwrap();
let said: Vec<&str> = c
.session
.messages
.iter()
.flat_map(|m| &m.content)
.filter_map(|c| match c {
Content::Text(t) => Some(t.as_str()),
_ => None,
})
.collect();
assert_eq!(said[0], marker_text(PI, "orig", None), "{said:#?}");
assert!(said.iter().any(|t| t.starts_with("fix it")), "{said:#?}");
// Only the quoted one is left.
let old = said.iter().filter(|t| t.contains("4bd6b5bd")).count();
assert_eq!(old, 1, "{said:#?}");
}
/// A session with nothing of the conversation in it can't be continued, in any harness: there
/// would be nothing to carry over, and in pi no prompt for the marker to ride on.
#[rstest]
crates/atuin-common/src/harnesstools/continuation/tools.rs
+12 -18
View File
@@ -4,6 +4,7 @@
//! target's API takes. A model reads a call to a tool it doesn't have, and its result, like any
//! other: it just can't make that call again.
use std::borrow::Cow;
use std::collections::HashSet;
use std::path::Path;
@@ -314,22 +315,10 @@ fn input_for(target: AnyHarness, input: Value) -> Value {
}
}
/// A result's output as text, which every target writes as a tool's output: the text of its
/// blocks (Claude Code's and pi's `text`, Codex's `input_text`), else its JSON.
pub(super) fn output_text(output: &Value) -> String {
match output {
Value::Null => UNCAPTURED_OUTPUT.to_owned(),
Value::String(text) => text.clone(),
Value::Array(blocks) => blocks
.iter()
.map(|block| match block {
Value::String(text) => text.clone(),
_ => block["text"].as_str().map_or_else(|| block.to_string(), str::to_owned),
})
.collect::<Vec<_>>()
.join("\n"),
other => other.to_string(),
}
/// A result's output as text, which every target writes as a tool's output (see
/// [`ToolResult::output_text`]).
pub(super) fn output_text(result: &ToolResult) -> String {
result.output_text().map_or_else(|| UNCAPTURED_OUTPUT.to_owned(), Cow::into_owned)
}
/// Call ids for the target, unique within the session.
@@ -435,7 +424,7 @@ pub(super) fn carry(
};
let result = ToolResult {
call: id,
output: Value::String(output_text(&result.output)),
output: Value::String(output_text(result)),
error: result.error,
};
(call, result)
@@ -505,7 +494,12 @@ mod tests {
#[case::uncaptured(Value::Null, UNCAPTURED_OUTPUT)]
#[case::structured(json!({"exit": 0}), r#"{"exit":0}"#)]
fn outputs_become_text(#[case] output: Value, #[case] want: &str) {
assert_eq!(output_text(&output), want);
let result = ToolResult {
call: ToolCallId::from("c1".to_owned()),
output,
error: false,
};
assert_eq!(output_text(&result), want);
}
#[rstest]
crates/atuin-common/src/harnesstools/json_hooks.rs
+45 -7
View File
@@ -55,7 +55,9 @@ pub async fn install(
/// changed.
///
/// Atuin hooks are removed wherever they sit, so reinstalling after an upgrade or a moved binary
/// never leaves a stale or duplicate hook behind.
/// never leaves a stale or duplicate hook behind. An entry already as it would be added stays
/// where it is: Codex trusts each hook by where it is, so moving it would move the user's after
/// it and leave them untrusted.
fn add_hook_entries(
hooks: &mut Map<String, Value>,
matcher: &str,
@@ -70,7 +72,15 @@ fn add_hook_entries(
.as_array_mut()
.ok_or(InstallHookError::Malformed("a hook event is not an array"))?;
for entry in entries.iter_mut() {
let installed = json!({
"matcher": matcher,
"hooks": [{"type": "command", "command": hook_command}],
});
let kept = entries.iter().position(|entry| *entry == installed);
for (at, entry) in entries.iter_mut().enumerate() {
if Some(at) == kept {
continue;
}
if let Some(entry_hooks) = entry.get_mut("hooks").and_then(Value::as_array_mut) {
entry_hooks.retain(|hook| {
!hook
@@ -84,10 +94,9 @@ fn add_hook_entries(
entry.get("hooks").and_then(Value::as_array).is_none_or(|hooks| !hooks.is_empty())
});
entries.push(json!({
"matcher": matcher,
"hooks": [{"type": "command", "command": hook_command}],
}));
if kept.is_none() {
entries.push(installed);
}
}
Ok(*hooks != before)
@@ -116,7 +125,7 @@ fn hook_command(executable: &Path, harness: &str) -> Result<String, InstallHookE
}
/// Whether `command` runs `atuin hook <harness>`, through any path to the atuin executable.
fn invokes_atuin_hook(command: &str, harness: &str) -> bool {
pub fn invokes_atuin_hook(command: &str, harness: &str) -> bool {
let Some(parts) = shlex::split(command) else {
return false;
};
@@ -178,6 +187,35 @@ mod tests {
);
}
/// Installed again, atuin's entry stays where it is, ahead of the user's added since.
#[rstest]
fn add_hook_entries_keeps_an_installed_entry_in_place() {
let command = "/opt/atuin/bin/atuin hook codex";
let installed = json!({
"matcher": "^Bash$",
"hooks": [{"type": "command", "command": command}],
});
let theirs = json!({"matcher": "", "hooks": [{"type": "command", "command": "echo hi"}]});
let stale = json!({"matcher": "^Bash$", "hooks": [{"type": "command", "command": "atuin hook codex"}]});
let mut hooks = json!({
"PreToolUse": [installed, theirs],
"PostToolUse": [stale, installed, theirs],
"PostToolUseFailure": [installed],
});
assert!(
add_hook_entries(hooks.as_object_mut().unwrap(), "^Bash$", "codex", command).unwrap()
);
assert_eq!(
hooks,
json!({
"PreToolUse": [installed, theirs],
"PostToolUse": [installed, theirs],
"PostToolUseFailure": [installed],
})
);
}
#[cfg(not(windows))]
#[rstest]
#[case::plain("/opt/atuin/bin/atuin", "/opt/atuin/bin/atuin hook codex")]
crates/atuin-common/src/harnesstools/opencode/mod.rs
+140 -4
View File
@@ -1,4 +1,4 @@
use std::path::PathBuf;
use std::path::{Path, PathBuf};
use super::rehydrate::{RehydrateError, RehydrateSession};
use super::resume::{self, CwdRequirement, ResumeError, ResumePlan, ResumeTarget};
@@ -18,6 +18,20 @@ impl Opencode {
}
}
/// The plugin for opencode 1.x, whose loaders (up to 1.3) call every export of a plugin file.
const PLUGIN_1: &str = include_str!("../../../res/harnesstools/opencode/atuin.ts");
/// The plugin for opencode 2, which loads a default export's `id` and `setup`.
const PLUGIN_2: &str = include_str!("../../../res/harnesstools/opencode/atuin.v2.ts");
/// The plugin for the opencode of major version `major`: no export shape loads in both. One
/// whose version can't be told (not installed yet) gets 1.x's, which npm still ships.
const fn plugin_for(major: Option<u64>) -> &'static str {
match major {
Some(major) if major >= 2 => PLUGIN_2,
_ => PLUGIN_1,
}
}
impl Harness for Opencode {
fn name(&self) -> &'static str {
"opencode"
@@ -25,8 +39,8 @@ impl Harness for Opencode {
/// Install the hooks on the current active user's machine
async fn install_hooks(&self) -> Result<PathBuf, InstallHookError> {
const PLUGIN_SOURCE: &str = include_str!("../../../res/harnesstools/opencode/atuin.ts");
const PLUGIN_NAME: &str = "atuin.ts";
let plugin = plugin_for(rehydrate::major_version(Path::new("opencode")).await.ok());
// Opencode is XDG-compliant: plugins live under $XDG_CONFIG_HOME/opencode,
// falling back to ~/.config/opencode.
@@ -42,13 +56,13 @@ impl Harness for Opencode {
let already_installed = tokio::fs::read_to_string(&extension_path)
.await
.is_ok_and(|existing| existing == PLUGIN_SOURCE);
.is_ok_and(|existing| existing == plugin);
if already_installed {
return Err(InstallHookError::AlreadyInstalled);
}
tokio::fs::write(&extension_path, PLUGIN_SOURCE).await?;
tokio::fs::write(&extension_path, plugin).await?;
Ok(extension_path)
}
@@ -70,3 +84,125 @@ impl Harness for Opencode {
rehydrate::rehydrate(session).await
}
}
#[cfg(test)]
mod tests {
use rstest::rstest;
use super::*;
/// Each opencode gets the plugin its loader takes: 1.x (and one whose version can't be told)
/// a named plugin function, the only export; 2 a default export with an `id` and `setup`.
#[rstest]
#[case::none(None, PLUGIN_1)]
#[case::one(Some(1), PLUGIN_1)]
#[case::two(Some(2), PLUGIN_2)]
#[case::three(Some(3), PLUGIN_2)]
fn each_opencode_gets_the_plugin_it_loads(#[case] major: Option<u64>, #[case] want: &str) {
assert_eq!(plugin_for(major), want);
let exports = |plugin: &str| -> Vec<String> {
plugin.lines().filter(|l| l.starts_with("export ")).map(str::to_owned).collect()
};
assert_eq!(exports(PLUGIN_1), [
"export const AtuinPlugin: Plugin = async ({ directory }) => {"
]);
assert_eq!(exports(PLUGIN_2), ["export default { id: \"atuin\", setup };"]);
}
/// Whether `node` is here and runs TypeScript as it is.
fn node_runs_typescript() -> bool {
std::process::Command::new("node")
.args(["-e", "process.exit(process.features.typescript ? 0 : 1)"])
.status()
.is_ok_and(|status| status.success())
}
/// A fake opencode 2 driving the plugin through its hooks as opencode calls them for a shell
/// call (`execute.before`, `create.before`, `execute.after`): one call per way a call ends.
const DRIVER: &str = r#"
import plugin from "./atuin.ts";
const hooks = {};
const add = (kind) => (name, fn) => { (hooks[`${kind}.${name}`] ??= []).push(fn); };
plugin.setup({ tool: { hook: add("tool") }, shell: { hook: add("shell") } });
const fire = async (name, payload) => { for (const fn of hooks[name] ?? []) await fn(payload); };
async function call(id, command, after) {
await fire("tool.execute.before", { tool: "shell", id, input: { command, description: `about ${id}` } });
const shell = { command, cwd: "." };
await fire("shell.create.before", shell);
shell.cwd = process.cwd();
await fire("tool.execute.after", { tool: "shell", id, input: { command }, ...after });
}
const done = (metadata) => ({ status: "completed", result: { content: [], metadata } });
await call("c1", "true", done({ status: "completed", truncated: false, exit: 0 }));
await call("c2", "false", done({ status: "completed", truncated: false, exit: 1 }));
await call("c3", "sleep 999", done({ status: "completed", truncated: false, timeout: true, signal: "SIGTERM" }));
await call("c4", "npm run dev", done({ status: "running", truncated: false, shellID: "sh_1" }));
await call("c5", "rm -rf x", { status: "error", error: new Error("denied") });
await fire("tool.execute.after", { tool: "read", id: "c6", status: "completed", result: {} });
"#;
/// The opencode 2 plugin records each shell command that ended, with its exit code (a
/// timeout's 124); not one left running in the background, nor one that was denied.
#[cfg(unix)]
#[rstest]
fn the_opencode_2_plugin_records_commands_that_ended() {
use std::os::unix::fs::PermissionsExt;
if !node_runs_typescript() {
eprintln!("skipped: no node that runs TypeScript");
return;
}
// next to the test binary: CI's temporary directory is mounted without exec
let exe = std::env::current_exe().unwrap();
let dir = tempfile::tempdir_in(exe.parent().unwrap()).unwrap();
let log = dir.path().join("atuin.log");
let stub = dir.path().join("atuin");
std::fs::write(
&stub,
format!(
"#!/bin/sh\necho \"$@\" >> {log}\n[ \"$2\" = start ] && echo h_1\nexit 0\n",
log = log.display()
),
)
.unwrap();
std::fs::set_permissions(&stub, std::fs::Permissions::from_mode(0o755)).unwrap();
std::fs::write(dir.path().join("atuin.ts"), PLUGIN_2).unwrap();
std::fs::write(dir.path().join("driver.mts"), DRIVER).unwrap();
let path = std::env::var("PATH").unwrap_or_default();
let ran = std::process::Command::new("node")
.arg("driver.mts")
.current_dir(dir.path())
.env("PATH", format!("{}:{path}", dir.path().display()))
.output()
.unwrap();
assert!(ran.status.success(), "{}", String::from_utf8_lossy(&ran.stderr));
let calls: Vec<String> = std::fs::read_to_string(&log)
.unwrap()
.lines()
.map(|line| match line.split_once(" --duration ") {
// how long it took is not the test's business
Some((before, nanos)) => {
assert!(nanos.parse::<u64>().is_ok(), "{line}");
format!("{before} --duration _")
}
None => line.to_owned(),
})
.collect();
let start = |id: &str, command: &str| {
format!(
"history start --author opencode --author-kind agent --intent about {id} -- \
{command}"
)
};
let end = |exit: i32| format!("history end h_1 --exit {exit} --duration _");
assert_eq!(calls, [
start("c1", "true"),
end(0),
start("c2", "false"),
end(1),
start("c3", "sleep 999"),
end(124),
]);
}
}
crates/atuin-common/src/harnesstools/opencode/rehydrate.rs
+80 -22
View File
@@ -1,6 +1,10 @@
//! An opencode session written back out from captured messages (see [`rehydrate`]), through
//! opencode's own `opencode import`.
//!
//! opencode 2 dropped `opencode import` for `opencode session import`, which takes a format of its
//! own (see `v2`); which one the installed opencode takes is told by its version. What follows
//! is opencode 1.x's.
//!
//! opencode keeps its sessions in one SQLite database, whose schema is its own business; `opencode
//! import <file>` takes the JSON `opencode export` writes and inserts it the way opencode itself
//! does (`cli/cmd/import.ts`): the session under its own id, each message and part under theirs.
@@ -58,6 +62,7 @@
//! as new rows. Their calls and results, rows apart, are joined into one tool part.
use std::collections::{HashMap, HashSet};
use std::ffi::OsStr;
use std::path::{Path, PathBuf};
use std::process::Stdio;
use std::time::Duration;
@@ -78,6 +83,8 @@ use crate::harnesstools::session::{
};
use crate::harnesstools::{AnyHarness, continuation};
mod v2;
/// How long `opencode import` may take before it is given up on.
const IMPORT_TIMEOUT: Duration = Duration::from_secs(120);
@@ -93,22 +100,25 @@ const QUIET_ENV: &[(&str, &str)] = &[
("OPENCODE_PURE", "1"),
];
/// Import `session` into opencode's database with `opencode import`, run from the session's
/// directory ([`RehydrateSession::cwd`], which becomes the session's directory and project), and
/// hand back the database: the path `locate` finds the session in.
/// Import `session` into opencode's database with `opencode import` (opencode 2: `opencode
/// session import`), into the session's directory ([`RehydrateSession::cwd`], which becomes the
/// session's directory and project), and hand back the database: the path `locate` finds the
/// session in.
///
/// The database is the one opencode resolves from this process's environment (`OPENCODE_DB`,
/// `XDG_DATA_HOME`, its channel databases; see `session::default_db`), and `opencode import` is
/// pointed at it explicitly (`OPENCODE_DB`), so the two cannot disagree. opencode must be on
/// `PATH`. It is never asked to call a model: import only writes the database.
/// `XDG_DATA_HOME`, its channel databases; see `session::default_db`), and opencode is pointed at
/// it explicitly (`OPENCODE_DB`), so the two cannot disagree: opencode 2 imports through a server
/// of its own (`--standalone`), not its background one, which keeps the environment it started
/// with. opencode must be on `PATH`. It is never asked to call a model: import only writes the
/// database.
///
/// Fails with [`RehydrateError::AlreadyExists`] when opencode already holds the session whole.
/// The export goes to opencode in one file, but opencode inserts it a row at a time (the session
/// first), so an import that fails part-way leaves the session partly written, where
/// [`locate`](crate::harnesstools::Harness::locate) finds it. Such a session is imported again:
/// import inserts each message and part only if its id is new, so what is there stays as it is
/// and what is missing is added. A session that has been worked on since (see `Held`) is left
/// alone, so that a message a revert deleted never comes back.
/// opencode 1.x inserts an import a row at a time (the session first), so one that fails part-way
/// leaves the session partly written, where [`locate`](crate::harnesstools::Harness::locate)
/// finds it. Such a session is imported again: import inserts each message and part only if its
/// id is new, so what is there stays as it is and what is missing is added. A session that has
/// been worked on since (see `Held`) is left alone, so that a message a revert deleted never comes
/// back. opencode 2 imports a session whole or not at all.
pub async fn rehydrate(session: &RehydrateSession) -> Result<PathBuf, RehydrateError> {
let db = session::default_db().ok_or(RehydrateError::NoDataDir)?;
import(Path::new("opencode"), &db, session).await
@@ -127,6 +137,7 @@ pub(crate) async fn import(
)));
}
let exported = export(session);
// a session opencode 2 holds is in none of the tables `held` reads: held whole
if session::locate(db, &session.id).await.is_some()
&& held(db, &session.id, &exported).await != Held::Partly
{
@@ -138,12 +149,46 @@ pub(crate) async fn import(
session.cwd.display()
)));
}
let stdout = run_import(program, db, &session.cwd, &exported).await?;
let stdout = if major_version(program).await? >= 2 {
let args = [
OsStr::new("session"),
OsStr::new("import"),
OsStr::new("--standalone"),
OsStr::new("--directory"),
session.cwd.as_os_str(),
];
run(program, &args, db, &session.cwd, &v2::export(session)).await?
} else {
run_import(program, db, &session.cwd, &exported).await?
};
session::locate(db, &session.id).await.ok_or_else(|| {
RehydrateError::Other(format!("opencode import did not write the session: {stdout}"))
})
}
/// The major version of the `opencode` executable `program`, as `opencode --version` prints it
/// (`1.18.32`; opencode 2: `opencode v2.0.24`).
pub(crate) async fn major_version(program: &Path) -> Result<u64, RehydrateError> {
let mut command = tokio::process::Command::new(program);
command
.arg("--version")
.envs(QUIET_ENV.iter().copied())
.stdin(Stdio::null())
.stderr(Stdio::null())
.kill_on_drop(true);
let output = tokio::time::timeout(IMPORT_TIMEOUT, command.output())
.await
.map_err(|_| RehydrateError::Other("opencode --version timed out".to_owned()))??;
let printed = String::from_utf8_lossy(&output.stdout);
printed
.split_whitespace()
.last()
.and_then(|version| version.trim_start_matches('v').split('.').next()?.parse().ok())
.ok_or_else(|| {
RehydrateError::Other(format!("can't tell opencode's version from {printed:?}"))
})
}
/// Run `opencode import` (the executable `program`) of `export` into `db`, from `cwd` (which
/// becomes the session's directory and project), handing back what it printed.
pub(crate) async fn run_import(
@@ -151,12 +196,23 @@ pub(crate) async fn run_import(
db: &Path,
cwd: &Path,
export: &Value,
) -> Result<String, RehydrateError> {
run(program, &[OsStr::new("import"), OsStr::new("--pure")], db, cwd, export).await
}
/// Run `opencode <args> <file>` (the executable `program`), `file` holding `export`, against `db`
/// from `cwd`, handing back what it printed.
async fn run(
program: &Path,
args: &[&OsStr],
db: &Path,
cwd: &Path,
export: &Value,
) -> Result<String, RehydrateError> {
let file = export_file(export)?;
let mut command = tokio::process::Command::new(program);
command
.arg("import")
.arg("--pure")
.args(args)
.arg(&*file)
.current_dir(cwd)
.env("OPENCODE_DB", db)
@@ -420,7 +476,7 @@ fn is_failure(row: &RehydrateMessage) -> bool {
&& matches!(row.content.as_slice(), [Content::Error(_)])
}
/// The row for the session itself (`<session>:title:<title>`, `<session>:session`), which the
/// The row for the session itself (`<session>:title:<digest>`, `<session>:session`), which the
/// session's own info stands for; for a fork, the original's too.
fn is_session_row(session: &RehydrateSession, row: &RehydrateMessage) -> bool {
let of = |id: &str| row.source_id.strip_prefix(id).is_some_and(|rest| rest.starts_with(':'));
@@ -1329,7 +1385,7 @@ pub(crate) mod tests {
row
}
const SES: &str = "ses_0d148d4b0001WQdGo06CVdqA3u";
pub const SES: &str = "ses_0d148d4b0001WQdGo06CVdqA3u";
const U1: &str = "msg_0d148d4b9001WQdGo06CVdqA3u";
const U2: &str = "msg_0d149b05a001c4o2UTJB1vX46n";
const A2: &str = "msg_0d149b2d5001gkfmc9bUMDXLJ5";
@@ -1337,7 +1393,7 @@ pub(crate) mod tests {
const TURN2: &str = "1790217663189:openai/gpt-x";
/// A session with every kind of row capture makes of opencode's parts.
fn every_kind() -> Vec<RehydrateMessage> {
pub fn every_kind() -> Vec<RehydrateMessage> {
let mut finish = answer(
"prt_0d148db35001Zcq9788cC4lx83",
1_790_217_608,
@@ -1357,7 +1413,7 @@ pub(crate) mod tests {
answer(A2, 1_790_217_663, U2, TURN2, vec![Content::Error("mock bad request".into())]);
failure.stop_reason = Some(StopReason::Error);
vec![
row(&format!("{SES}:title:Fix the build"), 1_790_217_606, Role::System, vec![]),
row(&session::title_id(SES, "Fix the build"), 1_790_217_606, Role::System, vec![]),
row("prt_0d148d4bf001JaxA9TwaMgU5vE", 1_790_217_606, Role::User, vec![Content::Text(
"fix the build".into(),
)]),
@@ -1578,9 +1634,10 @@ pub(crate) mod tests {
std::fs::write(
&path,
format!(
"#!/bin/sh\necho \"$@\" > {log}/args\npwd > {log}/cwd\necho \"$OPENCODE_DB \
$OPENCODE_DISABLE_MODELS_FETCH $OPENCODE_PURE\" > {log}/env\ncp \"$3\" \
{log}/export.json\necho boom >&2\nexit 3\n",
"#!/bin/sh\n[ \"$1\" = --version ] && echo 1.18.32 && exit 0\necho \"$@\" > \
{log}/args\npwd > {log}/cwd\necho \"$OPENCODE_DB $OPENCODE_DISABLE_MODELS_FETCH \
$OPENCODE_PURE\" > {log}/env\ncp \"$3\" {log}/export.json\necho boom >&2\nexit \
3\n",
log = dir.display()
),
)
@@ -1699,6 +1756,7 @@ pub(crate) mod tests {
&path,
format!(
"#!/bin/sh
[ \"$1\" = --version ] && echo 1.18.32 && exit 0
touch {}/ran
",
dir.display()
crates/atuin-common/src/harnesstools/opencode/rehydrate/v2.rs
+566
View File
@@ -0,0 +1,566 @@
//! The session opencode 2 imports (`opencode session import`): the JSON its `opencode session
//! export` writes (`SessionTransfer.Data`), the session's `info` and its `messages`, each one of
//! its `Session.Message.Info` whole.
//!
//! opencode 2 keeps a session as messages of its own kinds (its `session_message` rows): a
//! `user` prompt, text it injected (`synthetic`), each model call (`assistant`: its text and tool
//! calls, and how it ended), a `compaction`, and how a turn that did not succeed ended (`idle`).
//! Import inserts them under their own ids, in the order given, and refuses a session the
//! database holds already.
//!
//! # What survives
//!
//! - A row capture read from opencode 2, keyed on its message, is written back as that message,
//! under its id: re-captured, it reads back as itself.
//! - Other rows (opencode 1.x's parts, another harness's lines) are grouped into messages as
//! opencode 2 lays them out: a prompt's rows into one `user` message, each row of text the
//! harness injected into a `synthetic` one, and an assistant's rows into one message per model
//! call (the rows of one turn, up to the one that reports its usage). Their message ids are
//! minted, and each keeps the source ids of the rows it was written from in its `metadata`
//! ([`SOURCE`]): re-captured, it reads back as the first of them, and sync knows it holds the
//! rest. So does a message rows no message is written for (a result joined to its call, a
//! retry) went with.
//! - A fork's messages all get fresh ids (opencode's are keys of its whole database), and it
//! starts with a `synthetic` message holding its [marker](continuation::fork_marker_text),
//! which capture reads as its link.
//! - Reasoning (kept only as a marker), retries, and what capture keeps nothing of are left out.
//! A call without a result (one opencode was stopped in) is written failed, `aborted`, as
//! opencode 2 ends one.
//! - The session takes the directory it is imported into, its project, and its title. Its model
//! is left for opencode to pick, as for a new session: a model another harness ran on is no
//! provider's of opencode's.
use std::collections::{HashMap, HashSet};
use serde_json::{Map, Value, json};
use super::{
Turn, is_failure, is_session_row, join_results, millis, mint, step_tokens, title_of, turn_base,
};
use crate::harnesstools::opencode::session::SOURCE;
use crate::harnesstools::rehydrate::{
Flatten, RehydrateMessage, RehydrateSession, UNCAPTURED_OUTPUT, flatten_uncaptured_calls,
};
use crate::harnesstools::session::{Change, Content, Patch, Role, StopReason, ToolResult, ToolUse};
use crate::harnesstools::{AnyHarness, continuation};
/// What kind of opencode 2 message a draft is.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum Kind {
User,
Synthetic,
Assistant,
Compaction,
Idle,
}
/// A message being built from its rows.
struct Draft {
kind: Kind,
/// Its own id, for a row of opencode 2 (keyed on its message).
id: Option<String>,
/// The source ids of its rows, in order.
sources: Vec<String>,
/// The row a prompt's rows follow on from, which tells one prompt's rows from the next's.
parent: Option<String>,
/// The rows' turn (without step counts), which tells an assistant's rows from the next call's.
turn: Option<String>,
created: i64,
completed: i64,
model: Option<String>,
/// A prompt's or injected text's paragraphs; a compaction's summary.
text: Vec<String>,
/// An assistant's content (`text` and `tool` items).
content: Vec<Value>,
/// The model call reported its usage: the call is over.
tokens: Option<Value>,
finish: Option<&'static str>,
error: Option<Value>,
/// An `idle` message's outcome.
outcome: &'static str,
}
impl Draft {
fn new(kind: Kind, row: &RehydrateMessage) -> Self {
let at = millis(row.timestamp);
Self {
kind,
id: own_id(row),
sources: Vec::new(),
parent: row.parent_source_id.clone(),
turn: turn_base(row),
created: at,
completed: at,
model: row.model.clone(),
text: Vec::new(),
content: Vec::new(),
tokens: None,
finish: None,
error: None,
outcome: "failed",
}
}
/// Whether `row`, of `kind`, is more of this message.
fn holds(&self, kind: Kind, row: &RehydrateMessage) -> bool {
if self.kind != kind || self.id.is_some() || own_id(row).is_some() {
return false;
}
match kind {
Kind::User => self.parent == row.parent_source_id,
Kind::Assistant => self.tokens.is_none() && self.turn == turn_base(row),
Kind::Synthetic | Kind::Compaction | Kind::Idle => false,
}
}
}
/// The id of the opencode 2 message `row` is the whole of: a row capture read from opencode 2,
/// or a failed 1.x assistant message, both keyed on the message.
fn own_id(row: &RehydrateMessage) -> Option<String> {
let whole = row.source_id.starts_with("msg_") && !row.source_id.contains('/');
whole.then(|| row.source_id.clone())
}
/// What kind of message `row` goes into, and `None` for a row no message carries.
fn kind_of(row: &RehydrateMessage) -> Option<Kind> {
match (&row.role, row.content.as_slice()) {
(Role::User, content) if content.iter().any(|c| matches!(c, Content::Text(_))) => {
Some(Kind::User)
}
(Role::System, [Content::Text(_), ..]) => Some(Kind::Synthetic),
// how an opencode 2 turn that did not succeed ended
(Role::System, [])
if matches!(row.stop_reason, Some(StopReason::Error | StopReason::Aborted)) =>
{
Some(Kind::Idle)
}
(Role::Assistant, [Content::Summary(_)]) => Some(Kind::Compaction),
// a retry's error: the call went on
(Role::Assistant, [Content::Error(_)]) if !is_failure(row) => None,
(Role::Assistant, _) => Some(Kind::Assistant),
_ => None,
}
}
/// The JSON `opencode session export` writes for `session`, which `opencode session import`
/// reads back.
pub fn export(session: &RehydrateSession) -> Value {
// A note joins a text part of its own model call.
let rows = flatten_uncaptured_calls(&session.messages, &Flatten::Notes {
host: &|host, row| {
host.parent_source_id == row.parent_source_id && turn_base(host) == turn_base(row)
},
own: &|_| true,
});
let mut drafts: Vec<Draft> = Vec::new();
let joined = join_results(rows.clone());
let mut joined = joined.iter().peekable();
// rows no message is written for (a result joined to its call, a retry): held all the same,
// merged into the message they went with
let mut unwritten: Vec<String> = Vec::new();
// the draft each call went into, which its result's row goes with
let mut calls: HashMap<String, usize> = HashMap::new();
for original in &rows {
let Some(row) = joined.next_if(|row| row.source_id == original.source_id) else {
let call = match original.content.first() {
Some(Content::ToolResult(result)) => calls.get(result.call.as_ref()),
_ => None,
};
match call {
Some(&at) => drafts[at].sources.push(original.source_id.clone()),
None => unwritten.push(original.source_id.clone()),
}
continue;
};
if is_session_row(session, row) {
continue;
}
let Some(kind) = kind_of(row) else {
unwritten.push(row.source_id.clone());
continue;
};
if !drafts.last().is_some_and(|draft| draft.holds(kind, row)) {
drafts.push(Draft::new(kind, row));
}
let draft = drafts.last_mut().expect("a draft was just ensured");
draft.completed = draft.completed.max(millis(row.timestamp));
if draft.model.is_none() {
draft.model.clone_from(&row.model);
}
draft.sources.push(row.source_id.clone());
draft.sources.append(&mut unwritten);
add(draft, row);
for content in &row.content {
if let Content::ToolUse(call) = content {
calls.insert(call.id.to_string(), drafts.len() - 1);
}
}
}
if let Some(last) = drafts.last_mut() {
last.sources.append(&mut unwritten);
}
let mut used: HashSet<String> = HashSet::new();
// each message written, whether it kept its own id, and the rows it holds
let mut written: Vec<(Value, bool, Vec<String>)> = Vec::new();
// the rows of drafts no message is written for (a turn of reasoning alone), before any is
let mut orphaned: Vec<String> = Vec::new();
for (n, draft) in drafts.iter().enumerate() {
let id = match &draft.id {
Some(id) if session.fork_of.is_none() && used.insert(id.clone()) => id.clone(),
_ => mint("msg", draft.created, &format!("{}/message/{n}", session.id)),
};
let Some(message) = message(draft, &id) else {
// held all the same: by the message before it, else the one after
match written.last_mut() {
Some((_, _, sources)) => sources.extend(draft.sources.iter().cloned()),
None => orphaned.extend(draft.sources.iter().cloned()),
}
continue;
};
let mut sources = draft.sources.clone();
let after_first = sources.len().min(1);
sources.splice(after_first..after_first, std::mem::take(&mut orphaned));
written.push((message, draft.id.as_ref() == Some(&id), sources));
}
let mut messages: Vec<Value> = written
.into_iter()
.map(|(mut message, keyed, sources)| {
// keyed on the rows it was written from, so capture reads it back as them (see
// `SOURCE`); a fork's are new rows of a new session
if (!keyed || sources.len() > 1)
&& session.fork_of.is_none()
&& let Some((first, merged)) = sources.split_first()
{
message["metadata"] =
json!({SOURCE: {"session": session.id, "source": first, "merged": merged}});
}
message
})
.collect();
if let Some(of) = &session.fork_of {
let source = AnyHarness::from_name("opencode").expect("opencode is a harness");
let created = messages
.first()
.and_then(|m| m["time"]["created"].as_i64())
.unwrap_or_else(|| millis(session.started_at));
messages.insert(
0,
json!({
"id": mint("msg", created.saturating_sub(1), &format!("{}/fork", session.id)),
"type": "synthetic",
"time": {"created": created},
"text": continuation::fork_marker_text(source, &of.id, of.atuin_id.as_deref()),
}),
);
}
// how the last turn ended: failed or interrupted when its idle row says so
let outcome = match drafts.last() {
Some(draft) if draft.kind == Kind::Idle => draft.outcome,
_ => "succeeded",
};
let started = millis(session.started_at);
let updated = session.messages.iter().map(|m| millis(m.timestamp)).max().unwrap_or(started);
let updated = updated.max(started);
json!({
"info": {
"id": session.id,
// Import puts the session in the project of the directory it is imported into.
"projectID": "",
"title": session.title.clone().unwrap_or_else(|| title_of(session)),
"cost": 0,
"tokens": zero_tokens(),
// idle: a session that is not is one opencode waits on to finish its turn
"time": {"created": started, "updated": updated, "idle": updated},
"outcome": outcome,
"location": {"directory": session.cwd},
},
"messages": messages,
})
}
fn zero_tokens() -> Value {
json!({"input": 0, "output": 0, "reasoning": 0, "cache": {"read": 0, "write": 0}})
}
/// Add what `row` says to the message `draft` stands for.
fn add(draft: &mut Draft, row: &RehydrateMessage) {
let at = millis(row.timestamp);
for content in &row.content {
match content {
Content::Text(text) if draft.kind == Kind::Assistant => {
draft.content.push(json!({"type": "text", "text": text}));
}
Content::Text(text) | Content::Summary(text) => draft.text.push(text.clone()),
Content::ToolUse(call) => {
let result = row.content.iter().find_map(|c| match c {
Content::ToolResult(result) if result.call == call.id => Some(result),
_ => None,
});
let patch = row.content.iter().find_map(|c| match c {
Content::Patch(patch) if patch.call == call.id => Some(patch),
_ => None,
});
draft.content.push(tool(call, result, patch, at));
}
Content::Error(why) => draft.error = Some(error(why, row.stop_reason.as_ref())),
_ => {}
}
}
if draft.kind == Kind::Idle {
draft.outcome = match row.stop_reason {
Some(StopReason::Aborted) => "interrupted",
_ => "failed",
};
}
if let Some(usage) = &row.usage {
draft.tokens = Some(step_tokens(row, usage));
}
if row.usage.is_some() || row.stop_reason.is_some() {
draft.finish = Some(finish_of(row.stop_reason.as_ref()));
}
// opencode ends a call it was stopped in with an error saying so
if draft.kind == Kind::Assistant
&& row.stop_reason == Some(StopReason::Aborted)
&& draft.error.is_none()
{
draft.error = Some(error("aborted", row.stop_reason.as_ref()));
}
}
/// A model call's `finish` as opencode 2 takes it, which is one of its own few: an aborted call
/// finished in an `error`, and one of a reason it doesn't know `unknown`.
fn finish_of(stop: Option<&StopReason>) -> &'static str {
match stop {
None | Some(StopReason::EndTurn | StopReason::StopSequence) => "stop",
Some(StopReason::MaxTokens) => "length",
Some(StopReason::ToolUse) => "tool-calls",
Some(StopReason::Refusal) => "content-filter",
Some(StopReason::Error | StopReason::Aborted) => "error",
Some(StopReason::Other(other)) => match other.as_str() {
"stop" => "stop",
"length" => "length",
"tool-calls" => "tool-calls",
"content-filter" => "content-filter",
"error" => "error",
_ => "unknown",
},
}
}
/// The message `draft` is, under `id`; `None` for one with nothing to say.
fn message(draft: &Draft, id: &str) -> Option<Value> {
let created = draft_created(draft);
let text = draft.text.join("\n\n");
Some(match draft.kind {
Kind::User => {
json!({"id": id, "type": "user", "time": {"created": created}, "text": text, "files": []})
}
Kind::Synthetic => {
json!({"id": id, "type": "synthetic", "time": {"created": created}, "text": text})
}
Kind::Compaction => {
let mut message = json!({
"id": id,
"type": "compaction",
"time": {"created": created},
"status": "completed",
"reason": "auto",
"summary": text,
"recent": "",
"cost": 0,
"tokens": draft.tokens.clone().unwrap_or_else(zero_tokens),
});
// the model its call ran on, which its turn is keyed on
if let Some(Turn::Call {
provider, model, ..
}) = draft.turn.as_deref().and_then(Turn::of)
{
message["model"] = json!({"id": model, "providerID": provider});
}
message
}
Kind::Idle => {
json!({"id": id, "type": "idle", "time": {"created": created}, "outcome": draft.outcome})
}
Kind::Assistant => {
if draft.content.is_empty() && draft.error.is_none() && draft.tokens.is_none() {
return None;
}
let (provider, model) = provider_model(draft);
let mut message = json!({
"id": id,
"type": "assistant",
// opencode imports only a call that ended
"time": {"created": created, "completed": draft.completed.max(created)},
"agent": "build",
"model": {"id": model, "providerID": provider},
"content": draft.content,
"cost": 0,
"tokens": draft.tokens.clone().unwrap_or_else(zero_tokens),
});
if let Some(finish) = &draft.finish {
message["finish"] = json!(finish);
}
if let Some(error) = &draft.error {
message["error"] = error.clone();
}
message
}
})
}
/// When a message was created: an assistant's, as its turn says (the call's start, which the turn
/// key of every row of it repeats), else its first row's time.
fn draft_created(draft: &Draft) -> i64 {
match draft.turn.as_deref().and_then(Turn::of) {
Some(Turn::Call { created, .. }) if draft.kind == Kind::Assistant => created,
_ => draft.created,
}
}
/// The provider and model an assistant message ran on: its turn's, else what its rows say (a
/// model with no provider is taken for opencode's own).
fn provider_model(draft: &Draft) -> (String, String) {
if let Some(Turn::Call {
provider, model, ..
}) = draft.turn.as_deref().and_then(Turn::of)
{
return (provider.to_owned(), model.to_owned());
}
match draft.model.as_deref().map(|model| model.split_once('/').unwrap_or(("opencode", model))) {
Some((provider, model)) => (provider.to_owned(), model.to_owned()),
None => ("opencode".to_owned(), "unknown".to_owned()),
}
}
/// An error as opencode 2 reports one (`Session.StructuredError`), saying `why`.
fn error(why: &str, stop: Option<&StopReason>) -> Value {
let kind = match stop {
Some(StopReason::Aborted) => "aborted",
Some(StopReason::Refusal) => "provider.content-filter",
_ => "unknown",
};
json!({"type": kind, "message": why})
}
/// A tool item of an assistant message: `completed` or `error` with its result, failed as
/// `aborted` without one. A completed call whose patch capture kept gets the metadata opencode 2
/// keeps the diff in (`files[]`, `{file, patch, status}`).
fn tool(call: &ToolUse, result: Option<&ToolResult>, patch: Option<&Patch>, at: i64) -> Value {
// opencode's tool input is always an object.
let input = match &call.input {
Value::Object(_) => call.input.clone(),
other => Value::Object(Map::from_iter([("input".to_owned(), other.clone())])),
};
let state = match result {
// a command opencode 2 ran that failed: completed, its exit code in the metadata
Some(result) if result.error && call.name == "shell" && result.output.is_array() => json!({
"status": "completed",
"input": input,
"content": content_of(&result.output),
"metadata": {"exit": exit_code(&result.output)},
}),
Some(result) if result.error => json!({
"status": "error",
"input": input,
"error": {"type": "tool.execution", "message": text_of(&result.output)},
}),
Some(result) => {
let mut state = json!({
"status": "completed",
"input": input,
"content": content_of(&result.output),
});
if let Some(patch) = patch {
state["metadata"] = files(patch);
}
state
}
None => json!({
"status": "error",
"input": input,
"error": {"type": "aborted", "message": "Tool execution was interrupted"},
}),
};
json!({
"type": "tool",
"id": call.id.as_ref(),
"name": call.name,
"state": state,
"time": {"created": at, "ran": at, "completed": at},
})
}
/// A tool's output as opencode 2 keeps it: a list of content items, `{type: "text", text}` or
/// `{type: "file", uri, ..}`. A list read from opencode 2 is kept as it is; another harness's
/// blocks are made text (a text block's own, else a note of what it was: `[image]`), since
/// opencode refuses an import holding any item of another shape.
fn content_of(output: &Value) -> Value {
let Value::Array(items) = output else {
return json!([{"type": "text", "text": text_of(output)}]);
};
let items =
items.iter().map(|item| match (item["type"].as_str(), &item["text"], &item["uri"]) {
(Some("file"), _, Value::String(_)) => item.clone(),
// `text`, and another harness's text by another name (Codex's `input_text`)
(_, Value::String(text), _) => json!({"type": "text", "text": text}),
(kind, ..) => {
json!({"type": "text", "text": format!("[{}]", kind.unwrap_or("output"))})
}
});
Value::Array(items.collect())
}
/// The exit code a failed command's output reports (`Exited with code 2`, the shell tool's last
/// item), else 1.
fn exit_code(output: &Value) -> i64 {
output
.as_array()
.into_iter()
.flatten()
.filter_map(|item| {
item["text"].as_str()?.strip_prefix("Exited with code ")?.trim().parse().ok()
})
.next_back()
.unwrap_or(1)
}
/// A tool's output as one text.
fn text_of(output: &Value) -> String {
match output {
Value::String(text) => text.clone(),
Value::Null => UNCAPTURED_OUTPUT.to_owned(),
Value::Array(items) if items.iter().all(|item| item["text"].is_string()) => {
items.iter().filter_map(|item| item["text"].as_str()).collect::<Vec<_>>().join("\n")
}
other => other.to_string(),
}
}
/// The metadata opencode 2 keeps of a call that changed files (`edit`, `write`, `apply_patch`).
fn files(patch: &Patch) -> Value {
let files: Vec<Value> = patch
.files
.iter()
.map(|file| {
let status = match file.change {
Change::Add => "added",
Change::Delete => "deleted",
Change::Update => "modified",
};
json!({
"file": file.moved_to.as_deref().unwrap_or(&file.path),
"patch": file.unified(),
"status": status,
"additions": file.additions(),
"deletions": file.deletions(),
})
})
.collect();
json!({"files": files})
}
#[cfg(test)]
mod tests;
crates/atuin-common/src/harnesstools/opencode/rehydrate/v2/tests.rs
+647
View File
@@ -0,0 +1,647 @@
use std::collections::HashMap;
use std::path::{Path, PathBuf};
use rstest::rstest;
use serde_json::{Value, json};
use sqlx::sqlite::{SqliteConnectOptions, SqliteConnection};
use sqlx::{Connection, Executor};
use time::OffsetDateTime;
use super::*;
use crate::harnesstools::opencode::rehydrate::import;
use crate::harnesstools::opencode::rehydrate::tests::{SES, captured, every_kind};
use crate::harnesstools::opencode::session::projection::projected;
use crate::harnesstools::rehydrate::{ForkOf, RehydrateError, testing};
use crate::harnesstools::session::Usage;
/// The part of opencode 2.0.24's schema its import writes and capture reads.
const DDL: &[&str] = &[
"CREATE TABLE event_sequence (aggregate_id TEXT PRIMARY KEY, seq INTEGER NOT NULL, owner_id \
TEXT)",
"CREATE TABLE event (id TEXT PRIMARY KEY, aggregate_id TEXT NOT NULL, seq INTEGER NOT NULL, \
created INTEGER DEFAULT 0 NOT NULL, type TEXT NOT NULL, data TEXT NOT NULL)",
"CREATE TABLE session_v2 (id TEXT PRIMARY KEY, project_id TEXT NOT NULL, parent_id TEXT, \
fork_session_id TEXT, slug TEXT NOT NULL, directory TEXT NOT NULL, title TEXT, version TEXT \
NOT NULL, revert TEXT, time_created INTEGER NOT NULL, time_updated INTEGER NOT NULL)",
"CREATE TABLE session_message (id TEXT PRIMARY KEY, session_id TEXT NOT NULL, type TEXT NOT \
NULL, seq INTEGER NOT NULL, time_created INTEGER NOT NULL, time_updated INTEGER NOT NULL, \
data TEXT NOT NULL)",
];
/// The real export opencode 2.0.24 wrote of a session it ran (`opencode session export`; the
/// directory redacted): three turns of shell, read, write and edit calls, one shell command
/// failing.
const EXPORT: &str = include_str!("../../../../../tests/fixtures/opencode/v2_export.json");
async fn database(path: &Path) -> SqliteConnection {
let opts = SqliteConnectOptions::new().filename(path).create_if_missing(true);
let mut conn = SqliteConnection::connect_with(&opts).await.unwrap();
for ddl in DDL {
conn.execute(*ddl).await.unwrap();
}
conn
}
/// What `opencode session import` does with an export (opencode 2.0.24,
/// `SessionTransfer.import`): the session in the directory it is imported into, and each message
/// but an assistant's that never ended under its id, numbered in the order given, its data the
/// message without `id` and `type`.
async fn session_import(conn: &mut SqliteConnection, export: &Value, directory: &str) {
let info = &export["info"];
let id = info["id"].as_str().unwrap();
crate::db::query::<sqlx::Sqlite>(
"INSERT INTO session_v2 (id, project_id, slug, directory, title, version, time_created, \
time_updated) VALUES (?1, 'global', 'slug', ?2, ?3, '2.0.24', ?4, ?4)",
)
.bind(id)
.bind(directory)
.bind(info["title"].as_str())
.bind(info["time"]["created"].as_i64())
.execute(&mut *conn)
.await
.unwrap();
let messages: Vec<&Value> = export["messages"]
.as_array()
.unwrap()
.iter()
.filter(|m| m["type"] != "assistant" || !m["time"]["completed"].is_null())
.collect();
for (seq, message) in (1i64..).zip(&messages) {
let mut data = (*message).clone();
let data = data.as_object_mut().unwrap();
let id = data.remove("id").unwrap();
let kind = data.remove("type").unwrap();
crate::db::query::<sqlx::Sqlite>(
"INSERT INTO session_message (id, session_id, type, seq, time_created, time_updated, \
data) VALUES (?1, ?2, ?3, ?4, ?5, ?5, ?6)",
)
.bind(id.as_str())
.bind(info["id"].as_str())
.bind(kind.as_str())
.bind(seq)
.bind(message["time"]["created"].as_i64())
.bind(Value::Object(data.clone()).to_string())
.execute(&mut *conn)
.await
.unwrap();
}
crate::db::query::<sqlx::Sqlite>(
"INSERT INTO event_sequence (aggregate_id, seq) VALUES (?1, ?2)",
)
.bind(id)
.bind(i64::try_from(messages.len()).unwrap())
.execute(&mut *conn)
.await
.unwrap();
}
/// `export` imported into a new database, and captured back.
async fn imported(export: &Value) -> Vec<RehydrateMessage> {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("opencode.db");
let mut conn = database(&path).await;
session_import(&mut conn, export, "/here").await;
conn.close().await.unwrap();
captured(&path).await
}
fn session(id: &str, title: Option<&str>, messages: Vec<RehydrateMessage>) -> RehydrateSession {
RehydrateSession {
id: id.to_owned(),
title: title.map(str::to_owned),
cwd: PathBuf::from("/elsewhere/proj"),
original_cwd: Some(PathBuf::from("/work/proj")),
git_branch: None,
model: None,
started_at: messages.first().map_or(OffsetDateTime::UNIX_EPOCH, |m| m.timestamp),
messages,
fork_of: None,
}
}
fn is_session(id: &str, row: &RehydrateMessage) -> bool {
row.source_id.starts_with(&format!("{id}:"))
}
/// Each row's id, role and content, but for reasoning, which is not written back (capture keeps
/// only that there was some).
fn keys(rows: &[RehydrateMessage]) -> Vec<(String, Role, Vec<Content>)> {
rows.iter()
.map(|r| {
let mut content = r.content.clone();
content.retain(|c| !matches!(c, Content::ReasoningSummary { .. }));
(r.source_id.clone(), r.role.clone(), content)
})
.collect()
}
/// A session opencode 2.0.24 wrote, captured, written back and captured again, keys the same rows
/// in the same order, saying the same: every message is its own row.
#[rstest]
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn a_real_session_recaptures_as_the_same_rows() {
let real: Value = serde_json::from_str(EXPORT).unwrap();
let id = real["info"]["id"].as_str().unwrap();
let original = testing::synced_as(imported(&real).await, true);
assert!(original.len() > 10);
let written = export(&session(id, Some("octest"), original.clone()));
let again = imported(&written).await;
pretty_assertions::assert_eq!(keys(&again), keys(&original));
testing::assert_nothing_new(&original, again.iter().map(|r| r.source_id.as_str()));
let calls = |rows: &[RehydrateMessage]| -> Vec<_> {
rows.iter()
.filter(|r| !is_session(id, r))
.map(|r| (r.turn_id.clone(), r.usage, r.stop_reason.clone(), r.model.clone()))
.collect()
};
assert_eq!(calls(&again), calls(&original));
// the messages are those opencode exported, but for what capture keeps nothing of: the idle
// rows of turns that succeeded, and reasoning
let kinds = |export: &Value| -> Vec<(String, String)> {
export["messages"]
.as_array()
.unwrap()
.iter()
.filter(|m| m["type"] != "idle")
.map(|m| (m["id"].as_str().unwrap().to_owned(), m["type"].as_str().unwrap().to_owned()))
.collect()
};
assert_eq!(kinds(&written), kinds(&real));
let edit = |export: &Value| -> Value {
export["messages"]
.as_array()
.unwrap()
.iter()
.flat_map(|m| m["content"].as_array().into_iter().flatten())
.find(|item| item["name"] == "edit")
.map(|item| item["state"]["metadata"]["files"][0]["patch"].clone())
.unwrap()
};
// the diff, without the `Index:` header opencode's has
let (written, real) = (edit(&written), edit(&real));
assert!(real.as_str().unwrap().ends_with(written.as_str().unwrap()), "{written} vs {real}");
}
/// Rows opencode 1.x wrote (parts) are laid out as opencode 2's messages: the prompt, the text
/// opencode injected into it, the model call with its text, calls and usage, and the next prompt
/// with the call that failed.
#[rstest]
fn every_kind_of_1_x_row_goes_into_a_message() {
let written = export(&session(SES, Some("Fix the build"), every_kind()));
let messages = written["messages"].as_array().unwrap();
let kinds: Vec<&str> = messages.iter().map(|m| m["type"].as_str().unwrap()).collect();
assert_eq!(kinds, ["user", "synthetic", "assistant", "user", "assistant"]);
assert_eq!(messages[0]["text"], "fix the build");
assert_eq!(messages[1]["text"], "Called the Read tool");
let call = &messages[2];
assert_eq!(call["model"], json!({"id": "gpt-x", "providerID": "openai"}));
assert_eq!(call["time"]["created"], 1_790_217_606_979i64);
assert_eq!(call["finish"], "tool-calls");
assert_eq!(
call["tokens"],
json!({"input": 802, "output": 42, "reasoning": 10, "cache": {"read": 200, "write": 0}})
);
let content: Vec<(&str, &str)> = call["content"]
.as_array()
.unwrap()
.iter()
.map(|item| {
let what = item["text"].as_str().or_else(|| item["state"]["status"].as_str());
(item["type"].as_str().unwrap(), what.unwrap())
})
.collect();
// the retry's error is not one of the call's
assert_eq!(content, [("text", "On it."), ("tool", "completed"), ("tool", "error")]);
assert_eq!(call["content"][1]["state"]["content"], json!([{"type": "text", "text": "ok"}]));
assert_eq!(
call["content"][2]["state"]["error"],
json!({"type": "tool.execution", "message": "no such file"})
);
// a failed call keeps its message's id
let failed = &messages[4];
assert_eq!(failed["id"], "msg_0d149b2d5001gkfmc9bUMDXLJ5");
assert_eq!(failed["error"], json!({"type": "unknown", "message": "mock bad request"}));
assert_eq!(written["info"]["title"], "Fix the build");
// idle, or opencode waits for the turn it would be in to end
assert_eq!(written["info"]["time"]["idle"], written["info"]["time"]["updated"]);
assert_eq!(written["info"]["location"]["directory"], "/elsewhere/proj");
}
/// A session opencode 1.x wrote, restored on opencode 2, reads back under the ids it was synced
/// with: its rows are grouped into fewer messages than it had parts, each keyed on its first
/// row's id, so re-capturing it pushes nothing new.
#[rstest]
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn a_1_x_session_restored_on_2_recaptures_as_nothing_new() {
holds_every_row(every_kind()).await;
}
/// `rows` written back for opencode 2 as session [`SES`], and what sync makes of the copy: every
/// row of it re-captured is one synced already, and the copy is known to hold every row, those
/// merged into another's message too; the rows each merged row went into, by its id.
async fn holds_every_row(rows: Vec<RehydrateMessage>) -> HashMap<String, String> {
let written = export(&session(SES, Some("Fix the build"), rows.clone()));
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("opencode.db");
let mut conn = database(&path).await;
session_import(&mut conn, &written, "/here").await;
conn.close().await.unwrap();
let again = captured(&path).await;
let again: Vec<&str> = again
.iter()
.map(|r| r.source_id.as_str())
.filter(|id| !id.starts_with(&format!("{SES}:")))
.collect();
assert_eq!(again.len(), written["messages"].as_array().unwrap().len());
testing::assert_nothing_new(&rows, again.iter().copied());
let held = projected(&path, SES).await.unwrap().unwrap();
let mut known: Vec<&str> =
held.rows.iter().chain(held.merged.keys()).map(String::as_str).collect();
known.sort_unstable();
let mut synced: Vec<&str> = rows
.iter()
.map(|r| r.source_id.as_str())
.filter(|id| !id.starts_with(&format!("{SES}:")))
.collect();
synced.sort_unstable();
assert_eq!(known, synced);
assert_eq!(held.rows, again);
held.merged
}
fn said_by_assistant(id: &str, content: Vec<Content>) -> RehydrateMessage {
RehydrateMessage {
role: Role::Assistant,
content,
..prompt(id, None)
}
}
/// A row no message is written for is held all the same: a turn that left only reasoning (by
/// the message before it, else the one after), a retry, and a result kept apart from its call
/// (by its call's message).
#[rstest]
#[case::a_reasoning_only_turn(&["prt_u1", "prt_r", "prt_u2"])]
#[case::a_retry_only_turn(&["prt_u1", "prt_e", "prt_u2"])]
#[case::a_reasoning_only_turn_first(&["prt_r", "prt_u1"])]
#[case::a_result_apart(&["prt_u1", "prt_c", "prt_res", "prt_u2"])]
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn a_row_no_message_is_written_for_is_held(#[case] ids: &[&str]) {
let call = || ToolUse {
id: "call_1".to_owned().into(),
name: "shell".to_owned(),
input: json!({"command": "ls"}),
};
let rows = ids
.iter()
.map(|id| match *id {
"prt_r" => said_by_assistant(id, vec![Content::ReasoningSummary { tokens: None }]),
"prt_e" => said_by_assistant(id, vec![Content::Error("overloaded".to_owned())]),
"prt_c" => said_by_assistant(id, vec![Content::ToolUse(call())]),
"prt_res" => RehydrateMessage {
role: Role::Tool,
..said_by_assistant(id, vec![Content::ToolResult(ToolResult {
call: "call_1".to_owned().into(),
output: json!("files"),
error: false,
})])
},
_ => prompt(id, None),
})
.collect();
let merged = holds_every_row(rows).await;
if ids.contains(&"prt_res") {
assert_eq!(merged.get("prt_res").map(String::as_str), Some("prt_c"));
}
}
/// A message atuin wrote names the session it wrote it for: opencode's own fork, which copies it
/// whole into another session, has a row of its own, read under its own id.
#[rstest]
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn a_copy_of_a_restored_message_in_another_session_is_its_own() {
let mut written = export(&session(SES, None, every_kind()));
written["info"]["id"] = json!("ses_0d148d4b0001WQdGo06CVdqAff");
for (n, message) in written["messages"].as_array_mut().unwrap().iter_mut().enumerate() {
message["id"] = json!(format!("msg_copy{n}"));
}
let again = imported(&written).await;
let ids: Vec<&str> =
again.iter().map(|r| r.source_id.as_str()).filter(|id| !id.starts_with("ses_")).collect();
assert!(!ids.is_empty());
assert!(ids.iter().all(|id| id.starts_with("msg_copy")), "{ids:?}");
}
/// A fork's messages get ids of their own, and it starts with its marker, which links it to the
/// session it was forked from once captured. Its calls, a compaction's too, are the original's,
/// so they are not counted again.
#[rstest]
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn a_fork_is_linked_to_its_original() {
let mut real: Value = serde_json::from_str(EXPORT).unwrap();
// a completed compaction as opencode 2 writes one (`v2.json`)
real["messages"].as_array_mut().unwrap().push(json!({
"id": "msg_11e4bc8960010Ypn4F9uXZ2kQe",
"type": "compaction",
"time": {"created": 1_791_509_645_470i64},
"status": "completed",
"reason": "auto",
"model": {"id": "mock-model", "providerID": "mock"},
"summary": "## Objective\n- mock objective\n",
"recent": "",
"cost": 0.000_923,
"tokens": {"input": 801, "output": 41, "reasoning": 10, "cache": {"read": 200, "write": 0}},
}));
let original = testing::synced_as(imported(&real).await, true);
let mut fork = session("ses_ee1b7b1dfffe46nQL0NuU68Gzz", None, original.clone());
fork.fork_of = Some(ForkOf {
id: real["info"]["id"].as_str().unwrap().to_owned(),
atuin_id: None,
path: None,
});
let written = export(&fork);
let ids = |export: &Value| -> Vec<String> {
export["messages"]
.as_array()
.unwrap()
.iter()
.map(|m| m["id"].as_str().unwrap().to_owned())
.collect()
};
let theirs = ids(&real);
assert!(ids(&written).iter().all(|id| id.starts_with("msg_") && !theirs.contains(id)));
assert_eq!(ids(&written), ids(&export(&fork)), "a retried import mints the same ids");
let again = imported(&written).await;
let first = again.iter().find(|row| !is_session(&fork.id, row)).unwrap();
assert_eq!(first.role, Role::System);
let [Content::Text(marker)] = first.content.as_slice() else {
panic!("no marker: {first:?}");
};
assert!(continuation::is_fork_marker(marker), "{marker}");
assert!(marker.contains(theirs_session(&real)), "{marker}");
let calls = |rows: &[RehydrateMessage]| -> Vec<_> {
rows.iter().filter_map(|r| r.turn_id.clone().zip(r.usage)).collect()
};
assert!(calls(&original).iter().any(|(turn, _)| turn.contains(":mock/mock-model#")));
assert_eq!(calls(&again), calls(&original));
}
fn theirs_session(export: &Value) -> &str {
export["info"]["id"].as_str().unwrap()
}
fn prompt(id: &str, parent: Option<&str>) -> RehydrateMessage {
RehydrateMessage {
source_id: id.to_owned(),
parent_source_id: parent.map(str::to_owned),
timestamp: OffsetDateTime::from_unix_timestamp(1_790_000_000).unwrap(),
role: Role::User,
content: vec![Content::Text(format!("said {id}"))],
model: None,
usage: None,
stop_reason: None,
turn_id: None,
cwd: None,
git_branch: None,
}
}
/// Rows of one prompt (following on from the same row) are one message; prompts that follow on
/// from each other (another harness's consecutive lines) stay apart.
#[rstest]
#[case::one_prompt(&[("a", None), ("b", None)], &["said a\n\nsaid b"])]
#[case::two_prompts(&[("a", None), ("b", Some("a"))], &["said a", "said b"])]
fn a_prompt_is_the_rows_that_follow_on_from_one_row(
#[case] rows: &[(&str, Option<&str>)],
#[case] texts: &[&str],
) {
let rows = rows.iter().map(|(id, parent)| prompt(id, *parent)).collect();
let written = export(&session(SES, None, rows));
let got: Vec<&str> = written["messages"]
.as_array()
.unwrap()
.iter()
.map(|m| m["text"].as_str().unwrap())
.collect();
assert_eq!(got, texts);
}
/// A call's `finish` is one opencode 2's import takes (`Session.Message.Assistant.finish`), whatever
/// the call ended for; one that was stopped also says so in its error, as opencode writes it.
#[rstest]
#[case::end_turn(Some(StopReason::EndTurn), "stop")]
#[case::stop_sequence(Some(StopReason::StopSequence), "stop")]
#[case::max_tokens(Some(StopReason::MaxTokens), "length")]
#[case::tool_use(Some(StopReason::ToolUse), "tool-calls")]
#[case::refusal(Some(StopReason::Refusal), "content-filter")]
#[case::error(Some(StopReason::Error), "error")]
#[case::aborted(Some(StopReason::Aborted), "error")]
#[case::another_harness_s(Some(StopReason::Other("pause_turn".into())), "unknown")]
#[case::opencode_s_own(Some(StopReason::Other("length".into())), "length")]
#[case::none(None, "stop")]
fn a_call_finishes_as_opencode_2_takes_it(#[case] stop: Option<StopReason>, #[case] want: &str) {
const ALLOWED: [&str; 6] =
["stop", "length", "tool-calls", "content-filter", "error", "unknown"];
let mut row = said_by_assistant("prt_a", vec![Content::Text("hi".to_owned())]);
row.usage = Some(Usage::default());
row.stop_reason = stop.clone();
let written = export(&session(SES, None, vec![row]));
let call = &written["messages"][0];
assert_eq!(call["finish"], want);
assert!(ALLOWED.contains(&want));
let aborted = stop == Some(StopReason::Aborted);
assert_eq!(call["error"]["type"] == "aborted", aborted, "{call}");
}
/// The session ended as its last turn did: failed or interrupted when its idle row says so.
#[rstest]
#[case::succeeded(None, "succeeded")]
#[case::failed(Some(StopReason::Error), "failed")]
#[case::interrupted(Some(StopReason::Aborted), "interrupted")]
fn the_session_ended_as_its_last_turn_did(#[case] stop: Option<StopReason>, #[case] want: &str) {
let mut rows = vec![prompt("msg_u", None)];
if let Some(stop) = stop {
let mut idle = prompt("msg_i", None);
idle.role = Role::System;
idle.content.clear();
idle.stop_reason = Some(stop);
rows.push(idle);
}
let written = export(&session(SES, None, rows));
assert_eq!(written["info"]["outcome"], want);
}
/// A model call capture kept no result of is written failed, as opencode 2 ends a call it was
/// stopped in: import keeps only ended calls.
#[rstest]
fn a_call_without_a_result_is_written_aborted() {
let mut row = RehydrateMessage {
source_id: "prt_1".to_owned(),
parent_source_id: None,
timestamp: OffsetDateTime::from_unix_timestamp(1_790_000_000).unwrap(),
role: Role::Assistant,
content: vec![Content::ToolUse(ToolUse {
id: "call_1".to_owned().into(),
name: "shell".to_owned(),
input: json!({"command": "sleep 100"}),
})],
model: Some("opencode/big-pickle".to_owned()),
usage: None,
stop_reason: None,
turn_id: None,
cwd: None,
git_branch: None,
};
row.usage = Some(Usage::default());
let written = export(&session(SES, None, vec![row]));
let call = &written["messages"][0];
assert_eq!(call["model"], json!({"id": "big-pickle", "providerID": "opencode"}));
assert!(call["time"]["completed"].is_i64());
assert_eq!(call["content"][0]["state"]["status"], "error");
assert_eq!(call["content"][0]["state"]["error"]["type"], "aborted");
}
/// A tool's output as opencode 2 keeps it: items of its own (`text`, `file`) as they are, another
/// harness's text blocks (whatever their type) as text, and anything else (a Claude image, a `tool_reference`) as a note
/// of it, which opencode's import takes. Only a shell call is written completed with an exit code
/// for a failure.
#[rstest]
#[case::opencode_s_own(
"shell",
json!([{"type": "text", "text": "out"}, {"type": "file", "uri": "file:///a.png", "mime": "image/png"}]),
false,
json!({"status": "completed", "content": [
{"type": "text", "text": "out"},
{"type": "file", "uri": "file:///a.png", "mime": "image/png"},
]}),
)]
#[case::a_claude_image(
"Read",
json!([
{"type": "text", "text": "see", "cache_control": {"type": "ephemeral"}},
{"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": "iVBORw0KGgo="}},
{"type": "tool_reference", "tool_name": "x"},
]),
false,
json!({"status": "completed", "content": [
{"type": "text", "text": "see"},
{"type": "text", "text": "[image]"},
{"type": "text", "text": "[tool_reference]"},
]}),
)]
#[case::codex_s_text(
"exec",
json!([{"type": "input_text", "text": "ran"}]),
false,
json!({"status": "completed", "content": [{"type": "text", "text": "ran"}]}),
)]
#[case::a_failed_command(
"shell",
json!([{"type": "text", "text": "no\n"}, {"type": "text", "text": "Exited with code 2"}]),
true,
json!({"status": "completed", "content": [
{"type": "text", "text": "no\n"},
{"type": "text", "text": "Exited with code 2"},
], "metadata": {"exit": 2}}),
)]
#[case::another_tool_s_failure(
"Read",
json!([{"type": "text", "text": "no such file"}]),
true,
json!({"status": "error", "error": {"type": "tool.execution", "message": "no such file"}}),
)]
fn a_tool_s_output_is_written_as_opencode_2_keeps_it(
#[case] name: &str,
#[case] output: Value,
#[case] error: bool,
#[case] want: Value,
) {
let call = ToolUse {
id: "call_1".to_owned().into(),
name: name.to_owned(),
input: json!({}),
};
let result = ToolResult {
call: "call_1".to_owned().into(),
output,
error,
};
let mut state = tool(&call, Some(&result), None, 0)["state"].clone();
state.as_object_mut().unwrap().remove("input");
assert_eq!(state, want);
}
/// A directory programs can be run from, next to the test binary: CI's temporary directory is
/// mounted without exec, so a fake program written there could not be run.
#[cfg(unix)]
fn exec_dir() -> tempfile::TempDir {
let exe = std::env::current_exe().unwrap();
tempfile::tempdir_in(exe.parent().unwrap()).unwrap()
}
/// A fake opencode 2: says its version, and records how an import was run, then fails it.
#[cfg(unix)]
fn fake_opencode(dir: &Path) -> (tempfile::TempDir, PathBuf) {
use std::os::unix::fs::PermissionsExt;
let bin = exec_dir();
let path = bin.path().join("opencode");
std::fs::write(
&path,
format!(
"#!/bin/sh\n[ \"$1\" = --version ] && echo 'opencode v2.0.24' && exit 0\necho \"$@\" \
> {log}/args\npwd > {log}/cwd\necho \"$OPENCODE_DB\" > {log}/env\nfor file; do :; \
done\ncp \"$file\" {log}/export.json\necho boom >&2\nexit 3\n",
log = dir.display()
),
)
.unwrap();
std::fs::set_permissions(&path, std::fs::Permissions::from_mode(0o755)).unwrap();
(bin, path)
}
/// opencode 2 imports with `opencode session import`, through a server of its own on the
/// database it was given, into the session's directory, the export in opencode 2's format.
#[cfg(unix)]
#[rstest]
#[tokio::test]
async fn opencode_2_imports_its_own_export() {
let dir = tempfile::tempdir().unwrap();
let (_bin, program) = fake_opencode(dir.path());
let db = dir.path().join("opencode.db");
let mut s = session(SES, Some("t"), every_kind());
s.cwd = dir.path().canonicalize().unwrap();
let err = import(&program, &db, &s).await.unwrap_err();
assert!(matches!(&err, RehydrateError::Other(why) if why.contains("boom")), "{err}");
let read = |name: &str| std::fs::read_to_string(dir.path().join(name)).unwrap();
let args = read("args");
let expected = format!("session import --standalone --directory {} ", s.cwd.display());
assert!(args.starts_with(&expected), "{args}");
assert_eq!(read("cwd").trim(), s.cwd.to_str().unwrap());
assert_eq!(read("env").trim(), db.display().to_string());
let sent: Value = serde_json::from_str(&read("export.json")).unwrap();
assert_eq!(sent, export(&s));
let file = args.trim().rsplit(' ').next().unwrap();
assert!(!Path::new(file).exists(), "the export file is removed");
}
/// A session opencode 2 holds already is not imported again.
#[rstest]
#[tokio::test]
async fn a_session_opencode_2_holds_is_not_imported_again() {
let dir = tempfile::tempdir().unwrap();
let db = dir.path().join("opencode.db");
let mut conn = database(&db).await;
session_import(&mut conn, &export(&session(SES, None, every_kind())), "/here").await;
conn.close().await.unwrap();
let missing = dir.path().join("no-such-opencode");
let err = import(&missing, &db, &session(SES, None, vec![])).await.unwrap_err();
assert!(matches!(err, RehydrateError::AlreadyExists(path) if path == db));
}
crates/atuin-common/src/harnesstools/opencode/session.rs
+165 -62
View File
@@ -103,7 +103,8 @@ use tokio::sync::{Mutex, MutexGuard, watch};
use typed_builder::TypedBuilder;
use crate::db::sqlite::observe::{
ObserveConfig, ReplayBehavior, RowAppendedEvent, SqliteObserver, TableSchema, Tailable,
ObserveConfig, ReplayBehavior, RowAppendedEvent, SqliteObserver, SqliteTableObserver,
TableSchema, Tailable,
};
use crate::db::{query_as, query_scalar};
use crate::harnesstools::opencode::Opencode;
@@ -122,6 +123,11 @@ use crate::utils::{env_nonempty, home_dir};
pub(crate) mod projection;
mod v2;
/// The key of an opencode 2 message's `metadata` under which atuin keeps, in a message it wrote
/// (a restore), the source ids of the rows it was written from: `{source, merged}`, the first and
/// the rest. Capture reads the message back as the first, so it pushes nothing new.
pub(crate) const SOURCE: &str = "atuin";
#[derive(Debug, Clone, TypedBuilder)]
pub struct OpencodeSessions {
#[builder(default, setter(strip_option, into))]
@@ -350,39 +356,56 @@ impl Listener for OpencodeListener {
/// loses nothing: the session is offered again, under the same id, with its aggregate's next
/// row, and resumes at the row it had reached.
///
/// A failure of the tail is reported and then carried on from. Only a database that cannot
/// be observed at all ends this stream, and with it every session it has offered.
/// A failure of the tail is reported and then carried on from, and so is a database with no
/// event log to tail yet: it is reported once and probed again on every poll. The stream
/// ends only with the tail, and with it every session it has offered.
fn watch(self) -> impl Stream<Item = Result<OpencodeSession, WatchError>> + Send + 'static {
let db = self.db;
let replay = self.replay;
async_stream::stream! {
let reads = Arc::new(Reads::new(db.clone()));
let observe = |source| WatchError::Observe { db: db.clone(), source };
// opencode 1.x's sessions are tailed from its event log; opencode 2.0 creates the
// table too and never writes to it, so its sessions are polled for (see
// `v2::changes`)
if let Err(err) = reads.require_event_table().await {
yield Err(err);
return;
}
let mut events = match SqliteObserver::new(&db)
.append::<EventRow>(
ObserveConfig::builder().replay(replay).poll_interval(POLL).build(),
)
.await
{
Ok(events) => events,
Err(err) => {
yield Err(observe(err));
return;
}
};
// opencode 1.x's sessions are tailed from its event log, once there is one: a
// database opencode has not created its tables in yet, or is migrating, has none.
// opencode 2.0 creates the table too and never writes to it, so its sessions are
// polled for meanwhile all the same (see `v2::changes`).
let mut events: Option<SqliteTableObserver<RowAppendedEvent<EventRow>>> = None;
let mut probe = tokio::time::interval(POLL);
let mut waited = false;
let changes = v2::changes(db.clone(), replay);
futures::pin_mut!(changes);
let mut sessions: HashMap<Aggregate, Live> = HashMap::new();
loop {
tokio::select! {
next = events.next() => {
_ = probe.tick(), if events.is_none() => {
let tail = match reads.require_event_table().await {
// a log that appeared while it was waited for is new, rows and all
Ok(()) => SqliteObserver::new(&db)
.append::<EventRow>(
ObserveConfig::builder()
.replay(if waited { ReplayBehavior::All } else { replay })
.poll_interval(POLL)
.build(),
)
.await
.map_err(observe),
Err(err) => Err(err),
};
match tail {
Ok(tail) => events = Some(tail),
Err(err) if !waited => {
waited = true;
yield Err(err);
}
Err(err) => tracing::debug!(%err, "opencode's event log still cannot be tailed"),
}
}
next = async {
match events.as_mut() {
Some(events) => events.next().await,
None => std::future::pending().await,
}
} => {
let row = match next {
Some(Ok(RowAppendedEvent(row))) => row,
// The tail takes the table up again by itself and the rows it could
@@ -729,7 +752,7 @@ impl Reader {
Err(err) => Some(Err(MessageError::from(err))),
},
EventKind::Session => {
let message = OpencodeMessage::session(data)?;
let message = OpencodeMessage::session(data, true)?;
let title = message.title();
if title.is_some() && title == self.title {
return None;
@@ -1245,19 +1268,11 @@ impl Reads {
/// log, or a database caught mid-migration); the consumer decides when to retry.
async fn require_event_table(self: &Arc<Self>) -> Result<(), WatchError> {
let probe = self
.read("event table probe", |conn| {
Box::pin(async move {
query_scalar::<Sqlite, i64>(
"SELECT 1 FROM sqlite_master WHERE type = 'table' AND name = 'event'",
)
.fetch_optional(conn)
.await
})
})
.read("event table probe", |conn| Box::pin(async move { logged(conn).await }))
.await;
// a probe that could not run is inconclusive: whatever is wrong with the database, the
// observer's own connection will report it
if probe.is_none_or(|found| found.is_some()) {
if probe.is_none_or(|found| found) {
return Ok(());
}
Err(WatchError::Io(io::Error::new(
@@ -1274,6 +1289,9 @@ impl Reads {
///
/// opencode indexes `(aggregate_id, seq)` uniquely, so this is a seek rather than a scan: the
/// cast that rebuilds a text id sits on the bound value, which leaves the index usable.
///
/// A database with no event log holds no rows of it (opencode 2.0's sessions are read from
/// elsewhere; see `v2`).
async fn page(self: &Arc<Self>, aggregate: &Aggregate, from: i64) -> Option<Vec<PageRow>> {
const TEXT: &str = "SELECT id, seq, type, data FROM event WHERE aggregate_id = CAST(?1 AS \
TEXT) AND seq >= ?2 ORDER BY seq ASC LIMIT ?3";
@@ -1285,12 +1303,16 @@ impl Reads {
};
self.read("session page read", move |conn| {
Box::pin(async move {
query_as::<Sqlite, PageRow>(sql)
let page = query_as::<Sqlite, PageRow>(sql)
.bind(bytes)
.bind(from)
.bind(i64::try_from(PAGE).expect("PAGE fits an i64"))
.fetch_all(conn)
.await
.fetch_all(&mut *conn)
.await;
match page {
Err(_) if !logged(conn).await? => Ok(Vec::new()),
page => page,
}
})
})
.await
@@ -1298,7 +1320,7 @@ impl Reads {
/// Every aggregate the event log holds a row this module reads for, in the order the rows
/// were written. The kinds mirror [`EventRow::classify`]: an aggregate whose every row is
/// one this module ignores is no session to a consumer.
/// one this module ignores is no session to a consumer. None without an event log.
async fn aggregates(self: &Arc<Self>) -> Option<Vec<Aggregate>> {
const SQL: &str = "SELECT aggregate_id, min(rowid) AS first FROM event WHERE type LIKE \
'message.updated.%' OR type LIKE 'message.part.updated.%' OR type LIKE \
@@ -1306,7 +1328,12 @@ impl Reads {
'session.next.%' GROUP BY aggregate_id ORDER BY first ASC";
let rows: Vec<AggregateRow> = self
.read("aggregate scan", |conn| {
Box::pin(async move { query_as::<Sqlite, AggregateRow>(SQL).fetch_all(conn).await })
Box::pin(async move {
if !logged(conn).await? {
return Ok(Vec::new());
}
query_as::<Sqlite, AggregateRow>(SQL).fetch_all(conn).await
})
})
.await?;
Some(rows.into_iter().filter_map(|row| row.aggregate).collect())
@@ -1513,7 +1540,7 @@ impl Session for OpencodeSession {
match reads.backlog(&aggregate).await {
Some(backlog) => {
if let Some(message) = backlog.session.and_then(|info| {
OpencodeMessage::session(serde_json::json!({"info": info}))
OpencodeMessage::session(serde_json::json!({"info": info}), false)
}) {
yield Ok(message);
}
@@ -1615,6 +1642,17 @@ impl Session for OpencodeSession {
}
}
/// Whether the database has an `event` table: not before opencode creates its tables, nor
/// mid-migration.
async fn logged(conn: &mut SqliteConnection) -> Result<bool, sqlx::Error> {
query_scalar::<Sqlite, i64>(
"SELECT 1 FROM sqlite_master WHERE type = 'table' AND name = 'event'",
)
.fetch_optional(conn)
.await
.map(|found| found.is_some())
}
/// Whether every one of `tables` is there to read. opencode has kept its `session`, `message` and
/// `part` projections since it moved to SQLite, before it kept an event log, so a database with
/// an `event` table has them; one without is a database this module was not given by opencode,
@@ -2081,6 +2119,18 @@ pub(crate) fn tool_patch(call: &ToolCallId, metadata: &Value) -> Option<Patch> {
})
}
/// The id of session `session`'s info row while its title is `title`.
pub(crate) fn title_id(session: &str, title: &str) -> String {
format!("{session}:title:{:016x}", xxhash_rust::xxh3::xxh3_64(title.as_bytes()))
}
/// Whether a completed tool call ran a command that failed: opencode's shell tool completes
/// whatever the command's exit code, which it keeps in the call's metadata (`exit`), and marks
/// one it killed for running too long (`timeout`).
fn exited_non_zero(metadata: &Value) -> bool {
metadata["exit"].as_i64().is_some_and(|exit| exit != 0) || metadata["timeout"] == true
}
/// A model call's usage from its token counts (a `step-finish` part's or a
/// `session.next.step.ended` row's `tokens`), `None` when there are none.
///
@@ -2155,12 +2205,22 @@ impl OpencodeMessage {
/// A session's `session.created.1` or `session.updated.1` payload, as a row carrying its
/// title and the session it was spawned from. `None` for a payload with no info.
fn session(mut data: Value) -> Option<Self> {
///
/// Stamped when it `changed`, its `time.updated` -- a `session.updated.1` row's own time, or
/// a rename seen while following the session -- so that a new title outranks the rows before
/// it. Otherwise, the session read as it stands, when the session was created: when it was
/// last updated moves with every prompt, and would place the row wherever it was last read.
fn session(mut data: Value, changed: bool) -> Option<Self> {
let info = data.get_mut("info").filter(|info| info.is_object()).map(Value::take)?;
let time = &info["time"];
let (created, updated) =
(epoch_millis(&info["time"]["created"]), epoch_millis(&info["time"]["updated"]));
Some(Self {
role: Role::System,
time: epoch_millis(&time["updated"]).or_else(|| epoch_millis(&time["created"])),
time: if changed {
updated.or(created)
} else {
created.or(updated)
},
body: Body::Session(info),
info: None,
})
@@ -2253,7 +2313,7 @@ impl OpencodeMessage {
content.push(Content::ToolResult(ToolResult {
call,
output: state["output"].clone(),
error: false,
error: exited_non_zero(&state["metadata"]),
}));
content.extend(patch.map(Content::Patch));
}
@@ -2333,15 +2393,16 @@ impl OpencodeMessage {
}
impl Message for OpencodeMessage {
/// A part's id; a failed message's id; for a session's info, its id and title, so that a
/// title is delivered once however many rows repeat it -- and one of opencode 2.0's sessions
/// has none until one is generated.
/// A part's id; a failed message's id; for a session's info, its id and a digest of its
/// title, so that a title is delivered once however many rows repeat it -- and one of
/// opencode 2.0's sessions has none until one is generated. A digest, not the title: an id is
/// never redacted, and a title says what the session was about.
fn id(&self) -> Option<MessageId> {
let id = match &self.body {
Body::Part(value) | Body::Raw(value) => value["id"].as_str()?.to_owned(),
Body::Failure { id, .. } | Body::Next(Next { id, .. }) => id.clone(),
Body::Session(info) => match info["title"].as_str() {
Some(title) => format!("{}:title:{title}", info["id"].as_str()?),
Some(title) => title_id(info["id"].as_str()?, title),
None => format!("{}:session", info["id"].as_str()?),
},
};
@@ -2534,6 +2595,28 @@ mod tests {
));
}
/// opencode's bash tool completes a command that failed, with its exit code in the metadata.
#[rstest]
#[case::succeeded(0, false)]
#[case::failed(1, true)]
fn a_command_that_exited_non_zero_is_an_error(#[case] exit: i64, #[case] error: bool) {
let m = part_message(serde_json::json!({
"type": "tool",
"callID": "call_1",
"tool": "bash",
"state": {
"status": "completed",
"input": {"command": "false"},
"output": "",
"metadata": {"output": "", "exit": exit, "description": "fail"},
},
}));
assert!(matches!(
m.content().as_slice(),
[Content::ToolUse(_), Content::ToolResult(r)] if r.error == error
));
}
#[rstest]
fn normalizes_an_errored_tool_part() {
let m = part_message(serde_json::json!({
@@ -2894,18 +2977,26 @@ mod tests {
#[rstest]
#[tokio::test]
async fn a_db_without_an_event_table_is_reported_as_not_found() {
async fn a_db_without_an_event_table_is_watched_until_it_has_one() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("opencode.db");
drop(Sqlite::builder(path.as_os_str()).open().await.unwrap());
// what SQLite leaves the moment opencode first opens its database
std::fs::write(&path, b"").unwrap();
let listener = OpencodeSessions::builder().db(path).build().listener().unwrap();
let listener = OpencodeSessions::builder().db(path.clone()).build().listener().unwrap();
let mut stream = listener.watch().boxed();
assert!(matches!(
stream.next().await,
Some(Err(WatchError::Io(err))) if err.kind() == io::ErrorKind::NotFound
));
assert!(stream.next().await.is_none());
// reported once, and then waited out: opencode creates its tables and writes a session
let db = event_db(&path).await;
insert_event(&db, "evt_1", "ses_1", "message.updated.1", &role_event("msg_1", "user"))
.await;
let offered = tokio::time::timeout(Duration::from_secs(10), stream.next())
.await
.expect("the session was not offered within 10s");
assert!(matches!(offered, Some(Ok(session)) if session.id().to_string() == "ses_1"));
}
#[rstest]
@@ -4063,10 +4154,13 @@ mod tests {
.collect()
.await;
let [session] = <[_; 1]>::try_from(sessions).unwrap();
query::<sqlx::Sqlite>("DROP TABLE event")
.execute(&mut *db.pool().acquire().await.unwrap())
.await
.unwrap();
// a log that is there but can't be read (a missing one has no rows to read)
for sql in ["DROP TABLE event", "CREATE TABLE event (id TEXT PRIMARY KEY)"] {
query::<sqlx::Sqlite>(sql)
.execute(&mut *db.pool().acquire().await.unwrap())
.await
.unwrap();
}
let read: Vec<_> = session.read().collect().await;
assert!(
@@ -4306,13 +4400,13 @@ mod tests {
let titles: Vec<(Option<MessageId>, String)> =
messages.iter().filter_map(|m| Some((m.id(), m.title()?.text?))).collect();
assert_eq!(titles, [
(Some(MessageId::from("ses_R:title:New session".to_owned())), "New session".into()),
(Some(MessageId::from(title_id("ses_R", "New session"))), "New session".into()),
(
Some(MessageId::from("ses_R:title:Fix the flaky test".to_owned())),
Some(MessageId::from(title_id("ses_R", "Fix the flaky test"))),
"Fix the flaky test".into()
),
]);
let session = find(&messages, "ses_R:title:New session");
let session = find(&messages, &title_id("ses_R", "New session"));
assert_eq!(session.role(), Role::System);
assert!(session.content().is_empty());
assert_eq!(session.cwd(), Some(PathBuf::from("/work/proj")));
@@ -4586,14 +4680,15 @@ mod tests {
"msg_0d149b2d5001gkfmc9bUMDXLJ5",
// from the log, which opencode continued after the wipe
"prt_0d154c3f0001VPSmfCQt4r6LWN",
"ses_W:title:Mock title number 1",
title_id("ses_W", "Mock title number 1").as_str(),
"prt_0d154c996001Lr5h4i3jpsGN60",
"prt_0d154c99a001dtKmLh8SYe2pad",
"prt_0d154c9a5001gvwUYJAY42itXy",
]);
let title = find(&messages, "ses_W:title:Mock title number 1");
let title = find(&messages, &title_id("ses_W", "Mock title number 1"));
assert_eq!(title.cwd(), Some(PathBuf::from("/work/proj")));
// when its title changed, after the rows before it, not when the session was created
assert_eq!(title.timestamp().unwrap().unix_timestamp(), 1_790_218_388);
let prompt = find(&messages, "prt_0d148d4bf001JaxA9TwaMgU5vE");
@@ -4688,6 +4783,11 @@ mod tests {
child.title().and_then(|t| t.text).as_deref(),
Some("mock subtask (@general subagent)")
);
// the session as it stands, at its start: when it was last updated moves with reads
assert_eq!(
child.timestamp().unwrap().unix_timestamp_nanos(),
1_790_217_674_726_000_000
);
}
/// A session whose log holds its `session.created` row is read from the log alone, even
@@ -4707,7 +4807,10 @@ mod tests {
let messages = read_all(&path).await;
let delivered = ids(&messages);
assert_eq!(delivered.first().map(String::as_str), Some("ses_W:title:New session"));
assert_eq!(
delivered.first().map(String::as_str),
Some(title_id("ses_W", "New session").as_str())
);
assert!(
!delivered.contains(&"prt_0d148d4bf001JaxA9TwaMgU5vE".to_owned()),
"the projection was read for a session the log holds from its start"
crates/atuin-common/src/harnesstools/opencode/session/projection.rs
+73 -5
View File
@@ -1,8 +1,9 @@
//! One session as opencode's projection holds it now (its `session`, `message` and `part`
//! tables), keyed as capture keys it: what `opencode --session` continues, for sync (see
//! [`crate::harnesstools::sync`]). The event log capture follows also holds what a revert
//! removed since; the projection does not.
//! tables, and opencode 2.0's `session_v2` and `session_message`), keyed as capture keys it: what
//! `opencode --session` continues, for sync (see [`crate::harnesstools::sync`]). The event log
//! capture follows also holds what a revert removed since; the projection does not.
use std::collections::HashMap;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use std::time::Duration;
@@ -11,7 +12,8 @@ use sqlx::{Connection, Row, Sqlite};
use xxhash_rust::xxh3::Xxh3;
use super::{
Aggregate, Message, PartProjection, Reader, Reads, SqliteConnectOptions, SqliteConnection, text,
Aggregate, Message, PartProjection, Reader, Reads, SqliteConnectOptions, SqliteConnection,
text, v2,
};
use crate::db::{query, query_as, query_scalar};
@@ -36,6 +38,13 @@ pub struct Projected {
/// When opencode last updated the session or any of its messages or parts (`time_updated`,
/// milliseconds since the Unix epoch); `None` when that can't be read.
pub updated: Option<i64>,
/// opencode 2.0 holds the session (in `session_v2`; a session it copied from 1.x too), whose
/// rows follow 1.x's in [`rows`](Self::rows) and [`messages`](Self::messages). No import
/// adds to a session 2.0 holds.
pub next: bool,
/// Rows the session holds merged into another's (an opencode 2 message atuin wrote from
/// several), each with the row it went into.
pub merged: HashMap<String, String>,
}
/// A message of the projection, as sync needs to know it.
@@ -60,7 +69,8 @@ impl Projected {
}
}
/// Session `id` in the database `db`; `None` when its `session` table has no such session.
/// Session `id` in the database `db`; `None` when neither its `session` table nor opencode 2.0's
/// `session_v2` has such a session.
pub async fn projected(db: &Path, id: &str) -> Result<Option<Projected>, sqlx::Error> {
let opts = SqliteConnectOptions::new()
.filename(db)
@@ -76,6 +86,62 @@ async fn read(
conn: &mut SqliteConnection,
db: &Path,
id: &str,
) -> Result<Option<Projected>, sqlx::Error> {
// A database only opencode 2.0 ever wrote has no 1.x tables.
let legacy = if super::projected(conn, &["session", "message", "part"]).await? {
legacy(conn, db, id).await?
} else {
None
};
let Some(held) = v2::held(conn, id).await? else {
return Ok(legacy);
};
let mut content = Xxh3::new();
let mut found = legacy.unwrap_or_else(|| Projected {
directory: PathBuf::new(),
created: held.created.unwrap_or_default(),
reverting: false,
rows: Vec::new(),
messages: Vec::new(),
content: 0,
updated: None,
next: false,
merged: HashMap::new(),
});
content.update(&found.content.to_le_bytes());
// where 2.0 has the session now, which a move changes
if let Some(directory) = held.directory {
found.directory = PathBuf::from(directory);
}
found.reverting |= held.reverting;
found.updated = found.updated.max(held.updated);
found.next = true;
for row in held.rows {
digest(&mut content, b'v', Some(&row.id), row.created, row.data.as_deref());
let assistant = row.kind == "assistant";
if assistant || row.kind == "user" {
found.messages.push(ProjectedMessage {
id: row.id,
created: row.created.unwrap_or_default(),
assistant,
parent: None,
completed: row.settled,
});
}
if let Some(into) = &row.delivered {
found.merged.extend(row.merged.into_iter().map(|merged| (merged, into.clone())));
}
found.rows.extend(row.delivered);
}
found.content = content.digest();
Ok(Some(found))
}
/// Session `id` as opencode 1.x's `session`, `message` and `part` tables hold it.
async fn legacy(
conn: &mut SqliteConnection,
db: &Path,
id: &str,
) -> Result<Option<Projected>, sqlx::Error> {
let Some(session) =
query::<Sqlite>("SELECT directory, time_created FROM session WHERE id = ?1")
@@ -159,6 +225,8 @@ async fn read(
messages,
content: content.digest(),
updated,
next: false,
merged: HashMap::new(),
}))
}
crates/atuin-common/src/harnesstools/opencode/session/v2.rs
+172 -14
View File
@@ -80,8 +80,8 @@ use sqlx::{Connection, Row, Sqlite};
use xxhash_rust::xxh3::xxh3_64;
use super::{
Body, MessageInfo, Next, OpencodeMessage, POLL, Reads, epoch_millis, file_identity, json, text,
tokens, usage_of,
Body, MessageInfo, Next, OpencodeMessage, POLL, Reads, SOURCE, epoch_millis, exited_non_zero,
file_identity, json, text, tokens, tool_patch, usage_of,
};
use crate::db::sqlite::TransientResultCode;
use crate::db::sqlite::observe::{ObserveError, ReplayBehavior};
@@ -89,7 +89,7 @@ use crate::db::{query_as, query_scalar};
use crate::harnesstools::session::model::{
Content, Role, StopReason, ToolCallId, ToolResult, ToolUse,
};
use crate::harnesstools::session::{Checkpoint, MessageError};
use crate::harnesstools::session::{Checkpoint, Message as _, MessageError};
use crate::os::fs::FdIdentity;
/// The first line of the prompt a subagent session starts with (`tool/plugin/subagent.ts`,
@@ -245,7 +245,7 @@ impl Reader {
self.delivered.insert(seq);
} else {
self.mark = Some(Mark::of(seq, id));
self.delivered = self.delivered.split_off(&(seq + 1));
self.delivered.retain(|&delivered| delivered > seq);
}
}
}
@@ -303,7 +303,13 @@ impl Pass<'_> {
self.state = state;
self.stage = Stage::Rows;
if reader.info.as_ref().is_none_or(|said| !said.says(&info)) {
let message = info.message();
// A title read with the session from its start is the one it has; one
// read past rows delivered before may be a rename since.
let renamed = match &reader.info {
Some(said) => said.title != info.title,
None => reader.mark.is_some(),
};
let message = info.message(renamed);
reader.info = Some(info);
return Some((reader.mark, Ok(message)));
}
@@ -458,8 +464,9 @@ impl SessionInfo {
/// The session info as the message a 1.x `session.updated.1` row would be: its title, its
/// directory, and the session it came from -- a subagent's parent, or the session a fork
/// copied.
fn message(&self) -> OpencodeMessage {
/// copied. Stamped when the session was last updated if it was `renamed` since it was last
/// read, else when it was created (see `OpencodeMessage::session`).
fn message(&self, renamed: bool) -> OpencodeMessage {
let mut info = serde_json::json!({
"id": self.id,
"title": self.title,
@@ -473,7 +480,7 @@ impl SessionInfo {
} else if let Some(fork) = &self.fork {
info["forkedFrom"] = Value::from(fork.as_str());
}
OpencodeMessage::session(serde_json::json!({ "info": info }))
OpencodeMessage::session(serde_json::json!({ "info": info }), renamed)
.expect("the info is an object")
}
}
@@ -587,6 +594,7 @@ impl Full {
return Some(Err(MessageError::from(err)));
}
};
let session = info.id.as_str();
let time = epoch_millis(&data["time"]["created"]).or(self.created);
let said = |role: Role, content: Vec<Content>| Next {
role,
@@ -596,7 +604,7 @@ impl Full {
stop: None,
};
let text = |key: &str| data[key].as_str().unwrap_or_default().to_owned();
let (next, info) = match self.kind.as_str() {
let (mut next, info) = match self.kind.as_str() {
"user" => {
let prompt = text("text");
let role = if prompt.starts_with(SUBAGENT_PROMPT) {
@@ -648,6 +656,10 @@ impl Full {
}));
}
};
// a message atuin wrote from rows of its own is read back as the first of them
if let Some(source) = written_from(&data, session)["source"].as_str() {
source.clone_into(&mut next.id);
}
Some(Ok(OpencodeMessage {
role: next.role.clone(),
body: Body::Next(next),
@@ -687,11 +699,15 @@ impl Full {
input: state["input"].clone(),
}));
match state["status"].as_str() {
Some("completed") => content.push(Content::ToolResult(ToolResult {
call,
output: state["content"].clone(),
error: false,
})),
Some("completed") => {
let patch = tool_patch(&call, &state["metadata"]);
content.push(Content::ToolResult(ToolResult {
call,
output: state["content"].clone(),
error: exited_non_zero(&state["metadata"]),
}));
content.extend(patch.map(Content::Patch));
}
Some("error") => content.push(Content::ToolResult(ToolResult {
call,
output: Value::from(error_text(&state["error"])),
@@ -742,6 +758,18 @@ impl Full {
}
}
/// What a message atuin wrote of session `session` says it was written from (see [`SOURCE`]),
/// `Null` for any other: opencode's own fork copies a message whole, `metadata` and all, into
/// another session, where it is a row of that session's own.
fn written_from<'a>(data: &'a Value, session: &str) -> &'a Value {
let written = &data["metadata"][SOURCE];
if written["session"] == session {
written
} else {
&Value::Null
}
}
/// The tables a read of the 2.0 layout looks for.
async fn tables(conn: &mut SqliteConnection) -> Result<Vec<String>, sqlx::Error> {
query_scalar::<Sqlite, String>(
@@ -752,6 +780,136 @@ async fn tables(conn: &mut SqliteConnection) -> Result<Vec<String>, sqlx::Error>
.await
}
/// A 2.0 session as opencode holds it now, for sync (see `projection`).
#[derive(Debug)]
pub(super) struct Held {
pub(super) directory: Option<String>,
pub(super) created: Option<i64>,
/// When opencode last updated the session or any of its rows.
pub(super) updated: Option<i64>,
/// A revert is pending.
pub(super) reverting: bool,
/// Its rows, in `seq` order.
pub(super) rows: Vec<HeldRow>,
}
/// A row of a [`Held`] session.
#[derive(Debug)]
pub(super) struct HeldRow {
pub(super) id: String,
pub(super) kind: String,
pub(super) created: Option<i64>,
pub(super) data: Option<String>,
/// opencode is done writing it (see [`Reads::rows`]).
pub(super) settled: bool,
/// The id capture delivers the row under; `None` for a row it delivers nothing of: one with
/// no message, one not settled yet, or a copy of a 1.x message.
pub(super) delivered: Option<String>,
/// The other rows atuin wrote the message from, which it holds without a row of their own
/// (see [`SOURCE`]).
pub(super) merged: Vec<String>,
}
/// Session `id` as `conn`'s `session_v2` and `session_message` tables hold it; `None` when they
/// hold no such session, or aren't there.
pub(super) async fn held(
conn: &mut SqliteConnection,
id: &str,
) -> Result<Option<Held>, sqlx::Error> {
const PLAIN: &str = "SELECT id, seq, type, time_created, data, 0 AS migrated FROM \
session_message WHERE session_id = ?1 ORDER BY seq";
// as `Reads::row` tells a copy of a 1.x message
const LEGACY: &str = "SELECT m.id, m.seq, m.type, m.time_created, m.data, EXISTS (SELECT 1 \
FROM message l WHERE l.session_id = m.session_id AND (l.id = m.id OR \
(m.type = 'synthetic' AND substr(l.id, 1, 16) = substr(m.id, 1, 16)))) \
AS migrated FROM session_message m WHERE m.session_id = ?1 ORDER BY \
m.seq";
let tables = tables(&mut *conn).await?;
let has = |name: &str| tables.iter().any(|table| table == name);
if !has("session_v2") || !has("session_message") {
return Ok(None);
}
let Some(info) = query_as::<Sqlite, SessionInfo>(
"SELECT id, title, parent_id, fork_session_id, directory, time_created, time_updated FROM \
session_v2 WHERE id = ?1",
)
.bind(id)
.fetch_optional(&mut *conn)
.await?
else {
return Ok(None);
};
let revert =
query_scalar::<Sqlite, Option<String>>("SELECT revert FROM session_v2 WHERE id = ?1")
.bind(id)
.fetch_one(&mut *conn)
.await?;
let updated = query_scalar::<Sqlite, Option<i64>>(
"SELECT MAX(t) FROM (SELECT time_updated AS t FROM session_v2 WHERE id = ?1 UNION ALL \
SELECT time_updated FROM session_message WHERE session_id = ?1)",
)
.bind(id)
.fetch_one(&mut *conn)
.await?;
let legacy = has("message");
let rows = query_as::<Sqlite, Full>(if legacy {
LEGACY
} else {
PLAIN
})
.bind(id)
.fetch_all(&mut *conn)
.await?;
let state = State {
info: Some(info.clone()),
legacy,
..State::default()
};
let rows = rows
.into_iter()
.map(|full| {
let parsed = full.data.as_deref().and_then(|data| json(data).ok()).unwrap_or_default();
let settled = match full.kind.as_str() {
"assistant" => !parsed["time"]["completed"].is_null(),
"compaction" => parsed["status"] != "running",
_ => true,
};
let (id, kind, created, data) =
(full.id.clone(), full.kind.clone(), full.created, full.data.clone());
let merged = written_from(&parsed, &info.id)["merged"]
.as_array()
.into_iter()
.flatten()
.filter_map(|id| id.as_str().map(str::to_owned))
.collect();
let delivered = if settled && !full.migrated {
match full.message(&info, &state) {
Some(Ok(message)) => message.id().map(String::from),
_ => None,
}
} else {
None
};
HeldRow {
id,
kind,
created,
data,
settled,
delivered,
merged,
}
})
.collect();
Ok(Some(Held {
directory: info.directory,
created: info.created,
updated,
reverting: revert.is_some_and(|r| !r.trim().is_empty() && r != "null"),
rows,
}))
}
impl Reads {
/// What a pass needs to know of `session` before reading its rows (see [`State`]), checking
/// `mark` on the way.
crates/atuin-common/src/harnesstools/opencode/session/v2/tests.rs
+173 -7
View File
@@ -11,7 +11,7 @@ use rstest::{fixture, rstest};
use serde_json::{Value, json};
use tempfile::TempDir;
use super::super::{OpencodeSession, OpencodeSessions};
use super::super::{OpencodeSession, OpencodeSessions, title_id};
use super::*;
use crate::db::query;
use crate::db::sqlite::Sqlite;
@@ -383,7 +383,14 @@ fn line(message: &OpencodeMessage) -> String {
})
.collect::<Vec<_>>()
.join(" | ");
let id = message.id().map(String::from).unwrap_or_default();
let mut id = message.id().map(String::from).unwrap_or_default();
// a title's row is keyed on a digest of it, shown here as the title it digests
if let Some(title) = message.title().and_then(|title| title.text)
&& let Some((session, _)) = id.split_once(":title:")
{
assert_eq!(id, title_id(session, &title));
id = format!("{session}:title:{title}");
}
format!("{:?} {id}: {content}", message.role())
}
@@ -583,6 +590,55 @@ fn a_call_s_content_is_its_text_reasoning_and_tools() {
));
}
/// A shell call opencode 2.0.24 completed whose command failed: `completed`, its exit code only
/// in the metadata, or, killed for running too long, `timeout` and no exit code.
#[rstest]
#[case::succeeded(json!({"status": "completed", "truncated": false, "exit": 0}), false)]
#[case::failed(json!({"status": "completed", "truncated": false, "exit": 2}), true)]
#[case::timed_out(
json!({"status": "completed", "truncated": false, "timeout": true, "signal": "SIGTERM"}),
true
)]
fn a_shell_call_that_exited_non_zero_is_an_error(#[case] metadata: Value, #[case] error: bool) {
let mut data = assistant("", true);
data["content"] = json!([{"type": "tool", "id": "call_1", "name": "shell", "state": {
"status": "completed", "input": {"command": "ls /nonexistent"},
"content": [{"type": "text", "text": "ls: cannot access"}],
"metadata": metadata}}]);
let message = normalised("assistant", &data).unwrap();
assert!(matches!(
message.content().as_slice(),
[Content::ToolUse(_), Content::ToolResult(result)] if result.error == error
));
}
/// An edit opencode 2.0.24 made keeps its diff in the call's metadata (`files[]`:
/// `{file, patch, status}`), which is the call's patch.
#[rstest]
fn an_edit_s_diff_is_its_patch() {
let mut data = assistant("", true);
data["content"] = json!([{"type": "tool", "id": "call_1", "name": "edit", "state": {
"status": "completed",
"input": {"path": "greet.py", "oldString": "\"hello \"", "newString": "\"hi \""},
"content": [{"type": "text", "text": "Edited greet.py (1 replacement)"}],
"metadata": {"files": [{"file": "greet.py", "patch": "Index: greet.py\n===\n--- \
greet.py\n+++ greet.py\n@@ -1,2 +1,2 @@\n def greet(name):\n- return \"hello \" + \
name\n+ return \"hi \" + name\n", "status": "modified", "additions": 1,
"deletions": 1}], "truncated": false}}}]);
let message = normalised("assistant", &data).unwrap();
let content = message.content();
let [Content::ToolUse(_), Content::ToolResult(result), Content::Patch(patch)] =
content.as_slice()
else {
panic!("no patch: {content:?}");
};
assert!(!result.error);
assert_eq!(patch.call.as_ref(), "call_1");
assert_eq!(patch.files.len(), 1);
assert_eq!(patch.files[0].path, "greet.py");
assert_eq!((patch.files[0].additions(), patch.files[0].deletions()), (1, 1));
}
/// opencode 2.0's `input` leaves out the cached tokens and its `output` the reasoning ones
/// (`SessionUsage.tokens`: `nonCachedInputTokens`, `visibleOutputTokens`); reasoning is billed as
/// output (`calculateCost`), so it is added back.
@@ -647,20 +703,61 @@ fn a_row_was_written_where_the_session_was_then(#[case] seq: i64, #[case] cwd: &
#[rstest]
fn a_session_s_info_names_its_title_and_where_it_came_from() {
let mut info = state().info.unwrap();
let untitled = info.message();
let untitled = info.message(false);
assert_eq!(untitled.id().map(String::from).as_deref(), Some("ses_1:session"));
assert_eq!(untitled.title(), None);
assert_eq!(untitled.parent_session(), None);
info.title = Some("Fix it".to_owned());
info.fork = Some("ses_0".to_owned());
let titled = info.message();
assert_eq!(titled.id().map(String::from).as_deref(), Some("ses_1:title:Fix it"));
let titled = info.message(false);
// keyed on the title, which an id never says: ids are synced as they are, unredacted
let id = titled.id().map(String::from).unwrap();
assert_eq!(id, title_id("ses_1", "Fix it"));
assert!(!id.contains("Fix it"), "{id}");
assert_eq!(titled.title().and_then(|title| title.text).as_deref(), Some("Fix it"));
assert_eq!(titled.parent_session().map(String::from).as_deref(), Some("ses_0"));
assert_eq!(titled.parent_kind(), Some(ParentKind::Fork));
assert_eq!(titled.cwd(), Some(PathBuf::from("/work/proj")));
}
/// The session's info, read as it stands, is stamped when the session was created, not when
/// opencode last touched it (which every prompt bumps), so it sits at the session's start however
/// late it is read. A rename is stamped when opencode last touched the session, which the rename
/// did, so that it outranks the rows before it.
#[rstest]
#[case::as_it_stands(false, 1_000)]
#[case::renamed(true, 9_000)]
fn a_session_s_info_is_stamped_when_its_title_changed(#[case] renamed: bool, #[case] stamp: i128) {
let mut info = state().info.unwrap();
info.title = Some("Fix it".to_owned());
info.created = Some(1_000);
info.updated = Some(9_000);
let stamped = info.message(renamed).timestamp().unwrap();
assert_eq!(stamped.unix_timestamp_nanos(), stamp * 1_000_000);
}
/// A renamed session read from scratch says its title as of its start, however much later it is
/// read and whatever has happened in it since: two reads say the same.
#[rstest]
#[tokio::test]
async fn a_renamed_session_read_from_scratch_is_titled_at_its_start(#[future] db: Db) {
let db = db.await;
db.session("ses_1", Some("Old"), None).await;
db.append("ses_1", "msg_u", "user", user("hi")).await;
db.execute("UPDATE session_v2 SET title = 'Named', time_updated = 3000").await;
let first = read_all(&db.path).await;
db.append("ses_1", "msg_u2", "user", user("again")).await;
db.execute("UPDATE session_v2 SET time_updated = 5000").await;
let second = read_all(&db.path).await;
let title = |messages: &[(String, OpencodeMessage)]| {
let (_, title) = &messages[0];
(title.id().map(String::from), title.timestamp().unwrap().unix_timestamp_nanos())
};
assert_eq!(title(&first), (Some(title_id("ses_1", "Named")), 1_000_000_000));
assert_eq!(title(&second), title(&first));
}
proptest! {
/// Both places come back out of the checkpoint they were packed into.
#[test]
@@ -880,7 +977,8 @@ async fn replay_from_now_skips_preexisting_rows(#[future] db: Db) {
]);
}
/// A session opencode 2.0 renames is delivered its new title.
/// A session opencode 2.0 renames is delivered its new title, stamped when it was renamed: after
/// the rows before it, so that it is the session's title from then on.
#[rstest]
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn a_rename_is_delivered(#[future] db: Db) {
@@ -893,10 +991,12 @@ async fn a_rename_is_delivered(#[future] db: Db) {
db.execute("UPDATE session_v2 SET time_updated = 2000").await;
db.bump("ses_1").await;
quiet(&mut stream).await;
db.execute("UPDATE session_v2 SET title = 'Named'").await;
// as `session.renamed` projects: the title, and `time_updated` to when it was renamed
db.execute("UPDATE session_v2 SET title = 'Named', time_updated = 3000").await;
db.bump("ses_1").await;
let renamed = next_n(&mut stream, &stored, 1).await;
assert_eq!(renamed[0].1.title().and_then(|title| title.text).as_deref(), Some("Named"));
assert_eq!(renamed[0].1.timestamp().unwrap().unix_timestamp_nanos(), 3_000_000_000);
}
/// A database only opencode 2.0 ever wrote: no 1.x tables at all.
@@ -911,6 +1011,27 @@ async fn a_2_0_only_database_is_read_whole(#[future] db: Db) {
assert_eq!(lines(&read_all(&db.path).await), ["System ses_1:title:T: ", "User msg_u: hi"]);
}
/// A 2.0 database without the `event` table 1.x tailed (a later opencode that drops it) is read,
/// listed and watched all the same.
#[rstest]
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn a_database_without_an_event_log_is_captured(#[future] db: Db) {
let db = db.await;
db.execute("DROP TABLE event").await;
db.session("ses_1", Some("T"), None).await;
db.append("ses_1", "msg_u", "user", user("hi")).await;
assert_eq!(lines(&read_all(&db.path).await), ["System ses_1:title:T: ", "User msg_u: hi"]);
let stored = Stored::default();
let mut stream = events(&db.path, ReplayBehavior::All, &stored);
// the missing log is reported, once, and its sessions captured all the same
assert!(matches!(stream.next().await, Some(Err(CaptureError::Watch(_)))));
let got = next_n(&mut stream, &stored, 2).await;
assert_eq!(lines(&got), ["System ses_1:title:T: ", "User msg_u: hi"]);
db.append("ses_1", "msg_v", "user", user("again")).await;
assert_eq!(lines(&next_n(&mut stream, &stored, 1).await), ["User msg_v: again"]);
}
/// The whole live path against a database being written: two sessions, each prompting and
/// streaming calls that are rewritten many times over, one steering a prompt in mid-call. Every
/// row is delivered exactly once, as it settled.
@@ -1225,3 +1346,48 @@ async fn a_session_in_both_layouts_is_one_session(#[future] mixed: Db) {
]);
quiet(&mut stream).await;
}
/// Sync reads a session 2.0 copied from 1.x and went on with as capture does: every row capture
/// delivers of it, its 1.x parts and its 2.0 messages, and none of 2.0's copies; the tip is the
/// last 2.0 message.
#[rstest]
#[tokio::test]
async fn a_migrated_session_s_tip_is_its_last_2_0_row(#[future] mixed: Db) {
let db = mixed.await;
db.load(&fixture_json()["final"]).await;
let delivered: Vec<String> = read_all(&db.path)
.await
.into_iter()
.filter(|(session, _)| session == MIGRATED)
.filter_map(|(_, m)| m.id().map(String::from))
.filter(|id| !id.starts_with(&format!("{MIGRATED}:")))
.collect();
let held = super::super::projection::projected(&db.path, MIGRATED).await.unwrap().unwrap();
assert!(held.next);
let set = |ids: &[String]| ids.iter().cloned().collect::<HashSet<_>>();
assert_eq!(set(&held.rows), set(&delivered));
assert_eq!(held.rows.len(), delivered.len());
let last_2_0 = delivered.iter().rev().find(|id| id.starts_with("msg_"));
assert_eq!(held.rows.last(), last_2_0);
}
/// A row at the very end of the sequence is read like any other, without the sequence after it
/// overflowing.
#[rstest]
#[tokio::test]
async fn a_row_at_the_last_seq_is_read(#[future] db: Db) {
let db = db.await;
db.session("ses_1", Some("T"), None).await;
db.append("ses_1", "msg_u", "user", user("hi")).await;
db.execute(
"INSERT INTO session_message (id, session_id, type, seq, time_created, time_updated, \
data) VALUES ('msg_max', 'ses_1', 'user', 9223372036854775807, 1000, 1000, \
'{\"time\":{\"created\":1000},\"text\":\"last\",\"files\":[]}')",
)
.await;
assert_eq!(lines(&read_all(&db.path).await), [
"System ses_1:title:T: ",
"User msg_u: hi",
"User msg_max: last",
]);
}
crates/atuin-common/src/harnesstools/session/mod.rs
+2 -1
View File
@@ -13,7 +13,8 @@ pub use error::{CaptureError, MessageError, RuntimeError, WatchError};
use futures::{Stream, StreamExt, TryStreamExt};
pub use model::{
Content, MessageId, ParentKind, Role, SessionEvent, SessionId, StopReason, TitleChange,
TitleSource, ToolCallId, ToolResult, ToolUse, Usage, is_substantive, is_tool_output,
TitleSource, ToolCallId, ToolResult, ToolUse, UNREADABLE_SUMMARY, Usage, is_substantive,
is_tool_output,
};
pub use patch::{Change, FilePatch, Hunk, Patch};
use time::OffsetDateTime;
crates/atuin-common/src/harnesstools/session/model.rs
+29 -2
View File
@@ -48,6 +48,12 @@ pub enum Content {
Patch(Patch),
}
/// The [`Content::Summary`] of a compaction whose summary can't be read (Codex's, compacted by
/// the API, which keeps it encrypted): the conversation before it is no longer what the model
/// sees, and no summary says what it was.
pub const UNREADABLE_SUMMARY: &str =
"[the conversation before this was compacted; its summary can't be read]";
/// Human-readable breadcrumb shared by the transcript and `atuin ai session` rendering.
#[must_use]
pub fn reasoning_label(tokens: Option<u64>) -> String {
@@ -93,13 +99,26 @@ pub struct ToolResult {
}
impl ToolResult {
/// The output as display text: a JSON string's contents verbatim, any other JSON as encoded,
/// or `None` when capture did not keep it.
/// The output as display text: a JSON string's contents verbatim, a list of content blocks
/// as their text (Claude Code's and pi's `text`, Codex's `input_text`; a block without text
/// as encoded), any other JSON as encoded, or `None` when capture did not keep it.
#[must_use]
pub fn output_text(&self) -> Option<Cow<'_, str>> {
match &self.output {
serde_json::Value::Null => None,
serde_json::Value::String(text) => Some(Cow::Borrowed(text)),
serde_json::Value::Array(blocks) => Some(Cow::Owned(
blocks
.iter()
.map(|block| match block {
serde_json::Value::String(text) => text.clone(),
_ => {
block["text"].as_str().map_or_else(|| block.to_string(), str::to_owned)
}
})
.collect::<Vec<_>>()
.join("\n"),
)),
json => Some(Cow::Owned(json.to_string())),
}
}
@@ -266,6 +285,14 @@ mod tests {
#[rstest]
#[case::string(serde_json::json!("line1\nline2"), Some("line1\nline2"))]
#[case::structured(serde_json::json!({"exit": 0}), Some(r#"{"exit":0}"#))]
#[case::blocks(
serde_json::json!([{"type": "text", "text": "a"}, {"type": "input_text", "text": "b"}]),
Some("a\nb")
)]
#[case::blocks_without_text(
serde_json::json!([{"type": "text", "text": "a"}, {"type": "image"}]),
Some("a\n{\"type\":\"image\"}")
)]
#[case::uncaptured(serde_json::Value::Null, None)]
fn tool_result_output_text_unwraps_json_strings(
#[case] output: serde_json::Value,
crates/atuin-common/src/harnesstools/session/patch.rs
+26 -9
View File
@@ -391,13 +391,17 @@ fn hunk_header(line: &str) -> Option<Hunk> {
/// `hunks` as pi's `edit` tool writes its `details.diff` (`generateDiffString`): each line its
/// prefix, its line number padded to the widest, a space and its text (`+` lines numbered in the
/// new file, the rest in the old), with `...` for the lines skipped before and between hunks. Also
/// the first changed line, numbered in the new file (`firstChangedLine`).
/// the first changed line, numbered in the new file (`firstChangedLine`). Numbers past the last a
/// line number holds (a malformed header's) stay at it.
#[must_use]
pub fn pi_diff(hunks: &[Hunk]) -> (String, Option<u64>) {
let width = hunks
.iter()
// The last line each side shows.
.map(|h| (h.old_start + h.old_lines).max(h.new_start + h.new_lines).saturating_sub(1))
.map(|h| {
let old = h.old_start.saturating_add(h.old_lines);
old.max(h.new_start.saturating_add(h.new_lines)).saturating_sub(1)
})
.max()
.unwrap_or(0)
.to_string()
@@ -414,18 +418,21 @@ pub fn pi_diff(hunks: &[Hunk]) -> (String, Option<u64>) {
let number = match prefix {
"+" => {
first.get_or_insert(new);
new += 1;
new - 1
let number = new;
new = new.saturating_add(1);
number
}
"-" => {
first.get_or_insert(new);
old += 1;
old - 1
let number = old;
old = old.saturating_add(1);
number
}
" " => {
old += 1;
new += 1;
old - 1
let number = old;
old = old.saturating_add(1);
new = new.saturating_add(1);
number
}
_ => continue,
};
@@ -618,6 +625,16 @@ mod tests {
assert_eq!(first, Some(2));
}
/// A hunk header numbering lines past what a line number holds (a malformed or crafted one)
/// writes the last number there is rather than overflow.
#[rstest]
fn writes_pi_diffs_of_hunks_numbered_past_the_last_line() {
let hunks = vec![hunk((u64::MAX, u64::MAX), (u64::MAX, 2), &[" a", "-b", "+B", " c"])];
let (diff, first) = pi_diff(&hunks);
assert_eq!(diff.lines().count(), 5, "{diff}");
assert_eq!(first, Some(u64::MAX));
}
fn line() -> impl Strategy<Value = String> {
// Text a line could hold, header-like starts included.
prop_oneof![Just("--- x".to_owned()), Just("+++ y".to_owned()), "[ -~]{0,12}"]
crates/atuin-common/src/harnesstools/sync/ccode.rs
+30 -13
View File
@@ -630,8 +630,9 @@ pub(super) fn replace_to(
/// all but those that went on from `from` here. Those are the lines descending from it (by
/// `parentUuid`, or across a compaction boundary `logicalParentUuid`), and the lines that name
/// one of them (a `last-prompt`'s or a summary's `leafUuid`, a file history snapshot's
/// `messageId`); and any other line after the first of them, but for title lines
/// ([`TITLE_LINES`]), kept wherever they are.
/// `messageId`); a `queue-operation`, which copies the prompt it queued, going with the next line
/// that has an id (usually that prompt's); and any other line after the first of them, but for
/// title lines ([`TITLE_LINES`]), kept wherever they are.
fn shared_lines(bytes: &[u8], from: &str) -> Vec<u8> {
let parsed: Vec<Option<Value>> = jsonl_lines(bytes).map(parse).collect();
let mut parents: HashMap<&str, &str> = HashMap::new();
@@ -643,18 +644,34 @@ fn shared_lines(bytes: &[u8], from: &str) -> Vec<u8> {
}
let mut memo = HashMap::new();
let uuids: HashSet<&str> = parsed.iter().flatten().filter_map(|l| l["uuid"].as_str()).collect();
let mut lines = parsed.iter();
kept_lines(bytes, |_| {
let line = lines.next().and_then(Option::as_ref)?;
if TITLE_LINES.contains(&line["type"].as_str().unwrap_or_default()) {
return Some(true);
let mut keep: Vec<Option<bool>> = parsed
.iter()
.map(|line| {
let line = line.as_ref()?;
if TITLE_LINES.contains(&line["type"].as_str().unwrap_or_default()) {
return Some(true);
}
if let Some(uuid) = line["uuid"].as_str() {
return Some(!descends(&parents, &mut memo, uuid, from));
}
let named = line["leafUuid"].as_str().or_else(|| line["messageId"].as_str())?;
uuids.contains(named).then(|| !descends(&parents, &mut memo, named, from))
})
.collect();
// From the end, so each queue operation knows the next line with an id.
let mut next = None;
for (line, keep) in parsed.iter().zip(&mut keep).rev() {
let Some(line) = line else {
continue;
};
if line["uuid"].is_string() {
next = *keep;
} else if line["type"] == "queue-operation" && keep.is_none() {
*keep = next;
}
if let Some(uuid) = line["uuid"].as_str() {
return Some(!descends(&parents, &mut memo, uuid, from));
}
let named = line["leafUuid"].as_str().or_else(|| line["messageId"].as_str())?;
uuids.contains(named).then(|| !descends(&parents, &mut memo, named, from))
})
}
let mut keep = keep.into_iter();
kept_lines(bytes, |_| keep.next().flatten())
}
/// The directory the transcript last worked in.
crates/atuin-common/src/harnesstools/sync/ccode/tests.rs
+6 -1
View File
@@ -431,6 +431,11 @@ fn went_its_own_way() -> Vec<Value> {
a1,
r1,
json!({"type": "custom-title", "customTitle": "mine", "sessionId": SESSION}),
// Claude Code's queue bookkeeping for the prompt that follows, which it copies.
json!({"type": "queue-operation", "operation": "enqueue", "sessionId": SESSION,
"timestamp": "2026-09-18T10:00:04Z", "content": "queued here alone"}),
json!({"type": "queue-operation", "operation": "dequeue", "sessionId": SESSION,
"timestamp": "2026-09-18T10:00:04Z"}),
json!({"type": "file-history-snapshot", "messageId": "u2", "snapshot": {"files": {}}}),
user("u2", Some("r1"), 4),
assistant("a2", "u2", 5),
@@ -483,7 +488,7 @@ fn a_switch_keeps_the_shared_history_as_it_is_and_backs_the_copy_up(dir: TempDir
let written: Vec<Value> = lines[4..].iter().filter_map(|l| parse(l)).collect();
assert_eq!(written[0]["parentUuid"], "r1", "the branch hangs from where it left the copy");
assert!(written.iter().all(|l| l["sessionId"] == SESSION));
for gone in ["u2", "a2", "last-prompt", "file-history-snapshot"] {
for gone in ["u2", "a2", "last-prompt", "file-history-snapshot", "queue-operation"] {
let found = lines.iter().any(|l| memchr::memmem::find(l, gone.as_bytes()).is_some());
assert!(!found, "{gone} went on from where the branch left the copy");
}
crates/atuin-common/src/harnesstools/sync/codex.rs
+23 -18
View File
@@ -48,7 +48,7 @@ use super::{
capture_key, check_segment, continuing_from, jsonl_lines, merged_in, modified, past,
replace_jsonl, require_idle, resumes_from, switched_to,
};
use crate::harnesstools::codex::session::{CodexSession, default_root, resume_id, session_id_of};
use crate::harnesstools::codex::session::{CodexSession, resume_id, session_id_of};
use crate::harnesstools::codex::{Codex, rehydrate, state_db};
use crate::harnesstools::rehydrate::{RehydrateMessage, RehydrateSession};
use crate::harnesstools::resume::is_plain_name;
@@ -57,11 +57,12 @@ use crate::sync::BlockingPool;
impl SessionSync for Codex {
async fn local_tip(&self, id: &str) -> Result<Option<LocalTip>, SyncError> {
local_tip_in(&codex_home(), id).await
local_tip_in(&crate::harnesstools::codex::home(), id).await
}
async fn is_live(&self, id: &str, cwd: Option<&Path>) -> Liveness {
let (home, id, cwd) = (codex_home(), id.to_owned(), cwd.map(Path::to_path_buf));
let (home, id, cwd) =
(crate::harnesstools::codex::home(), id.to_owned(), cwd.map(Path::to_path_buf));
blocking_liveness(move || {
let thread = resume_id(&id);
let newest = is_plain_name(thread)
@@ -84,7 +85,7 @@ impl SessionSync for Codex {
lines: &[RehydrateMessage],
options: &AppendOptions<'_>,
) -> Result<AppendOutcome, SyncError> {
let (home, thread) = (codex_home(), resume_id(id).to_owned());
let (home, thread) = (crate::harnesstools::codex::home(), resume_id(id).to_owned());
let seen = appending(base, &thread, options.cwd);
let (at, of) = (home.clone(), id.to_owned());
require_idle(
@@ -102,7 +103,7 @@ impl SessionSync for Codex {
options: &AppendOptions<'_>,
backups: &Path,
) -> Result<ReplaceOutcome, SyncError> {
let (home, thread) = (codex_home(), resume_id(id).to_owned());
let (home, thread) = (crate::harnesstools::codex::home(), resume_id(id).to_owned());
let seen = appending(base, &thread, options.cwd);
let (at, of) = (home.clone(), id.to_owned());
require_idle(
@@ -112,12 +113,6 @@ impl SessionSync for Codex {
}
}
/// `$CODEX_HOME`, else `~/.codex`: where `sessions/` is.
fn codex_home() -> PathBuf {
let root = default_root();
root.parent().map_or_else(|| root.clone(), Path::to_path_buf)
}
/// Where Codex keeps its writer locks (codex-rs `WRITER_LOCK_DIR`).
const LOCKS: &str = "thread-writer-locks";
@@ -499,8 +494,18 @@ pub(super) async fn replace_in(
taken: taken.cloned(),
backups: backups.to_path_buf(),
};
let home = home.to_path_buf();
blocking(move || replace_rollout(&home, &switch, &rollout.rows)).await
let thread = switch.thread.clone();
let (rewritten, lock) = {
let home = home.to_path_buf();
blocking(move || replace_rollout(&home, &switch, &rollout.rows)).await?
};
// Projected from the lines the switch cut, Codex's history would show the branch left: it is
// forgotten before any Codex can load the thread again.
if let Err(err) = state_db::forget_history(home, &thread).await {
tracing::warn!(%err, thread, "Codex may go on showing the branch switched from");
}
drop(lock);
Ok(rewritten)
}
/// Why the rollout holding `bytes` can't be written out again along another branch, whatever the
@@ -545,12 +550,12 @@ pub(super) struct Switch {
/// what follows that line is what went on from it here), its `session_meta` and every line
/// before as they are; then the branch's rows past it, numbered on from the line kept last, as
/// [`append_in`] writes them. Codex continues from the last line, so the branch's head must be
/// the last row written, or merged into it.
/// the last row written, or merged into it. Hands back the thread's writer lock, still held.
pub(super) fn replace_rollout(
home: &Path,
switch: &Switch,
rows: &[(u64, String)],
) -> Result<ReplaceOutcome, SyncError> {
) -> Result<(ReplaceOutcome, Option<WriterLock>), SyncError> {
let Switch {
thread,
base,
@@ -608,14 +613,14 @@ pub(super) fn replace_rollout(
let merged = merged_in(&out, &tip.iter().cloned().collect());
resumes_from(tip.as_deref(), &merged, &end)?;
let (backup, warning) = replace_jsonl(path, base.stamp, &bytes, &out, backups)?;
drop(lock);
Ok(ReplaceOutcome {
let outcome = ReplaceOutcome {
native_path: path.clone(),
appended,
tip_source_id: tip,
backup,
warning,
})
};
Ok((outcome, lock))
}
/// The rollout `bytes`, up to the end of the line ending at `end` (and never short of its first
crates/atuin-common/src/harnesstools/sync/codex/tests.rs
+5 -1
View File
@@ -213,7 +213,8 @@ async fn their_branch(cut: usize) -> (Vec<RehydrateMessage>, String) {
/// content), drops what it went on with here, and appends the branch's lines past it, numbered on
/// from the line kept last; Codex resumes from the branch's head, and capture reads back nothing
/// the branch doesn't hold. In place, under the thread's writer lock where Codex keeps them, with
/// the rollout as it was kept whole outside Codex's sessions.
/// the rollout as it was kept whole outside Codex's sessions. Codex's history of the thread,
/// projected from the lines the switch cut, is forgotten.
#[rstest]
#[case::without_locks(false)]
#[case::under_the_writer_lock(true)]
@@ -229,6 +230,7 @@ async fn a_switch_writes_the_rollout_out_along_the_branch_in_place(
let (thread, path, base, _) = behind(dir.path(), 45).await;
let before = std::fs::read_to_string(&path).unwrap();
let (rows, head) = their_branch(30).await;
let history = state_db::tests::history(dir.path(), &[&thread]).await;
let outcome =
replace_in(dir.path(), &thread, &base, &rows, None, backups.path()).await.unwrap();
@@ -255,6 +257,8 @@ async fn a_switch_writes_the_rollout_out_along_the_branch_in_place(
if locks {
assert!(!dir.path().join(LOCKS).join(format!("{thread}.lock")).exists(), "let go");
}
let projected = state_db::tests::projected(&history, &thread).await;
assert_eq!(projected, 0, "Codex projects its history again, from the branch");
}
/// Nothing is written to a thread kept in several rollouts (reverted, or named by a segment of
crates/atuin-common/src/harnesstools/sync/opencode.rs
+15 -8
View File
@@ -15,7 +15,9 @@
//! and it is no fast-forward ([`SyncError::Unsupported`]: a copy that stopped mid-reply is forked).
//! Neither is a message that would land before the local last one (opencode would read it in among
//! the others by creation time and id), nor a session with a revert pending, which would remove
//! what is appended.
//! what is appended. Nor is anything appended to a session opencode 2.0 holds (its `session_v2`
//! and `session_message` rows, read after 1.x's for the tip): its `opencode session import` only
//! writes a session it doesn't hold yet. A copy with nothing to catch up resumes as it is.
//!
//! **Liveness.** opencode records no running session anywhere: any opencode running here in the
//! session's directory (or below it, or where that can't be told) may be writing it; so may one
@@ -26,7 +28,7 @@
//! reply alone is not: a copy restore or catch-up wrote from a reply synced half-way keeps it
//! unfinished for good.
use std::collections::{HashMap, HashSet};
use std::collections::HashSet;
use std::future::Future;
use std::path::{Path, PathBuf};
use std::time::{Duration, SystemTime};
@@ -40,8 +42,8 @@ use super::{
};
use crate::harnesstools::opencode::Opencode;
use crate::harnesstools::opencode::rehydrate::{export_minting, run_import};
use crate::harnesstools::opencode::session::default_db;
use crate::harnesstools::opencode::session::projection::{Projected, projected};
use crate::harnesstools::opencode::session::{self, default_db};
use crate::harnesstools::rehydrate::{RehydrateMessage, RehydrateSession};
impl SessionSync for Opencode {
@@ -89,6 +91,11 @@ impl SessionSync for Opencode {
}
}
/// Why a session opencode 2.0 holds is never caught up: its `opencode session import` only
/// writes sessions it doesn't hold yet.
const OPENCODE_2_APPEND: &str =
"opencode 2 can't add to a session it already holds, so it can't be caught up; fork it instead";
/// Why an opencode session is never switched to another branch.
pub(super) const OPENCODE_SWITCH: &str = "opencode keeps its sessions in a database, which can't \
be switched to another branch; fork instead";
@@ -175,17 +182,14 @@ pub(super) async fn local_tip_in(db: &Path, id: &str) -> Result<Option<LocalTip>
match projected(db, id).await {
Ok(Some(found)) => Ok(Some(LocalTip {
native_path: db.to_path_buf(),
known_source_ids: found.rows.iter().cloned().collect(),
merged: HashMap::new(),
known_source_ids: found.rows.iter().chain(found.merged.keys()).cloned().collect(),
merged: found.merged.clone(),
tip_source_id: found.rows.last().cloned(),
stamp: stamp(&found),
modified: updated(&found),
cwd: Some(found.directory.clone()),
unswitchable: Some(OPENCODE_SWITCH),
})),
Ok(None) if session::locate(db, id).await.is_some() => {
Err(SyncError::Unsupported("opencode 2.0 sessions can't be caught up yet"))
}
Ok(None) => Ok(None),
Err(err) => Err(SyncError::Other(format!("reading opencode's database: {err}"))),
}
@@ -223,6 +227,9 @@ where
tip_source_id: base.tip_source_id.clone(),
});
}
if current.next {
return Err(SyncError::Unsupported(OPENCODE_2_APPEND));
}
if !current.directory.is_dir() {
return Err(SyncError::Other(format!(
"{} is gone: opencode would move the session to wherever it is imported from",
crates/atuin-common/src/harnesstools/sync/opencode/tests.rs
+72
View File
@@ -1,3 +1,5 @@
use std::collections::HashMap;
use pretty_assertions::assert_eq;
use rstest::rstest;
use sqlx::Connection;
@@ -217,6 +219,8 @@ fn ending_in_msg_b() -> Projected {
messages: vec![message("msg_a", 1000), message("msg_b", 1000)],
content: 0,
updated: None,
next: false,
merged: HashMap::new(),
}
}
@@ -263,6 +267,8 @@ fn replies_to_msg_u(replies: &[&str]) -> Projected {
.collect(),
content: 0,
updated: None,
next: false,
merged: HashMap::new(),
}
}
@@ -420,3 +426,69 @@ async fn an_opencode_elsewhere_counts_when_named_or_writing(
assert_eq!(liveness(&procs, &appending(&base, &id, Some(&proj))), expected);
assert_eq!(liveness_in(Some(&db), &id, procs, Some(&proj)).await, expected);
}
const V2: &str = "ses_ee1b7b1dfffe46nQL0NuU68GcR";
const V2_REPLY: &str = "msg_11e48562d001LoisbfgKtTjN6N";
/// A database only opencode 2.0.24 wrote (no 1.x tables), holding one session working in `dir`:
/// a prompt, its reply, and the idle row a turn that succeeded ends with, as 2.0.24 stores them.
async fn v2_session_in(dir: &Path) -> PathBuf {
let db = dir.join("opencode.db");
let options = SqliteConnectOptions::new().filename(&db).create_if_missing(true);
SqliteConnection::connect_with(&options).await.unwrap().close().await.unwrap();
let user = r#"{"time":{"created":1791509417569},"text":"say OCALPHA","files":[]}"#;
let reply = r#"{"time":{"created":1791509419595,"completed":1791509421045},"agent":"build","model":{"id":"big-pickle","providerID":"opencode"},"content":[{"type":"text","text":"OCALPHA"}],"finish":"stop","tokens":{"input":208,"output":43,"reasoning":0,"cache":{"read":7424,"write":0}}}"#;
let idle = r#"{"time":{"created":1791509421052},"outcome":"succeeded"}"#;
execute(
&db,
&format!(
"CREATE TABLE session_v2 (id text PRIMARY KEY, project_id text NOT NULL, parent_id \
text, fork_session_id text, slug text NOT NULL, directory text NOT NULL, title text, \
version text NOT NULL, revert text, time_created integer NOT NULL, time_updated \
integer NOT NULL);
CREATE TABLE session_message (id text PRIMARY KEY, session_id text NOT NULL, type \
text NOT NULL, seq integer NOT NULL, time_created integer NOT NULL, time_updated \
integer NOT NULL, data text NOT NULL);
CREATE TABLE event_sequence (aggregate_id text PRIMARY KEY, seq integer NOT NULL, \
owner_id text);
INSERT INTO session_v2 VALUES ('{V2}', 'global', NULL, NULL, 'slug', '{dir}', \
'octest', '2.0.24', NULL, 1791509417520, 1791509421052);
INSERT INTO session_message VALUES
('msg_11e484e44001y9aBZ5Ki93Ka5M', '{V2}', 'user', 1, 1791509417569, 1791509417569, \
'{user}'),
('{V2_REPLY}', '{V2}', 'assistant', 2, 1791509419595, 1791509421045, '{reply}'),
('msg_11e485bfc001J3OIfsfEAZTPqW', '{V2}', 'idle', 3, 1791509421052, 1791509421052, \
'{idle}');
INSERT INTO event_sequence VALUES ('{V2}', 3, NULL);",
dir = dir.display(),
),
)
.await;
db
}
/// A session opencode 2.0 holds has a tip, read as capture reads it: resuming it with nothing to
/// catch up is no error, and catching it up is refused (opencode 2's import only writes a
/// session it doesn't hold).
#[rstest]
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn a_2_0_session_s_tip_is_the_last_row_capture_delivers(dir: TempDir) {
let db = v2_session_in(dir.path()).await;
let tip = local_tip_in(&db, V2).await.unwrap().unwrap();
let rows = synced(&db, V2).await;
assert_eq!(rows.len(), 2);
assert_eq!(
tip.known_source_ids,
rows.iter().map(|r| r.source_id.clone()).collect::<HashSet<_>>()
);
// the idle row after the reply is no row of capture's
assert_eq!(tip.tip_source_id.as_deref(), Some(V2_REPLY));
assert_eq!(tip.cwd.as_deref(), Some(dir.path()));
assert!(local_tip_in(&db, "ses_none").await.unwrap().is_none());
let caught_up = append(V2, &tip, &[]).await.unwrap();
assert!(caught_up.appended.is_empty());
assert_eq!(caught_up.tip_source_id.as_deref(), Some(V2_REPLY));
let err = append(V2, &tip, &[prompt(NEW, 1_791_600_000)]).await.unwrap_err();
assert!(matches!(err, SyncError::Unsupported(_)), "{err:?}");
}
crates/atuin-common/src/io/line_reader.rs
+192 -20
View File
@@ -1,6 +1,7 @@
//! Line reads of files that writers only append to.
use std::borrow::Borrow;
use std::collections::VecDeque;
use std::fs::File;
use std::io::{self, Read, Seek, SeekFrom};
use std::iter;
@@ -17,6 +18,19 @@ use crate::os::fs::FdIdentityExt;
/// file.
const READ_CHUNK_BYTES: u64 = 64 * 1024;
/// The longest line a reader hands out. One longer is passed over, unread past what this holds,
/// so a file of one huge or unterminated line can't hold that much memory (all of it, and
/// several times that once parsed): no transcript line is near this, the largest being a
/// prompt pasted whole or an image attached inline.
const MAX_LINE_BYTES: usize = 32 * 1024 * 1024;
/// The most buffer a reader keeps between lines: past it, the buffer a long line grew is let go.
const RETAINED_BYTES: usize = 256 * 1024;
/// How much of what precedes the read position a reader keeps, to tell a file rewritten in place
/// from one appended to ([`ReadLinesError::Rewritten`]).
const TAIL_BYTES: usize = 4 * 1024;
/// A complete line, without its newline.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Line {
@@ -67,8 +81,8 @@ impl Line {
/// Errors returned by [`ReadLines::lines`].
///
/// `Truncated` and `Replaced` have already restarted the reader: the next call reads the file from
/// its first byte.
/// `Truncated`, `Replaced` and `Rewritten` have already restarted the reader: the next call reads
/// the file from its first byte.
#[derive(Debug, thiserror::Error)]
pub enum ReadLinesError {
#[error(transparent)]
@@ -77,6 +91,10 @@ pub enum ReadLinesError {
Truncated,
#[error("the file was replaced; reading restarts from its start")]
Replaced,
/// Written over in place, to the read position or past it: what precedes the position is no
/// longer what was read.
#[error("the file was rewritten; reading restarts from its start")]
Rewritten,
}
/// Reads the complete lines appended to a file, each once.
@@ -92,9 +110,9 @@ pub trait ReadLines {
pub trait AsyncReadLines {
/// The complete lines past the last one handed out, read as the stream is pulled.
///
/// The stream ends at the end of the file or after its first error; a truncated or replaced
/// file is read on from its start. Lines read but not yet yielded are dropped with the stream
/// and not read again.
/// The stream ends at the end of the file or after its first error; a file truncated,
/// replaced or rewritten in place is read on from its start. Lines read but not yet yielded
/// are dropped with the stream and not read again.
fn lines(&mut self) -> impl Stream<Item = io::Result<Line>> + Send + '_;
}
@@ -112,6 +130,11 @@ struct Cursor {
identity: Option<FdIdentity>,
/// How many times the read started over from the file's start ([`Line::generation`]).
generation: u32,
/// The [`TAIL_BYTES`] before `offset` (fewer when `offset` is), as they were read; empty
/// until known.
tail: VecDeque<u8>,
/// Whether the line at `offset` is past [`MAX_LINE_BYTES`], and passed over up to its newline.
skipping: bool,
}
impl Cursor {
@@ -145,7 +168,8 @@ impl Cursor {
}))
}
/// Restart at byte 0 if `file` is not the file being read, or is shorter than what was read.
/// Restart at byte 0 if `file` is not the file being read, is shorter than what was read, or
/// no longer holds what was read just before the position.
fn restart_unless_appended(&mut self, file: &File) -> Result<(), ReadLinesError> {
let meta = file.metadata()?;
#[cfg(unix)]
@@ -159,6 +183,10 @@ impl Cursor {
ReadLinesError::Replaced
} else if meta.len() < self.position() {
ReadLinesError::Truncated
} else if !self.holds_tail(file)? {
// Written over in place to as long or longer (a harness migrating its transcript),
// where neither of the above can tell.
ReadLinesError::Rewritten
} else {
self.identity = identity;
return Ok(());
@@ -171,6 +199,44 @@ impl Cursor {
Err(restart)
}
/// Whether `file`, at least as long as the position, still holds the [tail](Self::tail)
/// before the offset, and the start of `pending` after it; a tail not known yet is taken from
/// it.
fn holds_tail(&mut self, mut file: &File) -> io::Result<bool> {
let len = usize::try_from(self.offset.min(TAIL_BYTES as u64)).map_err(io::Error::other)?;
// And the start of the line read in part past it.
let ahead = self.pending.len().min(TAIL_BYTES);
if len + ahead == 0 {
return Ok(true);
}
let mut found = vec![0; len + ahead];
file.seek(SeekFrom::Start(self.offset - len as u64))?;
file.read_exact(&mut found)?;
let (before, after) = found.split_at(len);
if after != &self.pending[..ahead] {
return Ok(false);
}
if self.tail.is_empty() {
self.tail = before.to_vec().into();
return Ok(true);
}
Ok(self.tail.make_contiguous() == before)
}
/// Keep the [tail](Self::tail) up to date with `consumed`, the bytes just handed out (or
/// passed over) up to `offset`.
fn keep_tail(&mut self, consumed: &[u8]) {
let from_start = self.offset == consumed.len() as u64;
if consumed.len() >= TAIL_BYTES {
self.tail.clear();
self.tail.extend(&consumed[consumed.len() - TAIL_BYTES..]);
} else if from_start || !self.tail.is_empty() {
self.tail.extend(consumed);
let excess = self.tail.len().saturating_sub(TAIL_BYTES);
self.tail.drain(..excess);
}
}
/// The next complete line, reading on from `file` as needed; `None` once it has none.
fn next_line(&mut self, file: &File) -> io::Result<Option<Line>> {
loop {
@@ -183,23 +249,63 @@ impl Cursor {
}
}
/// Hand out the first complete line in `pending`, if any, advancing `offset` past it.
/// Hand out the first complete line in `pending`, if any, advancing `offset` past it. A line
/// past [`MAX_LINE_BYTES`] is passed over, up to its newline.
fn split_line(&mut self) -> Option<Line> {
let Some(newline) = memchr::memchr(b'\n', &self.pending[self.scanned..]) else {
self.scanned = self.pending.len();
return None;
};
loop {
let Some(newline) = memchr::memchr(b'\n', &self.pending[self.scanned..]) else {
self.scanned = self.pending.len();
if self.pending.len() > MAX_LINE_BYTES {
self.skip_pending();
}
return None;
};
let len = self.scanned + newline + 1;
let mut bytes = self.pending.split_to(len).freeze();
bytes.truncate(len - 1);
let len = self.scanned + newline + 1;
let mut bytes = self.pending.split_to(len).freeze();
self.scanned = 0;
self.offset += u64::try_from(len).expect("a byte count fits u64");
self.keep_tail(&bytes);
// A line read whole leaves the buffer it grew to behind: let it go.
if self.pending.capacity() > RETAINED_BYTES {
self.pending = BytesMut::from(&self.pending[..]);
}
let skipping = std::mem::take(&mut self.skipping);
if skipping || len - 1 > MAX_LINE_BYTES {
if !skipping {
self.warn_too_long();
}
continue;
}
bytes.truncate(len - 1);
return Some(Line {
end: self.offset,
bytes,
generation: self.generation,
});
}
}
/// Pass over `pending`, the start of a line past [`MAX_LINE_BYTES`], and the rest of it as it
/// is read, releasing the memory it held.
fn skip_pending(&mut self) {
if !self.skipping {
self.warn_too_long();
}
self.skipping = true;
self.offset += u64::try_from(self.pending.len()).expect("a byte count fits u64");
self.pending = BytesMut::new();
self.scanned = 0;
self.offset += u64::try_from(len).expect("a byte count fits u64");
Some(Line {
end: self.offset,
bytes,
generation: self.generation,
})
// The bytes passed over are not kept, so what precedes the offset is not known.
self.tail.clear();
}
fn warn_too_long(&self) {
tracing::warn!(
near = self.offset,
max = MAX_LINE_BYTES,
"passing over a line longer than a reader reads"
);
}
/// Read up to [`READ_CHUNK_BYTES`] past the position into `pending`, returning how many.
@@ -437,6 +543,72 @@ mod tests {
assert_eq!(drain(reader.lines().unwrap()), [after_restarts(1, line(6, "xyzab"))]);
}
/// Written over in place (truncated and written again, as pi migrating a transcript does) to
/// as long as what was read or longer, which neither its identity nor its length tells.
#[rstest]
#[case::as_long(b"xx\nyyy\n", vec![line(3, "xx"), line(7, "yyy")])]
#[case::longer(b"x\nyy\nzzzz\n", vec![line(2, "x"), line(5, "yy"), line(10, "zzzz")])]
fn a_file_rewritten_in_place_restarts(
#[with(b"aa\nbbb\n")] file: TempFile,
#[case] rewritten: &[u8],
#[case] expected: Vec<Line>,
#[values(false, true)] resumed: bool,
) {
let mut reader = if resumed {
// Resumed past the lines, what precedes the offset is taken from its first read.
PathLineReader::at(&file.path, 7)
} else {
PathLineReader::new(&file.path)
};
drain(reader.lines().unwrap());
std::fs::write(&file.path, rewritten).unwrap();
assert!(matches!(reader.lines(), Err(ReadLinesError::Rewritten)));
let expected: Vec<Line> = expected.into_iter().map(|l| after_restarts(1, l)).collect();
assert_eq!(drain(reader.lines().unwrap()), expected);
}
/// Written over in place past the lines read, where a line was read in part: the part read is
/// not joined to another line's end.
#[rstest]
fn a_line_read_in_part_rewritten_in_place_restarts(#[with(b"aa\nbbb")] file: TempFile) {
let mut reader = PathLineReader::new(&file.path);
assert_eq!(drain(reader.lines().unwrap()), [line(3, "aa")]);
std::fs::write(&file.path, b"aa\nxyz\n").unwrap();
assert!(matches!(reader.lines(), Err(ReadLinesError::Rewritten)));
let again = [after_restarts(1, line(3, "aa")), after_restarts(1, line(7, "xyz"))];
assert_eq!(drain(reader.lines().unwrap()), again);
}
/// Written over in place with what was read, and more: read on, as if appended to.
#[rstest]
fn a_file_rewritten_with_what_was_read_reads_on(#[with(b"aa\nbbb\n")] file: TempFile) {
let mut reader = PathLineReader::new(&file.path);
drain(reader.lines().unwrap());
std::fs::write(&file.path, b"aa\nbbb\ncc\n").unwrap();
assert_eq!(drain(reader.lines().unwrap()), [line(10, "cc")]);
}
/// A line past the longest read is passed over, the buffer it took let go, and the lines after
/// it read as ever; so is one never terminated.
#[rstest]
fn a_line_past_the_longest_is_passed_over() {
let long = vec![b'm'; MAX_LINE_BYTES + 1];
let file = file(&[&b"a\n"[..], &long, b"\nb\n"].concat());
let mut reader = PathLineReader::new(&file.path);
let end = 2 + long.len() as u64 + 1;
assert_eq!(drain(reader.lines().unwrap()), [line(2, "a"), line(end + 2, "b")]);
append(&file.path, &long);
assert_eq!(drain(reader.lines().unwrap()), []);
assert!(reader.cursor.pending.capacity() < MAX_LINE_BYTES, "the line is not held");
append(&file.path, b"\nc\n");
let end = end + 2 + long.len() as u64 + 1;
assert_eq!(drain(reader.lines().unwrap()), [line(end + 2, "c")]);
}
#[cfg(unix)]
#[rstest]
fn a_file_replaced_at_the_path_restarts(#[with(b"1\n")] file: TempFile) {
crates/atuin-common/src/io/pooled_read_lines.rs
+10 -5
View File
@@ -81,14 +81,18 @@ enum LineBatchEnd {
impl LineBatch {
/// Pull lines from `reader` until they reach [`LINE_BATCH_BYTES`], the file ends, or a read fails;
/// a truncated or replaced file is read on from its start.
/// a file truncated, replaced or rewritten is read on from its start.
fn read(reader: &mut impl ReadLines) -> Self {
// A reader reports each truncation or replacement once, so this goes round again only
// for a file replaced again between two opens.
// A reader reports each truncation, replacement or rewrite once, so this goes round
// again only for a file replaced again between two opens.
let lines = loop {
match reader.lines() {
Ok(lines) => break lines,
Err(error @ (ReadLinesError::Truncated | ReadLinesError::Replaced)) => {
Err(
error @ (ReadLinesError::Truncated
| ReadLinesError::Replaced
| ReadLinesError::Rewritten),
) => {
tracing::debug!(%error, "reading from the start");
}
Err(ReadLinesError::Io(err)) => {
@@ -190,7 +194,8 @@ mod tests {
#[rstest]
fn a_restart_is_read_through_in_the_same_batch(
#[values(ReadLinesError::Truncated, ReadLinesError::Replaced)] restart: ReadLinesError,
#[values(ReadLinesError::Truncated, ReadLinesError::Replaced, ReadLinesError::Rewritten)]
restart: ReadLinesError,
) {
let mut reader = Scripted([Err(restart), Ok(vec![Ok(line("a"))])].into());
let batch = LineBatch::read(&mut reader);
crates/atuin-common/src/json/jsonl.rs
+33 -2
View File
@@ -36,7 +36,9 @@ pub trait JsonlExt: Stream<Item = io::Result<Line>> + Sized {
self.json_with(|bytes| crate::json::js::from_slice(bytes))
}
/// Each non-blank line decoded by `parse`, yielded with the line it came from.
/// Each non-blank line decoded by `parse`, yielded with the line it came from. A line whose
/// decoding panics is malformed, as one that fails to decode is, rather than the end of the
/// stream (and of whatever reads it, which reads it again).
fn json_with<T>(
self,
parse: impl Fn(&[u8]) -> serde_json::Result<T>,
@@ -44,7 +46,7 @@ pub trait JsonlExt: Stream<Item = io::Result<Line>> + Sized {
self.filter_map(move |line| {
let item = match line {
Ok(line) if line.bytes.trim_ascii().is_empty() => None,
Ok(line) => Some(match parse(&line.bytes) {
Ok(line) => Some(match decode(&parse, &line.bytes) {
Ok(value) => Ok((line, value)),
Err(source) => Err(JsonlError::Parse {
source,
@@ -60,6 +62,15 @@ pub trait JsonlExt: Stream<Item = io::Result<Line>> + Sized {
impl<S: Stream<Item = io::Result<Line>>> JsonlExt for S {}
/// `bytes` decoded by `parse`, a panic in it an error.
fn decode<T>(
parse: impl Fn(&[u8]) -> serde_json::Result<T>,
bytes: &[u8],
) -> serde_json::Result<T> {
std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| parse(bytes)))
.unwrap_or_else(|_| Err(serde::de::Error::custom("decoding the line panicked")))
}
#[cfg(test)]
mod tests {
use std::io::Write;
@@ -141,6 +152,26 @@ mod tests {
assert_eq!(got, expected);
}
/// A line whose decoding panics is malformed, and the lines after it are read on.
#[rstest]
#[tokio::test]
async fn a_line_whose_decoding_panics_is_malformed() {
let lines = vec![Ok(line(2, "!")), Ok(line(4, "1"))];
let got: Vec<Item> = stream::iter(lines)
.json_with(|bytes| {
assert_ne!(bytes, b"!", "a decoder bug");
serde_json::from_slice::<i64>(bytes)
})
.map(|item| match item {
Ok((line, value)) => Item::Value(line.end, value),
Err(JsonlError::Parse { end, .. }) => Item::Malformed(end),
Err(JsonlError::Io(_)) => Item::Unreadable,
})
.collect()
.await;
assert_eq!(got, [Item::Malformed(2), Item::Value(4, 1)]);
}
proptest! {
// Records written a line each, blank lines between some, come back in order, each with
// the offset just past its own line.
crates/atuin-common/src/secrets.rs
+98
View File
@@ -1043,6 +1043,81 @@ mod tests {
)]
// Five stars is not the marker; it is a value, and comes out as the marker.
#[case::five_stars_are_a_value("AWS_SECRET_ACCESS_KEY=*****", "AWS_SECRET_ACCESS_KEY=****")]
// A patch's lines (Codex `apply_patch`, opencode `patchText`, a unified diff) start with
// `+`, `-` or ` `: what is set on one is redacted as it would be on its own.
#[case::added_to_a_patch(
"*** Add File: config/database.yml\n+production:\n+ password: hunter2xyz\n+ host: db",
"*** Add File: config/database.yml\n+production:\n+ password: ****\n+ host: db"
)]
#[case::changed_in_a_patch(
"@@\n-spring.datasource.password=old\n+spring.datasource.password=hunter2xyz",
"@@\n-spring.datasource.password=****\n+spring.datasource.password=****"
)]
#[case::added_aws_credentials(
"+[default]\n+aws_secret_access_key = wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY",
"+[default]\n+aws_secret_access_key = ****"
)]
#[case::added_env_var("+API_KEY=abcdef1234567\n", "+API_KEY=****\n")]
#[case::added_list_item("+ - password: hunter2xyz", "+ - password: ****")]
// JSON text, its lines escaped (`\n`): a line it holds is still a line.
#[case::patch_in_json_text(
r#"{"patchText":"*** Begin Patch\n+password: hunter2xyz\n*** End Patch"}"#,
r#"{"patchText":"*** Begin Patch\n+password: ****\n*** End Patch"}"#
)]
#[case::setting_in_json_text(
r#"log: "first\npassword=hunter2xyz""#,
r#"log: "first\npassword=****""#
)]
#[case::env_var_in_json_text(r#""x\nAPI_KEY=abcdef123456""#, r#""x\nAPI_KEY=****""#)]
#[case::bearer_in_json_text(
r#""result\nBearer abcdefghijklmnopqrstuvwxyz""#,
r#""result\nBearer ****""#
)]
// JSON text clipped mid-string: its last line ends with the text.
#[case::clipped_json_text(r#""a\npassword: hunter2"#, r#""a\npassword: ****"#)]
#[case::command_in_json_text(
r#"{"cmd": "PGPASSWORD=hunter2xyz psql -h db -c select1"}"#,
r#"{"cmd": "PGPASSWORD=**** psql -h db -c select1"}"#
)]
#[case::password_flag_in_a_script(
"bash -lc 'docker login --password s3cret'",
"bash -lc 'docker login --password ****'"
)]
#[case::password_flag_in_json_text(
r#"{"cmd":"docker login --password s3cret\ndocker push x"}"#,
r#"{"cmd":"docker login --password ****\ndocker push x"}"#
)]
#[case::password_flag_then_command(
"docker login --password s3cret; docker push x",
"docker login --password ****; docker push x"
)]
#[case::password_flag_in_a_subshell(
"(docker login --password s3cret)&& echo",
"(docker login --password ****)&& echo"
)]
#[case::password_flag_then_and(
"docker login --password s3cret&& echo",
"docker login --password ****&& echo"
)]
#[case::clipped_json_line_in_output(
"\"a\\naws_secret_access_key = CXAT00ju\nmore output",
"\"a\\naws_secret_access_key = ****\nmore output"
)]
#[case::password_flag_quoted_in_prose("`--password=hunter2xyz`, or", "`--password=****`, or")]
#[case::jwt_in_json_text(
r#"{"data":"x\teyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIn0.dozjgNryP4J3jVmNHl0w5N_XgL0n3I9PlFUP0THsR8U"}"#,
r#"{"data":"x\t****"}"#
)]
// A password given on the command line.
#[case::password_flag_apart(
"docker login -u me --password hunter2xyz r.io",
"docker login -u me --password **** r.io"
)]
// AWS's own files and commands spell the name in lowercase.
#[case::lowercase_aws_assignment(
"export aws_secret_access_key=wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY",
"export aws_secret_access_key=****"
)]
fn redacts(#[case] input: &str, #[case] expected: &str) {
assert_eq!(&*redact(input), expected);
}
@@ -1627,6 +1702,29 @@ mod tests {
#[case::usage(r#"{"max_tokens": 4096, "input_tokens": 123456789}"#)]
#[case::count("tokens: 1500")]
#[case::prose("Set the token in your settings.")]
#[case::password_flag_reference("docker login --password \"$REGISTRY_PASSWORD\" r.io")]
#[case::password_from_stdin("docker login --password-stdin r.io")]
#[case::password_flag_then_flag("mysql --password --host db")]
#[case::password_flag_in_help(" -p, --password string Password")]
#[case::password_flag_in_click_help(" --password TEXT The password")]
#[case::password_flag_in_urfave_help(" --password value the password")]
#[case::passwd_flag_in_help(" --passwd FILE")]
#[case::password_flag_in_prose("use the --password flag, or see the --password option.")]
#[case::password_flag_listed("--password and --user")]
#[case::password_flag_in_a_subject("fix(secrets): redact --password flags and more")]
#[case::password_flag_reference_in_json_text(
r#"{"cmd":"docker login --password \"$REGISTRY_PASSWORD\" r.io"}"#
)]
#[case::setting_concatenated(r#"password = prompt + "x""#)]
#[case::setting_joined(r#"token = base + "/v1""#)]
#[case::setting_in_prose(r#"secret: see "docs""#)]
#[case::command_in_a_string(r#""VAULT_TOKEN=x vault status""#)]
#[case::password_flag_then_or("mysql --password || true")]
#[case::password_flag_then_redirect("mysql --password 2>&1 | tee log")]
#[case::password_flag_substituted("docker login --password `cat pw` r.io")]
#[case::password_flag_in_a_subshell("(mysql --password)")]
#[case::added_field("+ token: String,")]
#[case::short_lowercase_aws("aws configure get aws_secret_access_key --profile prod")]
fn stand_ins_are_kept(#[case] input: &str) {
assert_eq!(redact(input), input);
}
crates/atuin-common/src/secrets/structural.rs
+96 -8
View File
@@ -14,10 +14,21 @@ use super::REDACTED;
#[cfg(test)]
use super::tests::Test;
/// Where a word starts: at a word boundary, or just past an escaped newline or tab (`\n` in JSON
/// text), where there is none.
macro_rules! word_start {
() => {
r"(?:(?-u:\b)|\\[nrt])"
};
}
/// The words a name ends with when it holds a credential, as an environment variable spells them.
macro_rules! env_secret_name {
() => {
r"(?-u:\b)[A-Z0-9_]*(?:SECRET|TOKEN|PASSWORD|PASSWD|_PWD|API_?KEY|APIKEY|PRIVATE_KEY|ACCESS_KEY|SECRET_KEY|CREDENTIALS?|AUTH)S?(?-u:\b)"
concat!(
word_start!(),
r"[A-Z0-9_]*(?:SECRET|TOKEN|PASSWORD|PASSWD|_PWD|API_?KEY|APIKEY|PRIVATE_KEY|ACCESS_KEY|SECRET_KEY|CREDENTIALS?|AUTH)S?(?-u:\b)"
)
};
}
@@ -74,7 +85,10 @@ pub(super) static STRUCTURAL_PATTERNS: &[Pattern] = &[
// gitleaks jwt
Pattern {
name: "JSON web token",
regex: r"(?-u:\b)(?<secret>eyJ[A-Za-z0-9_-]{10,}\.eyJ[A-Za-z0-9_-]{10,}\.[A-Za-z0-9_-]{10,})",
regex: concat!(
word_start!(),
r"(?<secret>eyJ[A-Za-z0-9_-]{10,}\.eyJ[A-Za-z0-9_-]{10,}\.[A-Za-z0-9_-]{10,})"
),
prefilter: None,
#[cfg(test)]
tests: &[Test {
@@ -87,7 +101,7 @@ pub(super) static STRUCTURAL_PATTERNS: &[Pattern] = &[
},
Pattern {
name: "Bearer token",
regex: r"(?i-u:\bbearer)[ \t]+(?<secret>[A-Za-z0-9._~+/-]{16,}=*)",
regex: concat!(word_start!(), r"(?i-u:bearer)[ \t]+(?<secret>[A-Za-z0-9._~+/-]{16,}=*)"),
prefilter: None,
#[cfg(test)]
tests: &[Test {
@@ -140,7 +154,7 @@ pub(super) static STRUCTURAL_PATTERNS: &[Pattern] = &[
// A connection string's `password=` (keyword DSNs, JDBC and ODBC strings, query strings).
Pattern {
name: "Connection string password",
regex: r#"(?:^|[?&;\s])(?i-u:password|pwd)=(?<secret>"[^"\n]*"|'[^'\n]*'|[^&;\s"']+)"#,
regex: r#"(?:^|[?&;\s-])(?i-u:password|pwd)=(?<secret>"[^"\n]*"|'[^'\n]*'|[^&;\s"'`]*[^&;\s"'`,])"#,
prefilter: None,
#[cfg(test)]
tests: &[
@@ -152,13 +166,30 @@ pub(super) static STRUCTURAL_PATTERNS: &[Pattern] = &[
input: "Server=db;User Id=sa;Password=Hunter2!;",
redacted: "Server=db;User Id=sa;Password=****;",
},
Test {
input: "mysql -u root --password=hunter db",
redacted: "mysql -u root --password=**** db",
},
],
},
// AWS's own files and commands spell the name in lowercase (`aws configure set
// aws_secret_access_key ...`); the value is a key's 40 characters, so a space can separate them.
Pattern {
name: "AWS secret access key setting",
regex: r#"(?i-u:aws_secret_access_key)["']?(?:[ \t]*[:=][ \t]*|[ \t]+)["']?(?<secret>[A-Za-z0-9/+]{40,})"#,
prefilter: None,
#[cfg(test)]
tests: &[Test {
input: "aws configure set aws_secret_access_key \
wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY",
redacted: "aws configure set aws_secret_access_key ****",
}],
},
// gitleaks twilio-api-key: only redacted, as `SK` and 32 hex digits is too plain a shape to
// keep a command out of history.
Pattern {
name: "Twilio API key",
regex: r"(?-u:\b)(?<secret>SK[0-9a-fA-F]{32})(?-u:\b)",
regex: concat!(word_start!(), r"(?<secret>SK[0-9a-fA-F]{32})(?-u:\b)"),
prefilter: None,
#[cfg(test)]
tests: &[Test {
@@ -238,11 +269,37 @@ pub(super) static NAMED_PATTERNS: &[Pattern] = &[
},
],
},
// YAML, `.env`, INI and `.properties` lines: an unquoted value to the end of the line.
// A password given on the command line apart from its flag (`docker login --password ...`):
// quoted, or holding more than letters, as prose (`the --password flag`) and a help's
// placeholder (`--password TEXT`, `--password string`) don't. `--password=...` is redacted as
// a connection string's `password=`.
Pattern {
name: "Password flag",
regex: concat!(
r#"(?:^|[ \t])--passw(?:or)?d[ \t]+(?<secret>"(?:[^"\\\n]|\\.)*"|'[^'\n]*'"#,
r#"|\\"(?:[^"\\\n]|\\[^"])*\\""#,
r#"|(?:[^\s"'\\<>\[{()|&;`-][^\s"'\\|&;<>()`]*?)?"#,
r#"[^A-Za-z\s"'\\.,:;()\[\]{}<>|&`-][^\s"'\\|&;<>()`]*)(?:[\s"';&|)\\]|\z)"#
),
prefilter: None,
#[cfg(test)]
tests: &[
Test {
input: "docker login -u me --password hunter2xyz r.io",
redacted: "docker login -u me --password **** r.io",
},
Test {
input: "docker login -u me --password 'hunter 2' r.io",
redacted: "docker login -u me --password **** r.io",
},
],
},
// YAML, `.env`, INI and `.properties` lines: an unquoted value to the end of the line. Also
// such a line in a patch (`+`; a `-` is part of the name).
Pattern {
name: "Credential setting",
regex: concat!(
r#"(?m)^[ \t]*(?:-[ \t]+)?["']?"#,
r#"(?m)^\+?[ \t]*(?:-[ \t]+)?["']?"#,
code_secret_name!(),
r#"["']?[ \t]*[:=][ \t]*(?<secret>[^\s"'#|>&*!{}\[\],=][^\n#="]*?)[ \t]*$"#
),
@@ -266,6 +323,37 @@ pub(super) static NAMED_PATTERNS: &[Pattern] = &[
input: "\npassword: violetmeadow\n",
redacted: "\npassword: ****\n",
},
Test {
input: "\n+ password: hunter2xyz\n",
redacted: "\n+ password: ****\n",
},
],
},
// Such a line in JSON text: just past an escaped newline, and ending at the next one, the
// string's end, or, clipped, the line's.
Pattern {
name: "Credential setting in JSON text",
regex: concat!(
r#"(?m)\\n\+?(?:[ \t]|\\t)*(?:-[ \t]+)?["']?"#,
code_secret_name!(),
r#"["']?[ \t]*[:=][ \t]*(?<secret>[^\s"'#|>&*!{}\[\],=\\](?:[^\n#="\\]|\\(?:[^rn]|\z))*?)"#,
r#"[ \t]*(?:\\[rn]|"|$)"#
),
prefilter: None,
#[cfg(test)]
tests: &[
Test {
input: r#""a:\n password: hunter2xyz\n""#,
redacted: r#""a:\n password: ****\n""#,
},
Test {
input: r#""a\npassword: ab\"cd\"""#,
redacted: r#""a\npassword: ****""#,
},
Test {
input: r#""a:\n\tpassword: hunter2xyz""#,
redacted: r#""a:\n\tpassword: ****""#,
},
],
},
];
@@ -274,7 +362,7 @@ pub(super) static NAMED_PATTERNS: &[Pattern] = &[
/// (`$TOKEN`, `${{ secrets.X }}`, `os.environ[...]`), documentation's stand-ins (`<password>`,
/// `changeme`, `xxxx`), the marker itself, and numbers.
pub(super) fn is_placeholder(value: &str) -> bool {
let value = value.trim().trim_matches(|c| c == '"' || c == '\'').trim();
let value = value.trim().trim_matches(|c| c == '"' || c == '\'' || c == '\\').trim();
if value.is_empty() {
return true;
}
crates/atuin-common/tests/fixtures/codex/code-mode.jsonl
Vendored
+45
View File
@@ -0,0 +1,45 @@
{"timestamp":"2026-10-09T01:32:41.336Z","ordinal":0,"type":"session_meta","payload":{"session_id":"01a11e4a-7de0-77a3-9030-5535c9877485","id":"01a11e4a-7de0-77a3-9030-5535c9877485","timestamp":"2026-10-09T01:32:41.296Z","cwd":"/home/user/proj","runtime_workspace_roots":["/home/user/proj"],"originator":"codex_exec","cli_version":"0.160.1","source":"exec","thread_source":"user","model_provider":"openai","base_instructions":{"text":"You are Codex."},"history_mode":"paginated","context_window":{"window_id":"01a11e4a-7fc8-7513-ab63-89ed12fffc48"},"git":{"commit_hash":"35abdd3b185cd3ae5a0160940d66d934a323b950","branch":"feature-x"}}}
{"timestamp":"2026-10-09T01:32:41.336Z","ordinal":1,"type":"event_msg","payload":{"type":"task_started","turn_id":"01a11e4a-7fe0-7dc3-8a0c-6fe7d2fd4d48","root_turn_id":"01a11e4a-7fe0-7dc3-8a0c-6fe7d2fd4d48","started_at":1791509561,"model_context_window":258400,"collaboration_mode_kind":"default"}}
{"timestamp":"2026-10-09T01:32:42.332Z","ordinal":2,"type":"response_item","payload":{"type":"message","id":"msg_01a11e4a-83dc-7412-9ba9-32d29a1fa360","role":"developer","content":[{"type":"input_text","text":"<permissions instructions>trimmed</permissions instructions>"}],"internal_chat_message_metadata_passthrough":{"turn_id":"01a11e4a-7fe0-7dc3-8a0c-6fe7d2fd4d48","create_time":1791509562.3322098,"content_item_kinds":["host_skills.instructions","permissions.instructions","collaboration_mode.instructions","plugins.recommendations"]}},"metadata":{"client_authored":false,"mcp_attribution":{"status":"none"}}}
{"timestamp":"2026-10-09T01:32:42.332Z","ordinal":3,"type":"response_item","payload":{"type":"message","id":"msg_01a11e4a-83dc-7412-9ba9-32ea1f43f8db","role":"developer","content":[{"type":"input_text","text":"<permissions instructions>trimmed</permissions instructions>"}],"internal_chat_message_metadata_passthrough":{"turn_id":"01a11e4a-7fe0-7dc3-8a0c-6fe7d2fd4d48","create_time":1791509562.3322113,"content_item_kinds":["multi_agent.role_instructions"]}}}
{"timestamp":"2026-10-09T01:32:42.334Z","ordinal":4,"type":"response_item","payload":{"type":"message","id":"msg_01a11e4a-83dc-7412-9ba9-32f12d1d344c","role":"developer","content":[{"type":"input_text","text":"<permissions instructions>trimmed</permissions instructions>"}],"internal_chat_message_metadata_passthrough":{"turn_id":"01a11e4a-7fe0-7dc3-8a0c-6fe7d2fd4d48","create_time":1791509562.3322115,"content_item_kinds":["multi_agent.mode_instructions"]}}}
{"timestamp":"2026-10-09T01:32:42.335Z","ordinal":5,"type":"response_item","payload":{"type":"message","id":"msg_01a11e4a-83dc-7412-9ba9-330e14b35a97","role":"user","content":[{"type":"input_text","text":"<environment_context>\n <cwd>/home/user/proj</cwd>\n <shell>zsh</shell>\n <current_date>2026-10-09</current_date>\n <timezone>UTC</timezone>\n <filesystem><workspace_roots><root>/home/user/proj</root></workspace_roots><permission_profile type=\"managed\"><file_system type=\"restricted\"><entry access=\"read\"><special>:root</special></entry><entry access=\"write\"><path>/home/user/proj</path></entry><entry access=\"write\"><special>:slash_tmp</special></entry><entry access=\"write\"><special>:tmpdir</special></entry><entry access=\"read\"><path>/home/user/proj/.git</path></entry><entry access=\"read\"><path>/home/user/proj/.agents</path></entry><entry access=\"read\"><path>/home/user/proj/.codex</path></entry><entry access=\"read\"><path>/home/user/proj/.aws</path></entry></file_system></permission_profile></filesystem>\n</environment_context>"}],"internal_chat_message_metadata_passthrough":{"turn_id":"01a11e4a-7fe0-7dc3-8a0c-6fe7d2fd4d48","create_time":1791509562.3322117,"content_item_kinds":["environments.environment_context"]}}}
{"timestamp":"2026-10-09T01:32:42.338Z","ordinal":6,"type":"world_state","payload":{"full":true,"state":{}}}
{"timestamp":"2026-10-09T01:32:42.341Z","ordinal":7,"type":"turn_context","payload":{"turn_id":"01a11e4a-7fe0-7dc3-8a0c-6fe7d2fd4d48","cwd":"/home/user/proj","current_date":"2026-10-09","timezone":"UTC","approval_policy":"never","sandbox_policy":{"type":"workspace-write","network_access":false,"exclude_tmpdir_env_var":false,"exclude_slash_tmp":false},"model":"gpt-6-luna"}}
{"timestamp":"2026-10-09T01:32:42.364Z","ordinal":8,"type":"response_item","payload":{"type":"message","id":"msg_01a11e4a-83fc-7122-a811-7bc86975f81b","role":"user","content":[{"type":"input_text","text":"Marker word FOXTROTPEAR. Do these steps: 1) run the shell command 'git branch --show-current'; 2) print the file notes.txt with cat; 3) edit calc.py using apply_patch to add a function 'def mul(a, b): return a * b' at the end; 4) create a new file fox.txt containing 'foxtrotpear' using apply_patch; 5) run 'ls /nonexistent_foxtrotpear' (it will fail, that is expected). Then reply with one short sentence containing FOXTROTPEAR."}],"internal_chat_message_metadata_passthrough":{"turn_id":"01a11e4a-7fe0-7dc3-8a0c-6fe7d2fd4d48","create_time":1791509562.3646266,"content_item_kinds":["user.text"]}},"metadata":{"retained_source":{"id":{"message_id":"msg_01a11e4a-83fc-7122-a811-7bc86975f81b","turn_id":"01a11e4a-7fe0-7dc3-8a0c-6fe7d2fd4d48","role":"user"},"revision":"retained_01a11e4a-83fc-7122-a811-7be1c9cf7bac","complete":true},"client_authored":false,"user_input_order":0,"mcp_attribution":{"status":"none"}}}
{"timestamp":"2026-10-09T01:32:42.367Z","ordinal":9,"type":"event_msg","payload":{"type":"item_completed","thread_id":"01a11e4a-7de0-77a3-9030-5535c9877485","turn_id":"01a11e4a-7fe0-7dc3-8a0c-6fe7d2fd4d48","item":{"type":"UserMessage","id":"01a11e4a-83fc-7122-a811-7bd0f9961192","content":[{"type":"text","text":"Marker word FOXTROTPEAR. Do these steps: 1) run the shell command 'git branch --show-current'; 2) print the file notes.txt with cat; 3) edit calc.py using apply_patch to add a function 'def mul(a, b): return a * b' at the end; 4) create a new file fox.txt containing 'foxtrotpear' using apply_patch; 5) run 'ls /nonexistent_foxtrotpear' (it will fail, that is expected). Then reply with one short sentence containing FOXTROTPEAR.","text_elements":[]}]},"started_at_ms":1791509562367,"completed_at_ms":1791509562367}}
{"timestamp":"2026-10-09T01:32:43.572Z","ordinal":10,"type":"response_item","payload":{"type":"custom_tool_call","id":"ctc_0c4f036e98312819016ac8443b510c819aa96729ca97dd8c79","status":"completed","call_id":"call_71e14f1557fd45378ab5f7186973a336","name":"exec","input":"const r = await tools.exec_command({cmd:\"git branch --show-current\", max_output_tokens:200});\ntext(r.output);\n","internal_chat_message_metadata_passthrough":{"turn_id":"01a11e4a-7fe0-7dc3-8a0c-6fe7d2fd4d48","create_time":1791509563.3186536}}}
{"timestamp":"2026-10-09T01:32:43.619Z","ordinal":11,"type":"token_usage_record","payload":{"thread_id":"01a11e4a-7de0-77a3-9030-5535c9877485","turn_id":"01a11e4a-7fe0-7dc3-8a0c-6fe7d2fd4d48","session_id":"01a11e4a-7de0-77a3-9030-5535c9877485","root_turn_id":"01a11e4a-7fe0-7dc3-8a0c-6fe7d2fd4d48","response_id":"resp_0c4f036e98312819016ac8443a9904819aba4af4def6b7f7f6","usage":{"input_tokens":13962,"cached_input_tokens":11008,"cache_write_input_tokens":0,"output_tokens":36,"reasoning_output_tokens":0,"total_tokens":13998},"turn_token_usage":{"input_tokens":13962,"cached_input_tokens":11008,"cache_write_input_tokens":0,"output_tokens":36,"reasoning_output_tokens":0,"total_tokens":13998},"thread_token_usage":{"input_tokens":13962,"cached_input_tokens":11008,"cache_write_input_tokens":0,"output_tokens":36,"reasoning_output_tokens":0,"total_tokens":13998}}}
{"timestamp":"2026-10-09T01:32:43.784Z","ordinal":12,"type":"event_msg","payload":{"type":"item_completed","thread_id":"01a11e4a-7de0-77a3-9030-5535c9877485","turn_id":"01a11e4a-7fe0-7dc3-8a0c-6fe7d2fd4d48","item":{"type":"CommandExecution","id":"exec-9faa19c3-b2ae-4b9c-b5fe-15bc1d72ec73","process_id":"37846","command":["/bin/zsh","-lc","git branch --show-current"],"cwd":"file:///home/user/proj","parsed_cmd":[{"type":"unknown","cmd":"git branch --show-current"}],"source":"unified_exec_startup","status":"completed","stdout":"feature-x\n","stderr":"","aggregated_output":"feature-x\n","exit_code":0,"duration":{"secs":0,"nanos":6052},"formatted_output":"feature-x\n"},"started_at_ms":1791509563783,"completed_at_ms":1791509563783}}
{"timestamp":"2026-10-09T01:32:43.793Z","ordinal":13,"type":"response_item","payload":{"type":"custom_tool_call_output","id":"ctco_01a11e4a-8991-72e2-b9e0-eee8262d01af","call_id":"call_71e14f1557fd45378ab5f7186973a336","output":[{"type":"input_text","text":"Script completed\nWall time 0.2 seconds\nOutput:\n"},{"type":"input_text","text":"feature-x\n"}],"internal_chat_message_metadata_passthrough":{"turn_id":"01a11e4a-7fe0-7dc3-8a0c-6fe7d2fd4d48","create_time":1791509563.7938888}},"metadata":{"client_authored":false,"fallback_token_limit_override":12000}}
{"timestamp":"2026-10-09T01:32:43.796Z","ordinal":14,"type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"input_tokens":13962,"cached_input_tokens":11008,"cache_write_input_tokens":0,"output_tokens":36,"reasoning_output_tokens":0,"total_tokens":13998},"last_token_usage":{"input_tokens":13962,"cached_input_tokens":11008,"cache_write_input_tokens":0,"output_tokens":36,"reasoning_output_tokens":0,"total_tokens":13998},"model_context_window":258400},"rate_limits":null}}
{"timestamp":"2026-10-09T01:32:44.681Z","ordinal":15,"type":"response_item","payload":{"type":"custom_tool_call","id":"ctc_0c4f036e98312819016ac8443c7ea8819a840c742a30c9347d","status":"completed","call_id":"call_a37f1360c218494cbd5cdde0400b3e6b","name":"exec","input":"const r = await tools.exec_command({cmd:\"cat notes.txt\", max_output_tokens:1000});\ntext(r.output);\n","internal_chat_message_metadata_passthrough":{"turn_id":"01a11e4a-7fe0-7dc3-8a0c-6fe7d2fd4d48","create_time":1791509564.4966886}}}
{"timestamp":"2026-10-09T01:32:44.714Z","ordinal":16,"type":"token_usage_record","payload":{"thread_id":"01a11e4a-7de0-77a3-9030-5535c9877485","turn_id":"01a11e4a-7fe0-7dc3-8a0c-6fe7d2fd4d48","session_id":"01a11e4a-7de0-77a3-9030-5535c9877485","root_turn_id":"01a11e4a-7fe0-7dc3-8a0c-6fe7d2fd4d48","response_id":"resp_0c4f036e98312819016ac8443c0efc819a8f349bb3ec3b4b14","usage":{"input_tokens":14024,"cached_input_tokens":13056,"cache_write_input_tokens":0,"output_tokens":35,"reasoning_output_tokens":0,"total_tokens":14059},"turn_token_usage":{"input_tokens":27986,"cached_input_tokens":24064,"cache_write_input_tokens":0,"output_tokens":71,"reasoning_output_tokens":0,"total_tokens":28057},"thread_token_usage":{"input_tokens":27986,"cached_input_tokens":24064,"cache_write_input_tokens":0,"output_tokens":71,"reasoning_output_tokens":0,"total_tokens":28057}}}
{"timestamp":"2026-10-09T01:32:44.847Z","ordinal":17,"type":"event_msg","payload":{"type":"item_completed","thread_id":"01a11e4a-7de0-77a3-9030-5535c9877485","turn_id":"01a11e4a-7fe0-7dc3-8a0c-6fe7d2fd4d48","item":{"type":"CommandExecution","id":"exec-0cc28dd9-b1a1-4b58-8230-d2b26119bd6f","process_id":"25101","command":["/bin/zsh","-lc","cat notes.txt"],"cwd":"file:///home/user/proj","parsed_cmd":[{"type":"read","cmd":"cat notes.txt","name":"notes.txt","path":"notes.txt"}],"source":"unified_exec_startup","status":"completed","stdout":"notes: zephyrquartz\n","stderr":"","aggregated_output":"notes: zephyrquartz\n","exit_code":0,"duration":{"secs":0,"nanos":4949},"formatted_output":"notes: zephyrquartz\n"},"started_at_ms":1791509564847,"completed_at_ms":1791509564847}}
{"timestamp":"2026-10-09T01:32:44.855Z","ordinal":18,"type":"response_item","payload":{"type":"custom_tool_call_output","id":"ctco_01a11e4a-8db7-7f31-873d-5a748319e012","call_id":"call_a37f1360c218494cbd5cdde0400b3e6b","output":[{"type":"input_text","text":"Script completed\nWall time 0.2 seconds\nOutput:\n"},{"type":"input_text","text":"notes: zephyrquartz\n"}],"internal_chat_message_metadata_passthrough":{"turn_id":"01a11e4a-7fe0-7dc3-8a0c-6fe7d2fd4d48","create_time":1791509564.8550115}},"metadata":{"client_authored":false,"fallback_token_limit_override":12000}}
{"timestamp":"2026-10-09T01:32:44.857Z","ordinal":19,"type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"input_tokens":27986,"cached_input_tokens":24064,"cache_write_input_tokens":0,"output_tokens":71,"reasoning_output_tokens":0,"total_tokens":28057},"last_token_usage":{"input_tokens":14024,"cached_input_tokens":13056,"cache_write_input_tokens":0,"output_tokens":35,"reasoning_output_tokens":0,"total_tokens":14059},"model_context_window":258400},"rate_limits":null}}
{"timestamp":"2026-10-09T01:32:46.225Z","ordinal":20,"type":"response_item","payload":{"type":"custom_tool_call","id":"ctc_0c4f036e98312819016ac8443dbc14819a9878da2fa2ed8752","status":"completed","call_id":"call_875a60c88e0b4b6b91635c321bc9a227","name":"exec","input":"const r = await tools.exec_command({cmd:\"cat calc.py\", max_output_tokens:1000});\ntext(r.output);\n","internal_chat_message_metadata_passthrough":{"turn_id":"01a11e4a-7fe0-7dc3-8a0c-6fe7d2fd4d48","create_time":1791509565.7366378}}}
{"timestamp":"2026-10-09T01:32:46.303Z","ordinal":21,"type":"token_usage_record","payload":{"thread_id":"01a11e4a-7de0-77a3-9030-5535c9877485","turn_id":"01a11e4a-7fe0-7dc3-8a0c-6fe7d2fd4d48","session_id":"01a11e4a-7de0-77a3-9030-5535c9877485","root_turn_id":"01a11e4a-7fe0-7dc3-8a0c-6fe7d2fd4d48","response_id":"resp_0c4f036e98312819016ac8443d0dcc819a85ec57c40ad27682","usage":{"input_tokens":14089,"cached_input_tokens":13056,"cache_write_input_tokens":0,"output_tokens":35,"reasoning_output_tokens":0,"total_tokens":14124},"turn_token_usage":{"input_tokens":42075,"cached_input_tokens":37120,"cache_write_input_tokens":0,"output_tokens":106,"reasoning_output_tokens":0,"total_tokens":42181},"thread_token_usage":{"input_tokens":42075,"cached_input_tokens":37120,"cache_write_input_tokens":0,"output_tokens":106,"reasoning_output_tokens":0,"total_tokens":42181}}}
{"timestamp":"2026-10-09T01:32:46.391Z","ordinal":22,"type":"event_msg","payload":{"type":"item_completed","thread_id":"01a11e4a-7de0-77a3-9030-5535c9877485","turn_id":"01a11e4a-7fe0-7dc3-8a0c-6fe7d2fd4d48","item":{"type":"CommandExecution","id":"exec-5a9050eb-ceb4-406c-8a5a-2b015ef367fe","process_id":"23701","command":["/bin/zsh","-lc","cat calc.py"],"cwd":"file:///home/user/proj","parsed_cmd":[{"type":"read","cmd":"cat calc.py","name":"calc.py","path":"calc.py"}],"source":"unified_exec_startup","status":"completed","stdout":"def add(a, b):\n return a + b\n\ndef sub(a, b):\n return a - b\n","stderr":"","aggregated_output":"def add(a, b):\n return a + b\n\ndef sub(a, b):\n return a - b\n","exit_code":0,"duration":{"secs":0,"nanos":4749},"formatted_output":"def add(a, b):\n return a + b\n\ndef sub(a, b):\n return a - b\n"},"started_at_ms":1791509566391,"completed_at_ms":1791509566391}}
{"timestamp":"2026-10-09T01:32:46.400Z","ordinal":23,"type":"response_item","payload":{"type":"custom_tool_call_output","id":"ctco_01a11e4a-93c0-7d72-a0b0-f32d7793bc34","call_id":"call_875a60c88e0b4b6b91635c321bc9a227","output":[{"type":"input_text","text":"Script completed\nWall time 0.2 seconds\nOutput:\n"},{"type":"input_text","text":"def add(a, b):\n return a + b\n\ndef sub(a, b):\n return a - b\n"}],"internal_chat_message_metadata_passthrough":{"turn_id":"01a11e4a-7fe0-7dc3-8a0c-6fe7d2fd4d48","create_time":1791509566.4000916}},"metadata":{"client_authored":false,"fallback_token_limit_override":12000}}
{"timestamp":"2026-10-09T01:32:46.402Z","ordinal":24,"type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"input_tokens":42075,"cached_input_tokens":37120,"cache_write_input_tokens":0,"output_tokens":106,"reasoning_output_tokens":0,"total_tokens":42181},"last_token_usage":{"input_tokens":14089,"cached_input_tokens":13056,"cache_write_input_tokens":0,"output_tokens":35,"reasoning_output_tokens":0,"total_tokens":14124},"model_context_window":258400},"rate_limits":null}}
{"timestamp":"2026-10-09T01:32:48.816Z","ordinal":25,"type":"response_item","payload":{"type":"custom_tool_call","id":"ctc_0c4f036e98312819016ac84440118c819ab639e50c61a0d092","status":"completed","call_id":"call_b9854e37b255482a8e0760061b40aa72","name":"exec","input":"const patch = \"*** Begin Patch\\n*** Update File: calc.py\\n@@\\n def sub(a, b):\\n return a - b\\n+\\n+def mul(a, b):\\n+ return a * b\\n*** End Patch\";\nconst r = await tools.apply_patch(patch);\ntext(r);\n","internal_chat_message_metadata_passthrough":{"turn_id":"01a11e4a-7fe0-7dc3-8a0c-6fe7d2fd4d48","create_time":1791509568.0707211}}}
{"timestamp":"2026-10-09T01:32:48.892Z","ordinal":26,"type":"token_usage_record","payload":{"thread_id":"01a11e4a-7de0-77a3-9030-5535c9877485","turn_id":"01a11e4a-7fe0-7dc3-8a0c-6fe7d2fd4d48","session_id":"01a11e4a-7de0-77a3-9030-5535c9877485","root_turn_id":"01a11e4a-7fe0-7dc3-8a0c-6fe7d2fd4d48","response_id":"resp_0c4f036e98312819016ac8443f5a28819ab932f184d9f95e87","usage":{"input_tokens":14171,"cached_input_tokens":11008,"cache_write_input_tokens":0,"output_tokens":76,"reasoning_output_tokens":0,"total_tokens":14247},"turn_token_usage":{"input_tokens":56246,"cached_input_tokens":48128,"cache_write_input_tokens":0,"output_tokens":182,"reasoning_output_tokens":0,"total_tokens":56428},"thread_token_usage":{"input_tokens":56246,"cached_input_tokens":48128,"cache_write_input_tokens":0,"output_tokens":182,"reasoning_output_tokens":0,"total_tokens":56428}}}
{"timestamp":"2026-10-09T01:32:48.919Z","ordinal":27,"type":"event_msg","payload":{"type":"item_completed","thread_id":"01a11e4a-7de0-77a3-9030-5535c9877485","turn_id":"01a11e4a-7fe0-7dc3-8a0c-6fe7d2fd4d48","item":{"type":"FileChange","id":"exec-dbb70cd4-3cec-41f3-acb0-cce249b2aada","changes":{"/home/user/proj/calc.py":{"type":"update","unified_diff":"@@ -5 +5,4 @@\n return a - b\n+\n+def mul(a, b):\n+ return a * b\n","move_path":null}},"status":"completed","stdout":"Success. Updated the following files:\nM calc.py\n","stderr":""},"started_at_ms":1791509568835,"completed_at_ms":1791509568919}}
{"timestamp":"2026-10-09T01:32:48.928Z","ordinal":28,"type":"response_item","payload":{"type":"custom_tool_call_output","id":"ctco_01a11e4a-9da0-7bb0-9fd0-02cb8d24729e","call_id":"call_b9854e37b255482a8e0760061b40aa72","output":[{"type":"input_text","text":"Script completed\nWall time 0.1 seconds\nOutput:\n"},{"type":"input_text","text":"{}"}],"internal_chat_message_metadata_passthrough":{"turn_id":"01a11e4a-7fe0-7dc3-8a0c-6fe7d2fd4d48","create_time":1791509568.9282346}},"metadata":{"client_authored":false,"fallback_token_limit_override":12000}}
{"timestamp":"2026-10-09T01:32:48.930Z","ordinal":29,"type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"input_tokens":56246,"cached_input_tokens":48128,"cache_write_input_tokens":0,"output_tokens":182,"reasoning_output_tokens":0,"total_tokens":56428},"last_token_usage":{"input_tokens":14171,"cached_input_tokens":11008,"cache_write_input_tokens":0,"output_tokens":76,"reasoning_output_tokens":0,"total_tokens":14247},"model_context_window":258400},"rate_limits":null}}
{"timestamp":"2026-10-09T01:32:49.985Z","ordinal":30,"type":"response_item","payload":{"type":"custom_tool_call","id":"ctc_0c4f036e98312819016ac84441af48819ab8370c334234ff04","status":"completed","call_id":"call_e5f92d9061964574bde820232646737f","name":"exec","input":"const patch = \"*** Begin Patch\\n*** Add File: fox.txt\\n+foxtrotpear\\n*** End Patch\";\nconst r = await tools.apply_patch(patch);\ntext(r);\n","internal_chat_message_metadata_passthrough":{"turn_id":"01a11e4a-7fe0-7dc3-8a0c-6fe7d2fd4d48","create_time":1791509569.6951296}}}
{"timestamp":"2026-10-09T01:32:50.043Z","ordinal":31,"type":"token_usage_record","payload":{"thread_id":"01a11e4a-7de0-77a3-9030-5535c9877485","turn_id":"01a11e4a-7fe0-7dc3-8a0c-6fe7d2fd4d48","session_id":"01a11e4a-7de0-77a3-9030-5535c9877485","root_turn_id":"01a11e4a-7fe0-7dc3-8a0c-6fe7d2fd4d48","response_id":"resp_0c4f036e98312819016ac844413434819a843e971579bd83dd","usage":{"input_tokens":14271,"cached_input_tokens":13056,"cache_write_input_tokens":0,"output_tokens":47,"reasoning_output_tokens":0,"total_tokens":14318},"turn_token_usage":{"input_tokens":70517,"cached_input_tokens":61184,"cache_write_input_tokens":0,"output_tokens":229,"reasoning_output_tokens":0,"total_tokens":70746},"thread_token_usage":{"input_tokens":70517,"cached_input_tokens":61184,"cache_write_input_tokens":0,"output_tokens":229,"reasoning_output_tokens":0,"total_tokens":70746}}}
{"timestamp":"2026-10-09T01:32:50.075Z","ordinal":32,"type":"event_msg","payload":{"type":"item_completed","thread_id":"01a11e4a-7de0-77a3-9030-5535c9877485","turn_id":"01a11e4a-7fe0-7dc3-8a0c-6fe7d2fd4d48","item":{"type":"FileChange","id":"exec-df2a946a-e49c-453a-a5a1-9af8d26eeae8","changes":{"/home/user/proj/fox.txt":{"type":"add","content":"foxtrotpear\n"}},"status":"completed","stdout":"Success. Updated the following files:\nA fox.txt\n","stderr":""},"started_at_ms":1791509570006,"completed_at_ms":1791509570075}}
{"timestamp":"2026-10-09T01:32:50.084Z","ordinal":33,"type":"response_item","payload":{"type":"custom_tool_call_output","id":"ctco_01a11e4a-a224-78f0-90d4-42e5d940b074","call_id":"call_e5f92d9061964574bde820232646737f","output":[{"type":"input_text","text":"Script completed\nWall time 0.1 seconds\nOutput:\n"},{"type":"input_text","text":"{}"}],"internal_chat_message_metadata_passthrough":{"turn_id":"01a11e4a-7fe0-7dc3-8a0c-6fe7d2fd4d48","create_time":1791509570.0843356}},"metadata":{"client_authored":false,"fallback_token_limit_override":12000}}
{"timestamp":"2026-10-09T01:32:50.086Z","ordinal":34,"type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"input_tokens":70517,"cached_input_tokens":61184,"cache_write_input_tokens":0,"output_tokens":229,"reasoning_output_tokens":0,"total_tokens":70746},"last_token_usage":{"input_tokens":14271,"cached_input_tokens":13056,"cache_write_input_tokens":0,"output_tokens":47,"reasoning_output_tokens":0,"total_tokens":14318},"model_context_window":258400},"rate_limits":null}}
{"timestamp":"2026-10-09T01:32:53.373Z","ordinal":35,"type":"response_item","payload":{"type":"custom_tool_call","id":"ctc_0c4f036e98312819016ac8444523ac819a9416dff30f2d2df0","status":"completed","call_id":"call_3243009da42b4e1b916288e04872b89c","name":"exec","input":"const r = await tools.exec_command({cmd:\"ls /nonexistent_foxtrotpear\", max_output_tokens:200});\ntext(r.output);\n","internal_chat_message_metadata_passthrough":{"turn_id":"01a11e4a-7fe0-7dc3-8a0c-6fe7d2fd4d48","create_time":1791509573.1412773}}}
{"timestamp":"2026-10-09T01:32:53.409Z","ordinal":36,"type":"token_usage_record","payload":{"thread_id":"01a11e4a-7de0-77a3-9030-5535c9877485","turn_id":"01a11e4a-7fe0-7dc3-8a0c-6fe7d2fd4d48","session_id":"01a11e4a-7de0-77a3-9030-5535c9877485","root_turn_id":"01a11e4a-7fe0-7dc3-8a0c-6fe7d2fd4d48","response_id":"resp_0c4f036e98312819016ac84444bd04819ab6ae8cf2b532d641","usage":{"input_tokens":14342,"cached_input_tokens":14080,"cache_write_input_tokens":0,"output_tokens":40,"reasoning_output_tokens":0,"total_tokens":14382},"turn_token_usage":{"input_tokens":84859,"cached_input_tokens":75264,"cache_write_input_tokens":0,"output_tokens":269,"reasoning_output_tokens":0,"total_tokens":85128},"thread_token_usage":{"input_tokens":84859,"cached_input_tokens":75264,"cache_write_input_tokens":0,"output_tokens":269,"reasoning_output_tokens":0,"total_tokens":85128}}}
{"timestamp":"2026-10-09T01:32:53.550Z","ordinal":37,"type":"event_msg","payload":{"type":"item_completed","thread_id":"01a11e4a-7de0-77a3-9030-5535c9877485","turn_id":"01a11e4a-7fe0-7dc3-8a0c-6fe7d2fd4d48","item":{"type":"CommandExecution","id":"exec-4ba1d5ca-afce-484c-9151-6ffbca09fb0d","process_id":"70822","command":["/bin/zsh","-lc","ls /nonexistent_foxtrotpear"],"cwd":"file:///home/user/proj","parsed_cmd":[{"type":"list_files","cmd":"ls /nonexistent_foxtrotpear","path":"nonexistent_foxtrotpear"}],"source":"unified_exec_startup","status":"failed","stdout":"ls: cannot access '/nonexistent_foxtrotpear': No such file or directory\n","stderr":"","aggregated_output":"ls: cannot access '/nonexistent_foxtrotpear': No such file or directory\n","exit_code":2,"duration":{"secs":0,"nanos":7685},"formatted_output":"ls: cannot access '/nonexistent_foxtrotpear': No such file or directory\n"},"started_at_ms":1791509573549,"completed_at_ms":1791509573550}}
{"timestamp":"2026-10-09T01:32:53.559Z","ordinal":38,"type":"response_item","payload":{"type":"custom_tool_call_output","id":"ctco_01a11e4a-afb7-76a0-97a4-ec91e8a97d87","call_id":"call_3243009da42b4e1b916288e04872b89c","output":[{"type":"input_text","text":"Script completed\nWall time 0.2 seconds\nOutput:\n"},{"type":"input_text","text":"ls: cannot access '/nonexistent_foxtrotpear': No such file or directory\n"}],"internal_chat_message_metadata_passthrough":{"turn_id":"01a11e4a-7fe0-7dc3-8a0c-6fe7d2fd4d48","create_time":1791509573.5593219}},"metadata":{"client_authored":false,"fallback_token_limit_override":12000}}
{"timestamp":"2026-10-09T01:32:53.561Z","ordinal":39,"type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"input_tokens":84859,"cached_input_tokens":75264,"cache_write_input_tokens":0,"output_tokens":269,"reasoning_output_tokens":0,"total_tokens":85128},"last_token_usage":{"input_tokens":14342,"cached_input_tokens":14080,"cache_write_input_tokens":0,"output_tokens":40,"reasoning_output_tokens":0,"total_tokens":14382},"model_context_window":258400},"rate_limits":null}}
{"timestamp":"2026-10-09T01:32:54.430Z","ordinal":40,"type":"event_msg","payload":{"type":"item_completed","thread_id":"01a11e4a-7de0-77a3-9030-5535c9877485","turn_id":"01a11e4a-7fe0-7dc3-8a0c-6fe7d2fd4d48","item":{"type":"AgentMessage","id":"msg_0c4f036e98312819016ac844465580819a83cb4d71b2da161c","content":[{"type":"Text","text":"Done — FOXTROTPEAR."}],"phase":"final_answer"},"started_at_ms":1791509574357,"completed_at_ms":1791509574430}}
{"timestamp":"2026-10-09T01:32:54.435Z","ordinal":41,"type":"response_item","payload":{"type":"message","id":"msg_0c4f036e98312819016ac844465580819a83cb4d71b2da161c","role":"assistant","content":[{"type":"output_text","text":"Done — FOXTROTPEAR."}],"phase":"final_answer","internal_chat_message_metadata_passthrough":{"turn_id":"01a11e4a-7fe0-7dc3-8a0c-6fe7d2fd4d48","create_time":1791509574.3363957,"content_item_kinds":["unknown"]}},"metadata":{"retained_source":{"id":{"message_id":"msg_0c4f036e98312819016ac844465580819a83cb4d71b2da161c","turn_id":"01a11e4a-7fe0-7dc3-8a0c-6fe7d2fd4d48","role":"assistant"},"revision":"retained_01a11e4a-b323-7921-8620-64bdbfe31aba","complete":true},"client_authored":false,"user_input_order":1}}
{"timestamp":"2026-10-09T01:32:54.504Z","ordinal":42,"type":"token_usage_record","payload":{"thread_id":"01a11e4a-7de0-77a3-9030-5535c9877485","turn_id":"01a11e4a-7fe0-7dc3-8a0c-6fe7d2fd4d48","session_id":"01a11e4a-7de0-77a3-9030-5535c9877485","root_turn_id":"01a11e4a-7fe0-7dc3-8a0c-6fe7d2fd4d48","response_id":"resp_0c4f036e98312819016ac84445e2d8819a8ea656e403c0d737","usage":{"input_tokens":14424,"cached_input_tokens":13056,"cache_write_input_tokens":0,"output_tokens":12,"reasoning_output_tokens":0,"total_tokens":14436},"turn_token_usage":{"input_tokens":99283,"cached_input_tokens":88320,"cache_write_input_tokens":0,"output_tokens":281,"reasoning_output_tokens":0,"total_tokens":99564},"thread_token_usage":{"input_tokens":99283,"cached_input_tokens":88320,"cache_write_input_tokens":0,"output_tokens":281,"reasoning_output_tokens":0,"total_tokens":99564}}}
{"timestamp":"2026-10-09T01:32:54.508Z","ordinal":43,"type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"input_tokens":99283,"cached_input_tokens":88320,"cache_write_input_tokens":0,"output_tokens":281,"reasoning_output_tokens":0,"total_tokens":99564},"last_token_usage":{"input_tokens":14424,"cached_input_tokens":13056,"cache_write_input_tokens":0,"output_tokens":12,"reasoning_output_tokens":0,"total_tokens":14436},"model_context_window":258400},"rate_limits":null}}
{"timestamp":"2026-10-09T01:32:54.517Z","ordinal":44,"type":"event_msg","payload":{"type":"task_complete","turn_id":"01a11e4a-7fe0-7dc3-8a0c-6fe7d2fd4d48","last_agent_message":"Done — FOXTROTPEAR.","started_at":1791509561,"completed_at":1791509574,"duration_ms":13201,"time_to_first_token_ms":2031}}
crates/atuin-common/tests/fixtures/opencode/v2_export.json
Vendored
+778
View File
@@ -0,0 +1,778 @@
{
"info": {
"id": "ses_ee1b7b1dfffe46nQL0NuU68GcR",
"projectID": "7d479df3519baa04d706cdddd6c38a1b88cd34e3",
"model": {
"id": "big-pickle",
"providerID": "opencode",
"variant": "default"
},
"cost": 0,
"tokens": {
"input": 7475,
"output": 558,
"reasoning": 354,
"cache": {
"read": 82304,
"write": 0
}
},
"outcome": "succeeded",
"time": {
"created": 1791509417520,
"updated": 1791509629685,
"idle": 1791509645461
},
"title": "octest",
"location": {
"directory": "/work/proj"
}
},
"messages": [
{
"id": "msg_11e484e44001y9aBZ5Ki93Ka5M",
"time": {
"created": 1791509417569
},
"text": "\"Marker OCALPHA. Use the bash tool to run: echo ocalpha-out && ls. Then use the read tool to read notes.txt. Then reply with one short sentence containing the word OCALPHADONE.\"",
"files": [],
"type": "user"
},
{
"id": "msg_11e484e73001tHe41p17eurFKY",
"time": {
"created": 1791509417649,
"streamed": 1791509419466,
"completed": 1791509419547
},
"type": "assistant",
"agent": "build",
"model": {
"id": "big-pickle",
"providerID": "opencode",
"variant": "default"
},
"content": [
{
"type": "reasoning",
"text": "Let me follow the instructions.",
"state": {
"reasoningField": "reasoning_content"
},
"time": {
"created": 1791509419190,
"completed": 1791509419454
}
},
{
"type": "tool",
"id": "call_00_9wmbw0r7375sn5egtq8mhas6",
"name": "shell",
"executed": false,
"state": {
"status": "completed",
"input": {
"command": "echo ocalpha-out && ls"
},
"content": [
{
"type": "text",
"text": "ocalpha-out\ngreet.py\nnotes.txt\n"
}
],
"metadata": {
"status": "completed",
"truncated": false,
"exit": 0
}
},
"time": {
"created": 1791509419238,
"ran": 1791509419460,
"completed": 1791509419529
}
},
{
"type": "tool",
"id": "call_01_oxfcqov5rjteil044qz6otjl",
"name": "read",
"executed": false,
"state": {
"status": "completed",
"input": {
"path": "notes.txt"
},
"content": [
{
"type": "text",
"text": "Read file notes.txt, lines 1-2\n1: alpha line\n2: beta line"
}
],
"metadata": {
"truncated": false
}
},
"time": {
"created": 1791509419290,
"ran": 1791509419464,
"completed": 1791509419481
}
}
],
"snapshot": {
"start": "2d340ce5eeb94eb91fe0c17056193a24b37a999a",
"end": "2d340ce5eeb94eb91fe0c17056193a24b37a999a",
"files": []
},
"finish": "tool-calls",
"rawFinish": "tool_calls",
"cost": 0,
"tokens": {
"input": 5456,
"output": 71,
"reasoning": 6,
"cache": {
"read": 2048,
"write": 0
}
}
},
{
"id": "msg_11e48562d001LoisbfgKtTjN6N",
"time": {
"created": 1791509419595,
"streamed": 1791509421018,
"completed": 1791509421045
},
"type": "assistant",
"agent": "build",
"model": {
"id": "big-pickle",
"providerID": "opencode",
"variant": "default"
},
"content": [
{
"type": "text",
"text": "OCALPHADONE: the command printed `ocalpha-out`, the directory holds `greet.py` and `notes.txt`, and `notes.txt` contains \"alpha line\" and \"beta line\"."
}
],
"snapshot": {
"start": "2d340ce5eeb94eb91fe0c17056193a24b37a999a",
"end": "2d340ce5eeb94eb91fe0c17056193a24b37a999a",
"files": []
},
"finish": "stop",
"rawFinish": "stop",
"cost": 0,
"tokens": {
"input": 208,
"output": 43,
"reasoning": 0,
"cache": {
"read": 7424,
"write": 0
}
}
},
{
"id": "msg_11e485bfc001J3OIfsfEAZTPqW",
"time": {
"created": 1791509421052
},
"type": "idle",
"outcome": "succeeded"
},
{
"id": "msg_11e4a7b12001229UBmORgwKQSM",
"time": {
"created": 1791509560106
},
"text": "\"Marker OCBRAVO. Do these steps with tools, one at a time: 1) create a new file created.txt containing the single line OCCREATED. 2) edit greet.py replacing the string \\\"hello \\\" with \\\"OCEDITED hello \\\". 3) run the shell command: ls /nonexistent-ocfail. It will fail; that is expected. Then reply with one short sentence containing OCBRAVODONE.\"",
"files": [],
"type": "user"
},
{
"id": "msg_11e4a7b34001CaXaW6v4hhIuuZ",
"time": {
"created": 1791509560154,
"streamed": 1791509562068,
"completed": 1791509562109
},
"type": "assistant",
"agent": "build",
"model": {
"id": "big-pickle",
"providerID": "opencode",
"variant": "default"
},
"content": [
{
"type": "reasoning",
"text": "The user wants me to do steps one at a time with tools. Let me first create the file, then edit greet.py, then run the shell command. Let me check greet.py first? The instruction says edit greet.py replacing \"hello \" with \"OCEDITED hello \". Let me just do it.\n\nStep 1: create created.txt.",
"state": {
"reasoningField": "reasoning_content"
},
"time": {
"created": 1791509561418,
"completed": 1791509562060
}
},
{
"type": "tool",
"id": "call_00_mny7won4f7szrw7m19260lbs",
"name": "write",
"executed": false,
"state": {
"status": "completed",
"input": {
"path": "created.txt",
"content": "OCCREATED\n"
},
"content": [
{
"type": "text",
"text": "Created file successfully: created.txt"
}
],
"metadata": {
"truncated": false
}
},
"time": {
"created": 1791509561856,
"ran": 1791509562065,
"completed": 1791509562075
}
}
],
"snapshot": {
"start": "2d340ce5eeb94eb91fe0c17056193a24b37a999a",
"end": "9681da6b2f70128b8453a30f7202ca5ba46f75e8",
"files": [
"created.txt"
]
},
"finish": "tool-calls",
"rawFinish": "tool_calls",
"cost": 0,
"tokens": {
"input": 222,
"output": 58,
"reasoning": 70,
"cache": {
"read": 7552,
"write": 0
}
}
},
{
"id": "msg_11e4a830d001Uzcu2018zN2OJ2",
"time": {
"created": 1791509562150,
"streamed": 1791509564187,
"completed": 1791509564214
},
"type": "assistant",
"agent": "build",
"model": {
"id": "big-pickle",
"providerID": "opencode",
"variant": "default"
},
"content": [
{
"type": "reasoning",
"text": "Step 2: edit greet.py. Let me read it first? The edit tool requires exact text. The user says replace \"hello \" with \"OCEDITED hello \". I need to know exact content. Let me read greet.py first to get exact string. But the instruction says one at a time, tools. Reading is fine.\n\nActually let me just read it.",
"state": {
"reasoningField": "reasoning_content"
},
"time": {
"created": 1791509563598,
"completed": 1791509564178
}
},
{
"type": "tool",
"id": "call_00_tjc6jgeah1p8x0i8i58fluqk",
"name": "read",
"executed": false,
"state": {
"status": "completed",
"input": {
"path": "greet.py"
},
"content": [
{
"type": "text",
"text": "Read file greet.py, lines 1-2\n1: def greet(name):\n2: return \"hello \" + name"
}
],
"metadata": {
"truncated": false
}
},
"time": {
"created": 1791509564065,
"ran": 1791509564184,
"completed": 1791509564195
}
}
],
"snapshot": {
"start": "9681da6b2f70128b8453a30f7202ca5ba46f75e8",
"end": "9681da6b2f70128b8453a30f7202ca5ba46f75e8",
"files": []
},
"finish": "tool-calls",
"rawFinish": "tool_calls",
"cost": 0,
"tokens": {
"input": 240,
"output": 39,
"reasoning": 75,
"cache": {
"read": 7680,
"write": 0
}
}
},
{
"id": "msg_11e4a8b44001czDhsKOsToTGf3",
"time": {
"created": 1791509564280,
"streamed": 1791509566013,
"completed": 1791509566063
},
"type": "assistant",
"agent": "build",
"model": {
"id": "big-pickle",
"providerID": "opencode",
"variant": "default"
},
"content": [
{
"type": "tool",
"id": "call_00_qr44ecq9v2mu34w7df8bfyl7",
"name": "edit",
"executed": false,
"state": {
"status": "completed",
"input": {
"path": "greet.py",
"oldString": "\"hello \"",
"newString": "\"OCEDITED hello \""
},
"content": [
{
"type": "text",
"text": "Edited greet.py (1 replacement)"
}
],
"metadata": {
"files": [
{
"file": "greet.py",
"patch": "Index: greet.py\n===================================================================\n--- greet.py\n+++ greet.py\n@@ -1,2 +1,2 @@\n def greet(name):\n- return \"hello \" + name\n+ return \"OCEDITED hello \" + name\n",
"status": "modified",
"additions": 1,
"deletions": 1
}
],
"truncated": false
}
},
"time": {
"created": 1791509565808,
"ran": 1791509566010,
"completed": 1791509566018
}
}
],
"snapshot": {
"start": "9681da6b2f70128b8453a30f7202ca5ba46f75e8",
"end": "2455321451b4af4d863c3855b4cdb2b8947e6331",
"files": [
"greet.py"
]
},
"finish": "tool-calls",
"rawFinish": "tool_calls",
"cost": 0,
"tokens": {
"input": 264,
"output": 78,
"reasoning": 0,
"cache": {
"read": 7808,
"write": 0
}
}
},
{
"id": "msg_11e4a927d0014iEkQGd3okYmP6",
"time": {
"created": 1791509566102,
"streamed": 1791509567718,
"completed": 1791509567750
},
"type": "assistant",
"agent": "build",
"model": {
"id": "big-pickle",
"providerID": "opencode",
"variant": "default"
},
"content": [
{
"type": "tool",
"id": "call_00_fl3ofucqjfie07qtfvoq11ls",
"name": "shell",
"executed": false,
"state": {
"status": "completed",
"input": {
"command": "ls /nonexistent-ocfail"
},
"content": [
{
"type": "text",
"text": "ls: cannot access '/nonexistent-ocfail': No such file or directory\n"
},
{
"type": "text",
"text": "Exited with code 2"
}
],
"metadata": {
"status": "completed",
"truncated": false,
"exit": 2
}
},
"time": {
"created": 1791509567554,
"ran": 1791509567715,
"completed": 1791509567736
}
}
],
"snapshot": {
"start": "2455321451b4af4d863c3855b4cdb2b8947e6331",
"end": "2455321451b4af4d863c3855b4cdb2b8947e6331",
"files": []
},
"finish": "tool-calls",
"rawFinish": "tool_calls",
"cost": 0,
"tokens": {
"input": 104,
"output": 44,
"reasoning": 0,
"cache": {
"read": 8064,
"write": 0
}
}
},
{
"id": "msg_11e4a9914001nHuzXWhRL9f9gV",
"time": {
"created": 1791509567791,
"streamed": 1791509569332,
"completed": 1791509569343
},
"type": "assistant",
"agent": "build",
"model": {
"id": "big-pickle",
"providerID": "opencode",
"variant": "default"
},
"content": [
{
"type": "text",
"text": "OCBRAVODONE: created `created.txt` with \"OCCREATED\", changed greet.py to `return \"OCEDITED hello \" + name`, and the `ls /nonexistent-ocfail` command failed as expected."
}
],
"snapshot": {
"start": "2455321451b4af4d863c3855b4cdb2b8947e6331",
"end": "2455321451b4af4d863c3855b4cdb2b8947e6331",
"files": []
},
"finish": "stop",
"rawFinish": "stop",
"cost": 0,
"tokens": {
"input": 184,
"output": 51,
"reasoning": 0,
"cache": {
"read": 8064,
"write": 0
}
}
},
{
"id": "msg_11e4a9f43001DuO3wyTyu0FJGg",
"time": {
"created": 1791509569347
},
"type": "idle",
"outcome": "succeeded"
},
{
"id": "msg_11e4b88a1001rPpaETKG8fvrIJ",
"time": {
"created": 1791509629710
},
"text": "\"Marker OCCHARLIE. Run these shell commands one at a time, each as a separate tool call: echo occharlie-one; sleep 4; echo occharlie-two; sleep 4; echo occharlie-three. Then reply with one short sentence containing OCCHARLIEDONE.\"",
"files": [],
"type": "user"
},
{
"id": "msg_11e4b8b1b001sZQXcrCo3kLFDO",
"time": {
"created": 1791509629768,
"streamed": 1791509632666,
"completed": 1791509636773
},
"type": "assistant",
"agent": "build",
"model": {
"id": "big-pickle",
"providerID": "opencode",
"variant": "default"
},
"content": [
{
"type": "reasoning",
"text": "The user wants me to run these shell commands one at a time, each as a separate tool call: echo occharlie-one; sleep 4; echo occharlie-two; sleep 4; echo occharlie-three. Hmm, this is ambiguous. \"one at a time, each as a separate tool call\" \u2014 I think they mean three commands: `echo occharlie-one && sleep 4`, then `echo occharlie-two && sleep 4`, then `echo occharlie-three`. The semicolons separate the sequence but they say one at a time each as separate tool call. So I'll split into three calls.\n\nWait, \"one at a time\" implies sequential, not parallel. So I cannot make them in parallel. Since each has sleep 4, and they must be sequential, I'll do them one after another. That's 3 separate tool calls, each waiting for completion.\n\nLet me do call 1.",
"state": {
"reasoningField": "reasoning_content"
},
"time": {
"created": 1791509631177,
"completed": 1791509632651
}
},
{
"type": "tool",
"id": "call_00_hbd26xs90wkohszm6oc93xe9",
"name": "shell",
"executed": false,
"state": {
"status": "completed",
"input": {
"command": "echo occharlie-one; sleep 4"
},
"content": [
{
"type": "text",
"text": "occharlie-one\n"
}
],
"metadata": {
"status": "completed",
"truncated": false,
"exit": 0
}
},
"time": {
"created": 1791509632532,
"ran": 1791509632662,
"completed": 1791509636744
}
}
],
"snapshot": {
"start": "2455321451b4af4d863c3855b4cdb2b8947e6331",
"end": "2455321451b4af4d863c3855b4cdb2b8947e6331",
"files": []
},
"finish": "tool-calls",
"rawFinish": "tool_calls",
"cost": 0,
"tokens": {
"input": 176,
"output": 46,
"reasoning": 203,
"cache": {
"read": 8192,
"write": 0
}
}
},
{
"id": "msg_11e4ba6c6001EqbdVvOaJFiMaE",
"time": {
"created": 1791509636839,
"streamed": 1791509638597,
"completed": 1791509642655
},
"type": "assistant",
"agent": "build",
"model": {
"id": "big-pickle",
"providerID": "opencode",
"variant": "default"
},
"content": [
{
"type": "tool",
"id": "call_00_ufwu7lcsqa0c1sp4tezqba8l",
"name": "shell",
"executed": false,
"state": {
"status": "completed",
"input": {
"command": "echo occharlie-two; sleep 4"
},
"content": [
{
"type": "text",
"text": "occharlie-two\n"
}
],
"metadata": {
"status": "completed",
"truncated": false,
"exit": 0
}
},
"time": {
"created": 1791509638455,
"ran": 1791509638591,
"completed": 1791509642625
}
}
],
"snapshot": {
"start": "2455321451b4af4d863c3855b4cdb2b8947e6331",
"end": "2455321451b4af4d863c3855b4cdb2b8947e6331",
"files": []
},
"finish": "tool-calls",
"rawFinish": "tool_calls",
"cost": 0,
"tokens": {
"input": 315,
"output": 46,
"reasoning": 0,
"cache": {
"read": 8320,
"write": 0
}
}
},
{
"id": "msg_11e4bbdb9001OrfcTiVfOWxcAL",
"time": {
"created": 1791509642713,
"streamed": 1791509644045,
"completed": 1791509644085
},
"type": "assistant",
"agent": "build",
"model": {
"id": "big-pickle",
"providerID": "opencode",
"variant": "default"
},
"content": [
{
"type": "tool",
"id": "call_00_o80zozm0csfwz1kqrecxoabt",
"name": "shell",
"executed": false,
"state": {
"status": "completed",
"input": {
"command": "echo occharlie-three"
},
"content": [
{
"type": "text",
"text": "occharlie-three\n"
}
],
"metadata": {
"status": "completed",
"truncated": false,
"exit": 0
}
},
"time": {
"created": 1791509643879,
"ran": 1791509644041,
"completed": 1791509644065
}
}
],
"snapshot": {
"start": "2455321451b4af4d863c3855b4cdb2b8947e6331",
"end": "2455321451b4af4d863c3855b4cdb2b8947e6331",
"files": []
},
"finish": "tool-calls",
"rawFinish": "tool_calls",
"cost": 0,
"tokens": {
"input": 123,
"output": 42,
"reasoning": 0,
"cache": {
"read": 8576,
"write": 0
}
}
},
{
"id": "msg_11e4bc34c0015OkyREcrN5wc71",
"time": {
"created": 1791509644138,
"streamed": 1791509645428,
"completed": 1791509645453
},
"type": "assistant",
"agent": "build",
"model": {
"id": "big-pickle",
"providerID": "opencode",
"variant": "default"
},
"content": [
{
"type": "text",
"text": "OCCHARLIEDONE: ran the three commands in separate sequential calls, which printed `occharlie-one`, `occharlie-two`, and `occharlie-three`."
}
],
"snapshot": {
"start": "2455321451b4af4d863c3855b4cdb2b8947e6331",
"end": "2455321451b4af4d863c3855b4cdb2b8947e6331",
"files": []
},
"finish": "stop",
"rawFinish": "stop",
"cost": 0,
"tokens": {
"input": 183,
"output": 40,
"reasoning": 0,
"cache": {
"read": 8576,
"write": 0
}
}
},
{
"id": "msg_11e4bc895001t74Z6FxEfkstCP",
"time": {
"created": 1791509645461
},
"type": "idle",
"outcome": "succeeded"
}
]
}
crates/atuin-daemon/proto/ai/session.proto
+10
View File
@@ -142,6 +142,10 @@ message RebuildSessionsRequest {}
message RebuildSessionsReply {}
message ProjectSyncedRequest {}
message ProjectSyncedReply {}
message ImportSessionsEvent {
oneof event {
ImportSessionsProgress progress = 1;
@@ -177,4 +181,10 @@ service AiSession {
// the daemon starts). FAILED_PRECONDITION when the store failed to open or
// recover at startup.
rpc RebuildSessions(RebuildSessionsRequest) returns (RebuildSessionsReply);
// Project the records another process synced into the record store (`atuin
// sync`), as the daemon does those it syncs itself, once startup recovery is
// over. Returns once the projection is over; it runs to the end even if the
// caller stops waiting, and a failure is left in the daemon's log.
rpc ProjectSynced(ProjectSyncedRequest) returns (ProjectSyncedReply);
}
crates/atuin-daemon/src/client.rs
+19 -4
View File
@@ -27,8 +27,8 @@ 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, ImportSessionsEvent, ImportSessionsRequest,
ListSessionsRequest, RebuildSessionsRequest, SearchSessionsMatch, SearchSessionsRequest,
TailSessionsEvent, TailSessionsRequest,
ListSessionsRequest, ProjectSyncedRequest, RebuildSessionsRequest, SearchSessionsMatch,
SearchSessionsRequest, TailSessionsEvent, TailSessionsRequest,
};
use crate::grpc::history::pb::history_client::HistoryClient as HistoryServiceClient;
use crate::grpc::history::pb::{
@@ -474,6 +474,13 @@ pub struct AiClient {
}
impl AiClient {
/// The stub over `channel`, taking messages of any size: a message the daemon streams holds
/// one captured message, as big as the prompt or reply in it, past tonic's default 4 MiB.
/// The daemon is local and trusted (the server encodes any size by default).
fn service(channel: Channel) -> AiSessionServiceClient<Channel> {
AiSessionServiceClient::new(channel).max_decoding_message_size(usize::MAX)
}
#[cfg(unix)]
pub async fn new(path: PathBuf) -> Result<Self> {
let log_path = path.clone();
@@ -495,7 +502,7 @@ impl AiClient {
})?;
Ok(Self {
client: AiSessionServiceClient::new(channel),
client: Self::service(channel),
})
}
@@ -517,7 +524,7 @@ impl AiClient {
})?;
Ok(Self {
client: AiSessionServiceClient::new(channel),
client: Self::service(channel),
})
}
@@ -636,6 +643,14 @@ impl AiClient {
Ok(())
}
/// Have the daemon project the records this process synced into the record store. Returns
/// once the projection is over, which goes on if this stops waiting; a failure is only in the
/// daemon's log.
pub async fn project_synced(&mut self) -> Result<()> {
self.client.project_synced(ProjectSyncedRequest {}).await?;
Ok(())
}
pub async fn import_sessions(
&mut self,
harness: Option<HarnessKind>,
crates/atuin-daemon/src/grpc/ai/session/mod.rs
+11 -3
View File
@@ -15,9 +15,9 @@ use crate::grpc::ai::session::pb::ai_session_server::AiSession as GrpcService;
use crate::grpc::ai::session::pb::{
GetSessionEvent, GetSessionRequest, HarnessFilterRequest, ImportSessionsEvent,
ImportSessionsProgress, ImportSessionsRequest, ImportSessionsSummary, ListSessionsRequest,
RebuildSessionsReply, RebuildSessionsRequest, SearchSessionsMatch, SearchSessionsRequest,
SessionRefRequest, TailSessionsEvent, TailSessionsRequest, get_session_event,
import_sessions_event, tail_sessions_event,
ProjectSyncedReply, ProjectSyncedRequest, RebuildSessionsReply, RebuildSessionsRequest,
SearchSessionsMatch, SearchSessionsRequest, 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;
@@ -231,6 +231,14 @@ impl GrpcService for Service {
}
}
async fn project_synced(
&self,
_request: Request<ProjectSyncedRequest>,
) -> Result<Response<ProjectSyncedReply>, Status> {
self.capture.project_synced().await;
Ok(Response::new(ProjectSyncedReply {}))
}
async fn import_sessions(
&self,
request: Request<ImportSessionsRequest>,
crates/atuin-daemon/src/session_capture/engine/mod.rs
+343 -69
View File
@@ -1,4 +1,6 @@
use std::any::Any;
use std::collections::{HashMap, HashSet};
use std::panic::AssertUnwindSafe;
use std::sync::Arc;
use std::time::Duration;
@@ -6,11 +8,12 @@ use atuin_client::ai_session::{
Appended, DbError, HarnessKind, HarnessSession, Message, NativeSessionId, Session, SourceId,
};
use atuin_common::harnesstools::AnyHarness;
use atuin_common::harnesstools::session::any::AnySessions;
use atuin_common::harnesstools::session::{
AnyMessage, CaptureError, Checkpoint, RuntimeError, SessionEvent, SessionId, TitleChange,
};
use atuin_common::sync::BlockingPool;
use futures::{Stream, StreamExt};
use futures::{FutureExt, Stream, StreamExt};
use parking_lot::Mutex;
use tokio::task::JoinHandle;
use tokio::time::Instant;
@@ -27,8 +30,18 @@ const LISTENER_RETRY_MAX: Duration = Duration::from_secs(60);
/// How long a live session with rows waiting for a timestamp may stay quiet before they are
/// stored with capture time: a transcript of untimed lines alone (Claude Code title lines)
/// would otherwise hold them until the daemon stops, and never checkpoint past them.
///
/// Quiet means the harness's capture as a whole, not just the session: while it is busy (a
/// backlog of transcripts read in), the session's next, timed, line may be waiting its turn, and
/// would give the rows their time, so they wait. Lines read go first, so through a backlog that
/// always has the next line ready they wait until it drains; through lines that come now and
/// then (other sessions written live), up to [`UNTIMED_GRACE_BUSY`] times as long.
const UNTIMED_GRACE: Duration = Duration::from_secs(5);
/// How many times [`UNTIMED_GRACE`] rows wait for a timestamp at most while lines of other
/// sessions keep coming, when none is waiting to be handled (see [`UNTIMED_GRACE`]).
const UNTIMED_GRACE_BUSY: u32 = 12;
pub struct SessionCaptureEngine {
listeners: Vec<JoinHandle<()>>,
}
@@ -43,54 +56,95 @@ impl SessionCaptureEngine {
};
let kind = HarnessKind::from(harness);
let sink = sink.clone();
listeners.push(tokio::spawn(async move {
// A harness whose session directory does not exist yet (not installed, or never
// run) must not be dropped for the life of the daemon: retry with capped backoff
// until it appears, so capture starts without a restart.
let mut delay = LISTENER_RETRY_START;
let listener = loop {
match sessions.listener() {
Ok(listener) => break listener,
Err(RuntimeError::NotFound(_)) => {
tokio::time::sleep(delay).await;
delay = (delay * 2).min(LISTENER_RETRY_MAX);
}
Err(e) => {
tracing::warn!(
?e,
?kind,
"ai-session listener failed; capture disabled for this harness"
);
return;
}
}
};
// Sessions `events()` has just (re)opened, with the checkpoint each was handed:
// set before the new read yields anything, so its first event finds it here.
let opened: Opened = Arc::default();
let checkpoint = {
let (sink, opened) = (sink.clone(), opened.clone());
move |id: &SessionId| {
let (sink, opened, id) = (sink.clone(), opened.clone(), id.clone());
async move {
let from = checkpoint_of(&sink, kind, &id).await;
opened.lock().insert(id, from);
from
}
}
};
capture(kind, &sink, listener.events(checkpoint), &opened, UNTIMED_GRACE).await;
}));
let listen = move || listen(kind, sink.clone(), sessions.clone());
listeners.push(tokio::spawn(supervise(kind, listen, LISTENER_RETRY_START)));
}
Self { listeners }
}
}
/// Keep one harness's capture running for the life of the engine: `listen` again, after a backoff
/// growing from `retry` (capped at [`LISTENER_RETRY_MAX`]), whenever it ends, fails or panics.
///
/// A harness whose session directory does not exist yet (not installed, or never run) is
/// listened to once it appears, so capture starts without a restart. One whose listener ended
/// (its source gone, or no longer readable) or panicked would otherwise stop capturing until the
/// daemon restarts, and nothing would say so but a log line.
async fn supervise<F, L>(kind: HarnessKind, mut listen: F, retry: Duration)
where
F: FnMut() -> L,
L: Future<Output = Result<(), RuntimeError>>,
{
let mut delay = retry;
loop {
let started = Instant::now();
match AssertUnwindSafe(listen()).catch_unwind().await {
Ok(Err(RuntimeError::NotFound(_))) => {}
Ok(Err(e)) => {
tracing::warn!(?e, ?kind, ?delay, "ai-session listener failed; retrying");
}
Ok(Ok(())) => {
tracing::warn!(?kind, ?delay, "ai-session listener ended; restarting it");
}
Err(panic) => {
tracing::error!(
panic = panic_message(&*panic),
?kind,
?delay,
"ai-session capture panicked; restarting it"
);
}
}
// One that ran a good while before it stopped is not failing over and over.
if started.elapsed() > LISTENER_RETRY_MAX {
delay = retry;
}
tokio::time::sleep(delay).await;
delay = (delay * 2).min(LISTENER_RETRY_MAX);
}
}
/// Capture what one harness's `sessions` write until its listener ends.
async fn listen(
kind: HarnessKind,
sink: Arc<Sink>,
sessions: AnySessions,
) -> Result<(), RuntimeError> {
let listener = sessions.listener()?;
// Sessions `events()` has just (re)opened, with the checkpoint each was handed: set before
// the new read yields anything, so its first event finds it here.
let opened: Opened = Arc::default();
let checkpoint = {
let (sink, opened) = (sink.clone(), opened.clone());
move |id: &SessionId| {
let (sink, opened, id) = (sink.clone(), opened.clone(), id.clone());
async move {
let from = checkpoint_of(&sink, kind, &id).await;
opened.lock().insert(id, from);
from
}
}
};
capture(kind, &sink, listener.events(checkpoint), &opened, UNTIMED_GRACE).await;
Ok(())
}
/// What a panic said, when it said it as text.
fn panic_message(panic: &(dyn Any + Send)) -> &str {
panic
.downcast_ref::<&str>()
.copied()
.or_else(|| panic.downcast_ref::<String>().map(String::as_str))
.unwrap_or("(not text)")
}
/// Capture every event of one harness's listener until it ends. Rows waiting for a timestamp
/// are stored once their session has been quiet for `grace` (see [`UNTIMED_GRACE`]).
///
/// A line whose capture panics is passed over, as a malformed one is: the rows of the session's
/// next lines are stored, and checkpointed past it, so a restart doesn't panic on it again (but
/// for the last line of a session, read again each restart, and passed over again).
async fn capture(
kind: HarnessKind,
sink: &Sink,
@@ -106,9 +160,11 @@ async fn capture(
// Rows of a transcript's content before it was replaced whose append failed: no read can
// bring them back, so they are kept, and retried, and the checkpoint held, until stored.
let mut held: Held = HashMap::new();
// Sessions holding rows that wait for a timestamp: when they may be flushed, and the
// checkpoint past the last line seen.
// Sessions holding rows that wait for a timestamp: when their last line was seen, and the
// checkpoint past it.
let mut untimed: HashMap<SessionId, (Instant, Checkpoint)> = HashMap::new();
// When the listener last yielded anything (see [`UNTIMED_GRACE`]).
let mut last_event = Instant::now();
// Which read of its transcript each session's stream is in ([`Checkpoint::generation`]):
// one that changes under the same stream started over from the beginning of a transcript
// replaced or cut short under it, which the line reader says, as positions can't (the new
@@ -117,10 +173,17 @@ async fn capture(
let mut reads: HashMap<SessionId, u32> = HashMap::new();
loop {
let due = untimed.values().map(|(at, _)| *at).min();
// When a session's rows waiting for a timestamp may be stored with capture time.
let quiet_from = last_event;
let due_of =
move |seen: Instant| (quiet_from + grace).min(seen + grace * UNTIMED_GRACE_BUSY);
let due = untimed.values().map(|(seen, _)| due_of(*seen)).min();
tokio::select! {
// A line already read goes first, as it may give them their timestamp.
biased;
ev = events.next() => {
let Some(ev) = ev else { break };
last_event = Instant::now();
let SessionEvent { session, checkpoint, message } = match ev {
Ok(ev) => ev,
Err(e) => {
@@ -128,32 +191,55 @@ async fn capture(
continue;
}
};
let opened = opened.lock().remove(&session);
let read = reads.insert(session.clone(), checkpoint.generation);
if let Some(from) = opened {
let start = Start::of(from, checkpoint);
// A new read retries whatever a failed append lost.
stuck.remove(&session);
untimed.remove(&session);
warm(sink, &mut enricher, &session, start).await;
} else if read.is_some_and(|read| read != checkpoint.generation) {
restarted(sink, &mut enricher, &mut stuck, &mut held, &mut untimed, &session)
.await;
}
let rows = enricher.capture(&session, &message);
if rows.is_empty() {
if enricher.has_untimed(&session) {
untimed.insert(session, (Instant::now() + grace, checkpoint));
let line = async {
let opened = opened.lock().remove(&session);
#[cfg(test)]
if sink.hook(super::hooks::Point::Capturing).await == super::hooks::Fault::Panic
{
std::panic::panic_any(super::hooks::INJECTED_PANIC);
}
continue;
let read = reads.insert(session.clone(), checkpoint.generation);
if let Some(from) = opened {
let start = Start::of(from, checkpoint);
// A new read retries whatever a failed append lost.
stuck.remove(&session);
untimed.remove(&session);
warm(sink, &mut enricher, &session, start).await;
} else if read.is_some_and(|read| read != checkpoint.generation) {
restarted(sink, &mut enricher, &mut stuck, &mut held, &mut untimed, &session)
.await;
}
let rows = enricher.capture(&session, &message);
if rows.is_empty() {
if enricher.has_untimed(&session) {
untimed.insert(session.clone(), (Instant::now(), checkpoint));
}
return;
}
untimed.remove(&session);
store(sink, &enricher, &mut stuck, &mut held, &session, rows, checkpoint).await;
};
if let Err(panic) = AssertUnwindSafe(line).catch_unwind().await {
tracing::error!(
panic = panic_message(&*panic),
?kind,
%session,
"capturing an ai-session line panicked; passing over it"
);
// Its bookkeeping may be half set up (not warmed yet, or half captured): set
// up again from what the sidecar holds of the session, whichever read of its
// transcript this is, as the lines after this one follow what is stored. Rows
// waiting for a timestamp are stored first, as on a restart.
untimed.remove(&session);
let rows = enricher.finish(&session);
hold(sink, &mut held, &session, rows).await;
warm(sink, &mut enricher, &session, Start::Resumed).await;
}
untimed.remove(&session);
store(sink, &enricher, &mut stuck, &mut held, &session, rows, checkpoint).await;
}
() = tokio::time::sleep_until(due.unwrap_or_else(Instant::now)), if due.is_some() => {
let now = Instant::now();
let quiet: Vec<SessionId> =
untimed.iter().filter(|(_, (at, _))| *at <= now).map(|(s, _)| s.clone()).collect();
untimed.iter().filter(|(_, (seen, _))| due_of(*seen) <= now).map(|(s, _)| s.clone()).collect();
for session in quiet {
let Some((_, checkpoint)) = untimed.remove(&session) else { continue };
let rows = enricher.finish(&session);
@@ -468,6 +554,8 @@ impl Drop for SessionCaptureEngine {
#[cfg(test)]
mod tests {
use std::sync::atomic::{AtomicUsize, Ordering};
use rstest::rstest;
use super::*;
@@ -552,18 +640,201 @@ mod tests {
}))
.unwrap();
let sent = Instant::now();
let opened = Opened::default();
let run = capture(HarnessKind::ClaudeCode, &sink, rx, &opened, GRACE);
let check = async {
tokio::time::sleep(GRACE / 10).await;
assert!(sink.sidecar.get_session(&handle).await.unwrap().is_none(), "flushed early");
tokio::time::sleep(GRACE * 2).await;
let row = sink.sidecar.get_session(&handle).await.unwrap().expect("never stored");
let row = tokio::time::timeout(Duration::from_secs(30), async {
loop {
if let Some(row) = sink.sidecar.get_session(&handle).await.unwrap() {
return row;
}
tokio::time::sleep(GRACE / 10).await;
}
})
.await
.expect("never stored");
assert!(sent.elapsed() >= GRACE, "flushed early");
assert_eq!(row.title.as_deref(), Some("Draft"));
assert_eq!(sink.sidecar.checkpoint(&handle).await.unwrap(), Some(at(40)));
drop(tx);
};
tokio::join!(run, check);
// Checkpointed once the row is stored, by the time capture ends.
assert_eq!(sink.sidecar.checkpoint(&handle).await.unwrap(), Some(at(40)));
}
/// Rows waiting for a timestamp wait while the harness's capture is busy with other
/// sessions, however long a backlog takes past the grace period: the session's next line,
/// once it comes, gives them its time, rather than capture time.
#[rstest]
#[tokio::test]
async fn untimed_rows_wait_out_a_busy_capture_for_a_timestamp() {
const GRACE: Duration = Duration::from_millis(200);
let sink = sink().await;
let line = |raw: &str| AnyMessage::Ccode(serde_json::from_str(raw).unwrap());
let event = |session: &str, end, raw: &str| {
Ok(SessionEvent {
session: SessionId::from(session.to_owned()),
checkpoint: at(end),
message: line(raw),
})
};
let (tx, rx) = futures::channel::mpsc::unbounded();
tx.unbounded_send(event("s1", 40, &title("mine"))).unwrap();
let opened = Opened::default();
let run = capture(HarnessKind::ClaudeCode, &sink, rx, &opened, GRACE);
let feed = async {
// Another session's lines, well past the grace period.
for n in 0..8 {
tokio::time::sleep(GRACE / 2).await;
tx.unbounded_send(event("s2", 100 * (n + 1), &turn(&format!("o{n}"), None, 1)))
.unwrap();
}
tx.unbounded_send(event("s1", 200, &turn("u1", None, 7))).unwrap();
drop(tx);
};
tokio::join!(run, feed);
let handle = handle_of(HarnessKind::ClaudeCode, &SessionId::from("s1".to_owned()));
let rows: Vec<Message> =
futures::TryStreamExt::try_collect(sink.sidecar.messages(&handle)).await.unwrap();
let u1 = rows.iter().find(|row| row.source_id.as_ref() == "u1").expect("u1 stored");
assert_eq!(rows.len(), 2);
assert!(rows.iter().all(|row| row.timestamp == u1.timestamp), "{rows:?}");
}
/// Panics the handling of the first line capture handles.
#[derive(Debug, Default)]
struct PanicOnce(std::sync::atomic::AtomicBool);
impl super::super::hooks::Hooks for PanicOnce {
fn at(
&self,
point: super::super::hooks::Point,
) -> futures::future::BoxFuture<'_, super::super::hooks::Fault> {
use super::super::hooks::{Fault, Point};
let panic = point == Point::Capturing && !self.0.swap(true, Ordering::SeqCst);
Box::pin(async move {
if panic {
Fault::Panic
} else {
Fault::None
}
})
}
}
/// A line whose capture panics is passed over: the harness's capture goes on, and stores
/// and checkpoints past it.
#[rstest]
#[tokio::test]
async fn a_line_whose_capture_panics_is_passed_over() {
let mut sink = sink().await;
sink.hooks = Some(Arc::new(PanicOnce::default()));
let session = SessionId::from("s1".to_owned());
let handle = handle_of(HarnessKind::ClaudeCode, &session);
let (tx, rx) = futures::channel::mpsc::unbounded();
for (end, raw) in [(100, turn("u1", None, 1)), (200, turn("u2", Some("u1"), 2))] {
tx.unbounded_send(Ok(SessionEvent {
session: session.clone(),
checkpoint: at(end),
message: AnyMessage::Ccode(serde_json::from_str(&raw).unwrap()),
}))
.unwrap();
}
drop(tx);
capture(HarnessKind::ClaudeCode, &sink, rx, &Opened::default(), UNTIMED_GRACE).await;
assert_eq!(stored(&sink).await, ["u2"]);
assert_eq!(sink.sidecar.checkpoint(&handle).await.unwrap(), Some(at(200)));
}
/// A line whose capture panics before its resumed session is warmed up leaves the session to
/// be warmed up by the next: an id-less line identical to one stored before the restart is
/// stored under the next ordinal, not taken for that one and dropped.
#[rstest]
#[tokio::test]
async fn a_session_whose_line_panicked_before_warming_up_is_warmed_up_by_the_next(
#[values(0, 1)] generation: u32,
) {
let mut sink = sink().await;
let session = SessionId::from("s1".to_owned());
let handle = handle_of(HarnessKind::ClaudeCode, &session);
let line = |raw: &str| AnyMessage::Ccode(serde_json::from_str(raw).unwrap());
// Captured before the restart: a turn, then a title line.
let mut before = MessageEnricher::new(HarnessKind::ClaudeCode);
for raw in [turn("u1", None, 1), title("mine")] {
for row in
before.capture(&session, &line(&raw)).into_iter().chain(before.finish(&session))
{
sink.append(row).await.unwrap();
}
}
sink.sidecar.set_checkpoint(&handle, at(100)).await.unwrap();
sink.hooks = Some(Arc::new(PanicOnce::default()));
let (tx, rx) = futures::channel::mpsc::unbounded();
for (end, raw) in
[(150, title("mine")), (200, title("mine")), (300, turn("u2", Some("u1"), 2))]
{
tx.unbounded_send(Ok(SessionEvent {
session: session.clone(),
// In a read started over (the transcript replaced under it), or not.
checkpoint: Checkpoint {
generation,
..at(end)
},
message: line(&raw),
}))
.unwrap();
}
drop(tx);
let opened = Opened::default();
opened.lock().insert(session.clone(), Some(at(100)));
capture(HarnessKind::ClaudeCode, &sink, rx, &opened, UNTIMED_GRACE).await;
let ids = stored(&sink).await;
let titles = ids.iter().filter(|id| id.starts_with(SYNTHETIC)).count();
assert_eq!(titles, 2, "the title line after the panic is a new row: {ids:?}");
}
/// A harness's listener is started again, after a backoff, however it stopped: panicked,
/// failed (its sessions not there yet) or ended.
#[rstest]
#[tokio::test]
async fn a_listener_is_started_again_however_it_stops() {
let runs = Arc::new(AtomicUsize::new(0));
let listen = {
let runs = runs.clone();
move || {
let run = runs.fetch_add(1, Ordering::SeqCst);
async move {
match run {
0 => panic!("listener panicked"),
1 => Err(RuntimeError::NotFound("/nowhere".into())),
2 => Ok(()),
_ => futures::future::pending().await,
}
}
}
};
let supervised = supervise(HarnessKind::ClaudeCode, listen, Duration::from_millis(1));
let started = async {
while runs.load(Ordering::SeqCst) < 4 {
tokio::time::sleep(Duration::from_millis(5)).await;
}
};
tokio::time::timeout(Duration::from_secs(10), async {
tokio::select! {
() = supervised => unreachable!("supervision never ends"),
() = started => {}
}
})
.await
.expect("the listener was not started again");
}
/// Byte offsets and row seqs alike: only a first event strictly past the checkpoint handed
@@ -723,6 +994,9 @@ mod tests {
#[case::renamed_over(false, false)]
#[case::truncated_in_place(true, false)]
#[case::renamed_over_with_a_first_line_past_what_was_read(false, true)]
// Written over in place to longer than what was read (pi migrating a transcript on resume):
// neither the file's identity nor its length tells.
#[case::rewritten_in_place_past_what_was_read(true, true)]
#[tokio::test]
async fn a_followed_transcript_replaced_in_place_is_captured_once(
#[case] in_place: bool,
crates/atuin-daemon/src/session_capture/hooks.rs
+3
View File
@@ -15,6 +15,9 @@ pub enum Point {
ReplayIncomplete,
/// A replay has reprojected, before reporting. A panic fault panics it.
ReplayBeforeSettle,
/// Capture is handling a line, its session's read just (re)opened if it was, before its
/// bookkeeping is set up. A panic fault panics the handling.
Capturing,
/// Capture made its record pending, before pushing it. A panic fault panics the capture.
CapturePushing,
/// Capture pushed its record, before projecting it. A panic fault panics the capture.
crates/atuin-daemon/src/session_capture/message_enricher.rs
+26 -4
View File
@@ -72,9 +72,12 @@ impl MessageEnricher {
) {
let mut occurrences = HashMap::new();
for id in synthetic {
if let Some((hash, ordinal)) = parse_synthetic(id) {
// The last ordinal there is can't be counted past: no read made it, so it is malformed.
if let Some((hash, ordinal)) = parse_synthetic(id)
&& let Some(after) = ordinal.checked_add(1)
{
let next = occurrences.entry(hash).or_insert(0);
*next = (*next).max(ordinal + 1);
*next = (*next).max(after);
}
}
let mut replayed = BTreeMap::new();
@@ -119,7 +122,11 @@ impl MessageEnricher {
// A line keyed on its content: the reader cannot know which of identical lines it
// is, and names the first; the newest row of that content captured is the one
// (none: the line made no row, and there is nothing to follow).
Some((hash, _)) => state.occurrences.get(&hash).map(|&n| synthetic_id(hash, n - 1)),
Some((hash, _)) => state
.occurrences
.get(&hash)
.and_then(|n| n.checked_sub(1))
.map(|n| synthetic_id(hash, n)),
None => Some(follows),
};
}
@@ -241,7 +248,7 @@ fn build(
let hash = content_hash(session, m);
let ordinal = state.occurrences.entry(hash).or_insert(0);
let id = synthetic_id(hash, *ordinal);
*ordinal += 1;
*ordinal = ordinal.saturating_add(1);
id
}
};
@@ -747,6 +754,21 @@ mod tests {
assert_eq!(again[0].parent, None);
}
/// A stored id that reads as content-keyed at the last ordinal there is (a native id made to
/// look like one: Claude Code takes any `uuid`) can't be counted past: it is passed over as
/// malformed, and capture goes on.
#[rstest]
fn a_stored_synthetic_id_at_the_last_ordinal_is_passed_over() {
let mut n = MessageEnricher::new(HarnessKind::ClaudeCode);
let untimed = ccode(&serde_json::json!({"type": "ai-title", "aiTitle": "Renamed"}));
let (hash, _) = parse_synthetic(&MessageEnricher::source_id(&session(), &untimed)).unwrap();
n.resume(&session(), None, &[], None, &[synthetic_id(hash, u32::MAX)]);
let rows = rows(&mut n, &session(), &[untimed.clone(), untimed]);
let ids: Vec<_> = rows.into_iter().map(|m| m.source_id).collect();
assert_eq!(ids, [synthetic_id(hash, 0), synthetic_id(hash, 1)]);
}
/// Codex lines name no parent: each row names the row captured before it, a content-keyed
/// one included, passing over lines capture keeps no row of.
#[rstest]
crates/atuin-daemon/src/session_capture/mod.rs
+276 -23
View File
@@ -15,9 +15,9 @@ use std::sync::Arc;
use std::time::Duration;
use atuin_client::ai_session::{
AiSessionDatabase, AiSessionStore, Appended, AtuinSessionId, BuildError, DbError, HarnessKind,
HarnessSession, Message, PushError, ReprojectProgress, SearchTerms, Session, SessionFilter,
SessionMatch,
AiSessionDatabase, AiSessionStore, Appended, AtuinSessionId, BuildError, CaptureMark, DbError,
HarnessKind, HarnessSession, Message, PushError, ReprojectProgress, SearchTerms, Session,
SessionFilter, SessionMatch,
};
use atuin_client::record::sqlite_store::SqliteStore;
use atuin_common::encryption::paseto_v4::Key;
@@ -101,6 +101,8 @@ pub(crate) struct Sink {
/// Calls seen running a command whose output may carry a credential, and not answered yet:
/// their results' output is never kept (see [`Self::withhold_unsafe_output`]).
unsafe_calls: Arc<parking_lot::Mutex<UnsafeCalls>>,
/// Where capture last moved this host's watermark (see [`AiSessionStore::advance_past_capture`]).
capture_mark: Arc<parking_lot::Mutex<Option<CaptureMark>>>,
#[cfg(test)]
hooks: Option<Arc<dyn hooks::Hooks>>,
}
@@ -127,6 +129,7 @@ impl Sink {
tools,
policy: Arc::default(),
unsafe_calls: Arc::default(),
capture_mark: Arc::default(),
#[cfg(test)]
hooks: None,
}
@@ -155,6 +158,12 @@ impl Sink {
BroadcastStream::new(self.tail.subscribe())
}
/// Whether capture keeps the patches of calls (see
/// [`MessageEnricher::keeping_patches`](message_enricher::MessageEnricher::keeping_patches)).
pub(crate) fn keeps_patches(&self) -> bool {
self.tools != ToolCapture::Names
}
/// Persist and project `msg`, unless its logical message is already projected.
///
/// Only while the store is [`StoreState::Ready`]: this waits out a rebuild, and refuses with
@@ -162,24 +171,14 @@ impl Sink {
/// sidecar the dedup gate trusts may then be missing persisted messages.
///
/// Cancel-safe: dropped while it waits for the store or its locks, it has written nothing;
/// once it has them, its dedup check, push and projection run to the end in a task of their
/// own, which dropping this does not stop. (Stopped between its push and its projection, the
/// record would be in the record store but neither projected nor pending, so the next
/// capture of the line would find it missing from the sidecar and push it again.)
/// Whether capture keeps the patches of calls (see [`MessageEnricher::keeping_patches`](message_enricher::MessageEnricher::keeping_patches)).
pub(crate) fn keeps_patches(&self) -> bool {
self.tools != ToolCapture::Names
}
pub(crate) async fn append(&self, mut msg: Message) -> Result<Appended, AppendError> {
// Apply capture policy before either persistence path or the live tail. Keep structural
// rows (even with no content) so parent links and usage accounting remain intact.
/// once it has them, the rest (capture policy, dedup check, push and projection) runs to the
/// end in a task of its own, which dropping this does not stop. (Stopped between its push and
/// its projection, the record would be in the record store but neither projected nor
/// pending, so the next capture of the line would find it missing from the sidecar and push
/// it again; stopped once policy forgot a credential command waiting for its output, the
/// output appended again would be kept.)
pub(crate) async fn append(&self, msg: Message) -> Result<Appended, AppendError> {
let locks = self.lock_ready().await?;
// Under the locks, the sidecar holds every local row (see `withhold_unsafe_output`).
self.withhold_unsafe_output(&mut msg).await?;
sanitize(&mut msg, self.tools, &self.policy);
// Captured here, so on this host; a reproject reads the same from the record envelope.
msg.host = Some(self.records.host_id());
let sink = self.clone();
// The locks go with the task: released when it ends, however it ends (a panic unwinding
// included, which leaves a tokio mutex usable).
@@ -205,8 +204,8 @@ impl Sink {
/// is answered: a session's lines come here in order, live and imported. One made before the
/// daemon (re)started is found among the session's stored rows, the first time the session
/// comes in, which capture's locks (held) make trustworthy; so is where the session runs. A stored call whose input
/// wasn't captured (made while `ai.capture_tools` was off, or dropped) can't be checked, so
/// its output is withheld too.
/// wasn't captured (made while `ai.capture_tools` was off, or dropped), or whose written
/// content was withheld, can't be checked, so its output is withheld too.
async fn withhold_unsafe_output(&self, msg: &mut Message) -> Result<(), AppendError> {
if self.tools == ToolCapture::Names {
return Ok(());
@@ -223,6 +222,7 @@ impl Sink {
match block {
Content::ToolUse(call)
if call.input.is_null()
|| tools::is_withheld(&call.input)
|| ignored
|| self.policy.judge(&call, cwd.as_deref()) != Verdict::Keep =>
{
@@ -309,6 +309,31 @@ impl Sink {
(mut pending, local): CaptureLocks,
mut msg: Message,
) -> Result<Appended, AppendError> {
// Apply capture policy before either persistence path or the live tail. Keep structural
// rows (even with no content) so parent links and usage accounting remain intact.
// Under the locks, the sidecar holds every local row (see `withhold_unsafe_output`).
self.withhold_unsafe_output(&mut msg).await?;
let (tools, policy) = (self.tools, self.policy.clone());
let mut msg = if bulky(&msg) {
// Off the runtime's thread, which the daemon's every other request shares: shaping
// a large tool call can take a while.
tokio::task::spawn_blocking(move || {
sanitize(&mut msg, tools, &policy);
msg
})
.await
.map_err(|err| {
tracing::error!(?err, "sanitizing an ai-session message stopped unexpectedly");
AppendError::Aborted
})?
} else {
sanitize(&mut msg, tools, &policy);
msg
};
// Captured here, so on this host; a reproject reads the same from the record envelope.
msg.host = Some(self.records.host_id());
// As the sidecar would store it, so the record synced never names an instant past it.
msg.timestamp = msg.timestamp.clamp(Message::EARLIEST, Message::LATEST);
self.repair(&mut pending).await?;
// Dedup gate: if this logical message is already projected it is already in the record
// store too, so there is nothing to do. Stable source ids (see MessageEnricher::source_id)
@@ -327,6 +352,13 @@ impl Sink {
// it if so (see `repair`). Nothing else stops this task between its push and its
// projection.
*pending = Some(msg.clone());
// Where the last capture moved the watermark, or else the generation, read before the
// push: an invalidation after it keeps the watermark from moving past it.
let last = *self.capture_mark.lock();
let known = match last {
Some(known) => known,
None => CaptureMark::Unknown(self.sidecar.projection_generation().await?),
};
// Write the record store first: it is the synced source of truth and the sidecar is a
// pure projection of it. If the sidecar write fails afterwards a later rebuild repairs it;
// the reverse ordering could strand a message in the sidecar only -- lost on rebuild and
@@ -335,13 +367,20 @@ impl Sink {
if self.hook(hooks::Point::CapturePushing).await == hooks::Fault::Panic {
std::panic::panic_any(hooks::INJECTED_PANIC);
}
self.records.push(&msg).await?;
let pushed = self.records.push(&msg).await?;
#[cfg(test)]
if self.hook(hooks::Point::CapturePushed).await == hooks::Fault::Panic {
std::panic::panic_any(hooks::INJECTED_PANIC);
}
let appended = self.project_and_broadcast(&msg).await?;
*pending = None;
// So a restart doesn't replay what capture projected since the last; one that doesn't
// move (a reprojection running, say) leaves it to the next.
let moved = self.records.advance_past_capture(&self.sidecar, known, pushed).await;
*self.capture_mark.lock() = moved.unwrap_or_else(|err| {
tracing::warn!(?err, "failed to move the ai-session watermark past a capture");
None
});
drop(local);
drop(pending);
Ok(appended)
@@ -434,6 +473,49 @@ struct UnsafeCalls {
cwds: HashMap<HarnessSession, PathBuf>,
}
/// Whether `msg` holds enough text that [`sanitize`] may take a while: past [`BULKY`] bytes of it.
/// Most rows hold far less, and handing each to a blocking thread costs more than sanitizing it:
/// the thread's own allocator arena holds on to memory, tens of megabytes across a busy ingest.
fn bulky(msg: &Message) -> bool {
fn size(value: &serde_json::Value, left: &mut usize) {
match value {
serde_json::Value::String(text) => *left = left.saturating_sub(text.len()),
serde_json::Value::Array(items) => items.iter().for_each(|item| size(item, left)),
serde_json::Value::Object(object) => {
object.iter().for_each(|(key, item)| {
*left = left.saturating_sub(key.len());
size(item, left);
});
}
serde_json::Value::Null | serde_json::Value::Bool(_) | serde_json::Value::Number(_) => {
*left = left.saturating_sub(8);
}
}
}
let mut left = BULKY;
for block in &msg.content {
match block {
Content::Text(text) | Content::Summary(text) | Content::Error(text) => {
left = left.saturating_sub(text.len());
}
Content::ToolUse(tool) => size(&tool.input, &mut left),
Content::ToolResult(result) => size(&result.output, &mut left),
Content::Patch(patch) => {
let lines = patch.files.iter().flat_map(|f| &f.hunks).flat_map(|h| &h.lines);
left = left.saturating_sub(lines.map(String::len).sum());
}
Content::Reasoning(_) | Content::ReasoningSummary { .. } | Content::Other(_) => {}
}
if left == 0 {
return true;
}
}
false
}
/// See [`bulky`].
const BULKY: usize = 256 * 1024;
/// Retain conversation text, and of tool calls what `tools` keeps: their inputs and outputs,
/// unless `ai.capture_tools` is off, `policy` withholds them, or they are over the size limit
/// (then null, which keeps the wire format without storing them). Reasoning is kept only as a
@@ -751,6 +833,28 @@ impl AiHarnessSessionCapture {
answer.await.unwrap_or(Err(RebuildError::Aborted))
}
/// Project records another process synced into the record store (`atuin sync`), as the sync
/// worker does those it downloads ([`Recovery::project_synced`]), once startup recovery is
/// over. A nop facade has nothing to project.
///
/// Runs to the end once asked for, even if the caller stops waiting (a client disconnecting):
/// what it downloaded is not downloaded again, so nothing would ask a second time. A failure
/// is logged, not returned: the next projection retries what this one left out.
pub async fn project_synced(&self) {
if self.coordinator.is_none() {
return;
}
let recovery = self.recovery();
let sink = self.sink.clone();
let projecting = tokio::spawn(async move {
recovery.finished().await;
recovery.project_synced(&sink.records, &sink.sidecar).await;
});
if let Err(err) = projecting.await {
tracing::error!(?err, "projecting synced ai-session records failed");
}
}
/// A handle for waiting out startup recovery.
#[must_use]
pub fn recovery(&self) -> Recovery {
@@ -1221,6 +1325,72 @@ mod tests {
);
enabled.append(row(7, Role::User, vec![result("off")])).await.unwrap();
assert_last_output_withheld(&sink.sidecar, &rows[0].session).await;
// And a Codex script that wrote a credential file, stored without what it wrote (which
// leaves no call to judge again).
let script = Content::ToolUse(ToolUse {
id: ToolCallId::from("script".to_owned()),
name: "exec".to_owned(),
input: serde_json::json!(
"await tools.apply_patch(\"*** Begin Patch\\n*** Add File: .env\\n+A=1\\n*** End \
Patch\");"
),
});
enabled.append(row(8, Role::Assistant, vec![script])).await.unwrap();
let again = Sink::with_state(
sink.records.clone(),
sink.sidecar.clone(),
ToolCapture::Payloads,
watch::channel(StoreState::Ready).1,
);
again.append(row(9, Role::User, vec![result("script")])).await.unwrap();
assert_last_output_withheld(&sink.sidecar, &rows[0].session).await;
}
/// An append dropped once policy has forgotten the credential command waiting for its output
/// still stores the row, its output withheld: appended again, the output is not kept.
#[rstest]
#[tokio::test]
async fn an_append_dropped_once_its_output_is_withheld_still_withholds_it() {
let mut sink = Sink::new(mem_store().await, AiSessionDatabase::in_memory().await.unwrap());
sink.tools = ToolCapture::Payloads;
let row = |n: u32, role, content| {
let mut msg = sample_message();
msg.id = RecordId(atuin_common::utils::uuid_v7());
msg.source_id = SourceId::from(format!("line-{n}"));
msg.role = role;
msg.content = vec![content];
msg
};
let call = row(
1,
Role::Assistant,
Content::ToolUse(ToolUse {
id: ToolCallId::from("key".to_owned()),
name: "Bash".to_owned(),
input: serde_json::json!({ "command": "atuin key" }),
}),
);
let answer = row(
2,
Role::User,
Content::ToolResult(ToolResult {
call: ToolCallId::from("key".to_owned()),
output: serde_json::json!("abandon ability able about above absent"),
error: false,
}),
);
sink.append(call).await.unwrap();
let mut append = Box::pin(sink.append(answer.clone()));
while futures::poll!(&mut append).is_pending()
&& !sink.unsafe_calls.lock().pending.is_empty()
{
tokio::task::yield_now().await;
}
drop(append);
sink.append(answer.clone()).await.unwrap();
assert_last_output_withheld(&sink.sidecar, &answer.session).await;
}
async fn assert_last_output_withheld(sidecar: &AiSessionDatabase, session: &HarnessSession) {
@@ -1641,6 +1811,36 @@ mod tests {
assert!(sidecar.contains_message(&msg.session, &msg.source_id).await.unwrap());
}
/// Records another process synced in are projected once asked for, even if the asker stops
/// waiting at once: they are not downloaded again, so nothing would ask a second time.
#[rstest]
#[tokio::test]
async fn synced_records_are_projected_though_the_asker_stops_waiting() {
let records = mem_store().await;
let sidecar = AiSessionDatabase::in_memory().await.unwrap();
let capture = AiHarnessSessionCapture::open(
records.clone(),
sidecar.clone(),
false,
ToolCapture::Names,
CapturePolicy::default(),
BlockingPool::new(NonZeroUsize::MIN),
);
assert!(capture.ready().await);
let msg = sample_message();
records.push(&msg).await.unwrap();
// Polled once, then dropped, as a client disconnecting does to its handler.
let _ = tokio::time::timeout(Duration::ZERO, capture.project_synced()).await;
tokio::time::timeout(Duration::from_secs(5), async {
while !sidecar.contains_message(&msg.session, &msg.source_id).await.unwrap() {
tokio::time::sleep(Duration::from_millis(10)).await;
}
})
.await
.expect("the projection stopped with its asker");
}
/// A message of session `session`, from line `source`.
fn message_of(session: &str, source: &str) -> Message {
let mut msg = sample_message();
@@ -1650,6 +1850,36 @@ mod tests {
msg
}
/// Capture moves the watermark past each record it projects, so a restart replays none of
/// them; a capture while a reprojection runs leaves it to the reprojection.
#[rstest]
#[tokio::test]
async fn capture_moves_the_watermark_past_what_it_projects() {
let records = mem_store().await;
let sidecar = AiSessionDatabase::in_memory().await.unwrap();
// Startup recovery, before capture starts.
records.reproject(&sidecar).await.unwrap();
let sink = Sink::new(records.clone(), sidecar.clone());
for source in ["a", "b"] {
sink.append(message_of("s1", source)).await.unwrap();
}
assert_eq!(records.reproject(&sidecar).await.unwrap().replayed, 0);
let replaying = sidecar.lock_reprojection().await;
sink.append(message_of("s1", "c")).await.unwrap();
drop(replaying);
assert_eq!(records.reproject(&sidecar).await.unwrap().replayed, 1);
sink.append(message_of("s1", "d")).await.unwrap();
assert_eq!(records.reproject(&sidecar).await.unwrap().replayed, 0);
// An invalidation (watermarks cleared) after capture last moved it: not moved past it.
sidecar.clear_reproject_watermarks().await.unwrap();
sink.append(message_of("s1", "e")).await.unwrap();
assert_eq!(records.reproject(&sidecar).await.unwrap().replayed, 5);
sink.append(message_of("s1", "f")).await.unwrap();
assert_eq!(records.reproject(&sidecar).await.unwrap().replayed, 0);
}
/// A rebuild (after a purge deleted records) takes out what the deleted records projected and
/// replays the rest, reporting it is rebuilding meanwhile. Capture waits it out: its dedup gate
/// would otherwise push a message the half-rebuilt sidecar has not got back yet a second time.
@@ -2503,6 +2733,29 @@ mod tests {
assert_eq!(sink.append(msg).await.unwrap(), Appended::Duplicate);
}
/// A line naming an instant past the years a date can name is captured at the nearest one
/// stored, so the record synced never carries it.
#[rstest]
#[case::late(time::macros::datetime!(9999-12-31 23:59:59 -23:59), Message::LATEST)]
#[case::early(time::macros::datetime!(-9999-01-01 0:00 +23:59), Message::EARLIEST)]
#[tokio::test]
async fn capture_holds_a_timestamp_to_the_range_stored(
#[case] timestamp: OffsetDateTime,
#[case] want: OffsetDateTime,
) {
let sink = Sink::new(mem_store().await, AiSessionDatabase::in_memory().await.unwrap());
let mut sub = sink.subscribe();
let mut msg = sample_message();
msg.timestamp = timestamp;
sink.append(msg).await.unwrap();
let _started = sub.next().await.unwrap().unwrap();
let SessionTailEvent::Message(pushed) = sub.next().await.unwrap().unwrap() else {
panic!("expected message");
};
assert_eq!(pushed.timestamp, want);
}
/// A restored transcript writes a call captured without its input back as a note in its
/// row's text. Re-captured, that row comes back under the source id already synced, with
/// other content: it is a duplicate, never pushed again, and the synced row stands.
crates/atuin-daemon/src/session_capture/policy.rs
+63 -8
View File
@@ -15,8 +15,8 @@ use regex::RegexSet;
use super::tools;
/// What a call is kept with.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
/// What a call is kept with, from the most kept to the least.
#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord)]
pub(super) enum Verdict {
/// Its input and output, as [`ToolCapture`](super::ToolCapture) keeps them.
Keep,
@@ -39,6 +39,8 @@ pub struct CapturePolicy {
command_filter: CommandFilter,
files: SensitiveFiles,
redactor: Redactor,
/// `redactor` with the user's own patterns alone.
users: Redactor,
}
impl Default for CapturePolicy {
@@ -51,6 +53,7 @@ impl Default for CapturePolicy {
command_filter: CommandFilter::default(),
files: SensitiveFiles::default(),
redactor: Redactor::default(),
users: Redactor::new(Vec::new(), false),
}
}
}
@@ -65,6 +68,7 @@ impl CapturePolicy {
command_filter: settings.output.command_filter().clone(),
files: settings.security.sensitive_files(&settings.key_path),
redactor: settings.security.redactor(true),
users: settings.security.redactor(false),
}
}
@@ -73,9 +77,27 @@ impl CapturePolicy {
self.cwd_filter.is_match(&cwd.to_string_lossy())
}
/// What `call`, made in `cwd`, is kept with.
/// What `call`, made in `cwd`, is kept with: the least any of what it touches is kept with (a
/// Codex script runs several tools at once, see [`Access::all`]). A script whose calls
/// weren't all read can't be judged: it is dropped, what it ran and wrote with it. (The
/// commands it runs still reach shell history, filtered, through Codex's hooks.)
pub(super) fn judge(&self, call: &ToolUse, cwd: Option<&Path>) -> Verdict {
let access = Access::of(&call.name, &call.input);
let (accesses, whole) = Access::all(&call.name, &call.input);
if !whole {
return Verdict::Drop;
}
let verdict = accesses
.iter()
.map(|access| self.judge_access(access, cwd))
.max()
.unwrap_or(Verdict::Keep);
if verdict == Verdict::Keep && tools::runs_output_unsafe(&call.input) {
return Verdict::WithholdOutput;
}
verdict
}
fn judge_access(&self, access: &Access, cwd: Option<&Path>) -> Verdict {
let dir = access.workdir.as_deref().map(Path::new).or(cwd);
if let Some(command) = &access.command {
let ignored = self.history_filter.is_match(command)
@@ -93,10 +115,7 @@ impl CapturePolicy {
|| self.command_filter.is_match(command)
|| self.files.named_in_command(command, dir)
});
if unsafe_command
|| access.reads.iter().any(|path| self.files.matches(path, dir))
|| tools::runs_output_unsafe(&call.input)
{
if unsafe_command || access.reads.iter().any(|path| self.files.matches(path, dir)) {
return Verdict::WithholdOutput;
}
Verdict::Keep
@@ -107,11 +126,26 @@ impl CapturePolicy {
self.files.matches(path, cwd)
}
#[cfg(test)]
pub(super) fn with_redact_patterns(patterns: Vec<regex::Regex>) -> Self {
Self {
redactor: Redactor::new(patterns.clone(), true),
users: Redactor::new(patterns, false),
..Self::default()
}
}
/// `text` redacted, or `None` when that takes too long (see [`REDACT_BUDGET`]): then it is
/// better not kept.
pub(super) fn redact<'a>(&self, text: &'a str) -> Option<Cow<'a, str>> {
self.redactor.redact_within(text, REDACT_BUDGET)
}
/// [`Self::redact`] with the user's own patterns alone, which, unlike the built-in ones, can
/// match across lines, any length.
pub(super) fn redact_users<'a>(&self, text: &'a str) -> Option<Cow<'a, str>> {
self.users.redact_within(text, REDACT_BUDGET)
}
}
#[cfg(test)]
@@ -137,6 +171,7 @@ pub(super) mod tests {
Some(PathBuf::from("/home/me")),
),
redactor: Redactor::default(),
users: Redactor::new(Vec::new(), false),
}
}
@@ -168,6 +203,26 @@ pub(super) mod tests {
#[case::read_example("Read", json!({"file_path": ".env.example"}), Verdict::Keep)]
#[case::ordinary("Bash", json!({"command": "cargo test"}), Verdict::Keep)]
#[case::search("WebSearch", json!({"query": "rust"}), Verdict::Keep)]
// Codex 0.160 runs each tool through a script; what it runs is judged as if run directly.
#[case::script_history_filter("exec", json!("const r = await tools.exec_command({cmd:\"psql -c 'select 1'\",workdir:\"/w\"});\ntext(r.output);\n"), Verdict::Drop)]
#[case::script_command_filter("exec", json!("const r = await tools.exec_command({cmd:\"cat notes.md\", max_output_tokens:200});\ntext(r.output);\n"), Verdict::WithholdOutput)]
#[case::script_cwd_filter("exec", json!("await tools.exec_command({cmd:\"ls\",workdir:\"/secret/x\"})"), Verdict::Drop)]
#[case::script_secret("exec", json!("await tools.exec_command({cmd:\"atuin login -u me -p hunter2\"})"), Verdict::Drop)]
#[case::script_cat_dotenv("exec", json!("await tools.exec_command({cmd:'cat .env'})"), Verdict::WithholdOutput)]
#[case::script_patch_dotenv("exec", json!("const patch = \"*** Begin Patch\\n*** Add File: .env\\n+A=1\\n*** End Patch\";\nconst r = await tools.apply_patch(patch);\ntext(r);\n"), Verdict::WithholdWritten)]
#[case::script_read_key("exec", json!("await tools.view_image({path:\"~/.ssh/id_ed25519\"})"), Verdict::WithholdOutput)]
#[case::script_worst_call("exec", json!("await tools.exec_command({cmd:'cat .env'});\nawait tools.exec_command({cmd:'psql'});"), Verdict::Drop)]
#[case::script_computed("exec", json!("const c = `psql ${db}`;\nawait tools.exec_command({cmd: c});"), Verdict::Drop)]
// A script not read whole is dropped, whatever it may hide, filtered or not (from review and
// validation).
#[case::script_misread_regex("exec", json!("if (ok) /\"/.test(x); await tools.exec_command({cmd:\"psql -c 1\"}); const y = \"z\";"), Verdict::Drop)]
#[case::script_misread_division("exec", json!("let i = 0; i++ / 2; await tools.exec_command({cmd:\"psql\"});"), Verdict::Drop)]
#[case::script_joined_list("exec", json!("await tools.exec_command({cmd: [\"psql\", \"-c 1\"].join(\" \")});"), Verdict::Drop)]
#[case::script_appended("exec", json!("let c = \"ps\"; c += \"ql -c 1\"; await tools.exec_command({cmd: c});"), Verdict::Drop)]
#[case::script_patch_built("exec", json!("const name = (await tools.exec_command({cmd: \"ls\"})).output.trim();\nawait tools.apply_patch(`*** Begin Patch\\n*** Update File: ${name}\\n+X=3\\n*** End Patch`);"), Verdict::Drop)]
#[case::script_split_joined("exec", json!("await tools.exec_command({cmd: \"echo,hi\".split(\",\").join(\" \")});"), Verdict::Drop)]
#[case::script_computed_workdir("exec", json!("await tools.exec_command({cmd: \"ls\", workdir: base + \"/x\"});"), Verdict::Drop)]
#[case::script_ordinary("exec", json!("const r = await tools.exec_command({cmd:\"cargo test\"});\ntext(r.output);"), Verdict::Keep)]
fn what_a_call_is_kept_with(
policy: CapturePolicy,
#[case] name: &str,
crates/atuin-daemon/src/session_capture/tools.rs
+427 -24
View File
@@ -1,6 +1,10 @@
//! What capture keeps of a tool call's input and output (`ai.capture_tools`).
use std::ops::Range;
use std::sync::LazyLock;
use atuin_common::harnesstools::session::{FilePatch, Patch};
use regex::Regex;
use serde_json::{Map, Value};
use super::policy::CapturePolicy;
@@ -20,6 +24,20 @@ const SHORT: usize = 256;
/// payload.
const OVERSHOOT: usize = 8 * SHORT;
/// How many JSON texts deep `shape` reads JSON text as the JSON it is (a Codex call's arguments,
/// a JSON file `cat` in them); deeper, it is withheld, as text hiding what it holds from
/// redaction (its quotes escaped). Each level is shaped again until it fits, so past a level or
/// two the work multiplies for nothing a harness writes.
const JSON_TEXT_DEPTH: usize = 2;
/// How many times `shape` shapes JSON text again to fit it in its room before withholding it.
const JSON_TEXT_ATTEMPTS: usize = 4;
/// Bytes past each cut a long text is redacted with, and then cut without: a secret across a
/// cut is redacted whole, rather than its part on the kept side left behind unrecognised. Longer
/// than any one secret's text (a 4096-bit RSA key in PEM is about 3.2 KiB).
const SECRET_MARGIN: usize = 4 * 1024;
/// What capture keeps of a tool call.
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub enum ToolCapture {
@@ -48,7 +66,7 @@ impl ToolCapture {
Self::Names => *payload = Value::Null,
Self::Payloads => {
let mut budget = TOOL_PAYLOAD_LIMIT;
shape(payload, &mut budget, policy);
shape(payload, &mut budget, policy, 0);
// `shape` keeps the payload's shape, so a harness reads it back as the same kind
// of input or output, and holds it to the limit (give or take the short strings
// and first items it always keeps), but for an object of very many keys, which it
@@ -137,8 +155,9 @@ pub(super) fn runs_output_unsafe(input: &Value) -> bool {
/// `value` with every string redacted and, once `budget` bytes of it (as stored) are kept,
/// clipped (but never one [`SHORT`] enough to be a name), and the items of a list past it
/// dropped (but never its first); media blocks become a note. Its shape stays as it was, so a
/// harness reads it back as the same kind of input or output.
fn shape(value: &mut Value, budget: &mut usize, policy: &CapturePolicy) {
/// harness reads it back as the same kind of input or output. `depth` is how many JSON texts it
/// is in (see [`JSON_TEXT_DEPTH`]).
fn shape(value: &mut Value, budget: &mut usize, policy: &CapturePolicy, depth: usize) {
match value {
Value::Object(object) => match media_note(object) {
Some(note) => {
@@ -165,7 +184,7 @@ fn shape(value: &mut Value, budget: &mut usize, policy: &CapturePolicy) {
if let Value::String(text) = value {
redact_named(key, text, policy);
}
shape(value, budget, policy);
shape(value, budget, policy, depth);
}
}
},
@@ -176,20 +195,18 @@ fn shape(value: &mut Value, budget: &mut usize, policy: &CapturePolicy) {
if *budget == 0 && kept > 0 {
break;
}
shape(item, budget, policy);
shape(item, budget, policy, depth);
kept += 1;
}
items.truncate(kept);
}
Value::String(text) => {
if let Some(shaped) = shape_json_text(text, *budget, policy) {
*text = shaped;
let json = if depth < JSON_TEXT_DEPTH {
shape_json_text(text, *budget, policy, depth)
} else {
*text = match policy.redact(text) {
Some(redacted) => fit(&redacted, (*budget).max(SHORT)),
None => NOT_REDACTED.to_owned(),
};
}
json_text(text).map(|_| NOT_REDACTED.to_owned())
};
*text = json.unwrap_or_else(|| fit_redacted(text, (*budget).max(SHORT), policy));
*budget = budget.saturating_sub(stored_len(text) + 1);
}
Value::Bool(_) | Value::Number(_) | Value::Null => {
@@ -201,14 +218,20 @@ fn shape(value: &mut Value, budget: &mut usize, policy: &CapturePolicy) {
/// JSON text (a Codex function call's arguments, a legacy Codex output), shaped as the JSON it
/// is so it stays JSON: as text, a quoted value's quotes are escaped, which hides it from
/// redaction, and clipping would cut it mid-string. Kept as written when nothing in it changes;
/// `None` when `text` isn't a JSON object or list.
fn shape_json_text(text: &str, budget: usize, policy: &CapturePolicy) -> Option<String> {
/// [`NOT_REDACTED`] when it doesn't fit in its room within [`JSON_TEXT_ATTEMPTS`]; `None` when
/// `text` isn't a JSON object or list. `depth` is how many JSON texts `text` is in.
fn shape_json_text(
text: &str,
budget: usize,
policy: &CapturePolicy,
depth: usize,
) -> Option<String> {
let parsed = json_text(text)?;
let room = budget.max(SHORT);
let mut allowance = budget;
loop {
for _ in 0..JSON_TEXT_ATTEMPTS {
let mut shaped = parsed.clone();
shape(&mut shaped, &mut allowance.clone(), policy);
shape(&mut shaped, &mut allowance.clone(), policy, depth + 1);
let shaped = if shaped == parsed {
text.to_owned()
} else {
@@ -222,6 +245,7 @@ fn shape_json_text(text: &str, budget: usize, policy: &CapturePolicy) -> Option<
}
allowance = (allowance * room / size).min(allowance - allowance.div_ceil(8));
}
Some(NOT_REDACTED.to_owned())
}
/// `text` parsed, when it is a JSON object or list.
@@ -243,21 +267,190 @@ fn fit(text: &str, room: usize) -> String {
if text.len() <= SHORT {
return text.to_owned();
}
fit_with(room, |keep| Some(clip(text, keep))).expect("clipping never fails")
}
/// What `clip` keeps of about `keep` bytes, for the most `keep` whose result takes no more than
/// `room` bytes stored; `None` if `clip` fails.
fn fit_with(room: usize, clip: impl Fn(usize) -> Option<String>) -> Option<String> {
let mut keep = room;
loop {
let kept = clip(text, keep);
let kept = clip(keep)?;
let size = stored_len(&kept);
if size <= room || keep <= SHORT / 4 {
return kept;
return Some(kept);
}
keep = (keep * room / size).min(keep - keep.div_ceil(8));
}
}
/// `text` redacted as `policy` does and [`fit`] to `room`, or [`NOT_REDACTED`] when redacting
/// takes too long.
///
/// A text longer than could be kept is cut first, so it takes no longer to redact than one that
/// fits: only its start and end are redacted (see [`redacted_start`]), and the note between them
/// is never part of what is redacted. The user's own patterns, which can match across a cut, are
/// redacted in all of it first.
fn fit_redacted(text: &str, room: usize, policy: &CapturePolicy) -> String {
let fitted = if text.len() <= room + NOTE + 2 * SECRET_MARGIN {
policy.redact(text).map(|redacted| fit(&redacted, room))
} else {
policy.redact_users(text).and_then(|text| fit_clipped(&text, room, policy))
};
fitted.unwrap_or_else(|| NOT_REDACTED.to_owned())
}
/// [`fit`] of `text` redacted, redacting only its start and end, once (see [`redacted_start`]);
/// `None` when redacting takes too long.
fn fit_clipped(text: &str, room: usize, policy: &CapturePolicy) -> Option<String> {
let keys = private_keys(text);
let (head, head_of) = redacted_start(text, room / 2, policy)?;
let (tail, tail_of) = redacted_end(text, &keys, room - room / 2, policy)?;
if head_of + tail_of >= text.len() {
// Redaction took out most of it: there's nothing between them to leave out.
return policy.redact(text).map(|redacted| fit(&redacted, room));
}
let (head, tail) = ((head, head_of), (tail, tail_of));
fit_with(room, |keep| {
let (head, head_of) = trimmed_start(text, &head, keep / 2, policy)?;
let (tail, tail_of) = trimmed_end(text, &keys, &tail, keep - keep / 2, policy)?;
Some(noted(&head, text.len().saturating_sub(head_of + tail_of), &tail))
})
}
/// `start`, a [`redacted_start`] of `text`, cut to about `want` bytes, and how many bytes of
/// `text` that stands for: cut within, when what it cuts off is as written (so stands for as many
/// bytes of `text`), else redacted again for `want`.
fn trimmed_start(
text: &str,
(start, of): &(String, usize),
want: usize,
policy: &CapturePolicy,
) -> Option<(String, usize)> {
let cut = start.floor_char_boundary(want);
if start.len() - cut <= common_suffix(&text[..*of], start) {
return Some((start[..cut].to_owned(), of - (start.len() - cut)));
}
redacted_start(text, want, policy)
}
/// [`trimmed_start`] for a [`redacted_end`] of `text`.
fn trimmed_end(
text: &str,
keys: &[Range<usize>],
(end, of): &(String, usize),
want: usize,
policy: &CapturePolicy,
) -> Option<(String, usize)> {
let cut = end.ceil_char_boundary(end.len().saturating_sub(want));
if cut <= common_prefix(&text[text.len() - of..], end) {
return Some((end[cut..].to_owned(), of - cut));
}
redacted_end(text, keys, want, policy)
}
/// How many times [`redacted_start`] reads further for a start redaction shrank.
const REDACTED_ATTEMPTS: usize = 4;
/// How many times the bytes it is to hold [`redacted_start`] reads at most: a start redaction
/// shrank further holds less, rather than taking longer to redact than a capture may.
const REDACTED_GROWTH: usize = 8;
/// The start of `text`, redacted, about `want` bytes long, and how many bytes of `text` it stands
/// for (no more than half of it); `None` when redacting takes too long.
///
/// It is redacted with [`SECRET_MARGIN`] more past its cut, which is then cut off, so a secret
/// across the cut is redacted whole rather than its part on the kept side left behind
/// unrecognised: a secret across the end of what is redacted is unseen, so past the last
/// redaction, where it would be, what is cut off is all as written, mapping the cut back to the
/// text exactly. A start that redaction shrank is read further, so it still holds about `want`.
fn redacted_start(text: &str, want: usize, policy: &CapturePolicy) -> Option<(String, usize)> {
let limit = text.len() / 2;
let mut cut = text.floor_char_boundary(want.min(limit));
for attempt in 1..=REDACTED_ATTEMPTS {
let window = &text[..text.floor_char_boundary(cut + SECRET_MARGIN)];
let redacted = policy.redact(window)?;
let same = common_suffix(window, &redacted);
let drop = same.min(window.len() - cut);
let (kept, of) = (&redacted[..redacted.len() - drop], window.len() - drop);
if kept.len() >= want || cut >= limit || attempt == REDACTED_ATTEMPTS {
return Some((kept.to_owned(), of));
}
let more = (cut * want / kept.len().max(1)).min(want * REDACTED_GROWTH);
cut = text.floor_char_boundary(more.min(limit));
}
unreachable!("the last attempt returns")
}
/// [`redacted_start`] for the end of `text`. A private key it would hold without its start line
/// (which redaction finds a key by) is left out, up to its end line, or the end of `text` for a
/// key cut off before it; `keys` are `text`'s [`private_keys`].
fn redacted_end(
text: &str,
keys: &[Range<usize>],
want: usize,
policy: &CapturePolicy,
) -> Option<(String, usize)> {
let limit = text.len() - text.len() / 2;
let mut cut = text.ceil_char_boundary(text.len() - want.min(limit));
for attempt in 1..=REDACTED_ATTEMPTS {
let mut from = text.ceil_char_boundary(cut.saturating_sub(SECRET_MARGIN));
if let Some(key) = keys.iter().find(|key| key.start < from && key.end > from) {
from = key.end;
cut = cut.max(from);
}
let window = &text[from..];
let redacted = policy.redact(window)?;
let same = common_prefix(window, &redacted);
let drop = same.min(cut - from);
let (kept, of) = (&redacted[drop..], window.len() - drop);
if kept.len() >= want || text.len() - cut >= limit || attempt == REDACTED_ATTEMPTS {
return Some((kept.to_owned(), of));
}
let more = ((text.len() - cut) * want / kept.len().max(1)).min(want * REDACTED_GROWTH);
cut = text.ceil_char_boundary(text.len() - more.min(limit));
}
unreachable!("the last attempt returns")
}
/// Where the private keys in `text` are, as redaction finds them: from their start line past
/// their end line, or to the end of `text` for a key cut off before it.
fn private_keys(text: &str) -> Vec<Range<usize>> {
static KEY: LazyLock<Regex> = LazyLock::new(|| {
Regex::new(
r"-----BEGIN[ A-Z0-9_-]{0,100}PRIVATE KEY(?: BLOCK)?-----(?s:.+?)(?:-----END[ A-Z0-9_-]{0,100}PRIVATE KEY(?: BLOCK)?-----|\z)",
)
.expect("valid")
});
KEY.find_iter(text).map(|key| key.range()).collect()
}
/// How many bytes `a` and `b` start with alike, on a character boundary.
fn common_prefix(a: &str, b: &str) -> usize {
let mut n = a.bytes().zip(b.bytes()).take_while(|(x, y)| x == y).count();
while !a.is_char_boundary(n) {
n -= 1;
}
n
}
/// How many bytes `a` and `b` end with alike, on a character boundary.
fn common_suffix(a: &str, b: &str) -> usize {
let mut n = a.bytes().rev().zip(b.bytes().rev()).take_while(|(x, y)| x == y).count();
while !a.is_char_boundary(a.len() - n) {
n -= 1;
}
n
}
/// `text`, the value of `key`, redacted as the assignment it is: a secret is often known by the
/// name it is kept under (`{"AWS_SECRET_ACCESS_KEY": "..."}`), which the value alone doesn't say.
///
/// Only a value's start follows its name: of a long one, as much as capture could keep is
/// redacted so, and the rest dropped should that redact anything, as it may go on with what was.
/// Shaped, it is all redacted as any text is.
fn redact_named(key: &str, text: &mut String, policy: &CapturePolicy) {
let line = format!("{key}={text}");
let line = format!("{key}={}", &text[..text.floor_char_boundary(TOOL_PAYLOAD_LIMIT)]);
let Some(redacted) = policy.redact(&line) else {
NOT_REDACTED.clone_into(text);
return;
@@ -298,12 +491,29 @@ fn media_note(object: &Map<String, Value>) -> Option<Value> {
}
/// In place of a string that took too long to redact (see
/// [`REDACT_BUDGET`](atuin_common::secrets::REDACT_BUDGET)).
/// [`REDACT_BUDGET`](atuin_common::secrets::REDACT_BUDGET)), or JSON text too deep or too big to
/// shape as JSON (see [`JSON_TEXT_DEPTH`]).
pub(super) const NOT_REDACTED: &str = "[not captured: took too long to redact]";
/// In place of what a call writes to a file holding credentials.
const WITHHELD: &str = "[not captured: a file holding credentials]";
/// Whether `input` is what [`withhold_written`] leaves of one: it holds [`WITHHELD`], or is a
/// patch's headers alone. Judged again, it may no longer say what it wrote (a Codex script left
/// as [`WITHHELD`] makes no call).
pub(super) fn is_withheld(input: &Value) -> bool {
match input {
Value::String(text) => {
text.contains(WITHHELD)
|| (!text.trim().is_empty()
&& text.lines().all(|line| line.trim_start().starts_with("*** ")))
}
Value::Array(items) => items.iter().any(is_withheld),
Value::Object(object) => object.values().any(is_withheld),
Value::Null | Value::Bool(_) | Value::Number(_) => false,
}
}
/// `input`, a call writing to a file holding credentials, without what it writes: every string
/// but the paths it names, and of a patch, all but its headers (`*** Update File: .env`). It keeps
/// its shape, so a harness reads it back as the same call.
@@ -357,10 +567,14 @@ fn clip(text: &str, keep: usize) -> String {
if text.len() <= keep + NOTE {
return text.to_owned();
}
let head = text.floor_char_boundary(keep / 2);
let tail = text.ceil_char_boundary(text.len() - (keep - keep / 2));
let omitted = tail - head;
format!("{}\n[… {omitted} bytes not captured …]\n{}", &text[..head], &text[tail..])
let head = &text[..text.floor_char_boundary(keep / 2)];
let tail = &text[text.ceil_char_boundary(text.len() - (keep - keep / 2))..];
noted(head, text.len() - head.len() - tail.len(), tail)
}
/// `head` and `tail` around a note that `omitted` bytes between them were left out.
fn noted(head: &str, omitted: usize, tail: &str) -> String {
format!("{head}\n[… {omitted} bytes not captured …]\n{tail}")
}
#[cfg(test)]
@@ -401,6 +615,17 @@ mod tests {
assert!(kept.contains("AWS_SECRET_ACCESS_KEY"), "the name stays: {kept}");
}
/// A patch a call makes (Codex `apply_patch`'s text, opencode's `patchText`) has its lines'
/// secrets redacted, what it adds as well as what it removes.
#[rstest]
#[case::codex(json!("*** Begin Patch\n*** Add File: db.yml\n+production:\n+ password: hunter2xyz\n*** End Patch\n"))]
#[case::opencode(json!({"patchText": "*** Begin Patch\n*** Update File: app.properties\n@@\n-db.password=hunter1xyz\n+db.password=hunter2xyz\n*** End Patch"}))]
fn a_patchs_added_secrets_are_redacted(#[case] value: Value) {
let kept = kept(value).to_string();
assert!(!kept.contains("hunter"), "{kept}");
assert!(kept.contains("password"), "the name stays: {kept}");
}
#[rstest]
fn names_only_keeps_nothing() {
let mut value = json!({"command": "ls"});
@@ -434,6 +659,184 @@ mod tests {
assert!(kept.len() < TOOL_PAYLOAD_LIMIT + 64);
}
/// A long output is clipped before it is redacted, so one far too long to redact whole in
/// time (this one takes seconds) is still kept, its start and end redacted.
#[rstest]
fn a_long_output_is_clipped_not_dropped_for_taking_long_to_redact() {
let token = "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";
let text = format!("{token} {}{token}", "atuin login ".repeat(1 << 20));
let Value::String(kept) = kept(json!(text)) else {
panic!("still a string")
};
assert!(kept.contains("bytes not captured"), "{kept}");
assert!(kept.starts_with("**** atuin login") && kept.ends_with("atuin login ****"));
}
/// Wherever a long output is cut, no part of a secret across the cut is left behind.
#[rstest]
fn a_secret_across_a_cut_is_redacted_whole(#[values(0, 7, 20, 33)] shift: usize) {
let token = "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";
let text = format!("{}{}", " ".repeat(shift), format!("{token} ").repeat(100_000));
let Value::String(kept) = kept(json!(text)) else {
panic!("still a string")
};
assert!(kept.contains("bytes not captured"));
for part in kept.split([' ', '\n']).filter(|part| !part.is_empty()) {
assert!(
part.chars().all(|c| c == '*') || !token.contains(part) || part.len() < 4,
"{part:?} is part of a secret"
);
}
}
/// A private key longer than the margin a cut is redacted with, across the cut (an armored
/// PGP secret key is 5-12 KiB): none of it is kept, though the start line redaction finds it
/// by is on the other side.
#[rstest]
fn a_long_private_key_across_a_cut_is_not_kept(#[values(false, true)] at_the_end: bool) {
let line =
|i: usize| format!("lQdGBGVxyz{i:03}").repeat(6).chars().take(64).collect::<String>();
let body: String = (0..140).map(|i| line(i) + "\n").collect();
let key = format!(
"-----BEGIN PGP PRIVATE KEY BLOCK-----\n\n{body}=AbCd\n-----END PGP PRIVATE KEY \
BLOCK-----\n"
);
assert!(key.len() > 2 * SECRET_MARGIN);
// Spaces, not newlines, which take twice the room stored and move the cuts.
let (near, far) = ("x ".repeat(13_000), "log line of output ".repeat(6000));
let text = if at_the_end {
format!("{far}{key}{near}")
} else {
format!("{near}{key}{far}")
};
let Value::String(kept) = kept(json!(text)) else {
panic!("still a string")
};
assert!(kept.contains("bytes not captured"));
assert!(!kept.contains("lQdGBG"), "{kept}");
}
/// A match of the user's own patterns across a cut, whose start or end is in what is cut
/// out, still redacts what it covers of what is kept.
#[rstest]
#[case::into_the_end(100, 180, 190)]
#[case::out_of_the_start(10, 20, 60)]
fn a_users_pattern_across_a_cut_is_redacted(
#[case] open: usize,
#[case] secret: usize,
#[case] close: usize,
) {
let policy = CapturePolicy::with_redact_patterns(vec![
Regex::new(r"(?s)<secret>.*?</secret>").unwrap(),
]);
let mut text = "x ".repeat(100 * 1024);
for (at, part) in [(close, "</secret>"), (secret, "violet-7-meadow"), (open, "<secret>")] {
text.insert_str(at * 1024, part);
}
let mut value = json!(text);
ToolCapture::Payloads.apply(&mut value, &policy);
let kept = value.to_string();
assert!(kept.contains("bytes not captured"));
assert!(!kept.contains("violet-7-meadow"), "{kept}");
}
/// A private key's armored body, `lines` lines of it, without its end line, as output cut off
/// mid-key leaves it.
fn unended_key(lines: usize) -> String {
let line =
|i: usize| format!("lQdGBGVxyz{i:03}").repeat(6).chars().take(64).collect::<String>();
let body: String = (0..lines).map(|i| line(i) + "\n").collect();
format!("-----BEGIN PGP PRIVATE KEY BLOCK-----\n\n{body}")
}
/// A private key cut off before its end line, its start line before the window the end of a
/// clipped text is redacted in: none of it is kept, in a later item left little of the budget
/// as in a text whose key outruns the window.
#[rstest]
#[case::strings(json!(["log line of output ".repeat(3340), unended_key(200)]))]
#[case::string(json!(format!("{}{}", "log line of output ".repeat(6000), unended_key(700))))]
#[case::then_a_key(json!(format!(
"{}{}{}-----END PGP PRIVATE KEY BLOCK-----\nlog line of output",
"log line of output ".repeat(6000),
unended_key(700),
unended_key(20),
)))]
fn an_unended_private_key_before_a_cut_is_not_kept(#[case] value: Value) {
let kept = kept(value).to_string();
assert!(kept.contains("bytes not captured"), "{kept:.200}");
assert!(!kept.contains("lQdGBG"), "{kept}");
}
/// Text redaction shrinks a lot (a long run of tokens) is kept as much of as fits, and the
/// note says exactly how much of the text, as written, it left out: also when what is kept
/// takes more room stored than it is long (newlines, escaped), and is cut again to fit.
#[rstest]
#[case::fits_once_redacted(5_000)]
#[case::does_not(200_000)]
#[case::by_far(1_000_000)]
fn heavily_redacted_text_is_kept_as_much_as_fits(
#[case] tokens: usize,
#[values(" ", "\n")] separator: &str,
) {
let token = "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";
let text = format!("{token}{separator}").repeat(tokens);
let Value::String(kept) = kept(json!(text)) else {
panic!("still a string")
};
let redacted = format!("****{separator}").repeat(tokens);
if stored_len(&redacted) <= TOOL_PAYLOAD_LIMIT {
assert_eq!(kept, redacted);
return;
}
assert!(stored_len(&kept) > TOOL_PAYLOAD_LIMIT * 3 / 4, "kept {} bytes", kept.len());
let (head, rest) = kept.split_once("\n[… ").unwrap();
let (omitted, tail) = rest.split_once(" bytes not captured …]\n").unwrap();
// What a part kept stands for, as written: each `****` a token.
let written = |part: &str| part.len() + part.matches("****").count() * (token.len() - 4);
let omitted: usize = omitted.parse().unwrap();
assert_eq!(written(head) + omitted + written(tail), text.len());
}
/// JSON text nested in JSON text, many levels deep, is shaped as JSON only a level or two
/// down, and as text below: shaping each level again until it fits, as deep as it goes, took
/// half a minute (holding the daemon) for this 1.5 MB call.
#[rstest]
fn deeply_nested_json_text_is_shaped_in_bounded_time() {
let keys: Map<String, Value> =
(0..3000).map(|i| (format!("k{i:04}"), json!("v"))).collect();
let mut text = Value::Object(keys).to_string();
for _ in 0..6 {
text = json!([text]).to_string();
}
let started = std::time::Instant::now();
let kept = kept(json!({"command": text}));
assert!(started.elapsed() < std::time::Duration::from_secs(10), "{:?}", started.elapsed());
assert!(kept.to_string().len() <= TOOL_PAYLOAD_LIMIT + OVERSHOOT);
}
/// JSON text nested deeper than `shape` reads it as JSON is not kept as text, where its keys'
/// quotes are escaped and hide a secret from redaction.
#[rstest]
fn json_text_nested_too_deep_keeps_no_secret(#[values(0, 1, 2, 3, 4)] levels: usize) {
let mut text = json!({"client_secret": "violet-7-meadow"}).to_string();
for _ in 0..levels {
text = json!([text]).to_string();
}
let kept = kept(json!(text)).to_string();
assert!(!kept.contains("violet-7-meadow"), "{kept}");
}
/// JSON text too big to shape as JSON in its room (more keys than fit, which `shape` never
/// drops) is not kept as text either, where JSON text in it would hide a secret.
#[rstest]
fn json_text_too_big_to_shape_keeps_no_secret() {
let mut keys: Map<String, Value> =
(0..8000).map(|i| (format!("k{i:04}"), json!("v"))).collect();
keys.insert("a".to_owned(), json!(json!({"client_secret": "violet-7-meadow"}).to_string()));
let kept = kept(json!(Value::Object(keys).to_string())).to_string();
assert!(!kept.contains("violet-7-meadow"), "{kept:.200}");
}
/// The budget is shared: once early strings spend it, later long ones are clipped to little
/// more than their note, while short ones (the blocks' kinds) are kept.
#[rstest]
crates/atuin/src/command/client/daemon.rs
+68 -5
View File
@@ -274,12 +274,16 @@ fn remove_sockets(
/// Remove any socket left behind by a daemon that is no longer listening.
#[cfg(unix)]
fn remove_stale_socket_if_present(settings: &Settings) -> Result<(), RemoveSocketError> {
use std::os::unix::fs::FileTypeExt;
remove_sockets(settings, |socket_path| {
// A refused connection means the socket is left over from a daemon that is gone.
matches!(
StdUnixStream::connect(socket_path),
Err(e) if e.kind() == ErrorKind::ConnectionRefused
)
// A refused connection means the socket is left over from a daemon that is gone. Only a
// socket, though: connecting to any other file is refused too.
fs::symlink_metadata(socket_path).is_ok_and(|meta| meta.file_type().is_socket())
&& matches!(
StdUnixStream::connect(socket_path),
Err(e) if e.kind() == ErrorKind::ConnectionRefused
)
})
}
@@ -523,6 +527,34 @@ pub async fn rebuild_ai_sessions(settings: &Settings) -> Result<()> {
client.rebuild_sessions().await
}
/// How long to look for a running daemon to project synced AI sessions: a stopped daemon's socket
/// still accepts connections but never answers.
const PROJECT_SYNCED_PROBE_TIMEOUT: Duration = Duration::from_secs(2);
/// How long a sync waits for the daemon to project the AI sessions it downloaded. A large store
/// takes longer; the daemon carries on regardless.
const PROJECT_SYNCED_WAIT: Duration = Duration::from_secs(10);
/// Have a running daemon project the AI session records this process synced into the record
/// store, returning once they are listed or, after [`PROJECT_SYNCED_WAIT`], leaving the daemon to
/// finish on its own. Never starts or replaces one: a daemon started later, or replacing one too
/// old to be asked, projects them as it starts.
pub async fn project_synced_ai_sessions(settings: &Settings) -> Result<()> {
let probe = tokio::time::timeout(PROJECT_SYNCED_PROBE_TIMEOUT, probe(settings)).await;
let Ok(Probe::Ready(_)) = probe else {
return Ok(());
};
let mut client = atuin_daemon::AiClient::from_settings(settings).await?;
match tokio::time::timeout(PROJECT_SYNCED_WAIT, client.project_synced()).await {
Ok(Err(err)) if is_legacy_daemon_error(&err) => Ok(()),
Ok(projected) => projected,
Err(_) => {
eprintln!("Indexing AI sessions in the background");
Ok(())
}
}
}
pub async fn compact_store(settings: &Settings) -> Result<u64> {
let reply = try_with_restart(settings, async |client| client.compact_store().await).await?;
Ok(reply.rewritten)
@@ -680,6 +712,10 @@ async fn run(
let _pidfile_guard = PidfileGuard::acquire(&settings.daemon)?;
// A daemon that crashed left its socket behind, which would keep this one from listening.
#[cfg(unix)]
remove_stale_socket_if_present(&settings)?;
atuin_daemon::boot(settings, store, history_db).await?;
Ok(())
@@ -733,6 +769,33 @@ mod tests {
);
}
/// A socket nothing listens on is removed; one a daemon listens on, and a file that is no
/// socket (a `socket_path` set wrong), are not.
#[cfg(unix)]
#[rstest]
#[case::stale(false, true, false)]
#[case::listening(true, true, true)]
#[case::not_a_socket(false, false, true)]
fn only_a_stale_socket_is_removed(
#[case] listening: bool,
#[case] socket: bool,
#[case] kept: bool,
) {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("atuin.sock");
let listener = socket.then(|| std::os::unix::net::UnixListener::bind(&path).unwrap());
if !socket {
fs::write(&path, "not a socket").unwrap();
}
if !listening {
drop(listener);
}
let mut settings = Settings::default();
settings.daemon.socket_path = Some(path.clone());
remove_stale_socket_if_present(&settings).unwrap();
assert_eq!(path.exists(), kept);
}
#[cfg(unix)]
#[rstest]
fn default_paths_are_absolute() {
crates/atuin/src/command/client/hook.rs
+17 -2
View File
@@ -4,7 +4,7 @@ use std::str::FromStr;
use atuin_client::history::{AuthorKind, HistoryId};
use atuin_client::settings::Settings;
use atuin_common::harnesstools::{AnyHarness, Harness, InstallHookError};
use atuin_common::harnesstools::{AnyHarness, Harness, InstallHookError, codex};
use clap::{Parser, Subcommand};
use eyre::{Result, bail};
use tracing::instrument;
@@ -73,6 +73,11 @@ async fn install(agent_name: &str) -> Result<()> {
match harness.install_hooks().await {
Ok(path) => {
eprintln!("Atuin hooks installed for {}. Path: {}", harness.name(), path.display());
if harness.name() == "codex" {
// Codex runs only the hooks the user trusted (see `codex::hook_trust`).
let config = path.with_file_name("config.toml");
eprintln!("Trusted them in Codex's config: {}", config.display());
}
if let Some(hint) = reload_hint(harness.name()) {
eprintln!("{hint}");
}
@@ -114,7 +119,17 @@ async fn handle(agent_name: &str, settings: &Settings) -> Result<()> {
std::fs::write(id_file_path(&tool_use_id), history_id.to_string())?;
}
}
Some(HookEvent::End { tool_use_id, exit }) => {
Some(HookEvent::End {
tool_use_id,
exit,
transcript,
}) => {
let exit = exit
.or_else(|| {
let transcript = transcript.filter(|_| harness.name() == "codex")?;
codex::session::command_exit_code(&transcript, &tool_use_id)
})
.unwrap_or(0);
let id_path = id_file_path(&tool_use_id);
if let Ok(history_id) = std::fs::read_to_string(&id_path) {
crates/atuin/src/command/client/hook/event.rs
+67 -23
View File
@@ -4,6 +4,8 @@
//! an event we care about. If we don't know how to deserialize it/don't care for it, we need to
//! drop that on the floor. The [`HookEvent`] type represents agent events we know/care about.
use std::path::PathBuf;
use atuin_common::string::NonNulStr;
use serde_json::error::Category;
@@ -41,7 +43,10 @@ pub enum HookEvent {
/// A Bash command finished; close the matching history entry.
End {
tool_use_id: String,
exit: i64,
/// Its exit code, when the event says it.
exit: Option<i64>,
/// The session's transcript, where Codex records the exit code its event doesn't say.
transcript: Option<PathBuf>,
},
}
@@ -80,22 +85,27 @@ impl From<WireHookEvent> for Option<HookEvent> {
tool_use_id,
})
}
HookEventName::PostToolUse => {
// TODO(markovejnovic): Is it safe to assume that no exit code
// means "success"?
let exit =
wire.tool_response.and_then(|response| response.exit_code()).unwrap_or(0);
Some(HookEvent::End { tool_use_id, exit })
}
HookEventName::PostToolUse => Some(HookEvent::End {
tool_use_id,
exit: wire.tool_response.and_then(|response| response.exit_code()),
transcript: wire.transcript_path,
}),
// Failed: with the exit code its error starts with (`Exit code 2`), else 1.
HookEventName::PostToolUseFailure => Some(HookEvent::End {
tool_use_id,
exit: 1,
exit: Some(wire.error.as_deref().and_then(failure_exit_code).unwrap_or(1)),
transcript: wire.transcript_path,
}),
HookEventName::Other => None,
}
}
}
/// The exit code a Claude Code failure's `error` starts with (`Exit code <n>`), if it does.
fn failure_exit_code(error: &str) -> Option<i64> {
error.lines().next()?.strip_prefix("Exit code ")?.trim().parse().ok()
}
impl HookEvent {
/// Parse a raw hook payload (the JSON an agent writes to stdin) into a [`HookEvent`], or `None`
/// when there is nothing to record.
@@ -163,28 +173,28 @@ mod tests {
"tool_response": {"exitCode": 3, "stdout": "hello\n"},
"tool_use_id": "toolu_abc123"
}),
Some(HookEvent::End { tool_use_id: "toolu_abc123".into(), exit: 3 })
Some(HookEvent::End { tool_use_id: "toolu_abc123".into(), exit: Some(3), transcript: None })
)]
#[case::post_tool_use_without_exit_code_defaults_zero(
#[case::post_tool_use_without_exit_code_says_none(
json!({
"hook_event_name": "PostToolUse",
"tool_name": "Bash",
"tool_response": {},
"tool_use_id": "toolu_abc123"
}),
Some(HookEvent::End { tool_use_id: "toolu_abc123".into(), exit: 0 })
Some(HookEvent::End { tool_use_id: "toolu_abc123".into(), exit: None, transcript: None })
)]
// A null exitCode also defaults to 0.
#[case::null_exit_code_defaults_zero(
// Nor does a null exitCode.
#[case::null_exit_code_says_none(
json!({
"hook_event_name": "PostToolUse",
"tool_name": "Bash",
"tool_response": {"exitCode": null},
"tool_use_id": "toolu_abc123"
}),
Some(HookEvent::End { tool_use_id: "toolu_abc123".into(), exit: 0 })
Some(HookEvent::End { tool_use_id: "toolu_abc123".into(), exit: None, transcript: None })
)]
// PostToolUseFailure forces exit 1 and ignores tool_response entirely.
// PostToolUseFailure without an exit code in its error is exit 1, whatever tool_response says.
#[case::failure_forces_exit_one_ignoring_response(
json!({
"hook_event_name": "PostToolUseFailure",
@@ -193,7 +203,41 @@ mod tests {
"tool_response": {"exitCode": 0},
"tool_use_id": "toolu_abc123"
}),
Some(HookEvent::End { tool_use_id: "toolu_abc123".into(), exit: 1 })
Some(HookEvent::End { tool_use_id: "toolu_abc123".into(), exit: Some(1), transcript: None })
)]
// Claude Code's failure says a shell command's exit code first in its error (Claude Code
// 2.1.291, `PostToolUseFailure`'s `error`).
#[case::failure_with_its_exit_code(
json!({
"hook_event_name": "PostToolUseFailure",
"tool_name": "Bash",
"tool_input": {"command": "ls /nope"},
"tool_use_id": "toolu_abc123",
"error": "Exit code 2\nls: cannot access '/nope': No such file or directory",
"is_interrupt": false
}),
Some(HookEvent::End { tool_use_id: "toolu_abc123".into(), exit: Some(2), transcript: None })
)]
// Codex 0.160 says only the output, and where its transcript is (a real payload's fields).
#[case::codex_post_tool_use(
json!({
"session_id": "01a11e4d-75d7-70f2-8c71-7f1cc87cbe85",
"turn_id": "01a11e4d-761f-7073-955e-e96a0a319fc3",
"transcript_path": "/w/.codex/sessions/2026/10/09/rollout-x.jsonl",
"cwd": "/w",
"hook_event_name": "PostToolUse",
"model": "gpt-6-luna",
"permission_mode": "bypassPermissions",
"tool_name": "Bash",
"tool_input": {"command": "ls /nonexistent_juliettea"},
"tool_response": "ls: cannot access '/nonexistent_juliettea': No such file or directory\n",
"tool_use_id": "exec-55bd0a92-7032-4356-b380-ac1cc1ea66f4"
}),
Some(HookEvent::End {
tool_use_id: "exec-55bd0a92-7032-4356-b380-ac1cc1ea66f4".into(),
exit: None,
transcript: Some("/w/.codex/sessions/2026/10/09/rollout-x.jsonl".into()),
})
)]
// Non-Bash tools are never recorded.
#[case::non_bash_tool_skipped(
@@ -304,7 +348,7 @@ mod tests {
#[rstest]
#[case::string_tool_response(
r#"{"hook_event_name":"PostToolUse","tool_name":"Bash","tool_use_id":"example-call","session_id":"example-session","cwd":"/tmp","tool_input":{"command":"printf probe"},"tool_response":"probe"}"#,
Some(HookEvent::End { tool_use_id: "example-call".into(), exit: 0 })
Some(HookEvent::End { tool_use_id: "example-call".into(), exit: None, transcript: None })
)]
fn completion_accepts_string_tool_response(
#[case] input: &str,
@@ -388,14 +432,14 @@ mod tests {
prop_assert_eq!(
HookEvent::from_json_str(&input.to_string()).unwrap(),
Some(HookEvent::End { tool_use_id, exit })
Some(HookEvent::End { tool_use_id, exit: Some(exit), transcript: None })
);
}
/// `PostToolUseFailure` always records exit 1, whatever the response
/// claims.
/// `PostToolUseFailure` records exit 1, whatever the response claims,
/// unless its error says the exit code.
#[test]
fn failure_event_always_exits_one(
fn failure_without_an_exit_code_in_its_error_exits_one(
reported_exit in any::<i64>(),
tool_use_id in r"[^\p{Cc}]+",
) {
@@ -408,7 +452,7 @@ mod tests {
prop_assert_eq!(
HookEvent::from_json_str(&input.to_string()).unwrap(),
Some(HookEvent::End { tool_use_id, exit: 1 })
Some(HookEvent::End { tool_use_id, exit: Some(1), transcript: None })
);
}
crates/atuin/src/command/client/hook/wire.rs
+10
View File
@@ -3,6 +3,8 @@
//! Claude Code and Codex invoke `atuin hook <agent>` for each tool use and
//! pass the event as JSON on stdin.
use std::path::PathBuf;
use atuin_common::string::NonNulStr;
use serde::Deserialize;
@@ -31,6 +33,14 @@ pub struct WireHookEvent {
/// How the command finished. Present on `PostToolUse`; absent elsewhere.
#[serde(default)]
pub tool_response: Option<WireToolResponse>,
/// Why the command failed, on Claude Code's `PostToolUseFailure`: for a shell command that
/// exited non-zero, `Exit code <n>` and then its output.
#[serde(default)]
pub error: Option<String>,
/// The session's transcript, where Codex records a command's exit code (its
/// `tool_response` is only the output).
#[serde(default)]
pub transcript_path: Option<PathBuf>,
}
/// The lifecycle stage an event represents.
crates/atuin/src/command/client/store/pull.rs
+6 -1
View File
@@ -114,7 +114,12 @@ impl Pull {
println!("Downloaded {} records", downloaded.len());
crate::sync::build(settings, &store, db, Some(&downloaded)).await?;
// On --force the daemon rebuilds the AI sessions from scratch after this.
if self.force {
crate::sync::build(settings, &store, db, Some(&downloaded)).await?;
} else {
crate::command::client::sync::build(settings, &store, db, &downloaded).await?;
}
Ok(())
}
crates/atuin/src/command/client/sync.rs
+27 -3
View File
@@ -4,7 +4,7 @@ use atuin_client::record::sqlite_store::SqliteStore;
use atuin_client::record::sync::{ClientSource, SyncSession};
use atuin_client::settings::Settings;
use atuin_common::encryption::paseto_v4;
use atuin_domain::record::RecordTag;
use atuin_domain::record::{RecordId, RecordTag};
use clap::Subcommand;
use easy_cast::Conv;
use eyre::{Result, WrapErr};
@@ -13,6 +13,8 @@ use tracing::instrument;
mod status;
use crate::command::client::account;
#[cfg(feature = "daemon")]
use crate::command::client::daemon;
use crate::i18n::fl;
#[derive(Subcommand, Debug)]
@@ -94,7 +96,7 @@ async fn run(settings: &Settings, force: bool, db: &Sqlite, store: SqliteStore)
.await
.map_err(crate::print_error::format_sync_error)?;
crate::sync::build(settings, &store, db, Some(&downloaded)).await?;
build(settings, &store, db, &downloaded).await?;
println!("{}", fl!("sync-up-down", uploaded = uploaded, downloaded = downloaded.len()));
@@ -122,7 +124,7 @@ async fn run(settings: &Settings, force: bool, db: &Sqlite, store: SqliteStore)
.await
.map_err(crate::print_error::format_sync_error)?;
crate::sync::build(settings, &store, db, Some(&downloaded)).await?;
build(settings, &store, db, &downloaded).await?;
println!("{}", fl!("sync-up-down", uploaded = uploaded, downloaded = downloaded.len()));
}
@@ -134,3 +136,25 @@ async fn run(settings: &Settings, force: bool, db: &Sqlite, store: SqliteStore)
Ok(())
}
/// [`crate::sync::build`], then have a running daemon project the AI sessions `downloaded` brought,
/// even if building the other stores failed: it owns their index, and otherwise projects only what
/// it syncs itself, or finds as it starts.
pub(super) async fn build(
settings: &Settings,
store: &SqliteStore,
db: &Sqlite,
downloaded: &[RecordId],
) -> Result<()> {
let built = crate::sync::build(settings, store, db, Some(downloaded)).await;
#[cfg(feature = "daemon")]
if settings.daemon.enabled
&& !downloaded.is_empty()
&& let Err(e) = daemon::project_synced_ai_sessions(settings).await
{
eprintln!("Warning: failed to build ai sessions: {e}");
}
built
}
crates/atuin/tests/common/agents.rs
+39
View File
@@ -294,6 +294,33 @@ impl Transcript {
None
}
/// Rename the opencode session `title`, as `Session.setTitle` does: a `session.updated.1` row
/// with the session's info, `time.updated` bumped to now.
pub fn rename(&self, title: &str) {
assert_eq!(self.agent, Agent::Opencode, "only opencode's sessions are renamed here");
let (directory, created) = rusqlite::Connection::open(&self.path)
.unwrap()
.query_row(
"SELECT directory, time_created FROM session WHERE id = ?1",
[&self.id],
|row| Ok((row.get::<_, String>(0)?, row.get::<_, i64>(1)?)),
)
.unwrap();
let home = self.path.ancestors().nth(4).unwrap();
opencode_events(home, &[(
"session.updated.1",
json!({
"sessionID": self.id,
"info": {
"id": self.id,
"title": title,
"directory": directory,
"time": {"created": created, "updated": millis(tick())},
},
}),
)]);
}
/// The transcript as it stands.
pub fn read(&self) -> String {
fs::read_to_string(&self.path).unwrap()
@@ -646,6 +673,18 @@ fn opencode_events(home: &Path, events: &[(&str, Value)]) {
)
.unwrap();
}
"session.updated.1" => {
let info = &data["info"];
tx.execute(
"UPDATE session SET title = ?2, time_updated = ?3 WHERE id = ?1",
rusqlite::params![
session,
info["title"].as_str(),
info["time"]["updated"].as_i64()
],
)
.unwrap();
}
"message.updated.1" => {
let info = &data["info"];
let created = info["time"]["created"].as_i64();
crates/atuin/tests/common/ai.rs
+27 -13
View File
@@ -135,8 +135,8 @@ impl Machine {
}
/// A machine with the settings `top` and `[ai]` settings `ai`, syncing through the server at
/// `sync`, if any, as often as its daemon can.
fn configured(sync: Option<&str>, top: &str, ai: &str) -> Self {
/// `sync`, if any: its daemon syncing every so many seconds, or (0) never.
fn configured(sync: Option<(&str, u32)>, top: &str, ai: &str) -> Self {
let env = FreshEnv::new();
let bin = tempfile::Builder::new().prefix("bin-").tempdir_in(target_tmp()).unwrap();
install_fakes(bin.path(), &Agent::ALL);
@@ -144,10 +144,12 @@ impl Machine {
agent.install(env.home());
}
let (auto_sync, address, frequency) = match sync {
Some(address) => {
(true, format!("sync_address = \"{address}\"\n"), "sync_frequency = 1\n")
}
None => (false, String::new(), ""),
Some((address, frequency)) => (
true,
format!("sync_address = \"{address}\"\n"),
format!("sync_frequency = {frequency}\n"),
),
None => (false, String::new(), String::new()),
};
env.write_config(&format!(
"auto_sync = {auto_sync}\n{address}{top}[daemon]\nenabled = true\nautostart = \
@@ -162,18 +164,30 @@ impl Machine {
}
}
/// Two machines of one account on `server`, their daemons running.
/// Two machines of one account on `server`, their daemons running and syncing as often as
/// they can.
pub fn pair(server: &SyncServer) -> (Self, Self) {
let one = Self::syncing(server, None);
Self::pair_syncing(server, 1)
}
/// Two machines of one account on `server`, their daemons running: the first's syncing as
/// often as it can, the second's every `sync_frequency` seconds, or (0) never.
pub fn pair_syncing(server: &SyncServer, sync_frequency: u32) -> (Self, Self) {
let one = Self::syncing(server, None, 1);
let key = one.run(&["key", "--base64"]).trim().to_owned();
let two = Self::syncing(server, Some((&one, &key)));
let two = Self::syncing(server, Some((&one, &key)), sync_frequency);
(one, two)
}
/// A machine syncing through `server`: registering a new account, or logging into `login`'s
/// with its key. Logged in before its daemon starts, which keeps the key it starts with.
fn syncing(server: &SyncServer, login: Option<(&Self, &str)>) -> Self {
let mut machine = Self::configured(Some(&server.address), "", "capture_sessions = true\n");
/// A machine syncing through `server`, its daemon every `sync_frequency` seconds: registering
/// a new account, or logging into `login`'s with its key. Logged in before its daemon starts,
/// which keeps the key it starts with.
fn syncing(server: &SyncServer, login: Option<(&Self, &str)>, sync_frequency: u32) -> Self {
let mut machine = Self::configured(
Some((&server.address, sync_frequency)),
"",
"capture_sessions = true\n",
);
match login {
None => {
let user = machine.user();
crates/atuin/tests/e2e_ai_sessions.rs
+102
View File
@@ -164,6 +164,27 @@ fn capture_follows_a_transcript_exactly_once_across_restarts(
machine.wait_until_said(&transcript.id, &said(&turns));
}
/// An opencode session renamed after its conversation is listed, and found, under its new title
/// at once, not only once another message carries it.
#[rstest]
fn a_renamed_session_is_listed_under_its_new_title(machine: Machine, project: TempDir) {
let turns = [Turn::new()];
let transcript = Transcript::create(machine.home(), Agent::Opencode, project.path(), &turns);
machine.wait_until_said(&transcript.id, &said(&turns));
let title = format!("renamed {}", word());
transcript.rename(&title);
ai::wait_for("the session to be listed under its new title", || {
let listed = machine.session(&transcript.id).unwrap()["title"].clone();
if listed == title.as_str() {
Ok(())
} else {
Err(format!("it is listed as {listed}"))
}
});
assert!(finds(&machine, last_word(&title), &transcript.id), "the new title didn't find it");
}
/// Turning capture on in the config takes effect in the running daemon, which then captures what
/// the agents wrote while it was off.
#[rstest]
@@ -282,6 +303,9 @@ fn restored(agent: Agent) -> Restored {
let restored = Transcript::find(here.home(), agent, &original.id)
.unwrap_or_else(|| panic!("nothing restored: {stdout} {stderr}"));
assert_eq!(stdout.trim(), resume_line(&restored, project.path()), "{stderr}");
// Where it ran is here too, and the status says what was written all the same.
assert!(stderr.contains("atuin: restored the session from sync\n"), "{stderr}");
assert!(!stderr.lines().any(|line| line.contains("locked")), "{stderr}");
Restored {
_server: server,
project,
@@ -330,6 +354,27 @@ fn a_session_from_another_machine_restores_as_the_same_session(
there.wait_until_said(&original.id, &all);
}
/// A session `atuin sync` brings down is listed as soon as it returns, though the daemon running
/// beside it never syncs itself.
#[rstest]
fn a_session_synced_by_hand_is_listed_at_once(project: TempDir) {
let server = SyncServer::start();
let (there, here) = Machine::pair_syncing(&server, 0);
let turns = [Turn::new()];
let original = Transcript::create(there.home(), Agent::Claude, project.path(), &turns);
there.wait_until_said(&original.id, &said(&turns));
let downloaded = Regex::new(r"\d+/(\d+) up/down").unwrap();
ai::wait_for("a sync to download the session", || {
let synced = here.run(&["sync"]);
match downloaded.captures(&synced) {
Some(counts) if &counts[1] != "0" => Ok(()),
_ => Err(synced),
}
});
assert!(here.session(&original.id).is_some(), "not listed once synced");
}
/// A copy behind the session's other machine catches up with it as it resumes, saying so, and
/// once caught up resumes as it is.
#[rstest]
@@ -435,6 +480,53 @@ fn a_diverged_copy_resumes_only_as_chosen(
assert!(!text.contains(&ours.reply), "our branch is still in it:\n{text}");
}
/// Continuing a session that went different ways on two machines in another agent carries the
/// branch `--branch` names, never both interleaved; without one named, it continues nothing,
/// listing the branches to name.
#[rstest]
#[case::claude_to_codex(Agent::Claude, Agent::Codex)]
#[case::codex_to_claude(Agent::Codex, Agent::Claude)]
#[case::pi_to_codex(Agent::Pi, Agent::Codex)]
fn continuing_a_diverged_session_carries_one_branch(#[case] from: Agent, #[case] to: Agent) {
let Restored {
_server,
project,
here,
there,
turns,
original,
restored,
} = restored(from);
// Theirs is the newest.
let (ours, theirs) = (Turn::new(), Turn::new());
restored.append(project.path(), std::slice::from_ref(&ours));
original.append(project.path(), std::slice::from_ref(&theirs));
ai::wait("both branches", || {
let show = here.show(&original.id).to_string();
show.contains(&theirs.reply) && show.contains(&ours.reply)
});
let args = [original.id.as_str(), "--in", to.program(), "--print"];
let stderr = ai::failed(&here.resume(&args, here.home()));
assert!(
stderr.contains("it went 2 ways: name the branch to continue with `--branch`:"),
"{stderr}"
);
assert!(stderr.contains(&there.short_host()) && stderr.contains("this machine"), "{stderr}");
let theirs_by = there.short_host();
for (branch, kept, dropped) in [(&*theirs_by, &theirs, &ours), ("this", &ours, &theirs)] {
let args = [&args[..], &["--branch", branch]].concat();
let (stdout, stderr) = succeeded(&here.resume(&args, here.home()));
let continued = Transcript::find(here.home(), to, last_word(stdout.trim()))
.unwrap_or_else(|| panic!("nothing continued: {stdout} {stderr}"));
let text = continued.read();
let want = [said(&turns), said(std::slice::from_ref(kept))].concat();
assert!(ai::holds_in_order(&text, &want), "{branch:?}: not carried over:\n{text}");
assert!(!text.contains(&dropped.reply), "{branch:?}: the other branch came too:\n{text}");
}
}
/// Forking writes the session out as a new one of its agent, with the whole conversation, linked
/// to it as a fork, and resumes that; the original is left as it was.
#[rstest]
@@ -534,3 +626,13 @@ fn rebuilding_the_index_reproduces_it(machine: Machine, project: TempDir) {
// `ai session` commands wait for the rebuild.
assert_eq!(index(&machine), before);
}
/// A message bigger than gRPC takes by default (4 MiB: a pasted log, a huge file) reads back
/// whole.
#[rstest]
fn a_huge_message_reads_back(machine: Machine, project: TempDir) {
let mut turns = [Turn::new()];
turns[0].prompt = format!("{} {}", "x".repeat(9 << 19), turns[0].prompt);
let transcript = Transcript::create(machine.home(), Agent::Claude, project.path(), &turns);
machine.wait_until_said(&transcript.id, &said(&turns));
}
crates/atuin/tests/e2e_daemon.rs
+40 -21
View File
@@ -24,6 +24,32 @@ fn daemon() -> Daemon {
}
}
impl Daemon {
/// Start the daemon in the foreground, and wait until it is healthy.
async fn start_foreground(&mut self) {
self.foreground = Some(Process::spawn(self.env.atuin(&["daemon", "start", "--show-logs"])));
tokio::time::timeout(TIMEOUT, async {
loop {
let process = self.foreground.as_mut().unwrap();
assert!(
process.child.try_wait().unwrap().is_none(),
"daemon exited during startup: {}",
process.logs()
);
if let Ok(mut client) = HistoryClient::new(self.env.socket()).await
&& let Ok(status) = client.status().await
&& status.healthy
{
break;
}
tokio::time::sleep(Duration::from_millis(10)).await;
}
})
.await
.expect("daemon never became healthy");
}
}
impl Drop for Daemon {
fn drop(&mut self) {
let _ = Process::spawn(self.env.atuin(&["daemon", "stop"])).try_wait();
@@ -47,27 +73,7 @@ async fn fresh_daemon_serves_history(
// Check socket and startup-lock reuse after shutdown.
for _ in 0..2 {
if !autostart {
daemon.foreground =
Some(Process::spawn(daemon.env.atuin(&["daemon", "start", "--show-logs"])));
tokio::time::timeout(TIMEOUT, async {
loop {
let process = daemon.foreground.as_mut().unwrap();
assert!(
process.child.try_wait().unwrap().is_none(),
"daemon exited during startup: {}",
process.logs()
);
if let Ok(mut client) = HistoryClient::new(daemon.env.socket()).await
&& let Ok(status) = client.status().await
&& status.healthy
{
break;
}
tokio::time::sleep(Duration::from_millis(10)).await;
}
})
.await
.expect("daemon never became healthy");
daemon.start_foreground().await;
}
// Spawn all writers first to exercise concurrent autostart.
@@ -151,6 +157,19 @@ async fn history_end_does_not_autostart_daemon(daemon: Daemon, #[case] exit: i64
assert!(daemon.env.run(&["daemon", "status"]).contains("Daemon is not running"));
}
/// A daemon started in the foreground listens on the socket a crashed one left behind, which
/// nothing answers on.
#[rstest]
#[tokio::test]
async fn foreground_daemon_replaces_a_stale_socket(mut daemon: Daemon) {
daemon.env.write_config("local_timeout = 15\n[daemon]\nenabled = true\nautostart = false\n");
// Bound and closed without being removed, as a crash leaves it.
drop(std::os::unix::net::UnixListener::bind(daemon.env.socket()).unwrap());
assert!(daemon.env.socket().exists());
daemon.start_foreground().await;
}
/// With the daemon enabled but not running (and not allowed to start), history is saved locally
/// instead of being dropped (#3866). The fallback warns, unless `--hook` silences logging.
#[rstest]
docs/docs/ai/sessions.md
+10 -2
View File
@@ -52,6 +52,11 @@ to your agents' tool calls too:
[`sensitive_files`](../configuration/config.md#sensitive_files)) is never
captured: a call that reads one, or a command naming one, is kept without its
output, and a call that writes to one without what it wrote.
- Codex runs its tools through a script, and each command and patch in the
script is judged as above. A script Atuin can't read in full (one that builds
a command or a file name as it runs) is kept as the name of its tool only. The
commands it runs still reach your history, filtered, through
[Codex's hooks](../guide/agent-hooks.md#codex).
Text that can't be redacted quickly isn't kept. To keep more out, see
[Keeping Secrets Out of Atuin](../guide/excluding-commands.md).
@@ -310,8 +315,11 @@ another (`claude`, `codex`, `opencode` or `pi`). `atuin resume <id> --fork`
forks it. It catches the session up with sync as the picker does, and fails
when that needs a choice, saying which: `--as-is` resumes this copy unchanged,
`--switch` switches it to another branch, and `--fork` forks it from its newest
branch. `--branch` picks the branch to catch up to, to switch to, or to fork
from: `this`, `@<host id>`, or the start of the id the error lists. `--as-is`,
branch. `--branch` picks the branch to catch up to, to switch to, to fork from,
or to continue with `--in`: `this`, `@<host id>`, or the start of the id the
error lists. A session that went different ways on several machines needs
`--branch` with `--in`, and with `--switch` when more than one branch could be
switched to; without it, Atuin lists the branches to choose from. `--as-is`,
`--switch` and `--branch` need an id that names a single session:
when it names several, Atuin lists them and asks you to be more specific.
`--print` prints the command instead of running it.
docs/docs/guide/agent-hooks.md
+10 -6
View File
@@ -31,7 +31,7 @@ When `atuin hook install` runs, it writes the agent's config file or extension t
| Agent | Config file / extension |
|-------|-------------------------|
| Claude Code | `~/.claude/settings.json` |
| Codex | `~/.codex/hooks.json` |
| Codex | `$CODEX_HOME/hooks.json` (by default `~/.codex/hooks.json`), trusted in `$CODEX_HOME/config.toml` |
| opencode | `~/.config/opencode/plugins/atuin.ts` |
| pi | `~/.pi/agent/extensions/atuin.ts` |
@@ -96,7 +96,9 @@ This adds hook entries to `~/.claude/settings.json`. Claude Code calls `atuin ho
atuin hook install codex
```
This adds hook entries to `~/.codex/hooks.json`. Codex calls `atuin hook codex` on each Bash tool use matching `^Bash$`.
This adds hook entries to `$CODEX_HOME/hooks.json` (by default `~/.codex/hooks.json`). Codex calls `atuin hook codex` on each Bash tool use matching `^Bash$`.
Codex 0.160 and later only run hooks you have trusted, so the install also records Atuin's hooks as trusted in `$CODEX_HOME/config.toml` (`[hooks.state]` entries). It trusts only Atuin's own hooks, as installed, and leaves your other hooks and settings as they are. If you installed the hooks with an earlier version of Atuin, run the install again to trust them.
### opencode
@@ -106,11 +108,13 @@ atuin hook install opencode
This writes Atuin's plugin to `~/.config/opencode/plugins/atuin.ts`, or to `$XDG_CONFIG_HOME/opencode/plugins/atuin.ts` if you have set `XDG_CONFIG_HOME`.
Then restart opencode. The plugin records every `bash` tool command with author `opencode`, using the tool's description as the entry's intent when opencode supplies one.
Then restart opencode. The plugin records every shell tool command with author `opencode`, using the tool's description as the entry's intent when opencode supplies one.
Two details are worth knowing:
opencode 1.x and 2.x load plugins differently, so the install writes the plugin for the opencode it finds (`opencode --version`). After upgrading to opencode 2, run `atuin hook install opencode` again.
- Commands you deny at the permission prompt aren't recorded. The plugin waits until opencode is cleared to run a command before opening a history entry.
A few details are worth knowing:
- Commands you deny at the permission prompt aren't recorded. With opencode 2, a command is recorded once it has finished, so its time is when it ended, and its duration includes any wait at the permission prompt. A command that timed out is recorded with exit code 124, and one left running in the background isn't recorded.
- Commands you run yourself in opencode -- with the `!command` prompt prefix or in a terminal pane -- aren't recorded either, since opencode didn't run them. Use Atuin's own shell integration for those.
### pi
@@ -161,4 +165,4 @@ hooks.PostToolUse: already installed, skipping
hooks.PostToolUseFailure: already installed, skipping
```
For opencode and pi, reinstalling will also skip if the installed extension already matches the bundled version.
For opencode and pi, reinstalling will also skip if the installed extension already matches the one for your agent's version.