| use serde::Serialize; |
| use serde_json; |
| use std::collections::HashMap; |
| use std::fs; |
| use std::path::PathBuf; |
| use uuid::Uuid; |
|
|
| use crate::models::{ |
| Account, AccountIndex, AccountSummary, DeviceProfile, DeviceProfileVersion, QuotaData, |
| TokenData, |
| }; |
| use crate::modules; |
| use once_cell::sync::Lazy; |
| use std::sync::Mutex; |
|
|
| #[cfg(test)] |
| mod tests { |
| use super::*; |
| use std::collections::HashSet; |
| use std::sync::Mutex as StdMutex; |
|
|
| |
| static TEST_MUTEX: Lazy<StdMutex<()>> = Lazy::new(|| StdMutex::new(())); |
|
|
| struct TestDataDir { |
| path: PathBuf, |
| } |
|
|
| impl TestDataDir { |
| fn new() -> Self { |
| let temp_path = std::env::temp_dir().join(format!( |
| "antigravity_test_{}_{}", |
| std::process::id(), |
| std::time::SystemTime::now() |
| .duration_since(std::time::UNIX_EPOCH) |
| .unwrap() |
| .as_millis() |
| )); |
| fs::create_dir_all(&temp_path).expect("Failed to create temp dir"); |
| |
| Self { |
| path: temp_path, |
| } |
| } |
|
|
| fn path(&self) -> &PathBuf { |
| &self.path |
| } |
| } |
|
|
| impl Drop for TestDataDir { |
| fn drop(&mut self) { |
| let _ = fs::remove_dir_all(&self.path); |
| } |
| } |
|
|
| |
| fn write_corrupted_index(path: &PathBuf, content: &[u8]) { |
| let index_path = path.join("accounts.json"); |
| fs::write(&index_path, content).expect("Failed to write corrupted index"); |
| } |
|
|
| |
| fn create_account_file(path: &PathBuf, account_id: &str, email: &str) { |
| let accounts_dir = path.join("accounts"); |
| fs::create_dir_all(&accounts_dir).expect("Failed to create accounts dir"); |
| |
| let account = Account::new( |
| account_id.to_string(), |
| email.to_string(), |
| TokenData::new( |
| "test_access_token".to_string(), |
| "test_refresh_token".to_string(), |
| 3600, |
| Some(email.to_string()), |
| None, |
| None, |
| true, |
| ), |
| ); |
| |
| let content = serde_json::to_string_pretty(&account).expect("Failed to serialize account"); |
| let account_path = accounts_dir.join(format!("{}.json", account_id)); |
| fs::write(&account_path, content).expect("Failed to write account file"); |
| } |
|
|
| #[test] |
| fn test_load_account_index_with_bom_prefix() { |
| let _guard = TEST_MUTEX.lock().unwrap(); |
| let dir = TestDataDir::new(); |
|
|
| |
| let bom = [0xEF, 0xBB, 0xBF]; |
| let json = r#"{"version":"2.0","accounts":[],"current_account_id":null}"#; |
| let mut content = Vec::new(); |
| content.extend_from_slice(&bom); |
| content.extend_from_slice(json.as_bytes()); |
| |
| write_corrupted_index(dir.path(), &content); |
|
|
| let result = load_account_index_in_dir(dir.path()); |
| |
| |
| assert!(result.is_ok(), "BOM should be stripped and JSON should parse: {:?}", result); |
| let index = result.unwrap(); |
| assert!(index.accounts.is_empty()); |
| println!("BOM case: successfully loaded index after sanitization"); |
| } |
|
|
| #[test] |
| fn test_load_account_index_with_nul_prefix() { |
| let _guard = TEST_MUTEX.lock().unwrap(); |
| let dir = TestDataDir::new(); |
|
|
| |
| let nul = [0x00]; |
| let json = r#"{"version":"2.0","accounts":[],"current_account_id":null}"#; |
| let mut content = Vec::new(); |
| content.extend_from_slice(&nul); |
| content.extend_from_slice(json.as_bytes()); |
| |
| write_corrupted_index(dir.path(), &content); |
|
|
| let result = load_account_index_in_dir(dir.path()); |
| |
| |
| assert!(result.is_ok(), "NUL prefix should be stripped and JSON should parse: {:?}", result); |
| let index = result.unwrap(); |
| assert!(index.accounts.is_empty()); |
| println!("NUL prefix case: successfully loaded index after sanitization"); |
| } |
|
|
| #[test] |
| fn test_load_account_index_with_garbage_content() { |
| let _guard = TEST_MUTEX.lock().unwrap(); |
| let dir = TestDataDir::new(); |
|
|
| |
| write_corrupted_index(dir.path(), b"\0\0not json"); |
|
|
| let result = load_account_index_in_dir(dir.path()); |
| |
| |
| assert!(result.is_ok(), "Garbage content should trigger recovery and return Ok: {:?}", result); |
| let index = result.unwrap(); |
| assert!(index.accounts.is_empty(), "Recovered index should be empty when no account files exist"); |
| println!("Garbage content case: successfully recovered to empty index"); |
| } |
|
|
| #[test] |
| fn test_load_account_index_with_empty_file() { |
| let _guard = TEST_MUTEX.lock().unwrap(); |
| let dir = TestDataDir::new(); |
|
|
| |
| write_corrupted_index(dir.path(), b""); |
|
|
| let result = load_account_index_in_dir(dir.path()); |
| |
| |
| assert!(result.is_ok()); |
| let index = result.unwrap(); |
| assert!(index.accounts.is_empty()); |
| } |
|
|
| #[test] |
| fn test_load_account_index_with_whitespace_only() { |
| let _guard = TEST_MUTEX.lock().unwrap(); |
| let dir = TestDataDir::new(); |
|
|
| |
| write_corrupted_index(dir.path(), b" \n\t "); |
|
|
| let result = load_account_index_in_dir(dir.path()); |
| |
| |
| assert!(result.is_ok()); |
| let index = result.unwrap(); |
| assert!(index.accounts.is_empty()); |
| } |
|
|
| #[test] |
| fn test_missing_index_with_existing_accounts() { |
| let _guard = TEST_MUTEX.lock().unwrap(); |
| let dir = TestDataDir::new(); |
|
|
| |
| create_account_file(dir.path(), "test-id-1", "user1@example.com"); |
| create_account_file(dir.path(), "test-id-2", "user2@example.com"); |
|
|
| |
| let index_path = dir.path().join("accounts.json"); |
| assert!(!index_path.exists()); |
|
|
| |
| let result = load_account_index_in_dir(dir.path()); |
| assert!(result.is_ok(), "Should recover from accounts directory"); |
| let index = result.unwrap(); |
| assert_eq!(index.accounts.len(), 2, "Index should have 2 accounts recovered from accounts directory"); |
| |
| |
| let emails: Vec<_> = index.accounts.iter().map(|s| s.email.clone()).collect(); |
| assert!(emails.contains(&"user1@example.com".to_string())); |
| assert!(emails.contains(&"user2@example.com".to_string())); |
|
|
| |
| let accounts_dir = dir.path().join("accounts"); |
| let account_files: Vec<_> = fs::read_dir(&accounts_dir) |
| .unwrap() |
| .filter_map(|e| e.ok()) |
| .filter(|e| e.path().extension().map_or(false, |ext| ext == "json")) |
| .collect(); |
| assert_eq!(account_files.len(), 2, "Account files should still exist on disk"); |
| |
| println!("Missing index with existing accounts: successfully recovered {} accounts", index.accounts.len()); |
| } |
|
|
| #[test] |
| fn test_save_account_index_roundtrip() { |
| let _guard = TEST_MUTEX.lock().unwrap(); |
| let dir = TestDataDir::new(); |
|
|
| |
| let now = chrono::Utc::now().timestamp(); |
| let index = AccountIndex { |
| version: "2.0".to_string(), |
| accounts: vec![ |
| AccountSummary { |
| id: "acc-1".to_string(), |
| email: "user1@example.com".to_string(), |
| name: Some("User One".to_string()), |
| disabled: false, |
| proxy_disabled: false, |
| protected_models: HashSet::new(), |
| created_at: now, |
| last_used: now, |
| }, |
| AccountSummary { |
| id: "acc-2".to_string(), |
| email: "user2@example.com".to_string(), |
| name: None, |
| disabled: true, |
| proxy_disabled: true, |
| protected_models: HashSet::new(), |
| created_at: now - 100, |
| last_used: now - 50, |
| }, |
| ], |
| current_account_id: Some("acc-1".to_string()), |
| }; |
|
|
| |
| save_account_index_in_dir(dir.path(), &index).expect("Failed to save account index"); |
|
|
| |
| let loaded = load_account_index_in_dir(dir.path()).expect("Failed to load account index"); |
|
|
| |
| assert_eq!(loaded.accounts.len(), 2, "Should have 2 accounts"); |
| assert_eq!(loaded.current_account_id, Some("acc-1".to_string()), "current_account_id should match"); |
| |
| |
| let acc1 = loaded.accounts.iter().find(|a| a.id == "acc-1").expect("acc-1 should exist"); |
| assert_eq!(acc1.email, "user1@example.com"); |
| assert_eq!(acc1.name, Some("User One".to_string())); |
| assert!(!acc1.disabled); |
| assert!(!acc1.proxy_disabled); |
| |
| |
| let acc2 = loaded.accounts.iter().find(|a| a.id == "acc-2").expect("acc-2 should exist"); |
| assert_eq!(acc2.email, "user2@example.com"); |
| assert_eq!(acc2.name, None); |
| assert!(acc2.disabled); |
| assert!(acc2.proxy_disabled); |
|
|
| println!("save_account_index roundtrip: successfully saved and loaded index with {} accounts", loaded.accounts.len()); |
| } |
|
|
| #[test] |
| fn test_backup_created_on_parse_failure() { |
| let _guard = TEST_MUTEX.lock().unwrap(); |
| let dir = TestDataDir::new(); |
|
|
| |
| create_account_file(dir.path(), "recovered-acc", "recovered@example.com"); |
|
|
| |
| let garbage_content = b"this is not valid json { broken"; |
| write_corrupted_index(dir.path(), garbage_content); |
|
|
| |
| let index_path = dir.path().join("accounts.json"); |
| assert!(index_path.exists(), "accounts.json should exist"); |
|
|
| |
| let recovered = load_account_index_in_dir(dir.path()).expect("Should recover from accounts"); |
| assert_eq!(recovered.accounts.len(), 1, "Should recover 1 account"); |
| assert_eq!(recovered.accounts[0].email, "recovered@example.com"); |
| assert_eq!(recovered.current_account_id, Some("recovered-acc".to_string())); |
|
|
| |
| let data_dir = dir.path(); |
| let backup_files: Vec<_> = fs::read_dir(data_dir) |
| .unwrap() |
| .filter_map(|e| e.ok()) |
| .filter(|e| { |
| e.file_name() |
| .to_str() |
| .map_or(false, |name| name.starts_with("accounts.json.corrupt-")) |
| }) |
| .collect(); |
| |
| assert_eq!(backup_files.len(), 1, "Should have exactly one backup file"); |
| |
| |
| let backup_content = fs::read(&backup_files[0].path()).expect("Should be able to read backup file"); |
| assert_eq!(backup_content, garbage_content, "Backup should contain original corrupt content"); |
|
|
| println!("Backup creation on parse failure: successfully created backup"); |
| } |
|
|
| } |
|
|
| |
| static ACCOUNT_INDEX_LOCK: Lazy<Mutex<()>> = Lazy::new(|| Mutex::new(())); |
|
|
| |
| const DATA_DIR: &str = ".antigravity_tools"; |
| const ACCOUNTS_INDEX: &str = "accounts.json"; |
| const ACCOUNTS_DIR: &str = "accounts"; |
|
|
| |
| pub fn get_data_dir() -> Result<PathBuf, String> { |
| |
| if let Ok(env_path) = std::env::var("ABV_DATA_DIR") { |
| if !env_path.trim().is_empty() { |
| let data_dir = PathBuf::from(env_path); |
| if !data_dir.exists() { |
| fs::create_dir_all(&data_dir).map_err(|e| format!("failed_to_create_custom_data_dir: {}", e))?; |
| } |
| return Ok(data_dir); |
| } |
| } |
|
|
| let home = dirs::home_dir().ok_or("failed_to_get_home_dir")?; |
| let data_dir = home.join(DATA_DIR); |
|
|
| |
| if !data_dir.exists() { |
| fs::create_dir_all(&data_dir).map_err(|e| format!("failed_to_create_data_dir: {}", e))?; |
| } |
|
|
| Ok(data_dir) |
| } |
|
|
| |
| pub fn get_accounts_dir() -> Result<PathBuf, String> { |
| let data_dir = get_data_dir()?; |
| let accounts_dir = data_dir.join(ACCOUNTS_DIR); |
|
|
| if !accounts_dir.exists() { |
| fs::create_dir_all(&accounts_dir) |
| .map_err(|e| format!("failed_to_create_accounts_dir: {}", e))?; |
| } |
|
|
| Ok(accounts_dir) |
| } |
|
|
| |
| fn load_account_index_in_dir(data_dir: &PathBuf) -> Result<AccountIndex, String> { |
| let index_path = data_dir.join(ACCOUNTS_INDEX); |
|
|
| if !index_path.exists() { |
| crate::modules::logger::log_warn( |
| "Account index file not found, attempting recovery from accounts directory", |
| ); |
| let recovered = rebuild_index_from_accounts_in_dir(data_dir)?; |
| try_save_recovered_index(data_dir, &index_path, &recovered, None)?; |
| return Ok(recovered); |
| } |
|
|
| let raw_content = fs::read(&index_path) |
| .map_err(|e| format!("failed_to_read_account_index: {}", e))?; |
|
|
| |
| if raw_content.is_empty() { |
| crate::modules::logger::log_warn( |
| "Account index is empty, attempting recovery from accounts directory", |
| ); |
| let recovered = rebuild_index_from_accounts_in_dir(data_dir)?; |
| try_save_recovered_index(data_dir, &index_path, &recovered, None)?; |
| return Ok(recovered); |
| } |
|
|
| |
| let sanitized = sanitize_index_content(&raw_content); |
|
|
| |
| if sanitized.trim().is_empty() { |
| crate::modules::logger::log_warn( |
| "Account index is empty after sanitization, attempting recovery from accounts directory", |
| ); |
| let recovered = rebuild_index_from_accounts_in_dir(data_dir)?; |
| try_save_recovered_index(data_dir, &index_path, &recovered, None)?; |
| return Ok(recovered); |
| } |
|
|
| |
| match serde_json::from_str::<AccountIndex>(&sanitized) { |
| Ok(index) => { |
| crate::modules::logger::log_info(&format!( |
| "Successfully loaded index with {} accounts", |
| index.accounts.len() |
| )); |
| Ok(index) |
| } |
| Err(parse_err) => { |
| crate::modules::logger::log_error(&format!( |
| "Failed to parse account index: {}. Attempting recovery from accounts directory", |
| parse_err |
| )); |
| let recovered = rebuild_index_from_accounts_in_dir(data_dir)?; |
| try_save_recovered_index(data_dir, &index_path, &recovered, Some(&raw_content))?; |
| Ok(recovered) |
| } |
| } |
| } |
|
|
| |
| fn save_account_index_in_dir(data_dir: &PathBuf, index: &AccountIndex) -> Result<(), String> { |
| let index_path = data_dir.join(ACCOUNTS_INDEX); |
| |
| let temp_filename = format!("{}.tmp.{}", ACCOUNTS_INDEX, Uuid::new_v4()); |
| let temp_path = data_dir.join(&temp_filename); |
|
|
| let content = serde_json::to_string_pretty(index) |
| .map_err(|e| format!("failed_to_serialize_account_index: {}", e))?; |
|
|
| |
| if let Err(e) = fs::write(&temp_path, content) { |
| |
| let _ = fs::remove_file(&temp_path); |
| return Err(format!("failed_to_write_temp_index_file: {}", e)); |
| } |
|
|
| |
| if let Err(e) = atomic_replace_file(&temp_path, &index_path) { |
| |
| let _ = fs::remove_file(&temp_path); |
| return Err(format!("failed_to_replace_index_file: {}", e)); |
| } |
|
|
| Ok(()) |
| } |
|
|
| |
| fn rebuild_index_from_accounts_in_dir(data_dir: &PathBuf) -> Result<AccountIndex, String> { |
| let accounts_dir = data_dir.join(ACCOUNTS_DIR); |
| let mut summaries = Vec::new(); |
|
|
| if accounts_dir.exists() { |
| if let Ok(entries) = fs::read_dir(&accounts_dir) { |
| for entry in entries.filter_map(|e| e.ok()) { |
| let path = entry.path(); |
| if path.extension().map_or(false, |ext| ext == "json") { |
| if let Some(account_id) = path.file_stem().and_then(|s| s.to_str()) { |
| match load_account_at_path(&path) { |
| Ok(account) => { |
| summaries.push(AccountSummary { |
| id: account.id, |
| email: account.email, |
| name: account.name, |
| disabled: account.disabled, |
| proxy_disabled: account.proxy_disabled, |
| protected_models: account.protected_models, |
| created_at: account.created_at, |
| last_used: account.last_used, |
| }); |
| } |
| Err(e) => { |
| crate::modules::logger::log_warn(&format!( |
| "Failed to load account {} during recovery: {}", |
| account_id, e |
| )); |
| } |
| } |
| } |
| } |
| } |
| } |
| } |
|
|
| |
| summaries.sort_by(|a, b| { |
| b.last_used |
| .cmp(&a.last_used) |
| .then_with(|| a.email.cmp(&b.email)) |
| }); |
|
|
| let current_account_id = summaries.first().map(|s| s.id.clone()); |
|
|
| crate::modules::logger::log_info(&format!( |
| "Rebuilt index from accounts directory: {} accounts recovered", |
| summaries.len() |
| )); |
|
|
| Ok(AccountIndex { |
| version: "2.0".to_string(), |
| accounts: summaries, |
| current_account_id, |
| }) |
| } |
|
|
| |
| fn load_account_at_path(account_path: &PathBuf) -> Result<Account, String> { |
| let content = fs::read_to_string(account_path) |
| .map_err(|e| format!("failed_to_read_account_data: {}", e))?; |
| serde_json::from_str(&content).map_err(|e| format!("failed_to_parse_account_data: {}", e)) |
| } |
|
|
| |
| pub fn load_account_index() -> Result<AccountIndex, String> { |
| let data_dir = get_data_dir()?; |
| load_account_index_in_dir(&data_dir) |
| } |
|
|
| |
| fn sanitize_index_content(raw: &[u8]) -> String { |
| |
| let without_bom = if raw.starts_with(&[0xEF, 0xBB, 0xBF]) { |
| &raw[3..] |
| } else { |
| raw |
| }; |
|
|
| |
| let without_nul = without_bom |
| .iter() |
| .skip_while(|&&b| b == 0x00) |
| .copied() |
| .collect::<Vec<u8>>(); |
|
|
| |
| String::from_utf8_lossy(&without_nul).into_owned() |
| } |
|
|
| |
| fn try_save_recovered_index( |
| data_dir: &PathBuf, |
| _index_path: &PathBuf, |
| index: &AccountIndex, |
| corrupt_content: Option<&[u8]>, |
| ) -> Result<(), String> { |
| |
| if let Some(content) = corrupt_content { |
| let timestamp = chrono::Utc::now().timestamp(); |
| let backup_name = format!("accounts.json.corrupt-{}-{}", timestamp, Uuid::new_v4()); |
| let backup_path = data_dir.join(&backup_name); |
| if let Err(e) = fs::write(&backup_path, content) { |
| crate::modules::logger::log_warn(&format!( |
| "Failed to backup corrupt index to {}: {}", |
| backup_name, e |
| )); |
| } else { |
| crate::modules::logger::log_info(&format!( |
| "Backed up corrupt index to {}", |
| backup_name |
| )); |
| } |
| } |
|
|
| |
| match ACCOUNT_INDEX_LOCK.try_lock() { |
| Ok(_guard) => { |
| if let Err(e) = save_account_index_in_dir(data_dir, index) { |
| crate::modules::logger::log_warn(&format!( |
| "Failed to save recovered index: {}. Will retry on next load.", |
| e |
| )); |
| } else { |
| crate::modules::logger::log_info("Successfully saved recovered index"); |
| } |
| } |
| Err(_) => { |
| crate::modules::logger::log_warn( |
| "Could not acquire lock to save recovered index. Will retry on next load." |
| ); |
| } |
| } |
|
|
| Ok(()) |
| } |
|
|
| |
| pub fn save_account_index(index: &AccountIndex) -> Result<(), String> { |
| let data_dir = get_data_dir()?; |
| save_account_index_in_dir(&data_dir, index) |
| } |
|
|
| |
| #[cfg(target_os = "windows")] |
| fn atomic_replace_file(src: &PathBuf, dst: &PathBuf) -> Result<(), String> { |
| use std::os::windows::ffi::OsStrExt; |
|
|
| type Bool = i32; |
| type Dword = u32; |
|
|
| #[link(name = "Kernel32")] |
| extern "system" { |
| fn MoveFileExW(lp_existing_file_name: *const u16, lp_new_file_name: *const u16, dw_flags: Dword) -> Bool; |
| } |
|
|
| let src_wide: Vec<u16> = src |
| .as_os_str() |
| .encode_wide() |
| .chain(std::iter::once(0)) |
| .collect(); |
| let dst_wide: Vec<u16> = dst |
| .as_os_str() |
| .encode_wide() |
| .chain(std::iter::once(0)) |
| .collect(); |
|
|
| |
| |
| const MOVEFILE_REPLACE_EXISTING: u32 = 0x1; |
| const MOVEFILE_WRITE_THROUGH: u32 = 0x8; |
| let flags = MOVEFILE_REPLACE_EXISTING | MOVEFILE_WRITE_THROUGH; |
|
|
| let result = unsafe { MoveFileExW(src_wide.as_ptr(), dst_wide.as_ptr(), flags) }; |
| if result == 0 { |
| let err = std::io::Error::last_os_error(); |
| |
| let _ = fs::remove_file(src); |
| return Err(format!("MoveFileExW failed: {}", err)); |
| } |
|
|
| Ok(()) |
| } |
|
|
| |
| #[cfg(not(target_os = "windows"))] |
| fn atomic_replace_file(src: &PathBuf, dst: &PathBuf) -> Result<(), String> { |
| fs::rename(src, dst).map_err(|e| format!("rename failed: {}", e)) |
| } |
|
|
| |
| pub fn load_account(account_id: &str) -> Result<Account, String> { |
| let accounts_dir = get_accounts_dir()?; |
| let account_path = accounts_dir.join(format!("{}.json", account_id)); |
| load_account_at_path(&account_path) |
| } |
|
|
| |
| pub fn save_account(account: &Account) -> Result<(), String> { |
| let accounts_dir = get_accounts_dir()?; |
| let account_path = accounts_dir.join(format!("{}.json", account.id)); |
|
|
| let temp_filename = format!("{}.tmp.{}", account.id, Uuid::new_v4()); |
| let temp_path = accounts_dir.join(&temp_filename); |
|
|
| let content = serde_json::to_string_pretty(account) |
| .map_err(|e| format!("failed_to_serialize_account_data: {}", e))?; |
|
|
| if let Err(e) = std::fs::write(&temp_path, content) { |
| let _ = std::fs::remove_file(&temp_path); |
| return Err(format!("failed_to_write_temp_account_file: {}", e)); |
| } |
|
|
| if let Err(e) = atomic_replace_file(&temp_path, &account_path) { |
| let _ = std::fs::remove_file(&temp_path); |
| return Err(format!("failed_to_replace_account_file: {}", e)); |
| } |
|
|
| Ok(()) |
| } |
|
|
| |
| pub fn list_accounts() -> Result<Vec<Account>, String> { |
| crate::modules::logger::log_info("Listing accounts..."); |
| let index = load_account_index()?; |
| let mut accounts = Vec::new(); |
|
|
| for summary in &index.accounts { |
| match load_account(&summary.id) { |
| Ok(account) => accounts.push(account), |
| Err(e) => { |
| crate::modules::logger::log_error(&format!( |
| "Failed to load account {}: {}", |
| summary.id, e |
| )); |
| |
| |
| |
| } |
| } |
| } |
|
|
| Ok(accounts) |
| } |
|
|
| |
| pub fn add_account( |
| email: String, |
| name: Option<String>, |
| token: TokenData, |
| ) -> Result<Account, String> { |
| let _lock = ACCOUNT_INDEX_LOCK |
| .lock() |
| .map_err(|e| format!("failed_to_acquire_lock: {}", e))?; |
| let mut index = load_account_index()?; |
|
|
| |
| if index.accounts.iter().any(|s| s.email == email) { |
| return Err(format!("Account already exists: {}", email)); |
| } |
|
|
| |
| let account_id = Uuid::new_v4().to_string(); |
| let mut account = Account::new(account_id.clone(), email.clone(), token); |
| account.name = name.clone(); |
|
|
| |
| save_account(&account)?; |
|
|
| |
| index.accounts.push(AccountSummary { |
| id: account.id.clone(), |
| email: account.email.clone(), |
| name: account.name.clone(), |
| disabled: account.disabled, |
| proxy_disabled: account.proxy_disabled, |
| protected_models: account.protected_models.clone(), |
| created_at: account.created_at, |
| last_used: account.last_used, |
| }); |
|
|
| |
| if index.current_account_id.is_none() { |
| index.current_account_id = Some(account_id); |
| } |
|
|
| save_account_index(&index)?; |
|
|
| Ok(account) |
| } |
|
|
| |
| pub fn upsert_account( |
| email: String, |
| name: Option<String>, |
| token: TokenData, |
| ) -> Result<Account, String> { |
| let _lock = ACCOUNT_INDEX_LOCK |
| .lock() |
| .map_err(|e| format!("failed_to_acquire_lock: {}", e))?; |
| let mut index = load_account_index()?; |
|
|
| |
| let existing_account_id = index |
| .accounts |
| .iter() |
| .find(|s| s.email == email) |
| .map(|s| s.id.clone()); |
|
|
| if let Some(account_id) = existing_account_id { |
| |
| match load_account(&account_id) { |
| Ok(mut account) => { |
| let old_access_token = account.token.access_token.clone(); |
| let old_refresh_token = account.token.refresh_token.clone(); |
| account.token = token; |
| account.name = name.clone(); |
| |
| |
| if account.disabled |
| && (account.token.refresh_token != old_refresh_token |
| || account.token.access_token != old_access_token) |
| { |
| account.disabled = false; |
| account.disabled_reason = None; |
| account.disabled_at = None; |
| } |
| account.update_last_used(); |
| save_account(&account)?; |
|
|
| |
| if let Some(idx_summary) = index.accounts.iter_mut().find(|s| s.id == account_id) { |
| idx_summary.name = name; |
| save_account_index(&index)?; |
| } |
|
|
| return Ok(account); |
| } |
| Err(e) => { |
| crate::modules::logger::log_warn(&format!( |
| "Account {} file missing ({}), recreating...", |
| account_id, e |
| )); |
| |
| let mut account = Account::new(account_id.clone(), email.clone(), token); |
| account.name = name.clone(); |
| save_account(&account)?; |
|
|
| |
| if let Some(idx_summary) = index.accounts.iter_mut().find(|s| s.id == account_id) { |
| idx_summary.name = name; |
| save_account_index(&index)?; |
| } |
|
|
| return Ok(account); |
| } |
| } |
| } |
|
|
| |
| |
| |
|
|
| |
| drop(_lock); |
| add_account(email, name, token) |
| } |
|
|
| |
| pub fn delete_account(account_id: &str) -> Result<(), String> { |
| let _lock = ACCOUNT_INDEX_LOCK |
| .lock() |
| .map_err(|e| format!("failed_to_acquire_lock: {}", e))?; |
| let mut index = load_account_index()?; |
|
|
| |
| let original_len = index.accounts.len(); |
| index.accounts.retain(|s| s.id != account_id); |
|
|
| if index.accounts.len() == original_len { |
| return Err(format!("Account ID not found: {}", account_id)); |
| } |
|
|
| |
| if index.current_account_id.as_deref() == Some(account_id) { |
| index.current_account_id = index.accounts.first().map(|s| s.id.clone()); |
| } |
|
|
| save_account_index(&index)?; |
|
|
| |
| let accounts_dir = get_accounts_dir()?; |
| let account_path = accounts_dir.join(format!("{}.json", account_id)); |
|
|
| if account_path.exists() { |
| fs::remove_file(&account_path) |
| .map_err(|e| format!("failed_to_delete_account_file: {}", e))?; |
| } |
|
|
| |
| crate::proxy::server::trigger_account_delete(account_id); |
|
|
| Ok(()) |
| } |
|
|
| |
| pub fn delete_accounts(account_ids: &[String]) -> Result<(), String> { |
| let _lock = ACCOUNT_INDEX_LOCK |
| .lock() |
| .map_err(|e| format!("failed_to_acquire_lock: {}", e))?; |
| let mut index = load_account_index()?; |
|
|
| let accounts_dir = get_accounts_dir()?; |
|
|
| for account_id in account_ids { |
| |
| index.accounts.retain(|s| &s.id != account_id); |
|
|
| |
| if index.current_account_id.as_deref() == Some(account_id) { |
| index.current_account_id = None; |
| } |
|
|
| |
| let account_path = accounts_dir.join(format!("{}.json", account_id)); |
| if account_path.exists() { |
| let _ = fs::remove_file(&account_path); |
| } |
|
|
| |
| crate::proxy::server::trigger_account_delete(account_id); |
| } |
|
|
| |
| if index.current_account_id.is_none() { |
| index.current_account_id = index.accounts.first().map(|s| s.id.clone()); |
| } |
|
|
| save_account_index(&index) |
| } |
|
|
| |
| |
| pub fn reorder_accounts(account_ids: &[String]) -> Result<(), String> { |
| let _lock = ACCOUNT_INDEX_LOCK |
| .lock() |
| .map_err(|e| format!("failed_to_acquire_lock: {}", e))?; |
| let mut index = load_account_index()?; |
|
|
| |
| let id_to_summary: std::collections::HashMap<_, _> = index |
| .accounts |
| .iter() |
| .map(|s| (s.id.clone(), s.clone())) |
| .collect(); |
|
|
| |
| let mut new_accounts = Vec::new(); |
| for id in account_ids { |
| if let Some(summary) = id_to_summary.get(id) { |
| new_accounts.push(summary.clone()); |
| } |
| } |
|
|
| |
| for summary in &index.accounts { |
| if !account_ids.contains(&summary.id) { |
| new_accounts.push(summary.clone()); |
| } |
| } |
|
|
| index.accounts = new_accounts; |
|
|
| crate::modules::logger::log_info(&format!( |
| "Account order updated, {} accounts total", |
| index.accounts.len() |
| )); |
|
|
| save_account_index(&index) |
| } |
|
|
| |
| pub async fn switch_account( |
| account_id: &str, |
| integration: &(impl modules::integration::SystemIntegration + ?Sized), |
| ) -> Result<(), String> { |
| use crate::modules::oauth; |
|
|
| let index = { |
| let _lock = ACCOUNT_INDEX_LOCK |
| .lock() |
| .map_err(|e| format!("failed_to_acquire_lock: {}", e))?; |
| load_account_index()? |
| }; |
|
|
| |
| if !index.accounts.iter().any(|s| s.id == account_id) { |
| return Err(format!("Account not found: {}", account_id)); |
| } |
|
|
| let mut account = load_account(account_id)?; |
| crate::modules::logger::log_info(&format!( |
| "Switching to account: {} (ID: {})", |
| account.email, account.id |
| )); |
|
|
| |
| let fresh_token = match oauth::ensure_fresh_token(&account.token, Some(&account.id)).await { |
| Ok(token) => token, |
| Err(e) => { |
| if is_account_access_blocked_message(&e) { |
| mark_validation_blocked(&mut account, &e); |
| } |
| return Err(format_switch_refresh_error(&e)); |
| } |
| }; |
|
|
| |
| if fresh_token.access_token != account.token.access_token { |
| account.token = fresh_token.clone(); |
| save_account(&account)?; |
| } |
|
|
| ensure_enterprise_project_ready(&mut account).await?; |
|
|
| |
| if account.device_profile.is_none() { |
| crate::modules::logger::log_info(&format!( |
| "Account {} has no bound fingerprint, generating new one for isolation...", |
| account.email |
| )); |
| let new_profile = modules::device::generate_profile(); |
| apply_profile_to_account( |
| &mut account, |
| new_profile.clone(), |
| Some("auto_generated".to_string()), |
| true, |
| )?; |
| } |
|
|
| |
| integration.on_account_switch(&account).await?; |
|
|
| |
| { |
| let _lock = ACCOUNT_INDEX_LOCK |
| .lock() |
| .map_err(|e| format!("failed_to_acquire_lock: {}", e))?; |
| let mut index = load_account_index()?; |
| index.current_account_id = Some(account_id.to_string()); |
| save_account_index(&index)?; |
| } |
|
|
| account.update_last_used(); |
| save_account(&account)?; |
|
|
| crate::modules::logger::log_info(&format!( |
| "Account switch core logic completed: {}", |
| account.email |
| )); |
|
|
| Ok(()) |
| } |
|
|
| fn is_enterprise_client(client_key: Option<&str>) -> bool { |
| client_key |
| .map(str::trim) |
| .filter(|key| !key.is_empty()) |
| .map(|key| key.eq_ignore_ascii_case("antigravity_enterprise")) |
| .unwrap_or(false) |
| } |
|
|
| fn normalize_project_id(project_id: Option<&str>) -> Option<String> { |
| project_id |
| .map(str::trim) |
| .filter(|pid| !pid.is_empty()) |
| .map(ToOwned::to_owned) |
| } |
|
|
| async fn ensure_enterprise_project_ready(account: &mut Account) -> Result<(), String> { |
| if !is_enterprise_client(account.token.oauth_client_key.as_deref()) { |
| return Ok(()); |
| } |
|
|
| if normalize_project_id(account.token.project_id.as_deref()).is_some() { |
| return Ok(()); |
| } |
|
|
| crate::modules::logger::log_warn(&format!( |
| "Account {} is using enterprise OAuth client but missing project_id. Trying to resolve before switch...", |
| account.email |
| )); |
|
|
| match crate::proxy::project_resolver::fetch_project_id(&account.token.access_token).await { |
| Ok(project_id) => { |
| crate::modules::logger::log_info(&format!( |
| "Resolved enterprise project_id for {}: {}", |
| account.email, project_id |
| )); |
| account.token.project_id = Some(project_id); |
| save_account(account)?; |
| Ok(()) |
| } |
| Err(e) => { |
| crate::modules::logger::log_warn(&format!( |
| "Account {} is currently missing enterprise project_id and auto-resolve failed ({}). Allowing switch to proceed, but certain enterprise features may be limited.", |
| account.email, e |
| )); |
| Ok(()) |
| } |
| } |
| } |
|
|
| fn is_rate_limit_error(err: &crate::error::AppError) -> bool { |
| match err { |
| crate::error::AppError::Network(_, Some(status)) => *status == 429, |
| crate::error::AppError::Unknown(msg) |
| | crate::error::AppError::OAuth(msg) |
| | crate::error::AppError::Account(msg) |
| | crate::error::AppError::Config(msg) => { |
| let lower = msg.to_lowercase(); |
| lower.contains("429") |
| || lower.contains("too many requests") |
| || lower.contains("resource_exhausted") |
| || lower.contains("resource has been exhausted") |
| } |
| _ => false, |
| } |
| } |
|
|
| fn recover_cached_quota_on_rate_limit( |
| account: &Account, |
| err: &crate::error::AppError, |
| ) -> Option<QuotaData> { |
| if !is_rate_limit_error(err) { |
| return None; |
| } |
|
|
| let cached = account.quota.clone()?; |
| if cached.models.is_empty() { |
| return None; |
| } |
|
|
| Some(cached) |
| } |
|
|
| fn is_validation_required_error(err: &crate::error::AppError) -> bool { |
| let text = err.to_string().to_lowercase(); |
| text.contains("verify your account") |
| || text.contains("further action is required") |
| || text.contains("validation_url") |
| || text.contains("appeal_url") |
| || text.contains("validation required") |
| } |
|
|
| fn is_account_access_blocked_message(message: &str) -> bool { |
| let text = message.to_lowercase(); |
| text.contains("verify your account") |
| || text.contains("further action is required") |
| || text.contains("validation_url") |
| || text.contains("appeal_url") |
| || text.contains("validation required") |
| || text.contains("unauthorized_client") |
| || text.contains("invalid_client") |
| || text.contains("invalid_grant") |
| || text.contains("resource_exhausted") |
| || text.contains("resource has been exhausted") |
| } |
|
|
| fn format_switch_refresh_error(message: &str) -> String { |
| let lower = message.to_lowercase(); |
|
|
| if lower.contains("unauthorized_client") |
| || lower.contains("invalid_client") |
| || lower.contains("invalid_grant") |
| { |
| return format!( |
| "Token refresh failed: OAuth client is not authorized for this account. Please sign in again in Antigravity-Manager and complete authorization/verification. Raw error: {}", |
| message |
| ); |
| } |
|
|
| if lower.contains("verify your account") |
| || lower.contains("further action is required") |
| || lower.contains("validation_url") |
| || lower.contains("appeal_url") |
| || lower.contains("validation required") |
| { |
| return format!( |
| "Token refresh failed: account requires additional verification. Please finish verification in Antigravity, then retry account switch. Raw error: {}", |
| message |
| ); |
| } |
|
|
| if lower.contains("resource_exhausted") || lower.contains("resource has been exhausted") { |
| return format!( |
| "Token refresh failed: account is rate-limited or temporarily restricted (RESOURCE_EXHAUSTED). Please retry later. Raw error: {}", |
| message |
| ); |
| } |
|
|
| format!("Token refresh failed: {}", message) |
| } |
|
|
| fn format_rate_limit_block_reason(err: &crate::error::AppError) -> String { |
| format!( |
| "Account is temporarily rate-limited or risk-controlled (RESOURCE_EXHAUSTED). Please cool down and retry later. Raw error: {}", |
| err |
| ) |
| } |
|
|
| fn mark_validation_blocked(account: &mut Account, reason: &str) { |
| if account.validation_blocked |
| && account.validation_blocked_reason.as_deref() == Some(reason) |
| { |
| return; |
| } |
|
|
| account.validation_blocked = true; |
| account.validation_blocked_reason = Some(reason.to_string()); |
| if let Err(e) = save_account(account) { |
| crate::modules::logger::log_warn(&format!( |
| "Failed to persist validation_blocked state for {}: {}", |
| account.email, e |
| )); |
| } |
| } |
|
|
| fn clear_validation_blocked(account: &mut Account) { |
| if !account.validation_blocked { |
| return; |
| } |
|
|
| account.validation_blocked = false; |
| account.validation_blocked_until = None; |
| account.validation_blocked_reason = None; |
| account.validation_url = None; |
| if let Err(e) = save_account(account) { |
| crate::modules::logger::log_warn(&format!( |
| "Failed to clear validation_blocked state for {}: {}", |
| account.email, e |
| )); |
| } |
| } |
|
|
| |
| #[derive(Debug, Serialize)] |
| pub struct DeviceProfiles { |
| pub current_storage: Option<DeviceProfile>, |
| pub bound_profile: Option<DeviceProfile>, |
| pub history: Vec<DeviceProfileVersion>, |
| pub baseline: Option<DeviceProfile>, |
| } |
|
|
| pub fn get_device_profiles(account_id: &str) -> Result<DeviceProfiles, String> { |
| |
| let current = crate::modules::device::get_storage_path() |
| .ok() |
| .and_then(|path| crate::modules::device::read_profile(&path).ok()); |
| let account = load_account(account_id)?; |
| Ok(DeviceProfiles { |
| current_storage: current, |
| bound_profile: account.device_profile.clone(), |
| history: account.device_history.clone(), |
| baseline: crate::modules::device::load_global_original(), |
| }) |
| } |
|
|
| |
| pub fn bind_device_profile(account_id: &str, mode: &str) -> Result<DeviceProfile, String> { |
| use crate::modules::device; |
|
|
| let profile = match mode { |
| "capture" => device::read_profile(&device::get_storage_path()?)?, |
| "generate" => device::generate_profile(), |
| _ => return Err("mode must be 'capture' or 'generate'".to_string()), |
| }; |
|
|
| let mut account = load_account(account_id)?; |
| let _ = device::save_global_original(&profile); |
| apply_profile_to_account( |
| &mut account, profile.clone(), Some(mode.to_string()), true)?; |
|
|
| Ok(profile) |
| } |
|
|
| |
| pub fn bind_device_profile_with_profile( |
| account_id: &str, |
| profile: DeviceProfile, |
| label: Option<String>, |
| ) -> Result<DeviceProfile, String> { |
| let mut account = load_account(account_id)?; |
| let _ = crate::modules::device::save_global_original(&profile); |
| apply_profile_to_account(&mut account, profile.clone(), label, true)?; |
|
|
| Ok(profile) |
| } |
|
|
| fn apply_profile_to_account( |
| account: &mut Account, |
| profile: DeviceProfile, |
| label: Option<String>, |
| add_history: bool, |
| ) -> Result<(), String> { |
| account.device_profile = Some(profile.clone()); |
| if add_history { |
| |
| for h in account.device_history.iter_mut() { |
| h.is_current = false; |
| } |
| account.device_history.push(DeviceProfileVersion { |
| id: Uuid::new_v4().to_string(), |
| created_at: chrono::Utc::now().timestamp(), |
| label: label.unwrap_or_else(|| "generated".to_string()), |
| profile: profile.clone(), |
| is_current: true, |
| }); |
| } |
| save_account(account)?; |
| Ok(()) |
| } |
|
|
| |
| pub fn list_device_versions(account_id: &str) -> Result<DeviceProfiles, String> { |
| get_device_profiles(account_id) |
| } |
|
|
| |
| pub fn restore_device_version(account_id: &str, version_id: &str) -> Result<DeviceProfile, String> { |
| let mut account = load_account(account_id)?; |
|
|
| let target_profile = if version_id == "baseline" { |
| crate::modules::device::load_global_original().ok_or("Global original profile not found")? |
| } else if let Some(v) = account.device_history.iter().find(|v| v.id == version_id) { |
| v.profile.clone() |
| } else if version_id == "current" { |
| account |
| .device_profile |
| .clone() |
| .ok_or("No currently bound profile")? |
| } else { |
| return Err("Device profile version not found".to_string()); |
| }; |
|
|
| account.device_profile = Some(target_profile.clone()); |
| for h in account.device_history.iter_mut() { |
| h.is_current = h.id == version_id; |
| } |
| save_account(&account)?; |
| Ok(target_profile) |
| } |
|
|
| |
| pub fn delete_device_version(account_id: &str, version_id: &str) -> Result<(), String> { |
| if version_id == "baseline" { |
| return Err("Original profile cannot be deleted".to_string()); |
| } |
| let mut account = load_account(account_id)?; |
| if account |
| .device_history |
| .iter() |
| .any(|v| v.id == version_id && v.is_current) |
| { |
| return Err("Currently bound profile cannot be deleted".to_string()); |
| } |
| let before = account.device_history.len(); |
| account.device_history.retain(|v| v.id != version_id); |
| if account.device_history.len() == before { |
| return Err("Historical device profile not found".to_string()); |
| } |
| save_account(&account)?; |
| Ok(()) |
| } |
| |
| pub fn apply_device_profile(account_id: &str) -> Result<DeviceProfile, String> { |
| use crate::modules::device; |
| let mut account = load_account(account_id)?; |
| let profile = account |
| .device_profile |
| .clone() |
| .ok_or("Account has no bound device profile")?; |
| let storage_path = device::get_storage_path()?; |
| device::write_profile(&storage_path, &profile)?; |
| account.update_last_used(); |
| save_account(&account)?; |
| Ok(profile) |
| } |
|
|
| |
| pub fn restore_original_device() -> Result<String, String> { |
| if let Some(current_id) = get_current_account_id()? { |
| if let Ok(mut account) = load_account(¤t_id) { |
| if let Some(original) = crate::modules::device::load_global_original() { |
| account.device_profile = Some(original); |
| for h in account.device_history.iter_mut() { |
| h.is_current = false; |
| } |
| save_account(&account)?; |
| return Ok( |
| "Reset current account bound profile to original (not applied to storage)" |
| .to_string(), |
| ); |
| } |
| } |
| } |
| Err("Original profile not found, cannot restore".to_string()) |
| } |
|
|
| |
| pub fn get_current_account_id() -> Result<Option<String>, String> { |
| let index = load_account_index()?; |
| Ok(index.current_account_id) |
| } |
|
|
| |
| pub fn get_current_account() -> Result<Option<Account>, String> { |
| if let Some(id) = get_current_account_id()? { |
| Ok(Some(load_account(&id)?)) |
| } else { |
| Ok(None) |
| } |
| } |
|
|
| |
| pub fn set_current_account_id(account_id: &str) -> Result<(), String> { |
| let _lock = ACCOUNT_INDEX_LOCK |
| .lock() |
| .map_err(|e| format!("failed_to_acquire_lock: {}", e))?; |
| let mut index = load_account_index()?; |
| index.current_account_id = Some(account_id.to_string()); |
| save_account_index(&index) |
| } |
|
|
| |
| pub fn update_account_quota(account_id: &str, quota: QuotaData) -> Result<(), String> { |
| let mut account = load_account(account_id)?; |
| account.update_quota(quota); |
|
|
| |
| if let Ok(config) = crate::modules::config::load_app_config() { |
| if config.quota_protection.enabled { |
| if let Some(ref q) = account.quota { |
| let threshold = config.quota_protection.threshold_percentage as i32; |
|
|
| let mut group_min_percentage: HashMap<String, i32> = HashMap::new(); |
|
|
| for model in &q.models { |
| if let Some(std_id) = |
| crate::proxy::common::model_mapping::normalize_to_standard_id(&model.name) |
| { |
| let entry = group_min_percentage.entry(std_id).or_insert(100); |
| if model.percentage < *entry { |
| *entry = model.percentage; |
| } |
| } |
| } |
|
|
| for std_id in &config.quota_protection.monitored_models { |
| let min_pct = group_min_percentage.get(std_id).cloned().unwrap_or(100); |
|
|
| if min_pct <= threshold { |
| if !account.protected_models.contains(std_id) { |
| crate::modules::logger::log_info(&format!( |
| "[Quota] Triggering model protection: {} (Group: {} Min: {}% <= Thres: {}%)", |
| account.email, std_id, min_pct, threshold |
| )); |
| account.protected_models.insert(std_id.clone()); |
| } |
| } else { |
| if account.protected_models.contains(std_id) { |
| crate::modules::logger::log_info(&format!( |
| "[Quota] Model protection recovered: {} (Group: {} Min: {}% > Thres: {}%)", |
| account.email, std_id, min_pct, threshold |
| )); |
| account.protected_models.remove(std_id); |
| } |
| } |
| } |
|
|
| |
| if account.proxy_disabled |
| && account |
| .proxy_disabled_reason |
| .as_ref() |
| .map_or(false, |r| r == "quota_protection") |
| { |
| crate::modules::logger::log_info(&format!( |
| "[Quota] Migrating account {} from account-level to model-level protection", |
| account.email |
| )); |
| account.proxy_disabled = false; |
| account.proxy_disabled_reason = None; |
| account.proxy_disabled_at = None; |
| } |
| } |
| } |
| } |
| |
|
|
| |
| save_account(&account)?; |
|
|
| |
| { |
| let _lock = ACCOUNT_INDEX_LOCK |
| .lock() |
| .map_err(|e| format!("failed_to_acquire_lock: {}", e))?; |
| if let Ok(mut index) = load_account_index() { |
| if let Some(summary) = index.accounts.iter_mut().find(|a| a.id == account_id) { |
| summary.protected_models = account.protected_models.clone(); |
| let _ = save_account_index(&index); |
| } |
| } |
| } |
|
|
| |
| |
| crate::proxy::server::trigger_account_reload(account_id); |
|
|
| Ok(()) |
| } |
|
|
| |
| pub fn toggle_proxy_status( |
| account_id: &str, |
| enable: bool, |
| reason: Option<&str>, |
| ) -> Result<(), String> { |
| let _lock = ACCOUNT_INDEX_LOCK |
| .lock() |
| .map_err(|e| format!("failed_to_acquire_lock: {}", e))?; |
|
|
| let mut account = load_account(account_id)?; |
|
|
| account.proxy_disabled = !enable; |
| account.proxy_disabled_reason = if !enable { |
| reason.map(|s| s.to_string()) |
| } else { |
| None |
| }; |
| account.proxy_disabled_at = if !enable { |
| Some(chrono::Utc::now().timestamp()) |
| } else { |
| None |
| }; |
|
|
| save_account(&account)?; |
|
|
| |
| let mut index = load_account_index()?; |
| if let Some(summary) = index.accounts.iter_mut().find(|a| a.id == account_id) { |
| summary.proxy_disabled = !enable; |
| save_account_index(&index)?; |
| } |
|
|
| Ok(()) |
| } |
|
|
| |
| pub fn find_account_id_by_email(email: &str) -> Option<String> { |
| load_account_index().ok()?.accounts.into_iter() |
| .find(|a| a.email == email) |
| .map(|a| a.id) |
| } |
|
|
| pub fn mark_account_forbidden(account_id: &str, reason: &str) -> Result<(), String> { |
| let _lock = ACCOUNT_INDEX_LOCK |
| .lock() |
| .map_err(|e| format!("failed_to_acquire_lock: {}", e))?; |
|
|
| let mut account = load_account(account_id)?; |
|
|
| |
| if let Some(ref mut q) = account.quota { |
| q.is_forbidden = true; |
| q.forbidden_reason = Some(reason.to_string()); |
| } else { |
| account.quota = Some(crate::models::QuotaData { |
| models: Vec::new(), |
| last_updated: chrono::Utc::now().timestamp(), |
| subscription_tier: None, |
| is_forbidden: true, |
| forbidden_reason: Some(reason.to_string()), |
| model_forwarding_rules: std::collections::HashMap::new(), |
| }); |
| } |
|
|
| |
| account.proxy_disabled = true; |
| account.proxy_disabled_reason = Some(format!("Forbidden (403): {}", reason)); |
| account.proxy_disabled_at = Some(chrono::Utc::now().timestamp()); |
|
|
| save_account(&account)?; |
|
|
| |
| let mut index = load_account_index()?; |
| if let Some(summary) = index.accounts.iter_mut().find(|a| a.id == account_id) { |
| summary.proxy_disabled = true; |
| save_account_index(&index)?; |
| } |
|
|
| |
| crate::modules::log_bridge::emit_accounts_refreshed(); |
|
|
| Ok(()) |
| } |
|
|
| |
| pub fn export_accounts_by_ids(account_ids: &[String]) -> Result<crate::models::AccountExportResponse, String> { |
| use crate::models::{AccountExportItem, AccountExportResponse}; |
| |
| let accounts = list_accounts()?; |
| |
| let export_items: Vec<AccountExportItem> = accounts |
| .into_iter() |
| .filter(|acc| account_ids.contains(&acc.id)) |
| .map(|acc| AccountExportItem { |
| email: acc.email, |
| refresh_token: acc.token.refresh_token, |
| }) |
| .collect(); |
|
|
| Ok(AccountExportResponse { |
| accounts: export_items, |
| }) |
| } |
|
|
| |
| #[allow(dead_code)] |
| pub fn export_accounts() -> Result<Vec<(String, String)>, String> { |
| let accounts = list_accounts()?; |
| let mut exports = Vec::new(); |
|
|
| for account in accounts { |
| exports.push((account.email, account.token.refresh_token)); |
| } |
|
|
| Ok(exports) |
| } |
|
|
| |
| pub async fn fetch_quota_with_retry(account: &mut Account) -> crate::error::AppResult<QuotaData> { |
| use crate::error::AppError; |
| use crate::modules::oauth; |
|
|
| |
| let token = match oauth::ensure_fresh_token(&account.token, Some(&account.id)).await { |
| Ok(t) => t, |
| Err(e) => { |
| if e.contains("invalid_grant") { |
| modules::logger::log_error(&format!( |
| "Disabling account {} due to invalid_grant during token refresh (quota check)", |
| account.email |
| )); |
| account.disabled = true; |
| account.disabled_at = Some(chrono::Utc::now().timestamp()); |
| account.disabled_reason = Some(format!("invalid_grant: {}", e)); |
| let _ = save_account(account); |
| crate::proxy::server::trigger_account_reload(&account.id); |
| } |
| return Err(AppError::OAuth(e)); |
| } |
| }; |
|
|
| if token.access_token != account.token.access_token { |
| modules::logger::log_info(&format!("Time-based Token refresh: {}", account.email)); |
| account.token = token.clone(); |
|
|
| |
| let name = if account.name.is_none() |
| || account.name.as_ref().map_or(false, |n| n.trim().is_empty()) |
| { |
| match oauth::get_user_info(&token.access_token, Some(&account.id)).await { |
| Ok(user_info) => user_info.get_display_name(), |
| Err(_) => None, |
| } |
| } else { |
| account.name.clone() |
| }; |
|
|
| account.name = name.clone(); |
| upsert_account(account.email.clone(), name, token.clone()).map_err(AppError::Account)?; |
| } |
|
|
| |
| if account.name.is_none() || account.name.as_ref().map_or(false, |n| n.trim().is_empty()) { |
| modules::logger::log_info(&format!( |
| "Account {} missing display name, attempting to fetch...", |
| account.email |
| )); |
| |
| match oauth::get_user_info(&account.token.access_token, Some(&account.id)).await { |
| Ok(user_info) => { |
| let display_name = user_info.get_display_name(); |
| modules::logger::log_info(&format!( |
| "Successfully fetched display name: {:?}", |
| display_name |
| )); |
| account.name = display_name.clone(); |
| |
| if let Err(e) = |
| upsert_account(account.email.clone(), display_name, account.token.clone()) |
| { |
| modules::logger::log_warn(&format!("Failed to save display name: {}", e)); |
| } |
| } |
| Err(e) => { |
| modules::logger::log_warn(&format!("Failed to fetch display name: {}", e)); |
| } |
| } |
| } |
|
|
| |
| let result: crate::error::AppResult<(QuotaData, Option<String>)> = |
| modules::fetch_quota(&account.token.access_token, &account.email, Some(&account.id)).await; |
|
|
| |
| if let Ok((ref _q, ref project_id)) = result { |
| if project_id.is_some() && *project_id != account.token.project_id { |
| modules::logger::log_info(&format!( |
| "Detected project_id update ({}), saving...", |
| account.email |
| )); |
| account.token.project_id = project_id.clone(); |
| if let Err(e) = upsert_account( |
| account.email.clone(), |
| account.name.clone(), |
| account.token.clone(), |
| ) { |
| modules::logger::log_warn(&format!("Failed to sync project_id: {}", e)); |
| } |
| } |
| } |
|
|
| |
| if let Err(AppError::Network(_, status)) = result { |
| if let Some(code) = status { |
| if code == 401 { |
| modules::logger::log_warn(&format!( |
| "401 Unauthorized for {}, forcing refresh...", |
| account.email |
| )); |
|
|
| |
| let token_res = match oauth::refresh_access_token_with_client( |
| &account.token.refresh_token, |
| Some(&account.id), |
| account.token.oauth_client_key.as_deref(), |
| ).await { |
| Ok(t) => t, |
| Err(e) => { |
| if e.contains("invalid_grant") { |
| modules::logger::log_error(&format!( |
| "Disabling account {} due to invalid_grant during forced refresh (quota check)", |
| account.email |
| )); |
| account.disabled = true; |
| account.disabled_at = Some(chrono::Utc::now().timestamp()); |
| account.disabled_reason = Some(format!("invalid_grant: {}", e)); |
| let _ = save_account(account); |
| crate::proxy::server::trigger_account_reload(&account.id); |
| } |
| return Err(AppError::OAuth(e)); |
| } |
| }; |
|
|
| let new_token = TokenData::new( |
| token_res.access_token.clone(), |
| account.token.refresh_token.clone(), |
| token_res.expires_in, |
| account.token.email.clone(), |
| account.token.project_id.clone(), |
| None, |
| account.token.is_gcp_tos, |
| ) |
| .with_oauth_client_key( |
| token_res |
| .oauth_client_key |
| .clone() |
| .or_else(|| account.token.oauth_client_key.clone()), |
| ); |
|
|
| |
| let name = if account.name.is_none() |
| || account.name.as_ref().map_or(false, |n| n.trim().is_empty()) |
| { |
| match oauth::get_user_info(&token_res.access_token, Some(&account.id)).await { |
| Ok(user_info) => user_info.get_display_name(), |
| Err(_) => None, |
| } |
| } else { |
| account.name.clone() |
| }; |
|
|
| account.token = new_token.clone(); |
| account.name = name.clone(); |
| upsert_account(account.email.clone(), name, new_token.clone()) |
| .map_err(AppError::Account)?; |
|
|
| |
| let retry_result: crate::error::AppResult<(QuotaData, Option<String>)> = |
| modules::fetch_quota(&new_token.access_token, &account.email, Some(&account.id)).await; |
|
|
| |
| if let Ok((ref _q, ref project_id)) = retry_result { |
| if project_id.is_some() && *project_id != account.token.project_id { |
| modules::logger::log_info(&format!( |
| "Detected update of project_id after retry ({}), saving...", |
| account.email |
| )); |
| account.token.project_id = project_id.clone(); |
| let _ = upsert_account( |
| account.email.clone(), |
| account.name.clone(), |
| account.token.clone(), |
| ); |
| } |
| } |
|
|
| if let Err(AppError::Network(_, status)) = retry_result { |
| if let Some(code) = status { |
| if code == 403 { |
| let mut q = QuotaData::new(); |
| q.is_forbidden = true; |
| return Ok(q); |
| } |
| } |
| } |
|
|
| match retry_result { |
| Ok((q, _)) => { |
| clear_validation_blocked(account); |
| return Ok(q); |
| } |
| Err(e) => { |
| if is_validation_required_error(&e) { |
| mark_validation_blocked(account, &e.to_string()); |
| } |
| if let Some(cached) = recover_cached_quota_on_rate_limit(account, &e) { |
| mark_validation_blocked(account, &format_rate_limit_block_reason(&e)); |
| modules::logger::log_warn(&format!( |
| "Quota API rate-limited for {}, using cached model list as fallback", |
| account.email |
| )); |
| return Ok(cached); |
| } |
| return Err(e); |
| } |
| } |
| } |
| } |
| } |
|
|
| |
| match result { |
| Ok((q, _)) => { |
| clear_validation_blocked(account); |
| Ok(q) |
| } |
| Err(e) => { |
| if is_validation_required_error(&e) { |
| mark_validation_blocked(account, &e.to_string()); |
| } |
| if let Some(cached) = recover_cached_quota_on_rate_limit(account, &e) { |
| mark_validation_blocked(account, &format_rate_limit_block_reason(&e)); |
| modules::logger::log_warn(&format!( |
| "Quota API rate-limited for {}, using cached model list as fallback", |
| account.email |
| )); |
| return Ok(cached); |
| } |
| Err(e) |
| } |
| } |
| } |
|
|
| #[derive(Serialize)] |
| pub struct RefreshStats { |
| pub total: usize, |
| pub success: usize, |
| pub failed: usize, |
| pub details: Vec<String>, |
| } |
|
|
| |
| pub async fn refresh_all_quotas_logic() -> Result<RefreshStats, String> { |
| use futures::future::join_all; |
| use std::sync::Arc; |
| use tokio::sync::Semaphore; |
|
|
| const MAX_CONCURRENT: usize = 5; |
| let start = std::time::Instant::now(); |
|
|
| crate::modules::logger::log_info(&format!( |
| "Starting batch refresh of all account quotas (Concurrent mode, max: {})", |
| MAX_CONCURRENT |
| )); |
| let accounts = list_accounts()?; |
|
|
| let semaphore = Arc::new(Semaphore::new(MAX_CONCURRENT)); |
|
|
| let tasks: Vec<_> = accounts |
| .into_iter() |
| .filter(|account| { |
| |
| |
| |
| |
| if let Some(ref q) = account.quota { |
| if q.is_forbidden { |
| crate::modules::logger::log_info(&format!( |
| " - Skipping {} (Forbidden)", |
| account.email |
| )); |
| return false; |
| } |
| } |
| true |
| }) |
| .map(|mut account| { |
| let email = account.email.clone(); |
| let account_id = account.id.clone(); |
| let permit = semaphore.clone(); |
| async move { |
| let _guard = permit.acquire().await.unwrap(); |
| crate::modules::logger::log_info(&format!(" - Processing {}", email)); |
| match fetch_quota_with_retry(&mut account).await { |
| Ok(quota) => { |
| if let Err(e) = update_account_quota(&account_id, quota) { |
| let msg = format!("Account {}: Save quota failed - {}", email, e); |
| crate::modules::logger::log_error(&msg); |
| Err(msg) |
| } else { |
| crate::modules::logger::log_info(&format!(" Success {}", email)); |
| Ok(()) |
| } |
| } |
| Err(e) => { |
| let msg = format!("Account {}: Fetch quota failed - {}", email, e); |
| crate::modules::logger::log_error(&msg); |
| Err(msg) |
| } |
| } |
| } |
| }) |
| .collect(); |
|
|
| let total = tasks.len(); |
| let results = join_all(tasks).await; |
|
|
| let mut success = 0; |
| let mut failed = 0; |
| let mut details = Vec::new(); |
|
|
| for result in results { |
| match result { |
| Ok(()) => success += 1, |
| Err(msg) => { |
| failed += 1; |
| details.push(msg); |
| } |
| } |
| } |
|
|
| let elapsed = start.elapsed(); |
| crate::modules::logger::log_info(&format!( |
| "Batch refresh completed: {} success, {} failed, took: {}ms", |
| success, |
| failed, |
| elapsed.as_millis() |
| )); |
|
|
| |
| |
| |
| |
| |
|
|
| Ok(RefreshStats { |
| total, |
| success, |
| failed, |
| details, |
| }) |
| } |
|
|
| |
| |
| pub async fn check_and_trigger_warmup_for_recovered_models() { |
| let accounts = match list_accounts() { |
| Ok(acc) => acc, |
| Err(_) => return, |
| }; |
|
|
| |
| let app_config = match crate::modules::config::load_app_config() { |
| Ok(cfg) => cfg, |
| Err(_) => return, |
| }; |
|
|
| if !app_config.scheduled_warmup.enabled { |
| return; |
| } |
|
|
| crate::modules::logger::log_info(&format!( |
| "[Warmup] Checking {} accounts for recovered models after quota refresh...", |
| accounts.len() |
| )); |
|
|
| for account in accounts { |
| |
| if account.disabled || account.proxy_disabled { |
| continue; |
| } |
|
|
| |
| crate::modules::scheduler::trigger_warmup_for_account(&account).await; |
| } |
| } |
|
|