opi_coding_agent/

interactive.rs

//! Interactive TUI mode using opi-tui for terminal rendering.
//!
//! The agent prompt runs in a spawned tokio task while the TUI render loop
//! continues to poll crossterm events and redraw at ~20 fps. Agent callbacks
//! update shared `TuiState`, which the render loop reads each frame.

use std::io;
use std::sync::{Arc, Mutex};

use crossterm::{
    event::{self, Event, KeyCode, KeyEventKind, KeyModifiers},
    terminal::{self, EnterAlternateScreen, LeaveAlternateScreen},
};
use ratatui::prelude::*;

use opi_agent::event::AgentEvent;
use opi_agent::loop_types::AgentError;
use opi_agent::message::AgentMessage;
use opi_ai::message::{AssistantContent, Message};
use opi_ai::stream::AssistantStreamEvent;
use opi_tui::terminal_image::{
    CapabilitySource, TerminalGraphicsProtocol, detect_graphics_protocol,
};
use opi_tui::{
    AppState, Key, KeyCombo, Keybindings, Message as TuiMessage, Role as TuiRole, SelectListState,
    Shell, Theme, ToolCallStatus, resolve_theme,
};
use opi_tui::{ImageData, ImagePayload, MediaType as TuiMediaType};

use crate::harness::CodingHarness;

/// Shared state mutated by the agent callback and read by the TUI render loop.
struct TuiState {
    messages: Vec<TuiMessage>,
    input_text: String,
    app_state: AppState,
    model: String,
    active_tool: Option<(String, String, ToolCallStatus)>,
    /// True when a TextDelta has been received for the current streaming cycle.
    /// Prevents MessageEnd from pushing a duplicate text message.
    streaming_started: bool,
    theme: Theme,
    keybindings: Keybindings,
    total_tokens: u64,
    cost_usd: Option<f64>,
    graphics_protocol: TerminalGraphicsProtocol,
    picker: Option<PickerOverlay>,
}

#[derive(Clone)]
struct PickerOverlay {
    kind: PickerKind,
    title: String,
    state: SelectListState,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum PickerKind {
    Model,
    Session,
    Branch,
}

#[derive(Debug, PartialEq, Eq)]
enum PickerAction {
    SelectModel(String),
    SelectSession(String),
    SelectBranch(String),
    Cancel,
}

pub async fn run_interactive_tui(
    harness: CodingHarness,
    model: String,
    theme_name: &str,
    keybindings: Keybindings,
) -> Result<(), Box<dyn std::error::Error>> {
    let theme = resolve_interactive_theme(&harness, theme_name);
    if theme.name != theme_name {
        eprintln!("opi: warning: unknown theme {theme_name:?}, using default");
    }
    let graphics_protocol = detect_graphics_protocol(
        std::env::var("TERM").ok().as_deref(),
        std::env::var("TERM_PROGRAM").ok().as_deref(),
        std::env::var("TERM_FEATURES").ok().as_deref(),
        &CapabilitySource::EnvVars,
    );
    let state = Arc::new(Mutex::new(TuiState {
        messages: Vec::new(),
        input_text: String::new(),
        app_state: AppState::Idle,
        model: model.clone(),
        active_tool: None,
        streaming_started: false,
        theme,
        keybindings,
        total_tokens: 0,
        cost_usd: None,
        graphics_protocol,
        picker: None,
    }));

    // Wire agent events into shared state before wrapping harness
    let state_clone = state.clone();
    let mut harness = harness;
    harness.subscribe(Box::new(move |event| {
        let mut s = state_clone.lock().unwrap();
        match event {
            AgentEvent::MessageStart { .. } => {
                s.app_state = AppState::Streaming;
                s.streaming_started = false;
            }
            AgentEvent::MessageUpdate {
                assistant_event, ..
            } => {
                if let AssistantStreamEvent::TextDelta { delta, .. } = assistant_event.as_ref() {
                    if !s.streaming_started {
                        s.messages
                            .push(TuiMessage::new(TuiRole::Assistant, delta.clone()));
                        s.streaming_started = true;
                    } else if let Some(msg) = s.messages.last_mut() {
                        msg.content.push_str(delta);
                    }
                }
            }
            AgentEvent::MessageEnd {
                message: AgentMessage::Llm(Message::Assistant(a)),
            } => {
                s.total_tokens += a.usage.total_tokens();
                for content in &a.content {
                    match content {
                        AssistantContent::Text { text } if !s.streaming_started => {
                            s.messages
                                .push(TuiMessage::new(TuiRole::Assistant, text.clone()));
                        }
                        AssistantContent::ToolCall { tool_call } => {
                            s.active_tool = Some((
                                tool_call.name.clone(),
                                tool_call.arguments.clone(),
                                ToolCallStatus::Running,
                            ));
                        }
                        _ => {}
                    }
                }
                s.streaming_started = false;
            }
            AgentEvent::ToolExecutionStart {
                tool_name, args, ..
            } => {
                s.app_state = AppState::ToolExecuting;
                s.active_tool = Some((
                    tool_name.clone(),
                    format!("{args}"),
                    ToolCallStatus::Running,
                ));
            }
            AgentEvent::ToolExecutionEnd {
                tool_name,
                is_error,
                details,
                result,
                ..
            } => {
                // Render diff for edit tool results that have before/after details.
                if !is_error
                    && tool_name == "edit"
                    && let Some(d) = details
                    && let (Some(path), Some(before), Some(after)) =
                        (d.get("path"), d.get("before"), d.get("after"))
                {
                    let path_str = path.as_str().unwrap_or("unknown");
                    let before_str = before.as_str().unwrap_or("");
                    let after_str = after.as_str().unwrap_or("");
                    s.messages
                        .push(TuiMessage::diff(path_str, before_str, after_str));
                }
                // Extract image content from tool result.
                let protocol = s.graphics_protocol;
                if let Some(content_arr) = result.as_array() {
                    for item in content_arr {
                        if item.get("type").and_then(|v| v.as_str()) == Some("image")
                            && let Some(source) = item.get("source")
                        {
                            let bytes = if source.get("type").and_then(|v| v.as_str())
                                == Some("bytes")
                            {
                                source
                                    .get("data")
                                    .and_then(|v| v.as_array())
                                    .map(|arr| {
                                        arr.iter()
                                            .filter_map(|v| v.as_u64().map(|n| n as u8))
                                            .collect::<Vec<u8>>()
                                    })
                                    .unwrap_or_default()
                            } else if source.get("type").and_then(|v| v.as_str()) == Some("base64")
                            {
                                use base64::Engine;
                                source
                                    .get("data")
                                    .and_then(|v| v.as_str())
                                    .and_then(|d| {
                                        base64::engine::general_purpose::STANDARD.decode(d).ok()
                                    })
                                    .unwrap_or_default()
                            } else {
                                vec![]
                            };
                            if !bytes.is_empty() {
                                let media_type = item.get("media_type").and_then(|v| v.as_str());
                                let tui_media = match media_type {
                                    Some("image/jpeg") => TuiMediaType::Jpeg,
                                    Some("image/gif") => TuiMediaType::Gif,
                                    Some("image/webp") => TuiMediaType::WebP,
                                    _ => TuiMediaType::Png,
                                };
                                let image_data = ImageData {
                                    bytes,
                                    media_type: tui_media,
                                    width: None,
                                    height: None,
                                };
                                s.messages.push(TuiMessage::image(
                                    TuiRole::Tool,
                                    ImagePayload {
                                        data: image_data,
                                        protocol,
                                    },
                                ));
                            }
                        }
                    }
                }
                if let Some((name, args, _)) = &s.active_tool
                    && name == tool_name
                {
                    let status = if *is_error {
                        ToolCallStatus::Error("failed".into())
                    } else {
                        ToolCallStatus::Success
                    };
                    s.active_tool = Some((name.clone(), args.clone(), status));
                }
                s.app_state = AppState::Streaming;
            }
            AgentEvent::AgentEnd { .. } => {
                s.app_state = AppState::Idle;
                s.active_tool = None;
            }
            AgentEvent::TurnStart => {
                s.app_state = AppState::Thinking;
            }
            AgentEvent::CompactionStart { reason } => {
                s.messages.push(TuiMessage::new(
                    TuiRole::System,
                    format!("[compaction started: {reason:?}]"),
                ));
            }
            AgentEvent::CompactionEnd {
                reason,
                result,
                aborted,
                error_message,
            } => {
                let summary = if *aborted {
                    format!(
                        "[compaction aborted ({reason:?}): {}]",
                        error_message.clone().unwrap_or_default()
                    )
                } else if let Some(r) = result {
                    format!(
                        "[compaction done ({reason:?}): {} -> {} tokens]",
                        r.tokens_before, r.tokens_after
                    )
                } else {
                    format!("[compaction done ({reason:?})]")
                };
                s.messages.push(TuiMessage::new(TuiRole::System, summary));
            }
            AgentEvent::SessionPersistError { message } => {
                s.messages.push(TuiMessage::new(
                    TuiRole::System,
                    format!("[session persist error: {message}]"),
                ));
            }
            _ => {}
        }
    }));

    let harness = Arc::new(tokio::sync::Mutex::new(harness));

    // Setup terminal
    terminal::enable_raw_mode()?;
    let mut stdout = io::stdout();
    crossterm::execute!(stdout, EnterAlternateScreen)?;
    let backend = CrosstermBackend::new(stdout);
    let mut terminal = Terminal::new(backend)?;

    // Main TUI loop
    let result = tui_event_loop(&mut terminal, &harness, &state).await;

    // Restore terminal
    terminal::disable_raw_mode()?;
    crossterm::execute!(terminal.backend_mut(), LeaveAlternateScreen)?;
    terminal.show_cursor()?;

    result
}

async fn tui_event_loop(
    terminal: &mut Terminal<CrosstermBackend<io::Stdout>>,
    harness: &Arc<tokio::sync::Mutex<CodingHarness>>,
    state: &Arc<Mutex<TuiState>>,
) -> Result<(), Box<dyn std::error::Error>> {
    let mut pending: Option<tokio::task::JoinHandle<Result<Vec<AgentMessage>, AgentError>>> = None;
    let mut cancel_token = harness.lock().await.cancel_token();

    loop {
        // Render current state
        {
            let s = state.lock().unwrap();
            let shell = build_shell(&s);
            terminal.draw(|frame| frame.render_widget(shell, frame.area()))?;
        }

        // Check if pending prompt finished (non-blocking)
        if let Some(handle) = &mut pending
            && handle.is_finished()
        {
            match handle.await {
                Ok(Ok(_messages)) => {
                    let mut s = state.lock().unwrap();
                    s.app_state = AppState::Idle;
                }
                Ok(Err(AgentError::Cancelled)) => {
                    let mut s = state.lock().unwrap();
                    s.app_state = AppState::Idle;
                }
                Ok(Err(e)) => {
                    let mut s = state.lock().unwrap();
                    s.messages
                        .push(TuiMessage::new(TuiRole::System, format!("error: {e}")));
                    s.app_state = AppState::Idle;
                }
                Err(e) => {
                    let mut s = state.lock().unwrap();
                    s.messages
                        .push(TuiMessage::new(TuiRole::System, format!("error: {e}")));
                    s.app_state = AppState::Idle;
                }
            }

            // Refresh cost from the harness session (pricing lookup may yield
            // a number; if the model isn't in the table we leave it as-is).
            {
                let h = harness.lock().await;
                if let Some(session) = h.session()
                    && let Some(cost) = session.cost_summary()
                {
                    state.lock().unwrap().cost_usd = Some(cost.total_cost());
                }
            }

            // Refresh cancel token — Agent::maybe_reset_cancel() creates a new one
            // after cancellation, so the old token would be stale.
            cancel_token = harness.lock().await.cancel_token();
            pending = None;
        }

        // Poll for terminal events (non-blocking with timeout)
        if event::poll(std::time::Duration::from_millis(50))?
            && let Event::Key(key) = event::read()?
        {
            if key.kind != KeyEventKind::Press {
                continue;
            }
            let kb = state.lock().unwrap().keybindings.clone();
            if let Some(action) = {
                let mut s = state.lock().unwrap();
                handle_picker_key(&mut s, key.code)
            } {
                match action {
                    PickerAction::SelectModel(model) => {
                        let mut h = harness.lock().await;
                        h.set_model(model.clone());
                        let mut s = state.lock().unwrap();
                        s.model = model.clone();
                        s.messages.push(TuiMessage::new(
                            TuiRole::System,
                            format!("[model switched: {model}]"),
                        ));
                    }
                    PickerAction::SelectSession(session_id) => {
                        let result = {
                            let mut h = harness.lock().await;
                            h.resume_session_id(&session_id)
                        };
                        let mut s = state.lock().unwrap();
                        match result {
                            Ok(count) => s.messages.push(TuiMessage::new(
                                TuiRole::System,
                                format!("[session resumed: {session_id}, {count} messages]"),
                            )),
                            Err(e) => s.messages.push(TuiMessage::new(
                                TuiRole::System,
                                format!("[session resume failed: {e}]"),
                            )),
                        }
                    }
                    PickerAction::SelectBranch(tip_id) => {
                        let result = {
                            let mut h = harness.lock().await;
                            h.resume_session_branch_tip(&tip_id)
                        };
                        let mut s = state.lock().unwrap();
                        match result {
                            Ok(count) => s.messages.push(TuiMessage::new(
                                TuiRole::System,
                                format!("[branch selected: {tip_id}, {count} messages]"),
                            )),
                            Err(e) => s.messages.push(TuiMessage::new(
                                TuiRole::System,
                                format!("[branch select failed: {e}]"),
                            )),
                        }
                    }
                    PickerAction::Cancel => {}
                }
                continue;
            }

            if matches_key_combo(key.code, key.modifiers, &kb.submit) {
                // Ignore submit while agent is running
                if pending.is_some() {
                    continue;
                }

                let input = {
                    let mut s = state.lock().unwrap();
                    let text = s.input_text.trim().to_string();
                    s.input_text.clear();
                    text
                };

                if input == "exit" || input == "quit" {
                    // Cancel any pending task on exit
                    if let Some(handle) = pending.take() {
                        cancel_token.cancel();
                        let _ = handle.await;
                    }
                    return Ok(());
                }
                if input.is_empty() {
                    continue;
                }

                if input == "/model" {
                    let items = {
                        let h = harness.lock().await;
                        h.model_picker_items()
                    };
                    let mut s = state.lock().unwrap();
                    if items.is_empty() {
                        s.messages.push(TuiMessage::new(
                            TuiRole::System,
                            "[model picker: no models available]",
                        ));
                    } else {
                        s.picker = Some(PickerOverlay {
                            kind: PickerKind::Model,
                            title: "Select model".into(),
                            state: SelectListState::new(items),
                        });
                    }
                    continue;
                }

                if input == "/session" {
                    let dir = crate::session_cli::session_dir();
                    let items = crate::picker::session_picker_items(&dir).unwrap_or_default();
                    let mut s = state.lock().unwrap();
                    if items.is_empty() {
                        s.messages.push(TuiMessage::new(
                            TuiRole::System,
                            "[session picker: no sessions available]",
                        ));
                    } else {
                        s.picker = Some(PickerOverlay {
                            kind: PickerKind::Session,
                            title: "Resume session".into(),
                            state: SelectListState::new(items),
                        });
                    }
                    continue;
                }

                if input == "/branch" {
                    let items_result = {
                        let h = harness.lock().await;
                        h.branch_picker_items()
                    };
                    let mut s = state.lock().unwrap();
                    match items_result {
                        Ok(items) if items.is_empty() => {
                            s.messages.push(TuiMessage::new(
                                TuiRole::System,
                                "[branch picker: no branches available]",
                            ));
                        }
                        Ok(items) => {
                            s.picker = Some(PickerOverlay {
                                kind: PickerKind::Branch,
                                title: "Select branch".into(),
                                state: SelectListState::new(items),
                            });
                        }
                        Err(e) => {
                            s.messages.push(TuiMessage::new(
                                TuiRole::System,
                                format!("[branch picker failed: {e}]"),
                            ));
                        }
                    }
                    continue;
                }

                if let Some(rest) = input.strip_prefix("/image ") {
                    let path = rest.trim();
                    if path.is_empty() {
                        let mut s = state.lock().unwrap();
                        s.messages.push(TuiMessage::new(
                            TuiRole::System,
                            "[/image: usage: /image <path>]".to_string(),
                        ));
                    } else {
                        let image_path = std::path::PathBuf::from(path);
                        let max_bytes = {
                            let h = harness.lock().await;
                            h.config().defaults.max_image_bytes
                        };
                        match crate::image::load_image_with_limit(&image_path, max_bytes) {
                            Ok(img) => {
                                harness.lock().await.queue_images(vec![img]);
                                let mut s = state.lock().unwrap();
                                s.messages.push(TuiMessage::new(
                                    TuiRole::System,
                                    format!("[image queued: {}]", image_path.display()),
                                ));
                            }
                            Err(e) => {
                                let mut s = state.lock().unwrap();
                                s.messages.push(TuiMessage::new(
                                    TuiRole::System,
                                    format!("[/image error: {e}]"),
                                ));
                            }
                        }
                    }
                    continue;
                }

                // Add user message to display
                {
                    let mut s = state.lock().unwrap();
                    s.messages
                        .push(TuiMessage::new(TuiRole::User, input.clone()));
                    s.app_state = AppState::Thinking;
                }

                // Spawn agent prompt in background task
                let h = harness.clone();
                let handle = tokio::spawn(async move {
                    let mut h = h.lock().await;
                    let pending = h.take_pending_images();
                    if pending.is_empty() {
                        h.prompt(&input).await
                    } else {
                        let mut content = vec![opi_ai::message::InputContent::Text { text: input }];
                        content.extend(pending);
                        h.prompt_with_content(content).await
                    }
                });
                pending = Some(handle);
            } else if matches_key_combo(key.code, key.modifiers, &kb.abort) {
                if pending.is_some() {
                    cancel_token.cancel();
                } else {
                    return Ok(());
                }
            } else if matches_key_combo(key.code, key.modifiers, &kb.new_line) {
                if pending.is_none() {
                    state.lock().unwrap().input_text.push('\n');
                }
            } else {
                match key.code {
                    KeyCode::Char(c) if pending.is_none() => {
                        state.lock().unwrap().input_text.push(c);
                    }
                    KeyCode::Backspace if pending.is_none() => {
                        state.lock().unwrap().input_text.pop();
                    }
                    _ => {}
                }
            }
        }
    }
}

fn resolve_interactive_theme(harness: &CodingHarness, theme_name: &str) -> Theme {
    harness
        .resolve_theme(theme_name)
        .unwrap_or_else(|_| resolve_theme(theme_name))
}

fn build_shell(s: &TuiState) -> Shell {
    let mut shell = Shell::new(s.model.clone())
        .input_text(s.input_text.clone())
        .state(s.app_state)
        .theme(s.theme.clone());

    if s.total_tokens > 0 {
        shell = shell.token_count(s.total_tokens);
    }

    if let Some(cost) = s.cost_usd {
        shell = shell.cost_usd(cost);
    }

    if !s.messages.is_empty() {
        shell = shell.messages(s.messages.clone());
    }

    if let Some((name, args, status)) = &s.active_tool {
        shell = shell.active_tool(name.clone(), args.clone(), status.clone());
    }

    if let Some(picker) = &s.picker {
        shell = shell.picker(picker.title.clone(), picker.state.clone());
    }

    shell
}

fn handle_picker_key(s: &mut TuiState, code: KeyCode) -> Option<PickerAction> {
    let picker = s.picker.as_mut()?;
    match code {
        KeyCode::Esc => {
            s.picker = None;
            Some(PickerAction::Cancel)
        }
        KeyCode::Enter => {
            let item = picker.state.confirm().cloned();
            let kind = picker.kind;
            s.picker = None;
            match (kind, item) {
                (PickerKind::Model, Some(item)) => Some(PickerAction::SelectModel(item.id)),
                (PickerKind::Session, Some(item)) => Some(PickerAction::SelectSession(item.id)),
                (PickerKind::Branch, Some(item)) => Some(PickerAction::SelectBranch(item.id)),
                (_, None) => Some(PickerAction::Cancel),
            }
        }
        KeyCode::Down => {
            picker.state.move_down();
            None
        }
        KeyCode::Up => {
            picker.state.move_up();
            None
        }
        KeyCode::PageDown => {
            picker.state.page_down(10);
            None
        }
        KeyCode::PageUp => {
            picker.state.page_up(10);
            None
        }
        KeyCode::Backspace => {
            let mut filter = picker.state.filter().to_string();
            filter.pop();
            picker.state.set_filter(filter);
            None
        }
        KeyCode::Char(c) => {
            let mut filter = picker.state.filter().to_string();
            filter.push(c);
            picker.state.set_filter(filter);
            None
        }
        _ => None,
    }
}

fn matches_key_combo(code: KeyCode, modifiers: KeyModifiers, combo: &KeyCombo) -> bool {
    let key_matches = match (code, &combo.key) {
        (KeyCode::Enter, Key::Enter) => true,
        (KeyCode::Esc, Key::Escape) => true,
        (KeyCode::Tab, Key::Tab) => true,
        (KeyCode::Backspace, Key::Backspace) => true,
        (KeyCode::Char(c), Key::Char(expected)) => c == *expected,
        _ => false,
    };
    if !key_matches {
        return false;
    }
    combo.modifiers.alt == modifiers.contains(KeyModifiers::ALT)
        && combo.modifiers.ctrl == modifiers.contains(KeyModifiers::CONTROL)
        && combo.modifiers.shift == modifiers.contains(KeyModifiers::SHIFT)
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::config::OpiConfig;
    use crate::resource::{DiscoveryLayer, ResourceDiscoveryLayers};
    use opi_ai::test_support::MockProvider;
    use opi_tui::SelectItem;

    fn state_with_picker(kind: PickerKind) -> TuiState {
        TuiState {
            messages: Vec::new(),
            input_text: String::new(),
            app_state: AppState::Idle,
            model: "mock:old".into(),
            active_tool: None,
            streaming_started: false,
            theme: Theme::default(),
            keybindings: Keybindings::default(),
            total_tokens: 0,
            cost_usd: None,
            graphics_protocol: TerminalGraphicsProtocol::Fallback,
            picker: Some(PickerOverlay {
                kind,
                title: "Pick".into(),
                state: SelectListState::new(vec![SelectItem {
                    id: "mock:new".into(),
                    display: "New".into(),
                    metadata: "mock".into(),
                }]),
            }),
        }
    }

    #[test]
    fn model_picker_enter_returns_selected_model() {
        let mut state = state_with_picker(PickerKind::Model);
        let action = handle_picker_key(&mut state, KeyCode::Enter);
        assert_eq!(action, Some(PickerAction::SelectModel("mock:new".into())));
        assert!(state.picker.is_none());
    }

    #[test]
    fn session_picker_enter_returns_selected_session() {
        let mut state = state_with_picker(PickerKind::Session);
        let action = handle_picker_key(&mut state, KeyCode::Enter);
        assert_eq!(action, Some(PickerAction::SelectSession("mock:new".into())));
        assert!(state.picker.is_none());
    }

    #[test]
    fn branch_picker_enter_returns_selected_tip() {
        let mut state = state_with_picker(PickerKind::Branch);
        let action = handle_picker_key(&mut state, KeyCode::Enter);
        assert_eq!(action, Some(PickerAction::SelectBranch("mock:new".into())));
        assert!(state.picker.is_none());
    }

    #[test]
    fn interactive_theme_resolver_uses_harness_discovered_themes() {
        let tmp = tempfile::tempdir().unwrap();
        let theme_dir = tmp.path().join("operator-theme");
        std::fs::create_dir_all(&theme_dir).unwrap();
        std::fs::write(
            theme_dir.join("theme.toml"),
            r##"
name = "operator-theme"
description = "Operator theme"

[colors]
role_user = "Red"
status_bg = "#1a1a2e"
"##,
        )
        .unwrap();

        let provider = MockProvider::new("mock", vec![]);
        let harness = CodingHarness::builder(
            Box::new(provider),
            "mock:mock-model".into(),
            OpiConfig::default(),
            tmp.path().to_path_buf(),
        )
        .resource_layers(ResourceDiscoveryLayers {
            themes: vec![DiscoveryLayer {
                root: theme_dir,
                subdirectory: None,
                precedence: 2,
            }],
            ..Default::default()
        })
        .build();

        let theme = resolve_interactive_theme(&harness, "operator-theme");

        assert_eq!(theme.name, "operator-theme");
    }
}