dot/tui/

mod.rs

pub mod app;
pub mod event;
pub mod input;
pub mod markdown;
pub mod theme;
pub mod tools;
pub mod ui;
pub mod ui_popups;
pub mod ui_tools;
pub mod widgets;

use std::sync::Arc;
use std::time::Instant;

use anyhow::Result;
use crossterm::{execute, terminal};
use tokio::sync::{Mutex, mpsc};

use crate::agent::{Agent, AgentProfile};
use crate::config::Config;
use crate::db::Db;
use crate::provider::Provider;
use crate::tools::ToolRegistry;

use app::{App, ChatMessage};
use event::{AppEvent, EventHandler};
use input::InputAction;
use widgets::{AgentEntry, SessionEntry, time_ago};

pub struct ExitInfo {
    pub conversation_id: String,
    pub title: Option<String>,
}

pub async fn run(
    config: Config,
    providers: Vec<Box<dyn Provider>>,
    db: Db,
    tools: ToolRegistry,
    profiles: Vec<AgentProfile>,
    cwd: String,
    resume_id: Option<String>,
) -> Result<()> {
    terminal::enable_raw_mode()?;
    let mut stdout = std::io::stderr();
    execute!(
        stdout,
        terminal::EnterAlternateScreen,
        crossterm::event::EnableMouseCapture,
        crossterm::event::EnableBracketedPaste
    )?;
    let backend = ratatui::backend::CrosstermBackend::new(stdout);
    let mut terminal = ratatui::Terminal::new(backend)?;

    let result = run_app(
        &mut terminal,
        config,
        providers,
        db,
        tools,
        profiles,
        cwd,
        resume_id,
    )
    .await;

    terminal::disable_raw_mode()?;
    execute!(
        std::io::stderr(),
        terminal::LeaveAlternateScreen,
        crossterm::event::DisableMouseCapture,
        crossterm::event::DisableBracketedPaste
    )?;
    terminal.show_cursor()?;

    if let Ok(ref info) = result {
        print_exit_screen(info);
    }

    result.map(|_| ())
}

fn print_exit_screen(info: &ExitInfo) {
    let title = info.title.as_deref().unwrap_or("untitled session");
    let id = &info.conversation_id;
    println!();
    println!("  \x1b[2mSession\x1b[0m   {}", title);
    println!("  \x1b[2mResume\x1b[0m    dot -s {}", id);
    println!();
}

#[allow(clippy::too_many_arguments)]
async fn run_app(
    terminal: &mut ratatui::Terminal<ratatui::backend::CrosstermBackend<std::io::Stderr>>,
    config: Config,
    providers: Vec<Box<dyn Provider>>,
    db: Db,
    tools: ToolRegistry,
    profiles: Vec<AgentProfile>,
    cwd: String,
    resume_id: Option<String>,
) -> Result<ExitInfo> {
    let model_name = providers[0].model().to_string();
    let provider_name = providers[0].name().to_string();
    let agent_name = profiles
        .first()
        .map(|p| p.name.clone())
        .unwrap_or_else(|| "dot".to_string());

    let history = db.get_user_message_history(500).unwrap_or_default();

    let agents_context = crate::context::AgentsContext::load(&cwd, &config.context);
    let agent = Arc::new(Mutex::new(Agent::new(
        providers,
        db,
        &config,
        tools,
        profiles,
        cwd,
        agents_context,
    )?));

    let context_window = {
        let agent_lock = agent.lock().await;
        let cw = agent_lock.fetch_context_window().await;
        if cw == 0 {
            tracing::warn!("Failed to fetch context window from API");
        }
        cw
    };

    if let Some(ref id) = resume_id {
        let mut agent_lock = agent.lock().await;
        match agent_lock.get_session(id) {
            Ok(conv) => {
                let _ = agent_lock.resume_conversation(&conv);
            }
            Err(e) => {
                tracing::warn!("Failed to resume session {}: {}", id, e);
            }
        }
    }

    let mut app = App::new(
        model_name,
        provider_name,
        agent_name,
        &config.theme.name,
        config.tui.vim_mode,
        context_window,
    );
    app.history = history;

    if let Some(ref id) = resume_id {
        let agent_lock = agent.lock().await;
        if let Ok(conv) = agent_lock.get_session(id) {
            app.conversation_title = conv.title.clone();
            for m in &conv.messages {
                let model = if m.role == "assistant" {
                    Some(conv.model.clone())
                } else {
                    None
                };
                app.messages.push(ChatMessage {
                    role: m.role.clone(),
                    content: m.content.clone(),
                    tool_calls: Vec::new(),
                    thinking: None,
                    model,
                });
            }
            app.scroll_to_bottom();
        }
        drop(agent_lock);
    }

    let mut events = EventHandler::new();
    let mut agent_rx: Option<mpsc::UnboundedReceiver<crate::agent::AgentEvent>> = None;

    loop {
        terminal.draw(|f| ui::draw(f, &mut app))?;

        let event = if let Some(ref mut rx) = agent_rx {
            tokio::select! {
                biased;
                agent_event = rx.recv() => {
                    match agent_event {
                        Some(ev) => {
                            app.handle_agent_event(ev);
                        }
                        None => {
                            if app.is_streaming {
                                app.is_streaming = false;
                            }
                            agent_rx = None;
                            if let Some(queued) = app.message_queue.pop_front() {
                                let (tx, rx) = mpsc::unbounded_channel();
                                agent_rx = Some(rx);
                                app.is_streaming = true;
                                app.streaming_started = Some(Instant::now());
                                app.current_response.clear();
                                app.current_thinking.clear();
                                app.current_tool_calls.clear();
                                app.error_message = None;
                                let agent_clone = Arc::clone(&agent);
                                tokio::spawn(async move {
                                    let mut agent = agent_clone.lock().await;
                                    let result = if queued.images.is_empty() {
                                        agent.send_message(&queued.text, tx).await
                                    } else {
                                        agent.send_message_with_images(&queued.text, queued.images, tx).await
                                    };
                                    if let Err(e) = result {
                                        tracing::error!("Agent send_message error: {}", e);
                                    }
                                });
                            }
                        }
                    }
                    continue;
                }
                ui_event = events.next() => {
                    match ui_event {
                        Some(ev) => ev,
                        None => break,
                    }
                }
            }
        } else {
            match events.next().await {
                Some(ev) => ev,
                None => break,
            }
        };

        if let LoopSignal::Quit = handle_event(&mut app, &agent, event, &mut agent_rx).await {
            break;
        }
    }

    let mut agent_lock = agent.lock().await;
    let conversation_id = agent_lock.conversation_id().to_string();
    let title = agent_lock.conversation_title();
    agent_lock.cleanup_if_empty();
    drop(agent_lock);

    Ok(ExitInfo {
        conversation_id,
        title,
    })
}

enum LoopSignal {
    Continue,
    Quit,
    CancelStream,
}

async fn dispatch_action(
    app: &mut App,
    agent: &Arc<Mutex<Agent>>,
    action: InputAction,
    agent_rx: &mut Option<mpsc::UnboundedReceiver<crate::agent::AgentEvent>>,
) -> LoopSignal {
    match action {
        InputAction::Quit => return LoopSignal::Quit,
        InputAction::CancelStream => {
            *agent_rx = None;
            app.is_streaming = false;
            app.streaming_started = None;
            if !app.current_response.is_empty() || !app.current_tool_calls.is_empty() {
                let thinking = if app.current_thinking.is_empty() {
                    None
                } else {
                    Some(std::mem::take(&mut app.current_thinking))
                };
                app.messages.push(app::ChatMessage {
                    role: "assistant".to_string(),
                    content: std::mem::take(&mut app.current_response),
                    tool_calls: std::mem::take(&mut app.current_tool_calls),
                    thinking,
                    model: Some(app.model_name.clone()),
                });
            } else {
                app.current_response.clear();
                app.current_thinking.clear();
                app.current_tool_calls.clear();
            }
            app.pending_tool_name = None;
            // Drop pending question/permission to unblock agent
            app.pending_question = None;
            app.pending_permission = None;
            app.error_message = Some("cancelled".to_string());
            return LoopSignal::CancelStream;
        }
        InputAction::SendMessage(msg) => {
            let images: Vec<(String, String)> = app
                .take_attachments()
                .into_iter()
                .map(|a| (a.media_type, a.data))
                .collect();

            let (tx, rx) = mpsc::unbounded_channel();
            *agent_rx = Some(rx);

            let agent_clone = Arc::clone(agent);
            tokio::spawn(async move {
                let mut agent = agent_clone.lock().await;
                let result = if images.is_empty() {
                    agent.send_message(&msg, tx).await
                } else {
                    agent.send_message_with_images(&msg, images, tx).await
                };
                if let Err(e) = result {
                    tracing::error!("Agent send_message error: {}", e);
                }
            });
        }
        InputAction::NewConversation => {
            let mut agent_lock = agent.lock().await;
            match agent_lock.new_conversation() {
                Ok(()) => app.clear_conversation(),
                Err(e) => {
                    app.error_message = Some(format!("failed to start new conversation: {e}"))
                }
            }
        }
        InputAction::OpenModelSelector => {
            let agent_lock = agent.lock().await;
            let grouped = agent_lock.fetch_all_models().await;
            let current_provider = agent_lock.current_provider_name().to_string();
            let current_model = agent_lock.current_model().to_string();
            drop(agent_lock);
            app.model_selector
                .open(grouped, &current_provider, &current_model);
        }
        InputAction::OpenAgentSelector => {
            let agent_lock = agent.lock().await;
            let entries: Vec<AgentEntry> = agent_lock
                .agent_profiles()
                .iter()
                .map(|p| AgentEntry {
                    name: p.name.clone(),
                    description: p.description.clone(),
                })
                .collect();
            let current = agent_lock.current_agent_name().to_string();
            drop(agent_lock);
            app.agent_selector.open(entries, &current);
        }
        InputAction::OpenSessionSelector => {
            let agent_lock = agent.lock().await;
            let current_id = agent_lock.conversation_id().to_string();
            let sessions = agent_lock.list_sessions().unwrap_or_default();
            drop(agent_lock);
            let entries: Vec<SessionEntry> = sessions
                .into_iter()
                .map(|s| {
                    let title = if let Some(t) = &s.title {
                        t.clone()
                    } else if s.id == current_id {
                        app.conversation_title
                            .clone()
                            .unwrap_or_else(|| "new conversation".to_string())
                    } else {
                        "untitled".to_string()
                    };
                    SessionEntry {
                        id: s.id.clone(),
                        title,
                        subtitle: format!("{} · {}", time_ago(&s.updated_at), s.provider),
                    }
                })
                .collect();
            app.session_selector.open(entries);
        }
        InputAction::ResumeSession { id } => {
            let mut agent_lock = agent.lock().await;
            match agent_lock.get_session(&id) {
                Ok(conv) => {
                    let title = conv.title.clone();
                    let conv_model = conv.model.clone();
                    let messages_for_ui: Vec<(String, String)> = conv
                        .messages
                        .iter()
                        .map(|m| (m.role.clone(), m.content.clone()))
                        .collect();
                    match agent_lock.resume_conversation(&conv) {
                        Ok(()) => {
                            drop(agent_lock);
                            app.clear_conversation();
                            app.conversation_title = title;
                            for (role, content) in messages_for_ui {
                                let model = if role == "assistant" {
                                    Some(conv_model.clone())
                                } else {
                                    None
                                };
                                app.messages.push(ChatMessage {
                                    role,
                                    content,
                                    tool_calls: Vec::new(),
                                    thinking: None,
                                    model,
                                });
                            }
                            app.scroll_to_bottom();
                        }
                        Err(e) => {
                            drop(agent_lock);
                            app.error_message = Some(format!("failed to resume session: {e}"));
                        }
                    }
                }
                Err(e) => {
                    drop(agent_lock);
                    app.error_message = Some(format!("session not found: {e}"));
                }
            }
        }
        InputAction::SelectModel { provider, model } => {
            let mut agent_lock = agent.lock().await;
            agent_lock.set_active_provider(&provider, &model);
            let cw = agent_lock.context_window();
            if cw > 0 {
                app.context_window = cw;
            } else {
                app.context_window = agent_lock.fetch_context_window().await;
            }
        }
        InputAction::SelectAgent { name } => {
            let mut agent_lock = agent.lock().await;
            agent_lock.switch_agent(&name);
            app.model_name = agent_lock.current_model().to_string();
            app.provider_name = agent_lock.current_provider_name().to_string();
            let cw = agent_lock.context_window();
            if cw > 0 {
                app.context_window = cw;
            } else {
                app.context_window = agent_lock.fetch_context_window().await;
            }
        }
        InputAction::ScrollUp(n) => app.scroll_up(n),
        InputAction::ScrollDown(n) => app.scroll_down(n),
        InputAction::ScrollToTop => app.scroll_to_top(),
        InputAction::ScrollToBottom => app.scroll_to_bottom(),
        InputAction::ClearConversation => app.clear_conversation(),
        InputAction::ToggleThinking => {
            app.thinking_expanded = !app.thinking_expanded;
        }
        InputAction::OpenThinkingSelector => {
            let level = app.thinking_level();
            app.thinking_selector.open(level);
        }
        InputAction::SetThinkingLevel(budget) => {
            let mut agent_lock = agent.lock().await;
            agent_lock.set_thinking_budget(budget);
        }
        InputAction::CycleThinkingLevel => {
            let next = app.thinking_level().next();
            let budget = next.budget_tokens();
            app.thinking_budget = budget;
            let mut agent_lock = agent.lock().await;
            agent_lock.set_thinking_budget(budget);
        }
        InputAction::TruncateToMessage(idx) => {
            app.messages.truncate(idx + 1);
            app.current_response.clear();
            app.current_thinking.clear();
            app.current_tool_calls.clear();
            app.scroll_to_bottom();
            let mut agent_lock = agent.lock().await;
            agent_lock.truncate_messages(idx + 1);
        }
        InputAction::ForkFromMessage(idx) => {
            let fork_messages: Vec<(String, String, Option<String>)> = app.messages[..=idx]
                .iter()
                .map(|m| (m.role.clone(), m.content.clone(), m.model.clone()))
                .collect();
            let mut agent_lock = agent.lock().await;
            match agent_lock.fork_conversation(idx + 1) {
                Ok(()) => {
                    drop(agent_lock);
                    app.clear_conversation();
                    for (role, content, model) in fork_messages {
                        app.messages.push(app::ChatMessage {
                            role,
                            content,
                            tool_calls: Vec::new(),
                            thinking: None,
                            model,
                        });
                    }
                    app.scroll_to_bottom();
                }
                Err(e) => {
                    drop(agent_lock);
                    app.error_message = Some(format!("fork failed: {e}"));
                }
            }
        }
        InputAction::AnswerQuestion(answer) => {
            app.messages.push(ChatMessage {
                role: "user".to_string(),
                content: answer,
                tool_calls: Vec::new(),
                thinking: None,
                model: None,
            });
            app.scroll_to_bottom();
        }
        InputAction::AnswerPermission(_) | InputAction::None => {}
    }
    LoopSignal::Continue
}

async fn handle_event(
    app: &mut App,
    agent: &Arc<Mutex<Agent>>,
    event: AppEvent,
    agent_rx: &mut Option<mpsc::UnboundedReceiver<crate::agent::AgentEvent>>,
) -> LoopSignal {
    let action = match event {
        AppEvent::Key(key) => input::handle_key(app, key),
        AppEvent::Mouse(mouse) => input::handle_mouse(app, mouse),
        AppEvent::Paste(text) => input::handle_paste(app, text),
        AppEvent::Tick => {
            app.tick_count = app.tick_count.wrapping_add(1);
            app.animate_scroll();
            return LoopSignal::Continue;
        }
        AppEvent::Agent(ev) => {
            app.handle_agent_event(ev);
            return LoopSignal::Continue;
        }
        AppEvent::Resize(_, _) => return LoopSignal::Continue,
    };
    dispatch_action(app, agent, action, agent_rx).await
}