use std::io;use std::sync::LazyLock;use crate::key_hint;use crate::key_hint::KeyBinding;use crate::key_hint::KeyBindingListExt;use codex_model_provider_info::DEFAULT_LMSTUDIO_PORT;use codex_model_provider_info::DEFAULT_OLLAMA_PORT;use codex_model_provider_info::LMSTUDIO_OSS_PROVIDER_ID;use codex_model_provider_info::OLLAMA_OSS_PROVIDER_ID;use crossterm::event::Event;use crossterm::event::KeyCode;use crossterm::event::KeyEvent;use crossterm::event::KeyEventKind;use crossterm::event::KeyModifiers;use crossterm::event::{self};use crossterm::execute;use crossterm::terminal::EnterAlternateScreen;use crossterm::terminal::LeaveAlternateScreen;use crossterm::terminal::disable_raw_mode;use crossterm::terminal::enable_raw_mode;use ratatui::Terminal;use ratatui::backend::CrosstermBackend;use ratatui::buffer::Buffer;use ratatui::layout::Alignment;use ratatui::layout::Constraint;use ratatui::layout::Direction;use ratatui::layout::Layout;use ratatui::layout::Margin;use ratatui::layout::Rect;use ratatui::prelude::*;use ratatui::style::Color;use ratatui::style::Modifier;use ratatui::style::Style;use ratatui::text::Line;use ratatui::text::Span;use ratatui::widgets::Paragraph;use ratatui::widgets::Widget;use ratatui::widgets::WidgetRef;use ratatui::widgets::Wrap;use std::time::Duration;#[derive(Clone)]struct ProviderOption { name: String, status: ProviderStatus,}#[derive(Clone)]enum ProviderStatus { Running, NotRunning, Unknown,}/// Options displayed in the *select* mode.////// The `key` is matched case-insensitively.struct SelectOption { label: Line<'static>, description: &'static str, key: KeyCode, provider_id: &'static str,}static OSS_SELECT_OPTIONS: LazyLock<Vec<SelectOption>> = LazyLock::new(|| { vec![ SelectOption { label: Line::from(vec!["L".underlined(), "M Studio".into()]), description: "Local LM Studio server (default port 1234)", key: KeyCode::Char('l'), provider_id: LMSTUDIO_OSS_PROVIDER_ID, }, SelectOption { label: Line::from(vec!["O".underlined(), "llama".into()]), description: "Local Ollama server (Responses API, default port 11434)", key: KeyCode::Char('o'), provider_id: OLLAMA_OSS_PROVIDER_ID, }, ]});// This startup wizard runs before the main TUI runtime keymap is available, so// it mirrors the built-in horizontal list defaults instead of reading config.// The shared matcher still covers raw C0 Ctrl-H/Ctrl-L terminal reports.const MOVE_LEFT_KEYS: [KeyBinding; 2] = [ key_hint::plain(KeyCode::Left), key_hint::ctrl(KeyCode::Char('h')),];const MOVE_RIGHT_KEYS: [KeyBinding; 2] = [ key_hint::plain(KeyCode::Right), key_hint::ctrl(KeyCode::Char('l')),];pub struct OssSelectionWidget<'a> { select_options: &'a Vec<SelectOption>, confirmation_prompt: Paragraph<'a>, /// Currently selected index in *select* mode. selected_option: usize, /// Set to `true` once a decision has been sent – the parent view can then /// remove this widget from its queue. done: bool, selection: Option<String>,}impl OssSelectionWidget<'_> { fn new(lmstudio_status: ProviderStatus, ollama_status: ProviderStatus) -> io::Result<Self> { let providers = vec![ ProviderOption { name: "LM Studio".to_string(), status: lmstudio_status, }, ProviderOption { name: "Ollama (Responses)".to_string(), status: ollama_status.clone(), }, ProviderOption { name: "Ollama (Chat)".to_string(), status: ollama_status, }, ]; let mut contents: Vec<Line> = vec![ Line::from(vec![ "? ".fg(Color::Blue), "Select an open-source provider".bold(), ]), Line::from(""), Line::from(" Choose which local AI server to use for your session."), Line::from(""), ]; // Add status indicators for each provider for provider in &providers { let (status_symbol, status_color) = get_status_symbol_and_color(&provider.status); contents.push(Line::from(vec![ Span::raw(" "), Span::styled(status_symbol, Style::default().fg(status_color)), Span::raw(format!(" {} ", provider.name)), ])); } contents.push(Line::from("")); contents.push(Line::from(" ● Running ○ Not Running").add_modifier(Modifier::DIM)); contents.push(Line::from("")); contents.push( Line::from(" Press Enter to select • Ctrl+C to exit").add_modifier(Modifier::DIM), ); let confirmation_prompt = Paragraph::new(contents).wrap(Wrap { trim: false }); Ok(Self { select_options: &OSS_SELECT_OPTIONS, confirmation_prompt, selected_option: 0, done: false, selection: None, }) } fn get_confirmation_prompt_height(&self, width: u16) -> u16 { // Should cache this for last value of width. self.confirmation_prompt.line_count(width) as u16 } /// Process a `KeyEvent` coming from crossterm. Always consumes the event /// while the modal is visible. /// Process a key event originating from crossterm. As the modal fully /// captures input while visible, we don't need to report whether the event /// was consumed—callers can assume it always is. pub fn handle_key_event(&mut self, key: KeyEvent) -> Option<String> { if key.kind == KeyEventKind::Press { self.handle_select_key(key); } if self.done { self.selection.clone() } else { None } } /// Normalize a key for comparison. /// - For `KeyCode::Char`, converts to lowercase for case-insensitive matching. /// - Other key codes are returned unchanged. fn normalize_keycode(code: KeyCode) -> KeyCode { match code { KeyCode::Char(c) => KeyCode::Char(c.to_ascii_lowercase()), other => other, } } fn handle_select_key(&mut self, key_event: KeyEvent) { match key_event { KeyEvent { code: KeyCode::Char('c'), modifiers, .. } if modifiers.contains(KeyModifiers::CONTROL) => { self.send_decision("__CANCELLED__".to_string()); } _ if MOVE_LEFT_KEYS.is_pressed(key_event) => { self.selected_option = (self.selected_option + self.select_options.len() - 1) % self.select_options.len(); } _ if MOVE_RIGHT_KEYS.is_pressed(key_event) => { self.selected_option = (self.selected_option + 1) % self.select_options.len(); } KeyEvent { code: KeyCode::Enter, .. } => { let opt = &self.select_options[self.selected_option]; self.send_decision(opt.provider_id.to_string()); } KeyEvent { code: KeyCode::Esc, .. } => { self.send_decision(LMSTUDIO_OSS_PROVIDER_ID.to_string()); } KeyEvent { code, .. } => { let other = code; let normalized = Self::normalize_keycode(other); if let Some(opt) = self .select_options .iter() .find(|opt| Self::normalize_keycode(opt.key) == normalized) { self.send_decision(opt.provider_id.to_string()); } } } } fn send_decision(&mut self, selection: String) { self.selection = Some(selection); self.done = true; } /// Returns `true` once the user has made a decision and the widget no /// longer needs to be displayed. pub fn is_complete(&self) -> bool { self.done } pub fn desired_height(&self, width: u16) -> u16 { self.get_confirmation_prompt_height(width) + self.select_options.len() as u16 }}impl WidgetRef for &OssSelectionWidget<'_> { fn render_ref(&self, area: Rect, buf: &mut Buffer) { let prompt_height = self.get_confirmation_prompt_height(area.width); let [prompt_chunk, response_chunk] = Layout::default() .direction(Direction::Vertical) .constraints([Constraint::Length(prompt_height), Constraint::Min(0)]) .areas(area); let lines: Vec<Line> = self .select_options .iter() .enumerate() .map(|(idx, opt)| { let style = if idx == self.selected_option { Style::new().bg(Color::Cyan).fg(Color::Black) } else { Style::new().bg(Color::DarkGray) }; opt.label.clone().alignment(Alignment::Center).style(style) }) .collect(); let [title_area, button_area, description_area] = Layout::vertical([ Constraint::Length(1), Constraint::Length(1), Constraint::Min(0), ]) .areas(response_chunk.inner(Margin::new(1, 0))); Line::from("Select provider?").render(title_area, buf); self.confirmation_prompt.clone().render(prompt_chunk, buf); let areas = Layout::horizontal( lines .iter() .map(|l| Constraint::Length(l.width() as u16 + 2)), ) .spacing(1) .split(button_area); for (idx, area) in areas.iter().enumerate() { let line = &lines[idx]; line.render(*area, buf); } Line::from(self.select_options[self.selected_option].description) .style(Style::new().italic().fg(Color::DarkGray)) .render(description_area.inner(Margin::new(1, 0)), buf); }}fn get_status_symbol_and_color(status: &ProviderStatus) -> (&'static str, Color) { match status { ProviderStatus::Running => ("●", Color::Green), ProviderStatus::NotRunning => ("○", Color::Red), ProviderStatus::Unknown => ("?", Color::Yellow), }}pub(crate) struct OssProviderSelection { pub(crate) provider: String, pub(crate) manually_selected: bool,}pub async fn select_oss_provider() -> io::Result<OssProviderSelection> { // Check provider statuses first let lmstudio_status = check_lmstudio_status().await; let ollama_status = check_ollama_status().await; // Autoselect if only one is running match (&lmstudio_status, &ollama_status) { (ProviderStatus::Running, ProviderStatus::NotRunning) => { let provider = LMSTUDIO_OSS_PROVIDER_ID.to_string(); return Ok(OssProviderSelection { provider, manually_selected: false, }); } (ProviderStatus::NotRunning, ProviderStatus::Running) => { let provider = OLLAMA_OSS_PROVIDER_ID.to_string(); return Ok(OssProviderSelection { provider, manually_selected: false, }); } _ => { // Both running or both not running - show UI } } let mut widget = OssSelectionWidget::new(lmstudio_status, ollama_status)?; enable_raw_mode()?; let mut stdout = io::stdout(); execute!(stdout, EnterAlternateScreen)?; let backend = CrosstermBackend::new(stdout); let mut terminal = Terminal::new(backend)?; let result = loop { terminal.draw(|f| { (&widget).render_ref(f.area(), f.buffer_mut()); })?; if let Event::Key(key_event) = event::read()? && let Some(selection) = widget.handle_key_event(key_event) { break Ok(OssProviderSelection { provider: selection, manually_selected: true, }); } }; disable_raw_mode()?; execute!(terminal.backend_mut(), LeaveAlternateScreen)?; result}async fn check_lmstudio_status() -> ProviderStatus { match check_port_status(DEFAULT_LMSTUDIO_PORT).await { Ok(true) => ProviderStatus::Running, Ok(false) => ProviderStatus::NotRunning, Err(_) => ProviderStatus::Unknown, }}async fn check_ollama_status() -> ProviderStatus { match check_port_status(DEFAULT_OLLAMA_PORT).await { Ok(true) => ProviderStatus::Running, Ok(false) => ProviderStatus::NotRunning, Err(_) => ProviderStatus::Unknown, }}async fn check_port_status(port: u16) -> io::Result<bool> { let client = reqwest::Client::builder() .timeout(Duration::from_secs(2)) .build() .map_err(io::Error::other)?; let url = format!("http://localhost:{port}"); match client.get(&url).send().await { Ok(response) => Ok(response.status().is_success()), Err(_) => Ok(false), // Connection failed = not running }}#[cfg(test)]mod tests { use super::*; #[test] fn ctrl_h_l_move_provider_selection() { let mut widget = OssSelectionWidget::new(ProviderStatus::Unknown, ProviderStatus::Unknown) .expect("widget should initialize"); assert_eq!(widget.selected_option, 0); widget.handle_key_event(KeyEvent::new(KeyCode::Char('l'), KeyModifiers::CONTROL)); assert_eq!(widget.selected_option, 1); widget.handle_key_event(KeyEvent::new(KeyCode::Char('h'), KeyModifiers::CONTROL)); assert_eq!(widget.selected_option, 0); }}