rust_genai/
files.rs

1//! Files API surface.
2
3use std::path::Path;
4use std::sync::Arc;
5use std::time::{Duration, Instant};
6
7use crate::client::{Backend, ClientInner};
8use crate::error::{Error, Result};
9use crate::upload;
10#[cfg(test)]
11use crate::upload::CHUNK_SIZE;
12use rust_genai_types::enums::FileState;
13use rust_genai_types::files::{
14    DownloadFileConfig, File, ListFilesConfig, ListFilesResponse, UploadFileConfig,
15};
16use serde_json::Value;
17
18#[derive(Clone)]
19pub struct Files {
20    pub(crate) inner: Arc<ClientInner>,
21}
22
23impl Files {
24    pub(crate) const fn new(inner: Arc<ClientInner>) -> Self {
25        Self { inner }
26    }
27
28    /// 上传文件（直接上传字节数据）。
29    ///
30    /// # Errors
31    /// 当配置无效、请求失败或响应解析失败时返回错误。
32    pub async fn upload(&self, data: Vec<u8>, mime_type: impl Into<String>) -> Result<File> {
33        let config = UploadFileConfig {
34            mime_type: Some(mime_type.into()),
35            ..UploadFileConfig::default()
36        };
37        self.upload_with_config(data, config).await
38    }
39
40    /// 上传文件（自定义配置）。
41    ///
42    /// # Errors
43    /// 当配置无效、请求失败或响应解析失败时返回错误。
44    pub async fn upload_with_config(
45        &self,
46        data: Vec<u8>,
47        config: UploadFileConfig,
48    ) -> Result<File> {
49        ensure_gemini_backend(&self.inner)?;
50
51        let mime_type = config
52            .mime_type
53            .clone()
54            .ok_or_else(|| Error::InvalidConfig {
55                message: "mime_type is required when uploading raw bytes".into(),
56            })?;
57        let size_bytes = data.len() as u64;
58        let file = build_upload_file(config, size_bytes, &mime_type);
59        let upload_url = self
60            .start_resumable_upload(file, size_bytes, &mime_type, None)
61            .await?;
62        self.upload_bytes(&upload_url, &data).await
63    }
64
65    /// 从文件路径上传。
66    ///
67    /// # Errors
68    /// 当文件无效、请求失败或响应解析失败时返回错误。
69    pub async fn upload_from_path(&self, path: impl AsRef<Path>) -> Result<File> {
70        self.upload_from_path_with_config(path, UploadFileConfig::default())
71            .await
72    }
73
74    /// 从文件路径上传（自定义配置）。
75    ///
76    /// # Errors
77    /// 当文件无效、请求失败或响应解析失败时返回错误。
78    pub async fn upload_from_path_with_config(
79        &self,
80        path: impl AsRef<Path>,
81        mut config: UploadFileConfig,
82    ) -> Result<File> {
83        ensure_gemini_backend(&self.inner)?;
84
85        let path = path.as_ref();
86        let metadata = tokio::fs::metadata(path).await?;
87        if !metadata.is_file() {
88            return Err(Error::InvalidConfig {
89                message: format!("{} is not a valid file path", path.display()),
90            });
91        }
92
93        let size_bytes = metadata.len();
94        let mime_type = config.mime_type.take().unwrap_or_else(|| {
95            mime_guess::from_path(path)
96                .first_or_octet_stream()
97                .essence_str()
98                .to_string()
99        });
100
101        let file_name = path.file_name().and_then(|name| name.to_str());
102        let file = build_upload_file(config, size_bytes, &mime_type);
103        let upload_url = self
104            .start_resumable_upload(file, size_bytes, &mime_type, file_name)
105            .await?;
106        let mut file_handle = tokio::fs::File::open(path).await?;
107        self.upload_reader(&upload_url, &mut file_handle, size_bytes)
108            .await
109    }
110
111    /// 下载文件（返回字节内容）。
112    ///
113    /// # Errors
114    /// 当请求失败或响应解析失败时返回错误。
115    pub async fn download(&self, name_or_uri: impl AsRef<str>) -> Result<Vec<u8>> {
116        ensure_gemini_backend(&self.inner)?;
117
118        let file_name = normalize_file_name(name_or_uri.as_ref())?;
119        let url = build_file_download_url(&self.inner, &file_name);
120        let request = self.inner.http.get(url);
121        let response = self.inner.send(request).await?;
122        if !response.status().is_success() {
123            return Err(Error::ApiError {
124                status: response.status().as_u16(),
125                message: response.text().await.unwrap_or_default(),
126            });
127        }
128        let bytes = response.bytes().await?;
129        Ok(bytes.to_vec())
130    }
131
132    #[allow(unused_variables)]
133    /// 下载文件（自定义配置）。
134    ///
135    /// # Errors
136    /// 当请求失败或响应解析失败时返回错误。
137    pub async fn download_with_config(
138        &self,
139        name_or_uri: impl AsRef<str>,
140        _config: DownloadFileConfig,
141    ) -> Result<Vec<u8>> {
142        self.download(name_or_uri).await
143    }
144
145    /// 列出文件。
146    ///
147    /// # Errors
148    /// 当请求失败或响应解析失败时返回错误。
149    pub async fn list(&self) -> Result<ListFilesResponse> {
150        self.list_with_config(ListFilesConfig::default()).await
151    }
152
153    /// 列出文件（自定义配置）。
154    ///
155    /// # Errors
156    /// 当请求失败或响应解析失败时返回错误。
157    pub async fn list_with_config(&self, config: ListFilesConfig) -> Result<ListFilesResponse> {
158        ensure_gemini_backend(&self.inner)?;
159        let url = build_files_list_url(&self.inner, &config)?;
160        let request = self.inner.http.get(url);
161        let response = self.inner.send(request).await?;
162        if !response.status().is_success() {
163            return Err(Error::ApiError {
164                status: response.status().as_u16(),
165                message: response.text().await.unwrap_or_default(),
166            });
167        }
168        Ok(response.json::<ListFilesResponse>().await?)
169    }
170
171    /// 列出所有文件（自动翻页）。
172    ///
173    /// # Errors
174    /// 当请求失败或响应解析失败时返回错误。
175    pub async fn all(&self) -> Result<Vec<File>> {
176        self.all_with_config(ListFilesConfig::default()).await
177    }
178
179    /// 列出所有文件（带配置，自动翻页）。
180    ///
181    /// # Errors
182    /// 当请求失败或响应解析失败时返回错误。
183    pub async fn all_with_config(&self, mut config: ListFilesConfig) -> Result<Vec<File>> {
184        let mut files = Vec::new();
185        loop {
186            let response = self.list_with_config(config.clone()).await?;
187            if let Some(items) = response.files {
188                files.extend(items);
189            }
190            match response.next_page_token {
191                Some(token) if !token.is_empty() => {
192                    config.page_token = Some(token);
193                }
194                _ => break,
195            }
196        }
197        Ok(files)
198    }
199
200    /// 获取文件元数据。
201    ///
202    /// # Errors
203    /// 当请求失败或响应解析失败时返回错误。
204    pub async fn get(&self, name_or_uri: impl AsRef<str>) -> Result<File> {
205        ensure_gemini_backend(&self.inner)?;
206
207        let file_name = normalize_file_name(name_or_uri.as_ref())?;
208        let url = build_file_url(&self.inner, &file_name);
209        let request = self.inner.http.get(url);
210        let response = self.inner.send(request).await?;
211        if !response.status().is_success() {
212            return Err(Error::ApiError {
213                status: response.status().as_u16(),
214                message: response.text().await.unwrap_or_default(),
215            });
216        }
217        Ok(response.json::<File>().await?)
218    }
219
220    /// 删除文件。
221    ///
222    /// # Errors
223    /// 当请求失败或响应解析失败时返回错误。
224    pub async fn delete(&self, name_or_uri: impl AsRef<str>) -> Result<()> {
225        ensure_gemini_backend(&self.inner)?;
226
227        let file_name = normalize_file_name(name_or_uri.as_ref())?;
228        let url = build_file_url(&self.inner, &file_name);
229        let request = self.inner.http.delete(url);
230        let response = self.inner.send(request).await?;
231        if !response.status().is_success() {
232            return Err(Error::ApiError {
233                status: response.status().as_u16(),
234                message: response.text().await.unwrap_or_default(),
235            });
236        }
237        Ok(())
238    }
239
240    /// 轮询直到文件状态变为 ACTIVE。
241    ///
242    /// # Errors
243    /// 当请求失败、文件失败或超时返回错误。
244    pub async fn wait_for_active(
245        &self,
246        name_or_uri: impl AsRef<str>,
247        config: WaitForFileConfig,
248    ) -> Result<File> {
249        ensure_gemini_backend(&self.inner)?;
250
251        let start = Instant::now();
252        loop {
253            let file = self.get(name_or_uri.as_ref()).await?;
254            match file.state {
255                Some(FileState::Active) => return Ok(file),
256                Some(FileState::Failed) => {
257                    return Err(Error::ApiError {
258                        status: 500,
259                        message: "File processing failed".into(),
260                    })
261                }
262                _ => {}
263            }
264
265            if let Some(timeout) = config.timeout {
266                if start.elapsed() >= timeout {
267                    return Err(Error::Timeout {
268                        message: "Timed out waiting for file to become ACTIVE".into(),
269                    });
270                }
271            }
272
273            tokio::time::sleep(config.poll_interval).await;
274        }
275    }
276
277    async fn start_resumable_upload(
278        &self,
279        file: File,
280        size_bytes: u64,
281        mime_type: &str,
282        file_name: Option<&str>,
283    ) -> Result<String> {
284        let url = build_files_upload_url(&self.inner);
285        let mut request = self
286            .inner
287            .http
288            .post(url)
289            .header("X-Goog-Upload-Protocol", "resumable")
290            .header("X-Goog-Upload-Command", "start")
291            .header(
292                "X-Goog-Upload-Header-Content-Length",
293                size_bytes.to_string(),
294            )
295            .header("X-Goog-Upload-Header-Content-Type", mime_type);
296
297        if let Some(file_name) = file_name {
298            request = request.header("X-Goog-Upload-File-Name", file_name);
299        }
300
301        let body = serde_json::json!({ "file": file });
302        let request = request.json(&body);
303        let response = self.inner.send(request).await?;
304        if !response.status().is_success() {
305            return Err(Error::ApiError {
306                status: response.status().as_u16(),
307                message: response.text().await.unwrap_or_default(),
308            });
309        }
310
311        let upload_url = response
312            .headers()
313            .get("x-goog-upload-url")
314            .and_then(|value| value.to_str().ok())
315            .ok_or_else(|| Error::Parse {
316                message: "Missing x-goog-upload-url header".into(),
317            })?;
318
319        Ok(upload_url.to_string())
320    }
321
322    async fn upload_bytes(&self, upload_url: &str, data: &[u8]) -> Result<File> {
323        let validate_status = |status: &str| {
324            if status != "active" {
325                return Err(Error::Parse {
326                    message: format!("Unexpected upload status: {status}"),
327                });
328            }
329            Ok(())
330        };
331
332        upload::upload_bytes_with(
333            data,
334            |chunk, offset, finalize| self.send_upload_chunk(upload_url, chunk, offset, finalize),
335            validate_status,
336            "Upload finished without final response",
337        )
338        .await
339    }
340
341    async fn upload_reader(
342        &self,
343        upload_url: &str,
344        reader: &mut tokio::fs::File,
345        total_size: u64,
346    ) -> Result<File> {
347        let validate_status = |status: &str| {
348            if status != "active" {
349                return Err(Error::Parse {
350                    message: format!("Unexpected upload status: {status}"),
351                });
352            }
353            Ok(())
354        };
355
356        upload::upload_reader_with(
357            reader,
358            total_size,
359            |chunk, offset, finalize| self.send_upload_chunk(upload_url, chunk, offset, finalize),
360            validate_status,
361            "Upload finished without final response",
362        )
363        .await
364    }
365
366    async fn send_upload_chunk(
367        &self,
368        upload_url: &str,
369        chunk: Vec<u8>,
370        offset: u64,
371        finalize: bool,
372    ) -> Result<(String, Option<File>)> {
373        let command = if finalize {
374            "upload, finalize"
375        } else {
376            "upload"
377        };
378        let chunk_len = chunk.len();
379        let response = self
380            .inner
381            .http
382            .post(upload_url)
383            .header("X-Goog-Upload-Command", command)
384            .header("X-Goog-Upload-Offset", offset.to_string())
385            .header("Content-Length", chunk_len.to_string())
386            .body(chunk)
387            .send()
388            .await?;
389
390        if !response.status().is_success() {
391            return Err(Error::ApiError {
392                status: response.status().as_u16(),
393                message: response.text().await.unwrap_or_default(),
394            });
395        }
396
397        let upload_status = response
398            .headers()
399            .get("x-goog-upload-status")
400            .and_then(|value| value.to_str().ok())
401            .ok_or_else(|| Error::Parse {
402                message: "Missing x-goog-upload-status header".into(),
403            })?
404            .to_string();
405
406        let body = response.bytes().await?;
407        if body.is_empty() {
408            return Ok((upload_status, None));
409        }
410
411        let value: Value = serde_json::from_slice(&body)?;
412        let file_value = value.get("file").cloned().unwrap_or(value);
413        let file: File = serde_json::from_value(file_value)?;
414
415        Ok((upload_status, Some(file)))
416    }
417}
418
419#[derive(Debug, Clone)]
420pub struct WaitForFileConfig {
421    pub poll_interval: Duration,
422    pub timeout: Option<Duration>,
423}
424
425impl Default for WaitForFileConfig {
426    fn default() -> Self {
427        Self {
428            poll_interval: Duration::from_secs(2),
429            timeout: Some(Duration::from_secs(300)),
430        }
431    }
432}
433
434#[cfg(test)]
435fn finalize_upload(status: &str, file: Option<File>) -> Result<File> {
436    upload::finalize_upload(status, file)
437}
438
439fn ensure_gemini_backend(inner: &ClientInner) -> Result<()> {
440    if inner.config.backend == Backend::VertexAi {
441        return Err(Error::InvalidConfig {
442            message: "Files API is only supported in Gemini API".into(),
443        });
444    }
445    Ok(())
446}
447
448fn build_upload_file(config: UploadFileConfig, size_bytes: u64, mime_type: &str) -> File {
449    let mut file = File::default();
450    if let Some(name) = config.name {
451        file.name = Some(normalize_upload_name(&name));
452    }
453    file.display_name = config.display_name;
454    file.mime_type = Some(mime_type.to_string());
455    file.size_bytes = Some(size_bytes.to_string());
456    file
457}
458
459fn normalize_upload_name(name: &str) -> String {
460    if name.starts_with("files/") {
461        name.to_string()
462    } else {
463        format!("files/{name}")
464    }
465}
466
467fn normalize_file_name(value: &str) -> Result<String> {
468    if value.starts_with("http://") || value.starts_with("https://") {
469        let marker = "files/";
470        let start = value.find(marker).ok_or_else(|| Error::InvalidConfig {
471            message: format!("Could not find 'files/' in URI: {value}"),
472        })?;
473        let suffix = &value[start + marker.len()..];
474        let name: String = suffix
475            .chars()
476            .take_while(|c| c.is_ascii_lowercase() || c.is_ascii_digit() || *c == '-')
477            .collect();
478        if name.is_empty() {
479            return Err(Error::InvalidConfig {
480                message: format!("Could not extract file name from URI: {value}"),
481            });
482        }
483        Ok(name)
484    } else if value.starts_with("files/") {
485        Ok(value.trim_start_matches("files/").to_string())
486    } else {
487        Ok(value.to_string())
488    }
489}
490
491fn build_files_upload_url(inner: &ClientInner) -> String {
492    let base = &inner.api_client.base_url;
493    let version = &inner.api_client.api_version;
494    format!("{base}upload/{version}/files")
495}
496
497fn build_files_list_url(inner: &ClientInner, config: &ListFilesConfig) -> Result<String> {
498    let base = &inner.api_client.base_url;
499    let version = &inner.api_client.api_version;
500    let url = format!("{base}{version}/files");
501    add_list_query_params(&url, config)
502}
503
504fn build_file_url(inner: &ClientInner, name: &str) -> String {
505    let base = &inner.api_client.base_url;
506    let version = &inner.api_client.api_version;
507    format!("{base}{version}/files/{name}")
508}
509
510fn build_file_download_url(inner: &ClientInner, name: &str) -> String {
511    let base = &inner.api_client.base_url;
512    let version = &inner.api_client.api_version;
513    format!("{base}{version}/files/{name}:download?alt=media")
514}
515
516fn add_list_query_params(url: &str, config: &ListFilesConfig) -> Result<String> {
517    let mut url = reqwest::Url::parse(url).map_err(|err| Error::InvalidConfig {
518        message: err.to_string(),
519    })?;
520    {
521        let mut pairs = url.query_pairs_mut();
522        if let Some(page_size) = config.page_size {
523            pairs.append_pair("pageSize", &page_size.to_string());
524        }
525        if let Some(page_token) = &config.page_token {
526            pairs.append_pair("pageToken", page_token);
527        }
528    }
529    Ok(url.to_string())
530}
531
532#[cfg(test)]
533mod tests {
534    use super::*;
535    use crate::client::Client;
536    use crate::test_support::test_client_inner;
537    use serde_json::json;
538    use wiremock::matchers::{method, path};
539    use wiremock::{Mock, MockServer, Request, Respond, ResponseTemplate};
540
541    #[test]
542    fn test_normalize_file_name() {
543        assert_eq!(normalize_file_name("files/abc-123").unwrap(), "abc-123");
544        assert_eq!(normalize_file_name("abc-123").unwrap(), "abc-123");
545        assert_eq!(
546            normalize_file_name("https://example.com/files/abc-123?foo=bar").unwrap(),
547            "abc-123"
548        );
549    }
550
551    #[test]
552    fn test_build_urls() {
553        let client = Client::new("test-key").unwrap();
554        let files = client.files();
555        let url = build_files_upload_url(&files.inner);
556        assert_eq!(
557            url,
558            "https://generativelanguage.googleapis.com/upload/v1beta/files"
559        );
560    }
561
562    #[test]
563    fn test_normalize_upload_and_list_params() {
564        assert_eq!(normalize_upload_name("files/abc"), "files/abc");
565        assert_eq!(normalize_upload_name("abc"), "files/abc");
566        assert!(normalize_file_name("https://example.com/no-files").is_err());
567        assert!(normalize_file_name("https://example.com/files/?x").is_err());
568
569        let url = add_list_query_params(
570            "https://example.com/files",
571            &ListFilesConfig {
572                page_size: Some(3),
573                page_token: Some("t".to_string()),
574            },
575        )
576        .unwrap();
577        assert!(url.contains("pageSize=3"));
578        assert!(url.contains("pageToken=t"));
579    }
580
581    #[test]
582    fn test_build_upload_file_and_finalize_errors() {
583        let file = build_upload_file(
584            UploadFileConfig {
585                name: Some("abc".to_string()),
586                display_name: Some("d".to_string()),
587                ..Default::default()
588            },
589            5,
590            "text/plain",
591        );
592        assert_eq!(file.name.as_deref(), Some("files/abc"));
593        assert_eq!(file.size_bytes.as_deref(), Some("5"));
594
595        let err = finalize_upload("active", None).unwrap_err();
596        assert!(matches!(err, Error::Parse { .. }));
597        let err = finalize_upload("final", None).unwrap_err();
598        assert!(matches!(err, Error::Parse { .. }));
599    }
600
601    #[test]
602    fn test_ensure_gemini_backend_error() {
603        let vertex = test_client_inner(Backend::VertexAi);
604        let err = ensure_gemini_backend(&vertex).unwrap_err();
605        assert!(matches!(err, Error::InvalidConfig { .. }));
606    }
607
608    #[tokio::test]
609    async fn test_start_resumable_upload_and_send_chunk_errors() {
610        let server = MockServer::start().await;
611        Mock::given(method("POST"))
612            .and(path("/upload/v1beta/files"))
613            .respond_with(ResponseTemplate::new(200))
614            .mount(&server)
615            .await;
616
617        let client = Client::builder()
618            .api_key("test-key")
619            .base_url(server.uri())
620            .build()
621            .unwrap();
622        let files = client.files();
623        let file = build_upload_file(UploadFileConfig::default(), 1, "text/plain");
624        let err = files
625            .start_resumable_upload(file, 1, "text/plain", None)
626            .await
627            .unwrap_err();
628        assert!(matches!(err, Error::Parse { .. }));
629
630        Mock::given(method("POST"))
631            .and(path("/upload-chunk"))
632            .respond_with(ResponseTemplate::new(200))
633            .mount(&server)
634            .await;
635        let err = files
636            .send_upload_chunk(
637                &format!("{}/upload-chunk", server.uri()),
638                Vec::new(),
639                0,
640                true,
641            )
642            .await
643            .unwrap_err();
644        assert!(matches!(err, Error::Parse { .. }));
645
646        Mock::given(method("POST"))
647            .and(path("/upload-fail"))
648            .respond_with(ResponseTemplate::new(400).set_body_string("bad"))
649            .mount(&server)
650            .await;
651        let err = files
652            .send_upload_chunk(
653                &format!("{}/upload-fail", server.uri()),
654                Vec::new(),
655                0,
656                true,
657            )
658            .await
659            .unwrap_err();
660        assert!(matches!(err, Error::ApiError { .. }));
661    }
662
663    #[tokio::test]
664    async fn test_files_upload_errors() {
665        let client = Client::new("test-key").unwrap();
666        let files = client.files();
667
668        let err = files
669            .upload_with_config(vec![1, 2, 3], UploadFileConfig::default())
670            .await
671            .unwrap_err();
672        assert!(matches!(err, Error::InvalidConfig { .. }));
673
674        let temp_dir = std::env::temp_dir().join("rust_genai_files_test_dir");
675        let _ = tokio::fs::create_dir_all(&temp_dir).await;
676        let err = files
677            .upload_from_path_with_config(&temp_dir, UploadFileConfig::default())
678            .await
679            .unwrap_err();
680        assert!(matches!(err, Error::InvalidConfig { .. }));
681        let _ = tokio::fs::remove_dir_all(&temp_dir).await;
682    }
683
684    #[tokio::test]
685    async fn test_start_resumable_upload_error_response() {
686        let server = MockServer::start().await;
687        Mock::given(method("POST"))
688            .and(path("/upload/v1beta/files"))
689            .respond_with(ResponseTemplate::new(500).set_body_string("boom"))
690            .mount(&server)
691            .await;
692
693        let client = Client::builder()
694            .api_key("test-key")
695            .base_url(server.uri())
696            .build()
697            .unwrap();
698        let files = client.files();
699        let file = build_upload_file(UploadFileConfig::default(), 1, "text/plain");
700        let err = files
701            .start_resumable_upload(file, 1, "text/plain", None)
702            .await
703            .unwrap_err();
704        assert!(matches!(err, Error::ApiError { .. }));
705    }
706
707    #[tokio::test]
708    async fn test_upload_bytes_empty_and_status_errors() {
709        let server = MockServer::start().await;
710
711        Mock::given(method("POST"))
712            .and(path("/upload-empty"))
713            .respond_with(
714                ResponseTemplate::new(200)
715                    .insert_header("x-goog-upload-status", "final")
716                    .set_body_json(json!({
717                        "file": {"name": "files/empty", "state": "ACTIVE"}
718                    })),
719            )
720            .mount(&server)
721            .await;
722
723        Mock::given(method("POST"))
724            .and(path("/upload-bad"))
725            .respond_with(
726                ResponseTemplate::new(200).insert_header("x-goog-upload-status", "paused"),
727            )
728            .mount(&server)
729            .await;
730
731        let client = Client::builder()
732            .api_key("test-key")
733            .base_url(server.uri())
734            .build()
735            .unwrap();
736        let files = client.files();
737
738        let file = files
739            .upload_bytes(&format!("{}/upload-empty", server.uri()), &[])
740            .await
741            .unwrap();
742        assert_eq!(file.name.as_deref(), Some("files/empty"));
743
744        let data = vec![0u8; CHUNK_SIZE + 1];
745        let err = files
746            .upload_bytes(&format!("{}/upload-bad", server.uri()), &data)
747            .await
748            .unwrap_err();
749        assert!(matches!(err, Error::Parse { .. }));
750    }
751
752    #[tokio::test]
753    async fn test_upload_reader_empty_file() {
754        let server = MockServer::start().await;
755        Mock::given(method("POST"))
756            .and(path("/upload-empty-file"))
757            .respond_with(
758                ResponseTemplate::new(200)
759                    .insert_header("x-goog-upload-status", "final")
760                    .set_body_json(json!({
761                        "file": {"name": "files/empty-file", "state": "ACTIVE"}
762                    })),
763            )
764            .mount(&server)
765            .await;
766
767        let client = Client::builder()
768            .api_key("test-key")
769            .base_url(server.uri())
770            .build()
771            .unwrap();
772        let files = client.files();
773        let temp_path = std::env::temp_dir().join("rust_genai_empty_upload_file");
774        let _ = tokio::fs::write(&temp_path, &[]).await;
775        let mut handle = tokio::fs::File::open(&temp_path).await.unwrap();
776
777        let file = files
778            .upload_reader(
779                &format!("{}/upload-empty-file", server.uri()),
780                &mut handle,
781                0,
782            )
783            .await
784            .unwrap();
785        assert_eq!(file.name.as_deref(), Some("files/empty-file"));
786        let _ = tokio::fs::remove_file(&temp_path).await;
787    }
788
789    #[tokio::test]
790    async fn test_upload_bytes_and_reader_active_then_final() {
791        #[derive(Clone)]
792        struct UploadResponder;
793
794        impl Respond for UploadResponder {
795            fn respond(&self, request: &Request) -> ResponseTemplate {
796                let finalize = request
797                    .headers
798                    .get("x-goog-upload-command")
799                    .and_then(|value| value.to_str().ok())
800                    .is_some_and(|value| value.contains("finalize"));
801                if finalize {
802                    ResponseTemplate::new(200)
803                        .insert_header("x-goog-upload-status", "final")
804                        .set_body_json(json!({
805                            "file": {"name": "files/final", "state": "ACTIVE"}
806                        }))
807                } else {
808                    ResponseTemplate::new(200).insert_header("x-goog-upload-status", "active")
809                }
810            }
811        }
812
813        let server = MockServer::start().await;
814        Mock::given(method("POST"))
815            .and(path("/upload-active"))
816            .respond_with(UploadResponder)
817            .mount(&server)
818            .await;
819
820        let client = Client::builder()
821            .api_key("test-key")
822            .base_url(server.uri())
823            .build()
824            .unwrap();
825        let files = client.files();
826        let data = vec![0u8; CHUNK_SIZE + 1];
827        let file = files
828            .upload_bytes(&format!("{}/upload-active", server.uri()), &data)
829            .await
830            .unwrap();
831        assert_eq!(file.name.as_deref(), Some("files/final"));
832
833        Mock::given(method("POST"))
834            .and(path("/upload-reader"))
835            .respond_with(UploadResponder)
836            .mount(&server)
837            .await;
838        let temp_path = std::env::temp_dir().join("rust_genai_reader_active");
839        let _ = tokio::fs::write(&temp_path, vec![0u8; CHUNK_SIZE + 1]).await;
840        let mut handle = tokio::fs::File::open(&temp_path).await.unwrap();
841        let file = files
842            .upload_reader(
843                &format!("{}/upload-reader", server.uri()),
844                &mut handle,
845                (CHUNK_SIZE + 1) as u64,
846            )
847            .await
848            .unwrap();
849        assert_eq!(file.name.as_deref(), Some("files/final"));
850        let _ = tokio::fs::remove_file(&temp_path).await;
851    }
852
853    #[tokio::test]
854    async fn test_upload_with_config_and_mime_guess() {
855        let server = MockServer::start().await;
856        Mock::given(method("POST"))
857            .and(path("/upload/v1beta/files"))
858            .respond_with(
859                ResponseTemplate::new(200)
860                    .insert_header("x-goog-upload-url", format!("{}/upload-ok", server.uri())),
861            )
862            .mount(&server)
863            .await;
864        Mock::given(method("POST"))
865            .and(path("/upload-ok"))
866            .respond_with(
867                ResponseTemplate::new(200)
868                    .insert_header("x-goog-upload-status", "final")
869                    .set_body_json(json!({
870                        "file": {"name": "files/ok", "state": "ACTIVE"}
871                    })),
872            )
873            .mount(&server)
874            .await;
875
876        let client = Client::builder()
877            .api_key("test-key")
878            .base_url(server.uri())
879            .build()
880            .unwrap();
881        let files = client.files();
882        let file = files.upload(vec![1, 2, 3], "text/plain").await.unwrap();
883        assert_eq!(file.name.as_deref(), Some("files/ok"));
884
885        let temp_path = std::env::temp_dir().join("rust_genai_upload_guess.txt");
886        let _ = tokio::fs::write(&temp_path, b"hello").await;
887        let file = files
888            .upload_from_path_with_config(&temp_path, UploadFileConfig::default())
889            .await
890            .unwrap();
891        assert_eq!(file.name.as_deref(), Some("files/ok"));
892        let _ = tokio::fs::remove_file(&temp_path).await;
893    }
894
895    #[tokio::test]
896    async fn test_wait_for_active_timeout_after_sleep() {
897        let server = MockServer::start().await;
898        Mock::given(method("GET"))
899            .and(path("/v1beta/files/slow"))
900            .respond_with(ResponseTemplate::new(200).set_body_json(json!({
901                "name": "files/slow",
902                "state": "PROCESSING"
903            })))
904            .mount(&server)
905            .await;
906
907        let client = Client::builder()
908            .api_key("test-key")
909            .base_url(server.uri())
910            .build()
911            .unwrap();
912        let files = client.files();
913        let err = files
914            .wait_for_active(
915                "slow",
916                WaitForFileConfig {
917                    poll_interval: Duration::from_millis(1),
918                    timeout: Some(Duration::from_millis(2)),
919                },
920            )
921            .await
922            .unwrap_err();
923        assert!(matches!(err, Error::Timeout { .. }));
924    }
925
926    #[test]
927    fn test_add_list_query_params_invalid_url() {
928        let err = add_list_query_params("http://[::1", &ListFilesConfig::default()).unwrap_err();
929        assert!(matches!(err, Error::InvalidConfig { .. }));
930    }
931}
rust_genai/files.rs

rust_genai/
files.rs
