use std::collections::HashMap;
use std::sync::Arc;
use reqwest::Client;
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct ClientCacheKey {
pub provider: String,
pub api_key: Option<String>,
pub api_base: Option<String>,
}
impl ClientCacheKey {
pub fn new(provider: &str, api_key: Option<&str>, api_base: Option<&str>) -> Self {
Self {
provider: provider.to_owned(),
api_key: api_key.map(|s| s.to_owned()),
api_base: api_base.map(|s| s.to_owned()),
}
}
}
#[derive(Debug, Clone)]
pub struct ClientConfig {
pub default_headers: HashMap<String, String>,
pub timeout_secs: u64,
}
impl Default for ClientConfig {
fn default() -> Self {
Self {
default_headers: HashMap::new(),
timeout_secs: 120,
}
}
}
#[derive(Debug, Clone)]
pub struct ClientRegistry {
cache: HashMap<ClientCacheKey, Arc<Client>>,
config: ClientConfig,
}
impl ClientRegistry {
pub fn new() -> Self {
Self {
cache: HashMap::new(),
config: ClientConfig::default(),
}
}
pub fn with_config(config: ClientConfig) -> Self {
Self {
cache: HashMap::new(),
config,
}
}
pub fn get_or_create(
&mut self,
provider: &str,
api_key: Option<&str>,
api_base: Option<&str>,
) -> Arc<Client> {
let cache_key = ClientCacheKey::new(provider, api_key, api_base);
if let Some(client) = self.cache.get(&cache_key) {
return Arc::clone(client);
}
let client = Arc::new(self.build_client(provider, api_key));
self.cache.insert(cache_key, Arc::clone(&client));
client
}
pub fn contains(&self, provider: &str, api_key: Option<&str>, api_base: Option<&str>) -> bool {
let cache_key = ClientCacheKey::new(provider, api_key, api_base);
self.cache.contains_key(&cache_key)
}
pub fn clear(&mut self) {
self.cache.clear();
}
pub fn len(&self) -> usize {
self.cache.len()
}
pub fn is_empty(&self) -> bool {
self.cache.is_empty()
}
fn build_client(&self, provider: &str, api_key: Option<&str>) -> Client {
let mut headers = reqwest::header::HeaderMap::new();
for (key, value) in &self.config.default_headers {
if let (Ok(name), Ok(val)) = (
reqwest::header::HeaderName::from_bytes(key.as_bytes()),
reqwest::header::HeaderValue::from_str(value),
) {
headers.insert(name, val);
}
}
let is_anthropic = provider.eq_ignore_ascii_case("anthropic");
let is_oauth_token = api_key
.map(|k| k.to_ascii_lowercase().starts_with("sk-ant-oat"))
.unwrap_or(false);
if let Some(key) = api_key {
if is_anthropic && !is_oauth_token {
if let Ok(val) = reqwest::header::HeaderValue::from_str(key)
&& let Ok(name) = reqwest::header::HeaderName::from_bytes(b"x-api-key")
{
headers.insert(name, val);
}
} else {
if let Ok(val) = reqwest::header::HeaderValue::from_str(&format!("Bearer {key}")) {
headers.insert(reqwest::header::AUTHORIZATION, val);
}
}
}
if is_anthropic
&& let Ok(val) = reqwest::header::HeaderValue::from_str("2023-06-01")
&& let Ok(name) = reqwest::header::HeaderName::from_bytes(b"anthropic-version")
{
headers.insert(name, val);
}
if is_anthropic && is_oauth_token {
headers.insert(
reqwest::header::USER_AGENT,
reqwest::header::HeaderValue::from_static("claude-cli/2.1.75"),
);
if let Ok(name) = reqwest::header::HeaderName::from_bytes(b"x-app") {
headers.insert(name, reqwest::header::HeaderValue::from_static("cli"));
}
headers.insert(
reqwest::header::ACCEPT,
reqwest::header::HeaderValue::from_static("application/json"),
);
if let Ok(name) = reqwest::header::HeaderName::from_bytes(
b"anthropic-dangerous-direct-browser-access",
) {
headers.insert(name, reqwest::header::HeaderValue::from_static("true"));
}
if let Ok(name) = reqwest::header::HeaderName::from_bytes(b"anthropic-beta") {
headers.insert(
name,
reqwest::header::HeaderValue::from_static(
"claude-code-20250219,oauth-2025-04-20,fine-grained-tool-streaming-2025-05-14,interleaved-thinking-2025-05-14"
),
);
}
}
headers.insert(
reqwest::header::CONTENT_TYPE,
reqwest::header::HeaderValue::from_static("application/json"),
);
Client::builder()
.default_headers(headers)
.timeout(std::time::Duration::from_secs(self.config.timeout_secs))
.build()
.unwrap_or_else(|err| {
tracing::warn!(%err, provider, "failed to build HTTP client with custom config, falling back to default");
Client::new()
})
}
}
impl Default for ClientRegistry {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_cache_key_equality() {
let k1 = ClientCacheKey::new("openai", Some("sk-123"), None);
let k2 = ClientCacheKey::new("openai", Some("sk-123"), None);
assert_eq!(k1, k2);
}
#[test]
fn test_get_or_create_caches() {
let mut registry = ClientRegistry::new();
let c1 = registry.get_or_create("openai", Some("sk-123"), None);
let c2 = registry.get_or_create("openai", Some("sk-123"), None);
assert!(Arc::ptr_eq(&c1, &c2));
assert_eq!(registry.len(), 1);
}
#[test]
fn test_different_keys_separate_clients() {
let mut registry = ClientRegistry::new();
let c1 = registry.get_or_create("openai", Some("sk-111"), None);
let c2 = registry.get_or_create("openai", Some("sk-222"), None);
assert!(!Arc::ptr_eq(&c1, &c2));
assert_eq!(registry.len(), 2);
}
#[test]
fn test_clear() {
let mut registry = ClientRegistry::new();
registry.get_or_create("openai", Some("sk-123"), None);
assert!(!registry.is_empty());
registry.clear();
assert!(registry.is_empty());
}
}