1use std::future::Future;
2use std::pin::Pin;
3use std::sync::Arc;
4
5use crate::types::{ImagesModel, Model, ProviderEnv, ProviderHeaders};
6
7#[derive(Debug, Clone)]
8pub struct ModelAuth {
9 pub api_key: Option<String>,
10 pub headers: Option<ProviderHeaders>,
11 pub base_url: Option<String>,
12}
13
14#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
15pub struct ApiKeyCredential {
16 #[serde(rename = "type")]
17 pub kind: String,
18 pub key: Option<String>,
19 #[serde(default, skip_serializing_if = "Option::is_none")]
20 pub env: Option<ProviderEnv>,
21}
22
23impl ApiKeyCredential {
24 pub fn new(key: impl Into<String>) -> Self {
25 Self {
26 kind: "api_key".to_string(),
27 key: Some(key.into()),
28 env: None,
29 }
30 }
31}
32
33#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
34pub struct OAuthCredential {
35 #[serde(rename = "type")]
36 pub kind: String,
37 pub access: String,
38 pub refresh: String,
39 pub expires: i64,
40 #[serde(default, skip_serializing_if = "Option::is_none")]
41 pub account_id: Option<String>,
42 #[serde(default, skip_serializing_if = "Option::is_none")]
43 pub enterprise_url: Option<String>,
44 #[serde(default, skip_serializing_if = "Option::is_none")]
45 pub available_model_ids: Option<Vec<String>>,
46}
47
48#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
49#[serde(tag = "type")]
50pub enum Credential {
51 #[serde(rename = "api_key")]
52 ApiKey(ApiKeyCredential),
53 #[serde(rename = "oauth")]
54 OAuth(OAuthCredential),
55}
56
57#[async_trait::async_trait]
58pub trait CredentialStore: Send + Sync {
59 async fn read(&self, provider_id: &str) -> Option<Credential>;
60 async fn modify(
61 &self,
62 provider_id: &str,
63 f: Box<dyn FnOnce(Option<Credential>) -> Pin<Box<dyn Future<Output = Option<Credential>> + Send>> + Send>,
64 ) -> Option<Credential>;
65 async fn delete(&self, provider_id: &str);
66}
67
68#[async_trait::async_trait]
69pub trait AuthContext: Send + Sync {
70 async fn env(&self, name: &str) -> Option<String>;
71 async fn file_exists(&self, path: &str) -> bool;
72}
73
74#[derive(Debug, Clone)]
75pub struct AuthResult {
76 pub auth: ModelAuth,
77 pub env: Option<ProviderEnv>,
78 pub source: Option<String>,
79}
80
81#[derive(Debug, Clone)]
82pub enum AuthPrompt {
83 Text {
84 message: String,
85 placeholder: Option<String>,
86 },
87 Secret {
88 message: String,
89 placeholder: Option<String>,
90 },
91 Select {
92 message: String,
93 options: Vec<AuthSelectOption>,
94 },
95 ManualCode {
96 message: String,
97 placeholder: Option<String>,
98 },
99}
100
101#[derive(Debug, Clone)]
102pub struct AuthSelectOption {
103 pub id: String,
104 pub label: String,
105 pub description: Option<String>,
106}
107
108#[derive(Debug, Clone)]
109pub enum AuthEvent {
110 AuthUrl {
111 url: String,
112 instructions: Option<String>,
113 },
114 DeviceCode {
115 user_code: String,
116 verification_uri: String,
117 interval_seconds: Option<u32>,
118 expires_in_seconds: Option<u32>,
119 },
120 Progress {
121 message: String,
122 },
123}
124
125#[async_trait::async_trait]
126pub trait AuthLoginCallbacks: Send + Sync {
127 async fn prompt(&self, prompt: AuthPrompt) -> anyhow::Result<String>;
128 fn notify(&self, event: AuthEvent);
129}
130
131pub type ApiKeyResolveFn =
132 Arc<dyn Fn(AuthResolveInput) -> Pin<Box<dyn Future<Output = Option<AuthResult>> + Send>> + Send + Sync>;
133pub type ApiKeyLoginFn = Arc<
134 dyn Fn(Arc<dyn AuthLoginCallbacks>) -> Pin<Box<dyn Future<Output = anyhow::Result<ApiKeyCredential>> + Send>>
135 + Send
136 + Sync,
137>;
138
139pub struct AuthResolveInput {
140 pub model: AuthModel,
141 pub ctx: Arc<dyn AuthContext>,
142 pub credential: Option<ApiKeyCredential>,
143}
144
145#[derive(Clone)]
146pub enum AuthModel {
147 Chat(Model),
148 Images(ImagesModel),
149}
150
151#[derive(Clone)]
152pub struct ApiKeyAuth {
153 pub name: String,
154 pub resolve: ApiKeyResolveFn,
155 pub login: Option<ApiKeyLoginFn>,
156}
157
158pub type OAuthLoginFn = Arc<
159 dyn Fn(Arc<dyn AuthLoginCallbacks>) -> Pin<Box<dyn Future<Output = anyhow::Result<OAuthCredential>> + Send>>
160 + Send
161 + Sync,
162>;
163pub type OAuthRefreshFn =
164 Arc<dyn Fn(OAuthCredential) -> Pin<Box<dyn Future<Output = anyhow::Result<OAuthCredential>> + Send>> + Send + Sync>;
165pub type OAuthToAuthFn =
166 Arc<dyn Fn(OAuthCredential) -> Pin<Box<dyn Future<Output = anyhow::Result<ModelAuth>> + Send>> + Send + Sync>;
167
168#[derive(Clone)]
169pub struct OAuthAuth {
170 pub name: String,
171 pub login: OAuthLoginFn,
172 pub refresh: OAuthRefreshFn,
173 pub to_auth: OAuthToAuthFn,
174}
175
176#[derive(Clone, Default)]
177pub struct ProviderAuth {
178 pub api_key: Option<ApiKeyAuth>,
179 pub oauth: Option<OAuthAuth>,
180}