From 877e28dde668faef774a461a4ff4719613e507fa Mon Sep 17 00:00:00 2001 From: Tuan-Lung Wang Date: Tue, 17 Mar 2026 13:27:32 -0700 Subject: [PATCH] Add fast answer interleaving --- codex-rs/tui/src/answer_interleave.rs | 578 ++++++++++++++++++ codex-rs/tui/src/app.rs | 30 + codex-rs/tui/src/app_event.rs | 11 + codex-rs/tui/src/chatwidget.rs | 319 +++++++++- ...ts__final_answer_interleave_committed.snap | 9 + ...ests__final_answer_interleave_pending.snap | 9 + codex-rs/tui/src/chatwidget/tests.rs | 99 +++ codex-rs/tui/src/lib.rs | 1 + 8 files changed, 1026 insertions(+), 30 deletions(-) create mode 100644 codex-rs/tui/src/answer_interleave.rs create mode 100644 codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__final_answer_interleave_committed.snap create mode 100644 codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__final_answer_interleave_pending.snap diff --git a/codex-rs/tui/src/answer_interleave.rs b/codex-rs/tui/src/answer_interleave.rs new file mode 100644 index 0000000000..f5299bbb2b --- /dev/null +++ b/codex-rs/tui/src/answer_interleave.rs @@ -0,0 +1,578 @@ +use std::path::PathBuf; +use std::sync::Arc; +use std::sync::Mutex; +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::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; +use tokio::time::timeout; +use tokio_stream::StreamExt; + +use crate::app_event::AppEvent; +use crate::app_event_sender::AppEventSender; +use crate::history_cell::HistoryCell; +use crate::markdown; +use crate::wrapping::RtOptions; +use crate::wrapping::adaptive_wrap_lines; + +const ANSWER_INTERLEAVE_MODEL: &str = "gpt-5.1-codex-mini"; +const ANSWER_INTERLEAVE_TIMEOUT: Duration = Duration::from_secs(8); +const ANSWER_INTERLEAVE_STATUS: &str = "restructuring answer..."; +const ANSWER_INTERLEAVE_AGENT_LABEL: &str = "tui_answer_interleave"; + +#[derive(Debug, Clone)] +pub(crate) struct AnswerInterleaveRequest { + pub(crate) request_id: u64, + pub(crate) user_prompt: String, + pub(crate) final_answer: String, +} + +#[derive(Debug)] +struct LiveAnswerState { + markdown: String, + interleaving_pending: bool, +} + +#[derive(Debug, Clone)] +pub(crate) struct LiveAnswerHandle { + state: Arc>, +} + +impl LiveAnswerHandle { + pub(crate) fn new(cwd: PathBuf) -> (Self, LiveAnswerCell) { + let state = Arc::new(Mutex::new(LiveAnswerState { + markdown: String::new(), + interleaving_pending: false, + })); + ( + Self { + state: Arc::clone(&state), + }, + LiveAnswerCell { state, cwd }, + ) + } + + pub(crate) fn push_delta(&self, delta: &str) { + #[expect(clippy::unwrap_used)] + let mut state = self.state.lock().unwrap(); + state.markdown.push_str(delta); + } + + pub(crate) fn set_markdown(&self, markdown: String) { + #[expect(clippy::unwrap_used)] + let mut state = self.state.lock().unwrap(); + state.markdown = markdown; + } + + pub(crate) fn set_interleaving_pending(&self, interleaving_pending: bool) { + #[expect(clippy::unwrap_used)] + let mut state = self.state.lock().unwrap(); + state.interleaving_pending = interleaving_pending; + } + + pub(crate) fn markdown(&self) -> String { + #[expect(clippy::unwrap_used)] + let state = self.state.lock().unwrap(); + state.markdown.clone() + } +} + +#[derive(Debug)] +pub(crate) struct LiveAnswerCell { + state: Arc>, + cwd: PathBuf, +} + +impl HistoryCell for LiveAnswerCell { + fn display_lines(&self, width: u16) -> Vec> { + self.render(width) + } + + fn transcript_lines(&self, width: u16) -> Vec> { + self.render(width) + } +} + +impl LiveAnswerCell { + fn render(&self, width: u16) -> Vec> { + #[expect(clippy::unwrap_used)] + let state = self.state.lock().unwrap(); + let mut lines = Vec::new(); + markdown::append_markdown( + &state.markdown, + Some(width as usize), + Some(self.cwd.as_path()), + &mut lines, + ); + if state.interleaving_pending { + if !lines.is_empty() { + lines.push(Line::from("")); + } + lines.push(Line::from(ANSWER_INTERLEAVE_STATUS.dim())); + } + adaptive_wrap_lines( + &lines, + RtOptions::new(width as usize) + .initial_indent("• ".dim().into()) + .subsequent_indent(" ".into()), + ) + } +} + +pub(crate) struct AnswerInterleaveManager { + auth_manager: Arc, + models_manager: Arc, + app_tx: AppEventSender, + config: Config, +} + +impl AnswerInterleaveManager { + 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: AnswerInterleaveRequest) { + 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( + ANSWER_INTERLEAVE_TIMEOUT, + interleave_answer( + auth_manager, + models_manager, + config, + &request.user_prompt, + &request.final_answer, + ), + ) + .await; + + let result = match result { + Ok(result) => result, + Err(_) => Err("answer interleaving timed out".to_string()), + } + .or_else(|err| { + heuristic_interleave_answer(&request.user_prompt, &request.final_answer).ok_or(err) + }); + + app_tx.send(AppEvent::AnswerInterleaveResult { + request_id: request.request_id, + result, + }); + }); + } +} + +async fn interleave_answer( + auth_manager: Arc, + models_manager: Arc, + config: Config, + user_prompt: &str, + final_answer: &str, +) -> Result { + let model_info = models_manager + .get_model_info(ANSWER_INTERLEAVE_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(), + ANSWER_INTERLEAVE_MODEL, + ANSWER_INTERLEAVE_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( + ANSWER_INTERLEAVE_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( + ANSWER_INTERLEAVE_AGENT_LABEL.to_string(), + )), + config.model_verbosity, + false, + false, + false, + None, + ); + let mut session = model_client.new_session(); + let prompt = PromptBuilder::new(user_prompt, final_answer).build(); + let mut stream = session + .stream( + &prompt, + &model_info, + &telemetry, + Some(ReasoningEffortConfig::None), + ReasoningSummaryConfig::None, + Some(codex_protocol::config_types::ServiceTier::Fast), + 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(_) => {} + } + } + + let output = output.trim().to_string(); + if output.is_empty() { + Err("answer interleaving returned an empty response".to_string()) + } else { + Ok(post_process_interleaved_answer(user_prompt, final_answer, &output).unwrap_or(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> { + user_prompt: &'a str, + final_answer: &'a str, +} + +impl<'a> PromptBuilder<'a> { + fn new(user_prompt: &'a str, final_answer: &'a str) -> Self { + Self { + user_prompt, + final_answer, + } + } + + fn build(self) -> Prompt { + let request = format!( + "You are formatting a Codex final answer for display in a CLI transcript.\n\ + Rewrite the assistant answer so it interleaves the user's questions with the answer.\n\ + Rules:\n\ + - Preserve the original meaning and facts.\n\ + - Do not invent information.\n\ + - If the user prompt contains multiple questions, numbered items, or bullet points, you must rewrite it into interleaved question and answer pairs.\n\ + - Keep the user's numbering or bullets when present.\n\ + - For each question or sub-question, show the prompt text, then a line starting with `A:`.\n\ + - Keep trailing sections like `Assumptions:` or warnings after the paired answers when they do not map cleanly to a single question.\n\ + - Output markdown only. Do not add a preamble or code fence.\n\ + - Only return the original answer unchanged when the prompt is clearly asking a single question.\n\n\ + \n{user_prompt}\n\n\n\ + \n{final_answer}\n\n", + user_prompt = self.user_prompt, + final_answer = self.final_answer, + ); + 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 + } +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum ListKind { + Numbered, + Bulleted, +} + +#[derive(Debug, PartialEq, Eq)] +struct ListItem { + marker: String, + text: String, +} + +#[derive(Debug, PartialEq, Eq)] +struct ParsedList { + items: Vec, + remainder: Option, +} + +fn post_process_interleaved_answer( + user_prompt: &str, + final_answer: &str, + model_output: &str, +) -> Option { + if is_meaningfully_interleaved(user_prompt, model_output) { + return None; + } + + heuristic_interleave_answer(user_prompt, final_answer) +} + +fn is_meaningfully_interleaved(user_prompt: &str, answer: &str) -> bool { + if !has_multiple_prompt_items(user_prompt) { + return true; + } + + answer + .lines() + .any(|line| line.trim_start().starts_with("A:")) +} + +fn has_multiple_prompt_items(user_prompt: &str) -> bool { + parse_simple_list(user_prompt).is_some_and(|parsed| parsed.items.len() > 1) +} + +fn heuristic_interleave_answer(user_prompt: &str, final_answer: &str) -> Option { + let questions = parse_simple_list(user_prompt)?; + let answers = parse_simple_list(final_answer)?; + if questions.items.len() < 2 || questions.items.len() != answers.items.len() { + return None; + } + + let mut lines = Vec::with_capacity(questions.items.len() * 2 + 2); + for (question, answer) in questions.items.iter().zip(answers.items.iter()) { + lines.push(format!("{} {}", question.marker, question.text)); + lines.push(format!("A: {}", answer.text)); + } + if let Some(remainder) = answers.remainder + && !remainder.trim().is_empty() + { + lines.push(String::new()); + lines.push(remainder); + } + + Some(lines.join("\n")) +} + +fn parse_simple_list(text: &str) -> Option { + let mut items = Vec::new(); + let mut kind = None; + let mut started = false; + let mut remainder_start = None; + let lines: Vec<&str> = text.lines().collect(); + + for (idx, line) in lines.iter().enumerate() { + let trimmed = line.trim(); + if trimmed.is_empty() { + if started { + remainder_start = lines[idx + 1..] + .iter() + .position(|candidate| !candidate.trim().is_empty()) + .map(|offset| idx + 1 + offset); + break; + } + continue; + } + + let Some((line_kind, marker, item_text)) = parse_list_item_line(trimmed) else { + if started { + remainder_start = Some(idx); + break; + } + 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(), + }); + } + + if items.is_empty() { + return None; + } + + let remainder = remainder_start.map(|start| lines[start..].join("\n").trim().to_string()); + + Some(ParsedList { items, remainder }) +} + +fn parse_list_item_line(line: &str) -> Option<(ListKind, String, &str)> { + if let Some((marker, text)) = parse_numbered_item_line(line) { + return Some((ListKind::Numbered, marker, text)); + } + if let Some(text) = line.strip_prefix("- ").or_else(|| line.strip_prefix("* ")) { + return Some((ListKind::Bulleted, line[..1].to_string(), text.trim())); + } + None +} + +fn parse_numbered_item_line(line: &str) -> Option<(String, &str)> { + let marker_end = line + .char_indices() + .take_while(|(_, ch)| ch.is_ascii_digit()) + .last() + .map(|(idx, _)| idx + 1)?; + let marker = line.get(..marker_end)?; + let remainder = line.get(marker_end..)?; + let (separator, text) = remainder.split_at_checked(1)?; + if !matches!(separator, "." | ")") { + return None; + } + let text = text.trim_start(); + if text.is_empty() { + return None; + } + Some((format!("{marker}{separator}"), text)) +} + +#[cfg(test)] +mod tests { + use pretty_assertions::assert_eq; + + use super::heuristic_interleave_answer; + use super::parse_simple_list; + use super::post_process_interleaved_answer; + + #[test] + fn heuristic_interleave_pairs_numbered_questions_and_answers() { + let result = heuristic_interleave_answer( + "1. what's my name?\n2. what are we doing?", + "1. Your name is Tuan-Lung.\n2. We are testing interleaving.\n\nAssumptions:\n- This is inferred.", + ); + + assert_eq!( + result, + Some( + "1. what's my name?\n\ + A: Your name is Tuan-Lung.\n\ + 2. what are we doing?\n\ + A: We are testing interleaving.\n\n\ + Assumptions:\n\ + - This is inferred." + .to_string() + ) + ); + } + + #[test] + fn heuristic_interleave_returns_none_for_mismatched_shapes() { + let result = heuristic_interleave_answer( + "1. what's my name?\n2. what are we doing?", + "Your name is Tuan-Lung.", + ); + + assert_eq!(result, None); + } + + #[test] + fn post_process_uses_heuristic_when_model_returns_raw_list_answer() { + let result = post_process_interleaved_answer( + "1. what's my name?\n2. what are we doing?", + "1. Your name is Tuan-Lung.\n2. We are testing interleaving.", + "1. Your name is Tuan-Lung.\n2. We are testing interleaving.", + ); + + assert_eq!( + result, + Some( + "1. what's my name?\n\ + A: Your name is Tuan-Lung.\n\ + 2. what are we doing?\n\ + A: We are testing interleaving." + .to_string() + ) + ); + } + + #[test] + fn parse_simple_list_preserves_remainder() { + let parsed = parse_simple_list("1. one\n2. two\n\nTrailing section"); + + assert_eq!( + parsed.map(|parsed| parsed.remainder), + Some(Some("Trailing section".to_string())) + ); + } + + #[test] + fn heuristic_interleave_supports_prompt_preamble_before_numbered_list() { + let result = heuristic_interleave_answer( + "For testing purpose, let me ask you a few questions:\n1. how are you?\n2. what's your favorite model?", + "1. I'm doing fine.\n2. I don't have a favorite.", + ); + + assert_eq!( + result, + Some( + "1. how are you?\n\ + A: I'm doing fine.\n\ + 2. what's your favorite model?\n\ + A: I don't have a favorite." + .to_string() + ) + ); + } +} diff --git a/codex-rs/tui/src/app.rs b/codex-rs/tui/src/app.rs index 898a45fa05..a405347128 100644 --- a/codex-rs/tui/src/app.rs +++ b/codex-rs/tui/src/app.rs @@ -1,3 +1,4 @@ +use crate::answer_interleave::AnswerInterleaveManager; use crate::app_backtrack::BacktrackState; use crate::app_event::AppEvent; use crate::app_event::ExitMode; @@ -675,6 +676,7 @@ pub(crate) struct App { runtime_sandbox_policy_override: Option, pub(crate) file_search: FileSearchManager, + answer_interleave: AnswerInterleaveManager, pub(crate) transcript_cells: Vec>, @@ -2141,6 +2143,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 answer_interleave = AnswerInterleaveManager::new( + auth_manager.clone(), + thread_manager.get_models_manager(), + app_event_tx.clone(), + config.clone(), + ); #[cfg(not(debug_assertions))] let upgrade_version = crate::updates::get_upgrade_version(&config); @@ -2157,6 +2165,7 @@ impl App { runtime_approval_policy_override: None, runtime_sandbox_policy_override: None, file_search, + answer_interleave, enhanced_keys_supported, transcript_cells: Vec::new(), overlay: None, @@ -2703,6 +2712,13 @@ impl App { AppEvent::FileSearchResult { query, matches } => { self.chat_widget.apply_file_search_result(query, matches); } + AppEvent::StartAnswerInterleave(request) => { + self.answer_interleave.start_request(request); + } + AppEvent::AnswerInterleaveResult { request_id, result } => { + self.chat_widget + .on_answer_interleave_result(request_id, result); + } AppEvent::RateLimitSnapshotFetched(snapshot) => { self.chat_widget.on_rate_limit_snapshot(Some(snapshot)); } @@ -6342,6 +6358,12 @@ guardian_approval = true CodexAuth::from_api_key("Test API Key"), ); let file_search = FileSearchManager::new(config.cwd.clone(), app_event_tx.clone()); + let answer_interleave = AnswerInterleaveManager::new( + auth_manager.clone(), + server.get_models_manager(), + app_event_tx.clone(), + config.clone(), + ); let model = codex_core::test_support::get_model_offline(config.model.as_deref()); let session_telemetry = test_session_telemetry(&config, model.as_str()); @@ -6358,6 +6380,7 @@ guardian_approval = true runtime_approval_policy_override: None, runtime_sandbox_policy_override: None, file_search, + answer_interleave, transcript_cells: Vec::new(), overlay: None, deferred_history_lines: Vec::new(), @@ -6401,6 +6424,12 @@ guardian_approval = true CodexAuth::from_api_key("Test API Key"), ); let file_search = FileSearchManager::new(config.cwd.clone(), app_event_tx.clone()); + let answer_interleave = AnswerInterleaveManager::new( + auth_manager.clone(), + server.get_models_manager(), + app_event_tx.clone(), + config.clone(), + ); let model = codex_core::test_support::get_model_offline(config.model.as_deref()); let session_telemetry = test_session_telemetry(&config, model.as_str()); @@ -6418,6 +6447,7 @@ guardian_approval = true runtime_approval_policy_override: None, runtime_sandbox_policy_override: None, file_search, + answer_interleave, transcript_cells: Vec::new(), overlay: None, deferred_history_lines: Vec::new(), diff --git a/codex-rs/tui/src/app_event.rs b/codex-rs/tui/src/app_event.rs index e2ed046690..e903a7fd9c 100644 --- a/codex-rs/tui/src/app_event.rs +++ b/codex-rs/tui/src/app_event.rs @@ -18,6 +18,7 @@ use codex_protocol::protocol::Event; use codex_protocol::protocol::RateLimitSnapshot; use codex_utils_approval_presets::ApprovalPreset; +use crate::answer_interleave::AnswerInterleaveRequest; use crate::bottom_pane::ApprovalRequest; use crate::bottom_pane::StatusLineItem; use crate::history_cell::HistoryCell; @@ -128,6 +129,16 @@ pub(crate) enum AppEvent { matches: Vec, }, + /// Kick off an asynchronous fast-model pass that restructures a final answer into an + /// interleaved question/answer display. + StartAnswerInterleave(AnswerInterleaveRequest), + + /// Result of a completed answer interleaving request. + AnswerInterleaveResult { + request_id: u64, + result: Result, + }, + /// Result of refreshing rate limits RateLimitSnapshotFetched(RateLimitSnapshot), diff --git a/codex-rs/tui/src/chatwidget.rs b/codex-rs/tui/src/chatwidget.rs index c692f26590..14902ba714 100644 --- a/codex-rs/tui/src/chatwidget.rs +++ b/codex-rs/tui/src/chatwidget.rs @@ -38,6 +38,9 @@ use std::time::Duration; use std::time::Instant; use self::realtime::PendingSteerCompareKey; +use crate::answer_interleave::AnswerInterleaveRequest; +use crate::answer_interleave::LiveAnswerCell; +use crate::answer_interleave::LiveAnswerHandle; use crate::app_event::RealtimeAudioDeviceKind; #[cfg(all(not(target_os = "linux"), feature = "voice-input"))] use crate::audio_device::list_realtime_audio_device_names; @@ -90,6 +93,7 @@ use codex_protocol::config_types::Settings; use codex_protocol::config_types::WindowsSandboxLevel; use codex_protocol::items::AgentMessageContent; use codex_protocol::items::AgentMessageItem; +use codex_protocol::items::TurnItem; use codex_protocol::models::MessagePhase; use codex_protocol::models::local_image_label_text; use codex_protocol::parse_command::ParsedCommand; @@ -118,6 +122,7 @@ use codex_protocol::protocol::GuardianAssessmentEvent; use codex_protocol::protocol::GuardianAssessmentStatus; use codex_protocol::protocol::ImageGenerationBeginEvent; use codex_protocol::protocol::ImageGenerationEndEvent; +use codex_protocol::protocol::ItemStartedEvent; use codex_protocol::protocol::ListCustomPromptsResponseEvent; use codex_protocol::protocol::ListSkillsResponseEvent; use codex_protocol::protocol::McpListToolsResponseEvent; @@ -667,6 +672,11 @@ pub(crate) struct ChatWidget { adaptive_chunking: AdaptiveChunkingPolicy, // Stream lifecycle controller stream_controller: Option, + agent_message_stream_kind: Option, + live_answer_handle: Option, + pending_answer_interleave: Option, + deferred_turn_complete: Option, + next_answer_interleave_request_id: u64, // Stream lifecycle controller for proposed plan output. plan_stream_controller: Option, // Latest completed user-visible Codex output that `/copy` should place on the clipboard. @@ -907,6 +917,25 @@ struct PendingSteer { compare_key: PendingSteerCompareKey, } +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum AgentMessageStreamKind { + Unknown, + Commentary, + FinalAnswer, +} + +#[derive(Debug)] +struct PendingAnswerInterleave { + request_id: u64, + handle: LiveAnswerHandle, + raw_answer: String, +} + +#[derive(Debug)] +struct DeferredTurnComplete { + from_replay: bool, +} + pub(crate) fn create_initial_user_message( text: Option, local_image_paths: Vec, @@ -1560,6 +1589,156 @@ impl ChatWidget { self.request_redraw(); } + fn current_agent_message_stream_kind(&self) -> AgentMessageStreamKind { + self.agent_message_stream_kind + .unwrap_or(AgentMessageStreamKind::Unknown) + } + + fn current_user_prompt_for_interleave(&self) -> Option { + self.last_rendered_user_message_event + .as_ref() + .map(|event| event.message.trim().to_string()) + .filter(|message| !message.is_empty()) + } + + fn current_session_cwd(&self) -> &Path { + self.current_cwd + .as_deref() + .unwrap_or(self.config.cwd.as_path()) + } + + fn prepare_for_new_assistant_output(&mut self) { + self.flush_unified_exec_wait_streak(); + self.flush_active_cell(); + + if self.needs_final_message_separator && self.had_work_activity { + let elapsed_seconds = self + .bottom_pane + .status_widget() + .map(super::status_indicator_widget::StatusIndicatorWidget::elapsed_seconds) + .map(|current| self.worked_elapsed_from(current)); + self.add_to_history(history_cell::FinalMessageSeparator::new( + elapsed_seconds, + None, + )); + self.needs_final_message_separator = false; + self.had_work_activity = false; + } else if self.needs_final_message_separator { + self.needs_final_message_separator = false; + } + } + + fn ensure_live_answer_handle(&mut self) -> LiveAnswerHandle { + if let Some(handle) = self.live_answer_handle.as_ref() { + return handle.clone(); + } + if let Some(pending) = self.pending_answer_interleave.as_ref() { + return pending.handle.clone(); + } + + if self.stream_controller.is_some() { + self.flush_answer_stream_with_separator(); + } + self.prepare_for_new_assistant_output(); + + let (handle, cell) = LiveAnswerHandle::new(self.current_session_cwd().to_path_buf()); + self.live_answer_handle = Some(handle.clone()); + self.active_cell = Some(Box::new(cell)); + self.bump_active_cell_revision(); + handle + } + + fn handle_final_answer_delta(&mut self, delta: String) { + if delta.is_empty() { + return; + } + let handle = self.ensure_live_answer_handle(); + handle.push_delta(&delta); + self.bump_active_cell_revision(); + self.request_redraw(); + } + + fn finalize_live_answer_without_interleave(&mut self, raw_answer: String) { + let handle = self.ensure_live_answer_handle(); + handle.set_interleaving_pending(false); + handle.set_markdown(raw_answer.clone()); + self.live_answer_handle = None; + self.pending_answer_interleave = None; + self.last_copyable_output = Some(raw_answer); + self.bump_active_cell_revision(); + self.flush_active_cell(); + self.request_redraw(); + } + + fn start_answer_interleave(&mut self, raw_answer: String) { + let Some(user_prompt) = self.current_user_prompt_for_interleave() else { + self.finalize_live_answer_without_interleave(raw_answer); + return; + }; + + self.next_answer_interleave_request_id = + self.next_answer_interleave_request_id.wrapping_add(1); + let request_id = self.next_answer_interleave_request_id; + let handle = self.ensure_live_answer_handle(); + handle.set_markdown(raw_answer.clone()); + handle.set_interleaving_pending(true); + self.pending_answer_interleave = Some(PendingAnswerInterleave { + request_id, + handle, + raw_answer: raw_answer.clone(), + }); + self.bump_active_cell_revision(); + self.app_event_tx + .send(AppEvent::StartAnswerInterleave(AnswerInterleaveRequest { + request_id, + user_prompt, + final_answer: raw_answer, + })); + self.request_redraw(); + } + + fn on_agent_message_item_started(&mut self, item: AgentMessageItem) { + self.agent_message_stream_kind = Some(match item.phase { + Some(MessagePhase::Commentary) => AgentMessageStreamKind::Commentary, + Some(MessagePhase::FinalAnswer) => AgentMessageStreamKind::FinalAnswer, + None => AgentMessageStreamKind::Unknown, + }); + } + + pub(crate) fn on_answer_interleave_result( + &mut self, + request_id: u64, + result: Result, + ) { + let Some(pending) = self.pending_answer_interleave.take() else { + return; + }; + if pending.request_id != request_id { + self.pending_answer_interleave = Some(pending); + return; + } + + let final_answer = match result { + Ok(answer) if !answer.trim().is_empty() => answer, + Ok(_) => pending.raw_answer.clone(), + Err(err) => { + tracing::warn!("answer interleave failed: {err}"); + pending.raw_answer.clone() + } + }; + + pending.handle.set_interleaving_pending(false); + pending.handle.set_markdown(final_answer.clone()); + self.live_answer_handle = None; + self.last_copyable_output = Some(final_answer.clone()); + self.bump_active_cell_revision(); + self.flush_active_cell(); + if let Some(deferred) = self.deferred_turn_complete.take() { + self.finish_task_complete(Some(final_answer), deferred.from_replay); + } + self.request_redraw(); + } + fn finalize_completed_assistant_message(&mut self, message: Option<&str>) { // If we have a stream_controller, the finalized message payload is redundant because the // visible content has already been accumulated through deltas. @@ -1579,7 +1758,12 @@ impl ChatWidget { } fn on_agent_message_delta(&mut self, delta: String) { - self.handle_streaming_delta(delta); + match self.current_agent_message_stream_kind() { + AgentMessageStreamKind::Unknown | AgentMessageStreamKind::Commentary => { + self.handle_streaming_delta(delta) + } + AgentMessageStreamKind::FinalAnswer => self.handle_final_answer_delta(delta), + } } fn on_plan_delta(&mut self, delta: String) { @@ -1699,6 +1883,10 @@ impl ChatWidget { self.plan_item_active = false; self.adaptive_chunking.reset(); self.plan_stream_controller = None; + self.agent_message_stream_kind = None; + self.live_answer_handle = None; + self.pending_answer_interleave = None; + self.deferred_turn_complete = None; self.turn_runtime_metrics = RuntimeMetricsSummary::default(); self.session_telemetry.reset_runtime_metrics(); self.bottom_pane.clear_quit_shortcut_hint(); @@ -1715,6 +1903,15 @@ impl ChatWidget { } fn on_task_complete(&mut self, last_agent_message: Option, from_replay: bool) { + if self.pending_answer_interleave.is_some() && !from_replay { + self.deferred_turn_complete = Some(DeferredTurnComplete { from_replay }); + self.request_redraw(); + return; + } + self.finish_task_complete(last_agent_message, from_replay); + } + + fn finish_task_complete(&mut self, last_agent_message: Option, from_replay: bool) { self.submit_pending_steers_after_interrupt = false; if let Some(message) = last_agent_message.as_ref() && !message.trim().is_empty() @@ -3000,16 +3197,61 @@ impl ChatWidget { /// Commentary completion sets a deferred restore flag so the status row /// returns once stream queues are idle. Final-answer completion (or absent /// phase for legacy models) clears the flag to preserve historical behavior. - fn on_agent_message_item_completed(&mut self, item: AgentMessageItem) { + fn on_agent_message_item_completed(&mut self, item: AgentMessageItem, from_replay: bool) { let mut message = String::new(); for content in &item.content { match content { AgentMessageContent::Text { text } => message.push_str(text), } } - self.finalize_completed_assistant_message( - (!message.is_empty()).then_some(message.as_str()), - ); + let live_answer_markdown = self + .live_answer_handle + .as_ref() + .map(LiveAnswerHandle::markdown); + if from_replay || self.is_review_mode { + self.finalize_completed_assistant_message((!message.is_empty()).then_some(&message)); + } else { + match item.phase { + Some(MessagePhase::Commentary) => { + if self.live_answer_handle.is_some() { + self.finalize_live_answer_without_interleave(message.clone()); + } else { + self.finalize_completed_assistant_message( + (!message.is_empty()).then_some(message.as_str()), + ); + } + } + Some(MessagePhase::FinalAnswer) => { + let final_answer = if message.is_empty() { + live_answer_markdown.unwrap_or_default() + } else { + message.clone() + }; + if !final_answer.is_empty() && self.thread_id.is_some() { + self.start_answer_interleave(final_answer); + } else { + self.finalize_live_answer_without_interleave(final_answer); + } + } + None => { + if self.live_answer_handle.is_some() { + let final_answer = if message.is_empty() { + live_answer_markdown.unwrap_or_default() + } else { + message.clone() + }; + self.finalize_live_answer_without_interleave(final_answer); + } else if !message.is_empty() && self.thread_id.is_some() { + self.start_answer_interleave(message.clone()); + } else { + self.finalize_completed_assistant_message( + (!message.is_empty()).then_some(message.as_str()), + ); + } + } + } + } + self.agent_message_stream_kind = None; self.pending_status_indicator_restore = match item.phase { // Models that don't support preambles only output AgentMessageItems on turn completion. Some(MessagePhase::FinalAnswer) | None => false, @@ -3097,29 +3339,8 @@ impl ChatWidget { #[inline] fn handle_streaming_delta(&mut self, delta: String) { - // Before streaming agent content, flush any active exec cell group. - self.flush_unified_exec_wait_streak(); - self.flush_active_cell(); - if self.stream_controller.is_none() { - // If the previous turn inserted non-stream history (exec output, patch status, MCP - // calls), render a separator before starting the next streamed assistant message. - if self.needs_final_message_separator && self.had_work_activity { - let elapsed_seconds = self - .bottom_pane - .status_widget() - .map(super::status_indicator_widget::StatusIndicatorWidget::elapsed_seconds) - .map(|current| self.worked_elapsed_from(current)); - self.add_to_history(history_cell::FinalMessageSeparator::new( - elapsed_seconds, - None, - )); - self.needs_final_message_separator = false; - self.had_work_activity = false; - } else if self.needs_final_message_separator { - // Reset the flag even if we don't show separator (no work was done) - self.needs_final_message_separator = false; - } + self.prepare_for_new_assistant_output(); self.stream_controller = Some(StreamController::new( self.last_rendered_width.get().map(|w| w.saturating_sub(2)), &self.config.cwd, @@ -3585,6 +3806,11 @@ impl ChatWidget { rate_limit_poller: None, adaptive_chunking: AdaptiveChunkingPolicy::default(), stream_controller: None, + agent_message_stream_kind: None, + live_answer_handle: None, + pending_answer_interleave: None, + deferred_turn_complete: None, + next_answer_interleave_request_id: 0, plan_stream_controller: None, last_copyable_output: None, running_commands: HashMap::new(), @@ -3771,6 +3997,11 @@ impl ChatWidget { rate_limit_poller: None, adaptive_chunking: AdaptiveChunkingPolicy::default(), stream_controller: None, + agent_message_stream_kind: None, + live_answer_handle: None, + pending_answer_interleave: None, + deferred_turn_complete: None, + next_answer_interleave_request_id: 0, plan_stream_controller: None, last_copyable_output: None, running_commands: HashMap::new(), @@ -3949,6 +4180,11 @@ impl ChatWidget { rate_limit_poller: None, adaptive_chunking: AdaptiveChunkingPolicy::default(), stream_controller: None, + agent_message_stream_kind: None, + live_answer_handle: None, + pending_answer_interleave: None, + deferred_turn_complete: None, + next_answer_interleave_request_id: 0, plan_stream_controller: None, last_copyable_output: None, running_commands: HashMap::new(), @@ -4792,7 +5028,14 @@ impl ChatWidget { } fn flush_active_cell(&mut self) { + let flushed_live_answer = self + .active_cell + .as_ref() + .is_some_and(|cell| cell.as_any().is::()); if let Some(active) = self.active_cell.take() { + if flushed_live_answer { + self.live_answer_handle = None; + } self.needs_final_message_separator = true; self.app_event_tx.send(AppEvent::InsertHistoryCell(active)); } @@ -4810,8 +5053,16 @@ impl ChatWidget { .active_cell .as_ref() .is_some_and(|c| c.as_any().is::()); + let keep_interleaving_answer_active = self.live_answer_handle.is_some() + && self + .active_cell + .as_ref() + .is_some_and(|c| c.as_any().is::()); - if !keep_placeholder_header_active && !cell.display_lines(u16::MAX).is_empty() { + if !keep_placeholder_header_active + && !keep_interleaving_answer_active + && !cell.display_lines(u16::MAX).is_empty() + { // Only break exec grouping if the cell renders visible lines. self.flush_active_cell(); self.needs_final_message_separator = true; @@ -5222,6 +5473,15 @@ impl ChatWidget { self.on_agent_message(message) } EventMsg::AgentMessage(AgentMessageEvent { .. }) => {} + EventMsg::ItemStarted(ItemStartedEvent { item, .. }) + if !from_replay && matches!(item, TurnItem::AgentMessage(_)) => + { + let TurnItem::AgentMessage(item) = item else { + unreachable!("guard ensures agent message item"); + }; + self.on_agent_message_item_started(item); + } + EventMsg::ItemStarted(_) => {} EventMsg::AgentMessageDelta(AgentMessageDeltaEvent { delta }) => { self.on_agent_message_delta(delta) } @@ -5397,7 +5657,6 @@ impl ChatWidget { } } EventMsg::RawResponseItem(_) - | EventMsg::ItemStarted(_) | EventMsg::AgentMessageContentDelta(_) | EventMsg::ReasoningContentDelta(_) | EventMsg::ReasoningRawContentDelta(_) @@ -5462,7 +5721,7 @@ impl ChatWidget { self.on_plan_item_completed(plan_item.text.clone()); } if let codex_protocol::items::TurnItem::AgentMessage(item) = item { - self.on_agent_message_item_completed(item); + self.on_agent_message_item_completed(item, from_replay); } } } diff --git a/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__final_answer_interleave_committed.snap b/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__final_answer_interleave_committed.snap new file mode 100644 index 0000000000..a749516cea --- /dev/null +++ b/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__final_answer_interleave_committed.snap @@ -0,0 +1,9 @@ +--- +source: tui/src/chatwidget/tests.rs +assertion_line: 4199 +expression: rendered +--- +• 1. what's my name? + A: Your name is Tuan-Lung. + 2. what are we doing? + A: We are testing interleaving. diff --git a/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__final_answer_interleave_pending.snap b/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__final_answer_interleave_pending.snap new file mode 100644 index 0000000000..6334225f8d --- /dev/null +++ b/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__final_answer_interleave_pending.snap @@ -0,0 +1,9 @@ +--- +source: tui/src/chatwidget/tests.rs +assertion_line: 4192 +expression: pending +--- +• 1. Your name is Tuan-Lung. + 2. We are testing interleaving. + + restructuring answer... diff --git a/codex-rs/tui/src/chatwidget/tests.rs b/codex-rs/tui/src/chatwidget/tests.rs index ffc571288a..08e9a6bd6f 100644 --- a/codex-rs/tui/src/chatwidget/tests.rs +++ b/codex-rs/tui/src/chatwidget/tests.rs @@ -5,6 +5,7 @@ //! changes show up as stable, reviewable diffs. use super::*; +use crate::answer_interleave::AnswerInterleaveRequest; use crate::app_event::AppEvent; use crate::app_event::ExitMode; #[cfg(all(not(target_os = "linux"), feature = "voice-input"))] @@ -1864,6 +1865,11 @@ async fn make_chatwidget_manual( rate_limit_poller: None, adaptive_chunking: crate::streaming::chunking::AdaptiveChunkingPolicy::default(), stream_controller: None, + agent_message_stream_kind: None, + live_answer_handle: None, + pending_answer_interleave: None, + deferred_turn_complete: None, + next_answer_interleave_request_id: 0, plan_stream_controller: None, pending_guardian_review_status: PendingGuardianReviewStatus::default(), last_copyable_output: None, @@ -2031,6 +2037,26 @@ fn drain_insert_history( out } +fn take_answer_interleave_request( + rx: &mut tokio::sync::mpsc::UnboundedReceiver, +) -> AnswerInterleaveRequest { + loop { + match rx.try_recv() { + Ok(AppEvent::StartAnswerInterleave(request)) => return request, + Ok(AppEvent::InsertHistoryCell(_)) => { + panic!("unexpected history insert before answer interleave result") + } + Ok(_) => continue, + Err(TryRecvError::Empty) => { + panic!("expected answer interleave request but queue was empty") + } + Err(TryRecvError::Disconnected) => { + panic!("expected answer interleave request but channel closed") + } + } + } +} + fn lines_to_single_string(lines: &[ratatui::text::Line<'static>]) -> String { let mut s = String::new(); for line in lines { @@ -4133,6 +4159,79 @@ async fn live_legacy_agent_message_after_item_completed_does_not_duplicate_assis assert!(drain_insert_history(&mut rx).is_empty()); } +#[tokio::test] +async fn explicit_final_answer_requests_interleaving_and_commits_result() { + let (mut chat, mut rx, _op_rx) = make_chatwidget_manual(None).await; + chat.thread_id = Some(ThreadId::new()); + complete_user_message( + &mut chat, + "user-1", + "1. what's my name?\n2. what are we doing?", + ); + let _ = drain_insert_history(&mut rx); + + complete_assistant_message( + &mut chat, + "msg-live", + "1. Your name is Tuan-Lung.\n2. We are testing interleaving.", + Some(MessagePhase::FinalAnswer), + ); + + let pending = lines_to_single_string( + &chat + .active_cell + .as_ref() + .expect("live answer cell") + .display_lines(80), + ); + assert_snapshot!("final_answer_interleave_pending", pending); + + let request = take_answer_interleave_request(&mut rx); + assert_eq!( + request.user_prompt, + "1. what's my name?\n2. what are we doing?" + ); + assert!(drain_insert_history(&mut rx).is_empty()); + chat.on_answer_interleave_result( + request.request_id, + Ok("1. what's my name?\nA: Your name is Tuan-Lung.\n2. what are we doing?\nA: We are testing interleaving.".into()), + ); + + let inserted = drain_insert_history(&mut rx); + assert_eq!(inserted.len(), 1); + let rendered = lines_to_single_string(&inserted[0]); + assert_snapshot!("final_answer_interleave_committed", rendered); +} + +#[tokio::test] +async fn legacy_final_answer_without_phase_still_requests_interleaving() { + let (mut chat, mut rx, _op_rx) = make_chatwidget_manual(None).await; + chat.thread_id = Some(ThreadId::new()); + complete_user_message( + &mut chat, + "user-legacy", + "1. what's my name?\n2. what are we doing?", + ); + let _ = drain_insert_history(&mut rx); + + complete_assistant_message( + &mut chat, + "msg-legacy", + "1. Your name is Tuan-Lung.\n2. We are testing interleaving.", + None, + ); + + let request = take_answer_interleave_request(&mut rx); + assert_eq!( + request.user_prompt, + "1. what's my name?\n2. what are we doing?" + ); + assert_eq!( + request.final_answer, + "1. Your name is Tuan-Lung.\n2. We are testing interleaving." + ); +} + #[test] fn rendered_user_message_event_from_inputs_matches_flattened_user_message_shape() { let local_image = PathBuf::from("/tmp/local.png"); diff --git a/codex-rs/tui/src/lib.rs b/codex-rs/tui/src/lib.rs index 43230eb2d4..3d1f8d5f44 100644 --- a/codex-rs/tui/src/lib.rs +++ b/codex-rs/tui/src/lib.rs @@ -66,6 +66,7 @@ use tracing_subscriber::prelude::*; use uuid::Uuid; mod additional_dirs; +mod answer_interleave; mod app; mod app_backtrack; mod app_event;