Skip to main content

rs_adk/auth/
config.rs

1//! Auth configuration and tool argument types.
2
3use serde::{Deserialize, Serialize};
4
5use super::credential::AuthCredential;
6use super::schemes::AuthScheme;
7
8/// Configuration binding a security scheme to credentials.
9#[derive(Debug, Clone, Serialize, Deserialize)]
10pub struct AuthConfig {
11    /// The security scheme that describes how authentication is performed.
12    pub auth_scheme: AuthScheme,
13    /// The raw (user-provided) credential before any exchange.
14    pub raw_auth_credential: Option<AuthCredential>,
15    /// The credential after an exchange (e.g. code-for-token swap).
16    pub exchanged_auth_credential: Option<AuthCredential>,
17    /// Key used to look up/store the credential in a credential service.
18    pub credential_key: String,
19}
20
21/// Arguments passed to a tool when authentication is required.
22#[derive(Debug, Clone, Serialize, Deserialize)]
23pub struct AuthToolArguments {
24    /// The function call ID this auth request is associated with.
25    pub function_call_id: String,
26    /// The auth configuration describing what credential is needed.
27    pub auth_config: AuthConfig,
28}
29
30#[cfg(test)]
31mod tests {
32    use super::*;
33    use crate::auth::credential::{AuthCredentialType, OAuth2Auth};
34    use crate::auth::schemes::OAuthGrantType;
35
36    #[test]
37    fn auth_config_roundtrip() {
38        let config = AuthConfig {
39            auth_scheme: AuthScheme::OAuth2 {
40                grant_type: Some(OAuthGrantType::AuthorizationCode),
41                authorization_url: Some("https://example.com/authorize".into()),
42                token_url: Some("https://example.com/token".into()),
43                scopes: None,
44            },
45            raw_auth_credential: Some(AuthCredential {
46                auth_type: AuthCredentialType::OAuth2,
47                resource_ref: None,
48                api_key: None,
49                http: None,
50                oauth2: Some(OAuth2Auth {
51                    client_id: Some("client-123".into()),
52                    client_secret: Some("secret-456".into()),
53                    auth_uri: None,
54                    token_uri: None,
55                    redirect_uri: None,
56                    auth_code: Some("auth-code-789".into()),
57                    access_token: None,
58                    refresh_token: None,
59                    expires_at: None,
60                    scopes: Some(vec!["openid".into()]),
61                    auth_response_uri: None,
62                }),
63                service_account: None,
64            }),
65            exchanged_auth_credential: None,
66            credential_key: "my-oauth-cred".into(),
67        };
68
69        let json = serde_json::to_string_pretty(&config).unwrap();
70        let parsed: AuthConfig = serde_json::from_str(&json).unwrap();
71
72        assert_eq!(parsed.credential_key, "my-oauth-cred");
73        assert!(parsed.raw_auth_credential.is_some());
74        assert!(parsed.exchanged_auth_credential.is_none());
75
76        let raw = parsed.raw_auth_credential.unwrap();
77        assert_eq!(raw.auth_type, AuthCredentialType::OAuth2);
78        let oauth2 = raw.oauth2.unwrap();
79        assert_eq!(oauth2.client_id.as_deref(), Some("client-123"));
80        assert_eq!(oauth2.auth_code.as_deref(), Some("auth-code-789"));
81    }
82
83    #[test]
84    fn auth_tool_arguments_roundtrip() {
85        let args = AuthToolArguments {
86            function_call_id: "fc-001".into(),
87            auth_config: AuthConfig {
88                auth_scheme: AuthScheme::ApiKey {
89                    location: "header".into(),
90                    name: "X-API-Key".into(),
91                },
92                raw_auth_credential: Some(AuthCredential {
93                    auth_type: AuthCredentialType::ApiKey,
94                    resource_ref: None,
95                    api_key: Some("my-api-key".into()),
96                    http: None,
97                    oauth2: None,
98                    service_account: None,
99                }),
100                exchanged_auth_credential: None,
101                credential_key: "api-key-cred".into(),
102            },
103        };
104
105        let json = serde_json::to_string(&args).unwrap();
106        let parsed: AuthToolArguments = serde_json::from_str(&json).unwrap();
107
108        assert_eq!(parsed.function_call_id, "fc-001");
109        assert_eq!(parsed.auth_config.credential_key, "api-key-cred");
110        let raw = parsed.auth_config.raw_auth_credential.unwrap();
111        assert_eq!(raw.api_key.as_deref(), Some("my-api-key"));
112    }
113
114    #[test]
115    fn auth_config_with_exchanged_credential() {
116        let config = AuthConfig {
117            auth_scheme: AuthScheme::Http {
118                scheme: "bearer".into(),
119                bearer_format: Some("JWT".into()),
120            },
121            raw_auth_credential: None,
122            exchanged_auth_credential: Some(AuthCredential {
123                auth_type: AuthCredentialType::OAuth2,
124                resource_ref: None,
125                api_key: None,
126                http: None,
127                oauth2: Some(OAuth2Auth {
128                    client_id: None,
129                    client_secret: None,
130                    auth_uri: None,
131                    token_uri: None,
132                    redirect_uri: None,
133                    auth_code: None,
134                    access_token: Some("ya29.exchanged-token".into()),
135                    refresh_token: Some("1//refresh".into()),
136                    expires_at: Some(1700000000),
137                    scopes: None,
138                    auth_response_uri: None,
139                }),
140                service_account: None,
141            }),
142            credential_key: "exchanged-cred".into(),
143        };
144
145        let json = serde_json::to_string(&config).unwrap();
146        let parsed: AuthConfig = serde_json::from_str(&json).unwrap();
147
148        assert!(parsed.raw_auth_credential.is_none());
149        let exchanged = parsed.exchanged_auth_credential.unwrap();
150        assert_eq!(
151            exchanged.oauth2.as_ref().unwrap().access_token.as_deref(),
152            Some("ya29.exchanged-token")
153        );
154    }
155}