elph_ai/auth/oauth/
registry.rs

1//! Global OAuth provider registry matching pi-ai `utils/oauth/index.ts`.
2
3use std::collections::HashMap;
4use std::sync::{Arc, RwLock};
5
6use crate::auth::helpers::lazy_oauth;
7use crate::auth::oauth::{anthropic_oauth_loader, github_copilot_oauth_loader, openai_codex_oauth_loader};
8use crate::auth::types::{AuthLoginCallbacks, ModelAuth, OAuthAuth, OAuthCredential};
9use crate::models::catalog::GITHUB_COPILOT_MODELS;
10use crate::types::Model;
11
12pub type OAuthProviderId = String;
13
14pub type OAuthModifyModelsFn = Arc<dyn Fn(Vec<Model>, &OAuthCredential) -> Vec<Model> + Send + Sync>;
15
16#[derive(Clone)]
17pub struct OAuthProviderInterface {
18    pub id: OAuthProviderId,
19    pub name: String,
20    pub auth: OAuthAuth,
21    pub get_api_key: Arc<dyn Fn(&OAuthCredential) -> String + Send + Sync>,
22    pub modify_models: Option<OAuthModifyModelsFn>,
23}
24
25fn anthropic_provider() -> OAuthProviderInterface {
26    OAuthProviderInterface {
27        id: "anthropic".to_string(),
28        name: "Anthropic (Claude Pro/Max)".to_string(),
29        auth: lazy_oauth("Anthropic (Claude Pro/Max)", anthropic_oauth_loader()),
30        get_api_key: Arc::new(|c| c.access.clone()),
31        modify_models: None,
32    }
33}
34
35fn github_copilot_provider() -> OAuthProviderInterface {
36    OAuthProviderInterface {
37        id: "github-copilot".to_string(),
38        name: "GitHub Copilot".to_string(),
39        auth: lazy_oauth("GitHub Copilot", github_copilot_oauth_loader()),
40        get_api_key: Arc::new(|c| c.access.clone()),
41        modify_models: Some(Arc::new(modify_github_copilot_models)),
42    }
43}
44
45fn openai_codex_provider() -> OAuthProviderInterface {
46    OAuthProviderInterface {
47        id: "openai-codex".to_string(),
48        name: "OpenAI (ChatGPT Plus/Pro)".to_string(),
49        auth: lazy_oauth("OpenAI (ChatGPT Plus/Pro)", openai_codex_oauth_loader()),
50        get_api_key: Arc::new(|c| c.access.clone()),
51        modify_models: None,
52    }
53}
54
55fn built_in_providers() -> Vec<OAuthProviderInterface> {
56    vec![anthropic_provider(), github_copilot_provider(), openai_codex_provider()]
57}
58
59fn modify_github_copilot_models(models: Vec<Model>, credential: &OAuthCredential) -> Vec<Model> {
60    let enterprise_domain = credential
61        .enterprise_url
62        .as_deref()
63        .and_then(crate::auth::oauth::normalize_domain);
64    let base_url =
65        crate::auth::oauth::get_github_copilot_base_url(Some(&credential.access), enterprise_domain.as_deref());
66    models
67        .into_iter()
68        .map(|mut model| {
69            model.base_url = base_url.clone();
70            model
71        })
72        .collect()
73}
74
75static REGISTRY: once_cell::sync::Lazy<RwLock<HashMap<String, OAuthProviderInterface>>> =
76    once_cell::sync::Lazy::new(|| {
77        let mut map = HashMap::new();
78        for provider in built_in_providers() {
79            map.insert(provider.id.clone(), provider);
80        }
81        RwLock::new(map)
82    });
83
84pub fn get_oauth_provider(id: &str) -> Option<OAuthProviderInterface> {
85    REGISTRY.read().ok()?.get(id).cloned()
86}
87
88pub fn register_oauth_provider(provider: OAuthProviderInterface) {
89    if let Ok(mut registry) = REGISTRY.write() {
90        registry.insert(provider.id.clone(), provider);
91    }
92}
93
94pub fn unregister_oauth_provider(id: &str) {
95    let Ok(mut registry) = REGISTRY.write() else {
96        return;
97    };
98    if let Some(built_in) = built_in_providers().into_iter().find(|p| p.id == id) {
99        registry.insert(id.to_string(), built_in);
100        return;
101    }
102    registry.remove(id);
103}
104
105pub fn reset_oauth_providers() {
106    if let Ok(mut registry) = REGISTRY.write() {
107        registry.clear();
108        for provider in built_in_providers() {
109            registry.insert(provider.id.clone(), provider);
110        }
111    }
112}
113
114pub fn get_oauth_providers() -> Vec<OAuthProviderInterface> {
115    REGISTRY
116        .read()
117        .map(|registry| registry.values().cloned().collect())
118        .unwrap_or_default()
119}
120
121pub async fn refresh_oauth_token(provider_id: &str, credential: OAuthCredential) -> anyhow::Result<OAuthCredential> {
122    let provider =
123        get_oauth_provider(provider_id).ok_or_else(|| anyhow::anyhow!("Unknown OAuth provider: {provider_id}"))?;
124    (provider.auth.refresh)(credential).await
125}
126
127pub struct OAuthApiKeyResult {
128    pub new_credentials: OAuthCredential,
129    pub api_key: String,
130}
131
132pub async fn get_oauth_api_key(
133    provider_id: &str,
134    mut credential: OAuthCredential,
135) -> anyhow::Result<OAuthApiKeyResult> {
136    let provider =
137        get_oauth_provider(provider_id).ok_or_else(|| anyhow::anyhow!("Unknown OAuth provider: {provider_id}"))?;
138
139    if chrono::Utc::now().timestamp_millis() >= credential.expires {
140        credential = (provider.auth.refresh)(credential).await?;
141    }
142
143    let api_key = (provider.get_api_key)(&credential);
144    Ok(OAuthApiKeyResult {
145        new_credentials: credential,
146        api_key,
147    })
148}
149
150pub async fn oauth_provider_login(
151    provider_id: &str,
152    callbacks: Arc<dyn AuthLoginCallbacks>,
153) -> anyhow::Result<OAuthCredential> {
154    let provider =
155        get_oauth_provider(provider_id).ok_or_else(|| anyhow::anyhow!("Unknown OAuth provider: {provider_id}"))?;
156    (provider.auth.login)(callbacks).await
157}
158
159pub async fn oauth_provider_to_auth(provider_id: &str, credential: OAuthCredential) -> anyhow::Result<ModelAuth> {
160    let provider =
161        get_oauth_provider(provider_id).ok_or_else(|| anyhow::anyhow!("Unknown OAuth provider: {provider_id}"))?;
162    (provider.auth.to_auth)(credential).await
163}
164
165pub fn oauth_provider_modify_models(provider_id: &str, models: Vec<Model>, credential: &OAuthCredential) -> Vec<Model> {
166    let Some(provider) = get_oauth_provider(provider_id) else {
167        return models;
168    };
169    provider
170        .modify_models
171        .as_ref()
172        .map(|modify| modify(models.clone(), credential))
173        .unwrap_or(models)
174}
175
176pub fn builtin_oauth_provider_ids() -> Vec<&'static str> {
177    vec!["anthropic", "github-copilot", "openai-codex"]
178}
179
180pub fn github_copilot_catalog_models() -> Vec<Model> {
181    GITHUB_COPILOT_MODELS.iter().cloned().collect()
182}
elph_ai/auth/oauth/registry.rs

elph_ai/auth/oauth/
registry.rs
