Hare/yoi
1
1
Fork 0
Code Actions Packages Releases Activity

fix: restore console run accounting from persisted logs

Browse Source
This commit is contained in:
Hare
2026-10-09 02:27:27 +09:00
parent e5ea7544d4
commit f1287f27e7
19 changed files with 1023 additions and 39 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/protocol/src/lib.rs
+35 -2
View File
@@ -1054,6 +1054,23 @@ pub enum SessionSnapshotEntryData {
source: RunResumeSource,
active_run_turn_count: usize,
},
/// Snapshot-only replay of persisted invocation/usage records. These are
/// accounting metadata, never model context or conversation rows.
Invoke {
trigger: InvokeKind,
},
Usage {
input_tokens: u64,
cache_read_input_tokens: u64,
output_tokens: u64,
},
/// Display-only accounting reconstructed from persisted run/usage records.
RunStats {
elapsed_ms: u64,
requests: u64,
upload_tokens: u64,
output_tokens: u64,
},
/// Durable terminal for an intentional cancellation. Consumers must not
/// present this lifecycle result as an execution failure.
RunCancelled,
@@ -1177,6 +1194,9 @@ pub enum Event {
/// following `UserMessage` / `SystemItem` event.
InvokeStart {
kind: InvokeKind,
/// Timestamp of the persisted Invoke, not the client's attach time.
#[serde(default, skip_serializing_if = "Option::is_none")]
timestamp_ms: Option<u64>,
},
/// One AgentTurn boundary opened. An AgentTurn is a maximal run of
/// LLM generation calls whose input messages are identical (i.e.
@@ -1293,6 +1313,9 @@ pub enum Event {
/// actually paid full price to send on this request, which is what
/// the TUI status line accumulates per turn.
Usage {
/// Present for usage delivered from the committed session-log lane.
#[serde(default, skip_serializing_if = "Option::is_none")]
timestamp_ms: Option<u64>,
input_tokens: Option<u64>,
output_tokens: Option<u64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
@@ -2335,19 +2358,29 @@ mod tests {
InvokeKind::SystemReminder,
InvokeKind::Wakeup,
] {
let event = Event::InvokeStart { kind };
let event = Event::InvokeStart {
kind,
timestamp_ms: Some(1234),
};
let json = serde_json::to_string(&event).unwrap();
let parsed: serde_json::Value = serde_json::from_str(&json).unwrap();
assert_eq!(parsed["event"], "invoke_start");
let decoded: Event = serde_json::from_str(&json).unwrap();
match decoded {
Event::InvokeStart { kind: k } => assert_eq!(k, kind),
Event::InvokeStart {
kind: k,
timestamp_ms,
} => {
assert_eq!(k, kind);
assert_eq!(timestamp_ms, Some(1234));
}
other => panic!("expected InvokeStart, got {other:?}"),
}
}
let parsed: serde_json::Value = serde_json::from_str(
&serde_json::to_string(&Event::InvokeStart {
kind: InvokeKind::UserSend,
timestamp_ms: None,
})
.unwrap(),
)
crates/session-store/src/public_snapshot.rs
+472 -8
View File
@@ -1,3 +1,5 @@
mod run_stats;
use base64::{Engine as _, engine::general_purpose::URL_SAFE_NO_PAD};
use protocol::{
Segment, SessionContentPart, SessionConversationTurn, SessionEntryProvenance,
@@ -850,6 +852,17 @@ fn read_history_turns_backward(
let mut pending_entries = Vec::new();
let mut seen_entries = HashSet::new();
let mut corrections = HashMap::<String, SessionSnapshotEntryData>::new();
// Accounting needs the whole Invoke range, including the marker immediately
// before the oldest user entry on this page. The existing scan budgets also
// bound this suffix; no unbounded second pass over the Session is needed.
let mut accounting_records = Vec::new();
let finish = |mut turns: Vec<PositionedHistoryTurn>, records: &[LogEntry]| {
let stats = run_stats::project(records.iter().rev());
for turn in &mut turns {
run_stats::append_to(&mut turn.turn.entries, &stats);
}
turns
};
// Carry only bounded log positions, never correction text or image bodies,
// across pages. Re-read and validate the referenced records from adopted
@@ -994,11 +1007,58 @@ fn read_history_turns_backward(
&mut pending_entries,
&mut turns,
) {
return Ok(turns);
return Ok(finish(turns, &accounting_records));
}
continue;
}
// TurnEnd is the lineage boundary, but Worker persists final usage and
// the run outcome immediately after it. For a compacted ancestor, read
// only this accounting tail, never adopt later conversation records.
if before_offset == adopted_end
&& segment_index > 0
&& matches!(
lineage[segment_index - 1].origin.as_ref().map(|o| o.kind),
Some(LineageOriginKind::Compact)
)
{
let mut offset = adopted_end;
let mut tail = Vec::new();
while offset < segment.file_len {
let (record, bytes) = store
.read_next_log_record_read_only_bounded(
session_id,
segment.segment_id,
offset,
limits.max_scan_bytes.saturating_sub(*scanned_bytes),
)
.map_err(map_history_store_error)?;
*scanned_bytes = scanned_bytes
.checked_add(bytes)
.ok_or(RetainedHistoryReadError::ResourceLimit)?;
let Some(record) = record else { break };
*scanned_entries = scanned_entries
.checked_add(persisted_entry_units(std::slice::from_ref(&record.entry)))
.ok_or(RetainedHistoryReadError::ResourceLimit)?;
if *scanned_entries > limits.max_entries {
return Err(RetainedHistoryReadError::ResourceLimit);
}
offset = record.end_offset;
match record.entry {
entry @ LogEntry::LlmUsage { .. } => tail.push(entry),
LogEntry::Extension { .. } => {}
entry @ (LogEntry::RunCompleted { .. }
| LogEntry::RunYielded { .. }
| LogEntry::RunCancelled { .. }
| LogEntry::RunErrored { .. }) => {
tail.push(entry);
break;
}
_ => break,
}
}
accounting_records.extend(tail.into_iter().rev());
}
let mut reader = store
.open_retained_segment_reader(session_id, segment.segment_id, before_offset)
.map_err(map_history_store_error)?;
@@ -1025,6 +1085,16 @@ fn read_history_turns_backward(
}
continue;
}
let at_invoke = matches!(record.entry, LogEntry::Invoke { .. });
accounting_records.push(record.entry.clone());
if turns.len() >= turn_limit {
if at_invoke {
return Ok(finish(turns, &accounting_records));
}
// Older records here supply accounting only, never additional
// conversation turns beyond the requested page boundary.
continue;
}
if capture_tool_correction(&record.entry, &mut corrections) {
let LogEntry::ToolResultCorrected { entry, .. } = &record.entry else {
unreachable!()
@@ -1044,7 +1114,7 @@ fn read_history_turns_backward(
&project_history_record(session_id, segment.segment_id, &record.entry),
&corrections,
);
if collect_history_entries_backward(
collect_history_entries_backward(
&entries,
segment.segment_id,
record.start_offset,
@@ -1054,12 +1124,11 @@ fn read_history_turns_backward(
&mut seen_entries,
&mut pending_entries,
&mut turns,
) {
return Ok(turns);
}
);
}
if segment_index + 1 == lineage.len()
if turns.len() < turn_limit
&& segment_index + 1 == lineage.len()
&& collect_history_entries_backward(
&corrected_history_entries(&segment.seed_entries, &corrections),
segment.segment_id,
@@ -1072,10 +1141,10 @@ fn read_history_turns_backward(
&mut turns,
)
{
return Ok(turns);
return Ok(finish(turns, &accounting_records));
}
}
Ok(turns)
Ok(finish(turns, &accounting_records))
}
fn capture_tool_correction(
@@ -1322,6 +1391,148 @@ pub fn project_current_session_snapshot(log: &[LogEntry]) -> SessionSnapshot {
project_session_snapshot(session_id.unwrap_or_else(SessionId::nil), log)
}
fn run_accounting_entries(session_id: &SessionId, log: &[LogEntry]) -> Vec<SessionSnapshotEntry> {
let start = log
.iter()
.rposition(|entry| matches!(entry, LogEntry::Invoke { .. }))
.unwrap_or(0);
log.iter()
.enumerate()
.skip(start)
.filter_map(|(index, record)| {
let (ts, data) = match record {
LogEntry::Invoke { ts, trigger } => {
(*ts, SessionSnapshotEntryData::Invoke { trigger: *trigger })
}
LogEntry::LlmUsage {
ts,
input_total_tokens,
cache_read_tokens,
output_tokens,
..
} => (
*ts,
SessionSnapshotEntryData::Usage {
input_tokens: *input_total_tokens,
cache_read_input_tokens: *cache_read_tokens,
output_tokens: *output_tokens,
},
),
_ => return None,
};
Some(legacy_entry(session_id, None, index, 0, ts, data))
})
.collect()
}
/// Reconstruct snapshot accounting from existing compacted ancestors, without a
/// second accumulator or checkpoint. The live suffix is the caller's frozen
/// mirror prefix; never reread the active segment after subscribing.
pub fn project_current_session_snapshot_with_accounting(
store: &dyn crate::Store,
log: &[LogEntry],
) -> SessionSnapshot {
let mut snapshot = project_current_session_snapshot(log);
let session_id = log
.iter()
.find_map(|entry| match entry {
LogEntry::AnnotatedSegmentStart { session_id, .. } => Some(*session_id),
_ => None,
})
.unwrap_or_else(SessionId::nil);
snapshot
.entries
.extend(run_accounting_entries(&session_id, log));
if log
.iter()
.any(|entry| matches!(entry, LogEntry::Invoke { .. }))
{
return snapshot;
}
let mut source = log.to_vec();
let mut ancestors = Vec::new();
let mut visited = HashSet::new();
let mut count = log.len();
for _ in 0..DEFAULT_RETAINED_HISTORY_MAX_SEGMENTS {
let Some((session_id, origin)) = source.iter().find_map(|entry| match entry {
LogEntry::AnnotatedSegmentStart {
session_id,
compacted_from: Some(origin),
..
} => Some((*session_id, origin.clone())),
_ => None,
}) else {
break;
};
if !visited.insert(origin.segment_id) {
break;
}
let Ok(mut entries) = store.read_all(session_id, origin.segment_id) else {
break;
};
count = count.saturating_add(entries.len());
if count > DEFAULT_RETAINED_HISTORY_MAX_ENTRIES {
break;
}
// Compaction adopts through TurnEnd plus its trailing usage/outcome,
// never through a subsequent Invoke or unrelated conversation.
if let Some(end) = entries.iter().position(|entry| {
matches!(entry, LogEntry::TurnEnd { turn_count, .. } if *turn_count == origin.at_turn_index)
}) {
let mut cutoff = end + 1;
while cutoff < entries.len()
&& matches!(entries[cutoff],
LogEntry::LlmUsage { .. } | LogEntry::Extension { .. }
| LogEntry::RunCompleted { .. } | LogEntry::RunYielded { .. }
| LogEntry::RunErrored { .. } | LogEntry::RunCancelled { .. })
{
cutoff += 1;
}
entries.truncate(cutoff);
} else if matches!(entries.first(), Some(LogEntry::AnnotatedSegmentStart {
compacted_from: Some(inherited), ..
}) if origin.at_turn_index <= inherited.at_turn_index) {
// Consecutive compactions can adopt only the inherited seed.
entries.truncate(1);
} else {
break;
}
let complete = entries
.iter()
.any(|entry| matches!(entry, LogEntry::Invoke { .. }));
source = entries.clone();
ancestors.push(entries);
if complete {
break;
}
}
if !ancestors.is_empty() {
let combined: Vec<LogEntry> = ancestors
.into_iter()
.rev()
.flatten()
.chain(log.iter().cloned())
.collect();
snapshot.entries.retain(|entry| {
!matches!(
entry.data,
SessionSnapshotEntryData::Invoke { .. } | SessionSnapshotEntryData::Usage { .. }
)
});
let session_id = log
.iter()
.find_map(|entry| match entry {
LogEntry::AnnotatedSegmentStart { session_id, .. } => Some(*session_id),
_ => None,
})
.unwrap_or_else(SessionId::nil);
snapshot
.entries
.extend(run_accounting_entries(&session_id, &combined));
}
snapshot
}
/// Project the current durable segment into the only public session-history
/// representation. Append-log records remain an internal persistence format.
pub fn project_session_snapshot(session_id: SessionId, log: &[LogEntry]) -> SessionSnapshot {
@@ -1523,6 +1734,7 @@ pub(crate) fn project_session_snapshot_for_segment(
}
}
run_stats::append_to(&mut entries, &run_stats::project(log));
SessionSnapshot {
pending_submissions: protocol::PendingSubmissionsSnapshot::default(),
entries,
@@ -1800,6 +2012,258 @@ mod tests {
LoggedWorkerSubject, Store, WorkerMetadataStore,
};
fn usage(ts: u64) -> LogEntry {
LogEntry::LlmUsage {
ts,
history_len: 2,
input_total_tokens: 2500,
cache_read_tokens: 1000,
cache_write_tokens: 0,
output_tokens: 200,
}
}
fn complete(ts: u64) -> LogEntry {
LogEntry::RunCompleted {
ts,
interrupted: false,
result: agen::EngineResult::Finished,
active_run_turn_count: None,
}
}
fn stats_entries(entries: &[SessionSnapshotEntry]) -> Vec<SessionSnapshotEntry> {
entries
.iter()
.filter(|entry| matches!(entry.data, SessionSnapshotEntryData::RunStats { .. }))
.cloned()
.collect()
}
#[test]
fn run_stats_restore_identically_from_snapshot_and_history_pages() {
let session_id = SessionId::now_v7();
let segment_id = SegmentId::now_v7();
let mut log = vec![segment_start(session_id, vec![], None, None)];
for turn in 1..=6 {
append_turn(&mut log, turn);
log.push(usage(turn as u64 * 10 + 4));
log.push(complete(turn as u64 * 10 + 8));
}
let expected = stats_entries(&project_current_session_snapshot(&log).entries);
assert_eq!(expected.len(), 6);
assert_eq!(
expected[0].data,
SessionSnapshotEntryData::RunStats {
elapsed_ms: 8,
requests: 1,
upload_tokens: 1500,
output_tokens: 200,
}
);
let (_root, path) =
persist_history_fixture("stats", session_id, segment_id, vec![(segment_id, log)]);
let mut cursor = None;
let mut actual = Vec::new();
loop {
let page = read_retained_session_history_page(
&path,
"stats",
cursor.as_deref(),
Some(1),
RetainedHistoryReadLimits::default(),
)
.unwrap();
assert_eq!(page.turns.len(), 1);
actual.splice(0..0, stats_entries(&page.turns[0].entries));
cursor = page.next_cursor;
if !page.has_more {
break;
}
}
assert_eq!(actual, expected);
}
#[test]
fn run_stats_span_compaction_without_counting_seed_history_twice() {
let session_id = SessionId::now_v7();
let source_id = SegmentId::now_v7();
let active_id = SegmentId::now_v7();
let mut source = vec![segment_start(session_id, vec![], None, None)];
append_turn(&mut source, 1);
source.push(usage(15));
source.push(LogEntry::RunYielded {
ts: 16,
entry_id: None,
reason: protocol::RunYieldReason::Compaction,
active_run_turn_count: 1,
});
let active = vec![
segment_start(
session_id,
vec![user_message(1), assistant_message(1, "final")],
None,
Some(SegmentOrigin {
segment_id: source_id,
at_turn_index: 1,
}),
),
LogEntry::RunResumed {
ts: 20,
entry_id: None,
source: protocol::RunResumeSource::Compaction,
active_run_turn_count: 1,
},
usage(25),
LogEntry::AnnotatedAssistantItem {
ts: 28,
entry: assistant_message(1, "after-compact"),
},
complete(30),
];
assert!(
stats_entries(&project_current_session_snapshot(&active).entries).is_empty(),
"a partial segment must not invent the missing start or traffic"
);
let (_root, path) = persist_history_fixture(
"stats",
session_id,
active_id,
vec![(source_id, source), (active_id, active)],
);
let page = read_retained_session_history_page(
&path,
"stats",
None,
Some(1),
RetainedHistoryReadLimits::default(),
)
.unwrap();
let stats = stats_entries(&page.turns[0].entries);
assert_eq!(stats.len(), 1);
assert_eq!(stats[0].entry_id, "run-stats:assistant-1-after-compact");
assert_eq!(
stats[0].data,
SessionSnapshotEntryData::RunStats {
elapsed_ms: 20,
requests: 2,
upload_tokens: 3000,
output_tokens: 400,
}
);
}
#[test]
fn reconnect_accounting_replays_existing_usage_across_compaction() {
let session_id = SessionId::now_v7();
let source_id = SegmentId::now_v7();
let active_id = SegmentId::now_v7();
let mut source = vec![segment_start(session_id, vec![], None, None)];
append_turn(&mut source, 1);
source.push(usage(15));
let active = vec![
segment_start(
session_id,
vec![],
None,
Some(SegmentOrigin {
segment_id: source_id,
at_turn_index: 1,
}),
),
usage(25),
];
let (_root, path) = persist_history_fixture(
"accounting",
session_id,
active_id,
vec![(source_id, source), (active_id, active.clone())],
);
let store = WorkerSessionStore::new(path.join("session")).unwrap();
let snapshot = project_current_session_snapshot_with_accounting(&store, &active);
let records: Vec<_> = snapshot
.entries
.iter()
.filter(|entry| {
matches!(
entry.data,
SessionSnapshotEntryData::Invoke { .. }
| SessionSnapshotEntryData::Usage { .. }
)
})
.collect();
assert_eq!(records.len(), 3);
assert_eq!(
records
.iter()
.map(|entry| entry.timestamp)
.collect::<Vec<_>>(),
[10, 15, 25]
);
assert!(matches!(
records[0].data,
SessionSnapshotEntryData::Invoke {
trigger: InvokeKind::UserSend
}
));
assert!(matches!(
records[1].data,
SessionSnapshotEntryData::Usage {
input_tokens: 2500,
cache_read_input_tokens: 1000,
output_tokens: 200
}
));
assert!(matches!(
records[2].data,
SessionSnapshotEntryData::Usage {
input_tokens: 2500,
cache_read_input_tokens: 1000,
output_tokens: 200
}
));
// A second compaction can adopt only the first one's inherited seed.
// Its unadopted usage must not leak into the restored totals.
let next = vec![
segment_start(
session_id,
vec![],
None,
Some(SegmentOrigin {
segment_id: active_id,
at_turn_index: 1,
}),
),
usage(35),
];
let snapshot = project_current_session_snapshot_with_accounting(&store, &next);
let timestamps: Vec<_> = snapshot
.entries
.iter()
.filter_map(|entry| {
matches!(
entry.data,
SessionSnapshotEntryData::Invoke { .. }
| SessionSnapshotEntryData::Usage { .. }
)
.then_some(entry.timestamp)
})
.collect();
assert_eq!(timestamps, [10, 15, 35]);
}
#[test]
fn run_stats_do_not_complete_an_active_or_unmeasured_run() {
let mut log = Vec::new();
append_turn(&mut log, 1);
log.push(usage(15));
assert!(stats_entries(&project_current_session_snapshot(&log).entries).is_empty());
log.pop();
log.push(complete(20));
assert!(stats_entries(&project_current_session_snapshot(&log).entries).is_empty());
}
fn history_message(
entry_id: impl Into<String>,
role: LoggedRole,
crates/session-store/src/public_snapshot/run_stats.rs
+133
View File
@@ -0,0 +1,133 @@
//! Display-only accounting; never inserted into model history or written back
//! to the journal. Both snapshot and bounded history paging use this projection.
use std::collections::HashMap;
use protocol::{SessionEntryProvenance, SessionSnapshotEntry, SessionSnapshotEntryData};
use crate::LogEntry;
#[derive(Default)]
struct Activity {
started_at: u64,
requests: u64,
upload: u64,
output: u64,
anchor: Option<String>,
}
pub(super) fn project<'a>(
records: impl IntoIterator<Item = &'a LogEntry>,
) -> HashMap<String, SessionSnapshotEntry> {
let mut active: Option<Activity> = None;
let mut stats = HashMap::new();
for record in records {
match record {
LogEntry::Invoke { ts, .. } => {
active = Some(Activity {
started_at: *ts,
..Activity::default()
})
}
LogEntry::LlmUsage {
input_total_tokens,
cache_read_tokens,
output_tokens,
..
} => {
if let Some(run) = &mut active {
run.requests += 1;
run.upload = run
.upload
.saturating_add(input_total_tokens.saturating_sub(*cache_read_tokens));
run.output = run.output.saturating_add(*output_tokens);
}
}
LogEntry::AnnotatedUserInput { history, .. } => {
if let Some(run) = &mut active {
for entry in history {
if super::project_item(&entry.metadata.entry_id.0, &entry.item).is_some() {
run.anchor = Some(entry.metadata.entry_id.0.clone());
}
}
}
}
LogEntry::AnnotatedAssistantItem { entry, .. }
| LogEntry::AnnotatedToolResult { entry, .. } => {
if let Some(run) = &mut active
&& super::project_item(&entry.metadata.entry_id.0, &entry.item).is_some()
{
run.anchor = Some(entry.metadata.entry_id.0.clone());
}
}
LogEntry::AnnotatedSystemItem { entry, .. } => {
if let Some(run) = &mut active {
run.anchor = Some(entry.metadata.entry_id.0.clone());
}
}
LogEntry::RunCompleted { ts, result, .. }
if !matches!(result, agen::EngineResult::Yielded) =>
{
finish(&active, *ts, &mut stats);
// Pause/resume retains the Invoke's cumulative accounting, just
// like the live TUI. Terminal outcomes cannot leak traffic into
// a later resumed segment whose Invoke is unavailable.
if !matches!(result, agen::EngineResult::Paused) {
active = None;
}
}
LogEntry::RunCancelled { ts, .. } | LogEntry::RunErrored { ts, .. } => {
finish(&active, *ts, &mut stats);
active = None;
}
// Compaction seeds repeat retained entries, not fresh traffic.
// A page traverses adopted lineage in journal order and retains the
// open Invoke across these checkpoints. A standalone new segment
// without that Invoke cannot claim complete accounting.
_ => {}
}
}
stats
}
fn finish(
active: &Option<Activity>,
ended_at: u64,
stats: &mut HashMap<String, SessionSnapshotEntry>,
) {
let Some(run) = active else { return };
let Some(anchor) = &run.anchor else { return };
// No usage records means unavailable measurement, not known zero traffic.
if run.requests == 0 {
return;
}
stats.insert(
anchor.clone(),
SessionSnapshotEntry {
entry_id: format!("run-stats:{anchor}"),
timestamp: ended_at,
provenance: SessionEntryProvenance::LegacyUnknown,
derived_from: vec![anchor.clone()],
data: SessionSnapshotEntryData::RunStats {
elapsed_ms: ended_at.saturating_sub(run.started_at),
requests: run.requests,
upload_tokens: run.upload,
output_tokens: run.output,
},
},
);
}
pub(super) fn append_to(
entries: &mut Vec<SessionSnapshotEntry>,
stats: &HashMap<String, SessionSnapshotEntry>,
) {
let mut with_stats = Vec::with_capacity(entries.len());
for entry in std::mem::take(entries) {
let summary = stats.get(&entry.entry_id);
with_stats.push(entry);
if let Some(summary) = summary {
with_stats.push(summary.clone());
}
}
*entries = with_stats;
}
crates/tui/src/app.rs
+8 -1
View File
@@ -1741,6 +1741,7 @@ impl App {
input_tokens,
output_tokens,
cache_read_input_tokens,
..
} => {
if let Some(input_tokens) = input_tokens {
self.session_context_tokens = input_tokens;
@@ -3000,7 +3001,10 @@ impl App {
self.apply_system_item(data);
}
}
SessionSnapshotEntryData::RunYielded { .. }
SessionSnapshotEntryData::Invoke { .. }
| SessionSnapshotEntryData::Usage { .. }
| SessionSnapshotEntryData::RunStats { .. }
| SessionSnapshotEntryData::RunYielded { .. }
| SessionSnapshotEntryData::RunResumed { .. }
| SessionSnapshotEntryData::RunCancelled => {
// Durable logical-Run transitions are intentionally not
@@ -3981,6 +3985,7 @@ mod completion_flow_tests {
});
app.handle_worker_event(Event::TurnStart { turn: 1 });
app.handle_worker_event(Event::Usage {
timestamp_ms: None,
input_tokens: Some(100),
output_tokens: Some(0),
cache_read_input_tokens: Some(40),
@@ -4210,6 +4215,7 @@ mod completion_flow_tests {
let mut app = App::new("test".into());
app.handle_worker_event(Event::TurnStart { turn: 1 });
app.handle_worker_event(Event::InvokeStart {
timestamp_ms: None,
kind: protocol::InvokeKind::UserSend,
});
app.handle_worker_event(Event::RunEnd {
@@ -5156,6 +5162,7 @@ mod completion_flow_tests {
let mut app = App::new("test".into());
app.handle_worker_event(Event::Usage {
timestamp_ms: None,
input_tokens: Some(42_000),
output_tokens: Some(9),
cache_read_input_tokens: Some(40_000),
crates/worker/examples/worker_protocol.rs
+1
View File
@@ -83,6 +83,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
input_tokens,
output_tokens,
cache_read_input_tokens,
..
} => {
println!(
"[usage] in={} (cache_read={}) out={}",
crates/worker/src/controller.rs
+6 -9
View File
@@ -178,7 +178,10 @@ impl WorkerHandle {
(entries, entry_rx, in_flight)
};
let mut session =
session_store::public_snapshot::project_current_session_snapshot(&entries);
session_store::public_snapshot::project_current_session_snapshot_with_accounting(
self.artifact_store.as_ref(),
&entries,
);
session.pending_submissions = self
.pending_activations
.lock()
@@ -1286,14 +1289,8 @@ pub(crate) fn wire_event_bridges_on_engine<C, St>(
});
});
let tx = working_event_tx.clone();
worker.on_usage(move |event| {
let _ = tx.send(Event::Usage {
input_tokens: event.input_tokens,
output_tokens: event.output_tokens,
cache_read_input_tokens: event.cache_read_input_tokens,
});
});
// Usage is published from committed LlmUsage records through the session-log
// lane. Snapshot and live usage therefore share one exact, gap-free boundary.
let tx = working_event_tx.clone();
worker.on_error(move |event| {
crates/worker/src/internal_worker.rs
+5 -1
View File
@@ -494,7 +494,11 @@ impl InternalWorkerSessionHandle {
(entries, snapshot_from_guard(&guard))
};
InternalWorkerSessionSnapshot {
session: session_store::public_snapshot::project_current_session_snapshot(&entries),
session:
session_store::public_snapshot::project_current_session_snapshot_with_accounting(
&self.store,
&entries,
),
greeting: self.greeting.as_ref().map(|greeting| {
greeting
.read()
crates/worker/src/ipc/protocol_session.rs
+16 -1
View File
@@ -96,7 +96,22 @@ pub fn live_log_entry_event(entry: LogEntry) -> Option<Event> {
item: value,
})
}
LogEntry::Invoke { trigger, .. } => Some(Event::InvokeStart { kind: trigger }),
LogEntry::Invoke { ts, trigger } => Some(Event::InvokeStart {
kind: trigger,
timestamp_ms: Some(ts),
}),
LogEntry::LlmUsage {
ts,
input_total_tokens,
cache_read_tokens,
output_tokens,
..
} => Some(Event::Usage {
timestamp_ms: Some(ts),
input_tokens: Some(input_total_tokens),
cache_read_input_tokens: Some(cache_read_tokens),
output_tokens: Some(output_tokens),
}),
other => {
// `SegmentLogSink::is_live_relevant` keeps non-live-relevant
// variants off the broadcast lane; reaching here means the two are
crates/worker/src/segment_log_sink.rs
+49
View File
@@ -124,6 +124,7 @@ impl SegmentLogSink {
| LogEntry::ToolResultCorrected { .. }
| LogEntry::AnnotatedSystemItem { .. }
| LogEntry::Invoke { .. }
| LogEntry::LlmUsage { .. }
| LogEntry::RunYielded { .. }
| LogEntry::RunResumed { .. }
| LogEntry::RunCancelled { .. }
@@ -357,6 +358,54 @@ mod tests {
assert!(rx.try_recv().is_err());
}
#[test]
fn persisted_usage_is_split_once_between_snapshot_and_live_even_with_equal_timestamps() {
let sink = SegmentLogSink::new();
sink.publish(LogEntry::Invoke {
ts: 100,
trigger: protocol::InvokeKind::UserSend,
});
let usage = LogEntry::LlmUsage {
ts: 200,
history_len: 2,
input_total_tokens: 2000,
cache_read_tokens: 1000,
cache_write_tokens: 0,
output_tokens: 100,
};
sink.publish(usage.clone());
let (snapshot, mut rx) = sink.subscribe_with_snapshot();
assert_eq!(
snapshot
.iter()
.filter(|e| matches!(e, LogEntry::LlmUsage { .. }))
.count(),
1
);
assert!(rx.try_recv().is_err());
sink.publish(usage);
let event =
crate::ipc::protocol_session::live_log_entry_event(rx.try_recv().unwrap()).unwrap();
assert!(matches!(
event,
protocol::Event::Usage {
timestamp_ms: Some(200),
input_tokens: Some(2000),
..
}
));
assert!(rx.try_recv().is_err());
let (reconnected, mut next_rx) = sink.subscribe_with_snapshot();
assert_eq!(
reconnected
.iter()
.filter(|e| matches!(e, LogEntry::LlmUsage { .. }))
.count(),
2
);
assert!(next_rx.try_recv().is_err());
}
#[test]
fn reset_with_initial_clears_and_broadcasts() {
let sink = SegmentLogSink::new();
crates/worker/src/session_capture.rs
+3
View File
@@ -294,6 +294,9 @@ impl SessionCapture {
// Observation deliberately excludes system items and
// controller errors from model-visible session evidence.
SessionSnapshotEntryData::SystemItem { .. }
| SessionSnapshotEntryData::Invoke { .. }
| SessionSnapshotEntryData::Usage { .. }
| SessionSnapshotEntryData::RunStats { .. }
| SessionSnapshotEntryData::RunYielded { .. }
| SessionSnapshotEntryData::RunResumed { .. }
| SessionSnapshotEntryData::RunCancelled
crates/worker/tests/controller_test.rs
+64
View File
@@ -1556,6 +1556,70 @@ async fn snapshot_includes_user_input_for_in_flight_turn() {
}
}
#[tokio::test]
async fn reconnect_snapshot_restores_live_run_accounting_from_engine_events() {
let client = MockClient::new(vec![
LlmEvent::tool_use_start(0, "hold-call", "Hold"),
LlmEvent::tool_input_delta(0, "{}"),
LlmEvent::tool_use_stop(0),
LlmEvent::Usage(UsageEvent {
input_tokens: Some(25_000),
output_tokens: Some(300),
cache_read_input_tokens: Some(20_000),
..Default::default()
}),
LlmEvent::Status(StatusEvent {
status: ResponseStatus::Completed,
}),
]);
let mut worker = make_worker(client).await;
worker
.engine_mut()
.register_tool(hanging_tool_definition("Hold"));
let handle = spawn_controller(worker).await;
let mut events = handle.sink.subscribe_with_snapshot().1;
handle
.send(Method::submit_text(
protocol::new_submission_request_id(),
"keep running",
))
.await
.unwrap();
let observed = tokio::time::timeout(std::time::Duration::from_secs(2), async {
loop {
if let LogEntry::LlmUsage {
ts,
input_total_tokens,
cache_read_tokens,
output_tokens,
..
} = events.recv().await.unwrap()
{
break (ts, input_total_tokens, cache_read_tokens, output_tokens);
}
}
})
.await
.expect("live accounting");
assert_eq!((observed.1, observed.2, observed.3), (25_000, 20_000, 300));
for _ in 0..2 {
let Event::Snapshot { session, .. } = handle.snapshot_event() else {
panic!("expected snapshot")
};
let usage: Vec<_> = session
.entries
.iter()
.filter(|e| matches!(e.data, protocol::SessionSnapshotEntryData::Usage { .. }))
.collect();
assert_eq!(usage.len(), 1);
assert_eq!(usage[0].timestamp, observed.0);
assert!(session.entries.iter().any(|e| matches!(
e.data,
protocol::SessionSnapshotEntryData::Invoke { .. }
) && e.timestamp > 0));
}
}
#[tokio::test]
async fn attach_snapshot_includes_current_status() {
let client = MockClient::sequential(vec![MockResponse::Hang(simple_text_events())]);