toast_api/

unified_cli.rs

use crate::utils::extract_org_id_from_cookie;
use anyhow::{anyhow, Context, Result};
use ctrlc;
use std::fs;
use std::future::Future;
use std::io::{self, Write};
use std::path::Path;
use std::pin::Pin;
use std::process::{self, Command, Stdio};
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;

use crate::api::{Attachment, Claude, Session as ClaudeSession};
use crate::config::{HAIKU_MODEL, MAX_INTERNAL_ITERS, OPUS_MODEL, SONNET_MODEL, SYSTEM_PROMPT};
use crate::deepseek::{DeepSeek, Session as DeepSeekSession};
use crate::utils::{extract_commands, prettify};
use log::debug;

/// Unified CLI arguments for both Claude and DeepSeek
#[derive(Debug)]
pub struct UnifiedArgs {
    pub use_deepseek: bool,
    pub use_opus: bool,
    pub use_haiku: bool,
}

/// Execute a shell command and capture its output
fn execute_command(command: &str) -> Result<String> {
    let result = Command::new("sh")
        .arg("-c")
        .arg(command)
        .stdout(Stdio::piped())
        .stderr(Stdio::piped())
        .spawn()?;

    let output = result.wait_with_output()?;
    let mut msg = String::new();

    if !output.stdout.is_empty() {
        msg.push_str("=== STDOUT ===\n");
        msg.push_str(&String::from_utf8_lossy(&output.stdout));
        msg.push('\n');
    }

    if !output.stderr.is_empty() {
        msg.push_str("=== STDERR ===\n");
        msg.push_str(&String::from_utf8_lossy(&output.stderr));
        msg.push('\n');
    }

    msg.push_str(&format!(
        "Exit code: {}",
        output.status.code().unwrap_or(-1)
    ));

    Ok(msg)
}

/// Read files into attachments for Claude
fn collect_claude_attachments(paths: &[&str]) -> Result<Vec<Attachment>> {
    const LIMIT: usize = 5;
    const SIZE_LIMIT: u64 = 10 * 1024 * 1024;

    if paths.len() > LIMIT {
        return Err(anyhow!("cannot attach more than {LIMIT} files"));
    }

    let mut atts = Vec::new();

    for p in paths {
        if let Ok(meta) = fs::metadata(p) {
            if meta.len() > SIZE_LIMIT {
                eprintln!("Warning: file {p} is larger than 10 MB, skipping");
                continue;
            }

            if let Ok(content) = fs::read_to_string(p) {
                atts.push(Attachment {
                    file_name: Path::new(p)
                        .file_name()
                        .unwrap_or_default()
                        .to_string_lossy()
                        .into(),
                    size: meta.len(),
                    content,
                });
            } else {
                eprintln!("Warning: couldn't read file {p}");
            }
        } else {
            eprintln!("Warning: couldn't access file {p}");
        }
    }

    Ok(atts)
}

/// Run the unified CLI application with provider selection
pub async fn run(args: UnifiedArgs) -> Result<()> {
    // Set up Ctrl-C handler
    let running = Arc::new(AtomicBool::new(true));
    {
        let running = running.clone();
        ctrlc::set_handler(move || {
            running.store(false, Ordering::SeqCst);
            println!("\nGoodbye!");
            process::exit(0);
        })?;
    }

    if args.use_deepseek {
        run_deepseek(args, running).await
    } else {
        run_claude(args, running).await
    }
}

/// Run the CLI with DeepSeek provider
async fn run_deepseek(args: UnifiedArgs, running: Arc<AtomicBool>) -> Result<()> {
    // Load session values from config files
    let config_dir = dirs::config_dir()
        .ok_or_else(|| anyhow!("Could not determine config directory"))?
        .join("toast")
        .join("deepseek");

    // Create config directory if it doesn't exist
    if !config_dir.exists() {
        fs::create_dir_all(&config_dir)?;
    }

    let auth_token_path = config_dir.join("auth_token");
    let cookies_path = config_dir.join("cookies.json");

    // Check auth token
    let auth_token = if auth_token_path.exists() {
        fs::read_to_string(&auth_token_path)
            .context(format!(
                "Failed to read auth token from {auth_token_path:?}"
            ))?
            .trim()
            .to_string()
    } else {
        return Err(anyhow!(
            "Auth token file not found at {:?}\n\nTo get your DeepSeek auth token:\n1. Go to chat.deepseek.com in your browser\n2. Open Developer Tools (F12)\n3. Go to Network tab\n4. Look for Authorization header in any request\n5. Save the token part (without 'Bearer ') to this file",
            auth_token_path
        ));
    };

    // Check cookies
    let cookies = if cookies_path.exists() {
        serde_json::from_str(
            &fs::read_to_string(&cookies_path)
                .context(format!("Failed to read cookies from {cookies_path:?}"))?,
        )?
    } else {
        return Err(anyhow!(
            "Cookies file not found at {:?}\n\nDeepSeek requires Cloudflare cookies.\nUse the deepseek4free library to generate them.",
            cookies_path
        ));
    };

    let session = DeepSeekSession {
        auth_token,
        cookies,
    };

    // Determine model based on flags - R1 is the reasoning model with thinking enabled
    let model = if args.use_opus {
        "deepseek-r1" // Use R1 reasoning model for opus
    } else if args.use_haiku {
        "deepseek-lite"
    } else {
        "deepseek-r1" // Default to R1 reasoning model
    };

    let mut deepseek = DeepSeek::new(session)?;

    let stdin = io::stdin();
    let mut stdout = io::stdout();

    // Track if system prompt has been sent
    let mut system_prompt_sent = false;

    // Create a new chat session
    println!("Starting new DeepSeek chat session...");
    let chat_id = match deepseek.create_chat_session().await {
        Ok(id) => {
            println!("Session started with DeepSeek!\n");
            id
        }
        Err(e) => {
            return Err(anyhow!("Failed to create DeepSeek chat session: {}", e));
        }
    };

    // Enable detailed thinking for reasoning model
    let thinking_mode = if model == "deepseek-r1" {
        crate::deepseek::ThinkingMode::Detailed
    } else {
        crate::deepseek::ThinkingMode::Simple
    };
    let search_mode = crate::deepseek::SearchMode::Disabled;

    // Main chat loop - simplified like working deepseek_cli.rs
    while running.load(Ordering::SeqCst) {
        print!("You: ");
        stdout.flush()?;

        let mut buf = String::new();
        match stdin.read_line(&mut buf) {
            Ok(0) => {
                // EOF reached, exit gracefully
                println!("\nGoodbye!");
                break;
            }
            Ok(_) => {
                let input = buf.trim_end();

                // Check for empty input or exit commands
                if input.is_empty() {
                    continue;
                }

                if input.eq_ignore_ascii_case("/exit")
                    || input.eq_ignore_ascii_case("exit")
                    || input == "x"
                {
                    break;
                }

                // Send message to API - let DeepSeek handle all commands in its response
                print!("DeepSeek: ");
                stdout.flush()?;

                debug!("Sending to DeepSeek API...");

                // Include system prompt only on first message
                let system_prompt = if !system_prompt_sent {
                    system_prompt_sent = true;
                    Some(SYSTEM_PROMPT)
                } else {
                    None
                };

                match deepseek
                    .chat_completion(
                        &chat_id,
                        input,
                        None,
                        thinking_mode,
                        search_mode,
                        system_prompt,
                    )
                    .await
                {
                    Ok(response) => {
                        debug!("Got response, length: {}", response.len());
                        println!("{}", prettify(&response));

                        // Process commands in the response
                        process_deepseek_commands(
                            &mut deepseek,
                            &chat_id,
                            &response,
                            thinking_mode,
                            search_mode,
                        )
                        .await?;
                    }
                    Err(e) => {
                        debug!("DeepSeek API error: {e}");
                        eprintln!("\nError: {e}");
                    }
                }
                println!();
            }
            Err(e) => {
                eprintln!("Failed to read input: {e}");
                break;
            }
        }
    }

    Ok(())
}

/// Process commands in DeepSeek's response
async fn process_deepseek_commands(
    deepseek: &mut DeepSeek,
    chat_id: &str,
    response: &str,
    thinking_mode: crate::deepseek::ThinkingMode,
    search_mode: crate::deepseek::SearchMode,
) -> Result<()> {
    process_deepseek_commands_internal(deepseek, chat_id, response, thinking_mode, search_mode, 0)
        .await
}

fn process_deepseek_commands_internal<'a>(
    deepseek: &'a mut DeepSeek,
    chat_id: &'a str,
    response: &'a str,
    thinking_mode: crate::deepseek::ThinkingMode,
    search_mode: crate::deepseek::SearchMode,
    depth: usize,
) -> Pin<Box<dyn Future<Output = Result<()>> + 'a>> {
    Box::pin(async move {
        // Limit recursion depth
        const MAX_DEPTH: usize = 20;
        if depth >= MAX_DEPTH {
            println!(
                "Maximum command processing depth reached ({MAX_DEPTH}). Returning to user."
            );
            return Ok(());
        }

        // Extract read_file and exec commands
        let (reads, execs) = extract_commands(response);

        if reads.is_empty() && execs.is_empty() {
            return Ok(());
        }

        // Short pause before processing commands
        tokio::time::sleep(std::time::Duration::from_millis(500)).await;

        // Process file reads
        if !reads.is_empty() {
            let mut file_contents = Vec::new();

            for path in &reads {
                match fs::read_to_string(path) {
                    Ok(content) => {
                        file_contents.push(format!("=== File: {path} ===\n{content}"));
                    }
                    Err(e) => {
                        file_contents.push(format!("Error reading file {path}: {e}"));
                    }
                }
            }

            let file_message = format!(
                "Here are the contents of the files you requested:\n\n{}",
                file_contents.join("\n\n")
            );

            // print!("Sending file contents... ");
            io::stdout().flush()?;

            match deepseek
                .chat_completion(
                    chat_id,
                    &file_message,
                    None,
                    thinking_mode,
                    search_mode,
                    None,
                )
                .await
            {
                Ok(response) => {
                    println!("Done!");
                    println!("DeepSeek: {}", prettify(&response));

                    // Process next level of commands
                    process_deepseek_commands_internal(
                        deepseek,
                        chat_id,
                        &response,
                        thinking_mode,
                        search_mode,
                        depth + 1,
                    )
                    .await?;
                }
                Err(e) => {
                    println!("Error: {e}");
                }
            }
        }

        // Process exec commands
        if !execs.is_empty() {
            for cmd in &execs {
                println!("\nExecuting: {cmd}");

                match execute_command(cmd) {
                    Ok(output) => {
                        println!("{output}");

                        print!("Sending command results... ");
                        io::stdout().flush()?;

                        let cmd_message =
                            format!("Command executed: {cmd}\n\nOutput:\n{output}");

                        match deepseek
                            .chat_completion(
                                chat_id,
                                &cmd_message,
                                None,
                                thinking_mode,
                                search_mode,
                                None,
                            )
                            .await
                        {
                            Ok(response) => {
                                println!("Done!");
                                println!("DeepSeek: {}", prettify(&response));

                                // Process next level of commands
                                process_deepseek_commands_internal(
                                    deepseek,
                                    chat_id,
                                    &response,
                                    thinking_mode,
                                    search_mode,
                                    depth + 1,
                                )
                                .await?;
                            }
                            Err(e) => {
                                println!("Error: {e}");
                            }
                        }
                    }
                    Err(e) => {
                        println!("Error executing command: {e}");
                    }
                }
            }
        }

        Ok(())
    })
}

/// Run the CLI with Claude provider
async fn run_claude(args: UnifiedArgs, running: Arc<AtomicBool>) -> Result<()> {
    // Load session values from config files
    let config_dir = dirs::config_dir()
        .ok_or_else(|| anyhow!("Could not determine config directory"))?
        .join("toast");

    let cookie_path = config_dir.join("cookie");
    let org_id_path = config_dir.join("org_id");

    // Check if config directory exists, if not create it and provide instructions
    if !config_dir.exists() {
        fs::create_dir_all(&config_dir).context(format!(
            "Failed to create config directory at {config_dir:?}"
        ))?;
        return Err(anyhow!(
            "Configuration directory created at {:?}\n\nPlease create a cookie file with your Claude cookie", 
            config_dir,
        ));
    }

    // Check and load cookie
    let cookie = if cookie_path.exists() {
        fs::read_to_string(&cookie_path)
            .context(format!("Failed to read cookie from {cookie_path:?}"))?
            .trim()
            .to_string()
    } else {
        return Err(anyhow!("Cookie file not found at {:?}", cookie_path,));
    };

    // Check and load org_id, or extract from cookie if file doesn't exist
    let org_id = if org_id_path.exists() {
        fs::read_to_string(&org_id_path)
            .context(format!(
                "Failed to read organization ID from {org_id_path:?}"
            ))?
            .trim()
            .to_string()
    } else {
        // Try to extract org_id from cookie
        if let Some(extracted_org_id) = extract_org_id_from_cookie(&cookie) {
            // Save the extracted org_id to the file for future use
            fs::write(&org_id_path, &extracted_org_id).context(format!(
                "Failed to write organization ID to {org_id_path:?}"
            ))?;
            println!(
                "Extracted organization ID from cookie and saved to {org_id_path:?}"
            );
            extracted_org_id
        } else {
            return Err(anyhow!(
                "Organization ID file not found at {:?} and couldn't extract it from cookie.",
                org_id_path,
            ));
        }
    };

    let user_agent =
        "Mozilla/5.0 (Macintosh; Intel Mac OS X 10.15; rv:137.0) Gecko/20100101 Firefox/137.0"
            .to_string();

    let session = ClaudeSession {
        cookie,
        user_agent,
        organization_id: org_id,
    };

    // Determine model based on flags
    let model: &str = if args.use_opus {
        OPUS_MODEL
    } else if args.use_haiku {
        HAIKU_MODEL
    } else {
        SONNET_MODEL
    };

    let claude = Claude::new(session.clone(), model)?;
    println!("Starting new Claude chat session using model: {model}");

    let stdin = io::stdin();
    let mut stdout = io::stdout();
    let mut chat_id = String::new();
    let mut system_prompt_sent = false;

    while running.load(Ordering::SeqCst) {
        print!("You: ");
        stdout.flush()?;
        let mut buf = String::new();
        stdin.read_line(&mut buf)?;
        let input = buf.trim_end();
        if input.is_empty() {
            continue;
        }
        if input.eq_ignore_ascii_case("/exit") || input.eq_ignore_ascii_case("exit") || input == "x"
        {
            if !chat_id.is_empty() {
                claude.delete_chat(&chat_id).await.ok();
            }
            break;
        }

        // Initialize chat
        if chat_id.is_empty() {
            chat_id = claude.create_chat().await.context("creating chat")?;
        }

        // Handle exec commands
        if let Some(caps) = crate::utils::EXEC_RE.captures(input) {
            let cmd = caps[1].to_string();
            if !system_prompt_sent {
                claude
                    .send_message(&chat_id, SYSTEM_PROMPT, &[])
                    .await
                    .context("sending system prompt")?;
                system_prompt_sent = true;
            }

            match execute_command(&cmd) {
                Ok(output) => {
                    let msg = format!("Command executed: {cmd}\n\n{output}");
                    let ans = claude.send_message(&chat_id, &msg, &[]).await?;
                    println!("Claude:\n{}", prettify(&ans));
                    process_claude_commands(&claude, &chat_id, &ans).await?;
                }
                Err(e) => {
                    eprintln!("Warning: command execution failed: {e}");
                    let msg = format!("Command execution failed: {e}");
                    let ans = claude.send_message(&chat_id, &msg, &[]).await?;
                    println!("Claude:\n{}", prettify(&ans));
                }
            }
            continue;
        }

        // Handle read_file commands
        if let Some(caps) = crate::utils::READ_RE.captures(input) {
            let paths: Vec<String> = caps[1].split_whitespace().map(String::from).collect();
            let path_refs: Vec<&str> = paths.iter().map(String::as_str).collect();
            if !system_prompt_sent {
                claude
                    .send_message(&chat_id, SYSTEM_PROMPT, &[])
                    .await
                    .context("sending system prompt")?;
                system_prompt_sent = true;
            }

            let rest = input.strip_prefix(&caps[0]).unwrap_or("").trim();
            let attachments = collect_claude_attachments(&path_refs).unwrap_or_default();
            let ans = claude
                .send_message(&chat_id, rest, &attachments)
                .await
                .context("sending user message")?;

            println!("Claude:\n{}", prettify(&ans));
            process_claude_commands(&claude, &chat_id, &ans).await?;
        } else {
            // Regular message
            if !system_prompt_sent {
                claude
                    .send_message(&chat_id, SYSTEM_PROMPT, &[])
                    .await
                    .context("sending system prompt")?;
                system_prompt_sent = true;
            }

            let ans = claude
                .send_message(&chat_id, input, &[])
                .await
                .context("sending user message")?;

            println!("Claude:\n{}", prettify(&ans));
            process_claude_commands(&claude, &chat_id, &ans).await?;
        }
    }

    Ok(())
}

/// Process Claude's responses for internal tool commands
async fn process_claude_commands(claude: &Claude, chat_id: &str, response: &str) -> Result<()> {
    process_claude_commands_internal(claude, chat_id, response, 0).await
}

fn process_claude_commands_internal<'a>(
    claude: &'a Claude,
    chat_id: &'a str,
    response: &'a str,
    depth: usize,
) -> Pin<Box<dyn Future<Output = Result<()>> + 'a>> {
    Box::pin(async move {
        // Avoid infinite recursion
        if depth >= MAX_INTERNAL_ITERS {
            println!("Max internal iterations reached, returning to user.");
            return Ok(());
        }

        let (reads, execs) = extract_commands(response);
        if reads.is_empty() && execs.is_empty() {
            return Ok(());
        }

        if !reads.is_empty() {
            let atts =
                collect_claude_attachments(&reads.iter().map(String::as_str).collect::<Vec<_>>())
                    .unwrap_or_default();

            match claude
                .send_message(chat_id, "read_file response:", &atts)
                .await
            {
                Ok(resp) => {
                    println!("Claude:\n{}", prettify(&resp));
                    return process_claude_commands_internal(claude, chat_id, &resp, depth + 1)
                        .await;
                }
                Err(e) => {
                    return Err(e);
                }
            }
        }

        if !execs.is_empty() {
            let mut outputs = String::new();

            for cmd in &execs {
                match execute_command(cmd) {
                    Ok(output) => outputs.push_str(&output),
                    Err(e) => outputs.push_str(&format!("Command execution failed: {e}")),
                }
                outputs.push_str("\n\n---\n\n");
            }

            match claude.send_message(chat_id, &outputs, &[]).await {
                Ok(resp) => {
                    println!("Claude:\n{}", prettify(&resp));
                    return process_claude_commands_internal(claude, chat_id, &resp, depth + 1)
                        .await;
                }
                Err(e) => {
                    return Err(e);
                }
            }
        }

        Ok(())
    })
}