toast_api/

agent_cli.rs

//! Agent-based CLI implementation

use crate::agent::{AgentConfig, AgentSession};
use crate::api::{Claude, Session as ClaudeSession};
use crate::deepseek::{DeepSeek, Session as DeepSeekSession};
use crate::utils::prettify;
use anyhow::{anyhow, Context, Result};
use std::io::{self, Read, Write};
use std::sync::atomic::Ordering;
use std::sync::mpsc;
use std::sync::Arc;
use std::thread;
use std::time::Duration;

#[cfg(unix)]
fn setup_raw_mode() -> Result<libc::termios> {
    use std::os::unix::io::AsRawFd;
    eprintln!("DEBUG: setup_raw_mode called");
    
    let stdin_fd = io::stdin().as_raw_fd();
    
    let original_termios = unsafe {
        let mut t: libc::termios = std::mem::zeroed();
        libc::tcgetattr(stdin_fd, &mut t);
        t
    };
    
    let mut raw_termios = original_termios;
    raw_termios.c_lflag &= !(libc::ICANON | libc::ECHO);
    raw_termios.c_cc[libc::VMIN] = 0;
    raw_termios.c_cc[libc::VTIME] = 1;
    
    unsafe {
        libc::tcsetattr(stdin_fd, libc::TCSANOW, &raw_termios);
    }
    
    eprintln!("DEBUG: raw mode set");
    Ok(original_termios)
}

#[cfg(not(unix))]
fn setup_raw_mode() -> Result<libc::termios> {
    Err(anyhow!("Raw mode not supported on this platform"))
}

#[cfg(unix)]
fn restore_raw_mode(original_termios: libc::termios) -> Result<()> {
    use std::os::unix::io::AsRawFd;
    
    let stdin_fd = io::stdin().as_raw_fd();
    unsafe {
        libc::tcsetattr(stdin_fd, libc::TCSANOW, &original_termios);
    }
    eprintln!("DEBUG: raw mode restored");
    Ok(())
}

#[cfg(not(unix))]
fn restore_raw_mode(_: libc::termios) -> Result<()> {
    Ok(())
}

#[cfg(unix)]
fn read_char() -> Result<Option<u8>> {
    let mut buf = [0u8; 1];
    let mut stdin = io::stdin();
    match stdin.read_exact(&mut buf) {
        Ok(_) => Ok(Some(buf[0])),
        Err(e) if e.kind() == io::ErrorKind::WouldBlock => Ok(None),
        Err(e) => Err(anyhow!("Read error: {}", e)),
    }
}

#[cfg(not(unix))]
fn read_char() -> Result<Option<u8>> {
    Ok(None)
}

fn setup_interrupt_monitor(interrupt_flag: Arc<std::sync::atomic::AtomicBool>) -> mpsc::Receiver<()> {
    let (tx, rx) = mpsc::channel();
    
    thread::spawn(move || {
        loop {
            if let Ok(Some(ch)) = read_char() {
                if ch == 27 { // ESC key
                    interrupt_flag.store(true, Ordering::Relaxed);
                    eprintln!("\n🛑 Interrupt detected (ESC)!");
                    let _ = tx.send(());
                    break;
                }
            }
            thread::sleep(Duration::from_millis(10));
        }
    });
    
    rx
}

pub async fn run_agent_cli(
    use_deepseek: bool,
    use_opus: bool,
    use_haiku: bool,
) -> Result<()> {
    println!("🤖 Starting Toast Agent...\n");

    let config = AgentConfig::default();
    let mut session = AgentSession::new(config);

    if use_deepseek {
        run_with_deepseek(&mut session, use_opus, use_haiku).await
    } else {
        run_with_claude(&mut session, use_opus, use_haiku).await
    }
}

async fn run_with_claude(
    session: &mut AgentSession,
    use_opus: bool,
    use_haiku: bool,
) -> Result<()> {
    let config_dir = dirs::config_dir()
        .ok_or_else(|| anyhow!("Could not determine config directory"))?
        .join("toast");

    let cookie = std::fs::read_to_string(config_dir.join("cookie"))
        .context("Failed to read cookie")?
        .trim()
        .to_string();

    let org_id = if let Ok(id) = std::fs::read_to_string(config_dir.join("org_id")) {
        id.trim().to_string()
    } else {
        crate::utils::extract_org_id_from_cookie(&cookie)
            .ok_or_else(|| anyhow!("Could not extract org_id from cookie"))?
    };

    let claude_session = ClaudeSession {
        cookie,
        user_agent: "Mozilla/5.0 (Macintosh; Intel Mac OS X 10.15; rv:137.0) Gecko/20100101 Firefox/137.0".to_string(),
        organization_id: org_id,
    };

    let model = if use_opus {
        crate::config::OPUS_MODEL
    } else if use_haiku {
        crate::config::HAIKU_MODEL
    } else {
        crate::config::SONNET_MODEL
    };

    let claude = Claude::new(claude_session, model)?;
    println!("Connected to Claude ({model})");
    println!("💡 Press ESC during tool execution to interrupt\n");

    let chat_id = claude.create_chat().await.context("Failed to create chat")?;
    
    let system_prompt = session.agent().get_system_prompt();
    claude.send_message(&chat_id, &system_prompt, &[]).await
        .context("Failed to send system prompt")?;

    let original_termios = setup_raw_mode()?;

    let interrupt_flag = session.interrupt_flag();
    let _interrupt_rx = setup_interrupt_monitor(interrupt_flag);

    let mut stdout = io::stdout();

    loop {
        print!("You: ");
        stdout.flush()?;

        let mut input = String::new();
        
        loop {
            match read_char() {
                Ok(Some(ch)) => {
                    if ch == b'\n' || ch == b'\r' {
                        println!();
                        break;
                    } else if ch == 27 {
                        println!("\n⏹️ Cancelled input");
                        input.clear();
                        break;
                    } else if ch == 127 {
                        if !input.is_empty() {
                            input.pop();
                            print!("\x08 \x08");
                            stdout.flush()?;
                        }
                    } else if ch >= 32 && ch < 127 {
                        input.push(ch as char);
                        print!("{}", ch as char);
                        stdout.flush()?;
                    }
                }
                Ok(None) => {
                    thread::sleep(Duration::from_millis(10));
                }
                Err(_) => break,
            }
        }

        let input = input.trim();
        if input.is_empty() {
            continue;
        }
        if matches!(input, "exit" | "quit" | "/exit" | "x") {
            break;
        }

        let response = claude.send_message(&chat_id, input, &[]).await
            .context("Failed to send message")?;

        println!("\nClaude: {}\n", prettify(&response));

        process_agent_response_claude(session, &claude, &chat_id, &response).await?;
        
        session.agent().reset();
    }

    restore_raw_mode(original_termios)?;

    claude.delete_chat(&chat_id).await.ok();
    println!("\n👋 Goodbye!");
    Ok(())
}

async fn run_with_deepseek(
    session: &mut AgentSession,
    use_opus: bool,
    use_haiku: bool,
) -> Result<()> {
    let config_dir = dirs::config_dir()
        .ok_or_else(|| anyhow!("Could not determine config directory"))?
        .join("toast")
        .join("deepseek");

    let auth_token = std::fs::read_to_string(config_dir.join("auth_token"))
        .context("Failed to read auth token")?
        .trim()
        .to_string();

    let cookies = serde_json::from_str(
        &std::fs::read_to_string(config_dir.join("cookies.json"))
            .context("Failed to read cookies")?
    )?;

    let deepseek_session = DeepSeekSession { auth_token, cookies };
    let mut deepseek = DeepSeek::new(deepseek_session)?;

    let model = if use_opus {
        "deepseek-r1"
    } else if use_haiku {
        "deepseek-lite"
    } else {
        "deepseek-r1"
    };

    println!("Connected to DeepSeek ({model})");
    println!("💡 Press ESC during tool execution to interrupt\n");

    let chat_id = deepseek.create_chat_session().await
        .context("Failed to create chat session")?;

    let thinking_mode = if model == "deepseek-r1" {
        crate::deepseek::ThinkingMode::Detailed
    } else {
        crate::deepseek::ThinkingMode::Simple
    };
    let search_mode = crate::deepseek::SearchMode::Disabled;

    let system_prompt = session.agent().get_system_prompt();
    
    let original_termios = setup_raw_mode()?;

    let interrupt_flag = session.interrupt_flag();
    let _interrupt_rx = setup_interrupt_monitor(interrupt_flag);

    let mut stdout = io::stdout();
    let mut first_message = true;

    loop {
        print!("You: ");
        stdout.flush()?;

        let mut input = String::new();
        
        loop {
            match read_char() {
                Ok(Some(ch)) => {
                    if ch == b'\n' || ch == b'\r' {
                        println!();
                        break;
                    } else if ch == 27 {
                        println!("\n⏹️ Cancelled input");
                        input.clear();
                        break;
                    } else if ch == 127 {
                        if !input.is_empty() {
                            input.pop();
                            print!("\x08 \x08");
                            stdout.flush()?;
                        }
                    } else if ch >= 32 && ch < 127 {
                        input.push(ch as char);
                        print!("{}", ch as char);
                        stdout.flush()?;
                    }
                }
                Ok(None) => {
                    thread::sleep(Duration::from_millis(10));
                }
                Err(_) => break,
            }
        }

        let input = input.trim();
        if input.is_empty() {
            continue;
        }
        if matches!(input, "exit" | "quit" | "/exit" | "x") {
            break;
        }

        let system_prompt_opt = if first_message {
            first_message = false;
            Some(system_prompt.as_str())
        } else {
            None
        };

        let response = deepseek.chat_completion(
            &chat_id,
            input,
            None,
            thinking_mode,
            search_mode,
            system_prompt_opt,
        ).await.context("Failed to send message")?;

        println!("\nDeepSeek: {}\n", prettify(&response));

        process_agent_response_deepseek(
            session,
            &mut deepseek,
            &chat_id,
            &response,
            thinking_mode,
            search_mode,
        ).await?;
        
        session.agent().reset();
    }

    restore_raw_mode(original_termios)?;

    println!("\n👋 Goodbye!");
    Ok(())
}

async fn process_agent_response_claude(
    session: &mut AgentSession,
    claude: &Claude,
    chat_id: &str,
    response: &str,
) -> Result<()> {
    let tool_results = session.agent().process_tool_calls(response).await?;
    
    if !tool_results.is_empty() {
        let mut result_message = String::from("Tool execution results:\n\n");
        for (tool_name, output) in tool_results {
            result_message.push_str(&format!("[{tool_name}]\n{output}\n\n"));
        }

        let follow_up = claude.send_message(chat_id, &result_message, &[]).await
            .context("Failed to send tool results")?;

        println!("Claude: {}\n", prettify(&follow_up));

        Box::pin(process_agent_response_claude(session, claude, chat_id, &follow_up)).await?;
    }

    Ok(())
}

async fn process_agent_response_deepseek(
    session: &mut AgentSession,
    deepseek: &mut DeepSeek,
    chat_id: &str,
    response: &str,
    thinking_mode: crate::deepseek::ThinkingMode,
    search_mode: crate::deepseek::SearchMode,
) -> Result<()> {
    let tool_results = session.agent().process_tool_calls(response).await?;
    
    if !tool_results.is_empty() {
        let mut result_message = String::from("Tool execution results:\n\n");
        for (tool_name, output) in tool_results {
            result_message.push_str(&format!("[{tool_name}]\n{output}\n\n"));
        }

        let follow_up = deepseek.chat_completion(
            chat_id,
            &result_message,
            None,
            thinking_mode,
            search_mode,
            None,
        ).await.context("Failed to send tool results")?;

        println!("DeepSeek: {}\n", prettify(&follow_up));

        Box::pin(process_agent_response_deepseek(
            session,
            deepseek,
            chat_id,
            &follow_up,
            thinking_mode,
            search_mode,
        )).await?;
    }

    Ok(())
}