openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-09-05 15:18:41 +00:00
Code Issues Packages Projects Releases Wiki Activity

refactor(core): simplify next prompt suggestion flow

Browse Source
This commit is contained in:
Felipe Coury
2026-05-23 16:36:56 -03:00
parent 7ac197318e
commit b9672997ab
4 changed files with 127 additions and 172 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
codex-rs/core/src/codex_thread.rs
View File

@@ -154,11 +154,11 @@ impl CodexThread {
&self,
cancellation_token: CancellationToken,
) -> CodexResult<Option<String>> {
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.

90
codex-rs/core/src/next_prompt_suggestion.rs
View File

@@ -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<String>> {
) -> Option<String> {
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<i64>,
model_context_window: Option<i64>,
) -> 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.

109
codex-rs/core/src/next_prompt_suggestion_tests.rs
View File

@@ -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 {

96
codex-rs/core/src/session/turn_context.rs
View File

@@ -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<TurnContext> {
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<TurnContext> {