codex/codex-rs/tui/src/model_migration.rs

//! Model upgrade copy and choices with shared picker presentation.
//! Mandatory and optional migration outcomes retain their existing keyboard policy.

use crate::bottom_pane::picker_option_list;
use crate::bottom_pane::render_menu_surface;
use crate::key_hint;
use crate::markdown_render::render_markdown_text_with_width;
use crate::render::Insets;
use crate::render::renderable::ColumnRenderable;
use crate::render::renderable::FlexRenderable;
use crate::render::renderable::Renderable;
use crate::render::renderable::RenderableExt as _;
use crate::render::renderable::RenderableItem;
use crate::tui::FrameRequester;
use crate::tui::Tui;
use crate::tui::TuiEvent;
use crossterm::event::KeyCode;
use crossterm::event::KeyEvent;
use crossterm::event::KeyEventKind;
use crossterm::event::KeyModifiers;
use ratatui::layout::Rect;
use ratatui::prelude::Stylize as _;
use ratatui::prelude::Widget;
use ratatui::text::Line;
use ratatui::text::Span;
use ratatui::widgets::Clear;
use ratatui::widgets::Paragraph;
use ratatui::widgets::WidgetRef;
use ratatui::widgets::Wrap;
use tokio_stream::StreamExt;

/// Outcome of the migration prompt.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(crate) enum ModelMigrationOutcome {
    Accepted,
    Rejected,
    Exit,
}

#[derive(Clone)]
pub(crate) struct ModelMigrationCopy {
    pub heading: Vec<Span<'static>>,
    pub content: Vec<Line<'static>>,
    pub can_opt_out: bool,
    pub markdown: Option<String>,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum MigrationMenuOption {
    TryNewModel,
    UseExistingModel,
}

impl MigrationMenuOption {
    fn all() -> [Self; 2] {
        [Self::TryNewModel, Self::UseExistingModel]
    }

    fn label(self) -> &'static str {
        match self {
            Self::TryNewModel => "Try new model",
            Self::UseExistingModel => "Use existing model",
        }
    }
}

#[allow(clippy::too_many_arguments)]
pub(crate) fn migration_copy_for_models(
    current_model: &str,
    target_model: &str,
    model_link: Option<String>,
    migration_copy: Option<String>,
    migration_markdown: Option<String>,
    target_display_name: String,
    target_description: Option<String>,
    can_opt_out: bool,
) -> ModelMigrationCopy {
    if let Some(migration_markdown) = migration_markdown {
        return ModelMigrationCopy {
            heading: Vec::new(),
            content: Vec::new(),
            can_opt_out,
            markdown: Some(fill_migration_markdown(
                &migration_markdown,
                current_model,
                target_model,
            )),
        };
    }

    let heading_text = Span::from(format!(
        "Codex just got an upgrade. Introducing {target_display_name}."
    ))
    .bold();
    let description_line: Line<'static>;
    if let Some(migration_copy) = &migration_copy {
        description_line = Line::from(migration_copy.clone());
    } else {
        description_line = target_description
            .filter(|desc| !desc.is_empty())
            .map(Line::from)
            .unwrap_or_else(|| {
                Line::from(format!(
                    "{target_display_name} is recommended for better performance and reliability."
                ))
            });
    }

    let mut content = vec![];
    if migration_copy.is_none() {
        content.push(Line::from(format!(
            "We recommend switching from {current_model} to {target_model}."
        )));
        content.push(Line::from(""));
    }

    if let Some(model_link) = model_link {
        content.push(Line::from(vec![
            format!("{description_line} Learn more about {target_display_name} at ").into(),
            model_link.fg(crate::style::accent_color()).underlined(),
        ]));
        content.push(Line::from(""));
    } else {
        content.push(description_line);
        content.push(Line::from(""));
    }

    if can_opt_out {
        content.push(Line::from(format!(
            "You can continue using {current_model} if you prefer."
        )));
    }

    ModelMigrationCopy {
        heading: vec![heading_text],
        content,
        can_opt_out,
        markdown: None,
    }
}

pub(crate) async fn run_model_migration_prompt(
    tui: &mut Tui,
    copy: ModelMigrationCopy,
) -> std::io::Result<ModelMigrationOutcome> {
    let alt = AltScreenGuard::enter(tui);
    let mut screen = ModelMigrationScreen::new(alt.tui.frame_requester(), copy);

    alt.tui.draw(u16::MAX, |frame| {
        frame.render_widget_ref(&screen, frame.area());
    })?;

    alt.tui.discard_pending_input_before_interactive_screen()?;
    let events = alt.tui.event_stream();
    tokio::pin!(events);

    while !screen.is_done() {
        if let Some(event) = events.next().await {
            let _ = alt.tui.screen_size_for_event(&event);
            match event {
                TuiEvent::Key(key_event) => screen.handle_key(key_event),
                TuiEvent::Paste(_) | TuiEvent::FocusLost | TuiEvent::Mouse(_) => {}
                TuiEvent::Draw | TuiEvent::Resume | TuiEvent::Resize(_) | TuiEvent::FocusGained => {
                    let _ = alt.tui.draw(u16::MAX, |frame| {
                        frame.render_widget_ref(&screen, frame.area());
                    });
                }
            }
        } else {
            screen.accept();
            break;
        }
    }

    Ok(screen.outcome())
}

struct ModelMigrationScreen {
    request_frame: FrameRequester,
    copy: ModelMigrationCopy,
    done: bool,
    outcome: ModelMigrationOutcome,
    highlighted_option: MigrationMenuOption,
}

impl ModelMigrationScreen {
    fn new(request_frame: FrameRequester, copy: ModelMigrationCopy) -> Self {
        Self {
            request_frame,
            copy,
            done: false,
            outcome: ModelMigrationOutcome::Accepted,
            highlighted_option: MigrationMenuOption::TryNewModel,
        }
    }

    fn finish_with(&mut self, outcome: ModelMigrationOutcome) {
        self.outcome = outcome;
        self.done = true;
        self.request_frame.schedule_frame();
    }

    fn accept(&mut self) {
        self.finish_with(ModelMigrationOutcome::Accepted);
    }

    fn reject(&mut self) {
        self.finish_with(ModelMigrationOutcome::Rejected);
    }

    fn exit(&mut self) {
        self.finish_with(ModelMigrationOutcome::Exit);
    }

    fn confirm_selection(&mut self) {
        if self.copy.can_opt_out {
            match self.highlighted_option {
                MigrationMenuOption::TryNewModel => self.accept(),
                MigrationMenuOption::UseExistingModel => self.reject(),
            }
        } else {
            self.accept();
        }
    }

    fn highlight_option(&mut self, option: MigrationMenuOption) {
        if self.highlighted_option != option {
            self.highlighted_option = option;
            self.request_frame.schedule_frame();
        }
    }

    fn handle_key(&mut self, key_event: KeyEvent) {
        if key_event.kind == KeyEventKind::Release {
            return;
        }

        if is_ctrl_exit_combo(key_event) {
            self.exit();
            return;
        }

        if self.copy.can_opt_out {
            self.handle_menu_key(key_event.code);
        } else if matches!(key_event.code, KeyCode::Esc | KeyCode::Enter) {
            self.accept();
        }
    }

    fn is_done(&self) -> bool {
        self.done
    }

    fn outcome(&self) -> ModelMigrationOutcome {
        self.outcome
    }
}

impl WidgetRef for &ModelMigrationScreen {
    fn render_ref(&self, area: ratatui::layout::Rect, buf: &mut ratatui::buffer::Buffer) {
        Clear.render(area, buf);

        let mut body = ColumnRenderable::new();
        if let Some(markdown) = self.copy.markdown.as_ref() {
            self.render_markdown_content(markdown, area.width, &mut body);
        } else {
            body.push(
                Paragraph::new(Line::from(self.copy.heading.clone()))
                    .wrap(Wrap { trim: false })
                    .inset(Insets::vh(/*v*/ 0, /*h*/ 2)),
            );
            body.push(Line::from(""));
            self.render_content(&mut body);
        }
        let mut column = FlexRenderable::new();
        column.push(/*flex*/ 1, RenderableItem::Borrowed(&""));
        column.push(/*flex*/ 1, RenderableItem::Borrowed(&body));
        if self.copy.can_opt_out {
            let selected_index = match self.highlighted_option {
                MigrationMenuOption::TryNewModel => 0,
                MigrationMenuOption::UseExistingModel => 1,
            };
            column.push(
                /*flex*/ 0,
                picker_option_list(
                    MigrationMenuOption::all()
                        .into_iter()
                        .map(|option| option.label().to_string())
                        .collect(),
                    selected_index,
                ),
            );
        }
        let mut hints = key_hint::key_label_spans(&format!(
            "{}/{}",
            key_hint::plain(KeyCode::Enter).display_label(),
            key_hint::plain(KeyCode::Esc).display_label()
        ));
        hints.push(if self.copy.can_opt_out {
            " confirm · ".dim()
        } else {
            " continue · ".dim()
        });
        hints.push(key_hint::ctrl(KeyCode::Char('c')).into());
        hints.push(" quit".dim());
        column.push(
            /*flex*/ 0,
            Paragraph::new(Line::from(hints))
                .wrap(Wrap { trim: false })
                .inset(Insets::vh(/*v*/ 0, /*h*/ 2)),
        );
        column.push(/*flex*/ 1, RenderableItem::Borrowed(&""));
        let panel = Rect {
            height: column.desired_height(area.width).min(area.height),
            ..area
        };
        render_menu_surface(panel, buf);
        column.render(panel, buf);
    }
}

impl ModelMigrationScreen {
    fn handle_menu_key(&mut self, code: KeyCode) {
        match code {
            KeyCode::Up | KeyCode::Char('k') => {
                self.highlight_option(MigrationMenuOption::TryNewModel);
            }
            KeyCode::Down | KeyCode::Char('j') => {
                self.highlight_option(MigrationMenuOption::UseExistingModel);
            }
            KeyCode::Char('1') => {
                self.highlight_option(MigrationMenuOption::TryNewModel);
                self.accept();
            }
            KeyCode::Char('2') => {
                self.highlight_option(MigrationMenuOption::UseExistingModel);
                self.reject();
            }
            KeyCode::Enter | KeyCode::Esc => self.confirm_selection(),
            _ => {}
        }
    }

    fn render_content(&self, column: &mut ColumnRenderable) {
        self.render_lines(&self.copy.content, column);
    }

    fn render_lines(&self, lines: &[Line<'static>], column: &mut ColumnRenderable) {
        for line in lines {
            column.push(
                Paragraph::new(line.clone())
                    .wrap(Wrap { trim: false })
                    .inset(Insets::vh(/*v*/ 0, /*h*/ 2)),
            );
        }
    }

    fn render_markdown_content(
        &self,
        markdown: &str,
        area_width: u16,
        column: &mut ColumnRenderable,
    ) {
        let horizontal_inset = 2;
        let content_width = area_width.saturating_sub(horizontal_inset * 2);
        let wrap_width = (content_width > 0).then_some(content_width as usize);
        let rendered = render_markdown_text_with_width(markdown, wrap_width);
        for line in rendered.lines {
            column.push(
                Paragraph::new(line)
                    .wrap(Wrap { trim: false })
                    .inset(Insets::vh(/*v*/ 0, horizontal_inset)),
            );
        }
    }
}

// Render the prompt on the terminal's alternate screen so exiting or cancelling
// does not leave a large blank region in the normal scrollback. This does not
// change the prompt's appearance – only where it is drawn.
struct AltScreenGuard<'a> {
    tui: &'a mut Tui,
}

impl<'a> AltScreenGuard<'a> {
    fn enter(tui: &'a mut Tui) -> Self {
        let _ = tui.enter_alt_screen();
        Self { tui }
    }
}

impl Drop for AltScreenGuard<'_> {
    fn drop(&mut self) {
        let _ = self.tui.leave_alt_screen();
    }
}

fn is_ctrl_exit_combo(key_event: KeyEvent) -> bool {
    key_event.modifiers.contains(KeyModifiers::CONTROL)
        && matches!(key_event.code, KeyCode::Char('c') | KeyCode::Char('d'))
}

fn fill_migration_markdown(template: &str, current_model: &str, target_model: &str) -> String {
    template
        .replace("{model_from}", current_model)
        .replace("{model_to}", target_model)
}

#[cfg(test)]
mod tests {
    use super::ModelMigrationCopy;
    use super::ModelMigrationScreen;
    use super::migration_copy_for_models;
    use crate::custom_terminal::Terminal;
    use crate::test_backend::VT100Backend;
    use crate::tui::FrameRequester;
    use crossterm::event::KeyCode;
    use crossterm::event::KeyEvent;
    use crossterm::event::KeyModifiers;
    use insta::assert_snapshot;
    use pretty_assertions::assert_eq;
    use ratatui::layout::Rect;
    use ratatui::widgets::FrameExt;

    #[test]
    fn prompt_snapshot() {
        let width: u16 = 60;
        let height: u16 = 28;
        let backend = VT100Backend::new(width, height);
        let mut terminal = Terminal::with_options(backend).expect("terminal");
        terminal.set_viewport_area(Rect::new(0, 0, width, height));

        let screen = ModelMigrationScreen::new(
            FrameRequester::test_dummy(),
            migration_copy_for_models(
                "gpt-5.1-codex-mini",
                "gpt-5.1-codex-max",
                /*model_link*/ None,
                Some(
                    "Upgrade to gpt-5.2-codex for the latest and greatest agentic coding model."
                        .to_string(),
                ),
                /*migration_markdown*/ None,
                "gpt-5.1-codex-max".to_string(),
                Some("Codex-optimized flagship for deep and fast reasoning.".to_string()),
                /*can_opt_out*/ true,
            ),
        );

        {
            let mut frame = terminal.get_frame();
            frame.render_widget_ref(&screen, frame.area());
        }
        terminal.flush().expect("flush");

        assert_snapshot!("model_migration_prompt", terminal.backend());
    }

    #[test]
    fn prompt_snapshot_gpt5_family() {
        let backend = VT100Backend::new(/*width*/ 65, /*height*/ 22);
        let mut terminal = Terminal::with_options(backend).expect("terminal");
        terminal.set_viewport_area(Rect::new(0, 0, 65, 22));

        let screen = ModelMigrationScreen::new(
            FrameRequester::test_dummy(),
            migration_copy_for_models(
                "gpt-5",
                "gpt-5.1",
                Some("https://www.codex.com/models/gpt-5.1".to_string()),
                /*migration_copy*/ None,
                /*migration_markdown*/ None,
                "gpt-5.1".to_string(),
                Some("Broad world knowledge with strong general reasoning.".to_string()),
                /*can_opt_out*/ false,
            ),
        );
        {
            let mut frame = terminal.get_frame();
            frame.render_widget_ref(&screen, frame.area());
        }
        terminal.flush().expect("flush");
        assert_snapshot!("model_migration_prompt_gpt5_family", terminal.backend());
    }

    #[test]
    fn escape_key_accepts_prompt() {
        let mut screen = ModelMigrationScreen::new(
            FrameRequester::test_dummy(),
            migration_copy_for_models(
                "gpt-old",
                "gpt-new",
                Some("https://www.codex.com/models/gpt-new".to_string()),
                /*migration_copy*/ None,
                /*migration_markdown*/ None,
                "gpt-new".to_string(),
                Some("Latest recommended model for better performance.".to_string()),
                /*can_opt_out*/ true,
            ),
        );

        // Simulate pressing Escape
        screen.handle_key(KeyEvent::new(
            KeyCode::Esc,
            crossterm::event::KeyModifiers::NONE,
        ));
        assert!(screen.is_done());
        // Esc should not be treated as Exit – it accepts like Enter.
        assert!(matches!(
            screen.outcome(),
            super::ModelMigrationOutcome::Accepted
        ));
    }

    #[test]
    fn selecting_use_existing_model_rejects_upgrade() {
        let mut screen = ModelMigrationScreen::new(
            FrameRequester::test_dummy(),
            migration_copy_for_models(
                "gpt-old",
                "gpt-new",
                Some("https://www.codex.com/models/gpt-new".to_string()),
                /*migration_copy*/ None,
                /*migration_markdown*/ None,
                "gpt-new".to_string(),
                Some("Latest recommended model for better performance.".to_string()),
                /*can_opt_out*/ true,
            ),
        );

        screen.handle_key(KeyEvent::new(
            KeyCode::Down,
            crossterm::event::KeyModifiers::NONE,
        ));
        screen.handle_key(KeyEvent::new(
            KeyCode::Enter,
            crossterm::event::KeyModifiers::NONE,
        ));

        assert!(screen.is_done());
        assert!(matches!(
            screen.outcome(),
            super::ModelMigrationOutcome::Rejected
        ));
    }

    #[test]
    fn markdown_prompt_keeps_long_url_tail_visible_when_narrow() {
        let long_url = "https://example.test/api/v1/projects/alpha-team/releases/2026-02-17/builds/1234567890/artifacts/reports/performance/summary/detail/with/a/very/long/path/tail42";
        let screen = ModelMigrationScreen::new(
            FrameRequester::test_dummy(),
            ModelMigrationCopy {
                heading: Vec::new(),
                content: Vec::new(),
                can_opt_out: false,
                markdown: Some(long_url.to_string()),
            },
        );

        let backend = VT100Backend::new(/*width*/ 40, /*height*/ 16);
        let mut terminal = Terminal::with_options(backend).expect("terminal");
        terminal.set_viewport_area(Rect::new(0, 0, 40, 16));

        {
            let mut frame = terminal.get_frame();
            frame.render_widget_ref(&screen, frame.area());
        }
        terminal.flush().expect("flush");

        let rendered = terminal.backend().to_string();
        assert!(
            rendered.contains("tail42"),
            "expected wrapped markdown URL tail to remain visible, got:\n{rendered}"
        );
    }

    #[test]
    fn migration_choice_and_hints_stay_visible_in_a_short_viewport() {
        let mut screen = ModelMigrationScreen::new(
            FrameRequester::test_dummy(),
            migration_copy_for_models(
                "gpt-old",
                "gpt-new",
                /*model_link*/ None,
                /*migration_copy*/ None,
                Some("# Introducing {model_to}\n\nA faster coding model with better reasoning.\n\nYou can keep using {model_from} if you prefer.".to_string()),
                "GPT New".to_string(),
                /*target_description*/ None,
                /*can_opt_out*/ true,
            ),
        );
        screen.handle_key(KeyEvent::new(KeyCode::Down, KeyModifiers::NONE));
        let (width, height) = (28, 12);
        let mut terminal =
            ratatui::Terminal::new(VT100Backend::new(width, height)).expect("terminal");
        terminal
            .draw(|frame| frame.render_widget_ref(&screen, frame.area()))
            .expect("render resized model migration picker");
        let rendered = terminal.backend().to_string();
        assert!(rendered.contains("› 2. Use existing model"));
        assert!(rendered.contains("enter/esc confirm"));
        assert!(!screen.is_done());
        assert_snapshot!(format!("model_picker_selected_{width}x{height}"), rendered);
        screen.handle_key(KeyEvent::new(KeyCode::Esc, KeyModifiers::NONE));
        assert_eq!(screen.outcome(), super::ModelMigrationOutcome::Rejected);
    }
}