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, Write};

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

    // Create agent with default config
    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<()> {
    // Load Claude configuration
    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})\n");

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

    // Main interaction loop
    let stdin = io::stdin();
    let mut stdout = io::stdout();

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

        let mut input = String::new();
        match stdin.read_line(&mut input) {
            Ok(0) => break, // EOF
            Ok(_) => {
                let input = input.trim();
                if input.is_empty() {
                    continue;
                }
                if matches!(input, "exit" | "quit" | "/exit" | "x") {
                    break;
                }

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

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

                // Process any tool calls in the response
                process_agent_response_claude(session, &claude, &chat_id, &response).await?;
                
                // Reset iteration count for next user interaction
                session.agent().reset();
            }
            Err(e) => {
                eprintln!("Error reading input: {e}");
                break;
            }
        }
    }

    // Clean up
    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<()> {
    // Load DeepSeek configuration
    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})\n");

    // Create chat session
    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;

    // Send initial system prompt
    let system_prompt = session.agent().get_system_prompt();
    
    // Main interaction loop
    let stdin = io::stdin();
    let mut stdout = io::stdout();
    let mut first_message = true;

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

        let mut input = String::new();
        match stdin.read_line(&mut input) {
            Ok(0) => break, // EOF
            Ok(_) => {
                let input = input.trim();
                if input.is_empty() {
                    continue;
                }
                if matches!(input, "exit" | "quit" | "/exit" | "x") {
                    break;
                }

                // Include system prompt only on first message
                let system_prompt_opt = if first_message {
                    first_message = false;
                    Some(system_prompt.as_str())
                } else {
                    None
                };

                // Send message to DeepSeek
                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 any tool calls in the response
                process_agent_response_deepseek(
                    session,
                    &mut deepseek,
                    &chat_id,
                    &response,
                    thinking_mode,
                    search_mode,
                ).await?;
                
                // Reset iteration count for next user interaction
                session.agent().reset();
            }
            Err(e) => {
                eprintln!("Error reading input: {e}");
                break;
            }
        }
    }

    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() {
        // Format tool results for Claude
        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"));
        }

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

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

        // Recursively process any new tool calls
        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() {
        // Format tool results for DeepSeek
        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"));
        }

        // Send results back to DeepSeek
        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));

        // Recursively process any new tool calls
        Box::pin(process_agent_response_deepseek(
            session,
            deepseek,
            chat_id,
            &follow_up,
            thinking_mode,
            search_mode,
        )).await?;
    }

    Ok(())
}