rust_genai/
client.rs

1//! Client configuration and transport layer.
2
3use std::collections::HashMap;
4use std::path::{Path, PathBuf};
5use std::sync::Arc;
6use std::time::Duration;
7
8use reqwest::header::{HeaderMap, HeaderName, HeaderValue, AUTHORIZATION};
9use reqwest::{Client as HttpClient, Proxy};
10use tokio::sync::OnceCell;
11
12use crate::auth::OAuthTokenProvider;
13use crate::error::{Error, Result};
14use google_cloud_auth::credentials::{
15    Builder as AuthBuilder, CacheableResource, Credentials as GoogleCredentials,
16};
17use http::Extensions;
18
19/// Gemini 客户端。
20#[derive(Clone)]
21pub struct Client {
22    inner: Arc<ClientInner>,
23}
24
25pub(crate) struct ClientInner {
26    pub http: HttpClient,
27    pub config: ClientConfig,
28    pub api_client: ApiClient,
29    pub(crate) auth_provider: Option<AuthProvider>,
30}
31
32/// 客户端配置。
33#[derive(Debug, Clone)]
34pub struct ClientConfig {
35    /// API 密钥（Gemini API）。
36    pub api_key: Option<String>,
37    /// 后端选择。
38    pub backend: Backend,
39    /// Vertex AI 配置。
40    pub vertex_config: Option<VertexConfig>,
41    /// HTTP 配置。
42    pub http_options: HttpOptions,
43    /// 认证信息。
44    pub credentials: Credentials,
45    /// OAuth scopes（服务账号/ADC 使用）。
46    pub auth_scopes: Vec<String>,
47}
48
49/// 后端选择。
50#[derive(Debug, Clone, Copy, PartialEq, Eq)]
51pub enum Backend {
52    GeminiApi,
53    VertexAi,
54}
55
56/// 认证方式。
57#[derive(Debug, Clone)]
58pub enum Credentials {
59    /// API Key（Gemini API）。
60    ApiKey(String),
61    /// OAuth 用户凭据。
62    OAuth {
63        client_secret_path: PathBuf,
64        token_cache_path: Option<PathBuf>,
65    },
66    /// Application Default Credentials (ADC)。
67    ApplicationDefault,
68}
69
70/// Vertex AI 配置。
71#[derive(Debug, Clone)]
72pub struct VertexConfig {
73    pub project: String,
74    pub location: String,
75    pub credentials: Option<VertexCredentials>,
76}
77
78/// Vertex AI 认证占位。
79#[derive(Debug, Clone)]
80pub struct VertexCredentials {
81    pub access_token: Option<String>,
82}
83
84/// HTTP 配置。
85#[derive(Debug, Clone, Default)]
86pub struct HttpOptions {
87    pub timeout: Option<u64>,
88    pub proxy: Option<String>,
89    pub headers: HashMap<String, String>,
90    pub base_url: Option<String>,
91    pub api_version: Option<String>,
92}
93
94impl Client {
95    /// 创建新客户端（Gemini API）。
96    ///
97    /// # Errors
98    /// 当配置无效或构建客户端失败时返回错误。
99    pub fn new(api_key: impl Into<String>) -> Result<Self> {
100        Self::builder()
101            .api_key(api_key)
102            .backend(Backend::GeminiApi)
103            .build()
104    }
105
106    /// 从环境变量创建客户端。
107    ///
108    /// # Errors
109    /// 当环境变量缺失或构建客户端失败时返回错误。
110    pub fn from_env() -> Result<Self> {
111        let api_key = std::env::var("GEMINI_API_KEY")
112            .or_else(|_| std::env::var("GOOGLE_API_KEY"))
113            .map_err(|_| Error::InvalidConfig {
114                message: "GEMINI_API_KEY or GOOGLE_API_KEY not found".into(),
115            })?;
116        let mut builder = Self::builder().api_key(api_key);
117        if let Ok(base_url) =
118            std::env::var("GENAI_BASE_URL").or_else(|_| std::env::var("GEMINI_BASE_URL"))
119        {
120            if !base_url.trim().is_empty() {
121                builder = builder.base_url(base_url);
122            }
123        }
124        if let Ok(api_version) = std::env::var("GENAI_API_VERSION") {
125            if !api_version.trim().is_empty() {
126                builder = builder.api_version(api_version);
127            }
128        }
129        builder.build()
130    }
131
132    /// 创建 Vertex AI 客户端。
133    ///
134    /// # Errors
135    /// 当配置无效或构建客户端失败时返回错误。
136    pub fn new_vertex(project: impl Into<String>, location: impl Into<String>) -> Result<Self> {
137        Self::builder()
138            .backend(Backend::VertexAi)
139            .vertex_project(project)
140            .vertex_location(location)
141            .build()
142    }
143
144    /// 使用 OAuth 凭据创建客户端（默认读取 token.json）。
145    ///
146    /// # Errors
147    /// 当凭据路径无效或构建客户端失败时返回错误。
148    pub fn with_oauth(client_secret_path: impl AsRef<Path>) -> Result<Self> {
149        Self::builder()
150            .credentials(Credentials::OAuth {
151                client_secret_path: client_secret_path.as_ref().to_path_buf(),
152                token_cache_path: None,
153            })
154            .build()
155    }
156
157    /// 使用 Application Default Credentials 创建客户端。
158    ///
159    /// # Errors
160    /// 当构建客户端失败时返回错误。
161    pub fn with_adc() -> Result<Self> {
162        Self::builder()
163            .credentials(Credentials::ApplicationDefault)
164            .build()
165    }
166
167    /// 创建 Builder。
168    #[must_use]
169    pub fn builder() -> ClientBuilder {
170        ClientBuilder::default()
171    }
172
173    /// 访问 Models API。
174    #[must_use]
175    pub fn models(&self) -> crate::models::Models {
176        crate::models::Models::new(self.inner.clone())
177    }
178
179    /// 访问 Chats API。
180    #[must_use]
181    pub fn chats(&self) -> crate::chats::Chats {
182        crate::chats::Chats::new(self.inner.clone())
183    }
184
185    /// 访问 Files API。
186    #[must_use]
187    pub fn files(&self) -> crate::files::Files {
188        crate::files::Files::new(self.inner.clone())
189    }
190
191    /// 访问 `FileSearchStores` API。
192    #[must_use]
193    pub fn file_search_stores(&self) -> crate::file_search_stores::FileSearchStores {
194        crate::file_search_stores::FileSearchStores::new(self.inner.clone())
195    }
196
197    /// 访问 Documents API。
198    #[must_use]
199    pub fn documents(&self) -> crate::documents::Documents {
200        crate::documents::Documents::new(self.inner.clone())
201    }
202
203    /// 访问 Live API。
204    #[must_use]
205    pub fn live(&self) -> crate::live::Live {
206        crate::live::Live::new(self.inner.clone())
207    }
208
209    /// 访问 Live Music API。
210    #[must_use]
211    pub fn live_music(&self) -> crate::live_music::LiveMusic {
212        crate::live_music::LiveMusic::new(self.inner.clone())
213    }
214
215    /// 访问 Caches API。
216    #[must_use]
217    pub fn caches(&self) -> crate::caches::Caches {
218        crate::caches::Caches::new(self.inner.clone())
219    }
220
221    /// 访问 Batches API。
222    #[must_use]
223    pub fn batches(&self) -> crate::batches::Batches {
224        crate::batches::Batches::new(self.inner.clone())
225    }
226
227    /// 访问 Tunings API。
228    #[must_use]
229    pub fn tunings(&self) -> crate::tunings::Tunings {
230        crate::tunings::Tunings::new(self.inner.clone())
231    }
232
233    /// 访问 Operations API。
234    #[must_use]
235    pub fn operations(&self) -> crate::operations::Operations {
236        crate::operations::Operations::new(self.inner.clone())
237    }
238
239    /// 访问 `AuthTokens` API（Ephemeral Tokens）。
240    #[must_use]
241    pub fn auth_tokens(&self) -> crate::tokens::AuthTokens {
242        crate::tokens::AuthTokens::new(self.inner.clone())
243    }
244
245    /// 访问 Interactions API。
246    #[must_use]
247    pub fn interactions(&self) -> crate::interactions::Interactions {
248        crate::interactions::Interactions::new(self.inner.clone())
249    }
250
251    /// 访问 Deep Research。
252    #[must_use]
253    pub fn deep_research(&self) -> crate::deep_research::DeepResearch {
254        crate::deep_research::DeepResearch::new(self.inner.clone())
255    }
256}
257
258/// 客户端 Builder。
259#[derive(Default)]
260pub struct ClientBuilder {
261    api_key: Option<String>,
262    credentials: Option<Credentials>,
263    backend: Option<Backend>,
264    vertex_project: Option<String>,
265    vertex_location: Option<String>,
266    http_options: HttpOptions,
267    auth_scopes: Option<Vec<String>>,
268}
269
270impl ClientBuilder {
271    /// 设置 API Key（Gemini API）。
272    #[must_use]
273    pub fn api_key(mut self, key: impl Into<String>) -> Self {
274        self.api_key = Some(key.into());
275        self
276    }
277
278    /// 设置认证方式（OAuth/ADC/API Key）。
279    #[must_use]
280    pub fn credentials(mut self, credentials: Credentials) -> Self {
281        self.credentials = Some(credentials);
282        self
283    }
284
285    /// 设置后端（Gemini API 或 Vertex AI）。
286    #[must_use]
287    pub const fn backend(mut self, backend: Backend) -> Self {
288        self.backend = Some(backend);
289        self
290    }
291
292    /// 设置 Vertex AI 项目 ID。
293    #[must_use]
294    pub fn vertex_project(mut self, project: impl Into<String>) -> Self {
295        self.vertex_project = Some(project.into());
296        self
297    }
298
299    /// 设置 Vertex AI 区域。
300    #[must_use]
301    pub fn vertex_location(mut self, location: impl Into<String>) -> Self {
302        self.vertex_location = Some(location.into());
303        self
304    }
305
306    /// 设置请求超时（秒）。
307    #[must_use]
308    pub const fn timeout(mut self, secs: u64) -> Self {
309        self.http_options.timeout = Some(secs);
310        self
311    }
312
313    /// 设置代理。
314    #[must_use]
315    pub fn proxy(mut self, url: impl Into<String>) -> Self {
316        self.http_options.proxy = Some(url.into());
317        self
318    }
319
320    /// 增加默认 HTTP 头。
321    #[must_use]
322    pub fn header(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
323        self.http_options.headers.insert(key.into(), value.into());
324        self
325    }
326
327    /// 设置自定义基础 URL。
328    #[must_use]
329    pub fn base_url(mut self, base_url: impl Into<String>) -> Self {
330        self.http_options.base_url = Some(base_url.into());
331        self
332    }
333
334    /// 设置 API 版本。
335    #[must_use]
336    pub fn api_version(mut self, api_version: impl Into<String>) -> Self {
337        self.http_options.api_version = Some(api_version.into());
338        self
339    }
340
341    /// 设置 OAuth scopes。
342    #[must_use]
343    pub fn auth_scopes(mut self, scopes: Vec<String>) -> Self {
344        self.auth_scopes = Some(scopes);
345        self
346    }
347
348    /// 构建客户端。
349    ///
350    /// # Errors
351    /// 当配置不完整、参数无效或构建 HTTP 客户端失败时返回错误。
352    pub fn build(self) -> Result<Client> {
353        let Self {
354            api_key,
355            credentials,
356            backend,
357            vertex_project,
358            vertex_location,
359            http_options,
360            auth_scopes,
361        } = self;
362
363        let backend = Self::resolve_backend(
364            backend,
365            vertex_project.as_deref(),
366            vertex_location.as_deref(),
367        );
368        Self::validate_vertex_config(
369            backend,
370            vertex_project.as_deref(),
371            vertex_location.as_deref(),
372        )?;
373        let credentials = Self::resolve_credentials(backend, api_key.as_deref(), credentials)?;
374        let headers = Self::build_headers(&http_options, backend, &credentials)?;
375        let http = Self::build_http_client(&http_options, headers)?;
376
377        let auth_scopes = auth_scopes.unwrap_or_else(|| default_auth_scopes(backend));
378        let api_key = match &credentials {
379            Credentials::ApiKey(key) => Some(key.clone()),
380            _ => None,
381        };
382        let vertex_config = Self::build_vertex_config(backend, vertex_project, vertex_location)?;
383        let config = ClientConfig {
384            api_key,
385            backend,
386            vertex_config,
387            http_options,
388            credentials: credentials.clone(),
389            auth_scopes,
390        };
391
392        let auth_provider = build_auth_provider(&credentials)?;
393        let api_client = ApiClient::new(&config);
394
395        Ok(Client {
396            inner: Arc::new(ClientInner {
397                http,
398                config,
399                api_client,
400                auth_provider,
401            }),
402        })
403    }
404
405    fn resolve_backend(
406        backend: Option<Backend>,
407        vertex_project: Option<&str>,
408        vertex_location: Option<&str>,
409    ) -> Backend {
410        backend.unwrap_or_else(|| {
411            if vertex_project.is_some() || vertex_location.is_some() {
412                Backend::VertexAi
413            } else {
414                Backend::GeminiApi
415            }
416        })
417    }
418
419    fn validate_vertex_config(
420        backend: Backend,
421        vertex_project: Option<&str>,
422        vertex_location: Option<&str>,
423    ) -> Result<()> {
424        if backend == Backend::VertexAi && (vertex_project.is_none() || vertex_location.is_none()) {
425            return Err(Error::InvalidConfig {
426                message: "Project and location required for Vertex AI".into(),
427            });
428        }
429        Ok(())
430    }
431
432    fn resolve_credentials(
433        backend: Backend,
434        api_key: Option<&str>,
435        credentials: Option<Credentials>,
436    ) -> Result<Credentials> {
437        if credentials.is_some()
438            && api_key.is_some()
439            && !matches!(credentials, Some(Credentials::ApiKey(_)))
440        {
441            return Err(Error::InvalidConfig {
442                message: "API key cannot be combined with OAuth/ADC credentials".into(),
443            });
444        }
445
446        let credentials = match credentials {
447            Some(credentials) => credentials,
448            None => {
449                if let Some(api_key) = api_key {
450                    Credentials::ApiKey(api_key.to_string())
451                } else if backend == Backend::VertexAi {
452                    Credentials::ApplicationDefault
453                } else {
454                    return Err(Error::InvalidConfig {
455                        message: "API key or OAuth credentials required for Gemini API".into(),
456                    });
457                }
458            }
459        };
460
461        if backend == Backend::VertexAi && matches!(credentials, Credentials::ApiKey(_)) {
462            return Err(Error::InvalidConfig {
463                message: "Vertex AI does not support API key authentication".into(),
464            });
465        }
466
467        Ok(credentials)
468    }
469
470    fn build_headers(
471        http_options: &HttpOptions,
472        backend: Backend,
473        credentials: &Credentials,
474    ) -> Result<HeaderMap> {
475        let mut headers = HeaderMap::new();
476        for (key, value) in &http_options.headers {
477            let name =
478                HeaderName::from_bytes(key.as_bytes()).map_err(|_| Error::InvalidConfig {
479                    message: format!("Invalid header name: {key}"),
480                })?;
481            let value = HeaderValue::from_str(value).map_err(|_| Error::InvalidConfig {
482                message: format!("Invalid header value for {key}"),
483            })?;
484            headers.insert(name, value);
485        }
486
487        if backend == Backend::GeminiApi {
488            let api_key = match credentials {
489                Credentials::ApiKey(key) => key.as_str(),
490                _ => "",
491            };
492            let header_name = HeaderName::from_static("x-goog-api-key");
493            if !api_key.is_empty() && !headers.contains_key(&header_name) {
494                let mut header_value =
495                    HeaderValue::from_str(api_key).map_err(|_| Error::InvalidConfig {
496                        message: "Invalid API key value".into(),
497                    })?;
498                header_value.set_sensitive(true);
499                headers.insert(header_name, header_value);
500            }
501        }
502
503        Ok(headers)
504    }
505
506    fn build_http_client(http_options: &HttpOptions, headers: HeaderMap) -> Result<HttpClient> {
507        let mut http_builder = HttpClient::builder();
508        if let Some(timeout) = http_options.timeout {
509            http_builder = http_builder.timeout(Duration::from_secs(timeout));
510        }
511
512        if let Some(proxy_url) = &http_options.proxy {
513            let proxy = Proxy::all(proxy_url).map_err(|e| Error::InvalidConfig {
514                message: format!("Invalid proxy: {e}"),
515            })?;
516            http_builder = http_builder.proxy(proxy);
517        }
518
519        if !headers.is_empty() {
520            http_builder = http_builder.default_headers(headers);
521        }
522
523        Ok(http_builder.build()?)
524    }
525
526    fn build_vertex_config(
527        backend: Backend,
528        vertex_project: Option<String>,
529        vertex_location: Option<String>,
530    ) -> Result<Option<VertexConfig>> {
531        if backend != Backend::VertexAi {
532            return Ok(None);
533        }
534        let project = vertex_project.ok_or_else(|| Error::InvalidConfig {
535            message: "Project and location required for Vertex AI".into(),
536        })?;
537        let location = vertex_location.ok_or_else(|| Error::InvalidConfig {
538            message: "Project and location required for Vertex AI".into(),
539        })?;
540        Ok(Some(VertexConfig {
541            project,
542            location,
543            credentials: None,
544        }))
545    }
546}
547
548fn build_auth_provider(credentials: &Credentials) -> Result<Option<AuthProvider>> {
549    match credentials {
550        Credentials::ApiKey(_) => Ok(None),
551        Credentials::OAuth {
552            client_secret_path,
553            token_cache_path,
554        } => Ok(Some(AuthProvider::OAuth(Arc::new(
555            OAuthTokenProvider::from_paths(client_secret_path.clone(), token_cache_path.clone())?,
556        )))),
557        Credentials::ApplicationDefault => Ok(Some(AuthProvider::ApplicationDefault(Arc::new(
558            OnceCell::new(),
559        )))),
560    }
561}
562
563#[derive(Clone)]
564pub(crate) enum AuthProvider {
565    OAuth(Arc<OAuthTokenProvider>),
566    ApplicationDefault(Arc<OnceCell<Arc<GoogleCredentials>>>),
567}
568
569impl AuthProvider {
570    async fn headers(&self, scopes: &[&str]) -> Result<HeaderMap> {
571        match self {
572            Self::OAuth(provider) => {
573                let token = provider.token().await?;
574                let mut header =
575                    HeaderValue::from_str(&format!("Bearer {token}")).map_err(|_| Error::Auth {
576                        message: "Invalid OAuth access token".into(),
577                    })?;
578                header.set_sensitive(true);
579                let mut headers = HeaderMap::new();
580                headers.insert(AUTHORIZATION, header);
581                Ok(headers)
582            }
583            Self::ApplicationDefault(cell) => {
584                let credentials = cell
585                    .get_or_try_init(|| async {
586                        AuthBuilder::default()
587                            .with_scopes(scopes.iter().copied())
588                            .build()
589                            .map(Arc::new)
590                            .map_err(|err| Error::Auth {
591                                message: format!("ADC init failed: {err}"),
592                            })
593                    })
594                    .await?;
595                let headers = credentials
596                    .headers(Extensions::new())
597                    .await
598                    .map_err(|err| Error::Auth {
599                        message: format!("ADC header fetch failed: {err}"),
600                    })?;
601                match headers {
602                    CacheableResource::New { data, .. } => Ok(data),
603                    CacheableResource::NotModified => Err(Error::Auth {
604                        message: "ADC header fetch returned NotModified without cached headers"
605                            .into(),
606                    }),
607                }
608            }
609        }
610    }
611}
612
613impl ClientInner {
614    /// 发送请求并自动注入鉴权头。
615    ///
616    /// # Errors
617    /// 当请求构建、鉴权头获取或网络请求失败时返回错误。
618    pub async fn send(&self, request: reqwest::RequestBuilder) -> Result<reqwest::Response> {
619        let mut request = request.build()?;
620        if let Some(headers) = self.auth_headers().await? {
621            for (name, value) in &headers {
622                if request.headers().contains_key(name) {
623                    continue;
624                }
625                let mut value = value.clone();
626                if name == AUTHORIZATION {
627                    value.set_sensitive(true);
628                }
629                request.headers_mut().insert(name.clone(), value);
630            }
631        }
632        #[cfg(feature = "mcp")]
633        crate::mcp::append_mcp_usage_header(request.headers_mut())?;
634        Ok(self.http.execute(request).await?)
635    }
636
637    async fn auth_headers(&self) -> Result<Option<HeaderMap>> {
638        let Some(provider) = &self.auth_provider else {
639            return Ok(None);
640        };
641
642        let scopes: Vec<&str> = self.config.auth_scopes.iter().map(String::as_str).collect();
643        let headers = provider.headers(&scopes).await?;
644        Ok(Some(headers))
645    }
646}
647
648fn default_auth_scopes(backend: Backend) -> Vec<String> {
649    match backend {
650        Backend::VertexAi => vec!["https://www.googleapis.com/auth/cloud-platform".into()],
651        Backend::GeminiApi => vec![
652            "https://www.googleapis.com/auth/generative-language".into(),
653            "https://www.googleapis.com/auth/generative-language.retriever".into(),
654        ],
655    }
656}
657
658pub(crate) struct ApiClient {
659    pub base_url: String,
660    pub api_version: String,
661}
662
663impl ApiClient {
664    /// 创建 API 客户端配置。
665    pub fn new(config: &ClientConfig) -> Self {
666        let base_url = config.http_options.base_url.as_deref().map_or_else(
667            || match config.backend {
668                Backend::VertexAi => {
669                    let location = config
670                        .vertex_config
671                        .as_ref()
672                        .map_or("", |cfg| cfg.location.as_str());
673                    if location.is_empty() {
674                        "https://aiplatform.googleapis.com/".to_string()
675                    } else {
676                        format!("https://{location}-aiplatform.googleapis.com/")
677                    }
678                }
679                Backend::GeminiApi => "https://generativelanguage.googleapis.com/".to_string(),
680            },
681            normalize_base_url,
682        );
683
684        let api_version =
685            config
686                .http_options
687                .api_version
688                .clone()
689                .unwrap_or_else(|| match config.backend {
690                    Backend::VertexAi => "v1beta1".to_string(),
691                    Backend::GeminiApi => "v1beta".to_string(),
692                });
693
694        Self {
695            base_url,
696            api_version,
697        }
698    }
699}
700
701fn normalize_base_url(base_url: &str) -> String {
702    let mut value = base_url.trim().to_string();
703    if !value.ends_with('/') {
704        value.push('/');
705    }
706    value
707}
708
709#[cfg(test)]
710mod tests {
711    use super::*;
712    use crate::test_support::with_env;
713    use std::path::PathBuf;
714    use tempfile::tempdir;
715
716    #[test]
717    fn test_client_from_api_key() {
718        let client = Client::new("test-api-key").unwrap();
719        assert_eq!(client.inner.config.backend, Backend::GeminiApi);
720    }
721
722    #[test]
723    fn test_client_builder() {
724        let client = Client::builder()
725            .api_key("test-key")
726            .timeout(30)
727            .build()
728            .unwrap();
729        assert!(client.inner.config.api_key.is_some());
730    }
731
732    #[test]
733    fn test_vertex_ai_config() {
734        let client = Client::new_vertex("my-project", "us-central1").unwrap();
735        assert_eq!(client.inner.config.backend, Backend::VertexAi);
736        assert_eq!(
737            client.inner.api_client.base_url,
738            "https://us-central1-aiplatform.googleapis.com/"
739        );
740    }
741
742    #[test]
743    fn test_base_url_normalization() {
744        let client = Client::builder()
745            .api_key("test-key")
746            .base_url("https://example.com")
747            .build()
748            .unwrap();
749        assert_eq!(client.inner.api_client.base_url, "https://example.com/");
750    }
751
752    #[test]
753    fn test_from_env_reads_overrides() {
754        with_env(
755            &[
756                ("GEMINI_API_KEY", Some("env-key")),
757                ("GENAI_BASE_URL", Some("https://env.example.com")),
758                ("GENAI_API_VERSION", Some("v99")),
759                ("GOOGLE_API_KEY", None),
760            ],
761            || {
762                let client = Client::from_env().unwrap();
763                assert_eq!(client.inner.api_client.base_url, "https://env.example.com/");
764                assert_eq!(client.inner.api_client.api_version, "v99");
765            },
766        );
767    }
768
769    #[test]
770    fn test_from_env_ignores_empty_overrides() {
771        with_env(
772            &[
773                ("GEMINI_API_KEY", Some("env-key")),
774                ("GENAI_BASE_URL", Some("   ")),
775                ("GENAI_API_VERSION", Some("")),
776                ("GOOGLE_API_KEY", None),
777            ],
778            || {
779                let client = Client::from_env().unwrap();
780                assert_eq!(
781                    client.inner.api_client.base_url,
782                    "https://generativelanguage.googleapis.com/"
783                );
784                assert_eq!(client.inner.api_client.api_version, "v1beta");
785            },
786        );
787    }
788
789    #[test]
790    fn test_from_env_missing_key_errors() {
791        with_env(
792            &[
793                ("GEMINI_API_KEY", None),
794                ("GOOGLE_API_KEY", None),
795                ("GENAI_BASE_URL", None),
796            ],
797            || {
798                let result = Client::from_env();
799                assert!(result.is_err());
800            },
801        );
802    }
803
804    #[test]
805    fn test_from_env_google_api_key_fallback() {
806        with_env(
807            &[
808                ("GEMINI_API_KEY", None),
809                ("GOOGLE_API_KEY", Some("google-key")),
810            ],
811            || {
812                let client = Client::from_env().unwrap();
813                assert_eq!(client.inner.config.api_key.as_deref(), Some("google-key"));
814            },
815        );
816    }
817
818    #[test]
819    fn test_with_oauth_missing_client_secret_errors() {
820        let dir = tempdir().unwrap();
821        let secret_path = dir.path().join("missing_client_secret.json");
822        let err = Client::with_oauth(&secret_path).err().unwrap();
823        assert!(matches!(err, Error::InvalidConfig { .. }));
824    }
825
826    #[test]
827    fn test_with_adc_builds_client() {
828        let client = Client::with_adc().unwrap();
829        assert!(matches!(
830            client.inner.config.credentials,
831            Credentials::ApplicationDefault
832        ));
833    }
834
835    #[test]
836    fn test_builder_defaults_to_vertex_when_project_set() {
837        let client = Client::builder()
838            .vertex_project("proj")
839            .vertex_location("loc")
840            .build()
841            .unwrap();
842        assert_eq!(client.inner.config.backend, Backend::VertexAi);
843        assert!(matches!(
844            client.inner.config.credentials,
845            Credentials::ApplicationDefault
846        ));
847    }
848
849    #[test]
850    fn test_valid_proxy_is_accepted() {
851        let client = Client::builder()
852            .api_key("test-key")
853            .proxy("http://127.0.0.1:8888")
854            .build();
855        assert!(client.is_ok());
856    }
857
858    #[test]
859    fn test_vertex_requires_project_and_location() {
860        let result = Client::builder().backend(Backend::VertexAi).build();
861        assert!(result.is_err());
862    }
863
864    #[test]
865    fn test_api_key_with_oauth_is_invalid() {
866        let result = Client::builder()
867            .api_key("test-key")
868            .credentials(Credentials::OAuth {
869                client_secret_path: PathBuf::from("client_secret.json"),
870                token_cache_path: None,
871            })
872            .build();
873        assert!(result.is_err());
874    }
875
876    #[test]
877    fn test_missing_api_key_for_gemini_errors() {
878        let result = Client::builder().backend(Backend::GeminiApi).build();
879        assert!(result.is_err());
880    }
881
882    #[test]
883    fn test_invalid_header_name_is_rejected() {
884        let result = Client::builder()
885            .api_key("test-key")
886            .header("bad header", "value")
887            .build();
888        assert!(result.is_err());
889    }
890
891    #[test]
892    fn test_invalid_header_value_is_rejected() {
893        let result = Client::builder()
894            .api_key("test-key")
895            .header("x-test", "bad\nvalue")
896            .build();
897        assert!(result.is_err());
898    }
899
900    #[test]
901    fn test_invalid_api_key_value_is_rejected() {
902        let err = Client::builder().api_key("bad\nkey").build().err().unwrap();
903        assert!(
904            matches!(err, Error::InvalidConfig { message } if message.contains("Invalid API key value"))
905        );
906    }
907
908    #[test]
909    fn test_invalid_proxy_is_rejected() {
910        let result = Client::builder()
911            .api_key("test-key")
912            .proxy("not a url")
913            .build();
914        assert!(result.is_err());
915    }
916
917    #[test]
918    fn test_vertex_api_key_is_rejected() {
919        let result = Client::builder()
920            .backend(Backend::VertexAi)
921            .vertex_project("proj")
922            .vertex_location("loc")
923            .credentials(Credentials::ApiKey("key".into()))
924            .build();
925        assert!(result.is_err());
926    }
927
928    #[test]
929    fn test_default_auth_scopes() {
930        let gemini = default_auth_scopes(Backend::GeminiApi);
931        assert!(gemini.iter().any(|s| s.contains("generative-language")));
932
933        let vertex = default_auth_scopes(Backend::VertexAi);
934        assert!(vertex.iter().any(|s| s.contains("cloud-platform")));
935    }
936
937    #[test]
938    fn test_custom_auth_scopes_override_default() {
939        let client = Client::builder()
940            .api_key("test-key")
941            .auth_scopes(vec!["scope-1".to_string()])
942            .build()
943            .unwrap();
944        assert_eq!(client.inner.config.auth_scopes, vec!["scope-1".to_string()]);
945    }
946}
rust_genai/client.rs

rust_genai/
client.rs
