diff --git a/PLAN.md b/PLAN.md index 3b23e95..c95ef0a 100644 --- a/PLAN.md +++ b/PLAN.md @@ -181,8 +181,24 @@ it touches the transcript or the phone: not just the one that answered (added 2026-08-24, same reasoning as `UserMessage`). - `Status { state }` — idle / running / awaiting-input / compacting / exited. -- `UsageDelta { tokens }` — per-turn token counts where the dialect reports - them (both do). +- `UsageDelta { tokens, context }` — what a turn cost, and how much the + model was holding when it ended, where the dialect reports them (both do). + `context` is prompt plus both cache figures, taken from the **last + assistant message** rather than the turn's `result`: measured 2026-08-30 + against CLI 2.1.237, the result adds a turn's messages up, so its cache + read of 40,211 was the same conversation counted twice and no size the + model ever held. It is carried rather than summed by readers because it + goes *down* — a compaction replaces it with what the compaction reports, + and a clear leaves it unmeasured. `driver::context_after` is that rule, + and the phone folds with the same one (2026-08-30: this replaced a running + spend total, which could only climb and so kept reporting a context a + compaction or a clear had already taken away). + A session the server has no measurement of asks the CLI's own file + instead of waiting for a turn — `import::context_of`, the same three + fields the import list reads, in the background at load so a start never + waits on an ssh. A clear needs no special case: it gives the CLI a new + session id, so the lookup lands on a file with no usage in it and + answers "unknown", which is true. - `Error { message }`. Every event is appended to the session's transcript file with a sequence @@ -611,7 +627,9 @@ dev-updater (Kotlin 2.4.x, CMP 1.11.x, JDK 21). Screens: - Input bar: text, attach (camera/gallery/file), send — **always enabled**; mid-run sends become steering messages. - Top bar: model chip (tap to change), stop button while running, token - count, compact button (llama), overflow → delete. + count, compact button (llama), overflow → delete. (The count settled as + context held rather than tokens spent, and sits on the status row under + the transcript — see `UsageDelta` above.) 4. **Usage** — window bars for the 5-hour and weekly limits with reset times. 5. **Settings** — server address + token, hosts editor, llama model list editor. diff --git a/app/androidApp/src/main/kotlin/com/example/aiapp/Api.kt b/app/androidApp/src/main/kotlin/com/example/aiapp/Api.kt index ea99a32..b572694 100644 --- a/app/androidApp/src/main/kotlin/com/example/aiapp/Api.kt +++ b/app/androidApp/src/main/kotlin/com/example/aiapp/Api.kt @@ -150,9 +150,15 @@ data class SessionSummary( */ val notify: Boolean, /** - * Every token this session has spent, as the server counts it -- see `SessionEvent.UsageDelta`. + * How much context this session is holding, as the server last measured it -- see + * `SessionEvent.UsageDelta`. + * + * Null where nothing has been measured: a session that has not run a turn, a provider that does + * not report usage, or a clear nobody has run a turn since. That is not zero, and the status + * row says so in words rather than drawing an empty context for a conversation that may be + * nearly full. */ - val totalTokens: Long, + val contextTokens: Long?, /** * The longest edge an image should have when it reaches this session, or null where the * provider has no limit. @@ -178,7 +184,8 @@ private fun parseSession(session: JSONObject) = permissionMode = session.optString("permissionMode").ifEmpty { null }, imported = session.optBoolean("imported", false), notify = session.optBoolean("notify", true), - totalTokens = session.optLong("totalTokens", 0), + contextTokens = + if (session.has("contextTokens")) session.getLong("contextTokens") else null, maxImageEdge = session.optInt("maxImageEdge", 0).takeIf { it > 0 }, status = session.getString("status"), lastActivity = session.getDouble("lastActivity"), diff --git a/app/androidApp/src/main/kotlin/com/example/aiapp/Compaction.kt b/app/androidApp/src/main/kotlin/com/example/aiapp/Compaction.kt index d515496..7ed3e64 100644 --- a/app/androidApp/src/main/kotlin/com/example/aiapp/Compaction.kt +++ b/app/androidApp/src/main/kotlin/com/example/aiapp/Compaction.kt @@ -37,7 +37,14 @@ fun compactionSummary(item: TranscriptItem.CompactedNote): String { } } -private fun tokens(count: Long): String = "%,d".format(count) +/** + * A token count as a reader reads one. + * + * Shared with the status row rather than formatted at each: the divider and the row report the same + * quantity about the same moment, and one of them grouping its thousands while the other did not + * read as two different measurements. + */ +fun tokens(count: Long): String = "%,d".format(count) /** * What the working indicator says while a compaction is running. diff --git a/app/androidApp/src/main/kotlin/com/example/aiapp/Events.kt b/app/androidApp/src/main/kotlin/com/example/aiapp/Events.kt index 65d0c9c..3f612c5 100644 --- a/app/androidApp/src/main/kotlin/com/example/aiapp/Events.kt +++ b/app/androidApp/src/main/kotlin/com/example/aiapp/Events.kt @@ -113,14 +113,16 @@ sealed class SessionEvent { data class Settings(val model: String?, val permissionMode: String?) : SessionEvent() /** - * What a turn cost, and what the session has cost in total. + * What a turn cost, and how much the model was holding when it ended. * - * [total] is the server's running figure, carried on the event so a reader never adds up its - * own: a phone opens a session on the newest page of the transcript, so a sum it computed would - * be that page's share of the conversation wearing the whole conversation's label. Zero on - * entries recorded before the backend sent it. + * [context] is prompt plus both cache figures, measured by the backend from the turn's own + * usage. Carried on the event rather than summed by the reader, because it is not a sum: a + * conversation's context drops at a compaction and a clear, so adding turns up would report a + * figure the session stopped being true of. Null where the dialect did not say, and on entries + * recorded before the backend sent it -- which leaves the context unmeasured rather than + * unchanged. */ - data class UsageDelta(val tokens: Long, val total: Long) : SessionEvent() + data class UsageDelta(val tokens: Long, val context: Long?) : SessionEvent() /** * A compaction that finished, and how much context it recovered. @@ -235,7 +237,10 @@ fun parseSeqEvent(json: String): SeqEvent { permissionMode = body.optString("permissionMode").ifEmpty { null }, ) "usageDelta" -> - SessionEvent.UsageDelta(body.getLong("tokens"), body.optLong("total", 0)) + SessionEvent.UsageDelta( + body.getLong("tokens"), + if (body.has("context")) body.getLong("context") else null, + ) "compacted" -> SessionEvent.Compacted( preTokens = if (body.has("preTokens")) body.getLong("preTokens") else null, @@ -248,3 +253,29 @@ fun parseSeqEvent(json: String): SeqEvent { } return SeqEvent(seq = body.getLong("seq"), ts = body.getDouble("ts"), event = event) } + +/** + * The context after [event], given what it was before. + * + * The same rule the server folds with, because the screen has to keep up between page loads: the + * summary it opened with is a measurement from before this stream started, and every event that + * moves the figure arrives here. + * + * The two that lower it are the point. A clear takes the conversation away and a compaction + * replaces it with a summary, so a figure measured before either stopped being true at that moment + * -- and carrying it forward is how a session that had just been cleared went on reporting the + * context it no longer had. + * + * Null is "we don't know", which is a state each of them can reach: nothing measured yet, a + * compaction that finished without saying how much it recovered, or a clear nobody has run a turn + * since. + */ +fun contextAfter(current: Long?, event: SessionEvent): Long? = + when (event) { + // Falls back to what we had, so a turn the dialect reported no usage for is stale by a + // turn -- which every context figure is -- rather than unknown. + is SessionEvent.UsageDelta -> event.context ?: current + is SessionEvent.Compacted -> event.postTokens + is SessionEvent.Cleared -> null + else -> current + } diff --git a/app/androidApp/src/main/kotlin/com/example/aiapp/SessionScreen.kt b/app/androidApp/src/main/kotlin/com/example/aiapp/SessionScreen.kt index 4148fd5..41cce10 100644 --- a/app/androidApp/src/main/kotlin/com/example/aiapp/SessionScreen.kt +++ b/app/androidApp/src/main/kotlin/com/example/aiapp/SessionScreen.kt @@ -216,9 +216,20 @@ sealed class TranscriptItem { * Only ever consulted when the call is first folded in. That is what makes the name stable -- a run * keeps whatever it was called when it started, however many calls arrive at either end of it * afterwards. + * + * A question to the reader is in a run of its own, which is what puts it on the transcript as a row + * rather than inside a collapsed "Called 6 tools" card. Two things follow from being alone: it is + * always visible, since a run of one is drawn as itself rather than as a group; and the calls + * around it fall into a group before it and a group after it, so where the reader was asked + * something is legible in the shape of the transcript without opening anything. It ends the run + * before it as well as starting a fresh one after -- the moment somebody was asked is a boundary in + * the work, not a gap in the middle of one run. */ -private fun runIdFor(items: List, id: String): String = - (items.lastOrNull() as? TranscriptItem.ToolRun)?.runId ?: id +private fun runIdFor(items: List, id: String, tool: String): String { + val previous = items.lastOrNull() as? TranscriptItem.ToolRun ?: return id + if (tool == ASK_USER_QUESTION || previous.tool == ASK_USER_QUESTION) return id + return previous.runId +} /** * Puts a page of older items in front of the ones already loaded, healing whatever the page @@ -311,8 +322,15 @@ private fun adoptRun( earlier: List, later: List, ): List { - val joining = (later.firstOrNull() as? TranscriptItem.ToolRun)?.runId ?: return earlier - val tail = earlier.takeLastWhile { it is TranscriptItem.ToolRun } + val first = later.firstOrNull() as? TranscriptItem.ToolRun ?: return earlier + // A question is in a run of its own on both sides of the join, the same as it would be had + // the two pages been folded as one -- see `runIdFor`. Without this the heal would merge a + // group straight through the row the reader was asked something on. + if (first.tool == ASK_USER_QUESTION) return earlier + val joining = first.runId + val tail = earlier.takeLastWhile { + it is TranscriptItem.ToolRun && it.tool != ASK_USER_QUESTION + } if (tail.isEmpty()) return earlier return earlier.dropLast(tail.size) + tail.map { (it as TranscriptItem.ToolRun).copy(runId = joining) } @@ -338,7 +356,7 @@ fun foldEvent(items: List, entry: SeqEvent): List, entry: SeqEvent): List()) } var status by remember { mutableStateOf(summary.status) } - // Seeded from the row this screen was opened from, so a conversation that has spent - // something says so before any turn happens here. Zero used to mean "nothing yet" and - // "nothing in the page I loaded" at once, and the second one is most of the long sessions. - var totalTokens by remember(summary.id) { mutableLongStateOf(summary.totalTokens) } + // Seeded from the row this screen was opened from, so a conversation already under way says + // how much it is holding before any turn happens here. Null is "nobody has measured it", + // which is a different answer from an empty context and is drawn differently. + var contextTokens by remember(summary.id) { mutableStateOf(summary.contextTokens) } // When this screen saw the current compaction start, on this device's own clock, and how long // ago that is. See `compactingLabel`: null is the honest answer whenever the start was not // witnessed here, which is what opening a session that is already compacting looks like. @@ -612,14 +633,14 @@ fun SessionScreen( */ fun apply(entry: SeqEvent) { lastSeq.set(entry.seq) + // Before the rest, and for every event rather than only the usage ones: a compaction and + // a clear move this as much as a turn does, which is the whole reason it is a fold and + // not a running total. See `contextAfter`. + contextTokens = contextAfter(contextTokens, entry.event) when (val event = entry.event) { - // Taken, not accumulated: the server's running total is on the event, and adding - // up the deltas this screen happened to receive counted one page of a conversation - // and called it the whole. `max` because pages arrive in no guaranteed order and an - // older event's total is a smaller true answer, never a correction downwards; it - // also leaves the seeded figure alone for transcripts recorded before the backend - // sent a total at all. - is SessionEvent.UsageDelta -> totalTokens = maxOf(totalTokens, event.total) + // Nothing further: what it carries was folded into the context above, and what a + // turn cost is not something the transcript draws. + is SessionEvent.UsageDelta -> {} else -> { // What the session says it is set to now, which is the only thing that // says it: picking from either menu asks, and the answer comes back here. @@ -1230,7 +1251,11 @@ fun SessionScreen( ) } - SessionStatusRow(status = status, compactingFor = compactingFor, totalTokens = totalTokens) + SessionStatusRow( + status = status, + compactingFor = compactingFor, + contextTokens = contextTokens, + ) // Between the transcript and the box: above what is being typed, so the list does not // cover the thing the command is about, and below everything that explains it. @@ -1488,7 +1513,8 @@ private fun SessionStatusRow( status: String, /** Seconds since this device saw the compaction start; null if it did not see it. */ compactingFor: Long?, - totalTokens: Long, + /** Context the session is holding, or null where nothing has measured it. */ + contextTokens: Long?, modifier: Modifier = Modifier, ) { Row( @@ -1558,15 +1584,19 @@ private fun SessionStatusRow( modifier = Modifier.weight(1f), ) } - // Nothing rather than "0 tok" before anything has been spent: a total of zero is a fact - // about a conversation that has not started, and it is the one reading nobody needs. - if (totalTokens > 0) { - Text( - "$totalTokens tok", - style = MaterialTheme.typography.labelSmall, - color = MaterialTheme.colorScheme.onSurfaceVariant, - ) - } + // How full the session is, which is the number a reader is asking about -- how much room + // is left before the next compaction -- rather than what has been spent getting here. + // + // "unknown" in words, and always drawn. A context nobody has measured is not an empty + // one, and the two used to share an appearance: a session that had just been cleared, one + // whose provider never reports usage, and one that has not run a turn all showed nothing + // at all, which reads as a conversation with room to spare. It is the same reason the + // status word beside it names the quiet state instead of leaving the row blank. + Text( + contextTokens?.let { "context ${tokens(it)}" } ?: "context unknown", + style = MaterialTheme.typography.labelSmall, + color = MaterialTheme.colorScheme.onSurfaceVariant, + ) } } diff --git a/app/androidApp/src/main/kotlin/com/example/aiapp/ToolRows.kt b/app/androidApp/src/main/kotlin/com/example/aiapp/ToolRows.kt index d2c3570..4d53dd0 100644 --- a/app/androidApp/src/main/kotlin/com/example/aiapp/ToolRows.kt +++ b/app/androidApp/src/main/kotlin/com/example/aiapp/ToolRows.kt @@ -296,5 +296,10 @@ private fun PermissionAsk(ask: TranscriptItem.QuestionCard, onAnswer: (List, session_dir: PathBuf, } @@ -106,6 +122,7 @@ impl Translator { pending: HashMap::new(), asked: HashMap::new(), interrupting: false, + context: None, session_dir, } } @@ -221,8 +238,9 @@ impl Translator { .to_string(), }); } + let context = self.context.take(); if tokens > 0 { - events.push(Event::UsageDelta { tokens, total: 0 }); + events.push(Event::UsageDelta { tokens, context }); } events.push(Event::Status { state: SessionStatus::Idle, @@ -358,6 +376,14 @@ impl Translator { } fn translate_assistant(&mut self, message: &Value) -> Vec { + if let Some(usage) = message.get("usage") { + let field = |name: &str| usage.get(name).and_then(Value::as_u64).unwrap_or(0); + self.context = Some(context_tokens( + field("input_tokens"), + field("cache_creation_input_tokens"), + field("cache_read_input_tokens"), + )); + } let Some(content) = message.get("content").and_then(Value::as_array) else { return Vec::new(); }; @@ -1054,7 +1080,7 @@ mod tests { vec![ Event::UsageDelta { tokens: 182, - total: 0 + context: None }, Event::Status { state: SessionStatus::Idle @@ -1063,6 +1089,52 @@ mod tests { ); } + /// The context is the last assistant message's, not the result's. + /// + /// Real figures from a two-message haiku turn on 2.1.237, captured + /// 2026-08-30. The result adds the turn up -- its + /// `cache_read_input_tokens` of 40,211 is 14,259 and 25,952, the same + /// conversation counted twice -- so reading the context off it would + /// report a size the model never held, and by more the more tool calls + /// a turn makes. The last message's three input figures are what it + /// was holding when the turn ended. + #[test] + fn the_context_is_what_the_last_message_held_not_the_turn_added_up() { + let dir = tempfile::tempdir().expect("tempdir"); + let mut translator = Translator::new(dir.path().to_path_buf()); + let events = translate_lines( + &mut translator, + &[ + r#"{"type":"assistant","message":{"content":[],"usage":{"input_tokens":9,"cache_creation_input_tokens":11693,"cache_read_input_tokens":14259,"output_tokens":3}}}"#, + r#"{"type":"assistant","message":{"content":[],"usage":{"input_tokens":8,"cache_creation_input_tokens":171,"cache_read_input_tokens":25952,"output_tokens":2}}}"#, + r#"{"type":"result","subtype":"success","is_error":false,"session_id":"s","usage":{"input_tokens":17,"cache_creation_input_tokens":11864,"cache_read_input_tokens":40211,"output_tokens":156}}"#, + ], + ); + assert_eq!( + events.first(), + Some(&Event::UsageDelta { + tokens: 173, + context: Some(26_131), + }) + ); + + // Taken by that result, so a following turn whose messages carry no + // usage reports none rather than repeating this one's. + let events = translate_lines( + &mut translator, + &[ + r#"{"type":"result","subtype":"success","is_error":false,"session_id":"s","usage":{"input_tokens":4,"output_tokens":9}}"#, + ], + ); + assert_eq!( + events.first(), + Some(&Event::UsageDelta { + tokens: 13, + context: None, + }) + ); + } + #[test] fn a_compaction_reports_its_start_and_what_it_recovered() { // Real lines (trimmed) from a 2.1.237 session driven through diff --git a/server/src/session/driver.rs b/server/src/session/driver.rs index c388626..b195a7b 100644 --- a/server/src/session/driver.rs +++ b/server/src/session/driver.rs @@ -247,23 +247,25 @@ pub enum Event { }, /// Per-turn token counts, where the dialect reports them. UsageDelta { + /// What this turn cost: the tokens it was charged for. tokens: u64, - /// Every token this session has spent, this turn included. + /// What the model was holding when the turn ended -- see + /// [`context_tokens`] for what goes into it. /// - /// Filled in by the pump, not by drivers: a driver reports what its - /// turn cost, and only the pump sees all of them. Carried on the - /// event rather than left to be added up by whoever is reading, - /// because a reader has only *part* of the transcript -- a phone - /// opens a session on the newest page -- so a total it summed - /// itself would be the newest page's total wearing the whole - /// conversation's label. Worse when the page has no turn in it at - /// all: the count then reads zero, and a zero is drawn as nothing. + /// Carried on the event rather than summed by whoever is reading, + /// because it is not a sum: a conversation's context goes *down* + /// at a compaction and a clear, so adding turns up would report a + /// figure the session stopped being true of long ago. It is also + /// the number a reader is asking about -- how much room is left + /// before the next compaction -- rather than what has been spent + /// getting here. /// - /// Zero on entries written before this existed, which is why the - /// pump seeds its running total by adding up `tokens` at startup - /// rather than reading the last of these. - #[serde(default)] - total: u64, + /// `None` where the dialect did not say, which every reader has to + /// be able to draw: a turn whose usage the CLI omitted leaves the + /// context unmeasured rather than unchanged, and entries written + /// before this existed have no answer at all. + #[serde(default, skip_serializing_if = "Option::is_none")] + context: Option, }, /// A compaction that finished, and how much context it recovered. /// @@ -338,6 +340,51 @@ pub enum Event { }, } +/// How much the model was holding, from the three figures a turn reports. +/// +/// The input side only -- prompt plus both cache figures. A cached token +/// is cheaper but it is still one the model was given, so all three count; +/// output is left out because it is what the turn produced rather than +/// what continuing from here has to carry. +/// +/// One function so the definition cannot drift, because it is extracted in +/// two quite different ways: the live translators have the usage object +/// parsed, and `import::context_tokens` scans it out of a raw line without +/// parsing, since those files reach tens of megabytes. +pub fn context_tokens(input: u64, cache_creation: u64, cache_read: u64) -> u64 { + input + cache_creation + cache_read +} + +/// The context after `event`, given what it was before. +/// +/// The whole rule in one place, because three readers need the same +/// answer: the pump keeping a live session's figure, the transcript +/// seeding it at startup, and the phone folding the same events into what +/// it draws. Written here beside the events it reads so a fourth reader +/// finds it. +/// +/// The two that *lower* it are the point. A clear takes the conversation +/// away and a compaction replaces it with a summary, so a figure measured +/// before either stopped being true at that moment -- and carrying it +/// forward is how a session that had just been cleared went on reporting +/// the context it no longer had. +/// +/// `None` is "we don't know", which is a state each of them can reach: +/// nothing has been measured yet, a compaction finished without saying +/// how much it recovered, or a clear left a conversation nobody has +/// counted since. +pub fn context_after(current: Option, event: &Event) -> Option { + match event { + // `or`, so a turn the dialect reported no usage for leaves the last + // measurement standing: it is stale by a turn, which every context + // figure is, rather than wrong. + Event::UsageDelta { context, .. } => context.or(current), + Event::Compacted { post_tokens, .. } => *post_tokens, + Event::Cleared => None, + _ => current, + } +} + /// Something a session can be asked to do to itself. /// /// A closed set rather than a string, because the two that are not @@ -543,4 +590,76 @@ mod tests { }) ); } + + /// The two events that take the context *down* are the point of the + /// fold: a figure measured before a compaction or a clear stopped being + /// true at that moment, and carrying it forward is how a session that + /// had just been cleared went on reporting the context it no longer + /// had. + #[test] + fn a_compaction_and_a_clear_move_the_context_a_turn_cannot() { + let after = |current, event| context_after(current, &event); + + assert_eq!( + after( + Some(500), + Event::UsageDelta { + tokens: 12, + context: Some(30_100), + } + ), + Some(30_100) + ); + assert_eq!( + after( + Some(128_402), + Event::Compacted { + pre_tokens: Some(128_402), + post_tokens: Some(9_617), + trigger: Some("auto".to_string()), + } + ), + Some(9_617) + ); + assert_eq!(after(Some(9_617), Event::Cleared), None); + + // A compaction that did not say how much it recovered leaves the + // context unknown rather than stale: it definitely moved, and the + // one thing that is certainly wrong is the figure from before it. + assert_eq!( + after( + Some(128_402), + Event::Compacted { + pre_tokens: None, + post_tokens: None, + trigger: None, + } + ), + None + ); + + // A turn the dialect reported no context for is stale by a turn, + // which every context figure is, rather than unknown. + assert_eq!( + after( + Some(30_100), + Event::UsageDelta { + tokens: 12, + context: None, + } + ), + Some(30_100) + ); + + // Everything else leaves it alone. + assert_eq!( + after( + Some(30_100), + Event::Status { + state: SessionStatus::Idle, + } + ), + Some(30_100) + ); + } } diff --git a/server/src/session/echo.rs b/server/src/session/echo.rs index b97f1d7..ce23aa6 100644 --- a/server/src/session/echo.rs +++ b/server/src/session/echo.rs @@ -12,7 +12,10 @@ //! - `/question [text]` -- a question, exercising the answer path. //! - `/ask` -- an AskUserQuestion call: two questions on one tool call, //! with descriptions, a preview and a multi-select, which is the shape -//! that is awkward to get a real model to produce on demand. +//! that is awkward to get a real model to produce on demand. Wrapped in +//! a run of ordinary calls on each side, because being asked something +//! happens in the middle of work and the screen has to keep it out of +//! the collapsed group around it. //! - `/slow [seconds]` -- a turn that stays running (default 30), so states that only //! exist *while* something is happening can be looked at. //! - `/error [text]` -- a failure, which is otherwise awkward to cause. @@ -43,7 +46,7 @@ //! is the same every run. use std::path::{Path, PathBuf}; -use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; use std::sync::{Arc, Mutex}; use std::time::Duration; @@ -94,9 +97,35 @@ pub struct EchoDriver { /// way AskUserQuestion does, and the turn resumes when the last of /// them is answered rather than the first. pending_questions: Mutex>, + /// A pretend context, so the status row has something that behaves the + /// way a real one does: it grows with each turn, drops to what the + /// compaction says it recovered, and a clear leaves it unmeasured. The + /// numbers are invented like everything else here; what is real is + /// which way they move. + context: Arc, } impl EchoDriver { + /// A short run of ordinary calls, to sit either side of something. + /// + /// Three, because two is the fewest that groups and three makes it + /// obvious the group is a group -- and because the point of the + /// fixture is what a question looks like with work around it. + fn some_calls(&self, label: &str) { + for index in 0..3 { + let id = format!("echo-{label}-{index}-{}", super::random_hex()); + self.emit(Event::ToolStart { + id: id.clone(), + tool: "echo-tool".to_string(), + input: serde_json::json!({ "step": format!("{label} {index}") }), + }); + self.emit(Event::ToolEnd { + id, + output: format!("{label} step {index} finished"), + }); + } + } + /// An AskUserQuestion call, in the shape the CLI sends one. /// /// Two questions on one call, because that is where the display is @@ -169,6 +198,7 @@ impl EchoDriver { self.emit(Event::Status { state: SessionStatus::Running, }); + self.some_calls("before"); self.emit(Event::ToolStart { id: call.clone(), tool: "AskUserQuestion".to_string(), @@ -355,6 +385,7 @@ impl EchoDriver { .map(|rest| rest.trim().to_string()); let busy = Arc::clone(&self.busy); let queued = Arc::clone(&self.queued); + let context = Arc::clone(&self.context); let dir = self.session_dir.clone(); busy.store(true, Ordering::SeqCst); tokio::spawn(async move { @@ -480,9 +511,13 @@ impl EchoDriver { }); tokio::time::sleep(DELTA_DELAY).await; } + // A conversation gets bigger, so the pretend context does too: + // roughly a hundred tokens a turn plus the words themselves, + // which is enough to watch it climb between compactions. + let spent = text.split_whitespace().count() as u64; send(Event::UsageDelta { - tokens: text.split_whitespace().count() as u64, - total: 0, + tokens: spent, + context: Some(context.fetch_add(spent + 100, Ordering::SeqCst) + spent + 100), }); finish(); }); @@ -492,6 +527,7 @@ impl EchoDriver { let driver = Self { sink, pending_questions: Mutex::new(Vec::new()), + context: Arc::new(AtomicU64::new(0)), busy: Arc::new(AtomicBool::new(false)), queued: Arc::new(Mutex::new(Vec::new())), session_dir, @@ -702,6 +738,10 @@ impl Driver for EchoDriver { id: call, output: format!("answered: {answer}"), }); + // The work carries on where it left off, which is what makes + // the asked-here row a boundary with a group on each side + // rather than the last thing in the turn. + self.some_calls("after"); } else { self.emit(Event::AssistantText { delta: format!("You answered: {answer}"), @@ -750,12 +790,18 @@ impl Driver for EchoDriver { let sink = self.sink.clone(); let queued = Arc::clone(&self.queued); let busy = Arc::clone(&self.busy); + let context = Arc::clone(&self.context); busy.store(true, Ordering::SeqCst); tokio::spawn(async move { let _ = sink.send(Event::Status { state: SessionStatus::Compacting, }); tokio::time::sleep(COMPACT_TIME).await; + // What it says it recovered is what the pretend context becomes, + // so the figure on the status row and the one on the divider + // agree -- two numbers about the same moment disagreeing is the + // thing this rig exists to catch. + context.store(9_617, Ordering::SeqCst); let _ = sink.send(Event::Compacted { pre_tokens: Some(128_402), post_tokens: Some(9_617), @@ -770,6 +816,7 @@ impl Driver for EchoDriver { /// scroll behaviour and the transcript's shape can be exercised /// without spending a real session's context to produce one. fn clear(&self) { + self.context.store(0, Ordering::SeqCst); let _ = self.sink.send(Event::Cleared); } diff --git a/server/src/session/import.rs b/server/src/session/import.rs index e49bdc8..fb69684 100644 --- a/server/src/session/import.rs +++ b/server/src/session/import.rs @@ -21,7 +21,7 @@ use anyhow::{Context, Result}; use serde::Serialize; use serde_json::Value; -use super::driver::Event; +use super::driver::{self, Event}; use super::transport::{Launch, Transport}; /// How much of a transcript's tail is replayed into the phone's view. @@ -282,9 +282,9 @@ fn parse_row(line: &str) -> Option { /// The input tokens named in one `usage` object, added up. /// /// Prompt plus cache creation plus cache read: all three are context the -/// model was given, and a cached token is cheaper but not free. Output is -/// deliberately left out -- it is what the turn produced, not what -/// continuing from here has to carry. +/// model was given -- the definition is [`driver::context_tokens`]; this +/// is the same three figures dug out of a raw line rather than a parsed +/// one, because these files reach tens of megabytes. /// /// `None` for an empty blob, meaning no assistant turn has recorded usage. /// Missing individual fields count as zero, which is what an absent @@ -307,11 +307,11 @@ fn context_tokens(usage: &str) -> Option { }) .unwrap_or(0) }; - Some( - field("input_tokens") - + field("cache_creation_input_tokens") - + field("cache_read_input_tokens"), - ) + Some(driver::context_tokens( + field("input_tokens"), + field("cache_creation_input_tokens"), + field("cache_read_input_tokens"), + )) } /// The first line of what a person typed, short enough for a list row. @@ -688,6 +688,50 @@ pub async fn line_count(transport: &Transport, path: &str) -> Result { .with_context(|| format!("couldn't read a line count out of {out:?}")) } +/// What the CLI's own file says a session is holding, for a session this +/// server has no measurement of. +/// +/// A restarting server has been told nothing, and a session that has not +/// taken a turn since will not tell it -- so a conversation that is nearly +/// full reads as one nobody has counted until somebody sends a message to +/// it. The CLI records the figure on every assistant message, so it is +/// there to be read rather than waited for, and reading it is a +/// measurement rather than a guess: the same three fields, from the same +/// file, that the import list reports. +/// +/// A clear needs no special case here even though it makes the last usage +/// in a file stale. Clearing gives the CLI a *new* session id, which the +/// reader persists as the resume token, so this looks in a file that has +/// no usage in it yet and answers `None` -- which is the true answer. +/// +/// `None` for every way it cannot be read: no resume token, no file, a +/// machine that cannot be reached, or a file with no assistant turn in it. +/// Not knowing is a state the status row draws, so there is nothing to be +/// gained by inventing a number here. +pub async fn context_of(transport: &Transport, session_id: &str) -> Option { + // The id crosses as an argument rather than as script text: it comes + // from the CLI, but it reaches a shell on a machine that may not be + // this one, and the rule there is that data never becomes syntax. + let script = r#" +for f in "$HOME"/.claude/projects/*/"$1".jsonl; do + [ -f "$f" ] || continue + grep -o '"usage":{[^}]*' "$f" | tail -1 + exit 0 +done +"#; + let launch = Launch::new( + "sh", + vec![ + "-c".to_string(), + script.to_string(), + "sh".to_string(), + session_id.to_string(), + ], + None, + ); + context_tokens(&transport.capture(&launch).await.ok()?) +} + /// Events from the lines after `after`, which is a 0-based count of lines /// already accounted for. pub async fn replay_after( diff --git a/server/src/session/llama.rs b/server/src/session/llama.rs index ad10d91..f65b669 100644 --- a/server/src/session/llama.rs +++ b/server/src/session/llama.rs @@ -576,6 +576,10 @@ fn generate( let reader = std::io::BufReader::new(response.body_mut().as_reader()); let mut tokens = 0u64; + // The prompt side only, which is what the model is holding -- the same + // definition the other dialects report, so one word on the phone means + // one thing whichever kind of session it is. + let mut context = None; for line in std::io::BufRead::lines(reader) { if cancel.load(Ordering::Relaxed) { break; @@ -592,10 +596,19 @@ fn generate( let Ok(chunk) = serde_json::from_str::(payload) else { continue; }; - if let Some(usage) = chunk.get("usage").and_then(|u| u.get("total_tokens")) - && let Some(total) = usage.as_u64() - { - tokens = total; + if let Some(usage) = chunk.get("usage") { + if let Some(total) = usage + .get("total_tokens") + .and_then(serde_json::Value::as_u64) + { + tokens = total; + } + if let Some(prompt) = usage + .get("prompt_tokens") + .and_then(serde_json::Value::as_u64) + { + context = Some(prompt); + } } let delta = chunk .get("choices") @@ -611,7 +624,7 @@ fn generate( } } if tokens > 0 { - let _ = sink.send(Event::UsageDelta { tokens, total: 0 }); + let _ = sink.send(Event::UsageDelta { tokens, context }); } Ok(()) } @@ -721,7 +734,7 @@ mod tests { }, Event::UsageDelta { tokens: 12, - total: 12, + context: Some(12), }, ]); let messages = conversation(&path); diff --git a/server/src/session/mod.rs b/server/src/session/mod.rs index 20fa443..b53fc0d 100644 --- a/server/src/session/mod.rs +++ b/server/src/session/mod.rs @@ -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, - /// 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, /// 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, @@ -302,12 +309,12 @@ struct Shared { /// session was *launched* with, so reporting from it would show the /// mode a change had already replaced. permission_mode: Mutex>, - /// 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, + /// 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>, /// Whether this session's attention-wanting moments are announced. /// /// Mirrored out of the config so the pump can read it without taking @@ -425,7 +432,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, @@ -763,7 +770,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, @@ -1287,11 +1294,37 @@ 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), }); + // Nothing here has measured this session's context: the transcript + // predates the figure being recorded, or the last turn happened before + // this server was watching. The CLI wrote it down at the time, so ask + // its file rather than leaving the row saying "unknown" until somebody + // sends a message. In the background, because it is a file read on a + // machine that may be at the other end of an ssh connection, and a + // server start must not wait on one. + if provider.kind == DriverKind::ClaudeCli + && shared.context_tokens.lock().unwrap().is_none() + && let Some(session_id) = claude::read_resume_token(&dir) + { + let transport = Transport::for_setup(setup); + let shared = Arc::clone(&shared); + tokio::spawn(async move { + if let Some(context) = import::context_of(&transport, &session_id).await { + // Only if nothing else has answered in the meantime: a turn + // that finished while this was in flight measured the + // context after the one this read. + let mut held = shared.context_tokens.lock().unwrap(); + if held.is_none() { + *held = Some(context); + } + } + }); + } + // An imported session shares its transcript file with the CLI -- // `--resume` appends to the same one rather than forking, measured // rather than assumed -- so work done at a terminal belongs in this @@ -1420,20 +1453,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 @@ -2115,16 +2143,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"); @@ -2138,34 +2171,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 diff --git a/server/src/session/transcript.rs b/server/src/session/transcript.rs index 07b56d7..f2f841d 100644 --- a/server/src/session/transcript.rs +++ b/server/src/session/transcript.rs @@ -14,7 +14,7 @@ use std::path::Path; use anyhow::{Context, Result}; use serde::{Deserialize, Serialize}; -use super::driver::{Event, SessionStatus}; +use super::driver::{Event, SessionStatus, context_after}; /// One transcript line: an [`Event`] plus its position and time. The event /// is flattened so the wire shape stays one flat object. @@ -32,7 +32,7 @@ pub struct Transcript { next_seq: u64, last_status: Option, last_activity: Option, - total_tokens: u64, + context_tokens: Option, } impl Transcript { @@ -62,16 +62,12 @@ impl Transcript { next_seq: last_seq + 1, last_status, last_activity: existing.last().map(|entry| entry.ts), - // Added up rather than read off the newest entry: `total` is a - // later addition, so a transcript written before it has zero on - // every line while `tokens` was always there. - total_tokens: existing + // Folded rather than read off the newest usage entry: a clear + // or a compaction after it is what the answer is, and those + // events carry no usage of their own. + context_tokens: existing .iter() - .map(|entry| match entry.event { - Event::UsageDelta { tokens, .. } => tokens, - _ => 0, - }) - .sum(), + .fold(None, |current, entry| context_after(current, &entry.event)), }) } @@ -107,13 +103,16 @@ impl Transcript { self.last_activity } - /// Everything this session has spent, as of opening. + /// How much context the session was holding, as of opening. /// - /// Zero for an empty transcript, which is a session that has spent - /// nothing -- the one case where zero is the answer rather than the - /// absence of one. - pub fn total_tokens(&self) -> u64 { - self.total_tokens + /// `None` for a transcript nothing has been measured in -- a new + /// session, one whose dialect never reported usage, or one whose last + /// word on the subject was a clear. That is not zero, and it is why + /// this is an option: 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 { + self.context_tokens } /// Appends `event`, assigning it the next sequence number. Flushed per @@ -395,7 +394,7 @@ mod tests { }, Event::UsageDelta { tokens: 42, - total: 0, + context: Some(42), }, Event::Error { message: "boom".into(),