From f8603f3e003ba12c64d008ac85f76db646b9f013 Mon Sep 17 00:00:00 2001 From: Felipe Coury Date: Fri, 22 May 2026 23:58:02 -0300 Subject: [PATCH] fix(app-server): preserve immediate suggestion cancels --- .../schema/json/ClientRequest.json | 5 +- .../codex_app_server_protocol.schemas.json | 5 +- .../codex_app_server_protocol.v2.schemas.json | 5 +- .../v2/ThreadSuggestNextPromptParams.json | 5 +- .../v2/ThreadSuggestNextPromptParams.ts | 2 +- .../src/protocol/common.rs | 23 +- .../src/protocol/v2/thread.rs | 4 +- .../src/request_processors/turn_processor.rs | 136 +++++++- .../suite/v2/thread_suggest_next_prompt.rs | 309 ++++++++++++++++++ codex-rs/core/src/client.rs | 4 + codex-rs/core/src/next_prompt_suggestion.rs | 80 +++-- .../tests/suite/next_prompt_suggestion.rs | 48 +++ 12 files changed, 565 insertions(+), 61 deletions(-) diff --git a/codex-rs/app-server-protocol/schema/json/ClientRequest.json b/codex-rs/app-server-protocol/schema/json/ClientRequest.json index 859f32ef61..64dfcb3b5b 100644 --- a/codex-rs/app-server-protocol/schema/json/ClientRequest.json +++ b/codex-rs/app-server-protocol/schema/json/ClientRequest.json @@ -3903,10 +3903,7 @@ "properties": { "cancel": { "description": "When true, cancel the request identified by `cancellation_token` without sampling.\n\nThis requires `cancellation_token`.", - "type": [ - "boolean", - "null" - ] + "type": "boolean" }, "cancellationToken": { "description": "Optional client token used to cancel an older in-flight suggestion request.\n\nReusing a token cancels the older request before this request samples.", diff --git a/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json b/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json index 75831d5fe3..0fb12aca77 100644 --- a/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json +++ b/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json @@ -18118,10 +18118,7 @@ "properties": { "cancel": { "description": "When true, cancel the request identified by `cancellation_token` without sampling.\n\nThis requires `cancellation_token`.", - "type": [ - "boolean", - "null" - ] + "type": "boolean" }, "cancellationToken": { "description": "Optional client token used to cancel an older in-flight suggestion request.\n\nReusing a token cancels the older request before this request samples.", diff --git a/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.v2.schemas.json b/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.v2.schemas.json index d3de6da632..991564c023 100644 --- a/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.v2.schemas.json +++ b/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.v2.schemas.json @@ -15935,10 +15935,7 @@ "properties": { "cancel": { "description": "When true, cancel the request identified by `cancellation_token` without sampling.\n\nThis requires `cancellation_token`.", - "type": [ - "boolean", - "null" - ] + "type": "boolean" }, "cancellationToken": { "description": "Optional client token used to cancel an older in-flight suggestion request.\n\nReusing a token cancels the older request before this request samples.", diff --git a/codex-rs/app-server-protocol/schema/json/v2/ThreadSuggestNextPromptParams.json b/codex-rs/app-server-protocol/schema/json/v2/ThreadSuggestNextPromptParams.json index ad1299165d..7d361ae019 100644 --- a/codex-rs/app-server-protocol/schema/json/v2/ThreadSuggestNextPromptParams.json +++ b/codex-rs/app-server-protocol/schema/json/v2/ThreadSuggestNextPromptParams.json @@ -4,10 +4,7 @@ "properties": { "cancel": { "description": "When true, cancel the request identified by `cancellation_token` without sampling.\n\nThis requires `cancellation_token`.", - "type": [ - "boolean", - "null" - ] + "type": "boolean" }, "cancellationToken": { "description": "Optional client token used to cancel an older in-flight suggestion request.\n\nReusing a token cancels the older request before this request samples.", diff --git a/codex-rs/app-server-protocol/schema/typescript/v2/ThreadSuggestNextPromptParams.ts b/codex-rs/app-server-protocol/schema/typescript/v2/ThreadSuggestNextPromptParams.ts index f1f1dd94a7..7b87e1496a 100644 --- a/codex-rs/app-server-protocol/schema/typescript/v2/ThreadSuggestNextPromptParams.ts +++ b/codex-rs/app-server-protocol/schema/typescript/v2/ThreadSuggestNextPromptParams.ts @@ -25,4 +25,4 @@ cancellationToken?: string | null, * * This requires `cancellation_token`. */ -cancel?: boolean | null, }; +cancel?: boolean, }; diff --git a/codex-rs/app-server-protocol/src/protocol/common.rs b/codex-rs/app-server-protocol/src/protocol/common.rs index 7c4ab057b6..bc011d3741 100644 --- a/codex-rs/app-server-protocol/src/protocol/common.rs +++ b/codex-rs/app-server-protocol/src/protocol/common.rs @@ -587,8 +587,7 @@ client_request_definitions! { }, ThreadSuggestNextPrompt => "thread/suggestNextPrompt" { params: v2::ThreadSuggestNextPromptParams, - // Suggestions do not mutate thread state, and cancellation requests must not queue - // behind the in-flight sampling request they are trying to stop. + // Hidden prediction must never queue interactive thread work behind its sampling request. serialization: None, response: v2::ThreadSuggestNextPromptResponse, }, @@ -2036,6 +2035,26 @@ mod tests { "remote-control-clients" )) ); + + let suggest_next_prompt_cancel = ClientRequest::ThreadSuggestNextPrompt { + request_id: request_id(), + params: v2::ThreadSuggestNextPromptParams { + thread_id: "suggest-thread".to_string(), + cancellation_token: Some("suggestion-1".to_string()), + cancel: true, + }, + }; + assert_eq!(suggest_next_prompt_cancel.serialization_scope(), None); + + let suggest_next_prompt = ClientRequest::ThreadSuggestNextPrompt { + request_id: request_id(), + params: v2::ThreadSuggestNextPromptParams { + thread_id: "suggest-thread".to_string(), + cancellation_token: Some("suggestion-1".to_string()), + cancel: false, + }, + }; + assert_eq!(suggest_next_prompt.serialization_scope(), None); } #[test] diff --git a/codex-rs/app-server-protocol/src/protocol/v2/thread.rs b/codex-rs/app-server-protocol/src/protocol/v2/thread.rs index 2d2d7c69c9..4fd2ea839d 100644 --- a/codex-rs/app-server-protocol/src/protocol/v2/thread.rs +++ b/codex-rs/app-server-protocol/src/protocol/v2/thread.rs @@ -1187,8 +1187,8 @@ pub struct ThreadSuggestNextPromptParams { /// When true, cancel the request identified by `cancellation_token` without sampling. /// /// This requires `cancellation_token`. - #[ts(optional = nullable)] - pub cancel: Option, + #[serde(default, skip_serializing_if = "std::ops::Not::not")] + pub cancel: bool, } #[derive(Serialize, Deserialize, Debug, Clone, PartialEq, JsonSchema, TS)] diff --git a/codex-rs/app-server/src/request_processors/turn_processor.rs b/codex-rs/app-server/src/request_processors/turn_processor.rs index bb2ad1f590..3a50d513ee 100644 --- a/codex-rs/app-server/src/request_processors/turn_processor.rs +++ b/codex-rs/app-server/src/request_processors/turn_processor.rs @@ -12,7 +12,8 @@ pub(crate) struct TurnRequestProcessor { config: Arc, config_manager: ConfigManager, pending_thread_unloads: Arc>>, - pending_next_prompt_suggestions: Arc>>>, + pending_next_prompt_suggestions: + Arc>>, thread_state_manager: ThreadStateManager, thread_watch_manager: ThreadWatchManager, thread_list_state_permit: Arc, @@ -72,6 +73,54 @@ struct ThreadSettingsBuildParams { personality: Option, } +#[derive(Clone, Eq, Hash, PartialEq)] +struct NextPromptSuggestionCancellationKey { + thread_id: String, + cancellation_token: String, +} + +impl NextPromptSuggestionCancellationKey { + fn new(thread_id: &str, cancellation_token: &str) -> Self { + Self { + thread_id: thread_id.to_string(), + cancellation_token: cancellation_token.to_string(), + } + } +} + +enum NextPromptSuggestionCancellation { + Active(Arc), + /// One-shot marker for a cancel request that arrived before its matching + /// suggestion request reserved the token. + Cancelled { + expires_at: std::time::Instant, + }, + /// Short-lived record for a request that was already cancelled or completed. + /// + /// This lets late duplicate cancels become no-ops instead of creating a + /// stale one-shot cancellation for a future request that reuses the token. + Retired { + expires_at: std::time::Instant, + }, +} + +const NEXT_PROMPT_SUGGESTION_RETIRED_TOKEN_TTL: std::time::Duration = + std::time::Duration::from_secs(/*secs*/ 60); + +fn prune_expired_next_prompt_suggestions( + pending_next_prompt_suggestions: &mut HashMap< + NextPromptSuggestionCancellationKey, + NextPromptSuggestionCancellation, + >, +) { + let now = std::time::Instant::now(); + pending_next_prompt_suggestions.retain(|_, cancellation| match cancellation { + NextPromptSuggestionCancellation::Active(_) => true, + NextPromptSuggestionCancellation::Cancelled { expires_at } + | NextPromptSuggestionCancellation::Retired { expires_at } => *expires_at > now, + }); +} + impl TurnRequestProcessor { #[allow(clippy::too_many_arguments)] pub(crate) fn new( @@ -726,9 +775,8 @@ impl TurnRequestProcessor { cancellation_token, cancel, } = params; - let cancel = cancel.unwrap_or_default(); let suggestion_cancellation = self - .next_prompt_suggestion_cancellation(cancellation_token.as_deref(), cancel) + .next_prompt_suggestion_cancellation(&thread_id, cancellation_token.as_deref(), cancel) .await?; let Some(suggestion_cancellation) = suggestion_cancellation else { return Ok(ThreadSuggestNextPromptResponse { suggestion: None }); @@ -743,6 +791,7 @@ impl TurnRequestProcessor { } .await; self.clear_next_prompt_suggestion_cancellation( + &thread_id, cancellation_token.as_deref(), &suggestion_cancellation, ) @@ -753,6 +802,7 @@ impl TurnRequestProcessor { async fn next_prompt_suggestion_cancellation( &self, + thread_id: &str, cancellation_token: Option<&str>, cancel: bool, ) -> Result>, JSONRPCErrorError> { @@ -766,21 +816,71 @@ impl TurnRequestProcessor { }; let mut pending_next_prompt_suggestions = self.pending_next_prompt_suggestions.lock().await; - if let Some(existing) = pending_next_prompt_suggestions.remove(cancellation_token) { - existing.cancel(); - } + prune_expired_next_prompt_suggestions(&mut pending_next_prompt_suggestions); + let key = NextPromptSuggestionCancellationKey::new(thread_id, cancellation_token); if cancel { + match pending_next_prompt_suggestions.remove(&key) { + Some(NextPromptSuggestionCancellation::Active(existing)) => { + existing.cancel(); + pending_next_prompt_suggestions.insert( + key, + NextPromptSuggestionCancellation::Retired { + expires_at: std::time::Instant::now() + + NEXT_PROMPT_SUGGESTION_RETIRED_TOKEN_TTL, + }, + ); + } + Some(NextPromptSuggestionCancellation::Cancelled { expires_at }) => { + pending_next_prompt_suggestions.insert( + key, + NextPromptSuggestionCancellation::Cancelled { expires_at }, + ); + } + Some(NextPromptSuggestionCancellation::Retired { expires_at }) => { + pending_next_prompt_suggestions.insert( + key, + NextPromptSuggestionCancellation::Retired { expires_at }, + ); + } + None => { + pending_next_prompt_suggestions.insert( + key, + NextPromptSuggestionCancellation::Cancelled { + expires_at: std::time::Instant::now() + + NEXT_PROMPT_SUGGESTION_RETIRED_TOKEN_TTL, + }, + ); + } + } return Ok(None); } + match pending_next_prompt_suggestions.remove(&key) { + Some(NextPromptSuggestionCancellation::Active(existing)) => existing.cancel(), + Some(NextPromptSuggestionCancellation::Cancelled { .. }) => { + pending_next_prompt_suggestions.insert( + key, + NextPromptSuggestionCancellation::Retired { + expires_at: std::time::Instant::now() + + NEXT_PROMPT_SUGGESTION_RETIRED_TOKEN_TTL, + }, + ); + return Ok(None); + } + Some(NextPromptSuggestionCancellation::Retired { .. }) | None => {} + } + let cancellation = Arc::new(CancellationToken::new()); - pending_next_prompt_suggestions - .insert(cancellation_token.to_string(), Arc::clone(&cancellation)); + pending_next_prompt_suggestions.insert( + key, + NextPromptSuggestionCancellation::Active(Arc::clone(&cancellation)), + ); Ok(Some(cancellation)) } async fn clear_next_prompt_suggestion_cancellation( &self, + thread_id: &str, cancellation_token: Option<&str>, suggestion_cancellation: &Arc, ) { @@ -789,11 +889,25 @@ impl TurnRequestProcessor { }; let mut pending_next_prompt_suggestions = self.pending_next_prompt_suggestions.lock().await; + prune_expired_next_prompt_suggestions(&mut pending_next_prompt_suggestions); + let key = NextPromptSuggestionCancellationKey::new(thread_id, cancellation_token); if pending_next_prompt_suggestions - .get(cancellation_token) - .is_some_and(|current| Arc::ptr_eq(current, suggestion_cancellation)) + .get(&key) + .is_some_and(|current| match current { + NextPromptSuggestionCancellation::Active(current) => { + Arc::ptr_eq(current, suggestion_cancellation) + } + NextPromptSuggestionCancellation::Cancelled { .. } + | NextPromptSuggestionCancellation::Retired { .. } => false, + }) { - pending_next_prompt_suggestions.remove(cancellation_token); + pending_next_prompt_suggestions.insert( + key, + NextPromptSuggestionCancellation::Retired { + expires_at: std::time::Instant::now() + + NEXT_PROMPT_SUGGESTION_RETIRED_TOKEN_TTL, + }, + ); } } diff --git a/codex-rs/app-server/tests/suite/v2/thread_suggest_next_prompt.rs b/codex-rs/app-server/tests/suite/v2/thread_suggest_next_prompt.rs index e2563cd36c..67b96c0a5b 100644 --- a/codex-rs/app-server/tests/suite/v2/thread_suggest_next_prompt.rs +++ b/codex-rs/app-server/tests/suite/v2/thread_suggest_next_prompt.rs @@ -11,6 +11,7 @@ use codex_app_server_protocol::ThreadSuggestNextPromptResponse; use codex_protocol::models::ContentItem; use codex_protocol::models::ResponseItem; use core_test_support::responses; +use pretty_assertions::assert_eq; use serde_json::json; use std::path::Path; use tempfile::TempDir; @@ -152,6 +153,314 @@ async fn thread_suggest_next_prompt_cancel_stops_in_flight_sampling() -> Result< Ok(()) } +#[tokio::test] +async fn thread_suggest_next_prompt_cancel_is_scoped_to_thread() -> Result<()> { + let server = responses::start_mock_server().await; + let body = responses::sse(vec![ + responses::ev_response_created("resp-1"), + responses::ev_assistant_message("msg-1", "run the tests"), + responses::ev_completed("resp-1"), + ]); + let response_mock = responses::mount_response_once( + &server, + responses::sse_response(body) + .set_delay(std::time::Duration::from_millis(/*millis*/ 100)), + ) + .await; + + let codex_home = TempDir::new()?; + create_config_toml(codex_home.path(), &server.uri())?; + + let mut mcp = McpProcess::new(codex_home.path()).await?; + timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; + + let first_thread_req = mcp + .send_thread_start_request(ThreadStartParams { + model: Some("mock-model".to_string()), + ..Default::default() + }) + .await?; + let first_thread_resp: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(first_thread_req)), + ) + .await??; + let ThreadStartResponse { + thread: first_thread, + .. + } = to_response::(first_thread_resp)?; + inject_suggestion_history(&mut mcp, &first_thread.id).await?; + + let second_thread_req = mcp + .send_thread_start_request(ThreadStartParams { + model: Some("mock-model".to_string()), + ..Default::default() + }) + .await?; + let second_thread_resp: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(second_thread_req)), + ) + .await??; + let ThreadStartResponse { + thread: second_thread, + .. + } = to_response::(second_thread_resp)?; + + let suggestion_req = mcp + .send_raw_request( + "thread/suggestNextPrompt", + Some(json!({ + "threadId": first_thread.id, + "cancellationToken": "next-prompt-suggestion-shared-token", + })), + ) + .await?; + wait_for_request_count(&response_mock, /*expected*/ 1).await?; + + let cancel_req = mcp + .send_raw_request( + "thread/suggestNextPrompt", + Some(json!({ + "threadId": second_thread.id, + "cancellationToken": "next-prompt-suggestion-shared-token", + "cancel": true, + })), + ) + .await?; + let cancel_resp: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(cancel_req)), + ) + .await??; + let cancel_resp = to_response::(cancel_resp)?; + assert_eq!(cancel_resp.suggestion, None); + + let suggestion_resp: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(suggestion_req)), + ) + .await??; + let suggestion_resp = to_response::(suggestion_resp)?; + assert_eq!(suggestion_resp.suggestion.as_deref(), Some("run the tests")); + assert_eq!(response_mock.requests().len(), /*expected*/ 1); + + Ok(()) +} + +#[tokio::test] +async fn thread_suggest_next_prompt_cancel_before_start_suppresses_sampling() -> Result<()> { + let server = responses::start_mock_server().await; + let body = responses::sse(vec![ + responses::ev_response_created("resp-1"), + responses::ev_assistant_message("msg-1", "run the tests"), + responses::ev_completed("resp-1"), + ]); + let response_mock = responses::mount_sse_once(&server, body).await; + + let codex_home = TempDir::new()?; + create_config_toml(codex_home.path(), &server.uri())?; + + let mut mcp = McpProcess::new(codex_home.path()).await?; + timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; + + let thread_req = mcp + .send_thread_start_request(ThreadStartParams { + model: Some("mock-model".to_string()), + ..Default::default() + }) + .await?; + let thread_resp: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(thread_req)), + ) + .await??; + let ThreadStartResponse { thread, .. } = to_response::(thread_resp)?; + inject_suggestion_history(&mut mcp, &thread.id).await?; + + let cancel_req = mcp + .send_raw_request( + "thread/suggestNextPrompt", + Some(json!({ + "threadId": thread.id, + "cancellationToken": "next-prompt-suggestion-immediate-cancel", + "cancel": true, + })), + ) + .await?; + let cancel_resp: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(cancel_req)), + ) + .await??; + let cancel_resp = to_response::(cancel_resp)?; + assert_eq!(cancel_resp.suggestion, None); + + let suggestion_req = mcp + .send_raw_request( + "thread/suggestNextPrompt", + Some(json!({ + "threadId": thread.id, + "cancellationToken": "next-prompt-suggestion-immediate-cancel", + })), + ) + .await?; + let suggestion_resp: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(suggestion_req)), + ) + .await??; + let suggestion_resp = to_response::(suggestion_resp)?; + assert_eq!(suggestion_resp.suggestion, None); + assert_eq!(response_mock.requests().len(), /*expected*/ 0); + + let late_duplicate_cancel_req = mcp + .send_raw_request( + "thread/suggestNextPrompt", + Some(json!({ + "threadId": thread.id, + "cancellationToken": "next-prompt-suggestion-immediate-cancel", + "cancel": true, + })), + ) + .await?; + let late_duplicate_cancel_resp: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(late_duplicate_cancel_req)), + ) + .await??; + let late_duplicate_cancel_resp = + to_response::(late_duplicate_cancel_resp)?; + assert_eq!(late_duplicate_cancel_resp.suggestion, None); + + let reused_token_req = mcp + .send_raw_request( + "thread/suggestNextPrompt", + Some(json!({ + "threadId": thread.id, + "cancellationToken": "next-prompt-suggestion-immediate-cancel", + })), + ) + .await?; + let reused_token_resp: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(reused_token_req)), + ) + .await??; + let reused_token_resp = to_response::(reused_token_resp)?; + assert_eq!( + reused_token_resp.suggestion.as_deref(), + Some("run the tests") + ); + assert_eq!(response_mock.requests().len(), /*expected*/ 1); + + Ok(()) +} + +#[tokio::test] +async fn thread_suggest_next_prompt_late_cancel_does_not_suppress_reused_token() -> Result<()> { + let server = responses::start_mock_server().await; + let response_mock = responses::mount_sse_sequence( + &server, + vec![ + responses::sse(vec![ + responses::ev_response_created("resp-1"), + responses::ev_assistant_message("msg-1", "run the tests"), + responses::ev_completed("resp-1"), + ]), + responses::sse(vec![ + responses::ev_response_created("resp-2"), + responses::ev_assistant_message("msg-2", "commit changes"), + responses::ev_completed("resp-2"), + ]), + ], + ) + .await; + + let codex_home = TempDir::new()?; + create_config_toml(codex_home.path(), &server.uri())?; + + let mut mcp = McpProcess::new(codex_home.path()).await?; + timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; + + let thread_req = mcp + .send_thread_start_request(ThreadStartParams { + model: Some("mock-model".to_string()), + ..Default::default() + }) + .await?; + let thread_resp: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(thread_req)), + ) + .await??; + let ThreadStartResponse { thread, .. } = to_response::(thread_resp)?; + inject_suggestion_history(&mut mcp, &thread.id).await?; + + let first_suggestion_req = mcp + .send_raw_request( + "thread/suggestNextPrompt", + Some(json!({ + "threadId": thread.id, + "cancellationToken": "next-prompt-suggestion-late-cancel", + })), + ) + .await?; + let first_suggestion_resp: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(first_suggestion_req)), + ) + .await??; + let first_suggestion_resp = + to_response::(first_suggestion_resp)?; + assert_eq!( + first_suggestion_resp.suggestion.as_deref(), + Some("run the tests") + ); + + let cancel_req = mcp + .send_raw_request( + "thread/suggestNextPrompt", + Some(json!({ + "threadId": thread.id, + "cancellationToken": "next-prompt-suggestion-late-cancel", + "cancel": true, + })), + ) + .await?; + let cancel_resp: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(cancel_req)), + ) + .await??; + let cancel_resp = to_response::(cancel_resp)?; + assert_eq!(cancel_resp.suggestion, None); + + let second_suggestion_req = mcp + .send_raw_request( + "thread/suggestNextPrompt", + Some(json!({ + "threadId": thread.id, + "cancellationToken": "next-prompt-suggestion-late-cancel", + })), + ) + .await?; + let second_suggestion_resp: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(second_suggestion_req)), + ) + .await??; + let second_suggestion_resp = + to_response::(second_suggestion_resp)?; + assert_eq!( + second_suggestion_resp.suggestion.as_deref(), + Some("commit changes") + ); + assert_eq!(response_mock.requests().len(), /*expected*/ 2); + + Ok(()) +} + async fn inject_suggestion_history(mcp: &mut McpProcess, thread_id: &str) -> Result<()> { let items = vec![ message("user", "fix the bug"), diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index bceb250730..6385548a71 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -394,6 +394,10 @@ impl ModelClient { } } + /// Creates a streaming session isolated from the reusable turn WebSocket cache. + /// + /// Hidden best-effort requests use this so they cannot consume or reset the warm connection + /// reserved for the next real turn. pub(crate) fn new_isolated_session(&self) -> ModelClientSession { ModelClientSession { client: self.clone(), diff --git a/codex-rs/core/src/next_prompt_suggestion.rs b/codex-rs/core/src/next_prompt_suggestion.rs index d3db54a1d4..8cb98a1e95 100644 --- a/codex-rs/core/src/next_prompt_suggestion.rs +++ b/codex-rs/core/src/next_prompt_suggestion.rs @@ -43,7 +43,7 @@ mod next_prompt_suggestion_tests; const NEXT_PROMPT_SUGGESTION_TOKEN_HEADROOM: i64 = 1_024; const NEXT_PROMPT_SUGGESTION_MAX_HISTORY_TOKENS: usize = 4_096; const NEXT_PROMPT_SUGGESTION_MAX_OUTPUT_TOKENS: u64 = 32; -const NEXT_PROMPT_SUGGESTION_SAMPLE_TIMEOUT: Duration = Duration::from_secs(8); +const NEXT_PROMPT_SUGGESTION_SAMPLE_TIMEOUT: Duration = Duration::from_secs(/*secs*/ 8); #[derive(Clone, Copy)] struct HistorySnapshot { @@ -115,38 +115,55 @@ pub(crate) async fn suggest_next_prompt( return None; } let mut client_session = sess.services.model_client.new_isolated_session(); - let mut stream = match client_session - .stream( - &prompt, - &turn_context.model_info, - &turn_context.session_telemetry, - turn_context.reasoning_effort, - turn_context.reasoning_summary, - turn_context.config.service_tier.clone(), - /*turn_metadata_header*/ None, - &InferenceTraceContext::disabled(), - ) - .or_cancel(&cancellation_token) - .await - { - Ok(Ok(stream)) => stream, - Ok(Err(err)) => { - tracing::debug!( - error = ?err, - "next prompt suggestion failed before sampling started" - ); - return None; - } - Err(codex_async_utils::CancelErr::Cancelled) => { - tracing::debug!("next prompt suggestion canceled before sampling started"); - return None; + let sample_deadline = tokio::time::sleep(NEXT_PROMPT_SUGGESTION_SAMPLE_TIMEOUT); + tokio::pin!(sample_deadline); + let inference_trace = InferenceTraceContext::disabled(); + let mut stream = { + let stream = client_session + .stream( + &prompt, + &turn_context.model_info, + &turn_context.session_telemetry, + turn_context.reasoning_effort, + turn_context.reasoning_summary, + turn_context.config.service_tier.clone(), + /*turn_metadata_header*/ None, + &inference_trace, + ) + .or_cancel(&cancellation_token); + match tokio::select! { + result = stream => result, + _ = &mut sample_deadline => { + tracing::debug!("next prompt suggestion timed out before sampling started"); + return None; + } + } { + Ok(Ok(stream)) => stream, + Ok(Err(err)) => { + tracing::debug!( + error = ?err, + "next prompt suggestion failed before sampling started" + ); + return None; + } + Err(codex_async_utils::CancelErr::Cancelled) => { + tracing::debug!("next prompt suggestion canceled before sampling started"); + return None; + } } }; + if cancellation_token.is_cancelled() { + tracing::debug!("next prompt suggestion skipped after cancellation"); + client_session.reset_websocket_session(); + return None; + } + if !session_is_idle_for_suggestion(sess).await { + client_session.reset_websocket_session(); + return None; + } let mut streamed_text = String::new(); let mut completed_text = None; let mut latest_rate_limits = None; - let sample_deadline = tokio::time::sleep(NEXT_PROMPT_SUGGESTION_SAMPLE_TIMEOUT); - tokio::pin!(sample_deadline); let completed_response_id = loop { if !session_is_idle_for_suggestion(sess).await { client_session.reset_websocket_session(); @@ -161,7 +178,7 @@ pub(crate) async fn suggest_next_prompt( return None; } }, - _ = tokio::time::sleep(Duration::from_millis(100)) => continue, + _ = tokio::time::sleep(Duration::from_millis(/*millis*/ 100)) => continue, _ = &mut sample_deadline => { tracing::debug!("next prompt suggestion timed out while sampling"); client_session.reset_websocket_session(); @@ -225,6 +242,10 @@ pub(crate) async fn suggest_next_prompt( .await; } client_session.reset_websocket_session(); + if cancellation_token.is_cancelled() { + tracing::debug!("next prompt suggestion canceled after sampling completed"); + return None; + } if !session_history_matches_snapshot(sess, history_snapshot).await { tracing::debug!("next prompt suggestion skipped after history changed"); return None; @@ -410,6 +431,7 @@ fn has_unpaired_tool_flow(items: &[ResponseItem]) -> bool { client_tool_search_outputs.insert(call_id.clone()); } } + ResponseItem::ToolSearchOutput { execution, .. } if execution == "server" => {} ResponseItem::CustomToolCall { call_id, .. } => { custom_tool_calls.insert(call_id.clone()); } diff --git a/codex-rs/core/tests/suite/next_prompt_suggestion.rs b/codex-rs/core/tests/suite/next_prompt_suggestion.rs index 594ca451bb..a74589a015 100644 --- a/codex-rs/core/tests/suite/next_prompt_suggestion.rs +++ b/codex-rs/core/tests/suite/next_prompt_suggestion.rs @@ -8,8 +8,10 @@ use core_test_support::responses::ev_assistant_message; use core_test_support::responses::ev_completed; use core_test_support::responses::ev_completed_with_tokens; use core_test_support::responses::ev_response_created; +use core_test_support::responses::mount_response_sequence; use core_test_support::responses::mount_sse_sequence; use core_test_support::responses::sse; +use core_test_support::responses::sse_response; use core_test_support::responses::start_websocket_server; use core_test_support::skip_if_no_network; use core_test_support::streaming_sse::StreamingSseChunk; @@ -267,6 +269,52 @@ async fn suggest_next_prompt_drops_stale_result_when_real_turn_starts() -> Resul Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn suggest_next_prompt_times_out_during_stream_startup() -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = MockServer::start().await; + let first_turn_response = sse_response(sse(vec![ + ev_response_created("resp-1"), + ev_assistant_message("msg-1", "first turn complete"), + ev_completed("resp-1"), + ])); + let second_turn_response = sse_response(sse(vec![ + ev_response_created("resp-2"), + ev_assistant_message("msg-2", "second turn complete"), + ev_completed("resp-2"), + ])); + let delayed_suggestion_response = sse_response(sse(vec![ + ev_response_created("resp-suggestion"), + ev_assistant_message("msg-suggestion", "run the tests"), + ev_completed("resp-suggestion"), + ])) + .set_delay(Duration::from_secs(/*secs*/ 30)); + let responses = mount_response_sequence( + &server, + vec![ + first_turn_response, + second_turn_response, + delayed_suggestion_response, + ], + ) + .await; + + let test = test_codex().build(&server).await?; + test.submit_turn("first task").await?; + test.submit_turn("second task").await?; + + let suggestion = timeout( + Duration::from_secs(/*secs*/ 10), + test.codex.suggest_next_prompt(CancellationToken::new()), + ) + .await??; + assert_eq!(suggestion, None); + assert_eq!(responses.requests().len(), 3); + + Ok(()) +} + fn ev_rate_limits() -> serde_json::Value { json!({ "type": "codex.rate_limits",