From 6fb17e14c2d6cdfc8a00936fba26c35c3f029788 Mon Sep 17 00:00:00 2001 From: Felipe Coury Date: Sun, 31 May 2026 13:54:11 -0300 Subject: [PATCH] fix(core): bound next-prompt suggestion requests --- codex-rs/codex-api/src/common.rs | 5 +++ codex-rs/codex-api/tests/clients.rs | 2 + codex-rs/core/src/client.rs | 17 ++++--- codex-rs/core/src/client_common.rs | 4 ++ codex-rs/core/src/client_common_tests.rs | 4 ++ codex-rs/core/src/compact_remote.rs | 1 + codex-rs/core/src/compact_remote_v2.rs | 1 + codex-rs/core/src/next_prompt_suggestion.rs | 44 ++++++++++++++++++- .../core/src/next_prompt_suggestion_tests.rs | 40 +++++++++++++++++ codex-rs/core/src/session/turn.rs | 1 + .../tests/suite/next_prompt_suggestion.rs | 1 + 11 files changed, 112 insertions(+), 8 deletions(-) diff --git a/codex-rs/codex-api/src/common.rs b/codex-rs/codex-api/src/common.rs index 91b251c41f..3bb59a46a3 100644 --- a/codex-rs/codex-api/src/common.rs +++ b/codex-rs/codex-api/src/common.rs @@ -180,6 +180,8 @@ pub struct ResponsesApiRequest { pub stream: bool, pub include: Vec, #[serde(skip_serializing_if = "Option::is_none")] + pub max_output_tokens: Option, + #[serde(skip_serializing_if = "Option::is_none")] pub service_tier: Option, #[serde(skip_serializing_if = "Option::is_none")] pub prompt_cache_key: Option, @@ -203,6 +205,7 @@ impl From<&ResponsesApiRequest> for ResponseCreateWsRequest { store: request.store, stream: request.stream, include: request.include.clone(), + max_output_tokens: request.max_output_tokens, service_tier: request.service_tier.clone(), prompt_cache_key: request.prompt_cache_key.clone(), text: request.text.clone(), @@ -228,6 +231,8 @@ pub struct ResponseCreateWsRequest { pub stream: bool, pub include: Vec, #[serde(skip_serializing_if = "Option::is_none")] + pub max_output_tokens: Option, + #[serde(skip_serializing_if = "Option::is_none")] pub service_tier: Option, #[serde(skip_serializing_if = "Option::is_none")] pub prompt_cache_key: Option, diff --git a/codex-rs/codex-api/tests/clients.rs b/codex-rs/codex-api/tests/clients.rs index afaf82389f..2fe4c784e5 100644 --- a/codex-rs/codex-api/tests/clients.rs +++ b/codex-rs/codex-api/tests/clients.rs @@ -331,6 +331,7 @@ async fn streaming_client_retries_on_transport_error() -> Result<()> { store: false, stream: true, include: Vec::new(), + max_output_tokens: None, service_tier: None, prompt_cache_key: None, text: None, @@ -432,6 +433,7 @@ async fn azure_default_store_attaches_ids_and_headers() -> Result<()> { store: true, stream: true, include: Vec::new(), + max_output_tokens: None, service_tier: None, prompt_cache_key: None, text: None, diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index f6f42d9274..ebed7ea288 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -798,6 +798,7 @@ impl ModelClient { store: provider.is_azure_responses_endpoint(), stream: true, include, + max_output_tokens: prompt.max_output_tokens, service_tier, prompt_cache_key, text, @@ -1663,10 +1664,11 @@ impl ModelClientSession { } } - /// Permanently disables WebSockets for this Codex session and resets WebSocket state. + /// Disables WebSockets for this Codex session and resets WebSocket state. /// - /// This is used after exhausting the provider retry budget, to force subsequent requests onto - /// the HTTP transport. + /// Shared sessions permanently force subsequent requests onto HTTP after exhausting the + /// provider retry budget. Isolated sessions only reset their own WebSocket state, so hidden + /// best-effort requests cannot change transport behavior for real turns. /// /// Returns `true` if this call activated fallback, or `false` if fallback was already active. pub(crate) fn try_switch_fallback_transport( @@ -1674,9 +1676,12 @@ impl ModelClientSession { session_telemetry: &SessionTelemetry, model_info: &ModelInfo, ) -> bool { - let activated = self - .client - .force_http_fallback(session_telemetry, model_info); + let activated = match self.websocket_session_cache_mode { + WebsocketSessionCacheMode::Shared => self + .client + .force_http_fallback(session_telemetry, model_info), + WebsocketSessionCacheMode::Isolated => false, + }; self.websocket_session = WebsocketSession::default(); activated } diff --git a/codex-rs/core/src/client_common.rs b/codex-rs/core/src/client_common.rs index d775f9c01c..e08449ddaa 100644 --- a/codex-rs/core/src/client_common.rs +++ b/codex-rs/core/src/client_common.rs @@ -35,6 +35,9 @@ pub struct Prompt { /// Whether the Responses API should strictly validate `output_schema`. pub output_schema_strict: bool, + + /// Optional cap for generated output tokens. + pub max_output_tokens: Option, } impl Default for Prompt { @@ -47,6 +50,7 @@ impl Default for Prompt { personality: None, output_schema: None, output_schema_strict: true, + max_output_tokens: None, } } } diff --git a/codex-rs/core/src/client_common_tests.rs b/codex-rs/core/src/client_common_tests.rs index 45d4334222..7908f7f9b2 100644 --- a/codex-rs/core/src/client_common_tests.rs +++ b/codex-rs/core/src/client_common_tests.rs @@ -22,6 +22,7 @@ fn serializes_text_verbosity_when_set() { store: false, stream: true, include: vec![], + max_output_tokens: None, prompt_cache_key: None, service_tier: None, text: Some(TextControls { @@ -69,6 +70,7 @@ fn serializes_text_schema_with_strict_format() { store: false, stream: true, include: vec![], + max_output_tokens: None, prompt_cache_key: None, service_tier: None, text: Some(text_controls), @@ -130,6 +132,7 @@ fn omits_text_when_not_set() { store: false, stream: true, include: vec![], + max_output_tokens: None, prompt_cache_key: None, service_tier: None, text: None, @@ -153,6 +156,7 @@ fn serializes_flex_service_tier_when_set() { store: false, stream: true, include: vec![], + max_output_tokens: None, prompt_cache_key: None, service_tier: Some(ServiceTier::Flex.to_string()), text: None, diff --git a/codex-rs/core/src/compact_remote.rs b/codex-rs/core/src/compact_remote.rs index c2ef417b9b..1b26536bcf 100644 --- a/codex-rs/core/src/compact_remote.rs +++ b/codex-rs/core/src/compact_remote.rs @@ -199,6 +199,7 @@ async fn run_remote_compact_task_inner_impl( personality: turn_context.personality, output_schema: None, output_schema_strict: true, + max_output_tokens: None, }; let window_id = sess.services.model_client.current_window_id(); let turn_metadata_header = turn_context diff --git a/codex-rs/core/src/compact_remote_v2.rs b/codex-rs/core/src/compact_remote_v2.rs index eb10eb211d..ea7c2177df 100644 --- a/codex-rs/core/src/compact_remote_v2.rs +++ b/codex-rs/core/src/compact_remote_v2.rs @@ -216,6 +216,7 @@ async fn run_remote_compact_task_inner_impl( personality: turn_context.personality, output_schema: None, output_schema_strict: true, + max_output_tokens: None, }; let window_id = sess.services.model_client.current_window_id(); diff --git a/codex-rs/core/src/next_prompt_suggestion.rs b/codex-rs/core/src/next_prompt_suggestion.rs index 00cec6562f..13e2fa03a6 100644 --- a/codex-rs/core/src/next_prompt_suggestion.rs +++ b/codex-rs/core/src/next_prompt_suggestion.rs @@ -29,6 +29,7 @@ use codex_protocol::protocol::Event; use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::TokenCountEvent; use codex_rollout_trace::InferenceTraceContext; +use codex_utils_output_truncation::approx_token_count; use futures::StreamExt; use std::collections::HashSet; use std::time::Duration; @@ -40,6 +41,8 @@ use tokio_util::sync::CancellationToken; 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); #[derive(Clone, Copy)] @@ -80,7 +83,7 @@ pub(crate) async fn suggest_next_prompt( tracing::debug!("next prompt suggestion skipped for incomplete tool flow"); return None; } - let mut prompt_input = history.for_prompt(&turn_context.model_info.input_modalities); + let prompt_input = history.for_prompt(&turn_context.model_info.input_modalities); if !history_ends_at_assistant_response(&prompt_input) { tracing::debug!("next prompt suggestion skipped before assistant boundary"); return None; @@ -88,6 +91,12 @@ pub(crate) async fn suggest_next_prompt( if final_assistant_message_count(&prompt_input) < 2 { return None; } + let Some(mut prompt_input) = + bounded_recent_prompt_history(prompt_input, NEXT_PROMPT_SUGGESTION_MAX_HISTORY_TOKENS) + else { + tracing::debug!("next prompt suggestion skipped without bounded recent history"); + return None; + }; prompt_input.push(ContextualUserFragment::into( NextPromptSuggestionInstructions, )); @@ -100,6 +109,7 @@ pub(crate) async fn suggest_next_prompt( personality: turn_context.personality, output_schema: None, output_schema_strict: true, + max_output_tokens: Some(NEXT_PROMPT_SUGGESTION_MAX_OUTPUT_TOKENS), }; if !session_is_idle_for_suggestion(sess).await { return None; @@ -305,8 +315,38 @@ fn final_assistant_message_count(items: &[ResponseItem]) -> usize { .count() } +fn bounded_recent_prompt_history( + items: Vec, + max_tokens: usize, +) -> Option> { + let mut remaining_tokens = max_tokens; + let mut keep_from = items.len(); + for (index, item) in items.iter().enumerate().rev() { + let serialized = serde_json::to_string(item).ok()?; + let item_tokens = approx_token_count(&serialized); + if item_tokens > remaining_tokens { + break; + } + keep_from = index; + remaining_tokens = remaining_tokens.saturating_sub(item_tokens); + } + let bounded = items[keep_from..].to_vec(); + if bounded.is_empty() + || has_unpaired_tool_flow(&bounded) + || !history_ends_at_assistant_response(&bounded) + || final_assistant_message_count(&bounded) < 2 + { + return None; + } + Some(bounded) +} + fn history_ends_at_assistant_response(items: &[ResponseItem]) -> bool { - matches!(items.last(), Some(ResponseItem::Message { role, .. }) if role == "assistant") + matches!( + items.last(), + Some(ResponseItem::Message { role, phase, .. }) + if role == "assistant" && !matches!(phase, Some(MessagePhase::Commentary)) + ) } fn assistant_output_text(item: &ResponseItem) -> Option { diff --git a/codex-rs/core/src/next_prompt_suggestion_tests.rs b/codex-rs/core/src/next_prompt_suggestion_tests.rs index 4c1b474209..8ddb265bd9 100644 --- a/codex-rs/core/src/next_prompt_suggestion_tests.rs +++ b/codex-rs/core/src/next_prompt_suggestion_tests.rs @@ -1,4 +1,5 @@ use super::HistorySnapshot; +use super::bounded_recent_prompt_history; use super::filter_next_prompt_suggestion; use super::has_unpaired_tool_flow; use super::history_ends_at_assistant_response; @@ -76,6 +77,14 @@ fn history_boundary_requires_final_assistant_message() { ))); for tail in [ user, + ResponseItem::Message { + id: None, + role: "assistant".to_string(), + content: vec![ContentItem::OutputText { + text: "still working".to_string(), + }], + phase: Some(MessagePhase::Commentary), + }, ResponseItem::FunctionCallOutput { call_id: "call-1".to_string(), output: FunctionCallOutputPayload::from_text("done".to_string()), @@ -88,6 +97,37 @@ fn history_boundary_requires_final_assistant_message() { } } +#[test] +fn bounded_recent_prompt_history_keeps_recent_complete_turns() { + let message = |role: &str, text: &str| ResponseItem::Message { + id: None, + role: role.to_string(), + content: vec![if role == "assistant" { + ContentItem::OutputText { + text: text.to_string(), + } + } else { + ContentItem::InputText { + text: text.to_string(), + } + }], + phase: None, + }; + let recent = vec![ + message("user", "first recent task"), + message("assistant", "first recent answer"), + message("user", "second recent task"), + message("assistant", "second recent answer"), + ]; + let mut history = vec![message("user", &"old context ".repeat(1_000))]; + history.extend(recent.clone()); + + assert_eq!( + bounded_recent_prompt_history(history, /*max_tokens*/ 256), + Some(recent) + ); +} + #[test] fn final_assistant_message_count_ignores_commentary() { let assistant_message = |phase| ResponseItem::Message { diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index 42d609561e..f9d6c3db21 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -970,6 +970,7 @@ pub(crate) fn build_prompt( output_schema_strict: !crate::guardian::is_guardian_reviewer_source( &turn_context.session_source, ), + max_output_tokens: None, } } diff --git a/codex-rs/core/tests/suite/next_prompt_suggestion.rs b/codex-rs/core/tests/suite/next_prompt_suggestion.rs index 2bfea55909..594ca451bb 100644 --- a/codex-rs/core/tests/suite/next_prompt_suggestion.rs +++ b/codex-rs/core/tests/suite/next_prompt_suggestion.rs @@ -73,6 +73,7 @@ async fn suggest_next_prompt_samples_history_without_tools() -> Result<()> { let suggestion_body = suggestion_request.body_json(); assert_eq!(suggestion_body["tools"], json!([])); assert_eq!(suggestion_body["parallel_tool_calls"], false); + assert_eq!(suggestion_body["max_output_tokens"], json!(32)); let user_texts = suggestion_request.message_input_texts("user"); let suggestion_prompt = user_texts