From fc7113ba3caf2c5e21ee86ad842ad3c6e0566e42 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=9B=B7=E7=94=B5=E8=8A=BD=E8=A1=A3?= Date: Sun, 13 Sep 2026 07:24:07 -0400 Subject: [PATCH] fix(chat-state): reasoning estimate and calibration reseed on every rewrite A Reasoning item counted its text and its encrypted blob as two separate payloads, inflating the estimate for retained reasoning; the estimate now takes the larger of the two (base64 runs 4/3 over the raw bytes). Only compaction carried the provider-confirmed ratio into the reseeded total; a rewind or a non-compaction replace sprang the context meter back to the raw estimate. reseed_total_tokens applies the ratio (capped at the pre-rewrite total) to every rewrite, and the truncate path uses it too. Ported from upstream 1.0.1x (grok-build xai-chat-state actor/state.rs, mutations.rs, queries.rs). Claude-Session: https://claude.ai/code/session_01RUqV8HVGXcAETyYbZ2tPQ4 --- .../kigi-chat-state/src/actor/mutations.rs | 42 +++++++-------- .../kigi-chat-state/src/actor/queries.rs | 6 +-- .../kigi-chat-state/src/actor/state.rs | 6 +-- .../kigi-chat-state/src/actor/tests.rs | 52 +++++++++++++++++-- 4 files changed, 73 insertions(+), 33 deletions(-) diff --git a/crates/codegen/kigi-chat-state/src/actor/mutations.rs b/crates/codegen/kigi-chat-state/src/actor/mutations.rs index 8c422da..1935e1c 100644 --- a/crates/codegen/kigi-chat-state/src/actor/mutations.rs +++ b/crates/codegen/kigi-chat-state/src/actor/mutations.rs @@ -390,14 +390,24 @@ impl ChatStateActor { }); } - /// Replace the entire conversation, persist, re-estimate `total_tokens`, - /// and emit reset + token-update events. - /// - /// Compaction replaces carry the provider-side overhead forward as a - /// *ratio* (`base_estimate × provider_total ÷ estimate_at_last_response`, - /// capped at the pre-compaction total; `base_estimate` when that estimate is - /// 0) so the reseed neither springs back nor over-counts (see - /// `COMPACTION.md`). + /// Reseed after a rewrite: the raw estimate scaled by the provider-confirmed ratio, capped at the pre-rewrite total (see COMPACTION.md). + pub(super) fn reseed_total_tokens(&self, base_estimate: u64) -> u64 { + let pre_replace_total = self.state.total_tokens; + let mut estimated_tokens = + if pre_replace_total > 0 && self.state.estimate_at_last_response > 0 { + let ratio = pre_replace_total as f64 / self.state.estimate_at_last_response as f64; + (base_estimate as f64 * ratio).round() as u64 + } else { + base_estimate + }; + // Over-counting fires premature compaction; an under-count self-heals on the next response. + if pre_replace_total > 0 { + estimated_tokens = estimated_tokens.min(pre_replace_total); + } + estimated_tokens + } + + /// Replace the conversation, persist, reseed `total_tokens` and emit reset + token events; `is_compaction` only marks the turn capture. pub(super) fn replace_conversation( &mut self, items: Vec, @@ -407,28 +417,16 @@ impl ChatStateActor { if is_compaction && let Some(cap) = &mut self.state.turn_capture { cap.compaction_occurred = true; } - let pre_replace_total = self.state.total_tokens; // `harness_trace_buffer` / `harness_trace_turns` intentionally untouched: // the planner/verifier subagents ran, so their sealed trace turns survive // a conversation replace (same intent as the `TruncateToPromptIndex` arm). self.persistence.replace_history(&items); let base_estimate = super::state::estimate_conversation_tokens(&items); - let mut estimated_tokens = - if is_compaction && pre_replace_total > 0 && self.state.estimate_at_last_response > 0 { - let ratio = pre_replace_total as f64 / self.state.estimate_at_last_response as f64; - (base_estimate as f64 * ratio).round() as u64 - } else { - base_estimate - }; - // Compaction must never appear to increase usage. - if is_compaction && pre_replace_total > 0 { - estimated_tokens = estimated_tokens.min(pre_replace_total); - } + let estimated_tokens = self.reseed_total_tokens(base_estimate); self.state.conversation = items; self.state.estimated_tokens_since_model = 0; self.state.total_tokens = estimated_tokens; - self.state.estimate_at_last_response = - super::state::estimate_conversation_tokens(&self.state.conversation); + self.state.estimate_at_last_response = base_estimate; self.rebase_turn_capture_offset(); self.send_event(ChatStateEvent::ConversationReset { new_len: self.state.conversation.len(), diff --git a/crates/codegen/kigi-chat-state/src/actor/queries.rs b/crates/codegen/kigi-chat-state/src/actor/queries.rs index 8b3c69a..448d2d4 100644 --- a/crates/codegen/kigi-chat-state/src/actor/queries.rs +++ b/crates/codegen/kigi-chat-state/src/actor/queries.rs @@ -68,10 +68,10 @@ impl ChatStateActor { self.state.conversation.truncate(truncate_at); self.state.prompt_texts.truncate(target_prompt_index); self.state.prompt_index = target_prompt_index; - self.state.total_tokens = - super::state::estimate_conversation_tokens(&self.state.conversation); + let base_estimate = super::state::estimate_conversation_tokens(&self.state.conversation); + self.state.total_tokens = self.reseed_total_tokens(base_estimate); self.state.estimated_tokens_since_model = 0; - self.state.estimate_at_last_response = self.state.total_tokens; + self.state.estimate_at_last_response = base_estimate; self.persistence.replace_history(&self.state.conversation); diff --git a/crates/codegen/kigi-chat-state/src/actor/state.rs b/crates/codegen/kigi-chat-state/src/actor/state.rs index 54067c3..86a7a21 100644 --- a/crates/codegen/kigi-chat-state/src/actor/state.rs +++ b/crates/codegen/kigi-chat-state/src/actor/state.rs @@ -68,12 +68,10 @@ pub fn estimate_item_tokens(item: &ConversationItem) -> u64 { kigi_token_estimation::estimate_tokens(&b.text_summary()) } ConversationItem::Reasoning(r) => { - // Summary + content text follow the standard bytes-per-token - // estimate; encrypted blobs are base64 and don't survive - // tokenization 1:1, so estimate at len/4 as well. + // Text and ciphertext are the same reasoning twice: the larger wins (base64 is 4/3 over raw). let text_bytes = kigi_sampling_types::reasoning_item_text(r).len(); let enc_bytes = r.encrypted_content.as_deref().map(str::len).unwrap_or(0); - ((text_bytes + enc_bytes) as u64) / kigi_token_estimation::BYTES_PER_TOKEN + (text_bytes.max(enc_bytes * 3 / 4) as u64) / kigi_token_estimation::BYTES_PER_TOKEN } } } diff --git a/crates/codegen/kigi-chat-state/src/actor/tests.rs b/crates/codegen/kigi-chat-state/src/actor/tests.rs index 4463ff2..762628c 100644 --- a/crates/codegen/kigi-chat-state/src/actor/tests.rs +++ b/crates/codegen/kigi-chat-state/src/actor/tests.rs @@ -658,22 +658,66 @@ async fn compaction_reseed_without_provider_count_matches_plain_estimate() { } #[tokio::test] -async fn non_compaction_replace_does_not_carry_overhead() { +async fn non_compaction_replace_carries_confirmed_total() { + // A rewind/mode-switch replace scales from the provider-confirmed count, not the raw estimate. let h = TestHarness::new(); h.handle .push_user_message(ConversationItem::user("x".repeat(4000))); - h.handle.record_token_usage(51_000); + h.handle.record_token_usage(500); + // Estimate at last response = 1_000, confirmed = 500 → ratio 0.5. h.handle .replace_conversation(vec![ConversationItem::user("q".repeat(4000))]); let total = h.handle.get_total_tokens().await; assert_eq!( - total, 1_000, - "non-compaction replace (e.g. rewind) keeps the plain estimate" + total, 500, + "same-size replace carries the provider-confirmed count" ); } +#[tokio::test] +async fn replace_never_increases_total_tokens() { + // A growing replace is capped at the confirmed total; the under-count self-heals on the next usage. + let h = TestHarness::new(); + h.handle + .push_user_message(ConversationItem::user("x".repeat(4000))); + h.handle.record_token_usage(1_500); + + h.handle + .replace_conversation(vec![ConversationItem::user("q".repeat(40_000))]); + + assert_eq!(h.handle.get_total_tokens().await, 1_500); +} + +#[tokio::test] +async fn truncate_scales_from_confirmed_total() { + // Rewind path: TruncateToPromptIndex must use the same carry as replace_conversation. + let h = TestHarness::new(); + h.handle + .push_user_message(ConversationItem::user("x".repeat(4000))); + h.handle.increment_prompt_index(); + h.handle + .push_user_message(ConversationItem::user("y".repeat(4000))); + h.handle.increment_prompt_index(); + h.handle.record_token_usage(1_000); + // Estimate at last response = 2_000, confirmed = 1_000 → ratio 0.5. + + h.handle.truncate_to_prompt_index(1).await; + + // Keeps the first user item (raw estimate 1_000) → scaled to 500. + assert_eq!(h.handle.get_total_tokens().await, 500); +} + +#[test] +fn reasoning_estimate_takes_the_larger_of_text_and_ciphertext() { + // "thinking for r1" is 15 bytes; 800 base64 bytes are ~600 raw bytes. + let encrypted = reasoning_sibling("r1", Some(&"A".repeat(800))); + assert_eq!(crate::estimate_conversation_tokens(&[encrypted]), 150); + let plain = reasoning_sibling("r1", None); + assert_eq!(crate::estimate_conversation_tokens(&[plain]), 3); +} + #[tokio::test] async fn flush_calls_persistence_flush() { let mut h = TestHarness::new();