toast_api/
unified_cli.rs

1use anyhow::{anyhow, Context, Result};
2use ctrlc;
3use regex::Regex;
4use std::fs;
5use std::io::{self, Write};
6use std::path::Path;
7use std::process::{self, Command, Stdio};
8use std::sync::atomic::{AtomicBool, Ordering};
9use std::sync::Arc;
10use std::future::Future;
11use std::pin::Pin;
12
13use crate::api::{Attachment, Claude, Session as ClaudeSession};
14use crate::deepseek::{DeepSeek, Session as DeepSeekSession};
15use crate::config::{HAIKU_MODEL, MAX_INTERNAL_ITERS, OPUS_MODEL, SONNET_MODEL, SYSTEM_PROMPT};
16use crate::utils::{extract_commands, prettify};
17
18/// Unified CLI arguments for both Claude and DeepSeek
19#[derive(Debug)]
20pub struct UnifiedArgs {
21    pub use_deepseek: bool,
22    pub use_opus: bool,
23    pub use_haiku: bool,
24}
25
26/// Execute a shell command and capture its output
27fn execute_command(command: &str) -> Result<String> {
28    let result = Command::new("sh")
29        .arg("-c")
30        .arg(command)
31        .stdout(Stdio::piped())
32        .stderr(Stdio::piped())
33        .spawn()?;
34    
35    let output = result.wait_with_output()?;
36    let mut msg = String::new();
37    
38    if !output.stdout.is_empty() {
39        msg.push_str("=== STDOUT ===\n");
40        msg.push_str(&String::from_utf8_lossy(&output.stdout));
41        msg.push('\n');
42    }
43    
44    if !output.stderr.is_empty() {
45        msg.push_str("=== STDERR ===\n");
46        msg.push_str(&String::from_utf8_lossy(&output.stderr));
47        msg.push('\n');
48    }
49    
50    msg.push_str(&format!(
51        "Exit code: {}",
52        output.status.code().unwrap_or(-1)
53    ));
54    
55    Ok(msg)
56}
57
58/// Read files into attachments for Claude
59fn collect_claude_attachments(paths: &[&str]) -> Result<Vec<Attachment>> {
60    const LIMIT: usize = 5;
61    const SIZE_LIMIT: u64 = 10 * 1024 * 1024;
62    
63    if paths.len() > LIMIT {
64        return Err(anyhow!("cannot attach more than {LIMIT} files"));
65    }
66    
67    let mut atts = Vec::new();
68    
69    for p in paths {
70        if let Ok(meta) = fs::metadata(p) {
71            if meta.len() > SIZE_LIMIT {
72                eprintln!("Warning: file {p} is larger than 10 MB, skipping");
73                continue;
74            }
75            
76            if let Ok(content) = fs::read_to_string(p) {
77                atts.push(Attachment {
78                    file_name: Path::new(p)
79                        .file_name()
80                        .unwrap_or_default()
81                        .to_string_lossy()
82                        .into(),
83                    size: meta.len(),
84                    content,
85                });
86            } else {
87                eprintln!("Warning: couldn't read file {p}");
88            }
89        } else {
90            eprintln!("Warning: couldn't access file {p}");
91        }
92    }
93    
94    Ok(atts)
95}
96
97/// Run the unified CLI application with provider selection
98pub async fn run(args: UnifiedArgs) -> Result<()> {
99    // Set up Ctrl-C handler
100    let running = Arc::new(AtomicBool::new(true));
101    {
102        let running = running.clone();
103        ctrlc::set_handler(move || {
104            running.store(false, Ordering::SeqCst);
105            println!("\nGoodbye!");
106            process::exit(0);
107        })?;
108    }
109    
110    if args.use_deepseek {
111        run_deepseek(args, running).await
112    } else {
113        run_claude(args, running).await
114    }
115}
116
117/// Run the CLI with DeepSeek provider
118async fn run_deepseek(args: UnifiedArgs, running: Arc<AtomicBool>) -> Result<()> {
119    // Load session values from config files
120    let config_dir = dirs::config_dir()
121        .ok_or_else(|| anyhow!("Could not determine config directory"))?
122        .join("toast")
123        .join("deepseek");
124    
125    // Create config directory if it doesn't exist
126    if !config_dir.exists() {
127        fs::create_dir_all(&config_dir)?;
128    }
129    
130    let auth_token_path = config_dir.join("auth_token");
131    let cookies_path = config_dir.join("cookies.json");
132    
133    // Check auth token
134    let auth_token = if auth_token_path.exists() {
135        fs::read_to_string(&auth_token_path)
136            .context(format!("Failed to read auth token from {:?}", auth_token_path))?
137            .trim()
138            .to_string()
139    } else {
140        return Err(anyhow!(
141            "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",
142            auth_token_path
143        ));
144    };
145    
146    // Check cookies
147    let cookies = if cookies_path.exists() {
148        serde_json::from_str(&fs::read_to_string(&cookies_path)
149            .context(format!("Failed to read cookies from {:?}", cookies_path))?)?
150    } else {
151        return Err(anyhow!(
152            "Cookies file not found at {:?}\n\nDeepSeek requires Cloudflare cookies.\nUse the deepseek4free library to generate them.",
153            cookies_path
154        ));
155    };
156    
157    let session = DeepSeekSession {
158        auth_token,
159        cookies,
160    };
161    
162    // Determine model based on flags
163    let model = if args.use_opus {
164        "deepseek-coder"
165    } else if args.use_haiku {
166        "deepseek-lite"
167    } else {
168        "deepseek-chat" // Default model
169    };
170    
171    let mut deepseek = DeepSeek::new_with_model(session, model)?;
172    
173    let stdin = io::stdin();
174    let mut stdout = io::stdout();
175    
176    // Create a new chat session
177    println!("Starting new DeepSeek chat session...");
178    let chat_id = match deepseek.create_chat_session().await {
179        Ok(id) => {
180            println!("Session started with DeepSeek!\n");
181            id
182        }
183        Err(e) => {
184            return Err(anyhow!("Failed to create DeepSeek chat session: {}", e));
185        }
186    };
187    
188    // Set thinking mode to disabled for now
189    let thinking_mode = crate::deepseek::ThinkingMode::Disabled;
190    let search_mode = crate::deepseek::SearchMode::Disabled;
191    
192    // Main chat loop
193    while running.load(Ordering::SeqCst) {
194        print!("You: ");
195        stdout.flush()?;
196        
197        let mut buf = String::new();
198        stdin.read_line(&mut buf)?;
199        let input = buf.trim_end();
200        
201        // Check for empty input or exit commands
202        if input.is_empty() {
203            continue;
204        }
205        
206        if input.eq_ignore_ascii_case("/exit") || input.eq_ignore_ascii_case("exit") || input == "x" {
207            break;
208        }
209        
210        // Handle exec commands
211        if let Some(caps) = crate::utils::EXEC_RE.captures(input) {
212            let cmd = caps[1].to_string();
213            match execute_command(&cmd) {
214                Ok(output) => {
215                    let msg = format!("Command executed: {cmd}\n\n{output}");
216                    print!("DeepSeek: ");
217                    stdout.flush()?;
218                    
219                    match deepseek.chat_completion(&chat_id, &msg, None, thinking_mode.clone(), search_mode.clone()).await {
220                        Ok(response) => {
221                            println!("{}", prettify(&response));
222                            process_deepseek_commands(&mut deepseek, &chat_id, &response, thinking_mode.clone(), search_mode.clone()).await?;
223                        }
224                        Err(e) => {
225                            eprintln!("\nError: {}", e);
226                        }
227                    }
228                }
229                Err(e) => {
230                    eprintln!("Warning: command execution failed: {e}");
231                }
232            }
233            continue;
234        }
235        
236        // Handle read_file commands
237        if let Some(caps) = crate::utils::READ_RE.captures(input) {
238            let paths: Vec<String> = caps[1].split_whitespace().map(String::from).collect();
239            let mut file_contents = Vec::new();
240            
241            for path in &paths {
242                match fs::read_to_string(path) {
243                    Ok(content) => {
244                        file_contents.push(format!("=== File: {} ===\n{}", path, content));
245                    }
246                    Err(e) => {
247                        file_contents.push(format!("Error reading file {}: {}", path, e));
248                    }
249                }
250            }
251            
252            let file_message = format!("Here are the contents of the files you requested:\n\n{}", 
253                                      file_contents.join("\n\n"));
254            
255            print!("DeepSeek: ");
256            stdout.flush()?;
257            
258            match deepseek.chat_completion(&chat_id, &file_message, None, thinking_mode.clone(), search_mode.clone()).await {
259                Ok(response) => {
260                    println!("{}", prettify(&response));
261                    process_deepseek_commands(&mut deepseek, &chat_id, &response, thinking_mode.clone(), search_mode.clone()).await?;
262                }
263                Err(e) => {
264                    eprintln!("\nError: {}", e);
265                }
266            }
267            continue;
268        }
269        
270        // Regular message
271        print!("DeepSeek: ");
272        stdout.flush()?;
273        
274        match deepseek.chat_completion(&chat_id, input, None, thinking_mode.clone(), search_mode.clone()).await {
275            Ok(response) => {
276                println!("{}", prettify(&response));
277                process_deepseek_commands(&mut deepseek, &chat_id, &response, thinking_mode.clone(), search_mode.clone()).await?;
278            }
279            Err(e) => {
280                eprintln!("\nError: {}", e);
281            }
282        }
283        println!();
284    }
285    
286    Ok(())
287}
288
289/// Process commands in DeepSeek's response
290async fn process_deepseek_commands(
291    deepseek: &mut DeepSeek, 
292    chat_id: &str, 
293    response: &str,
294    thinking_mode: crate::deepseek::ThinkingMode,
295    search_mode: crate::deepseek::SearchMode,
296) -> Result<()> {
297    process_deepseek_commands_internal(deepseek, chat_id, response, thinking_mode, search_mode, 0).await
298}
299
300fn process_deepseek_commands_internal<'a>(
301    deepseek: &'a mut DeepSeek, 
302    chat_id: &'a str, 
303    response: &'a str,
304    thinking_mode: crate::deepseek::ThinkingMode,
305    search_mode: crate::deepseek::SearchMode,
306    depth: usize,
307) -> Pin<Box<dyn Future<Output = Result<()>> + 'a>> {
308    Box::pin(async move {
309        // Avoid infinite recursion
310        if depth >= MAX_INTERNAL_ITERS {
311            println!("Max internal iterations reached, returning to user.");
312            return Ok(());
313        }
314        
315        let (reads, execs) = extract_commands(response);
316        if reads.is_empty() && execs.is_empty() {
317            return Ok(());
318        }
319        
320        if !reads.is_empty() {
321            let mut file_contents = Vec::new();
322            
323            for path in &reads {
324                match fs::read_to_string(path) {
325                    Ok(content) => {
326                        file_contents.push(format!("=== File: {} ===\n{}", path, content));
327                    }
328                    Err(e) => {
329                        file_contents.push(format!("Error reading file {}: {}", path, e));
330                    }
331                }
332            }
333            
334            let file_message = format!("Here are the contents of the files you requested:\n\n{}", 
335                                      file_contents.join("\n\n"));
336            
337            match deepseek.chat_completion(chat_id, &file_message, None, thinking_mode.clone(), search_mode.clone()).await {
338                Ok(resp) => {
339                    println!("DeepSeek: {}", prettify(&resp));
340                    return process_deepseek_commands_internal(deepseek, chat_id, &resp, thinking_mode, search_mode, depth + 1).await;
341                }
342                Err(e) => {
343                    eprintln!("Error: {}", e);
344                    return Err(e);
345                }
346            }
347        }
348        
349        if !execs.is_empty() {
350            let mut outputs = String::new();
351            
352            for cmd in &execs {
353                match execute_command(cmd) {
354                    Ok(output) => outputs.push_str(&output),
355                    Err(e) => outputs.push_str(&format!("Command execution failed: {e}")),
356                }
357                outputs.push_str("\n\n---\n\n");
358            }
359            
360            match deepseek.chat_completion(chat_id, &outputs, None, thinking_mode.clone(), search_mode.clone()).await {
361                Ok(resp) => {
362                    println!("DeepSeek: {}", prettify(&resp));
363                    return process_deepseek_commands_internal(deepseek, chat_id, &resp, thinking_mode, search_mode, depth + 1).await;
364                }
365                Err(e) => {
366                    eprintln!("Error: {}", e);
367                    return Err(e);
368                }
369            }
370        }
371        
372        Ok(())
373    })
374}
375
376/// Extract organization ID from cookie string
377fn extract_org_id_from_cookie(cookie: &str) -> Option<String> {
378    let re = Regex::new(r"lastActiveOrg=([0-9a-f-]+)").ok()?;
379    re.captures(cookie)
380        .and_then(|caps| caps.get(1))
381        .map(|m| m.as_str().to_string())
382}
383
384/// Run the CLI with Claude provider
385async fn run_claude(args: UnifiedArgs, running: Arc<AtomicBool>) -> Result<()> {
386    // Load session values from config files
387    let config_dir = dirs::config_dir()
388        .ok_or_else(|| anyhow!("Could not determine config directory"))?
389        .join("toast");
390
391    let cookie_path = config_dir.join("cookie");
392    let org_id_path = config_dir.join("org_id");
393
394    // Check if config directory exists, if not create it and provide instructions
395    if !config_dir.exists() {
396        fs::create_dir_all(&config_dir).context(format!(
397            "Failed to create config directory at {:?}",
398            config_dir
399        ))?;
400        return Err(anyhow!(
401            "Configuration directory created at {:?}\n\nPlease create a cookie file with your Claude cookie", 
402            config_dir,
403        ));
404    }
405
406    // Check and load cookie
407    let cookie = if cookie_path.exists() {
408        fs::read_to_string(&cookie_path)
409            .context(format!("Failed to read cookie from {:?}", cookie_path))?
410            .trim()
411            .to_string()
412    } else {
413        return Err(anyhow!(
414            "Cookie file not found at {:?}",
415            cookie_path,
416        ));
417    };
418
419    // Check and load org_id, or extract from cookie if file doesn't exist
420    let org_id = if org_id_path.exists() {
421        fs::read_to_string(&org_id_path)
422            .context(format!(
423                "Failed to read organization ID from {:?}",
424                org_id_path
425            ))?
426            .trim()
427            .to_string()
428    } else {
429        // Try to extract org_id from cookie
430        if let Some(extracted_org_id) = extract_org_id_from_cookie(&cookie) {
431            // Save the extracted org_id to the file for future use
432            fs::write(&org_id_path, &extracted_org_id).context(format!(
433                "Failed to write organization ID to {:?}",
434                org_id_path
435            ))?;
436            println!(
437                "Extracted organization ID from cookie and saved to {:?}",
438                org_id_path
439            );
440            extracted_org_id
441        } else {
442            return Err(anyhow!(
443                "Organization ID file not found at {:?} and couldn't extract it from cookie.",
444                org_id_path,
445            ));
446        }
447    };
448
449    let user_agent =
450        "Mozilla/5.0 (Macintosh; Intel Mac OS X 10.15; rv:137.0) Gecko/20100101 Firefox/137.0"
451            .to_string();
452
453    let session = ClaudeSession {
454        cookie,
455        user_agent,
456        organization_id: org_id,
457    };
458
459    // Determine model based on flags
460    let model: &str = if args.use_opus {
461        OPUS_MODEL
462    } else if args.use_haiku {
463        HAIKU_MODEL
464    } else {
465        SONNET_MODEL
466    };
467
468    let claude = Claude::new(session.clone(), model)?;
469    println!("Starting new Claude chat session using model: {}", model);
470
471    let stdin = io::stdin();
472    let mut stdout = io::stdout();
473    let mut chat_id = String::new();
474    let mut system_prompt_sent = false;
475
476    while running.load(Ordering::SeqCst) {
477        print!("You: ");
478        stdout.flush()?;
479        let mut buf = String::new();
480        stdin.read_line(&mut buf)?;
481        let input = buf.trim_end();
482        if input == "" {
483            continue;
484        }
485        if input.eq_ignore_ascii_case("/exit") || input.eq_ignore_ascii_case("exit") || input == "x"
486        {
487            if !chat_id.is_empty() {
488                claude.delete_chat(&chat_id).await.ok();
489            }
490            break;
491        }
492
493        // Initialize chat
494        if chat_id.is_empty() {
495            chat_id = claude.create_chat().await.context("creating chat")?;
496        }
497
498        // Handle exec commands
499        if let Some(caps) = crate::utils::EXEC_RE.captures(input) {
500            let cmd = caps[1].to_string();
501            if !system_prompt_sent {
502                claude
503                    .send_message(&chat_id, SYSTEM_PROMPT, &[])
504                    .await
505                    .context("sending system prompt")?;
506                system_prompt_sent = true;
507            }
508            
509            match execute_command(&cmd) {
510                Ok(output) => {
511                    let msg = format!("Command executed: {cmd}\n\n{output}");
512                    let ans = claude.send_message(&chat_id, &msg, &[]).await?;
513                    println!("Claude:\n{}", prettify(&ans));
514                    process_claude_commands(&claude, &chat_id, &ans).await?;
515                }
516                Err(e) => {
517                    eprintln!("Warning: command execution failed: {e}");
518                    let msg = format!("Command execution failed: {e}");
519                    let ans = claude.send_message(&chat_id, &msg, &[]).await?;
520                    println!("Claude:\n{}", prettify(&ans));
521                }
522            }
523            continue;
524        }
525
526        // Handle read_file commands
527        if let Some(caps) = crate::utils::READ_RE.captures(input) {
528            let paths: Vec<String> = caps[1].split_whitespace().map(String::from).collect();
529            let path_refs: Vec<&str> = paths.iter().map(String::as_str).collect();
530            if !system_prompt_sent {
531                claude
532                    .send_message(&chat_id, SYSTEM_PROMPT, &[])
533                    .await
534                    .context("sending system prompt")?;
535                system_prompt_sent = true;
536            }
537            
538            let rest = input.strip_prefix(&caps[0]).unwrap_or("").trim();
539            let attachments = collect_claude_attachments(&path_refs).unwrap_or_default();
540            let ans = claude
541                .send_message(&chat_id, rest, &attachments)
542                .await
543                .context("sending user message")?;
544                
545            println!("Claude:\n{}", prettify(&ans));
546            process_claude_commands(&claude, &chat_id, &ans).await?;
547        } else {
548            // Regular message
549            if !system_prompt_sent {
550                claude
551                    .send_message(&chat_id, SYSTEM_PROMPT, &[])
552                    .await
553                    .context("sending system prompt")?;
554                system_prompt_sent = true;
555            }
556            
557            let ans = claude
558                .send_message(&chat_id, input, &[])
559                .await
560                .context("sending user message")?;
561                
562            println!("Claude:\n{}", prettify(&ans));
563            process_claude_commands(&claude, &chat_id, &ans).await?;
564        }
565    }
566    
567    Ok(())
568}
569
570/// Process Claude's responses for internal tool commands
571async fn process_claude_commands(claude: &Claude, chat_id: &str, response: &str) -> Result<()> {
572    process_claude_commands_internal(claude, chat_id, response, 0).await
573}
574
575fn process_claude_commands_internal<'a>(claude: &'a Claude, chat_id: &'a str, response: &'a str, depth: usize) -> Pin<Box<dyn Future<Output = Result<()>> + 'a>> {
576    Box::pin(async move {
577        // Avoid infinite recursion
578        if depth >= MAX_INTERNAL_ITERS {
579            println!("Max internal iterations reached, returning to user.");
580            return Ok(());
581        }
582        
583        let (reads, execs) = extract_commands(response);
584        if reads.is_empty() && execs.is_empty() {
585            return Ok(());
586        }
587        
588        if !reads.is_empty() {
589            let atts = collect_claude_attachments(&reads.iter().map(String::as_str).collect::<Vec<_>>())
590                .unwrap_or_default();
591                
592            match claude.send_message(chat_id, "read_file response:", &atts).await {
593                Ok(resp) => {
594                    println!("Claude:\n{}", prettify(&resp));
595                    return process_claude_commands_internal(claude, chat_id, &resp, depth + 1).await;
596                }
597                Err(e) => {
598                    return Err(e);
599                }
600            }
601        }
602        
603        if !execs.is_empty() {
604            let mut outputs = String::new();
605            
606            for cmd in &execs {
607                match execute_command(cmd) {
608                    Ok(output) => outputs.push_str(&output),
609                    Err(e) => outputs.push_str(&format!("Command execution failed: {e}")),
610                }
611                outputs.push_str("\n\n---\n\n");
612            }
613            
614            match claude.send_message(chat_id, &outputs, &[]).await {
615                Ok(resp) => {
616                    println!("Claude:\n{}", prettify(&resp));
617                    return process_claude_commands_internal(claude, chat_id, &resp, depth + 1).await;
618                }
619                Err(e) => {
620                    return Err(e);
621                }
622            }
623        }
624        
625        Ok(())
626    })
627}
toast_api/unified_cli.rs

toast_api/
unified_cli.rs
