gemini_adk_rs/artifacts/
gcs_service.rs

1//! GCS-backed artifact service.
2//!
3//! Feature-gated behind `gcs-artifacts`. Implements [`ArtifactService`] against
4//! Google Cloud Storage using the JSON HTTP API
5//! (`https://storage.googleapis.com/storage/v1/...` for metadata operations and
6//! `https://storage.googleapis.com/upload/...` for uploads). Authentication is
7//! via an OAuth2 bearer access token, supplied the same way as the rest of this
8//! crate's REST-backed services (a static token or a refresher closure).
9//!
10//! # Blob naming scheme
11//!
12//! This mirrors the ADK Python `GcsArtifactService`. The blob name depends on
13//! whether the filename carries a user namespace:
14//!
15//! - User-scoped (filename starts with `"user:"`):
16//!   `{app_name}/{user_id}/user/{filename}/{version}`
17//! - Session-scoped (regular filenames):
18//!   `{app_name}/{user_id}/{session_id}/{filename}/{version}`
19//!
20//! Versions are sequential integers per artifact. The first version is `0`
21//! (matching ADK on the wire); the next version is `max(existing) + 1`.
22//!
23//! # Trait mapping
24//!
25//! The crate's [`ArtifactService`] trait is session-scoped and does not carry a
26//! `user_id` argument, so a fixed `user_id` is held on the service (defaults to
27//! `"user"`, override with [`GcsArtifactService::user_id`]). The trait exposes
28//! 1-based version numbers (consistent with the in-memory and file services),
29//! while the GCS wire layout follows ADK's 0-based numbering. The two are
30//! mapped at the trait boundary: wire version `v` is surfaced as `v + 1`.
31
32use std::sync::Arc;
33
34use async_trait::async_trait;
35use serde::Deserialize;
36
37use super::{Artifact, ArtifactError, ArtifactMetadata, ArtifactService, now_secs};
38
39const STORAGE_BASE: &str = "https://storage.googleapis.com/storage/v1";
40const UPLOAD_BASE: &str = "https://storage.googleapis.com/upload/storage/v1";
41
42// ──────────────────────────────────────────────────────────────────────────────
43// Auth provider (mirrors session::vertex_ai::TokenProvider)
44// ──────────────────────────────────────────────────────────────────────────────
45
46/// How to supply a bearer token for GCS requests.
47enum TokenProvider {
48    /// No token configured — requests will fail with a clear message.
49    None,
50    /// A static, pre-fetched bearer token string.
51    Static(String),
52    /// A dynamic refresher: called before every request.
53    Refresher(Arc<dyn Fn() -> String + Send + Sync>),
54}
55
56impl TokenProvider {
57    /// Retrieve the current token, or return an error if none is configured.
58    fn get(&self) -> Result<String, ArtifactError> {
59        match self {
60            TokenProvider::None => Err(ArtifactError::Storage(
61                "missing auth token: call .with_token() or .with_token_refresher()".into(),
62            )),
63            TokenProvider::Static(t) => Ok(t.clone()),
64            TokenProvider::Refresher(f) => Ok(f()),
65        }
66    }
67}
68
69// ──────────────────────────────────────────────────────────────────────────────
70// DTO types — mirror GCS JSON shapes
71// ──────────────────────────────────────────────────────────────────────────────
72
73/// A GCS object resource (subset of fields we use).
74#[derive(Debug, Deserialize)]
75struct GcsObject {
76    /// Object name (the full blob path within the bucket).
77    name: String,
78}
79
80/// Response envelope for `objects.list`.
81#[derive(Debug, Default, Deserialize)]
82struct ListObjectsResponse {
83    #[serde(default)]
84    items: Vec<GcsObject>,
85    #[serde(rename = "nextPageToken", default)]
86    next_page_token: Option<String>,
87}
88
89// ──────────────────────────────────────────────────────────────────────────────
90// Pure helpers (blob naming) — mirror ADK's _get_blob_prefix / _get_blob_name
91// ──────────────────────────────────────────────────────────────────────────────
92
93/// Whether the filename carries a user namespace (starts with `"user:"`).
94fn file_has_user_namespace(filename: &str) -> bool {
95    filename.starts_with("user:")
96}
97
98/// Construct the blob name prefix (everything up to but excluding `/{version}`).
99fn blob_prefix(app_name: &str, user_id: &str, session_id: &str, filename: &str) -> String {
100    if file_has_user_namespace(filename) {
101        format!("{app_name}/{user_id}/user/{filename}")
102    } else {
103        format!("{app_name}/{user_id}/{session_id}/{filename}")
104    }
105}
106
107/// Construct the full blob name including the version suffix.
108fn blob_name(
109    app_name: &str,
110    user_id: &str,
111    session_id: &str,
112    filename: &str,
113    version: u64,
114) -> String {
115    format!(
116        "{}/{}",
117        blob_prefix(app_name, user_id, session_id, filename),
118        version
119    )
120}
121
122/// Parse the trailing `/{version}` integer from a blob name, if present.
123fn version_from_blob_name(name: &str) -> Option<u64> {
124    name.rsplit('/')
125        .next()
126        .and_then(|seg| seg.parse::<u64>().ok())
127}
128
129// ──────────────────────────────────────────────────────────────────────────────
130// Service struct
131// ──────────────────────────────────────────────────────────────────────────────
132
133/// GCS-backed artifact service.
134///
135/// Path format (on the wire, mirroring ADK):
136/// `{app_name}/{user_id}/{session_id}/{filename}/{version}` for session-scoped
137/// artifacts, and `{app_name}/{user_id}/user/{filename}/{version}` for
138/// user-namespaced filenames (those starting with `"user:"`).
139///
140/// # Quick start
141///
142/// ```rust,no_run
143/// # use gemini_adk_rs::artifacts::GcsArtifactService;
144/// let svc = GcsArtifactService::new("my-bucket", "my-app")
145///     .with_token("ya29.my-access-token")
146///     .user_id("alice");
147/// ```
148pub struct GcsArtifactService {
149    bucket: String,
150    app_name: String,
151    user_id: String,
152    client: reqwest::Client,
153    token_provider: TokenProvider,
154}
155
156impl GcsArtifactService {
157    /// Create a new GCS artifact service targeting the given bucket.
158    ///
159    /// No auth token is configured yet — call [`with_token`](Self::with_token)
160    /// or [`with_token_refresher`](Self::with_token_refresher) before issuing
161    /// requests, otherwise they will return
162    /// `ArtifactError::Storage("missing auth token")`.
163    pub fn new(bucket: impl Into<String>, app_name: impl Into<String>) -> Self {
164        Self {
165            bucket: bucket.into(),
166            app_name: app_name.into(),
167            user_id: "user".to_string(),
168            client: reqwest::Client::new(),
169            token_provider: TokenProvider::None,
170        }
171    }
172
173    /// Set a static bearer token (OAuth2 access token) for all requests.
174    pub fn with_token(mut self, token: impl Into<String>) -> Self {
175        self.token_provider = TokenProvider::Static(token.into());
176        self
177    }
178
179    /// Set a dynamic token refresher closure.
180    ///
181    /// The closure is invoked before every HTTP request, allowing the caller to
182    /// supply a freshly-refreshed token each time.
183    pub fn with_token_refresher(mut self, f: impl Fn() -> String + Send + Sync + 'static) -> Self {
184        self.token_provider = TokenProvider::Refresher(Arc::new(f));
185        self
186    }
187
188    /// Override the user ID used in blob paths (defaults to `"user"`).
189    ///
190    /// The crate's [`ArtifactService`] trait is session-scoped and carries no
191    /// `user_id`, so it is fixed on the service to complete ADK's
192    /// `{app_name}/{user_id}/{session_id}/...` blob layout.
193    pub fn user_id(mut self, user_id: impl Into<String>) -> Self {
194        self.user_id = user_id.into();
195        self
196    }
197
198    /// The bucket this service targets.
199    pub fn bucket(&self) -> &str {
200        &self.bucket
201    }
202
203    /// The application name prefix used in object paths.
204    pub fn app_name(&self) -> &str {
205        &self.app_name
206    }
207
208    // ── Internal HTTP helpers ──────────────────────────────────────────────
209
210    fn token(&self) -> Result<String, ArtifactError> {
211        self.token_provider.get()
212    }
213
214    /// List the wire versions (0-based integers) of an artifact, in ascending
215    /// order. Mirrors ADK's `_list_versions`.
216    async fn list_wire_versions(
217        &self,
218        session_id: &str,
219        filename: &str,
220    ) -> Result<Vec<u64>, ArtifactError> {
221        let prefix = format!(
222            "{}/",
223            blob_prefix(&self.app_name, &self.user_id, session_id, filename)
224        );
225        let mut versions: Vec<u64> = self
226            .list_object_names(&prefix)
227            .await?
228            .into_iter()
229            .filter_map(|name| version_from_blob_name(&name))
230            .collect();
231        versions.sort_unstable();
232        Ok(versions)
233    }
234
235    /// List all object names under a prefix, paginating through results.
236    async fn list_object_names(&self, prefix: &str) -> Result<Vec<String>, ArtifactError> {
237        let token = self.token()?;
238        let base = format!("{STORAGE_BASE}/b/{}/o", urlencode(&self.bucket));
239        let mut names = Vec::new();
240        let mut page_token: Option<String> = None;
241
242        loop {
243            let mut req = self
244                .client
245                .get(&base)
246                .header("Authorization", format!("Bearer {token}"))
247                .query(&[("prefix", prefix)]);
248            if let Some(tok) = &page_token {
249                req = req.query(&[("pageToken", tok.as_str())]);
250            }
251
252            let resp = req
253                .send()
254                .await
255                .map_err(|e| ArtifactError::Storage(format!("HTTP request failed: {e}")))?;
256            let parsed: ListObjectsResponse = parse_json(resp).await?.unwrap_or_default();
257            names.extend(parsed.items.into_iter().map(|o| o.name));
258
259            match parsed.next_page_token {
260                Some(t) if !t.is_empty() => page_token = Some(t),
261                _ => break,
262            }
263        }
264        Ok(names)
265    }
266
267    /// Upload a payload to the given blob name with the given content type.
268    async fn upload_blob(
269        &self,
270        blob: &str,
271        data: &[u8],
272        content_type: &str,
273    ) -> Result<(), ArtifactError> {
274        let token = self.token()?;
275        let url = format!("{UPLOAD_BASE}/b/{}/o", urlencode(&self.bucket));
276        let resp = self
277            .client
278            .post(&url)
279            .header("Authorization", format!("Bearer {token}"))
280            .header("Content-Type", content_type)
281            .query(&[("uploadType", "media"), ("name", blob)])
282            .body(data.to_vec())
283            .send()
284            .await
285            .map_err(|e| ArtifactError::Storage(format!("HTTP request failed: {e}")))?;
286        check_success(resp).await
287    }
288
289    /// Download a blob's payload and content type. Returns `Ok(None)` on 404.
290    async fn download_blob(&self, blob: &str) -> Result<Option<(Vec<u8>, String)>, ArtifactError> {
291        let token = self.token()?;
292        let url = format!(
293            "{STORAGE_BASE}/b/{}/o/{}",
294            urlencode(&self.bucket),
295            urlencode(blob)
296        );
297        let resp = self
298            .client
299            .get(&url)
300            .header("Authorization", format!("Bearer {token}"))
301            .query(&[("alt", "media")])
302            .send()
303            .await
304            .map_err(|e| ArtifactError::Storage(format!("HTTP request failed: {e}")))?;
305
306        let status = resp.status().as_u16();
307        if status == 404 {
308            return Ok(None);
309        }
310        if !(200..300).contains(&status) {
311            let body = resp.text().await.unwrap_or_else(|_| "<unreadable>".into());
312            return Err(ArtifactError::Storage(format!(
313                "GCS request failed [{status}]: {body}"
314            )));
315        }
316        let content_type = resp
317            .headers()
318            .get(reqwest::header::CONTENT_TYPE)
319            .and_then(|v| v.to_str().ok())
320            .map(|s| s.split(';').next().unwrap_or(s).trim().to_string())
321            .unwrap_or_else(|| "application/octet-stream".to_string());
322        let bytes = resp
323            .bytes()
324            .await
325            .map_err(|e| ArtifactError::Storage(format!("failed to read body: {e}")))?;
326        Ok(Some((bytes.to_vec(), content_type)))
327    }
328
329    /// Delete a single blob. A 404 is treated as success (idempotent delete).
330    async fn delete_blob(&self, blob: &str) -> Result<(), ArtifactError> {
331        let token = self.token()?;
332        let url = format!(
333            "{STORAGE_BASE}/b/{}/o/{}",
334            urlencode(&self.bucket),
335            urlencode(blob)
336        );
337        let resp = self
338            .client
339            .delete(&url)
340            .header("Authorization", format!("Bearer {token}"))
341            .send()
342            .await
343            .map_err(|e| ArtifactError::Storage(format!("HTTP request failed: {e}")))?;
344        let status = resp.status().as_u16();
345        if status == 404 || (200..300).contains(&status) {
346            return Ok(());
347        }
348        let body = resp.text().await.unwrap_or_else(|_| "<unreadable>".into());
349        Err(ArtifactError::Storage(format!(
350            "GCS request failed [{status}]: {body}"
351        )))
352    }
353}
354
355// ──────────────────────────────────────────────────────────────────────────────
356// HTTP response helpers
357// ──────────────────────────────────────────────────────────────────────────────
358
359/// Parse a JSON response, mapping 404 to `Ok(None)` and non-2xx to an error.
360async fn parse_json<T: for<'de> Deserialize<'de>>(
361    resp: reqwest::Response,
362) -> Result<Option<T>, ArtifactError> {
363    let status = resp.status().as_u16();
364    if status == 404 {
365        return Ok(None);
366    }
367    if (200..300).contains(&status) {
368        let parsed: T = resp
369            .json()
370            .await
371            .map_err(|e| ArtifactError::Storage(format!("failed to parse response: {e}")))?;
372        return Ok(Some(parsed));
373    }
374    let body = resp.text().await.unwrap_or_else(|_| "<unreadable>".into());
375    Err(ArtifactError::Storage(format!(
376        "GCS request failed [{status}]: {body}"
377    )))
378}
379
380/// Verify a response was 2xx (for void operations like upload).
381async fn check_success(resp: reqwest::Response) -> Result<(), ArtifactError> {
382    let status = resp.status().as_u16();
383    if (200..300).contains(&status) {
384        return Ok(());
385    }
386    let body = resp.text().await.unwrap_or_else(|_| "<unreadable>".into());
387    Err(ArtifactError::Storage(format!(
388        "GCS request failed [{status}]: {body}"
389    )))
390}
391
392/// Percent-encode a path segment for use in GCS object URLs.
393///
394/// GCS object names may contain `/`, which must be encoded as `%2F` when the
395/// object name appears as a single path segment (e.g. `objects.get`).
396fn urlencode(s: &str) -> String {
397    let mut out = String::with_capacity(s.len());
398    for b in s.bytes() {
399        match b {
400            b'A'..=b'Z' | b'a'..=b'z' | b'0'..=b'9' | b'-' | b'_' | b'.' | b'~' => {
401                out.push(b as char);
402            }
403            _ => out.push_str(&format!("%{b:02X}")),
404        }
405    }
406    out
407}
408
409// ──────────────────────────────────────────────────────────────────────────────
410// ArtifactService implementation
411// ──────────────────────────────────────────────────────────────────────────────
412
413#[async_trait]
414impl ArtifactService for GcsArtifactService {
415    async fn save(
416        &self,
417        session_id: &str,
418        artifact: Artifact,
419    ) -> Result<ArtifactMetadata, ArtifactError> {
420        let filename = &artifact.metadata.name;
421
422        // Next wire version = max(existing) + 1, or 0 if none. (ADK semantics.)
423        let versions = self.list_wire_versions(session_id, filename).await?;
424        let wire_version = versions.iter().copied().max().map(|v| v + 1).unwrap_or(0);
425
426        let blob = blob_name(
427            &self.app_name,
428            &self.user_id,
429            session_id,
430            filename,
431            wire_version,
432        );
433        self.upload_blob(&blob, &artifact.data, &artifact.metadata.mime_type)
434            .await?;
435
436        let mut metadata = artifact.metadata;
437        // Surface a 1-based version at the trait boundary for consistency with
438        // the in-memory and file services.
439        metadata.version = (wire_version + 1) as u32;
440        metadata.updated_at = now_secs();
441        if wire_version == 0 {
442            metadata.created_at = metadata.updated_at;
443        }
444        Ok(metadata)
445    }
446
447    async fn load(&self, session_id: &str, name: &str) -> Result<Option<Artifact>, ArtifactError> {
448        let versions = self.list_wire_versions(session_id, name).await?;
449        let Some(latest) = versions.iter().copied().max() else {
450            return Ok(None);
451        };
452        self.load_version(session_id, name, (latest + 1) as u32)
453            .await
454    }
455
456    async fn load_version(
457        &self,
458        session_id: &str,
459        name: &str,
460        version: u32,
461    ) -> Result<Option<Artifact>, ArtifactError> {
462        if version == 0 {
463            return Ok(None);
464        }
465        // Trait versions are 1-based; the wire layout is 0-based.
466        let wire_version = (version - 1) as u64;
467        let blob = blob_name(
468            &self.app_name,
469            &self.user_id,
470            session_id,
471            name,
472            wire_version,
473        );
474
475        let Some((data, content_type)) = self.download_blob(&blob).await? else {
476            return Ok(None);
477        };
478
479        let now = now_secs();
480        let size = data.len();
481        Ok(Some(Artifact {
482            metadata: ArtifactMetadata {
483                name: name.to_string(),
484                mime_type: content_type,
485                version,
486                size,
487                created_at: now,
488                updated_at: now,
489            },
490            data,
491        }))
492    }
493
494    async fn list(&self, session_id: &str) -> Result<Vec<ArtifactMetadata>, ArtifactError> {
495        use std::collections::BTreeSet;
496
497        // Session-scoped filenames.
498        let session_prefix = format!("{}/{}/{}/", self.app_name, self.user_id, session_id);
499        // User-namespaced filenames (no session_id).
500        let user_prefix = format!("{}/{}/user/", self.app_name, self.user_id);
501
502        let mut filenames: BTreeSet<String> = BTreeSet::new();
503
504        for name in self.list_object_names(&session_prefix).await? {
505            if let Some(rest) = name.strip_prefix(&session_prefix) {
506                // rest is `{filename}/{version}` (filename may contain slashes).
507                if let Some(idx) = rest.rfind('/') {
508                    filenames.insert(rest[..idx].to_string());
509                }
510            }
511        }
512        for name in self.list_object_names(&user_prefix).await? {
513            if let Some(rest) = name.strip_prefix(&user_prefix)
514                && let Some(idx) = rest.rfind('/')
515            {
516                // Re-attach the `user:` prefix so callers can round-trip the
517                // returned name back into load/load_version. The stored blob
518                // path keeps the literal `user:` filename segment.
519                filenames.insert(rest[..idx].to_string());
520            }
521        }
522
523        // Resolve latest-version metadata for each filename.
524        let mut result = Vec::with_capacity(filenames.len());
525        for filename in filenames {
526            if let Some(artifact) = self.load(session_id, &filename).await? {
527                result.push(artifact.metadata);
528            }
529        }
530        Ok(result)
531    }
532
533    async fn delete(&self, session_id: &str, name: &str) -> Result<(), ArtifactError> {
534        let versions = self.list_wire_versions(session_id, name).await?;
535        for wire_version in versions {
536            let blob = blob_name(
537                &self.app_name,
538                &self.user_id,
539                session_id,
540                name,
541                wire_version,
542            );
543            self.delete_blob(&blob).await?;
544        }
545        Ok(())
546    }
547}
548
549#[cfg(test)]
550mod tests {
551    use super::*;
552
553    #[test]
554    fn can_construct() {
555        let svc = GcsArtifactService::new("my-bucket", "my-app");
556        assert_eq!(svc.bucket(), "my-bucket");
557        assert_eq!(svc.app_name(), "my-app");
558    }
559
560    #[test]
561    fn default_user_id_is_user() {
562        let svc = GcsArtifactService::new("b", "a");
563        assert_eq!(svc.user_id, "user");
564    }
565
566    #[test]
567    fn user_id_override() {
568        let svc = GcsArtifactService::new("b", "a").user_id("alice");
569        assert_eq!(svc.user_id, "alice");
570    }
571
572    #[test]
573    fn session_scoped_blob_name() {
574        assert_eq!(
575            blob_name("app", "user", "sess1", "file.bin", 3),
576            "app/user/sess1/file.bin/3"
577        );
578    }
579
580    #[test]
581    fn user_namespaced_blob_name_ignores_session() {
582        // ADK: user-namespaced files use `/user/` and ignore session_id.
583        assert_eq!(
584            blob_name("app", "alice", "sess1", "user:prefs.json", 0),
585            "app/alice/user/user:prefs.json/0"
586        );
587    }
588
589    #[test]
590    fn blob_prefix_session_scoped() {
591        assert_eq!(blob_prefix("app", "u", "s", "doc.txt"), "app/u/s/doc.txt");
592    }
593
594    #[test]
595    fn blob_prefix_user_scoped() {
596        assert_eq!(
597            blob_prefix("app", "u", "s", "user:doc.txt"),
598            "app/u/user/user:doc.txt"
599        );
600    }
601
602    #[test]
603    fn detects_user_namespace() {
604        assert!(file_has_user_namespace("user:settings"));
605        assert!(!file_has_user_namespace("settings"));
606    }
607
608    #[test]
609    fn parses_version_suffix() {
610        assert_eq!(version_from_blob_name("app/u/s/file/7"), Some(7));
611        assert_eq!(version_from_blob_name("app/u/s/file/notanum"), None);
612        assert_eq!(version_from_blob_name("0"), Some(0));
613    }
614
615    #[test]
616    fn urlencode_encodes_slashes() {
617        assert_eq!(urlencode("a/b/c"), "a%2Fb%2Fc");
618        assert_eq!(urlencode("user:f.json"), "user%3Af.json");
619        assert_eq!(urlencode("plain-name_1.txt"), "plain-name_1.txt");
620    }
621
622    #[test]
623    fn missing_token_returns_storage_error() {
624        let svc = GcsArtifactService::new("b", "a");
625        let err = svc.token().unwrap_err();
626        assert!(matches!(err, ArtifactError::Storage(_)));
627        assert!(err.to_string().contains("missing auth token"));
628    }
629
630    #[test]
631    fn with_token_sets_provider() {
632        let svc = GcsArtifactService::new("b", "a").with_token("tok123");
633        assert_eq!(svc.token().unwrap(), "tok123");
634    }
635
636    #[test]
637    fn with_token_refresher_calls_closure() {
638        let svc =
639            GcsArtifactService::new("b", "a").with_token_refresher(|| "dynamic-token".to_string());
640        assert_eq!(svc.token().unwrap(), "dynamic-token");
641    }
642
643    #[test]
644    fn implements_artifact_service_trait() {
645        fn _assert_trait(_: &dyn ArtifactService) {}
646        let svc = GcsArtifactService::new("b", "a");
647        _assert_trait(&svc);
648    }
649}
gemini_adk_rs/artifacts/gcs_service.rs

gemini_adk_rs/artifacts/
gcs_service.rs
