diff --git a/codex-rs/core/src/codex_thread.rs b/codex-rs/core/src/codex_thread.rs index d5135899a7..536f0a2ab8 100644 --- a/codex-rs/core/src/codex_thread.rs +++ b/codex-rs/core/src/codex_thread.rs @@ -154,11 +154,11 @@ impl CodexThread { &self, cancellation_token: CancellationToken, ) -> CodexResult> { - crate::next_prompt_suggestion::suggest_next_prompt( + Ok(crate::next_prompt_suggestion::suggest_next_prompt( self.codex.session.as_ref(), cancellation_token, ) - .await + .await) } /// Wait until the underlying session loop has terminated. diff --git a/codex-rs/core/src/next_prompt_suggestion.rs b/codex-rs/core/src/next_prompt_suggestion.rs index 925cf8bfa1..a857ea1615 100644 --- a/codex-rs/core/src/next_prompt_suggestion.rs +++ b/codex-rs/core/src/next_prompt_suggestion.rs @@ -7,7 +7,7 @@ //! tools, mutate transcript state, or decide whether the TUI should render the //! result. //! -//! Suggestions are best-effort. The caller should treat `Ok(None)` as an +//! Suggestions are best-effort. The caller should treat `None` as an //! expected silent outcome for early conversations, active turns, incomplete //! tool flow, model silence, or filtered output. @@ -20,7 +20,6 @@ use crate::context_manager::ContextManager; use crate::session::session::Session; use codex_async_utils::OrCancelExt; use codex_features::Feature; -use codex_protocol::error::Result as CodexResult; use codex_protocol::models::ContentItem; use codex_protocol::models::ResponseItem; use codex_protocol::openai_models::ModelPreset; @@ -53,40 +52,40 @@ struct HistorySnapshot { /// The sample uses the prompt-visible history from `sess` plus one synthetic /// suggestion instruction. Active turns and histories with unmatched tool call /// pairs are suppressed before sampling because those states do not represent a -/// stable completed conversation boundary. Returning `Ok(None)` means there is +/// stable completed conversation boundary. Returning `None` means there is /// no suggestion worth showing, not that the request failed. pub(crate) async fn suggest_next_prompt( sess: &Session, cancellation_token: CancellationToken, -) -> CodexResult> { +) -> Option { if cancellation_token.is_cancelled() { tracing::debug!("next prompt suggestion skipped after cancellation"); - return Ok(None); + return None; } if !session_is_idle_for_suggestion(sess).await { - return Ok(None); + return None; } let started_at = Instant::now(); let mut turn_context = sess.new_lightweight_turn().await; prefer_fast_suggestion_profile(&mut turn_context); if !suggestion_prompt_fits_context_window(sess, &turn_context).await { - return Ok(None); + return None; } let history = sess.clone_history().await; let history_snapshot = HistorySnapshot::from_history(&history); if has_unpaired_tool_flow(history.raw_items()) { tracing::debug!("next prompt suggestion skipped for incomplete tool flow"); - return Ok(None); + return None; } let mut 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 Ok(None); + return None; } if assistant_message_count(&prompt_input) < 2 { - return Ok(None); + return None; } prompt_input.push(ContextualUserFragment::into( NextPromptSuggestionInstructions, @@ -102,7 +101,7 @@ pub(crate) async fn suggest_next_prompt( output_schema_strict: true, }; if !session_is_idle_for_suggestion(sess).await { - return Ok(None); + return None; } let mut client_session = sess.services.model_client.new_session(); let mut stream = match client_session @@ -125,11 +124,11 @@ pub(crate) async fn suggest_next_prompt( error = ?err, "next prompt suggestion failed before sampling started" ); - return Ok(None); + return None; } Err(codex_async_utils::CancelErr::Cancelled) => { tracing::debug!("next prompt suggestion canceled before sampling started"); - return Ok(None); + return None; } }; let mut streamed_text = String::new(); @@ -140,7 +139,7 @@ pub(crate) async fn suggest_next_prompt( let completed_response_id = loop { if !session_is_idle_for_suggestion(sess).await { client_session.reset_websocket_session(); - return Ok(None); + return None; } let Some(event) = (tokio::select! { event = stream.next().or_cancel(&cancellation_token) => match event { @@ -148,19 +147,19 @@ pub(crate) async fn suggest_next_prompt( Err(codex_async_utils::CancelErr::Cancelled) => { tracing::debug!("next prompt suggestion canceled while sampling"); client_session.reset_websocket_session(); - return Ok(None); + return None; } }, _ = tokio::time::sleep(Duration::from_millis(100)) => continue, _ = &mut sample_deadline => { tracing::debug!("next prompt suggestion timed out while sampling"); client_session.reset_websocket_session(); - return Ok(None); + return None; }, }) else { tracing::debug!("next prompt suggestion stream closed before completion"); client_session.reset_websocket_session(); - return Ok(None); + return None; }; let event = match event { Ok(event) => event, @@ -170,7 +169,7 @@ pub(crate) async fn suggest_next_prompt( "next prompt suggestion stream failed while sampling" ); client_session.reset_websocket_session(); - return Ok(None); + return None; } }; match event { @@ -226,10 +225,10 @@ pub(crate) async fn suggest_next_prompt( 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); + return None; } if !session_is_idle_for_suggestion(sess).await { - return Ok(None); + return None; } let raw = completed_text.unwrap_or(streamed_text); @@ -242,7 +241,7 @@ pub(crate) async fn suggest_next_prompt( has_suggestion = suggestion.is_some(), "next prompt suggestion sampled" ); - Ok(suggestion) + suggestion } impl HistorySnapshot { @@ -276,35 +275,23 @@ async fn suggestion_prompt_fits_context_window(sess: &Session, turn_context: &Tu tracing::debug!("next prompt suggestion skipped without model context window"); return false; }; - let estimated_token_count = sess.get_estimated_token_count(turn_context).await; - if suggestion_prompt_has_headroom(estimated_token_count, Some(model_context_window)) { - return true; + if let Some(estimated_token_count) = sess.get_estimated_token_count(turn_context).await + && !suggestion_prompt_has_headroom(estimated_token_count, model_context_window) + { + let suggestion_prompt_limit = + model_context_window.saturating_sub(NEXT_PROMPT_SUGGESTION_TOKEN_HEADROOM); + tracing::debug!( + estimated_token_count, + model_context_window, + suggestion_prompt_limit, + "next prompt suggestion skipped near context window" + ); + return false; } - let Some(estimated_token_count) = estimated_token_count else { - return true; - }; - let suggestion_prompt_limit = - model_context_window.saturating_sub(NEXT_PROMPT_SUGGESTION_TOKEN_HEADROOM); - - tracing::debug!( - estimated_token_count, - model_context_window, - suggestion_prompt_limit, - "next prompt suggestion skipped near context window" - ); - false + true } -fn suggestion_prompt_has_headroom( - estimated_token_count: Option, - model_context_window: Option, -) -> bool { - let Some(model_context_window) = model_context_window else { - return false; - }; - let Some(estimated_token_count) = estimated_token_count else { - return true; - }; +fn suggestion_prompt_has_headroom(estimated_token_count: i64, model_context_window: i64) -> bool { estimated_token_count < model_context_window.saturating_sub(NEXT_PROMPT_SUGGESTION_TOKEN_HEADROOM) } @@ -350,6 +337,7 @@ fn has_unpaired_tool_flow(items: &[ResponseItem]) -> bool { let mut custom_tool_outputs = HashSet::new(); let mut tool_search_calls = HashSet::new(); let mut tool_search_outputs = HashSet::new(); + let mut client_tool_search_outputs = HashSet::new(); for item in items { match item { @@ -365,12 +353,15 @@ fn has_unpaired_tool_flow(items: &[ResponseItem]) -> bool { } => { tool_search_calls.insert(call_id.clone()); } - ResponseItem::ToolSearchOutput { execution, .. } if execution == "server" => {} ResponseItem::ToolSearchOutput { call_id: Some(call_id), + execution, .. } => { tool_search_outputs.insert(call_id.clone()); + if execution != "server" { + client_tool_search_outputs.insert(call_id.clone()); + } } ResponseItem::CustomToolCall { call_id, .. } => { custom_tool_calls.insert(call_id.clone()); @@ -400,7 +391,8 @@ fn has_unpaired_tool_flow(items: &[ResponseItem]) -> bool { function_calls != function_outputs || custom_tool_calls != custom_tool_outputs - || tool_search_calls != tool_search_outputs + || !tool_search_calls.is_subset(&tool_search_outputs) + || !client_tool_search_outputs.is_subset(&tool_search_calls) } /// Selects the fastest supported reasoning effort for an ephemeral suggestion sample. diff --git a/codex-rs/core/src/next_prompt_suggestion_tests.rs b/codex-rs/core/src/next_prompt_suggestion_tests.rs index 196e6cfd7b..907d3e3add 100644 --- a/codex-rs/core/src/next_prompt_suggestion_tests.rs +++ b/codex-rs/core/src/next_prompt_suggestion_tests.rs @@ -12,47 +12,21 @@ use codex_utils_output_truncation::TruncationPolicy; use pretty_assertions::assert_eq; #[test] -fn filter_keeps_specific_prompt() { - assert_eq!( - filter_next_prompt_suggestion("run the tests"), - Some("run the tests".to_string()) - ); -} - -#[test] -fn filter_keeps_prompt_with_edge_whitespace() { - assert_eq!( - filter_next_prompt_suggestion(" run the tests\n"), - Some("run the tests".to_string()) - ); -} - -#[test] -fn filter_keeps_allowed_single_word_prompt() { - assert_eq!( - filter_next_prompt_suggestion("commit"), - Some("commit".to_string()) - ); -} - -#[test] -fn filter_keeps_code_identifier_prompt() { - assert_eq!( - filter_next_prompt_suggestion("set CODEX_HOME"), - Some("set CODEX_HOME".to_string()) - ); -} - -#[test] -fn filter_keeps_dotted_file_prompt() { - assert_eq!( - filter_next_prompt_suggestion("update Cargo.toml"), - Some("update Cargo.toml".to_string()) - ); - assert_eq!( - filter_next_prompt_suggestion("open app-server/README.md"), - Some("open app-server/README.md".to_string()) - ); +fn filter_keeps_valid_prompts() { + for (suggestion, expected) in [ + ("run the tests", "run the tests"), + (" run the tests\n", "run the tests"), + ("commit", "commit"), + ("set CODEX_HOME", "set CODEX_HOME"), + ("update Cargo.toml", "update Cargo.toml"), + ("open app-server/README.md", "open app-server/README.md"), + ] { + assert_eq!( + filter_next_prompt_suggestion(suggestion), + Some(expected.to_string()), + "expected {suggestion:?} to be retained" + ); + } } #[test] @@ -99,40 +73,24 @@ fn history_boundary_requires_final_assistant_message() { assert!(history_ends_at_assistant_response(std::slice::from_ref( &assistant ))); - assert!(!history_ends_at_assistant_response(&[assistant, user])); -} - -#[test] -fn history_boundary_rejects_tool_output_tail() { - assert!(!history_ends_at_assistant_response(&[ - ResponseItem::Message { - id: None, - role: "assistant".to_string(), - content: vec![ContentItem::OutputText { - text: "calling tool".to_string(), - }], - phase: None, - }, + for tail in [ + user, ResponseItem::FunctionCallOutput { call_id: "call-1".to_string(), output: FunctionCallOutputPayload::from_text("done".to_string()), }, - ])); -} - -#[test] -fn suggestion_prompt_skips_without_context_window() { - assert!(!suggestion_prompt_has_headroom( - /*estimated_token_count*/ Some(10), - /*model_context_window*/ None, - )); + ] { + assert!(!history_ends_at_assistant_response(&[ + assistant.clone(), + tail, + ])); + } } #[test] fn suggestion_prompt_skips_near_context_window() { assert!(!suggestion_prompt_has_headroom( - /*estimated_token_count*/ Some(127_100), - /*model_context_window*/ Some(128_000) + /*estimated_token_count*/ 127_100, /*model_context_window*/ 128_000 )); } @@ -175,6 +133,25 @@ fn server_tool_search_output_without_call_is_allowed() { }])); } +#[test] +fn completed_server_tool_search_flow_is_allowed() { + assert!(!has_unpaired_tool_flow(&[ + ResponseItem::ToolSearchCall { + id: None, + call_id: Some("call-1".to_string()), + status: Some("completed".to_string()), + execution: "server".to_string(), + arguments: serde_json::json!({"query": "tool"}), + }, + ResponseItem::ToolSearchOutput { + call_id: Some("call-1".to_string()), + status: "completed".to_string(), + execution: "server".to_string(), + tools: Vec::new(), + }, + ])); +} + #[test] fn client_tool_search_output_without_call_is_suppressed() { assert!(has_unpaired_tool_flow(&[ResponseItem::ToolSearchOutput { diff --git a/codex-rs/core/src/session/turn_context.rs b/codex-rs/core/src/session/turn_context.rs index 106d8c9efe..a057675088 100644 --- a/codex-rs/core/src/session/turn_context.rs +++ b/codex-rs/core/src/session/turn_context.rs @@ -820,56 +820,6 @@ impl Session { turn_context } - async fn new_lightweight_turn_from_configuration( - &self, - sub_id: String, - session_configuration: SessionConfiguration, - turn_environments: ResolvedTurnEnvironments, - ) -> Arc { - let primary_turn_environment = turn_environments.primary(); - let cwd = primary_turn_environment - .map(|turn_environment| turn_environment.cwd.clone()) - .unwrap_or_else(|| session_configuration.cwd.clone()); - let per_turn_config = Self::build_per_turn_config(&session_configuration, cwd.clone()); - let model_info = self - .services - .models_manager - .get_model_info( - session_configuration.collaboration_mode.model(), - &per_turn_config.to_models_manager_config(), - ) - .await; - Arc::new(Self::make_turn_context( - self.thread_id(), - self.session_id(), - Some(Arc::clone(&self.services.auth_manager)), - &self.services.session_telemetry, - session_configuration.provider.clone(), - &session_configuration, - self.services.user_shell.as_ref(), - self.services.shell_zsh_path.as_ref(), - self.services.main_execve_wrapper_exe.as_ref(), - per_turn_config, - model_info, - &self.services.models_manager, - self.services - .network_proxy - .load_full() - .as_ref() - .and_then(|started_proxy| { - Self::managed_network_proxy_active_for_permission_profile( - &session_configuration.permission_profile(), - ) - .then(|| started_proxy.proxy()) - }), - turn_environments, - cwd, - sub_id, - Arc::new(SkillLoadOutcome::default()), - /*goal_tools_supported*/ false, - )) - } - pub(crate) async fn maybe_emit_unknown_model_warning_for_turn(&self, tc: &TurnContext) { if tc.model_info.used_fallback_model_metadata { self.send_event( @@ -915,12 +865,48 @@ impl Session { } }; - self.new_lightweight_turn_from_configuration( - self.next_internal_sub_id(), - session_configuration, + let primary_turn_environment = turn_environments.primary(); + let cwd = primary_turn_environment + .map(|turn_environment| turn_environment.cwd.clone()) + .unwrap_or_else(|| session_configuration.cwd.clone()); + let per_turn_config = Self::build_per_turn_config(&session_configuration, cwd.clone()); + let model_info = self + .services + .models_manager + .get_model_info( + session_configuration.collaboration_mode.model(), + &per_turn_config.to_models_manager_config(), + ) + .await; + Arc::new(Self::make_turn_context( + self.thread_id(), + self.session_id(), + Some(Arc::clone(&self.services.auth_manager)), + &self.services.session_telemetry, + session_configuration.provider.clone(), + &session_configuration, + self.services.user_shell.as_ref(), + self.services.shell_zsh_path.as_ref(), + self.services.main_execve_wrapper_exe.as_ref(), + per_turn_config, + model_info, + &self.services.models_manager, + self.services + .network_proxy + .load_full() + .as_ref() + .and_then(|started_proxy| { + Self::managed_network_proxy_active_for_permission_profile( + &session_configuration.permission_profile(), + ) + .then(|| started_proxy.proxy()) + }), turn_environments, - ) - .await + cwd, + self.next_internal_sub_id(), + Arc::new(SkillLoadOutcome::default()), + /*goal_tools_supported*/ false, + )) } pub(crate) async fn new_default_turn_with_sub_id(&self, sub_id: String) -> Arc {