gemini_cli/rate_limits/
client.rs

1use std::path::{Path, PathBuf};
2use std::process::Command;
3
4use crate::json;
5
6pub struct UsageResponse {
7    pub body: String,
8}
9
10pub struct UsageRequest {
11    pub target_file: PathBuf,
12    pub refresh_on_401: bool,
13    pub endpoint: String,
14    pub api_version: String,
15    pub project: String,
16    pub connect_timeout_seconds: u64,
17    pub max_time_seconds: u64,
18}
19
20pub fn fetch_usage(request: &UsageRequest) -> Result<UsageResponse, String> {
21    let access_token = read_tokens(&request.target_file)?;
22    let mut response = send_request(request, &access_token)?;
23
24    if response.status == 401
25        && request.refresh_on_401
26        && let Ok(access_token) = read_tokens(&request.target_file)
27    {
28        response = send_request(request, &access_token)?;
29    }
30
31    if response.status != 200 {
32        let preview = response
33            .body
34            .chars()
35            .take(200)
36            .collect::<String>()
37            .replace(['\n', '\r'], " ");
38        if preview.is_empty() {
39            return Err(format!(
40                "gemini-rate-limits: POST {} failed (HTTP {})",
41                response.url, response.status
42            ));
43        }
44        return Err(format!(
45            "gemini-rate-limits: POST {} failed (HTTP {})\ngemini-rate-limits: body: {}",
46            response.url, response.status, preview
47        ));
48    }
49
50    Ok(UsageResponse {
51        body: response.body,
52    })
53}
54
55pub fn read_tokens(target_file: &Path) -> Result<String, String> {
56    let value = json::read_json(target_file).map_err(|err| err.to_string())?;
57    let access_token = json::string_at(&value, &["tokens", "access_token"])
58        .or_else(|| json::string_at(&value, &["access_token"]))
59        .ok_or_else(|| "missing access_token".to_string())?;
60    Ok(access_token)
61}
62
63fn send_request(request: &UsageRequest, access_token: &str) -> Result<HttpResponse, String> {
64    let url = build_usage_url(&request.endpoint, &request.api_version)?;
65    let payload = serde_json::json!({
66        "project": request.project,
67    })
68    .to_string();
69
70    let response = Command::new("curl")
71        .arg("-sS")
72        .arg("--connect-timeout")
73        .arg(request.connect_timeout_seconds.max(1).to_string())
74        .arg("--max-time")
75        .arg(request.max_time_seconds.max(1).to_string())
76        .arg("-X")
77        .arg("POST")
78        .arg("-H")
79        .arg("Accept: application/json")
80        .arg("-H")
81        .arg("Content-Type: application/json")
82        .arg("-H")
83        .arg("User-Agent: gemini-cli")
84        .arg("-H")
85        .arg(format!("Authorization: Bearer {access_token}"))
86        .arg("--data")
87        .arg(payload)
88        .arg(&url)
89        .arg("-w")
90        .arg("\n__HTTP_STATUS__:%{http_code}")
91        .output()
92        .map_err(|_| format!("gemini-rate-limits: request failed: {url}"))?;
93
94    if !response.status.success() {
95        return Err(format!("gemini-rate-limits: request failed: {url}"));
96    }
97
98    let response_text = String::from_utf8_lossy(&response.stdout).to_string();
99    let (body, status) = split_http_status_marker(&response_text);
100    if status == 0 {
101        return Err(format!("gemini-rate-limits: request failed: {url}"));
102    }
103
104    Ok(HttpResponse { status, body, url })
105}
106
107fn build_usage_url(endpoint: &str, api_version: &str) -> Result<String, String> {
108    let endpoint = endpoint.trim().trim_end_matches('/');
109    if endpoint.is_empty() || !(endpoint.starts_with("https://") || endpoint.starts_with("http://"))
110    {
111        return Err(format!(
112            "gemini-rate-limits: unsupported endpoint: {}",
113            endpoint
114        ));
115    }
116    let api_version = api_version.trim();
117    if api_version.is_empty() {
118        return Err("gemini-rate-limits: missing code assist api version".to_string());
119    }
120    Ok(format!("{endpoint}/{api_version}:retrieveUserQuota"))
121}
122
123fn split_http_status_marker(raw: &str) -> (String, u16) {
124    let marker = "__HTTP_STATUS__:";
125    if let Some(index) = raw.rfind(marker) {
126        let body = raw[..index]
127            .trim_end_matches('\n')
128            .trim_end_matches('\r')
129            .to_string();
130        let status_raw = raw[index + marker.len()..].trim();
131        let status = status_raw.parse::<u16>().unwrap_or(0);
132        (body, status)
133    } else {
134        (raw.to_string(), 0)
135    }
136}
137
138struct HttpResponse {
139    status: u16,
140    body: String,
141    url: String,
142}
gemini_cli/rate_limits/client.rs

gemini_cli/rate_limits/
client.rs
