adk-rs 0.6.0

Rust port of the Google Agent Development Kit (ADK).
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99

//! Credential refreshers — extend an expired credential's lifetime.

use async_trait::async_trait;
use std::collections::HashMap;
use std::sync::Arc;

use crate::auth::config::AuthConfig;
use crate::auth::credential::{AuthCredential, AuthCredentialType};
use crate::error::Result;

/// Refresh an expired credential.
#[async_trait]
pub trait CredentialRefresher: Send + Sync + std::fmt::Debug + 'static {
    /// Try to refresh `cred`. Returns the refreshed credential, or
    /// `Ok(None)` if no refresh is possible (e.g. no refresh token).
    async fn refresh(
        &self,
        config: &AuthConfig,
        cred: &AuthCredential,
    ) -> Result<Option<AuthCredential>>;
}

/// Maps credential types to refreshers.
#[derive(Default, Debug)]
pub struct RefresherRegistry {
    by_type: HashMap<AuthCredentialType, Arc<dyn CredentialRefresher>>,
}

impl RefresherRegistry {
    /// Empty registry.
    #[must_use]
    pub fn new() -> Self {
        Self::default()
    }

    /// Registry with the standard `OAuth2Refresher` for OAuth2 + OIDC.
    #[must_use]
    pub fn with_defaults() -> Self {
        let mut r = Self::new();
        r.register(AuthCredentialType::OAuth2, Arc::new(OAuth2Refresher));
        r.register(AuthCredentialType::OpenIdConnect, Arc::new(OAuth2Refresher));
        r
    }

    /// Replace the refresher for the given type.
    pub fn register(&mut self, ty: AuthCredentialType, refresher: Arc<dyn CredentialRefresher>) {
        self.by_type.insert(ty, refresher);
    }

    /// Look up a refresher.
    #[must_use]
    pub fn get(&self, ty: AuthCredentialType) -> Option<Arc<dyn CredentialRefresher>> {
        self.by_type.get(&ty).cloned()
    }
}

/// Refreshes an OAuth2 credential via its refresh token. Uses the standard
/// authorization-code flow's token endpoint (or the explicitly-set
/// `OAuth2Auth.token_uri`).
#[derive(Debug, Default)]
pub struct OAuth2Refresher;

#[async_trait]
impl CredentialRefresher for OAuth2Refresher {
    async fn refresh(
        &self,
        config: &AuthConfig,
        cred: &AuthCredential,
    ) -> Result<Option<AuthCredential>> {
        let Some(oauth2) = cred.oauth2.as_ref() else {
            return Ok(None);
        };
        let Some(refresh_token) = oauth2.refresh_token.as_deref() else {
            return Ok(None);
        };
        let mut populated = oauth2.clone();
        if let crate::auth::scheme::AuthScheme::OAuth2 { flows, .. } = &config.auth_scheme {
            if let Some(ac) = flows.authorization_code.as_ref() {
                if populated.auth_uri.is_none() {
                    populated.auth_uri.clone_from(&ac.authorization_url);
                }
                if populated.token_uri.is_none() {
                    populated.token_uri = Some(ac.token_url.clone());
                }
                if populated.redirect_uri.is_none() {
                    // PKCE refresh doesn't require a redirect_uri at the
                    // protocol level, but oauth2 crate's BasicClient does, so
                    // fall back to a no-op localhost value.
                    populated.redirect_uri = Some("http://localhost/__adk_refresh__".into());
                }
            }
        }
        let handler = crate::auth::handler::AuthHandler::from_oauth2(&populated)?;
        let tok = handler.refresh(refresh_token).await?;
        let mut new = oauth2.clone();
        tok.apply_to(&mut new);
        Ok(Some(AuthCredential::oauth2(new)))
    }
}