Report the context a session holds, not what it has spent

The number on the status row was a running total of tokens spent, so it
could only ever climb: a session compacted from 128k down to 10k, or
cleared outright, went on reporting the larger figure, and disagreed with
the divider directly above it saying what the compaction had recovered.

It now reports what the model is holding -- prompt plus both cache
figures -- folded through `driver::context_after`, which is the one rule
the pump, the transcript and the phone all use: a turn sets it, a
compaction replaces it with what the compaction measured, and a clear
leaves it unmeasured. Unmeasured says so in words, because an empty
context and one nobody has counted used to look identical.

Taken from the turn's last assistant message rather than its `result`:
measured against CLI 2.1.237, a two-message turn reported a cache read of
40,211, being 14,259 and 25,952 -- the same conversation counted twice,
and no size the model ever held.
This commit is contained in:
iris committed 2026-08-30 01:53:43 -04:00
1 parent 81c8a57181
commit 5e11b9da80
12 files changed
+442 -130

No files matched your search

+65 -41
View File
@@ -31,7 +31,7 @@ use crate::config::{
Config, DriverKind, ProviderConfig, SessionConfig, SetupConfig, SshConfig, TokenEntry,
};
use claude::ClaudeDriver;
use driver::{Driver, Event, EventSink, ImageRef, SessionCommand, SessionStatus};
use driver::{Driver, Event, EventSink, ImageRef, SessionCommand, SessionStatus, context_after};
use echo::EchoDriver;
use llama::LlamaDriver;
use transcript::{SeqEvent, Transcript};
@@ -143,9 +143,16 @@ pub struct SessionInfo {
pub imported: bool,
#[serde(skip_serializing_if = "Option::is_none")]
pub cwd: Option<PathBuf>,
/// Every token this session has spent, so a phone showing a total does
/// not have to add up a transcript it only holds part of.
pub total_tokens: u64,
/// How much context this session is holding, so a phone does not have
/// to fold a transcript it only holds part of.
///
/// Absent rather than zero where nothing has been measured -- a
/// session that has not run a turn, a dialect that does not report
/// usage, or a clear nobody has run a turn since. "Empty" and "we did
/// not find out" are different answers and the phone draws them
/// differently.
#[serde(skip_serializing_if = "Option::is_none")]
pub context_tokens: Option<u64>,
/// The longest edge an image should have by the time it gets here, or
/// absent where this provider has no limit -- see
/// [`DriverKind::max_image_edge`]. Absent rather than a large number,
@@ -280,12 +287,12 @@ struct Shared {
/// session was *launched* with, so reporting from it would show the
/// mode a change had already replaced.
permission_mode: Mutex<Option<String>>,
/// Every token this session has spent.
/// How much context this session is holding.
///
/// Kept here because only the pump sees every turn, and reported on the
/// session row so a phone opening a long conversation has the real
/// figure rather than the newest page's share of it.
total_tokens: Mutex<u64>,
/// Kept here because only the pump sees every event, and reported on
/// the session row so a phone opening a long conversation has the real
/// figure rather than whatever its newest page happens to mention.
context_tokens: Mutex<Option<u64>>,
/// Whether this session's attention-wanting moments are announced.
///
/// Mirrored out of the config so the pump can read it without taking
@@ -403,7 +410,7 @@ impl LiveSession {
title: self.shared.title.lock().unwrap().clone(),
model: self.shared.model.lock().unwrap().clone(),
permission_mode: self.shared.permission_mode.lock().unwrap().clone(),
total_tokens: *self.shared.total_tokens.lock().unwrap(),
context_tokens: *self.shared.context_tokens.lock().unwrap(),
notify: *self.shared.notify.lock().unwrap(),
max_image_edge: kind.and_then(DriverKind::max_image_edge),
imported,
@@ -741,7 +748,7 @@ impl SessionManager {
title: meta.title.clone(),
model: meta.model.clone(),
permission_mode: meta.permission_mode.clone(),
total_tokens: 0,
context_tokens: None,
max_image_edge: kind_of(&inner.config, &meta.setup, &meta.provider)
.and_then(DriverKind::max_image_edge),
notify: meta.notify,
@@ -1265,7 +1272,7 @@ fn launch(
last_activity: Mutex::new(transcript.last_activity().unwrap_or_else(now)),
model: Mutex::new(meta.model.clone()),
permission_mode: Mutex::new(meta.permission_mode.clone()),
total_tokens: Mutex::new(transcript.total_tokens()),
context_tokens: Mutex::new(transcript.context_tokens()),
notify: Mutex::new(meta.notify),
written: Mutex::new(0),
});
@@ -1398,20 +1405,15 @@ async fn pump(
// for where a user's message sits: where the session read it.
let event = match event {
Event::MessageTaken { id, text, images } => Event::UserMessage { id, text, images },
// The running total is the pump's to keep, for the reason the
// field gives: a driver knows what its own turn cost and
// nothing else does. Added here rather than at each driver so
// a new one cannot get it wrong by leaving it out.
Event::UsageDelta { tokens, .. } => {
let mut total = shared.total_tokens.lock().unwrap();
*total += tokens;
Event::UsageDelta {
tokens,
total: *total,
}
}
other => other,
};
// Where the session row's figure comes from. Kept here rather than
// at each driver because a clear and a compaction move it as much
// as a turn does, and only the pump sees all three.
{
let mut context = shared.context_tokens.lock().unwrap();
*context = context_after(*context, &event);
}
// Nothing changed, so there is nothing to record. Both of these
// repeat: an imported session reads the turn state off its file's
// newest record on every sync and mostly finds the answer it found
@@ -2038,16 +2040,21 @@ mod tests {
);
}
/// The total covers the whole conversation, not the part a reader holds.
/// The context figure follows the conversation down as well as up.
///
/// The bug this fixes was invisible in exactly the way that matters: a
/// phone opens a session on its newest page and used to add up the
/// `UsageDelta`s it found there, so a long conversation reported its
/// last few turns as the total -- and a page with no turn in it at all
/// reported nothing, since zero is drawn as blank. Both readings looked
/// like an answer.
/// It used to be a running total of what the session had spent, which
/// only ever climbs -- so a session that had just been compacted from
/// 128k to 10k, or cleared outright, went on reporting the larger
/// figure, and the number on the status row disagreed with the divider
/// directly above it. Turns raise it, a compaction replaces it with
/// what the compaction says it recovered, and a clear leaves it
/// unmeasured rather than guessing a small number.
///
/// The compaction leg is in `driver::tests` rather than here: echo
/// spends thirteen seconds on one so a person can watch the state, and
/// the rule both paths use is the same function.
#[tokio::test]
async fn the_token_total_covers_the_whole_conversation() {
async fn the_context_figure_follows_compactions_and_clears() {
let dir = tempfile::tempdir().expect("tempdir");
let config_path = dir.path().join("config.ron");
let data_dir = dir.path().join("sessions");
@@ -2061,34 +2068,51 @@ mod tests {
let info = manager.spawn_session(echo_spec()).expect("spawn");
let session = manager.session(&info.id).expect("live");
// Echo charges a token a word, so the arithmetic is checkable.
// Nothing measured yet, which is not the same as an empty context
// and is not reported as one.
assert_eq!(manager.sessions()[0].context_tokens, None);
// Echo's pretend context is a hundred a turn plus the words, so the
// arithmetic is checkable.
let mut rx = session.subscribe();
session.send_message("one two three".to_string(), Vec::new());
collect_turn(&mut rx).await;
session.send_message("four five".to_string(), Vec::new());
collect_turn(&mut rx).await;
let running = manager.sessions()[0].total_tokens;
assert_eq!(running, 5, "two turns of three and two words");
assert_eq!(
manager.sessions()[0].context_tokens,
Some(205),
"two turns of three and two words"
);
// The event carries it too, so a phone never has to add up its own.
let last_total = transcript::read_after(session.transcript_path(), 0)
// The event carries it too, so a phone never has to fold the part of
// the transcript it happens to hold.
let last_context = transcript::read_after(session.transcript_path(), 0)
.expect("transcript")
.iter()
.rev()
.find_map(|entry| match entry.event {
Event::UsageDelta { total, .. } => Some(total),
Event::UsageDelta { context, .. } => context,
_ => None,
})
.expect("a usage event");
assert_eq!(last_total, running);
assert_eq!(last_context, 205);
// And a restart picks it up from the file rather than starting over.
// A clear leaves it unmeasured: the conversation is gone, and how
// much is left is a thing nobody has counted.
session.run_command(SessionCommand::Clear);
collect_until(&mut rx, |event| matches!(event, Event::Cleared)).await;
assert_eq!(manager.sessions()[0].context_tokens, None);
// And a restart folds it back out of the file rather than starting
// over -- including the clear, which is why it is not the last
// usage event that decides.
drop(rx);
drop(session);
drop(manager);
let manager = SessionManager::new(config_path, data_dir.clone(), data_dir.join("models"))
.expect("manager restart");
assert_eq!(manager.sessions()[0].total_tokens, running);
assert_eq!(manager.sessions()[0].context_tokens, None);
}
/// Backdates every line in a transcript, so a restart has something to