toast_api/
cli.rs

1use anyhow::{anyhow, Context, Result};
2use clap::Parser;
3use ctrlc;
4use regex::Regex;
5use std::fs;
6use std::io::{self, Write};
7use std::path::Path;
8use std::process::{self, Command, Stdio};
9use std::sync::atomic::{AtomicBool, Ordering};
10use std::sync::Arc;
11
12use crate::api::{Attachment, Claude, Session};
13use crate::config::{HAIKU_MODEL, MAX_INTERNAL_ITERS, OPUS_MODEL, SONNET_MODEL, SYSTEM_PROMPT};
14use crate::utils::{extract_commands, prettify};
15
16/// CLI arguments for the default interface
17#[derive(Parser, Debug)]
18#[clap(author, version, about, long_about = None)]
19struct Args {
20    /// Use opus model
21    #[clap(long, conflicts_with_all = ["haiku", "custom_model"])]
22    opus: bool,
23
24    /// Use haiku model
25    #[clap(long, conflicts_with_all = ["opus", "custom_model"])]
26    haiku: bool,
27
28    /// Specify a custom model ID
29    #[clap(long, conflicts_with_all = ["opus", "haiku"])]
30    custom_model: Option<String>,
31}
32
33/// Run the CLI application
34pub fn run() -> Result<()> {
35    tokio::runtime::Builder::new_multi_thread()
36        .enable_all()
37        .build()
38        .unwrap()
39        .block_on(async_main())
40}
41
42fn get_config_help(file_name: &str) -> String {
43    let cookie_help = "To get your cookie:
441. Go to claude.ai in your browser
452. Open Developer Tools (F12 or right-click and select 'Inspect')
463. Go to the Network tab
474. Refresh the page
485. Click on any request to claude.ai
496. In the 'Headers' tab, find 'Request Headers'
507. Look for the 'Cookie' header
518. Copy the entire cookie value and save it to this folder with filename: cookie";
52
53    match file_name {
54        "cookie" => cookie_help.to_string(),
55        _ => format!("Configuration file {} is missing.", file_name),
56    }
57}
58
59/// Extract organization ID from cookie string
60pub fn extract_org_id_from_cookie(cookie: &str) -> Option<String> {
61    let re = Regex::new(r"lastActiveOrg=([0-9a-f-]+)").ok()?;
62    re.captures(cookie)
63        .and_then(|caps| caps.get(1))
64        .map(|m| m.as_str().to_string())
65}
66
67pub async fn async_main() -> Result<()> {
68    let args = Args::parse();
69
70    // Load session values from config files
71    let config_dir = dirs::config_dir()
72        .ok_or_else(|| anyhow!("Could not determine config directory"))?
73        .join("toast");
74
75    let cookie_path = config_dir.join("cookie");
76    let org_id_path = config_dir.join("org_id");
77
78    // Check if config directory exists, if not create it and provide instructions
79    if !config_dir.exists() {
80        fs::create_dir_all(&config_dir).context(format!(
81            "Failed to create config directory at {:?}",
82            config_dir
83        ))?;
84        return Err(anyhow!(
85            "Configuration directory created at {:?}\n\nPlease create the following files:\n\n1. cookie file:\n{}\n\n2", 
86            config_dir,
87            get_config_help("cookie"),
88        ));
89    }
90
91    // Check and load cookie
92    let cookie = if cookie_path.exists() {
93        fs::read_to_string(&cookie_path)
94            .context(format!("Failed to read cookie from {:?}", cookie_path))?
95            .trim()
96            .to_string()
97    } else {
98        return Err(anyhow!(
99            "Cookie file not found at {:?}\n\n{}",
100            cookie_path,
101            get_config_help("cookie")
102        ));
103    };
104
105    // Check and load org_id, or extract from cookie if file doesn't exist
106    let org_id = if org_id_path.exists() {
107        fs::read_to_string(&org_id_path)
108            .context(format!(
109                "Failed to read organization ID from {:?}",
110                org_id_path
111            ))?
112            .trim()
113            .to_string()
114    } else {
115        // Try to extract org_id from cookie
116        if let Some(extracted_org_id) = extract_org_id_from_cookie(&cookie) {
117            // Save the extracted org_id to the file for future use
118            fs::write(&org_id_path, &extracted_org_id).context(format!(
119                "Failed to write organization ID to {:?}",
120                org_id_path
121            ))?;
122            println!(
123                "Extracted organization ID from cookie and saved to {:?}",
124                org_id_path
125            );
126            extracted_org_id
127        } else {
128            return Err(anyhow!(
129                "Organization ID file not found at {:?} and couldn't extract it from cookie.\n\n{}",
130                org_id_path,
131                get_config_help("org_id")
132            ));
133        }
134    };
135
136    let user_agent =
137        "Mozilla/5.0 (Macintosh; Intel Mac OS X 10.15; rv:137.0) Gecko/20100101 Firefox/137.0"
138            .to_string();
139
140    let session = Session {
141        cookie,
142        user_agent,
143        organization_id: org_id,
144    };
145
146    // Determine model based on flags
147    let model: &str = if let Some(custom) = args.custom_model {
148        Box::leak(custom.into_boxed_str())
149    } else if args.opus {
150        OPUS_MODEL
151    } else if args.haiku {
152        HAIKU_MODEL
153    } else {
154        SONNET_MODEL
155    };
156    let claude = Claude::new(session.clone(), model)?;
157
158    // Ctrl-C handler
159    let running = Arc::new(AtomicBool::new(true));
160    {
161        let running = running.clone();
162        ctrlc::set_handler(move || {
163            running.store(false, Ordering::SeqCst);
164            println!("\nGoodbye!");
165            process::exit(0);
166        })?;
167    }
168
169    let stdin = io::stdin();
170    let mut stdout = io::stdout();
171    let mut chat_id = String::new();
172    let mut system_prompt_sent = false;
173
174    while running.load(Ordering::SeqCst) {
175        print!("You: ");
176        stdout.flush()?;
177        let mut buf = String::new();
178        stdin.read_line(&mut buf)?;
179        let input = buf.trim_end();
180        if input == "" {
181            continue;
182        }
183        if input.eq_ignore_ascii_case("/exit") || input.eq_ignore_ascii_case("exit") || input == "x"
184        {
185            if !chat_id.is_empty() {
186                claude.delete_chat(&chat_id).await.ok();
187            }
188            break;
189        }
190
191        // Initialize chat
192        if chat_id.is_empty() {
193            chat_id = claude.create_chat().await.context("creating chat")?;
194        }
195
196        // Handle exec commands
197        if let Some(caps) = crate::utils::EXEC_RE.captures(input) {
198            let cmd = caps[1].to_string();
199            if !system_prompt_sent {
200                claude
201                    .send_message(&chat_id, SYSTEM_PROMPT, &[])
202                    .await
203                    .context("sending system prompt")?;
204                system_prompt_sent = true;
205            }
206            run_exec(&claude, &chat_id, &cmd).await?;
207            continue;
208        }
209
210        // Handle read_file commands
211        if let Some(caps) = crate::utils::READ_RE.captures(input) {
212            let paths: Vec<String> = caps[1].split_whitespace().map(String::from).collect();
213            let path_refs: Vec<&str> = paths.iter().map(String::as_str).collect();
214            if !system_prompt_sent {
215                claude
216                    .send_message(&chat_id, SYSTEM_PROMPT, &[])
217                    .await
218                    .context("sending system prompt")?;
219                system_prompt_sent = true;
220            }
221            let rest = input.strip_prefix(&caps[0]).unwrap_or("").trim();
222            let attachments = collect_attachments(&path_refs).unwrap_or_default();
223            let ans = claude
224                .send_message(&chat_id, rest, &attachments)
225                .await
226                .context("sending user message")?;
227            println!("Claude:\n{}", prettify(&ans));
228            process_claude(&claude, &chat_id, ans).await?;
229        } else {
230            // Regular message
231            if !system_prompt_sent {
232                claude
233                    .send_message(&chat_id, SYSTEM_PROMPT, &[])
234                    .await
235                    .context("sending system prompt")?;
236                system_prompt_sent = true;
237            }
238            let ans = claude
239                .send_message(&chat_id, input, &[])
240                .await
241                .context("sending user message")?;
242            println!("Claude:\n{}", prettify(&ans));
243            process_claude(&claude, &chat_id, ans).await?;
244        }
245    }
246    Ok(())
247}
248
249/// Execute shell command and send output to Claude
250async fn run_exec(claude: &Claude, chat_id: &str, cmd: &str) -> Result<()> {
251    let out = match execute_command(cmd) {
252        Ok(output) => output,
253        Err(e) => {
254            eprintln!("Warning: command execution failed: {e}");
255            format!("Command execution failed: {e}")
256        }
257    };
258    let msg = format!("Command executed: {cmd}\n\n{out}");
259    let ans = claude.send_message(chat_id, &msg, &[]).await?;
260    println!("Claude:\n{}", prettify(&ans));
261    process_claude(claude, chat_id, ans).await
262}
263
264/// Execute a shell command and capture its output
265fn execute_command(command: &str) -> Result<String> {
266    let result = Command::new("sh")
267        .arg("-c")
268        .arg(command)
269        .stdout(Stdio::piped())
270        .stderr(Stdio::piped())
271        .spawn()?;
272    let output = result.wait_with_output()?;
273    let mut msg = String::new();
274    if !output.stdout.is_empty() {
275        msg.push_str("=== STDOUT ===\n");
276        msg.push_str(&String::from_utf8_lossy(&output.stdout));
277        msg.push('\n');
278    }
279    if !output.stderr.is_empty() {
280        msg.push_str("=== STDERR ===\n");
281        msg.push_str(&String::from_utf8_lossy(&output.stderr));
282        msg.push('\n');
283    }
284    msg.push_str(&format!(
285        "Exit code: {}",
286        output.status.code().unwrap_or(-1)
287    ));
288    Ok(msg)
289}
290
291/// Read files into attachments
292fn collect_attachments(paths: &[&str]) -> Result<Vec<Attachment>> {
293    const LIMIT: usize = 5;
294    const SIZE_LIMIT: u64 = 10 * 1024 * 1024;
295    if paths.len() > LIMIT {
296        return Err(anyhow!("cannot attach more than {LIMIT} files"));
297    }
298    let mut atts = Vec::new();
299    for p in paths {
300        if let Ok(meta) = fs::metadata(p) {
301            if meta.len() > SIZE_LIMIT {
302                eprintln!("Warning: file {p} is larger than 10 MB, skipping");
303                continue;
304            }
305            if let Ok(content) = fs::read_to_string(p) {
306                atts.push(Attachment {
307                    file_name: Path::new(p)
308                        .file_name()
309                        .unwrap_or_default()
310                        .to_string_lossy()
311                        .into(),
312                    size: meta.len(),
313                    content,
314                });
315            } else {
316                eprintln!("Warning: couldn't read file {p}");
317            }
318        } else {
319            eprintln!("Warning: couldn't access file {p}");
320        }
321    }
322    Ok(atts)
323}
324
325/// Process Claude's responses for internal tool commands
326async fn process_claude(claude: &Claude, chat_id: &str, mut ans: String) -> Result<()> {
327    for _ in 0..MAX_INTERNAL_ITERS {
328        let (reads, execs) = extract_commands(&ans);
329        if reads.is_empty() && execs.is_empty() {
330            return Ok(());
331        }
332        if !reads.is_empty() {
333            let atts = collect_attachments(&reads.iter().map(String::as_str).collect::<Vec<_>>())
334                .unwrap_or_default();
335            let ans2 = claude
336                .send_message(chat_id, "read_file response:", &atts)
337                .await?;
338            println!("Claude:\n{}", prettify(&ans2));
339            ans = ans2;
340            continue;
341        }
342        if !execs.is_empty() {
343            let mut outputs = String::new();
344            for cmd in &execs {
345                match execute_command(cmd) {
346                    Ok(output) => outputs.push_str(&output),
347                    Err(e) => outputs.push_str(&format!("Command execution failed: {e}")),
348                }
349                outputs.push_str("\n\n---\n\n");
350            }
351            let ans2 = claude.send_message(chat_id, &outputs, &[]).await?;
352            println!("Claude:\n{}", prettify(&ans2));
353            ans = ans2;
354            continue;
355        }
356    }
357    println!("Max internal iterations reached, returning to user.");
358    Ok(())
359}
toast_api/cli.rs

toast_api/
cli.rs
