From d022c24f68405c0251a055dc12ef112ba37a36fe Mon Sep 17 00:00:00 2001
From: Felipe Coury <felipe.coury@openai.com>
Date: Fri, 22 May 2026 21:46:12 -0300
Subject: [PATCH] fix(next-prompt): keep hidden usage live-only

---
 codex-rs/core/src/next_prompt_suggestion.rs   | 29 ++++++-------------
 .../tests/suite/next_prompt_suggestion.rs     | 12 ++------
 2 files changed, 11 insertions(+), 30 deletions(-)

diff --git a/codex-rs/core/src/next_prompt_suggestion.rs b/codex-rs/core/src/next_prompt_suggestion.rs
index c615d92e43..1edcdb653d 100644
--- a/codex-rs/core/src/next_prompt_suggestion.rs
+++ b/codex-rs/core/src/next_prompt_suggestion.rs
@@ -28,7 +28,6 @@ use codex_protocol::openai_models::ReasoningEffort;
 use codex_protocol::protocol::Event;
 use codex_protocol::protocol::EventMsg;
 use codex_protocol::protocol::TokenCountEvent;
-use codex_protocol::protocol::TokenUsageInfo;
 use codex_rollout_trace::InferenceTraceContext;
 use futures::StreamExt;
 use std::collections::HashSet;
@@ -177,22 +176,11 @@ pub(crate) async fn suggest_next_prompt(
                 latest_rate_limits = Some(snapshot.clone());
                 sess.record_rate_limits_info(snapshot).await;
             }
-            ResponseEvent::Completed {
-                response_id,
-                token_usage,
-                ..
-            } => {
-                let token_usage_info = TokenUsageInfo::new_or_append(
-                    &sess.token_usage_info().await,
-                    &token_usage,
-                    turn_context.model_context_window(),
-                );
-                let should_emit_token_count =
-                    token_usage_info.is_some() || latest_rate_limits.is_some();
-                if should_emit_token_count {
-                    // App-server clients key usage updates by turn id. Attribute hidden
-                    // suggestion sampling to the latest real turn instead of the ephemeral
-                    // lightweight turn used to configure the request.
+            ResponseEvent::Completed { response_id, .. } => {
+                if let Some(rate_limits) = latest_rate_limits {
+                    // App-server clients key updates by turn id. Attribute hidden
+                    // suggestion rate limits to the latest real turn instead of the
+                    // ephemeral lightweight turn used to configure the request.
                     if let Some(turn_id) = sess
                         .reference_context_item()
                         .await
@@ -201,12 +189,12 @@ pub(crate) async fn suggest_next_prompt(
                         sess.deliver_event_raw(Event {
                             id: turn_id,
                             msg: EventMsg::TokenCount(TokenCountEvent {
-                                info: token_usage_info,
-                                rate_limits: latest_rate_limits,
+                                info: None,
+                                rate_limits: Some(rate_limits),
                             }),
                         })
                         .await;
-                    } else if let Some(rate_limits) = latest_rate_limits {
+                    } else {
                         sess.send_event(
                             &turn_context,
                             EventMsg::TokenCount(TokenCountEvent {
@@ -227,6 +215,7 @@ pub(crate) async fn suggest_next_prompt(
             .send_response_processed(&completed_response_id)
             .await;
     }
+    client_session.reset_websocket_session();
     if !session_history_matches_snapshot(sess, history_snapshot).await {
         tracing::debug!("next prompt suggestion skipped after history changed");
         return Ok(None);
diff --git a/codex-rs/core/tests/suite/next_prompt_suggestion.rs b/codex-rs/core/tests/suite/next_prompt_suggestion.rs
index 5196e90821..796e714177 100644
--- a/codex-rs/core/tests/suite/next_prompt_suggestion.rs
+++ b/codex-rs/core/tests/suite/next_prompt_suggestion.rs
@@ -72,21 +72,13 @@ async fn suggest_next_prompt_samples_history_without_tools() -> Result<()> {
 
     let token_event = wait_for_event(&test.codex, |msg| {
         matches!(msg, EventMsg::TokenCount(ev)
-            if ev.info.as_ref().is_some_and(|info| info.last_token_usage.total_tokens == 33)
-                && ev.rate_limits.is_some())
+            if ev.info.is_none() && ev.rate_limits.is_some())
     })
     .await;
     let EventMsg::TokenCount(token_count) = token_event else {
         unreachable!("wait_for_event predicate only accepts TokenCount");
     };
-    assert_eq!(
-        token_count
-            .info
-            .expect("token usage should be recorded")
-            .last_token_usage
-            .total_tokens,
-        33
-    );
+    assert_eq!(token_count.info, None);
     assert_eq!(
         token_count
             .rate_limits