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

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<ConversationItem>,
@@ -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(),

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);
 

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
         }
     }
 }

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();