diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index 93bd025f7a..f6f42d9274 100644 --- a/codex-rs/core/src/client.rs +++ b/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>, } +#[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); diff --git a/codex-rs/core/src/next_prompt_suggestion.rs b/codex-rs/core/src/next_prompt_suggestion.rs index 9b19ca14e3..00cec6562f 100644 --- a/codex-rs/core/src/next_prompt_suggestion.rs +++ b/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() } diff --git a/codex-rs/core/src/next_prompt_suggestion_tests.rs b/codex-rs/core/src/next_prompt_suggestion_tests.rs index 907d3e3add..4c1b474209 100644 --- a/codex-rs/core/src/next_prompt_suggestion_tests.rs +++ b/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( diff --git a/codex-rs/core/tests/suite/next_prompt_suggestion.rs b/codex-rs/core/tests/suite/next_prompt_suggestion.rs index 9f4bf15349..0d3ab9f9c1 100644 --- a/codex-rs/core/tests/suite/next_prompt_suggestion.rs +++ b/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",