From 5f4a2b0fbb934dff2191fd2f0064a93ea954b106 Mon Sep 17 00:00:00 2001 From: Tuan-Lung Wang Date: Tue, 17 Mar 2026 15:44:10 -0700 Subject: [PATCH] Add fast draft completion popup --- codex-rs/tui/src/answer_interleave.rs | 80 +++-- codex-rs/tui/src/app.rs | 30 ++ codex-rs/tui/src/app_event.rs | 10 + codex-rs/tui/src/bottom_pane/chat_composer.rs | 305 ++++++++++++++++ .../src/bottom_pane/draft_completion_popup.rs | 165 +++++++++ codex-rs/tui/src/bottom_pane/mod.rs | 10 + ...mposer__tests__draft_completion_popup.snap | 13 + ...__tests__draft_completion_popup_error.snap | 13 + codex-rs/tui/src/chatwidget.rs | 24 +- codex-rs/tui/src/draft_completion.rs | 338 ++++++++++++++++++ codex-rs/tui/src/lib.rs | 1 + 11 files changed, 962 insertions(+), 27 deletions(-) create mode 100644 codex-rs/tui/src/bottom_pane/draft_completion_popup.rs create mode 100644 codex-rs/tui/src/bottom_pane/snapshots/codex_tui__bottom_pane__chat_composer__tests__draft_completion_popup.snap create mode 100644 codex-rs/tui/src/bottom_pane/snapshots/codex_tui__bottom_pane__chat_composer__tests__draft_completion_popup_error.snap create mode 100644 codex-rs/tui/src/draft_completion.rs diff --git a/codex-rs/tui/src/answer_interleave.rs b/codex-rs/tui/src/answer_interleave.rs index f5299bbb2b..c9c9f2acf1 100644 --- a/codex-rs/tui/src/answer_interleave.rs +++ b/codex-rs/tui/src/answer_interleave.rs @@ -12,11 +12,9 @@ use codex_core::default_client::originator; use codex_core::models_manager::manager::ModelsManager; use codex_otel::SessionTelemetry; use codex_protocol::ThreadId; -use codex_protocol::config_types::ReasoningSummary as ReasoningSummaryConfig; use codex_protocol::models::BaseInstructions; use codex_protocol::models::ContentItem; use codex_protocol::models::ResponseItem; -use codex_protocol::openai_models::ReasoningEffort as ReasoningEffortConfig; use codex_protocol::protocol::SessionSource; use ratatui::prelude::Stylize; use ratatui::text::Line; @@ -246,9 +244,9 @@ async fn interleave_answer( &prompt, &model_info, &telemetry, - Some(ReasoningEffortConfig::None), - ReasoningSummaryConfig::None, - Some(codex_protocol::config_types::ServiceTier::Fast), + None, + model_info.default_reasoning_summary, + None, None, ) .await @@ -412,8 +410,9 @@ fn parse_simple_list(text: &str) -> Option { let mut started = false; let mut remainder_start = None; let lines: Vec<&str> = text.lines().collect(); - - for (idx, line) in lines.iter().enumerate() { + let mut idx = 0; + while idx < lines.len() { + let line = lines[idx]; let trimmed = line.trim(); if trimmed.is_empty() { if started { @@ -423,31 +422,39 @@ fn parse_simple_list(text: &str) -> Option { .map(|offset| idx + 1 + offset); break; } + idx += 1; continue; } - let Some((line_kind, marker, item_text)) = parse_list_item_line(trimmed) else { - if started { - remainder_start = Some(idx); - break; + if let Some((line_kind, marker, item_text)) = parse_list_item_line(trimmed) { + match kind { + Some(existing_kind) if existing_kind != line_kind => { + remainder_start = Some(idx); + break; + } + None => kind = Some(line_kind), + Some(_) => {} } - continue; - }; - match kind { - Some(existing_kind) if existing_kind != line_kind => { - remainder_start = Some(idx); - break; - } - None => kind = Some(line_kind), - Some(_) => {} + started = true; + items.push(ListItem { + marker, + text: item_text.to_string(), + }); + idx += 1; + continue; } - started = true; - items.push(ListItem { - marker, - text: item_text.to_string(), - }); + if started { + if let Some(last_item) = items.last_mut() { + last_item.text.push('\n'); + last_item.text.push_str(trimmed); + } + idx += 1; + continue; + } + + idx += 1; } if items.is_empty() { @@ -575,4 +582,27 @@ mod tests { ) ); } + + #[test] + fn heuristic_interleave_keeps_continuation_lines_with_numbered_answers() { + let result = heuristic_interleave_answer( + "1. how are you doing?\n2. what's the biggest city in the world?\n3. who's the ceo of openai?", + "1. I'm doing fine.\n2. Jakarta is the largest city by population.\nSource: https://example.com/jakarta\n3. Sam Altman is the CEO of OpenAI.\nSource: https://example.com/openai", + ); + + assert_eq!( + result, + Some( + "1. how are you doing?\n\ + A: I'm doing fine.\n\ + 2. what's the biggest city in the world?\n\ + A: Jakarta is the largest city by population.\n\ + Source: https://example.com/jakarta\n\ + 3. who's the ceo of openai?\n\ + A: Sam Altman is the CEO of OpenAI.\n\ + Source: https://example.com/openai" + .to_string() + ) + ); + } } diff --git a/codex-rs/tui/src/app.rs b/codex-rs/tui/src/app.rs index a405347128..dce60da9a0 100644 --- a/codex-rs/tui/src/app.rs +++ b/codex-rs/tui/src/app.rs @@ -17,6 +17,7 @@ use crate::chatwidget::ExternalEditorState; use crate::chatwidget::ThreadInputState; use crate::cwd_prompt::CwdPromptAction; use crate::diff_render::DiffSummary; +use crate::draft_completion::DraftCompletionManager; use crate::exec_command::strip_bash_lc_and_escape; use crate::external_editor; use crate::file_search::FileSearchManager; @@ -676,6 +677,7 @@ pub(crate) struct App { runtime_sandbox_policy_override: Option, pub(crate) file_search: FileSearchManager, + draft_completion: DraftCompletionManager, answer_interleave: AnswerInterleaveManager, pub(crate) transcript_cells: Vec>, @@ -2143,6 +2145,12 @@ impl App { .maybe_prompt_windows_sandbox_enable(should_prompt_windows_sandbox_nux_at_startup); let file_search = FileSearchManager::new(config.cwd.clone(), app_event_tx.clone()); + let draft_completion = DraftCompletionManager::new( + auth_manager.clone(), + thread_manager.get_models_manager(), + app_event_tx.clone(), + config.clone(), + ); let answer_interleave = AnswerInterleaveManager::new( auth_manager.clone(), thread_manager.get_models_manager(), @@ -2165,6 +2173,7 @@ impl App { runtime_approval_policy_override: None, runtime_sandbox_policy_override: None, file_search, + draft_completion, answer_interleave, enhanced_keys_supported, transcript_cells: Vec::new(), @@ -2712,6 +2721,13 @@ impl App { AppEvent::FileSearchResult { query, matches } => { self.chat_widget.apply_file_search_result(query, matches); } + AppEvent::StartDraftCompletion(request) => { + self.draft_completion.start_request(request); + } + AppEvent::DraftCompletionResult { request_id, result } => { + self.chat_widget + .on_draft_completion_result(request_id, result); + } AppEvent::StartAnswerInterleave(request) => { self.answer_interleave.start_request(request); } @@ -6358,6 +6374,12 @@ guardian_approval = true CodexAuth::from_api_key("Test API Key"), ); let file_search = FileSearchManager::new(config.cwd.clone(), app_event_tx.clone()); + let draft_completion = DraftCompletionManager::new( + auth_manager.clone(), + server.get_models_manager(), + app_event_tx.clone(), + config.clone(), + ); let answer_interleave = AnswerInterleaveManager::new( auth_manager.clone(), server.get_models_manager(), @@ -6380,6 +6402,7 @@ guardian_approval = true runtime_approval_policy_override: None, runtime_sandbox_policy_override: None, file_search, + draft_completion, answer_interleave, transcript_cells: Vec::new(), overlay: None, @@ -6424,6 +6447,12 @@ guardian_approval = true CodexAuth::from_api_key("Test API Key"), ); let file_search = FileSearchManager::new(config.cwd.clone(), app_event_tx.clone()); + let draft_completion = DraftCompletionManager::new( + auth_manager.clone(), + server.get_models_manager(), + app_event_tx.clone(), + config.clone(), + ); let answer_interleave = AnswerInterleaveManager::new( auth_manager.clone(), server.get_models_manager(), @@ -6447,6 +6476,7 @@ guardian_approval = true runtime_approval_policy_override: None, runtime_sandbox_policy_override: None, file_search, + draft_completion, answer_interleave, transcript_cells: Vec::new(), overlay: None, diff --git a/codex-rs/tui/src/app_event.rs b/codex-rs/tui/src/app_event.rs index e903a7fd9c..4b4d3f7033 100644 --- a/codex-rs/tui/src/app_event.rs +++ b/codex-rs/tui/src/app_event.rs @@ -21,6 +21,7 @@ use codex_utils_approval_presets::ApprovalPreset; use crate::answer_interleave::AnswerInterleaveRequest; use crate::bottom_pane::ApprovalRequest; use crate::bottom_pane::StatusLineItem; +use crate::draft_completion::DraftCompletionRequest; use crate::history_cell::HistoryCell; use codex_core::config::types::ApprovalsReviewer; @@ -139,6 +140,15 @@ pub(crate) enum AppEvent { result: Result, }, + /// Kick off an asynchronous fast-model draft completion request. + StartDraftCompletion(DraftCompletionRequest), + + /// Result of a completed draft completion request. + DraftCompletionResult { + request_id: u64, + result: Result, String>, + }, + /// Result of refreshing rate limits RateLimitSnapshotFetched(RateLimitSnapshot), diff --git a/codex-rs/tui/src/bottom_pane/chat_composer.rs b/codex-rs/tui/src/bottom_pane/chat_composer.rs index c3b73659b7..08a2b29b5a 100644 --- a/codex-rs/tui/src/bottom_pane/chat_composer.rs +++ b/codex-rs/tui/src/bottom_pane/chat_composer.rs @@ -153,6 +153,7 @@ use super::chat_composer_history::HistoryEntry; use super::command_popup::CommandItem; use super::command_popup::CommandPopup; use super::command_popup::CommandPopupFlags; +use super::draft_completion_popup::DraftCompletionPopup; use super::file_search_popup::FileSearchPopup; use super::footer::CollaborationModeIndicator; use super::footer::FooterMode; @@ -258,6 +259,10 @@ pub enum InputResult { }, Command(SlashCommand), CommandWithArgs(SlashCommand, String, Vec), + RequestDraftCompletion { + request_id: u64, + draft: String, + }, None, } @@ -404,6 +409,7 @@ pub(crate) struct ChatComposer { collaboration_mode_indicator: Option, connectors_enabled: bool, fast_command_enabled: bool, + next_draft_completion_request_id: u64, personality_command_enabled: bool, realtime_conversation_enabled: bool, audio_device_selection_enabled: bool, @@ -430,6 +436,7 @@ struct ComposerMentionBinding { enum ActivePopup { None, Command(CommandPopup), + Completion(DraftCompletionPopup), File(FileSearchPopup), Skill(SkillPopup), } @@ -526,6 +533,7 @@ impl ChatComposer { collaboration_mode_indicator: None, connectors_enabled: false, fast_command_enabled: false, + next_draft_completion_request_id: 0, personality_command_enabled: false, realtime_conversation_enabled: false, audio_device_selection_enabled: false, @@ -666,6 +674,9 @@ impl ChatComposer { ActivePopup::Command(popup) => { Constraint::Max(popup.calculate_required_height(area.width)) } + ActivePopup::Completion(popup) => { + Constraint::Max(popup.calculate_required_height(area.width)) + } ActivePopup::File(popup) => Constraint::Max(popup.calculate_required_height()), ActivePopup::Skill(popup) => { Constraint::Max(popup.calculate_required_height(area.width)) @@ -1239,6 +1250,23 @@ impl ChatComposer { } } + pub(crate) fn on_draft_completion_result( + &mut self, + request_id: u64, + result: Result, String>, + ) { + let ActivePopup::Completion(popup) = &mut self.active_popup else { + return; + }; + if popup.request_id() != request_id { + return; + } + match result { + Ok(suggestions) => popup.set_suggestions(suggestions), + Err(err) => popup.set_error_message(err), + } + } + /// Show the transient "press again to quit" hint for `key`. /// /// The owner (`BottomPane`/`ChatWidget`) is responsible for scheduling a @@ -1332,6 +1360,9 @@ impl ChatComposer { let result = match &mut self.active_popup { ActivePopup::Command(_) => self.handle_key_event_with_slash_popup(key_event), + ActivePopup::Completion(_) => { + self.handle_key_event_with_draft_completion_popup(key_event) + } ActivePopup::File(_) => self.handle_key_event_with_file_popup(key_event), ActivePopup::Skill(_) => self.handle_key_event_with_skill_popup(key_event), ActivePopup::None => self.handle_key_event_without_popup(key_event), @@ -1348,6 +1379,9 @@ impl ChatComposer { /// Handle key event when the slash-command popup is visible. fn handle_key_event_with_slash_popup(&mut self, key_event: KeyEvent) -> (InputResult, bool) { + if is_draft_completion_shortcut(key_event) { + return (InputResult::None, true); + } if self.handle_shortcut_overlay_key(&key_event) { return (InputResult::None, true); } @@ -1550,6 +1584,70 @@ impl ChatComposer { p } + fn handle_key_event_with_draft_completion_popup( + &mut self, + key_event: KeyEvent, + ) -> (InputResult, bool) { + if is_draft_completion_shortcut(key_event) { + return (InputResult::None, true); + } + if self.handle_shortcut_overlay_key(&key_event) { + return (InputResult::None, true); + } + self.footer_mode = reset_mode_after_activity(self.footer_mode); + let ActivePopup::Completion(popup) = &mut self.active_popup else { + unreachable!(); + }; + + match key_event { + KeyEvent { + code: KeyCode::Up, .. + } + | KeyEvent { + code: KeyCode::Char('p'), + modifiers: KeyModifiers::CONTROL, + .. + } => { + popup.move_up(); + (InputResult::None, true) + } + KeyEvent { + code: KeyCode::Down, + .. + } + | KeyEvent { + code: KeyCode::Char('n'), + modifiers: KeyModifiers::CONTROL, + .. + } => { + popup.move_down(); + (InputResult::None, true) + } + KeyEvent { + code: KeyCode::Esc, .. + } => { + self.active_popup = ActivePopup::None; + (InputResult::None, true) + } + KeyEvent { + code: KeyCode::Enter, + modifiers: KeyModifiers::NONE, + .. + } => { + let selected = popup.selected_suggestion().map(str::to_string); + self.active_popup = ActivePopup::None; + if let Some(suggestion) = selected { + self.textarea.insert_str(&suggestion); + } + (InputResult::None, true) + } + input => { + self.active_popup = ActivePopup::None; + self.handle_key_event_without_popup(input) + } + } + } + /// Handle non-ASCII character input (often IME) while still supporting paste-burst detection. /// /// This handler exists because non-ASCII input often comes from IMEs, where characters can @@ -1636,6 +1734,9 @@ impl ChatComposer { /// Handle key events when file search popup is visible. fn handle_key_event_with_file_popup(&mut self, key_event: KeyEvent) -> (InputResult, bool) { + if is_draft_completion_shortcut(key_event) { + return (InputResult::None, true); + } if self.handle_shortcut_overlay_key(&key_event) { return (InputResult::None, true); } @@ -1759,6 +1860,9 @@ impl ChatComposer { } fn handle_key_event_with_skill_popup(&mut self, key_event: KeyEvent) -> (InputResult, bool) { + if is_draft_completion_shortcut(key_event) { + return (InputResult::None, true); + } if self.handle_shortcut_overlay_key(&key_event) { return (InputResult::None, true); } @@ -2778,6 +2882,25 @@ impl ChatComposer { } self.handle_input_basic(key_event) } + KeyEvent { + code: KeyCode::Char('o'), + modifiers: KeyModifiers::CONTROL, + kind: KeyEventKind::Press, + .. + } if self.popups_enabled() + && self.textarea.cursor() == self.textarea.text().len() + && !self.is_empty() => + { + self.next_draft_completion_request_id = + self.next_draft_completion_request_id.wrapping_add(1); + let request_id = self.next_draft_completion_request_id; + let draft = self.current_text_with_pending(); + self.active_popup = ActivePopup::Completion(DraftCompletionPopup::new(request_id)); + ( + InputResult::RequestDraftCompletion { request_id, draft }, + true, + ) + } KeyEvent { code: KeyCode::Tab, modifiers: KeyModifiers::NONE, @@ -3237,6 +3360,9 @@ impl ChatComposer { } pub(crate) fn sync_popups(&mut self) { + if matches!(self.active_popup, ActivePopup::Completion(_)) { + return; + } self.sync_slash_command_elements(); if !self.popups_enabled() { self.active_popup = ActivePopup::None; @@ -4154,6 +4280,7 @@ impl Renderable for ChatComposer { + match &self.active_popup { ActivePopup::None => footer_total_height, ActivePopup::Command(c) => c.calculate_required_height(width), + ActivePopup::Completion(c) => c.calculate_required_height(width), ActivePopup::File(c) => c.calculate_required_height(), ActivePopup::Skill(c) => c.calculate_required_height(width), } @@ -4172,6 +4299,9 @@ impl ChatComposer { ActivePopup::Command(popup) => { popup.render_ref(popup_rect, buf); } + ActivePopup::Completion(popup) => { + popup.render_ref(popup_rect, buf); + } ActivePopup::File(popup) => { popup.render_ref(popup_rect, buf); } @@ -4482,6 +4612,18 @@ fn prompt_selection_action( } } +fn is_draft_completion_shortcut(key_event: KeyEvent) -> bool { + matches!( + key_event, + KeyEvent { + code: KeyCode::Char('o'), + modifiers: KeyModifiers::CONTROL, + kind: KeyEventKind::Press, + .. + } + ) +} + impl Drop for ChatComposer { fn drop(&mut self) { // Stop any running spinner tasks. @@ -6537,6 +6679,160 @@ mod tests { } } + #[test] + fn draft_completion_popup_snapshot() { + use crossterm::event::KeyCode; + use crossterm::event::KeyEvent; + use crossterm::event::KeyModifiers; + + snapshot_composer_state_with_width("draft_completion_popup", 72, false, |composer| { + composer.set_text_content("hell".to_string(), Vec::new(), Vec::new()); + composer.move_cursor_to_end(); + let (result, _needs_redraw) = + composer.handle_key_event(KeyEvent::new(KeyCode::Char('o'), KeyModifiers::CONTROL)); + let InputResult::RequestDraftCompletion { request_id, .. } = result else { + panic!("expected draft completion request"); + }; + composer.on_draft_completion_result( + request_id, + Ok(vec![ + "o world".to_string(), + " there".to_string(), + " everyone".to_string(), + ]), + ); + }); + } + + #[test] + fn draft_completion_popup_error_snapshot() { + use crossterm::event::KeyCode; + use crossterm::event::KeyEvent; + use crossterm::event::KeyModifiers; + + snapshot_composer_state_with_width("draft_completion_popup_error", 72, false, |composer| { + composer.set_text_content("hell".to_string(), Vec::new(), Vec::new()); + composer.move_cursor_to_end(); + let (result, _needs_redraw) = + composer.handle_key_event(KeyEvent::new(KeyCode::Char('o'), KeyModifiers::CONTROL)); + let InputResult::RequestDraftCompletion { request_id, .. } = result else { + panic!("expected draft completion request"); + }; + composer.on_draft_completion_result( + request_id, + Err("draft completion timed out".to_string()), + ); + }); + } + + #[test] + fn draft_completion_enter_accepts_selected_suggestion() { + use crossterm::event::KeyCode; + use crossterm::event::KeyEvent; + use crossterm::event::KeyModifiers; + + let (tx, _rx) = unbounded_channel::(); + let sender = AppEventSender::new(tx); + let mut composer = ChatComposer::new( + true, + sender, + false, + "Ask Codex to do anything".to_string(), + false, + ); + composer.set_text_content("hell".to_string(), Vec::new(), Vec::new()); + composer.move_cursor_to_end(); + + let (result, _needs_redraw) = + composer.handle_key_event(KeyEvent::new(KeyCode::Char('o'), KeyModifiers::CONTROL)); + let InputResult::RequestDraftCompletion { request_id, draft } = result else { + panic!("expected draft completion request"); + }; + assert_eq!(draft, "hell"); + + composer.on_draft_completion_result( + request_id, + Ok(vec![ + "o world".to_string(), + " there".to_string(), + " everyone".to_string(), + ]), + ); + let _ = composer.handle_key_event(KeyEvent::new(KeyCode::Down, KeyModifiers::NONE)); + let (result, _needs_redraw) = + composer.handle_key_event(KeyEvent::new(KeyCode::Enter, KeyModifiers::NONE)); + + assert_eq!(result, InputResult::None); + assert_eq!(composer.current_text(), "hell there"); + assert!(matches!(composer.active_popup, ActivePopup::None)); + } + + #[test] + fn draft_completion_popup_dismisses_on_edit() { + use crossterm::event::KeyCode; + use crossterm::event::KeyEvent; + use crossterm::event::KeyModifiers; + + let (tx, _rx) = unbounded_channel::(); + let sender = AppEventSender::new(tx); + let mut composer = ChatComposer::new( + true, + sender, + false, + "Ask Codex to do anything".to_string(), + false, + ); + composer.set_text_content("hell".to_string(), Vec::new(), Vec::new()); + composer.move_cursor_to_end(); + + let (result, _needs_redraw) = + composer.handle_key_event(KeyEvent::new(KeyCode::Char('o'), KeyModifiers::CONTROL)); + let InputResult::RequestDraftCompletion { request_id, .. } = result else { + panic!("expected draft completion request"); + }; + composer.on_draft_completion_result( + request_id, + Ok(vec![ + "o world".to_string(), + " there".to_string(), + " everyone".to_string(), + ]), + ); + + let (result, _needs_redraw) = + composer.handle_key_event(KeyEvent::new(KeyCode::Char('!'), KeyModifiers::NONE)); + flush_after_paste_burst(&mut composer); + + assert_eq!(result, InputResult::None); + assert_eq!(composer.current_text(), "hell!"); + assert!(matches!(composer.active_popup, ActivePopup::None)); + } + + #[test] + fn draft_completion_shortcut_is_ignored_while_other_popup_is_active() { + use crossterm::event::KeyCode; + use crossterm::event::KeyEvent; + use crossterm::event::KeyModifiers; + + let (tx, _rx) = unbounded_channel::(); + let sender = AppEventSender::new(tx); + let mut composer = ChatComposer::new( + true, + sender, + false, + "Ask Codex to do anything".to_string(), + false, + ); + type_chars_humanlike(&mut composer, &['/', 'm', 'o']); + assert!(matches!(composer.active_popup, ActivePopup::Command(_))); + + let (result, _needs_redraw) = + composer.handle_key_event(KeyEvent::new(KeyCode::Char('o'), KeyModifiers::CONTROL)); + + assert_eq!(result, InputResult::None); + assert!(matches!(composer.active_popup, ActivePopup::Command(_))); + } + fn flush_after_paste_burst(composer: &mut ChatComposer) -> bool { std::thread::sleep(PasteBurst::recommended_active_flush_delay()); composer.flush_paste_burst_if_due() @@ -6600,6 +6896,9 @@ mod tests { InputResult::Queued { .. } => { panic!("expected command dispatch, but composer queued literal text") } + InputResult::RequestDraftCompletion { .. } => { + panic!("expected command dispatch, but composer requested draft completion") + } InputResult::None => panic!("expected Command result for '/init'"), } assert!(composer.textarea.is_empty(), "composer should be cleared"); @@ -7005,6 +7304,9 @@ mod tests { InputResult::Queued { .. } => { panic!("expected command dispatch after Tab completion, got literal queue") } + InputResult::RequestDraftCompletion { .. } => { + panic!("expected command dispatch after Tab completion, got draft completion") + } InputResult::None => panic!("expected Command result for '/diff'"), } assert!(composer.textarea.is_empty()); @@ -7172,6 +7474,9 @@ mod tests { InputResult::Queued { .. } => { panic!("expected command dispatch, but composer queued literal text") } + InputResult::RequestDraftCompletion { .. } => { + panic!("expected command dispatch, but composer requested draft completion") + } InputResult::None => panic!("expected Command result for '/mention'"), } assert!(composer.textarea.is_empty(), "composer should be cleared"); diff --git a/codex-rs/tui/src/bottom_pane/draft_completion_popup.rs b/codex-rs/tui/src/bottom_pane/draft_completion_popup.rs new file mode 100644 index 0000000000..853bf638fb --- /dev/null +++ b/codex-rs/tui/src/bottom_pane/draft_completion_popup.rs @@ -0,0 +1,165 @@ +use ratatui::buffer::Buffer; +use ratatui::layout::Rect; +use ratatui::prelude::Stylize; +use ratatui::text::Line; +use ratatui::text::Span; +use ratatui::widgets::Widget; +use ratatui::widgets::WidgetRef; +use unicode_width::UnicodeWidthStr; + +use crate::render::Insets; +use crate::render::RectExt; + +use super::popup_consts::MAX_POPUP_ROWS; +use super::scroll_state::ScrollState; +use super::selection_popup_common::GenericDisplayRow; +use super::selection_popup_common::measure_rows_height; +use super::selection_popup_common::render_rows; + +pub(crate) struct DraftCompletionPopup { + request_id: u64, + waiting: bool, + suggestions: Vec, + error_message: Option, + state: ScrollState, +} + +impl DraftCompletionPopup { + pub(crate) fn new(request_id: u64) -> Self { + let mut state = ScrollState::new(); + state.selected_idx = Some(0); + Self { + request_id, + waiting: true, + suggestions: Vec::new(), + error_message: None, + state, + } + } + + pub(crate) fn request_id(&self) -> u64 { + self.request_id + } + + pub(crate) fn set_suggestions(&mut self, suggestions: Vec) { + self.waiting = false; + self.suggestions = suggestions; + self.error_message = None; + let len = self.suggestions.len(); + self.state.clamp_selection(len); + self.state.ensure_visible(len, len.clamp(1, MAX_POPUP_ROWS)); + } + + pub(crate) fn set_error_message(&mut self, error_message: String) { + self.waiting = false; + self.suggestions.clear(); + self.error_message = Some(error_message); + self.state.selected_idx = None; + self.state.scroll_top = 0; + } + + pub(crate) fn move_up(&mut self) { + let len = self.suggestions.len(); + self.state.move_up_wrap(len); + self.state.ensure_visible(len, len.clamp(1, MAX_POPUP_ROWS)); + } + + pub(crate) fn move_down(&mut self) { + let len = self.suggestions.len(); + self.state.move_down_wrap(len); + self.state.ensure_visible(len, len.clamp(1, MAX_POPUP_ROWS)); + } + + pub(crate) fn selected_suggestion(&self) -> Option<&str> { + self.state + .selected_idx + .and_then(|idx| self.suggestions.get(idx)) + .map(String::as_str) + } + + pub(crate) fn calculate_required_height(&self, width: u16) -> u16 { + if let Some(message) = self.error_message.as_deref() { + let wrapped_height = + wrapped_message_lines(message, width.saturating_sub(2)).len() as u16; + return wrapped_height.max(1); + } + let rows = self.rows(); + measure_rows_height(&rows, &self.state, MAX_POPUP_ROWS, width.saturating_sub(2)) + } + + fn rows(&self) -> Vec { + self.suggestions + .iter() + .map(|suggestion| GenericDisplayRow { + name: suggestion.trim_start().to_string(), + name_prefix_spans: Vec::new(), + match_indices: None, + display_shortcut: None, + description: None, + category_tag: None, + wrap_indent: None, + is_disabled: false, + disabled_reason: None, + }) + .collect() + } +} + +impl WidgetRef for &DraftCompletionPopup { + fn render_ref(&self, area: Rect, buf: &mut Buffer) { + if let Some(message) = self.error_message.as_deref() { + render_wrapped_message(area, buf, message); + return; + } + + let rows = self.rows(); + let empty_message = if self.waiting { + "loading..." + } else { + "no suggestions" + }; + render_rows( + area.inset(Insets::tlbr(0, 2, 0, 0)), + buf, + &rows, + &self.state, + MAX_POPUP_ROWS, + empty_message, + ); + } +} + +fn render_wrapped_message(area: Rect, buf: &mut Buffer, message: &str) { + if area.height == 0 || area.width == 0 { + return; + } + + let lines = wrapped_message_lines(message, area.width); + for (idx, line) in lines.into_iter().take(area.height as usize).enumerate() { + Line::from(vec![Span::from(line).dim().italic()]).render( + Rect { + x: area.x, + y: area.y + idx as u16, + width: area.width, + height: 1, + }, + buf, + ); + } +} + +fn wrapped_message_lines(message: &str, width: u16) -> Vec { + let width = width.max(1) as usize; + textwrap::wrap(message, width) + .into_iter() + .map(|line| { + let text = line.into_owned(); + if UnicodeWidthStr::width(text.as_str()) > width { + let trimmed = text.trim().to_string(); + if trimmed.is_empty() { text } else { trimmed } + } else { + text + } + }) + .collect() +} diff --git a/codex-rs/tui/src/bottom_pane/mod.rs b/codex-rs/tui/src/bottom_pane/mod.rs index d15b03021e..df80395f6b 100644 --- a/codex-rs/tui/src/bottom_pane/mod.rs +++ b/codex-rs/tui/src/bottom_pane/mod.rs @@ -76,6 +76,7 @@ mod chat_composer; mod chat_composer_history; mod command_popup; pub mod custom_prompt_view; +mod draft_completion_popup; mod experimental_features_view; mod file_search_popup; mod footer; @@ -1099,6 +1100,15 @@ impl BottomPane { self.request_redraw(); } + pub(crate) fn on_draft_completion_result( + &mut self, + request_id: u64, + result: Result, String>, + ) { + self.composer.on_draft_completion_result(request_id, result); + self.request_redraw(); + } + pub(crate) fn attach_image(&mut self, path: PathBuf) { if self.view_stack.is_empty() { self.composer.attach_image(path); diff --git a/codex-rs/tui/src/bottom_pane/snapshots/codex_tui__bottom_pane__chat_composer__tests__draft_completion_popup.snap b/codex-rs/tui/src/bottom_pane/snapshots/codex_tui__bottom_pane__chat_composer__tests__draft_completion_popup.snap new file mode 100644 index 0000000000..75b91f933c --- /dev/null +++ b/codex-rs/tui/src/bottom_pane/snapshots/codex_tui__bottom_pane__chat_composer__tests__draft_completion_popup.snap @@ -0,0 +1,13 @@ +--- +source: tui/src/bottom_pane/chat_composer.rs +expression: terminal.backend() +--- +" " +"› hell " +" " +" " +" " +" " +" o world " +" there " +" everyone " diff --git a/codex-rs/tui/src/bottom_pane/snapshots/codex_tui__bottom_pane__chat_composer__tests__draft_completion_popup_error.snap b/codex-rs/tui/src/bottom_pane/snapshots/codex_tui__bottom_pane__chat_composer__tests__draft_completion_popup_error.snap new file mode 100644 index 0000000000..f32985e4b3 --- /dev/null +++ b/codex-rs/tui/src/bottom_pane/snapshots/codex_tui__bottom_pane__chat_composer__tests__draft_completion_popup_error.snap @@ -0,0 +1,13 @@ +--- +source: tui/src/bottom_pane/chat_composer.rs +expression: terminal.backend() +--- +" " +"› hell " +" " +" " +" " +" " +" " +" " +"draft completion timed out " diff --git a/codex-rs/tui/src/chatwidget.rs b/codex-rs/tui/src/chatwidget.rs index 14902ba714..c99cbaf535 100644 --- a/codex-rs/tui/src/chatwidget.rs +++ b/codex-rs/tui/src/chatwidget.rs @@ -47,6 +47,7 @@ use crate::audio_device::list_realtime_audio_device_names; use crate::bottom_pane::StatusLineItem; use crate::bottom_pane::StatusLinePreviewData; use crate::bottom_pane::StatusLineSetupView; +use crate::draft_completion::DraftCompletionRequest; use crate::status::RateLimitWindowDisplay; use crate::status::format_directory_display; use crate::status::format_tokens_compact; @@ -1594,7 +1595,7 @@ impl ChatWidget { .unwrap_or(AgentMessageStreamKind::Unknown) } - fn current_user_prompt_for_interleave(&self) -> Option { + fn last_user_prompt_text(&self) -> Option { self.last_rendered_user_message_event .as_ref() .map(|event| event.message.trim().to_string()) @@ -1671,7 +1672,7 @@ impl ChatWidget { } fn start_answer_interleave(&mut self, raw_answer: String) { - let Some(user_prompt) = self.current_user_prompt_for_interleave() else { + let Some(user_prompt) = self.last_user_prompt_text() else { self.finalize_live_answer_without_interleave(raw_answer); return; }; @@ -4450,6 +4451,16 @@ impl ChatWidget { InputResult::CommandWithArgs(cmd, args, text_elements) => { self.dispatch_command_with_args(cmd, args, text_elements); } + InputResult::RequestDraftCompletion { request_id, draft } => { + self.app_event_tx.send(AppEvent::StartDraftCompletion( + DraftCompletionRequest { + request_id, + draft, + last_user_turn: self.last_user_prompt_text(), + last_assistant_turn: self.last_copyable_output.clone(), + }, + )); + } InputResult::None => {} }, } @@ -8845,6 +8856,15 @@ impl ChatWidget { self.bottom_pane.on_file_search_result(query, matches); } + pub(crate) fn on_draft_completion_result( + &mut self, + request_id: u64, + result: Result, String>, + ) { + self.bottom_pane + .on_draft_completion_result(request_id, result); + } + /// Handles a Ctrl+C press at the chat-widget layer. /// /// The first press arms a time-bounded quit shortcut and shows a footer hint via the bottom diff --git a/codex-rs/tui/src/draft_completion.rs b/codex-rs/tui/src/draft_completion.rs new file mode 100644 index 0000000000..1841888524 --- /dev/null +++ b/codex-rs/tui/src/draft_completion.rs @@ -0,0 +1,338 @@ +use std::sync::Arc; +use std::time::Duration; + +use codex_core::AuthManager; +use codex_core::ModelClient; +use codex_core::Prompt; +use codex_core::ResponseEvent; +use codex_core::config::Config; +use codex_core::default_client::originator; +use codex_core::models_manager::manager::ModelsManager; +use codex_otel::SessionTelemetry; +use codex_protocol::ThreadId; +use codex_protocol::models::BaseInstructions; +use codex_protocol::models::ContentItem; +use codex_protocol::models::ResponseItem; +use codex_protocol::protocol::SessionSource; +use tokio::time::timeout; +use tokio_stream::StreamExt; + +use crate::app_event::AppEvent; +use crate::app_event_sender::AppEventSender; + +const DRAFT_COMPLETION_MODEL: &str = "galapagos-nano-alpha"; +const DRAFT_COMPLETION_TIMEOUT: Duration = Duration::from_secs(6); +const DRAFT_COMPLETION_AGENT_LABEL: &str = "tui_draft_completion"; +const MAX_SUGGESTIONS: usize = 3; + +#[derive(Debug, Clone)] +pub(crate) struct DraftCompletionRequest { + pub(crate) request_id: u64, + pub(crate) draft: String, + pub(crate) last_user_turn: Option, + pub(crate) last_assistant_turn: Option, +} + +pub(crate) struct DraftCompletionManager { + auth_manager: Arc, + models_manager: Arc, + app_tx: AppEventSender, + config: Config, +} + +impl DraftCompletionManager { + pub(crate) fn new( + auth_manager: Arc, + models_manager: Arc, + app_tx: AppEventSender, + config: Config, + ) -> Self { + Self { + auth_manager, + models_manager, + app_tx, + config, + } + } + + pub(crate) fn start_request(&self, request: DraftCompletionRequest) { + let auth_manager = Arc::clone(&self.auth_manager); + let models_manager = Arc::clone(&self.models_manager); + let app_tx = self.app_tx.clone(); + let config = self.config.clone(); + + tokio::spawn(async move { + let result = timeout( + DRAFT_COMPLETION_TIMEOUT, + complete_draft( + auth_manager, + models_manager, + config, + &request.draft, + request.last_user_turn.as_deref(), + request.last_assistant_turn.as_deref(), + ), + ) + .await; + + let result = match result { + Ok(Ok(suggestions)) => Ok(suggestions), + Ok(Err(err)) => { + tracing::warn!("draft completion failed: {err}"); + Err(err) + } + Err(_) => { + tracing::warn!("draft completion timed out"); + Err("draft completion timed out".to_string()) + } + }; + + app_tx.send(AppEvent::DraftCompletionResult { + request_id: request.request_id, + result, + }); + }); + } +} + +async fn complete_draft( + auth_manager: Arc, + models_manager: Arc, + config: Config, + draft: &str, + last_user_turn: Option<&str>, + last_assistant_turn: Option<&str>, +) -> Result, String> { + let model_info = models_manager + .get_model_info(DRAFT_COMPLETION_MODEL, &config) + .await; + let auth = auth_manager.auth_cached(); + let auth_mode = auth + .as_ref() + .map(codex_core::CodexAuth::auth_mode) + .map(codex_otel::TelemetryAuthMode::from); + let account_id = auth + .as_ref() + .and_then(codex_core::CodexAuth::get_account_id); + let account_email = auth + .as_ref() + .and_then(codex_core::CodexAuth::get_account_email); + let telemetry = SessionTelemetry::new( + ThreadId::new(), + DRAFT_COMPLETION_MODEL, + DRAFT_COMPLETION_MODEL, + account_id, + account_email, + auth_mode, + originator().value, + config.otel.log_user_prompt, + codex_core::terminal::user_agent(), + SessionSource::SubAgent(codex_protocol::protocol::SubAgentSource::Other( + DRAFT_COMPLETION_AGENT_LABEL.to_string(), + )), + ); + let model_client = ModelClient::new( + Some(auth_manager), + ThreadId::new(), + config.model_provider.clone(), + SessionSource::SubAgent(codex_protocol::protocol::SubAgentSource::Other( + DRAFT_COMPLETION_AGENT_LABEL.to_string(), + )), + config.model_verbosity, + false, + false, + false, + None, + ); + let mut session = model_client.new_session(); + let prompt = PromptBuilder::new(draft, last_user_turn, last_assistant_turn).build(); + let mut stream = session + .stream( + &prompt, + &model_info, + &telemetry, + None, + model_info.default_reasoning_summary, + None, + None, + ) + .await + .map_err(|err| err.to_string())?; + + let mut output = String::new(); + while let Some(event) = stream.next().await { + match event.map_err(|err| err.to_string())? { + ResponseEvent::OutputTextDelta(delta) => output.push_str(&delta), + ResponseEvent::OutputItemDone(item) if output.is_empty() => { + output.push_str(&response_item_text(&item)); + } + ResponseEvent::OutputItemDone(_) => {} + ResponseEvent::Completed { .. } => break, + ResponseEvent::Created + | ResponseEvent::OutputItemAdded(_) + | ResponseEvent::ServerModel(_) + | ResponseEvent::ServerReasoningIncluded(_) + | ResponseEvent::ReasoningSummaryDelta { .. } + | ResponseEvent::ReasoningContentDelta { .. } + | ResponseEvent::ReasoningSummaryPartAdded { .. } + | ResponseEvent::RateLimits(_) + | ResponseEvent::ModelsEtag(_) => {} + } + } + + parse_suggestions(draft, &output) +} + +fn response_item_text(item: &ResponseItem) -> String { + match item { + ResponseItem::Message { content, .. } => content + .iter() + .map(|content_item| match content_item { + ContentItem::InputText { text } | ContentItem::OutputText { text } => text.as_str(), + ContentItem::InputImage { .. } => "", + }) + .collect(), + _ => String::new(), + } +} + +struct PromptBuilder<'a> { + draft: &'a str, + last_user_turn: Option<&'a str>, + last_assistant_turn: Option<&'a str>, +} + +impl<'a> PromptBuilder<'a> { + fn new( + draft: &'a str, + last_user_turn: Option<&'a str>, + last_assistant_turn: Option<&'a str>, + ) -> Self { + Self { + draft, + last_user_turn, + last_assistant_turn, + } + } + + fn build(self) -> Prompt { + let request = format!( + "You are helping complete a draft in Codex CLI.\n\ + Return exactly three candidate continuations for the current draft.\n\ + Rules:\n\ + - Each suggestion must be only the text to append to the draft.\n\ + - Do not repeat the existing draft.\n\ + - Prefer concise, plausible continuations.\n\ + - Preserve the apparent tone and topic of the conversation.\n\ + - Return strict JSON: an array of exactly 3 strings and nothing else.\n\n\ + \n{}\n\n\n\ + \n{}\n\n\n\ + \n{}\n\n", + self.last_user_turn.unwrap_or(""), + self.last_assistant_turn.unwrap_or(""), + self.draft, + ); + let mut prompt = Prompt::default(); + prompt.input = vec![ResponseItem::Message { + id: None, + role: "user".to_string(), + content: vec![ContentItem::InputText { text: request }], + end_turn: None, + phase: None, + }]; + prompt.base_instructions = BaseInstructions::default(); + prompt + } +} + +fn parse_suggestions(draft: &str, output: &str) -> Result, String> { + let trimmed = output.trim(); + if trimmed.is_empty() { + return Err("draft completion returned an empty response".to_string()); + } + + let suggestions = if let Ok(parsed) = serde_json::from_str::>(trimmed) { + parsed + } else { + trimmed + .lines() + .map(str::trim) + .filter(|line| !line.is_empty()) + .map(|line| { + line.trim_start_matches(|ch: char| { + ch.is_ascii_digit() || matches!(ch, '.' | ')' | '-' | '*' | ' ') + }) + }) + .map(str::trim) + .filter(|line| !line.is_empty()) + .map(str::to_string) + .collect() + }; + + let mut sanitized = Vec::new(); + for suggestion in suggestions { + let suggestion = sanitize_suggestion(draft, &suggestion); + if suggestion.is_empty() || sanitized.contains(&suggestion) { + continue; + } + sanitized.push(suggestion); + if sanitized.len() == MAX_SUGGESTIONS { + break; + } + } + + if sanitized.is_empty() { + Err("draft completion returned no usable suggestions".to_string()) + } else { + Ok(sanitized) + } +} + +fn sanitize_suggestion(draft: &str, suggestion: &str) -> String { + let trimmed = suggestion + .trim_matches(|ch: char| matches!(ch, '\n' | '\r' | '\t' | '`')) + .trim_end(); + let suffix = trimmed.strip_prefix(draft).unwrap_or(trimmed); + suffix.trim_end().to_string() +} + +#[cfg(test)] +mod tests { + use pretty_assertions::assert_eq; + + use super::parse_suggestions; + + #[test] + fn parse_suggestions_strips_repeated_draft_prefix() { + let suggestions = + parse_suggestions("hello", "[\"hello world\", \" there\", \" everyone\"]") + .expect("parse suggestions"); + + assert_eq!( + suggestions, + vec![ + " world".to_string(), + " there".to_string(), + " everyone".to_string() + ] + ); + } + + #[test] + fn parse_suggestions_falls_back_to_plain_lines() { + let suggestions = parse_suggestions( + "draft", + "1. first suggestion\n2. second suggestion\n3. third suggestion", + ) + .expect("parse suggestions"); + + assert_eq!( + suggestions, + vec![ + "first suggestion".to_string(), + "second suggestion".to_string(), + "third suggestion".to_string() + ] + ); + } +} diff --git a/codex-rs/tui/src/lib.rs b/codex-rs/tui/src/lib.rs index 3d1f8d5f44..4fa52e2bc0 100644 --- a/codex-rs/tui/src/lib.rs +++ b/codex-rs/tui/src/lib.rs @@ -85,6 +85,7 @@ pub mod custom_terminal; mod cwd_prompt; mod debug_config; mod diff_render; +mod draft_completion; mod exec_cell; mod exec_command; mod external_editor;