toast_api/
utils.rs

1use anyhow::{anyhow, Result};
2use brotli;
3use flate2::read::{DeflateDecoder, GzDecoder};
4use rquest::Response;
5use serde::Deserialize;
6use std::io::Read;
7// For JSON parsing
8use chrono::Local;
9use lazy_static::lazy_static;
10use regex::Regex;
11use serde_json;
12
13/// Decode response body based on Content-Encoding
14pub async fn decode_body(res: Response) -> Result<Vec<u8>> {
15    // Extract the content encoding header as an owned String to avoid borrowing `res`
16    let encoding = res
17        .headers()
18        .get("Content-Encoding")
19        .and_then(|v| v.to_str().ok())
20        .unwrap_or("")
21        .to_string();
22
23    // Consume response to get bytes
24    let bytes_result = res.bytes().await;
25
26    // Handle possible empty body or connection errors
27    let buf = match bytes_result {
28        Ok(bytes) if bytes.is_empty() => {
29            return Err(anyhow!("Empty response body received from server"))
30        }
31        Ok(bytes) => bytes.to_vec(),
32        Err(e) => return Err(anyhow!("Failed to read response body: {}", e)),
33    };
34
35    // Process based on encoding
36    let out = match encoding.as_str() {
37        "gzip" => {
38            let mut d = GzDecoder::new(&buf[..]);
39            let mut v = Vec::new();
40            d.read_to_end(&mut v)?;
41            v
42        }
43        "deflate" => {
44            let mut d = DeflateDecoder::new(&buf[..]);
45            let mut v = Vec::new();
46            d.read_to_end(&mut v)?;
47            v
48        }
49        "br" => {
50            let mut v = Vec::new();
51            brotli::Decompressor::new(&buf[..], 4096).read_to_end(&mut v)?;
52            v
53        }
54        _ => buf,
55    };
56    Ok(out)
57}
58
59/// Parsed line from Claude stream
60#[derive(Debug, Deserialize)]
61struct SendLine {
62    #[serde(default)]
63    completion: String,
64    #[serde(default)]
65    error: Option<ClaudeErrorInfo>,
66}
67
68/// Info about errors from Claude
69#[derive(Debug, Deserialize)]
70struct ClaudeErrorInfo {
71    #[serde(default)]
72    #[allow(non_snake_case)]
73    r#type: String,
74    #[serde(default)]
75    message: String,
76    #[serde(default)]
77    resets_at: Option<i64>,
78}
79
80/// Parse the event-stream bytes into a completion string
81pub fn parse_stream(bytes: &[u8]) -> Result<String> {
82    let text = String::from_utf8_lossy(bytes);
83    let mut completions = String::new();
84    for line in text.lines() {
85        if let Some(start) = line.find('{') {
86            if let Ok(obj) = serde_json::from_str::<SendLine>(&line[start..]) {
87                if let Some(err) = obj.error {
88                    match err.r#type.as_str() {
89                        "message_rate_limit_error" => {
90                            let secs = err.resets_at.unwrap_or_default() - Local::now().timestamp();
91                            return Err(anyhow!(
92                                "Rate-limit: wait {} seconds – {}",
93                                secs,
94                                err.message
95                            ));
96                        }
97                        "overloaded_error" => {
98                            return Err(anyhow!("Server overloaded: {}", err.message));
99                        }
100                        _ => return Err(anyhow!("Claude error: {}", err.message)),
101                    }
102                }
103                completions.push_str(&obj.completion);
104            }
105        }
106    }
107    Ok(completions.trim().to_owned())
108}
109
110lazy_static! {
111    pub static ref READ_RE: Regex = Regex::new(r"(?i)^#?\s*read_file\s+(.+)").unwrap();
112    pub static ref EXEC_RE: Regex = Regex::new(r"(?i)^#?\s*exec\s+(.+)").unwrap();
113}
114
115/// Format a line with ANSI colors for read_file or exec commands
116pub fn format_tool_line(line: &str) -> String {
117    if READ_RE.is_match(line) {
118        format!("\x1b[34m{}\x1b[0m", line.trim_start_matches('#').trim())
119    } else if EXEC_RE.is_match(line) {
120        format!(
121            "\x1b[38;5;208m{}\x1b[0m",
122            line.trim_start_matches('#').trim()
123        )
124    } else {
125        line.to_owned()
126    }
127}
128
129/// Extract read_file and exec commands from a Claude answer
130pub fn extract_commands(answer: &str) -> (Vec<String>, Vec<String>) {
131    let mut reads = Vec::new();
132    let mut execs = Vec::new();
133    let mut lines = answer.lines().peekable();
134    while lines.peek().is_some() {
135        let line = lines.next().unwrap();
136        if let Some(caps) = READ_RE.captures(line) {
137            reads.extend(caps[1].split_whitespace().map(|s| s.to_string()));
138        } else if let Some(caps) = EXEC_RE.captures(line) {
139            let first = caps[1].to_string();
140            if let Some(hd) = first.split("<<").nth(1) {
141                let delim = hd.trim().trim_matches(|c| c == '\'' || c == '"');
142                let mut block = vec![first.clone()];
143                for l in lines.by_ref() {
144                    block.push(l.to_string());
145                    if l.trim() == delim {
146                        break;
147                    }
148                }
149                execs.push(block.join("\n"));
150            } else {
151                execs.push(first);
152            }
153        }
154    }
155    (reads, execs)
156}
157
158/// Apply ANSI formatting to an answer
159pub fn prettify(answer: &str) -> String {
160    answer
161        .lines()
162        .map(format_tool_line)
163        .collect::<Vec<_>>()
164        .join("\n")
165}
166
167/// Extract organization ID from Claude cookie string
168pub fn extract_org_id_from_cookie(cookie: &str) -> Option<String> {
169    let re = Regex::new(r"lastActiveOrg=([0-9a-f-]+)").ok()?;
170    re.captures(cookie)
171        .and_then(|caps| caps.get(1))
172        .map(|m| m.as_str().to_string())
173}
toast_api/utils.rs

toast_api/
utils.rs
