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