633 lines
21 KiB
Rust
633 lines
21 KiB
Rust
use std::fs::{File, OpenOptions};
|
|
use std::io::Write;
|
|
use std::ops::Range;
|
|
use std::os::unix::fs::OpenOptionsExt;
|
|
use std::path::Path;
|
|
|
|
use anyhow::{Context, Result};
|
|
use serde::Deserialize;
|
|
|
|
use super::driver::{Event, SessionStatus, context_after};
|
|
|
|
pub use event_model::SeqEvent;
|
|
|
|
pub struct Transcript {
|
|
file: File,
|
|
next_seq: u64,
|
|
last_status: Option<SessionStatus>,
|
|
last_activity: Option<f64>,
|
|
context_tokens: Option<u64>,
|
|
}
|
|
|
|
impl Transcript {
|
|
pub fn open(path: &Path) -> Result<Self> {
|
|
// One pass for all three answers. They are wanted at the same moment
|
|
// by the same caller, and reading the file again for each doubled the
|
|
// cost of starting every session.
|
|
let existing = read_after(path, 0)?;
|
|
let last_seq = existing.last().map(|entry| entry.seq).unwrap_or(0);
|
|
let last_status = existing.iter().rev().find_map(|entry| match entry.event {
|
|
Event::Status { state } => Some(state),
|
|
_ => None,
|
|
});
|
|
let file = OpenOptions::new()
|
|
.create(true)
|
|
.append(true)
|
|
.mode(0o600)
|
|
.open(path)
|
|
.with_context(|| format!("open transcript {}", path.display()))?;
|
|
Ok(Self {
|
|
file,
|
|
next_seq: last_seq + 1,
|
|
last_status,
|
|
last_activity: existing.last().map(|entry| entry.ts),
|
|
context_tokens: existing
|
|
.iter()
|
|
.fold(None, |current, entry| context_after(current, &entry.event)),
|
|
})
|
|
}
|
|
|
|
/// Read from the file rather than assumed, because a server that has just
|
|
/// restarted has been told nothing. Assuming idle claimed a session was
|
|
/// waiting for you when it had exited hours earlier.
|
|
///
|
|
/// `None` for a transcript that never carried a status.
|
|
pub fn last_status(&self) -> Option<SessionStatus> {
|
|
self.last_status
|
|
}
|
|
|
|
/// `None` for a transcript with no lines, which is a session that genuinely
|
|
/// has not done anything. Its caller answers with when the session was
|
|
/// created, not with the clock.
|
|
pub fn last_activity(&self) -> Option<f64> {
|
|
self.last_activity
|
|
}
|
|
|
|
/// `None` for a transcript nothing has been measured in. That is not zero:
|
|
/// a server that has just restarted has been told nothing, and answering
|
|
/// zero would draw an empty context for a conversation that may be nearly
|
|
/// full.
|
|
pub fn context_tokens(&self) -> Option<u64> {
|
|
self.context_tokens
|
|
}
|
|
|
|
/// Appends `event`, assigning it the next sequence number. Flushed per
|
|
/// event: each line is tiny, and the transcript is the source of truth a
|
|
/// crash must not lose the tail of.
|
|
pub fn append(&mut self, event: Event, ts: f64) -> Result<SeqEvent> {
|
|
let entry = SeqEvent {
|
|
seq: self.next_seq,
|
|
ts,
|
|
event,
|
|
};
|
|
let mut line = serde_json::to_string(&entry).context("serialize event")?;
|
|
line.push('\n');
|
|
self.file
|
|
.write_all(line.as_bytes())
|
|
.context("append to transcript")?;
|
|
self.next_seq += 1;
|
|
Ok(entry)
|
|
}
|
|
}
|
|
|
|
/// `after` is a floor: nothing at or below it is returned, and the page stops
|
|
/// there rather than at `limit`. A phone holding a cached run passes the end of
|
|
/// what it already has, so the page is exactly the gap and never overlaps its
|
|
/// copy -- an overlap it cannot store, since a coalesced event cannot be cut at
|
|
/// a seq inside its own delta run.
|
|
pub fn read_window(
|
|
path: &Path,
|
|
before: Option<u64>,
|
|
after: Option<u64>,
|
|
limit: usize,
|
|
coalesce: bool,
|
|
) -> Result<Vec<SeqEvent>> {
|
|
let Some(indexed) = Indexed::read(path)? else {
|
|
return Ok(Vec::new());
|
|
};
|
|
let end = match before {
|
|
Some(before) => indexed.first_at_or_after(before)?,
|
|
None => indexed.lines.len(),
|
|
};
|
|
let start = match after {
|
|
Some(after) => indexed.first_at_or_after(after.saturating_add(1))?,
|
|
None => 0,
|
|
};
|
|
let start = start.min(end);
|
|
if coalesce && before.is_some() {
|
|
indexed.parse_coalesced(start, end, limit)
|
|
} else {
|
|
indexed.parse(start.max(end.saturating_sub(limit))..end)
|
|
}
|
|
}
|
|
|
|
pub const CATCH_UP_LIMIT: usize = 200;
|
|
|
|
/// Two answers rather than one list, because they mean different things to the
|
|
/// screen holding the cursor: one continues what it has, the other replaces it.
|
|
/// Collapsing them would leave the client splicing a window onto rows it has no
|
|
/// way to know are no longer adjacent -- a seam that looks like ordinary output.
|
|
#[derive(Debug, Clone, PartialEq)]
|
|
pub enum CatchUp {
|
|
Continue(Vec<SeqEvent>),
|
|
Restart(Vec<SeqEvent>),
|
|
}
|
|
|
|
pub fn catch_up(path: &Path, after: u64, limit: usize) -> Result<CatchUp> {
|
|
let Some(indexed) = Indexed::read(path)? else {
|
|
return Ok(CatchUp::Continue(Vec::new()));
|
|
};
|
|
let end = indexed.lines.len();
|
|
let start = indexed.first_at_or_after(after.saturating_add(1))?;
|
|
if end - start > limit {
|
|
return Ok(CatchUp::Restart(indexed.parse(end - limit..end)?));
|
|
}
|
|
Ok(CatchUp::Continue(indexed.parse(start..end)?))
|
|
}
|
|
|
|
pub fn read_after(path: &Path, after: u64) -> Result<Vec<SeqEvent>> {
|
|
let Some(indexed) = Indexed::read(path)? else {
|
|
return Ok(Vec::new());
|
|
};
|
|
let start = indexed.first_at_or_after(after.saturating_add(1))?;
|
|
indexed.parse(start..indexed.lines.len())
|
|
}
|
|
|
|
struct Indexed<'a> {
|
|
path: &'a Path,
|
|
text: String,
|
|
lines: Vec<Range<usize>>,
|
|
}
|
|
|
|
impl<'a> Indexed<'a> {
|
|
fn read(path: &'a Path) -> Result<Option<Self>> {
|
|
let text = match std::fs::read_to_string(path) {
|
|
Ok(text) => text,
|
|
Err(err) if err.kind() == std::io::ErrorKind::NotFound => return Ok(None),
|
|
Err(err) => {
|
|
return Err(err).with_context(|| format!("read transcript {}", path.display()));
|
|
}
|
|
};
|
|
let mut lines = Vec::new();
|
|
let mut start = 0;
|
|
while start < text.len() {
|
|
let end = text[start..]
|
|
.find('\n')
|
|
.map(|at| start + at)
|
|
.unwrap_or(text.len());
|
|
if !text[start..end].trim().is_empty() {
|
|
lines.push(start..end);
|
|
}
|
|
start = end + 1;
|
|
}
|
|
Ok(Some(Self { path, text, lines }))
|
|
}
|
|
|
|
/// A bisection, which is only correct because the file is in sequence order.
|
|
/// A line that cannot be read is reported here rather than silently treated
|
|
/// as out of range, because the answer would be a window off by however much
|
|
/// of the file the bad line hid.
|
|
fn first_at_or_after(&self, seq: u64) -> Result<usize> {
|
|
let (mut low, mut high) = (0, self.lines.len());
|
|
while low < high {
|
|
let middle = (low + high) / 2;
|
|
if self.seq_at(middle)? < seq {
|
|
low = middle + 1;
|
|
} else {
|
|
high = middle;
|
|
}
|
|
}
|
|
Ok(low)
|
|
}
|
|
|
|
fn seq_at(&self, index: usize) -> Result<u64> {
|
|
#[derive(Deserialize)]
|
|
struct JustSeq {
|
|
seq: u64,
|
|
}
|
|
let line = &self.text[self.lines[index].clone()];
|
|
let entry: JustSeq = serde_json::from_str(line)
|
|
.with_context(|| format!("bad transcript line in {}", self.path.display()))?;
|
|
Ok(entry.seq)
|
|
}
|
|
|
|
fn parse(&self, range: Range<usize>) -> Result<Vec<SeqEvent>> {
|
|
self.lines[range]
|
|
.iter()
|
|
.map(|at| {
|
|
serde_json::from_str(&self.text[at.clone()])
|
|
.with_context(|| format!("bad transcript line in {}", self.path.display()))
|
|
})
|
|
.collect()
|
|
}
|
|
|
|
fn parse_one(&self, index: usize) -> Result<SeqEvent> {
|
|
serde_json::from_str(&self.text[self.lines[index].clone()])
|
|
.with_context(|| format!("bad transcript line in {}", self.path.display()))
|
|
}
|
|
|
|
fn parse_coalesced(&self, start: usize, end: usize, limit: usize) -> Result<Vec<SeqEvent>> {
|
|
let mut out: Vec<SeqEvent> = Vec::new();
|
|
let mut run: Option<(u64, f64, Vec<String>)> = None;
|
|
let flush = |run: &mut Option<(u64, f64, Vec<String>)>, out: &mut Vec<SeqEvent>| {
|
|
if let Some((seq, ts, mut deltas)) = run.take() {
|
|
deltas.reverse();
|
|
out.push(SeqEvent {
|
|
seq,
|
|
ts,
|
|
event: Event::AssistantText {
|
|
delta: deltas.concat(),
|
|
},
|
|
});
|
|
}
|
|
};
|
|
let mut index = end;
|
|
while index > start {
|
|
if out.len() >= limit && run.is_none() {
|
|
break;
|
|
}
|
|
index -= 1;
|
|
let entry = self.parse_one(index)?;
|
|
if let Event::AssistantText { delta } = entry.event {
|
|
match run {
|
|
Some((ref mut seq, ref mut ts, ref mut deltas)) => {
|
|
*seq = entry.seq;
|
|
*ts = entry.ts;
|
|
deltas.push(delta);
|
|
}
|
|
None => run = Some((entry.seq, entry.ts, vec![delta])),
|
|
}
|
|
} else {
|
|
flush(&mut run, &mut out);
|
|
out.push(entry);
|
|
}
|
|
}
|
|
flush(&mut run, &mut out);
|
|
out.reverse();
|
|
Ok(out)
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use crate::session::driver::{QuestionOption, SessionStatus};
|
|
|
|
fn text(delta: &str) -> Event {
|
|
Event::AssistantText {
|
|
delta: delta.to_string(),
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn assigns_increasing_seqs_and_replays_after_a_cursor() {
|
|
let dir = tempfile::tempdir().expect("tempdir");
|
|
let path = dir.path().join("transcript.jsonl");
|
|
|
|
let mut transcript = Transcript::open(&path).expect("open");
|
|
assert_eq!(transcript.append(text("a"), 1.0).expect("append").seq, 1);
|
|
assert_eq!(transcript.append(text("b"), 2.0).expect("append").seq, 2);
|
|
assert_eq!(transcript.append(text("c"), 3.0).expect("append").seq, 3);
|
|
|
|
let replay = read_after(&path, 1).expect("read");
|
|
assert_eq!(replay.len(), 2);
|
|
assert_eq!(replay[0].seq, 2);
|
|
assert_eq!(replay[0].event, text("b"));
|
|
assert_eq!(replay[1].seq, 3);
|
|
|
|
assert!(read_after(&path, 3).expect("read").is_empty());
|
|
}
|
|
|
|
#[test]
|
|
fn reopening_continues_the_numbering() {
|
|
let dir = tempfile::tempdir().expect("tempdir");
|
|
let path = dir.path().join("transcript.jsonl");
|
|
|
|
let mut transcript = Transcript::open(&path).expect("open");
|
|
transcript.append(text("a"), 1.0).expect("append");
|
|
transcript.append(text("b"), 2.0).expect("append");
|
|
drop(transcript);
|
|
|
|
let mut reopened = Transcript::open(&path).expect("reopen");
|
|
assert_eq!(reopened.append(text("c"), 3.0).expect("append").seq, 3);
|
|
}
|
|
|
|
#[test]
|
|
fn a_short_backlog_continues_and_a_long_one_restarts() {
|
|
let dir = tempfile::tempdir().expect("tempdir");
|
|
let path = dir.path().join("transcript.jsonl");
|
|
let mut transcript = Transcript::open(&path).expect("open");
|
|
for n in 0..10 {
|
|
transcript
|
|
.append(text(&n.to_string()), 0.0)
|
|
.expect("append");
|
|
}
|
|
|
|
let CatchUp::Continue(events) = catch_up(&path, 7, 5).expect("catch up") else {
|
|
panic!("a backlog of 3 should continue");
|
|
};
|
|
assert_eq!(events.len(), 3);
|
|
assert_eq!(events[0].seq, 8);
|
|
|
|
let CatchUp::Restart(events) = catch_up(&path, 0, 5).expect("catch up") else {
|
|
panic!("a backlog of 10 should restart");
|
|
};
|
|
assert_eq!(events.len(), 5);
|
|
assert_eq!(events[0].seq, 6);
|
|
assert_eq!(events[4].seq, 10);
|
|
|
|
assert!(matches!(
|
|
catch_up(&path, 5, 5).expect("catch up"),
|
|
CatchUp::Continue(_)
|
|
));
|
|
}
|
|
|
|
#[test]
|
|
fn reopening_reports_the_state_it_was_last_left_in() {
|
|
let dir = tempfile::tempdir().expect("tempdir");
|
|
let path = dir.path().join("transcript.jsonl");
|
|
|
|
assert_eq!(Transcript::open(&path).expect("open").last_status(), None);
|
|
|
|
let mut transcript = Transcript::open(&path).expect("open");
|
|
transcript
|
|
.append(
|
|
Event::Status {
|
|
state: SessionStatus::Running,
|
|
},
|
|
1.0,
|
|
)
|
|
.expect("append");
|
|
transcript
|
|
.append(
|
|
Event::Status {
|
|
state: SessionStatus::Exited,
|
|
},
|
|
2.0,
|
|
)
|
|
.expect("append");
|
|
transcript.append(text("trailing"), 3.0).expect("append");
|
|
drop(transcript);
|
|
|
|
let reopened = Transcript::open(&path).expect("reopen");
|
|
assert_eq!(reopened.last_status(), Some(SessionStatus::Exited));
|
|
assert_eq!(reopened.next_seq, 4);
|
|
}
|
|
|
|
#[test]
|
|
fn a_window_is_the_events_before_a_cursor_and_nothing_else() {
|
|
let dir = tempfile::tempdir().expect("tempdir");
|
|
let path = dir.path().join("transcript.jsonl");
|
|
let mut transcript = Transcript::open(&path).expect("open");
|
|
for n in 1..=10 {
|
|
transcript
|
|
.append(text(&n.to_string()), 0.0)
|
|
.expect("append");
|
|
}
|
|
|
|
let newest = read_window(&path, None, None, 3, false).expect("window");
|
|
assert_eq!(
|
|
newest.iter().map(|entry| entry.seq).collect::<Vec<_>>(),
|
|
[8, 9, 10]
|
|
);
|
|
|
|
let older = read_window(&path, Some(8), None, 3, false).expect("window");
|
|
assert_eq!(
|
|
older.iter().map(|entry| entry.seq).collect::<Vec<_>>(),
|
|
[5, 6, 7]
|
|
);
|
|
|
|
assert_eq!(
|
|
read_window(&path, None, None, 100, false)
|
|
.expect("window")
|
|
.len(),
|
|
10
|
|
);
|
|
|
|
assert!(
|
|
read_window(&path, Some(1), None, 3, false)
|
|
.expect("window")
|
|
.is_empty()
|
|
);
|
|
assert!(
|
|
read_window(&dir.path().join("nope.jsonl"), None, None, 3, false)
|
|
.expect("window")
|
|
.is_empty()
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn a_floor_stops_a_page_at_what_the_caller_already_holds() {
|
|
let dir = tempfile::tempdir().expect("tempdir");
|
|
let path = dir.path().join("transcript.jsonl");
|
|
let mut transcript = Transcript::open(&path).expect("open");
|
|
for n in 1..=10 {
|
|
transcript
|
|
.append(text(&n.to_string()), 0.0)
|
|
.expect("append");
|
|
}
|
|
|
|
let page = read_window(&path, Some(9), Some(5), 100, false).expect("window");
|
|
assert_eq!(
|
|
page.iter().map(|entry| entry.seq).collect::<Vec<_>>(),
|
|
[6, 7, 8]
|
|
);
|
|
|
|
let page = read_window(&path, Some(9), Some(2), 3, false).expect("window");
|
|
assert_eq!(
|
|
page.iter().map(|entry| entry.seq).collect::<Vec<_>>(),
|
|
[6, 7, 8]
|
|
);
|
|
|
|
assert!(
|
|
read_window(&path, Some(4), Some(9), 10, false)
|
|
.expect("window")
|
|
.is_empty()
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn a_floor_inside_a_delta_run_leaves_the_partial_run_it_cuts() {
|
|
let dir = tempfile::tempdir().expect("tempdir");
|
|
let path = dir.path().join("transcript.jsonl");
|
|
let mut transcript = Transcript::open(&path).expect("open");
|
|
for d in ["a", "b", "c", "d"] {
|
|
transcript.append(text(d), 0.0).expect("append"); // seq 1..4
|
|
}
|
|
transcript
|
|
.append(
|
|
Event::ToolStart {
|
|
id: "t".into(),
|
|
tool: "Bash".into(),
|
|
input: serde_json::Value::Null,
|
|
},
|
|
0.0,
|
|
)
|
|
.expect("append"); // seq 5
|
|
|
|
let rows = read_window(&path, Some(6), Some(2), 10, true).expect("window");
|
|
assert_eq!(rows.len(), 2);
|
|
assert!(matches!(
|
|
&rows[0],
|
|
SeqEvent { seq: 3, event: Event::AssistantText { delta }, .. } if delta == "cd"
|
|
));
|
|
assert!(matches!(
|
|
&rows[1],
|
|
SeqEvent {
|
|
seq: 5,
|
|
event: Event::ToolStart { .. },
|
|
..
|
|
}
|
|
));
|
|
|
|
let rows = read_window(&path, Some(6), None, 10, true).expect("window");
|
|
assert_eq!(rows.len(), 2);
|
|
assert!(matches!(
|
|
&rows[0],
|
|
SeqEvent { seq: 1, event: Event::AssistantText { delta }, .. } if delta == "abcd"
|
|
));
|
|
}
|
|
|
|
#[test]
|
|
fn coalescing_counts_rows_and_joins_delta_runs() {
|
|
let dir = tempfile::tempdir().expect("tempdir");
|
|
let path = dir.path().join("transcript.jsonl");
|
|
let mut transcript = Transcript::open(&path).expect("open");
|
|
for d in ["a", "b", "c"] {
|
|
transcript.append(text(d), 0.0).expect("append"); // seq 1..3
|
|
}
|
|
transcript
|
|
.append(
|
|
Event::ToolStart {
|
|
id: "t".into(),
|
|
tool: "Bash".into(),
|
|
input: serde_json::Value::Null,
|
|
},
|
|
0.0,
|
|
)
|
|
.expect("append"); // seq 4
|
|
for d in ["d", "e", "f"] {
|
|
transcript.append(text(d), 0.0).expect("append"); // seq 5..7
|
|
}
|
|
|
|
let rows = read_window(&path, Some(8), None, 3, true).expect("window");
|
|
assert_eq!(rows.len(), 3);
|
|
assert!(matches!(
|
|
&rows[0],
|
|
SeqEvent { seq: 1, event: Event::AssistantText { delta }, .. } if delta == "abc"
|
|
));
|
|
assert!(matches!(
|
|
&rows[1],
|
|
SeqEvent {
|
|
seq: 4,
|
|
event: Event::ToolStart { .. },
|
|
..
|
|
}
|
|
));
|
|
assert!(matches!(
|
|
&rows[2],
|
|
SeqEvent { seq: 5, event: Event::AssistantText { delta }, .. } if delta == "def"
|
|
));
|
|
|
|
let older = read_window(&path, Some(1), None, 3, true).expect("window");
|
|
assert!(older.is_empty());
|
|
|
|
let newest = read_window(&path, None, None, 2, true).expect("window");
|
|
assert_eq!(newest.iter().map(|e| e.seq).collect::<Vec<_>>(), [6, 7]);
|
|
}
|
|
|
|
#[test]
|
|
fn a_line_read_back_is_the_line_that_was_written() {
|
|
let dir = tempfile::tempdir().expect("tempdir");
|
|
let path = dir.path().join("transcript.jsonl");
|
|
let mut transcript = Transcript::open(&path).expect("open");
|
|
transcript
|
|
.append(text("hello"), 1788546972.6030757)
|
|
.expect("append");
|
|
drop(transcript);
|
|
|
|
let written = std::fs::read_to_string(&path).expect("read");
|
|
let entry = read_window(&path, None, None, 10, false).expect("window");
|
|
assert_eq!(
|
|
serde_json::to_string(&entry[0]).expect("serialize"),
|
|
written.trim()
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn a_missing_file_reads_as_empty() {
|
|
let dir = tempfile::tempdir().expect("tempdir");
|
|
assert!(
|
|
read_after(&dir.path().join("nope.jsonl"), 0)
|
|
.expect("read")
|
|
.is_empty()
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn round_trips_every_event_shape() {
|
|
let dir = tempfile::tempdir().expect("tempdir");
|
|
let path = dir.path().join("transcript.jsonl");
|
|
let events = vec![
|
|
Event::UserMessage {
|
|
id: None,
|
|
text: "hi".into(),
|
|
attachments: Vec::new(),
|
|
},
|
|
text("hello"),
|
|
Event::ToolStart {
|
|
id: "t1".into(),
|
|
tool: "bash".into(),
|
|
input: serde_json::json!({"command": "ls"}),
|
|
},
|
|
Event::ToolUpdate {
|
|
id: "t1".into(),
|
|
output: "partial".into(),
|
|
},
|
|
Event::ToolEnd {
|
|
id: "t1".into(),
|
|
output: "done".into(),
|
|
is_error: false,
|
|
},
|
|
Event::Image {
|
|
image: "img1".into(),
|
|
about: None,
|
|
},
|
|
Event::Question {
|
|
id: "q1".into(),
|
|
prompt: "Allow?".into(),
|
|
header: None,
|
|
options: vec![QuestionOption::plain("Yes"), QuestionOption::plain("No")],
|
|
multi_select: false,
|
|
about: None,
|
|
},
|
|
Event::Answered {
|
|
id: "q1".into(),
|
|
answers: vec!["Yes".into()],
|
|
},
|
|
Event::Status {
|
|
state: SessionStatus::Idle,
|
|
},
|
|
Event::UsageDelta {
|
|
tokens: 42,
|
|
context: Some(42),
|
|
},
|
|
Event::Error {
|
|
message: "boom".into(),
|
|
},
|
|
];
|
|
|
|
let mut transcript = Transcript::open(&path).expect("open");
|
|
for event in &events {
|
|
transcript.append(event.clone(), 0.0).expect("append");
|
|
}
|
|
|
|
let replayed: Vec<Event> = read_after(&path, 0)
|
|
.expect("read")
|
|
.into_iter()
|
|
.map(|entry| entry.event)
|
|
.collect();
|
|
assert_eq!(replayed, events);
|
|
}
|
|
}
|