fix(core): isolate next-prompt suggestion sampling

codex-rs/core/src/client.rs

@@ -238,6 +238,7 @@ pub struct ModelClient {
 pub struct ModelClientSession {
     client: ModelClient,
     websocket_session: WebsocketSession,
+    websocket_session_cache_mode: WebsocketSessionCacheMode,
     /// Turn state for sticky routing.
     ///
     /// This is an `OnceLock` that stores the turn state value received from the server
@@ -251,6 +252,12 @@ pub struct ModelClientSession {
     turn_state: Arc<OnceLock<String>>,
 }
 
+#[derive(Clone, Copy)]
+enum WebsocketSessionCacheMode {
+    Shared,
+    Isolated,
+}
+
 #[derive(Debug, Clone)]
 struct LastResponse {
     response_id: String,
@@ -382,6 +389,16 @@ impl ModelClient {
         ModelClientSession {
             client: self.clone(),
             websocket_session: self.take_cached_websocket_session(),
+            websocket_session_cache_mode: WebsocketSessionCacheMode::Shared,
+            turn_state: Arc::new(OnceLock::new()),
+        }
+    }
+
+    pub(crate) fn new_isolated_session(&self) -> ModelClientSession {
+        ModelClientSession {
+            client: self.clone(),
+            websocket_session: WebsocketSession::default(),
+            websocket_session_cache_mode: WebsocketSessionCacheMode::Isolated,
             turn_state: Arc::new(OnceLock::new()),
         }
     }
@@ -951,6 +968,12 @@ impl ModelClient {
 
 impl Drop for ModelClientSession {
     fn drop(&mut self) {
+        if matches!(
+            self.websocket_session_cache_mode,
+            WebsocketSessionCacheMode::Isolated
+        ) {
+            return;
+        }
         let websocket_session = std::mem::take(&mut self.websocket_session);
         self.client
             .store_cached_websocket_session(websocket_session);

codex-rs/core/src/next_prompt_suggestion.rs

@@ -21,6 +21,7 @@ use crate::session::session::Session;
 use codex_async_utils::OrCancelExt;
 use codex_features::Feature;
 use codex_protocol::models::ContentItem;
+use codex_protocol::models::MessagePhase;
 use codex_protocol::models::ResponseItem;
 use codex_protocol::openai_models::ModelPreset;
 use codex_protocol::openai_models::ReasoningEffort;
@@ -84,7 +85,7 @@ pub(crate) async fn suggest_next_prompt(
         tracing::debug!("next prompt suggestion skipped before assistant boundary");
         return None;
     }
-    if assistant_message_count(&prompt_input) < 2 {
+    if final_assistant_message_count(&prompt_input) < 2 {
         return None;
     }
     prompt_input.push(ContextualUserFragment::into(
@@ -103,7 +104,7 @@ pub(crate) async fn suggest_next_prompt(
     if !session_is_idle_for_suggestion(sess).await {
         return None;
     }
-    let mut client_session = sess.services.model_client.new_session();
+    let mut client_session = sess.services.model_client.new_isolated_session();
     let mut stream = match client_session
         .stream(
             &prompt,
@@ -201,15 +202,6 @@ pub(crate) async fn suggest_next_prompt(
                             }),
                         })
                         .await;
-                    } else {
-                        sess.send_event(
-                            &turn_context,
-                            EventMsg::TokenCount(TokenCountEvent {
-                                info: None,
-                                rate_limits: Some(rate_limits),
-                            }),
-                        )
-                        .await;
                     }
                 }
                 break response_id;
@@ -300,10 +292,16 @@ fn suggestion_prompt_has_headroom(estimated_token_count: i64, model_context_wind
         < model_context_window.saturating_sub(NEXT_PROMPT_SUGGESTION_TOKEN_HEADROOM)
 }
 
-fn assistant_message_count(items: &[ResponseItem]) -> usize {
+fn final_assistant_message_count(items: &[ResponseItem]) -> usize {
     items
         .iter()
-        .filter(|item| matches!(item, ResponseItem::Message { role, .. } if role == "assistant"))
+        .filter(|item| {
+            matches!(
+                item,
+                ResponseItem::Message { role, phase, .. }
+                    if role == "assistant" && !matches!(phase, Some(MessagePhase::Commentary))
+            )
+        })
         .count()
 }
 

codex-rs/core/src/next_prompt_suggestion_tests.rs

@@ -7,6 +7,7 @@ use super::suggestion_prompt_has_headroom;
 use crate::context_manager::ContextManager;
 use codex_protocol::models::ContentItem;
 use codex_protocol::models::FunctionCallOutputPayload;
+use codex_protocol::models::MessagePhase;
 use codex_protocol::models::ResponseItem;
 use codex_utils_output_truncation::TruncationPolicy;
 use pretty_assertions::assert_eq;
@@ -87,6 +88,26 @@ fn history_boundary_requires_final_assistant_message() {
     }
 }
 
+#[test]
+fn final_assistant_message_count_ignores_commentary() {
+    let assistant_message = |phase| ResponseItem::Message {
+        id: None,
+        role: "assistant".to_string(),
+        content: vec![ContentItem::OutputText {
+            text: "done".to_string(),
+        }],
+        phase,
+    };
+
+    assert_eq!(
+        super::final_assistant_message_count(&[
+            assistant_message(Some(MessagePhase::Commentary)),
+            assistant_message(Some(MessagePhase::FinalAnswer)),
+        ]),
+        1
+    );
+}
+
 #[test]
 fn suggestion_prompt_skips_near_context_window() {
     assert!(!suggestion_prompt_has_headroom(

codex-rs/core/tests/suite/next_prompt_suggestion.rs

@@ -3,6 +3,7 @@ use codex_protocol::protocol::ConversationStartParams;
 use codex_protocol::protocol::EventMsg;
 use codex_protocol::protocol::Op;
 use codex_protocol::protocol::RealtimeOutputModality;
+use codex_protocol::user_input::UserInput;
 use core_test_support::responses::ev_assistant_message;
 use core_test_support::responses::ev_completed;
 use core_test_support::responses::ev_completed_with_tokens;
@@ -11,10 +12,16 @@ use core_test_support::responses::mount_sse_sequence;
 use core_test_support::responses::sse;
 use core_test_support::responses::start_websocket_server;
 use core_test_support::skip_if_no_network;
+use core_test_support::streaming_sse::StreamingSseChunk;
+use core_test_support::streaming_sse::start_streaming_sse_server;
 use core_test_support::test_codex::test_codex;
 use core_test_support::wait_for_event;
 use pretty_assertions::assert_eq;
 use serde_json::json;
+use std::sync::Arc;
+use std::time::Duration;
+use tokio::sync::oneshot;
+use tokio::time::timeout;
 use tokio_util::sync::CancellationToken;
 use wiremock::MockServer;
 
@@ -182,6 +189,83 @@ async fn suggest_next_prompt_skips_active_realtime_conversation_without_request(
     Ok(())
 }
 
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn suggest_next_prompt_stops_when_real_turn_starts() -> Result<()> {
+    skip_if_no_network!(Ok(()));
+
+    let (release_suggestion_tx, release_suggestion_rx) = oneshot::channel();
+    let complete_response = |response_id, message_id, text| {
+        vec![StreamingSseChunk {
+            gate: None,
+            body: sse(vec![
+                ev_response_created(response_id),
+                ev_assistant_message(message_id, text),
+                ev_completed(response_id),
+            ]),
+        }]
+    };
+    let suggestion_chunks = vec![
+        StreamingSseChunk {
+            gate: None,
+            body: sse(vec![
+                ev_response_created("resp-suggestion"),
+                ev_assistant_message("msg-suggestion", "stale suggestion"),
+            ]),
+        },
+        StreamingSseChunk {
+            gate: Some(release_suggestion_rx),
+            body: sse(vec![ev_completed("resp-suggestion")]),
+        },
+    ];
+    let (server, _) = start_streaming_sse_server(vec![
+        complete_response("resp-1", "msg-1", "first turn complete"),
+        complete_response("resp-2", "msg-2", "second turn complete"),
+        suggestion_chunks,
+        complete_response("resp-3", "msg-3", "third turn complete"),
+    ])
+    .await;
+    let mut builder = test_codex();
+    let test = builder.build_with_streaming_server(&server).await?;
+    test.submit_turn("first task").await?;
+    test.submit_turn("second task").await?;
+
+    let codex = Arc::clone(&test.codex);
+    let suggestion_task =
+        tokio::spawn(async move { codex.suggest_next_prompt(CancellationToken::new()).await });
+    server.wait_for_request_count(/*count*/ 3).await;
+    test.codex
+        .submit(Op::UserInput {
+            items: vec![UserInput::Text {
+                text: "third task".to_string(),
+                text_elements: Vec::new(),
+            }],
+            environments: None,
+            final_output_json_schema: None,
+            responsesapi_client_metadata: None,
+            additional_context: Default::default(),
+            thread_settings: Default::default(),
+        })
+        .await?;
+    timeout(
+        Duration::from_secs(2),
+        server.wait_for_request_count(/*count*/ 4),
+    )
+    .await
+    .expect("real turn request should start");
+
+    let suggestion = timeout(Duration::from_secs(2), suggestion_task)
+        .await
+        .expect("suggestion task should stop after real turn starts")
+        .expect("suggestion task should not panic")?;
+    assert_eq!(suggestion, None);
+    let _ = release_suggestion_tx.send(());
+    wait_for_event(&test.codex, |msg| matches!(msg, EventMsg::TurnComplete(_))).await;
+    assert_eq!(server.requests().await.len(), 4);
+
+    server.shutdown().await;
+    Ok(())
+}
+
 fn ev_rate_limits() -> serde_json::Value {
     json!({
         "type": "codex.rate_limits",