feat: multi-account support and OAuth tokens management

Cargo.toml

@@ -56,6 +56,8 @@ once_cell = "1"            # Lazy statics
 async-trait = "0.1"        # Async trait support
 dirs = "5"                 # User directories
 regex = "1"                # Regular expressions
+uuid = { version = "1", features = ["v4", "serde"] }  # UUID generation
+atomicwrites = "0.4"       # Atomic file writes for crash-safe persistence
 
 # OAuth & Auth
 oauth2 = "4"               # OAuth 2.0 client
@@ -73,6 +75,7 @@ nix = { version = "0.29", features = ["signal"] }  # Unix signals
 tokio-test = "0.4"
 mockito = "1"
 criterion = "0.5"          # Benchmarking
+tempfile = "3"             # Temporary files for tests
 
 # Property Testing
 proptest = "1"

README.md

@@ -54,13 +54,17 @@ Claude Code → Claude Code Mux → Multiple AI Providers
 - 🌐 **Multi-Provider Support** - 18+ providers including OpenAI, Anthropic, Google Gemini/Vertex AI, Groq, ZenMux, etc.
 - ⚡️ **High Performance** - ~5MB RAM, <1ms routing overhead (Rust powered)
 - 🎯 **Unified API** - Full Anthropic Messages API compatibility
+- 🔀 **Multi-Account Load Balancing** - Distribute requests fairly across multiple accounts per provider
+- 📊 **Request Metrics** - Track usage statistics per account and provider
 
 ### 🚀 Advanced Features
 - 🔀 **Auto-mapping** - Regex-based model name transformation before routing (e.g., transform all `claude-*` to default model)
 - 🎯 **Background Detection** - Configurable regex patterns for background task detection
 - 🤖 **Multi-Agent Support** - Dynamic model switching via `CCM-SUBAGENT-MODEL` tags
 - 📊 **Live Testing** - Built-in test interface to verify routing and responses
 - ⚙️ **Centralized Settings** - Dedicated Settings tab for regex pattern management
+- 💾 **Smart Account Rotation** - Automatic cycling through accounts to balance load
+- 🔄 **Persistent Configuration** - Account state persists across service restarts
 
 ## Screenshots
 
@@ -1029,9 +1033,16 @@ If `zai` fails → automatically falls back to `openrouter`. **No manual interve
 - [Design Principles](docs/design-principles.md) - Claude Code Mux design philosophy and UX guidelines
 - [URL-based State Management](docs/url-state-management.md) - Admin UI URL-based state management pattern
 - [LocalStorage-based State Management](docs/localstorage-state-management.md) - Admin UI localStorage-based client state management
+- [Multi-Account Load Balancing](PHASE2_COMPLETION_REPORT.md) - Guide to setting up and using multiple accounts
+- [Build & Deployment](BUILD_SUMMARY.md) - Build instructions and deployment guide
 
 ## Changelog
 
+### Latest Features
+- **Multi-Account Load Balancing** - Distribute requests fairly across multiple accounts per provider
+- **Request Metrics** - Track usage and performance per account
+- **Smart Account Rotation** - Automatic cycling with state persistence across restarts
+
 See [CHANGELOG.md](CHANGELOG.md) for detailed release history or view [GitHub Releases](https://github.com/9j/claude-code-mux/releases) for downloads.
 
 ## Contributing

src/cli/mod.rs

@@ -1,7 +1,7 @@
 use serde::{Deserialize, Serialize};
 use std::path::PathBuf;
 use anyhow::{Context, Result};
-use crate::providers::ProviderConfig;
+use crate::providers::{ProviderConfig, ProviderAccount};
 
 /// Application configuration
 #[derive(Debug, Clone, Deserialize, Serialize)]
@@ -112,6 +112,9 @@ pub struct ModelMapping {
     pub provider: String,
     /// Actual model name to use with the provider
     pub actual_model: String,
+    /// Optional: Override provider-level load balance strategy for this mapping
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub load_balance_strategy: Option<crate::providers::LoadBalanceStrategy>,
 }
 
 impl ModelConfig {}
@@ -144,6 +147,9 @@ impl AppConfig {
         // Resolve environment variables
         config.resolve_env_vars()?;
 
+        // Normalize providers for multi-account support
+        config.normalize_providers()?;
+
         Ok(config)
     }
 
@@ -277,48 +283,235 @@ default = "placeholder-model"
                     }
                 }
             }
+
+            // Also resolve env vars in accounts if they exist
+            if let Some(ref mut accounts) = provider.accounts {
+                for account in accounts {
+                    if let Some(ref api_key) = account.api_key {
+                        if api_key.starts_with('$') {
+                            let env_var = &api_key[1..];
+                            if let Ok(value) = std::env::var(env_var) {
+                                account.api_key = Some(value);
+                            } else {
+                                anyhow::bail!("Environment variable {} not found for account {} in provider {}",
+                                    env_var, account.name, provider.name);
+                            }
+                        }
+                    }
+                }
+            }
+        }
+
+        Ok(())
+    }
+
+    /// Normalize provider configurations for multi-account support
+    /// - Ensures all providers have accounts (converts single-account format if needed)
+    /// - Generates UUIDs for accounts that don't have them
+    /// - Validates account configurations
+    fn normalize_providers(&mut self) -> Result<()> {
+        for provider in &mut self.providers {
+            // Skip disabled providers
+            if !provider.is_enabled() {
+                continue;
+            }
+
+            // If no accounts are configured, create one from the old single-account format
+            if provider.accounts.is_none() || provider.accounts.as_ref().unwrap().is_empty() {
+                let id = uuid::Uuid::new_v4().to_string();
+                let account = ProviderAccount {
+                    name: format!("{}-default", provider.name),
+                    id,
+                    api_key: provider.api_key.clone(),
+                    oauth_provider: provider.oauth_provider.clone(),
+                    enabled: true,
+                    priority: Some(1),
+                };
+                provider.accounts = Some(vec![account]);
+            } else {
+                // Ensure all accounts have IDs (for deterministic ordering)
+                if let Some(ref mut accounts) = provider.accounts {
+                    for account in accounts {
+                        if account.id.is_empty() {
+                            account.id = uuid::Uuid::new_v4().to_string();
+                        }
+                    }
+                }
+            }
         }
 
         Ok(())
     }
 }
 
-// TODO: Re-enable these tests by adding tempfile to dev-dependencies
-// #[cfg(test)]
-// mod tests {
-//     use super::*;
-//     use std::io::Write;
-//     use tempfile::NamedTempFile;
-//
-//     #[test]
-//     fn test_parse_toml_config() {
-//         let config_content = r#"
-// [server]
-// port = 3456
-// host = "127.0.0.1"
-// log_level = "info"
-//
-// [server.timeouts]
-// api_timeout_ms = 600000
-// connect_timeout_ms = 10000
-//
-// [litellm]
-// endpoint = "http://localhost:4000"
-// api_key = "anything"
-//
-// [router]
-// default = "default"
-// think = "think"
-//         "#;
-//
-//         let mut temp_file = NamedTempFile::new().unwrap();
-//         temp_file.write_all(config_content.as_bytes()).unwrap();
-//
-//         let config = AppConfig::from_file(&temp_file.path().to_path_buf()).unwrap();
-//
-//         assert_eq!(config.server.port, 3456);
-//         assert_eq!(config.litellm.endpoint, "http://localhost:4000");
-//         assert_eq!(config.litellm.api_key, "anything");
-//         assert_eq!(config.router.default, "default");
-//     }
-// }
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_parse_single_account_backward_compat() {
+        let config_content = r#"
+[server]
+port = 3456
+host = "127.0.0.1"
+
+[router]
+default = "claude-default"
+
+[[providers]]
+name = "openai"
+provider_type = "openai"
+auth_type = "apikey"
+api_key = "sk-test-key-123"
+models = []
+
+[[models]]
+name = "gpt-4"
+
+[[models.mappings]]
+priority = 1
+provider = "openai"
+actual_model = "gpt-4"
+"#;
+
+        let config: AppConfig = toml::from_str(config_content).unwrap();
+        assert_eq!(config.providers.len(), 1);
+        assert_eq!(config.providers[0].name, "openai");
+        assert_eq!(config.providers[0].api_key.as_ref().unwrap(), "sk-test-key-123");
+        assert!(config.providers[0].accounts.is_none());
+    }
+
+    #[test]
+    fn test_parse_multi_account_config() {
+        let config_content = r#"
+[server]
+port = 3456
+
+[router]
+default = "claude-default"
+
+[[providers]]
+name = "openai-multi"
+provider_type = "openai"
+models = []
+load_balance_strategy = "round_robin"
+
+[[providers.accounts]]
+name = "account-1"
+id = "uuid-123-account-1"
+api_key = "sk-test-1"
+enabled = true
+priority = 1
+
+[[providers.accounts]]
+name = "account-2"
+id = "uuid-123-account-2"
+api_key = "sk-test-2"
+enabled = true
+priority = 2
+
+[[models]]
+name = "gpt-4-multi"
+
+[[models.mappings]]
+priority = 1
+provider = "openai-multi"
+actual_model = "gpt-4"
+"#;
+
+        let config: AppConfig = toml::from_str(config_content).unwrap();
+        assert_eq!(config.providers.len(), 1);
+
+        let provider = &config.providers[0];
+        assert_eq!(provider.name, "openai-multi");
+        assert_eq!(provider.load_balance_strategy, Some(crate::providers::LoadBalanceStrategy::RoundRobin));
+
+        let accounts = provider.accounts.as_ref().unwrap();
+        assert_eq!(accounts.len(), 2);
+        assert_eq!(accounts[0].name, "account-1");
+        assert_eq!(accounts[0].id, "uuid-123-account-1");
+        assert_eq!(accounts[0].priority, Some(1));
+        assert_eq!(accounts[1].name, "account-2");
+        assert_eq!(accounts[1].priority, Some(2));
+    }
+
+    #[test]
+    fn test_normalize_providers_backward_compat() {
+        let config_content = r#"
+[router]
+default = "default"
+
+[[providers]]
+name = "openai"
+provider_type = "openai"
+api_key = "sk-test"
+models = []
+"#;
+
+        let mut config: AppConfig = toml::from_str(config_content).unwrap();
+        config.normalize_providers().unwrap();
+
+        // After normalization, single-account provider should have accounts created
+        let provider = &config.providers[0];
+        let accounts = provider.accounts.as_ref().unwrap();
+        assert_eq!(accounts.len(), 1);
+        assert_eq!(accounts[0].name, "openai-default");
+        assert_eq!(accounts[0].api_key.as_ref().unwrap(), "sk-test");
+        assert!(!accounts[0].id.is_empty()); // UUID should be generated
+    }
+
+    #[test]
+    fn test_normalize_providers_with_accounts() {
+        let config_content = r#"
+[router]
+default = "default"
+
+[[providers]]
+name = "openai-multi"
+provider_type = "openai"
+models = []
+
+[[providers.accounts]]
+name = "account-1"
+api_key = "sk-test-1"
+
+[[providers.accounts]]
+name = "account-2"
+api_key = "sk-test-2"
+"#;
+
+        let mut config: AppConfig = toml::from_str(config_content).unwrap();
+        config.normalize_providers().unwrap();
+
+        let provider = &config.providers[0];
+        let accounts = provider.accounts.as_ref().unwrap();
+        assert_eq!(accounts.len(), 2);
+
+        // All accounts should have IDs after normalization
+        for account in accounts {
+            assert!(!account.id.is_empty());
+        }
+    }
+
+    #[test]
+    fn test_model_mapping_with_strategy_override() {
+        let config_content = r#"
+[router]
+default = "default"
+
+[[models]]
+name = "gpt-4"
+
+[[models.mappings]]
+priority = 1
+provider = "openai"
+actual_model = "gpt-4"
+load_balance_strategy = "least_used"
+"#;
+
+        let config: AppConfig = toml::from_str(config_content).unwrap();
+        let model = &config.models[0];
+        let mapping = &model.mappings[0];
+        assert_eq!(mapping.load_balance_strategy, Some(crate::providers::LoadBalanceStrategy::LeastUsed));
+    }
+}