loki/src/repl/mod.rs

mod completer;
mod highlighter;
mod prompt;

use self::completer::ReplCompleter;
use self::highlighter::ReplHighlighter;
use self::prompt::ReplPrompt;

use crate::client::{call_chat_completions, call_chat_completions_streaming};
use crate::config::{
    AgentVariables, AssertState, Config, GlobalConfig, Input, LastMessage, StateFlags,
    macro_execute,
};
use crate::render::render_error;
use crate::utils::{
    AbortSignal, abortable_run_with_spinner, create_abort_signal, dimmed_text, set_text, temp_file,
};

use crate::mcp::McpRegistry;
use anyhow::{Context, Result, bail};
use crossterm::cursor::SetCursorStyle;
use fancy_regex::Regex;
use reedline::CursorConfig;
use reedline::{
    ColumnarMenu, EditCommand, EditMode, Emacs, KeyCode, KeyModifiers, Keybindings, Reedline,
    ReedlineEvent, ReedlineMenu, ValidationResult, Validator, Vi, default_emacs_keybindings,
    default_vi_insert_keybindings, default_vi_normal_keybindings,
};
use reedline::{MenuBuilder, Signal};
use std::sync::LazyLock;
use std::{env, mem, process};

const MENU_NAME: &str = "completion_menu";

static REPL_COMMANDS: LazyLock<[ReplCommand; 37]> = LazyLock::new(|| {
    [
        ReplCommand::new(".help", "Show this help guide", AssertState::pass()),
        ReplCommand::new(".info", "Show system info", AssertState::pass()),
        ReplCommand::new(
            ".edit config",
            "Modify configuration file",
            AssertState::False(StateFlags::AGENT),
        ),
        ReplCommand::new(".model", "Switch LLM model", AssertState::pass()),
        ReplCommand::new(
            ".prompt",
            "Set a temporary role using a prompt",
            AssertState::False(StateFlags::SESSION | StateFlags::AGENT),
        ),
        ReplCommand::new(
            ".role",
            "Create or switch to a role",
            AssertState::False(StateFlags::SESSION | StateFlags::AGENT),
        ),
        ReplCommand::new(
            ".info role",
            "Show role info",
            AssertState::True(StateFlags::ROLE),
        ),
        ReplCommand::new(
            ".edit role",
            "Modify current role",
            AssertState::TrueFalse(StateFlags::ROLE, StateFlags::SESSION),
        ),
        ReplCommand::new(
            ".save role",
            "Save current role to file",
            AssertState::TrueFalse(
                StateFlags::ROLE,
                StateFlags::SESSION_EMPTY | StateFlags::SESSION,
            ),
        ),
        ReplCommand::new(
            ".exit role",
            "Exit active role",
            AssertState::TrueFalse(StateFlags::ROLE, StateFlags::SESSION),
        ),
        ReplCommand::new(
            ".session",
            "Start or switch to a session",
            AssertState::False(StateFlags::SESSION_EMPTY | StateFlags::SESSION),
        ),
        ReplCommand::new(
            ".empty session",
            "Clear session messages",
            AssertState::True(StateFlags::SESSION),
        ),
        ReplCommand::new(
            ".compress session",
            "Compress session messages",
            AssertState::True(StateFlags::SESSION),
        ),
        ReplCommand::new(
            ".info session",
            "Show session info",
            AssertState::True(StateFlags::SESSION_EMPTY | StateFlags::SESSION),
        ),
        ReplCommand::new(
            ".edit session",
            "Modify current session",
            AssertState::True(StateFlags::SESSION_EMPTY | StateFlags::SESSION),
        ),
        ReplCommand::new(
            ".save session",
            "Save current session to file",
            AssertState::True(StateFlags::SESSION_EMPTY | StateFlags::SESSION),
        ),
        ReplCommand::new(
            ".exit session",
            "Exit active session",
            AssertState::True(StateFlags::SESSION_EMPTY | StateFlags::SESSION),
        ),
        ReplCommand::new(".agent", "Use an agent", AssertState::bare()),
        ReplCommand::new(
            ".starter",
            "Use a conversation starter",
            AssertState::True(StateFlags::AGENT),
        ),
        ReplCommand::new(
            ".edit agent-config",
            "Modify agent configuration file",
            AssertState::True(StateFlags::AGENT),
        ),
        ReplCommand::new(
            ".info agent",
            "Show agent info",
            AssertState::True(StateFlags::AGENT),
        ),
        ReplCommand::new(
            ".exit agent",
            "Leave agent",
            AssertState::True(StateFlags::AGENT),
        ),
        ReplCommand::new(
            ".rag",
            "Initialize or access RAG",
            AssertState::False(StateFlags::AGENT),
        ),
        ReplCommand::new(
            ".edit rag-docs",
            "Add or remove documents from an existing RAG",
            AssertState::TrueFalse(StateFlags::RAG, StateFlags::AGENT),
        ),
        ReplCommand::new(
            ".rebuild rag",
            "Rebuild RAG for document changes",
            AssertState::True(StateFlags::RAG),
        ),
        ReplCommand::new(
            ".sources rag",
            "Show citation sources used in last query",
            AssertState::True(StateFlags::RAG),
        ),
        ReplCommand::new(
            ".info rag",
            "Show RAG info",
            AssertState::True(StateFlags::RAG),
        ),
        ReplCommand::new(
            ".exit rag",
            "Leave RAG",
            AssertState::TrueFalse(StateFlags::RAG, StateFlags::AGENT),
        ),
        ReplCommand::new(".macro", "Execute a macro", AssertState::pass()),
        ReplCommand::new(
            ".file",
            "Include files, directories, URLs or commands",
            AssertState::pass(),
        ),
        ReplCommand::new(
            ".continue",
            "Continue previous response",
            AssertState::pass(),
        ),
        ReplCommand::new(
            ".regenerate",
            "Regenerate last response",
            AssertState::pass(),
        ),
        ReplCommand::new(".copy", "Copy last response", AssertState::pass()),
        ReplCommand::new(".set", "Modify runtime settings", AssertState::pass()),
        ReplCommand::new(
            ".delete",
            "Delete roles, sessions, RAGs, or agents",
            AssertState::pass(),
        ),
        ReplCommand::new(
            ".vault",
            "View or modify the Loki vault",
            AssertState::pass(),
        ),
        ReplCommand::new(".exit", "Exit REPL", AssertState::pass()),
    ]
});
static COMMAND_RE: LazyLock<Regex> = LazyLock::new(|| Regex::new(r"^\s*(\.\S*)\s*").unwrap());
static MULTILINE_RE: LazyLock<Regex> =
    LazyLock::new(|| Regex::new(r"(?s)^\s*:::\s*(.*)\s*:::\s*$").unwrap());

pub struct Repl {
    config: GlobalConfig,
    editor: Reedline,
    prompt: ReplPrompt,
    abort_signal: AbortSignal,
}

impl Repl {
    pub fn init(config: &GlobalConfig) -> Result<Self> {
        let editor = Self::create_editor(config)?;

        let prompt = ReplPrompt::new(config);
        let abort_signal = create_abort_signal();

        Ok(Self {
            config: config.clone(),
            editor,
            prompt,
            abort_signal,
        })
    }

    pub async fn run(&mut self) -> Result<()> {
        if AssertState::False(StateFlags::AGENT | StateFlags::RAG)
            .assert(self.config.read().state())
        {
            print!(
                r#"Welcome to {} {}
Type ".help" for additional help.
"#,
                env!("CARGO_CRATE_NAME"),
                env!("CARGO_PKG_VERSION"),
            )
        }

        loop {
            if self.abort_signal.aborted_ctrld() {
                break;
            }
            let sig = self.editor.read_line(&self.prompt);
            match sig {
                Ok(Signal::Success(line)) => {
                    self.abort_signal.reset();
                    match run_repl_command(&self.config, self.abort_signal.clone(), &line).await {
                        Ok(exit) => {
                            if exit {
                                break;
                            }
                        }
                        Err(err) => {
                            render_error(err);
                            println!()
                        }
                    }
                }
                Ok(Signal::CtrlC) => {
                    self.abort_signal.set_ctrlc();
                    println!("(To exit, press Ctrl+D or enter \".exit\")\n");
                }
                Ok(Signal::CtrlD) => {
                    self.abort_signal.set_ctrld();
                    break;
                }
                _ => {}
            }
        }
        self.config.write().exit_session()?;
        Ok(())
    }

    fn create_editor(config: &GlobalConfig) -> Result<Reedline> {
        let completer = ReplCompleter::new(config);
        let highlighter = ReplHighlighter::new(config);
        let menu = Self::create_menu();
        let edit_mode = Self::create_edit_mode(config);
        let cursor_config = CursorConfig {
            vi_insert: Some(SetCursorStyle::BlinkingBar),
            vi_normal: Some(SetCursorStyle::SteadyBlock),
            emacs: None,
        };
        let mut editor = Reedline::create()
            .with_completer(Box::new(completer))
            .with_highlighter(Box::new(highlighter))
            .with_menu(menu)
            .with_edit_mode(edit_mode)
            .with_cursor_config(cursor_config)
            .with_quick_completions(true)
            .with_partial_completions(true)
            .use_bracketed_paste(true)
            .with_validator(Box::new(ReplValidator))
            .with_ansi_colors(true);

        if let Ok(cmd) = config.read().editor() {
            let temp_file = temp_file("-repl-", ".md");
            let command = process::Command::new(cmd);
            editor = editor.with_buffer_editor(command, temp_file);
        }

        Ok(editor)
    }

    fn extra_keybindings(keybindings: &mut Keybindings) {
        keybindings.add_binding(
            KeyModifiers::NONE,
            KeyCode::Tab,
            ReedlineEvent::UntilFound(vec![
                ReedlineEvent::Menu(MENU_NAME.to_string()),
                ReedlineEvent::MenuNext,
            ]),
        );
        keybindings.add_binding(
            KeyModifiers::SHIFT,
            KeyCode::BackTab,
            ReedlineEvent::MenuPrevious,
        );
        keybindings.add_binding(
            KeyModifiers::CONTROL,
            KeyCode::Enter,
            ReedlineEvent::Edit(vec![EditCommand::InsertNewline]),
        );
        keybindings.add_binding(
            KeyModifiers::CONTROL,
            KeyCode::Char('j'),
            ReedlineEvent::Edit(vec![EditCommand::InsertNewline]),
        );
    }

    fn create_edit_mode(config: &GlobalConfig) -> Box<dyn EditMode> {
        let edit_mode: Box<dyn EditMode> = if config.read().keybindings == "vi" {
            let mut insert_keybindings = default_vi_insert_keybindings();
            Self::extra_keybindings(&mut insert_keybindings);
            Box::new(Vi::new(insert_keybindings, default_vi_normal_keybindings()))
        } else {
            let mut keybindings = default_emacs_keybindings();
            Self::extra_keybindings(&mut keybindings);
            Box::new(Emacs::new(keybindings))
        };
        edit_mode
    }

    fn create_menu() -> ReedlineMenu {
        let completion_menu = ColumnarMenu::default().with_name(MENU_NAME);
        ReedlineMenu::EngineCompleter(Box::new(completion_menu))
    }
}

#[derive(Debug, Clone)]
pub struct ReplCommand {
    name: &'static str,
    description: &'static str,
    state: AssertState,
}

impl ReplCommand {
    fn new(name: &'static str, desc: &'static str, state: AssertState) -> Self {
        Self {
            name,
            description: desc,
            state,
        }
    }

    fn is_valid(&self, flags: StateFlags) -> bool {
        self.state.assert(flags)
    }
}

/// A default validator which checks for mismatched quotes and brackets
struct ReplValidator;

impl Validator for ReplValidator {
    fn validate(&self, line: &str) -> ValidationResult {
        let line = line.trim();
        if line.starts_with(r#":::"#) && !line[3..].ends_with(r#":::"#) {
            ValidationResult::Incomplete
        } else {
            ValidationResult::Complete
        }
    }
}

pub async fn run_repl_command(
    config: &GlobalConfig,
    abort_signal: AbortSignal,
    mut line: &str,
) -> Result<bool> {
    if let Ok(Some(captures)) = MULTILINE_RE.captures(line)
        && let Some(text_match) = captures.get(1)
    {
        line = text_match.as_str();
    }
    match parse_command(line) {
        Some((cmd, args)) => match cmd {
            ".help" => {
                dump_repl_help();
            }
            ".info" => match args {
                Some("role") => {
                    let info = config.read().role_info()?;
                    print!("{info}");
                }
                Some("session") => {
                    let info = config.read().session_info()?;
                    print!("{info}");
                }
                Some("rag") => {
                    let info = config.read().rag_info()?;
                    print!("{info}");
                }
                Some("agent") => {
                    let info = config.read().agent_info()?;
                    print!("{info}");
                }
                Some(_) => unknown_command()?,
                None => {
                    let output = config.read().sysinfo()?;
                    print!("{output}");
                }
            },
            ".model" => match args {
                Some(name) => {
                    config.write().set_model(name)?;
                }
                None => println!("Usage: .model <name>"),
            },
            ".prompt" => match args {
                Some(text) => {
                    config.write().use_prompt(text)?;
                }
                None => println!("Usage: .prompt <text>..."),
            },
            ".role" => match args {
                Some(args) => match args.split_once(['\n', ' ']) {
                    Some((name, text)) => {
                        let role = config.read().retrieve_role(name.trim())?;
                        let input = Input::from_str(config, text, Some(role));
                        ask(config, abort_signal.clone(), input, false).await?;
                    }
                    None => {
                        let name = args;
                        if !Config::has_role(name) {
                            config.write().new_role(name)?;
                        }

                        Config::use_role_safely(config, name, abort_signal.clone()).await?;
                    }
                },
                None => println!(
                    r#"Usage:
    .role <name>                    # If the role exists, switch to it; otherwise, create a new role
    .role <name> [text]...          # Temporarily switch to the role, send the text, and switch back"#
                ),
            },
            ".session" => {
                Config::use_session_safely(config, args, abort_signal.clone()).await?;
                Config::maybe_autoname_session(config.clone());
            }
            ".rag" => {
                Config::use_rag(config, args, abort_signal.clone()).await?;
            }
            ".agent" => match split_first_arg(args) {
                Some((agent_name, args)) => {
                    let (new_args, _) = split_args_text(args.unwrap_or_default(), cfg!(windows));
                    let (session_name, variable_pairs) = match new_args.first() {
                        Some(name) if name.contains('=') => (None, new_args.as_slice()),
                        Some(name) => (Some(name.as_str()), &new_args[1..]),
                        None => (None, &[] as &[String]),
                    };
                    let variables: AgentVariables = variable_pairs
                        .iter()
                        .filter_map(|v| v.split_once('='))
                        .map(|(key, value)| (key.to_string(), value.to_string()))
                        .collect();
                    if variables.len() != variable_pairs.len() {
                        bail!("Some variable values are not key=value pairs");
                    }
                    if !variables.is_empty() {
                        config.write().agent_variables = Some(variables);
                    }
                    let ret =
                        Config::use_agent(config, agent_name, session_name, abort_signal.clone())
                            .await;
                    config.write().agent_variables = None;
                    ret?;
                }
                None => {
                    println!(r#"Usage: .agent <agent-name> [session-name] [key=value]..."#)
                }
            },
            ".starter" => match args {
                Some(id) => {
                    let mut text = None;
                    if let Some(agent) = config.read().agent.as_ref() {
                        for (i, value) in agent.conversation_starters().iter().enumerate() {
                            if (i + 1).to_string() == id {
                                text = Some(value.clone());
                            }
                        }
                    }
                    match text {
                        Some(text) => {
                            println!("{}", dimmed_text(&format!(">> {text}")));
                            let input = Input::from_str(config, &text, None);
                            ask(config, abort_signal.clone(), input, true).await?;
                        }
                        None => {
                            bail!("Invalid starter value");
                        }
                    }
                }
                None => {
                    let banner = config.read().agent_banner()?;
                    config.read().print_markdown(&banner)?;
                }
            },
            ".save" => match split_first_arg(args) {
                Some(("role", name)) => {
                    config.write().save_role(name)?;
                }
                Some(("session", name)) => {
                    config.write().save_session(name)?;
                }
                _ => {
                    println!(r#"Usage: .save <role|session> [name]"#)
                }
            },
            ".edit" => {
                if config.read().macro_flag {
                    bail!("Cannot perform this operation because you are in a macro")
                }
                match args {
                    Some("config") => {
                        config.read().edit_config()?;
                    }
                    Some("role") => {
                        let mut cfg = {
                            let mut guard = config.write();
                            mem::take(&mut *guard)
                        };

                        cfg.edit_role(abort_signal.clone()).await?;

                        {
                            let mut guard = config.write();
                            *guard = cfg;
                        }
                    }
                    Some("session") => {
                        config.write().edit_session()?;
                    }
                    Some("rag-docs") => {
                        Config::edit_rag_docs(config, abort_signal.clone()).await?;
                    }
                    Some("agent-config") => {
                        config.write().edit_agent_config()?;
                    }
                    _ => {
                        println!(r#"Usage: .edit <config|role|session|rag-docs|agent-config>"#)
                    }
                }
            }
            ".compress" => match args {
                Some("session") => {
                    abortable_run_with_spinner(
                        Config::compress_session(config),
                        "Compressing",
                        abort_signal.clone(),
                    )
                    .await?;
                    println!("✓ Successfully compressed the session.");
                }
                _ => {
                    println!(r#"Usage: .compress session"#)
                }
            },
            ".empty" => match args {
                Some("session") => {
                    config.write().empty_session()?;
                }
                _ => {
                    println!(r#"Usage: .empty session"#)
                }
            },
            ".rebuild" => match args {
                Some("rag") => {
                    Config::rebuild_rag(config, abort_signal.clone()).await?;
                }
                _ => {
                    println!(r#"Usage: .rebuild rag"#)
                }
            },
            ".sources" => match args {
                Some("rag") => {
                    let output = Config::rag_sources(config)?;
                    println!("{output}");
                }
                _ => {
                    println!(r#"Usage: .sources rag"#)
                }
            },
            ".macro" => match split_first_arg(args) {
                Some((name, extra)) => {
                    if !Config::has_macro(name) && extra.is_none() {
                        config.write().new_macro(name)?;
                    } else {
                        macro_execute(config, name, extra, abort_signal.clone()).await?;
                    }
                }
                None => println!("Usage: .macro <name> <text>..."),
            },
            ".file" => match args {
                Some(args) => {
                    let (files, text) = split_args_text(args, cfg!(windows));
                    let input = Input::from_files_with_spinner(
                        config,
                        text,
                        files,
                        None,
                        abort_signal.clone(),
                    )
                    .await?;
                    ask(config, abort_signal.clone(), input, true).await?;
                }
                None => println!(
                    r#"Usage: .file <file|dir|url|cmd|loader:resource|%%>... [-- <text>...]

.file /tmp/file.txt
.file src/ Cargo.toml -- analyze
.file https://example.com/file.txt -- summarize
.file https://example.com/image.png -- recognize text
.file `git diff` -- Generate git commit message
.file jina:https://example.com
.file %% -- translate last reply to english"#
                ),
            },
            ".continue" => {
                let LastMessage {
                    mut input, output, ..
                } = match config
                    .read()
                    .last_message
                    .as_ref()
                    .filter(|v| v.continuous && !v.output.is_empty())
                    .cloned()
                {
                    Some(v) => v,
                    None => bail!("Unable to continue the response"),
                };
                input.set_continue_output(&output);
                ask(config, abort_signal.clone(), input, true).await?;
            }
            ".regenerate" => {
                let LastMessage { mut input, .. } = match config
                    .read()
                    .last_message
                    .as_ref()
                    .filter(|v| v.continuous)
                    .cloned()
                {
                    Some(v) => v,
                    None => bail!("Unable to regenerate the response"),
                };
                input.set_regenerate();
                ask(config, abort_signal.clone(), input, true).await?;
            }
            ".set" => match args {
                Some(args) => {
                    Config::update(config, args, abort_signal).await?;
                }
                _ => {
                    println!("Usage: .set <key> <value>...")
                }
            },
            ".delete" => match args {
                Some(args) => {
                    Config::delete(config, args)?;
                }
                _ => {
                    println!("Usage: .delete <role|session|rag|macro|agent-data>")
                }
            },
            ".copy" => {
                let output = match config
                    .read()
                    .last_message
                    .as_ref()
                    .filter(|v| !v.output.is_empty())
                    .map(|v| v.output.clone())
                {
                    Some(v) => v,
                    None => bail!("No chat response to copy"),
                };
                set_text(&output).context("Failed to copy the last chat response")?;
            }
            ".exit" => match args {
                Some("role") => {
                    config.write().exit_role()?;
                    config.write().functions.clear_mcp_meta_functions();

                    let registry = config
                        .write()
                        .mcp_registry
                        .take()
                        .expect("MCP registry should exist");
                    let enabled_mcp_servers = if config.read().mcp_server_support {
                        config.read().enabled_mcp_servers.clone()
                    } else {
                        None
                    };
                    let registry =
                        McpRegistry::reinit(registry, enabled_mcp_servers, abort_signal.clone())
                            .await?;
                    if !registry.is_empty() {
                        config
                            .write()
                            .functions
                            .append_mcp_meta_functions(registry.list_started_servers());
                    }
                    config.write().mcp_registry = Some(registry);
                }
                Some("session") => {
                    if config.read().agent.is_some() {
                        config.write().exit_agent_session()?;
                        config.write().functions.clear_mcp_meta_functions();

                        let registry = config
                            .write()
                            .mcp_registry
                            .take()
                            .expect("MCP registry should exist");
                        let enabled_mcp_servers = if config.read().mcp_server_support {
                            config.read().enabled_mcp_servers.clone()
                        } else {
                            None
                        };
                        let registry = McpRegistry::reinit(
                            registry,
                            enabled_mcp_servers,
                            abort_signal.clone(),
                        )
                        .await?;
                        if !registry.is_empty() {
                            config
                                .write()
                                .functions
                                .append_mcp_meta_functions(registry.list_started_servers());
                        }
                        config.write().mcp_registry = Some(registry);
                    } else {
                        config.write().exit_session()?;
                    }
                }
                Some("rag") => {
                    config.write().exit_rag()?;
                }
                Some("agent") => {
                    config.write().exit_agent()?;
                    config.write().functions.clear_mcp_meta_functions();

                    let registry = config
                        .write()
                        .mcp_registry
                        .take()
                        .expect("MCP registry should exist");
                    let enabled_mcp_servers = if config.read().mcp_server_support {
                        config.read().enabled_mcp_servers.clone()
                    } else {
                        None
                    };
                    let registry =
                        McpRegistry::reinit(registry, enabled_mcp_servers, abort_signal.clone())
                            .await?;
                    if !registry.is_empty() {
                        config
                            .write()
                            .functions
                            .append_mcp_meta_functions(registry.list_started_servers());
                    }
                    config.write().mcp_registry = Some(registry);
                }
                Some(_) => unknown_command()?,
                None => {
                    return Ok(true);
                }
            },
            ".clear" => match args {
                Some("messages") => {
                    bail!("Use '.empty session' instead");
                }
                _ => unknown_command()?,
            },
            ".vault" => match split_first_arg(args) {
                Some(("add", name)) => {
                    if let Some(name) = name {
                        config.read().vault.add_secret(name)?;
                    } else {
                        println!("Usage: .vault add <name>");
                    }
                }
                Some(("get", name)) => {
                    if let Some(name) = name {
                        config.read().vault.get_secret(name, true)?;
                    } else {
                        println!("Usage: .vault get <name>");
                    }
                }
                Some(("update", name)) => {
                    if let Some(name) = name {
                        config.read().vault.update_secret(name)?;
                    } else {
                        println!("Usage: .vault update <name>");
                    }
                }
                Some(("delete", name)) => {
                    if let Some(name) = name {
                        config.read().vault.delete_secret(name)?;
                    } else {
                        println!("Usage: .vault delete <name>");
                    }
                }
                Some(("list", _)) => {
                    config.read().vault.list_secrets(true)?;
                }
                None | Some(_) => {
                    println!("Usage: .vault <add|get|update|delete|list> [name]")
                }
            },
            _ => unknown_command()?,
        },
        None => {
            if config
                .read()
                .agent
                .as_ref()
                .is_some_and(|a| a.continuation_count() > 0)
            {
                config.write().agent.as_mut().unwrap().reset_continuation();
            }
            let input = Input::from_str(config, line, None);
            ask(config, abort_signal.clone(), input, true).await?;
        }
    }

    if !config.read().macro_flag {
        println!();
    }

    Ok(false)
}

#[async_recursion::async_recursion]
async fn ask(
    config: &GlobalConfig,
    abort_signal: AbortSignal,
    mut input: Input,
    with_embeddings: bool,
) -> Result<()> {
    if input.is_empty() {
        return Ok(());
    }
    if with_embeddings {
        input.use_embeddings(abort_signal.clone()).await?;
    }
    while config.read().is_compressing_session() {
        tokio::time::sleep(std::time::Duration::from_millis(100)).await;
    }

    let client = input.create_client()?;
    config.write().before_chat_completion(&input)?;
    let (output, tool_results) = if input.stream() {
        call_chat_completions_streaming(&input, client.as_ref(), abort_signal.clone()).await?
    } else {
        call_chat_completions(&input, true, false, client.as_ref(), abort_signal.clone()).await?
    };
    config
        .write()
        .after_chat_completion(&input, &output, &tool_results)?;
    if !tool_results.is_empty() {
        ask(
            config,
            abort_signal,
            input.merge_tool_results(output, tool_results),
            false,
        )
        .await
    } else {
        let should_continue = {
            let cfg = config.read();
            if let Some(agent) = &cfg.agent {
                agent.auto_continue_enabled()
                    && agent.continuation_count() < agent.max_auto_continues()
                    && agent.todo_list().has_incomplete()
            } else {
                false
            }
        };

        if should_continue {
            let full_prompt = {
                let mut cfg = config.write();
                let agent = cfg.agent.as_mut().expect("agent checked above");
                agent.set_last_continuation_response(output.clone());
                agent.increment_continuation();
                let count = agent.continuation_count();
                let max = agent.max_auto_continues();

                let todo_state = agent.todo_list().render_for_model();
                let remaining = agent.todo_list().incomplete_count();
                let prompt = agent.continuation_prompt();

                let color = if cfg.light_theme() {
                    nu_ansi_term::Color::LightGray
                } else {
                    nu_ansi_term::Color::DarkGray
                };
                eprintln!(
                    "\n📋 {}",
                    color.italic().paint(format!(
                        "Auto-continuing ({count}/{max}): {remaining} incomplete todo(s) remain"
                    ))
                );

                format!("{prompt}\n\n{todo_state}")
            };
            let continuation_input = Input::from_str(config, &full_prompt, None);
            ask(config, abort_signal, continuation_input, false).await
        } else {
            if config
                .read()
                .agent
                .as_ref()
                .is_some_and(|a| a.continuation_count() > 0)
            {
                config.write().agent.as_mut().unwrap().reset_continuation();
            }
            Config::maybe_autoname_session(config.clone());

            let needs_compression = {
                let cfg = config.read();
                let compression_threshold = cfg.compression_threshold;
                cfg.session
                    .as_ref()
                    .is_some_and(|s| s.needs_compression(compression_threshold))
            };

            if needs_compression {
                let agent_can_continue_after_compress = {
                    let cfg = config.read();
                    cfg.agent.as_ref().is_some_and(|agent| {
                        agent.auto_continue_enabled()
                            && agent.continuation_count() < agent.max_auto_continues()
                            && agent.todo_list().has_incomplete()
                    })
                };

                {
                    let mut cfg = config.write();
                    if let Some(session) = cfg.session.as_mut() {
                        session.set_compressing(true);
                    }
                }

                let color = if config.read().light_theme() {
                    nu_ansi_term::Color::LightGray
                } else {
                    nu_ansi_term::Color::DarkGray
                };
                eprintln!("\n📢 {}", color.italic().paint("Compressing the session."),);

                if let Err(err) = Config::compress_session(config).await {
                    log::warn!("Failed to compress the session: {err}");
                }
                if let Some(session) = config.write().session.as_mut() {
                    session.set_compressing(false);
                }

                if agent_can_continue_after_compress {
                    let full_prompt = {
                        let mut cfg = config.write();
                        let agent = cfg.agent.as_mut().expect("agent checked above");
                        agent.increment_continuation();
                        let count = agent.continuation_count();
                        let max = agent.max_auto_continues();

                        let todo_state = agent.todo_list().render_for_model();
                        let remaining = agent.todo_list().incomplete_count();
                        let prompt = agent.continuation_prompt();

                        let color = if cfg.light_theme() {
                            nu_ansi_term::Color::LightGray
                        } else {
                            nu_ansi_term::Color::DarkGray
                        };
                        eprintln!(
                            "\n📋 {}",
                            color.italic().paint(format!(
                                "Auto-continuing after compression ({count}/{max}): {remaining} incomplete todo(s) remain"
                            ))
                        );

                        format!("{prompt}\n\n{todo_state}")
                    };
                    let continuation_input = Input::from_str(config, &full_prompt, None);
                    return ask(config, abort_signal, continuation_input, false).await;
                }
            } else {
                Config::maybe_compress_session(config.clone());
            }

            Ok(())
        }
    }
}

fn unknown_command() -> Result<()> {
    bail!(r#"Unknown command. Type ".help" for additional help."#);
}

fn dump_repl_help() {
    let head = REPL_COMMANDS
        .iter()
        .map(|cmd| format!("{:<24} {}", cmd.name, cmd.description))
        .collect::<Vec<String>>()
        .join("\n");
    println!(
        r###"{head}

Type ::: to start multi-line editing, type ::: to finish it.
Press Ctrl+O to open an editor for editing the input buffer.
Press Ctrl+C to cancel the response, Ctrl+D to exit the REPL."###,
    );
}

fn parse_command(line: &str) -> Option<(&str, Option<&str>)> {
    match COMMAND_RE.captures(line) {
        Ok(Some(captures)) => {
            let cmd = captures.get(1)?.as_str();
            let args = line[captures[0].len()..].trim();
            let args = if args.is_empty() { None } else { Some(args) };
            Some((cmd, args))
        }
        _ => None,
    }
}

fn split_first_arg(args: Option<&str>) -> Option<(&str, Option<&str>)> {
    args.map(|v| match v.split_once(' ') {
        Some((subcmd, args)) => (subcmd, Some(args.trim())),
        None => (v, None),
    })
}

pub fn split_args_text(line: &str, is_win: bool) -> (Vec<String>, &str) {
    let mut words = Vec::new();
    let mut word = String::new();
    let mut unbalance: Option<char> = None;
    let mut prev_char: Option<char> = None;
    let mut text_starts_at = None;
    let unquote_word = |word: &str| {
        if ((word.starts_with('"') && word.ends_with('"'))
            || (word.starts_with('\'') && word.ends_with('\'')))
            && word.len() >= 2
        {
            word[1..word.len() - 1].to_string()
        } else {
            word.to_string()
        }
    };
    let chars: Vec<char> = line.chars().collect();

    for (i, char) in chars.iter().cloned().enumerate() {
        match unbalance {
            Some(ub_char) if ub_char == char => {
                word.push(char);
                unbalance = None;
            }
            Some(_) => {
                word.push(char);
            }
            None => match char {
                ' ' | '\t' | '\r' | '\n' => {
                    if char == '\r' && chars.get(i + 1) == Some(&'\n') {
                        continue;
                    }
                    if let Some('\\') = prev_char.filter(|_| !is_win) {
                        word.push(char);
                    } else if !word.is_empty() {
                        if word == "--" {
                            word.clear();
                            text_starts_at = Some(i + 1);
                            break;
                        }
                        words.push(unquote_word(&word));
                        word.clear();
                    }
                }
                '\'' | '"' | '`' => {
                    word.push(char);
                    unbalance = Some(char);
                }
                '\\' => {
                    if is_win || prev_char.map(|c| c == '\\').unwrap_or_default() {
                        word.push(char);
                    }
                }
                _ => {
                    word.push(char);
                }
            },
        }
        prev_char = Some(char);
    }

    if !word.is_empty() && word != "--" {
        words.push(unquote_word(&word));
    }
    let text = match text_starts_at {
        Some(start) => &line[start..],
        None => "",
    };

    (words, text)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_process_command_line() {
        assert_eq!(parse_command(" ."), Some((".", None)));
        assert_eq!(parse_command(" .role"), Some((".role", None)));
        assert_eq!(parse_command(" .role  "), Some((".role", None)));
        assert_eq!(
            parse_command(" .set dry_run true"),
            Some((".set", Some("dry_run true")))
        );
        assert_eq!(
            parse_command(" .set dry_run true  "),
            Some((".set", Some("dry_run true")))
        );
        assert_eq!(
            parse_command(".prompt \nabc\n"),
            Some((".prompt", Some("abc")))
        );
    }

    #[test]
    fn test_split_args_text() {
        assert_eq!(split_args_text("", false), (vec![], ""));
        assert_eq!(
            split_args_text("file.txt", false),
            (vec!["file.txt".into()], "")
        );
        assert_eq!(
            split_args_text("file.txt --", false),
            (vec!["file.txt".into()], "")
        );
        assert_eq!(
            split_args_text("file.txt -- hello", false),
            (vec!["file.txt".into()], "hello")
        );
        assert_eq!(
            split_args_text("file.txt -- \thello", false),
            (vec!["file.txt".into()], "\thello")
        );
        assert_eq!(
            split_args_text("file.txt --\nhello", false),
            (vec!["file.txt".into()], "hello")
        );
        assert_eq!(
            split_args_text("file.txt --\r\nhello", false),
            (vec!["file.txt".into()], "hello")
        );
        assert_eq!(
            split_args_text("file.txt --\rhello", false),
            (vec!["file.txt".into()], "hello")
        );
        assert_eq!(
            split_args_text(r#"file1.txt 'file2.txt' "file3.txt""#, false),
            (
                vec!["file1.txt".into(), "file2.txt".into(), "file3.txt".into()],
                ""
            )
        );
        assert_eq!(
            split_args_text(r#"./file1.txt 'file1 - Copy.txt' file\ 2.txt"#, false),
            (
                vec![
                    "./file1.txt".into(),
                    "file1 - Copy.txt".into(),
                    "file 2.txt".into()
                ],
                ""
            )
        );
        assert_eq!(
            split_args_text(r#".\file.txt C:\dir\file.txt"#, true),
            (vec![".\\file.txt".into(), "C:\\dir\\file.txt".into()], "")
        );
    }
}