Add fast draft completion popup

This commit is contained in:
Tuan-Lung Wang
2026-03-17 15:44:10 -07:00
parent 877e28dde6
commit 5f4a2b0fbb
11 changed files with 962 additions and 27 deletions
+55 -25
View File
@@ -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<ParsedList> {
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<ParsedList> {
.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()
)
);
}
}
+30
View File
@@ -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<SandboxPolicy>,
pub(crate) file_search: FileSearchManager,
draft_completion: DraftCompletionManager,
answer_interleave: AnswerInterleaveManager,
pub(crate) transcript_cells: Vec<Arc<dyn HistoryCell>>,
@@ -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,
+10
View File
@@ -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<String, String>,
},
/// Kick off an asynchronous fast-model draft completion request.
StartDraftCompletion(DraftCompletionRequest),
/// Result of a completed draft completion request.
DraftCompletionResult {
request_id: u64,
result: Result<Vec<String>, String>,
},
/// Result of refreshing rate limits
RateLimitSnapshotFetched(RateLimitSnapshot),
@@ -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<TextElement>),
RequestDraftCompletion {
request_id: u64,
draft: String,
},
None,
}
@@ -404,6 +409,7 @@ pub(crate) struct ChatComposer {
collaboration_mode_indicator: Option<CollaborationModeIndicator>,
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<Vec<String>, 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::<AppEvent>();
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::<AppEvent>();
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::<AppEvent>();
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");
@@ -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<String>,
error_message: Option<String>,
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<String>) {
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<GenericDisplayRow> {
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<String> {
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()
}
+10
View File
@@ -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<Vec<String>, 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);
@@ -0,0 +1,13 @@
---
source: tui/src/bottom_pane/chat_composer.rs
expression: terminal.backend()
---
" "
"› hell "
" "
" "
" "
" "
" o world "
" there "
" everyone "
@@ -0,0 +1,13 @@
---
source: tui/src/bottom_pane/chat_composer.rs
expression: terminal.backend()
---
" "
"› hell "
" "
" "
" "
" "
" "
" "
"draft completion timed out "
+22 -2
View File
@@ -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<String> {
fn last_user_prompt_text(&self) -> Option<String> {
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<Vec<String>, 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
+338
View File
@@ -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<String>,
pub(crate) last_assistant_turn: Option<String>,
}
pub(crate) struct DraftCompletionManager {
auth_manager: Arc<AuthManager>,
models_manager: Arc<ModelsManager>,
app_tx: AppEventSender,
config: Config,
}
impl DraftCompletionManager {
pub(crate) fn new(
auth_manager: Arc<AuthManager>,
models_manager: Arc<ModelsManager>,
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<AuthManager>,
models_manager: Arc<ModelsManager>,
config: Config,
draft: &str,
last_user_turn: Option<&str>,
last_assistant_turn: Option<&str>,
) -> Result<Vec<String>, 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\
<last_user_turn>\n{}\n</last_user_turn>\n\n\
<last_assistant_turn>\n{}\n</last_assistant_turn>\n\n\
<current_draft>\n{}\n</current_draft>\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<Vec<String>, 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::<Vec<String>>(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()
]
);
}
}
+1
View File
@@ -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;