Skip to main content

yapi_ai/auth/
mod.rs

1//! Provider authentication: the `auth.json` credential store, OAuth sign-in
2//! and refresh, and the interaction a sign-in uses to talk to the user.
3//!
4//! Ports of `packages/ai/src/auth` in pi `v1.0.0`.
5
6pub mod anthropic;
7pub mod callback;
8pub mod chatgpt;
9pub mod codex;
10pub mod copilot;
11pub mod device;
12pub mod federation;
13pub mod google_adc;
14pub mod kimi;
15pub mod lock;
16pub mod meta;
17pub mod openrouter;
18pub mod pkce;
19pub mod radius;
20pub mod store;
21pub mod xai;
22
23use std::sync::Arc;
24
25use futures_util::future::BoxFuture;
26use serde::Deserialize;
27use tokio::sync::{mpsc, oneshot};
28use tokio_util::sync::CancellationToken;
29use yapi_types::auth::{ApiKeyCredential, OAuthCredential};
30
31pub use store::{CredentialKind, CredentialStore};
32
33/// Why a sign-in, refresh or credential change failed.
34#[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)]
35pub enum AuthError {
36    /// The user or the caller cancelled.
37    #[error("Login cancelled")]
38    Cancelled,
39    /// Anything else, with pi's message.
40    #[error("{0}")]
41    Failed(String),
42}
43
44impl AuthError {
45    pub(crate) fn failed(message: impl Into<String>) -> AuthError {
46        AuthError::Failed(message.into())
47    }
48
49    /// A sign-in fails when the OS random generator does, as Web Crypto
50    /// throws, rather than use a guessable verifier or state.
51    pub fn random(err: std::io::Error) -> AuthError {
52        AuthError::Failed(format!("The OS random generator failed: {err}"))
53    }
54}
55
56/// One choice of a select prompt.
57#[derive(Clone, Debug, PartialEq, Eq, Deserialize)]
58pub struct SelectOption {
59    /// Returned when chosen.
60    pub id: String,
61    /// Shown.
62    pub label: String,
63}
64
65/// A question a sign-in asks, deserialized from pi's `AuthPrompt` JSON.
66#[derive(Clone, Debug, PartialEq, Eq, Deserialize)]
67#[serde(
68    tag = "type",
69    rename_all = "snake_case",
70    rename_all_fields = "camelCase"
71)]
72pub enum AuthPrompt {
73    /// Free text.
74    Text {
75        /// The question.
76        message: String,
77        /// Example input.
78        placeholder: Option<String>,
79    },
80    /// Text that must not be echoed, such as a key.
81    Secret {
82        /// The question.
83        message: String,
84    },
85    /// One of several options. The answer is the option id.
86    Select {
87        /// The question.
88        message: String,
89        /// The choices.
90        options: Vec<SelectOption>,
91    },
92    /// A pasted authorization code or redirect URL. A sign-in may withdraw it
93    /// when the browser callback arrives first.
94    ManualCode {
95        /// The question.
96        message: String,
97        /// Example input.
98        placeholder: Option<String>,
99    },
100}
101
102impl AuthPrompt {
103    /// A [`AuthPrompt::Select`] of `(id, label)` options.
104    pub(crate) fn select(message: impl Into<String>, options: &[(&str, &str)]) -> AuthPrompt {
105        AuthPrompt::Select {
106            message: message.into(),
107            options: options
108                .iter()
109                .map(|(id, label)| SelectOption {
110                    id: (*id).to_owned(),
111                    label: (*label).to_owned(),
112                })
113                .collect(),
114        }
115    }
116}
117
118/// A labelled link in an info event.
119#[derive(Clone, Debug, PartialEq, Eq, Deserialize)]
120pub struct Link {
121    /// Shown.
122    #[serde(default)]
123    pub label: String,
124    /// Target.
125    pub url: String,
126}
127
128/// Something a sign-in tells the user, deserialized from pi's `AuthEvent`
129/// JSON.
130#[derive(Clone, Debug, PartialEq, Deserialize)]
131#[serde(
132    tag = "type",
133    rename_all = "snake_case",
134    rename_all_fields = "camelCase"
135)]
136pub enum AuthEvent {
137    /// Information with optional links.
138    Info {
139        /// Text.
140        message: String,
141        /// Links.
142        #[serde(default)]
143        links: Vec<Link>,
144    },
145    /// A URL to open in the browser.
146    AuthUrl {
147        /// The authorization URL.
148        url: String,
149        /// What to do there.
150        instructions: Option<String>,
151    },
152    /// A device code to enter at a verification page.
153    DeviceCode {
154        /// Code to enter.
155        user_code: String,
156        /// Page to enter it at.
157        verification_uri: String,
158        /// Polling interval.
159        interval_seconds: Option<f64>,
160        /// Lifetime of the code.
161        expires_in_seconds: Option<f64>,
162    },
163    /// Progress text.
164    Progress {
165        /// Text.
166        message: String,
167    },
168}
169
170/// A request from a sign-in to the UI.
171#[derive(Debug)]
172pub enum AuthRequest {
173    /// Ask a question. Dropping `reply` cancels the sign-in, and `cancel` fires
174    /// when the sign-in no longer needs the answer.
175    Prompt {
176        /// The question.
177        prompt: AuthPrompt,
178        /// Where the answer goes.
179        reply: oneshot::Sender<String>,
180        /// Fires when the prompt is withdrawn.
181        cancel: CancellationToken,
182    },
183    /// Show something.
184    Notify(AuthEvent),
185}
186
187/// The channel between a running sign-in and the UI that serves it.
188#[derive(Clone, Debug)]
189pub struct Interaction {
190    requests: mpsc::UnboundedSender<AuthRequest>,
191    cancel: CancellationToken,
192}
193
194impl Interaction {
195    /// An interaction cancelled by `cancel`, with the receiver the UI reads.
196    pub fn new(cancel: CancellationToken) -> (Interaction, mpsc::UnboundedReceiver<AuthRequest>) {
197        let (requests, receiver) = mpsc::unbounded_channel();
198        (Interaction { requests, cancel }, receiver)
199    }
200
201    /// Cancels the whole sign-in.
202    pub fn cancel(&self) -> &CancellationToken {
203        &self.cancel
204    }
205
206    /// Asks the UI and waits for the answer.
207    pub async fn prompt(&self, prompt: AuthPrompt) -> Result<String, AuthError> {
208        self.prompt_until(prompt, self.cancel.child_token()).await
209    }
210
211    /// Asks the UI. `cancel` withdraws the question.
212    pub async fn prompt_until(
213        &self,
214        prompt: AuthPrompt,
215        cancel: CancellationToken,
216    ) -> Result<String, AuthError> {
217        if self.cancel.is_cancelled() {
218            return Err(AuthError::Cancelled);
219        }
220        let (reply, answer) = oneshot::channel();
221        self.requests
222            .send(AuthRequest::Prompt {
223                prompt,
224                reply,
225                cancel: cancel.clone(),
226            })
227            .map_err(|_| AuthError::Cancelled)?;
228        match cancel.run_until_cancelled(answer).await {
229            Some(Ok(answer)) => Ok(answer),
230            _ => Err(AuthError::Cancelled),
231        }
232    }
233
234    /// Tells the UI something.
235    pub fn notify(&self, event: AuthEvent) {
236        let _ = self.requests.send(AuthRequest::Notify(event));
237    }
238
239    /// Fails with [`AuthError::Cancelled`] once cancelled.
240    pub fn check(&self) -> Result<(), AuthError> {
241        if self.cancel.is_cancelled() {
242            Err(AuthError::Cancelled)
243        } else {
244            Ok(())
245        }
246    }
247}
248
249/// Installation details some sign-ins need.
250#[derive(Clone, Debug, Default, PartialEq, Eq)]
251pub struct LoginOptions {
252    /// The installation's device id, a UUID kept in global settings.
253    pub device_id: Option<String>,
254}
255
256/// Credentials derived from an OAuth token for one request.
257#[derive(Clone, Debug, Default, PartialEq, Eq)]
258pub struct OAuthAuth {
259    /// Sent as the API key.
260    pub api_key: Option<String>,
261    /// Headers that carry the token instead. `None` removes a header.
262    pub headers: indexmap::IndexMap<String, Option<String>>,
263    /// Replaces the model's base URL, for providers whose endpoint depends on
264    /// the account.
265    pub base_url: Option<String>,
266}
267
268/// An account sign-in: log in, refresh, and derive request credentials.
269pub trait OAuthProvider: Send + Sync {
270    /// Method name, such as `Anthropic (Claude Pro/Max)`.
271    fn name(&self) -> &str;
272
273    /// Label of the login menu entry, when the provider sets one.
274    fn login_label(&self) -> Option<&str> {
275        None
276    }
277
278    /// Whether the account is a subscription, which the login menu labels.
279    fn is_subscription(&self) -> bool {
280        true
281    }
282
283    /// Runs the sign-in.
284    fn login<'a>(
285        &'a self,
286        interaction: &'a Interaction,
287        options: &'a LoginOptions,
288    ) -> BoxFuture<'a, Result<OAuthCredential, AuthError>>;
289
290    /// Exchanges the refresh token for a new credential.
291    fn refresh<'a>(
292        &'a self,
293        credential: &'a OAuthCredential,
294        cancel: &'a CancellationToken,
295    ) -> BoxFuture<'a, Result<OAuthCredential, AuthError>>;
296
297    /// Request credentials for a valid token.
298    fn to_auth<'a>(
299        &'a self,
300        credential: &'a OAuthCredential,
301    ) -> BoxFuture<'a, Result<OAuthAuth, AuthError>> {
302        Box::pin(async move {
303            Ok(OAuthAuth {
304                api_key: Some(credential.access.clone()),
305                ..OAuthAuth::default()
306            })
307        })
308    }
309}
310
311/// pi-ai's `ApiKeyAuth` of a provider an extension implements.
312pub trait KeyAuth: Send + Sync {
313    /// The method's name, such as "Anthropic API key".
314    fn name(&self) -> &str;
315
316    /// Whether [`KeyAuth::login`] asks for the credential. Without it the
317    /// key is set up outside yapi.
318    fn has_login(&self) -> bool;
319
320    /// Asks for the credential to store through `interaction`.
321    fn login<'a>(
322        &'a self,
323        interaction: &'a Interaction,
324    ) -> BoxFuture<'a, Result<ApiKeyCredential, AuthError>>;
325
326    /// Request credentials from `credential`, whose key is resolved, else
327    /// from the environment, or `None` when the provider is not configured.
328    fn resolve<'a>(
329        &'a self,
330        credential: Option<&'a ApiKeyCredential>,
331    ) -> BoxFuture<'a, Result<Option<crate::key_auth::Resolved>, String>>;
332}
333
334/// The built-in sign-in for a provider id, if yapi implements it.
335pub fn builtin_oauth(provider: &str) -> Option<Arc<dyn OAuthProvider>> {
336    match provider {
337        "anthropic" => Some(Arc::new(anthropic::AnthropicOAuth::default())),
338        "openai-codex" => Some(Arc::new(codex::CodexOAuth::default())),
339        "openai" => Some(Arc::new(chatgpt::ChatGptOAuth::default())),
340        "github-copilot" => Some(Arc::new(copilot::CopilotOAuth::default())),
341        "kimi-coding" => Some(Arc::new(kimi::KimiOAuth::default())),
342        "meta" => Some(Arc::new(meta::MetaOAuth::default())),
343        "xai" => Some(Arc::new(xai::XaiOAuth::default())),
344        "openrouter" => Some(Arc::new(openrouter::OpenRouterOAuth::default())),
345        "radius" => Some(Arc::new(radius::RadiusOAuth::new(
346            "Radius",
347            radius::DEFAULT_GATEWAY,
348        ))),
349        _ => None,
350    }
351}
352
353pub(crate) use yapi_types::time::now_ms;
354
355/// pi's `openBrowser`: opens `target` in the platform browser, best effort and
356/// without a shell. Returns at once, since the opener starts on its own thread.
357pub fn open_browser(target: &str) {
358    // `open` would read a target such as `-a…` as an option.
359    if target.starts_with('-') {
360        return;
361    }
362    let program = if cfg!(target_os = "macos") {
363        "open"
364    } else {
365        "xdg-open"
366    };
367    let mut command = std::process::Command::new(program);
368    command.arg(target);
369    launch(command);
370}
371
372/// Starts `command` without terminal input or output, best effort. Returns at
373/// once, since the command runs on its own thread.
374pub fn launch(mut command: std::process::Command) {
375    command
376        .stdin(std::process::Stdio::null())
377        .stdout(std::process::Stdio::null())
378        .stderr(std::process::Stdio::null());
379    // Waiting reaps the command when it exits.
380    std::thread::spawn(move || command.spawn().and_then(|mut child| child.wait()));
381}
382
383/// pi's `OAUTH_CALLBACK_HOST`: `PI_OAUTH_CALLBACK_HOST`, else `127.0.0.1`.
384pub(crate) fn callback_host() -> String {
385    std::env::var("PI_OAUTH_CALLBACK_HOST")
386        .ok()
387        .filter(|host| !host.trim().is_empty())
388        .unwrap_or_else(|| "127.0.0.1".into())
389}
390
391/// Splits pasted sign-in input into code and state: a redirect URL, `code#state`,
392/// a `code=` query, or a bare code. Port of `parseAuthorizationInput`.
393pub(crate) fn parse_authorization_input(input: &str) -> (Option<String>, Option<String>) {
394    let value = input.trim();
395    if value.is_empty() {
396        return (None, None);
397    }
398    if let Ok(url) = url::Url::parse(value) {
399        return (
400            callback::query(&url, "code"),
401            callback::query(&url, "state"),
402        );
403    }
404    if let Some((code, state)) = value.split_once('#') {
405        let state = state.split('#').next().unwrap_or_default();
406        return (Some(code.to_owned()), Some(state.to_owned()));
407    }
408    if value.contains("code=") {
409        let pairs: Vec<(String, String)> = url::form_urlencoded::parse(value.as_bytes())
410            .into_owned()
411            .collect();
412        let get = |name: &str| {
413            pairs
414                .iter()
415                .find(|(key, _)| key == name)
416                .map(|(_, value)| value.clone())
417        };
418        return (get("code"), get("state"));
419    }
420    (Some(value.to_owned()), None)
421}
422
423/// An `application/x-www-form-urlencoded` body.
424pub(crate) fn form(pairs: &[(&str, &str)]) -> String {
425    url::form_urlencoded::Serializer::new(String::new())
426        .extend_pairs(pairs)
427        .finish()
428}
429
430/// A POST of `fields` as a form that accepts JSON, without a timeout.
431pub fn post_form(url: impl reqwest::IntoUrl, fields: &[(&str, &str)]) -> reqwest::RequestBuilder {
432    crate::http::client()
433        .post(url)
434        .header("Content-Type", "application/x-www-form-urlencoded")
435        .header("Accept", "application/json")
436        .body(form(fields))
437}
438
439/// `Date.now() + expires_in * 1000 - margin`, in milliseconds.
440pub(crate) fn expiry(expires_in: f64, margin_ms: f64) -> u64 {
441    let value = now_ms() as f64 + expires_in * 1000.0 - margin_ms;
442    if value.is_finite() && value > 0.0 {
443        value as u64
444    } else {
445        0
446    }
447}
448
449/// A JSON object of string fields, in order.
450pub(crate) fn object(pairs: &[(&str, &str)]) -> serde_json::Value {
451    pairs.iter().copied().collect()
452}
453
454/// The body parsed as JSON, or `null` when it is not JSON.
455pub(crate) async fn json_body(response: reqwest::Response) -> serde_json::Value {
456    let bytes = response.bytes().await.unwrap_or_default();
457    serde_json::from_slice(&bytes).unwrap_or(serde_json::Value::Null)
458}
459
460/// The text of an HTTP error for messages: the body, else the status text.
461pub(crate) async fn error_text(response: reqwest::Response) -> String {
462    let status = response.status();
463    let text = response.text().await.unwrap_or_default();
464    if text.is_empty() {
465        status.canonical_reason().unwrap_or_default().to_owned()
466    } else {
467        text
468    }
469}
470
471/// Sends `request`, failing with [`AuthError::Cancelled`] on cancellation.
472pub(crate) async fn send(
473    request: reqwest::RequestBuilder,
474    cancel: &CancellationToken,
475) -> Result<reqwest::Response, AuthError> {
476    cancel
477        .run_until_cancelled(request.send())
478        .await
479        .ok_or(AuthError::Cancelled)?
480        .map_err(|err| AuthError::Failed(network_message(&err)))
481}
482
483/// A network error with its causes, like Node's `fetch failed` chain.
484pub(crate) fn network_message(err: &dyn std::error::Error) -> String {
485    let mut message = err.to_string();
486    let mut source = err.source();
487    while let Some(cause) = source {
488        message.push_str(": ");
489        message.push_str(&cause.to_string());
490        source = cause.source();
491    }
492    message
493}
494
495#[cfg(test)]
496mod tests {
497    use super::*;
498
499    /// pi's `AuthPrompt` and `AuthEvent` JSON, as extension sign-ins send it.
500    #[test]
501    fn reads_prompts_and_events() {
502        let prompt = |value| serde_json::from_value::<AuthPrompt>(value).unwrap();
503        assert_eq!(
504            prompt(
505                serde_json::json!({"type": "select", "message": "Team", "options": [{"id": "a", "label": "A", "description": "x"}]})
506            ),
507            AuthPrompt::select("Team", &[("a", "A")])
508        );
509        assert_eq!(
510            prompt(serde_json::json!({"type": "manual_code", "message": "Code"})),
511            AuthPrompt::ManualCode {
512                message: "Code".into(),
513                placeholder: None
514            }
515        );
516        let event = serde_json::from_value::<AuthEvent>(serde_json::json!({
517            "type": "device_code", "userCode": "U", "verificationUri": "https://v", "intervalSeconds": 5,
518        }))
519        .unwrap();
520        assert_eq!(
521            event,
522            AuthEvent::DeviceCode {
523                user_code: "U".into(),
524                verification_uri: "https://v".into(),
525                interval_seconds: Some(5.0),
526                expires_in_seconds: None,
527            }
528        );
529        assert!(
530            serde_json::from_value::<AuthEvent>(serde_json::json!({"type": "unknown"})).is_err()
531        );
532    }
533
534    #[test]
535    fn parses_pasted_input() {
536        let parse = |input| parse_authorization_input(input);
537        assert_eq!(
538            parse("http://localhost:1455/auth/callback?code=c1&state=s1"),
539            (Some("c1".into()), Some("s1".into()))
540        );
541        assert_eq!(parse("c2#s2"), (Some("c2".into()), Some("s2".into())));
542        assert_eq!(
543            parse("code=c3&state=s3"),
544            (Some("c3".into()), Some("s3".into()))
545        );
546        assert_eq!(parse(" c4 "), (Some("c4".into()), None));
547        assert_eq!(parse(""), (None, None));
548    }
549
550    #[tokio::test]
551    async fn prompts_answer_and_cancel() {
552        let cancel = CancellationToken::new();
553        let (interaction, mut requests) = Interaction::new(cancel.clone());
554        let asking = tokio::spawn({
555            let interaction = interaction.clone();
556            async move {
557                interaction
558                    .prompt(AuthPrompt::Secret {
559                        message: "Enter key".into(),
560                    })
561                    .await
562            }
563        });
564        let Some(AuthRequest::Prompt { reply, .. }) = requests.recv().await else {
565            panic!("expected a prompt");
566        };
567        reply.send("secret".into()).unwrap();
568        assert_eq!(asking.await.unwrap().as_deref(), Ok("secret"));
569        let dropped = tokio::spawn({
570            let interaction = interaction.clone();
571            async move {
572                interaction
573                    .prompt(AuthPrompt::Text {
574                        message: "?".into(),
575                        placeholder: None,
576                    })
577                    .await
578            }
579        });
580        drop(requests.recv().await);
581        assert_eq!(dropped.await.unwrap(), Err(AuthError::Cancelled));
582        cancel.cancel();
583        assert_eq!(interaction.check(), Err(AuthError::Cancelled));
584    }
585}