From be55e86c1dede43aa235e556132961ab36c30a02 Mon Sep 17 00:00:00 2001 From: Prajwal Date: Wed, 19 Aug 2026 22:08:12 +0530 Subject: [PATCH] feat(artifacts): add get_authenticated_url and get_signed_url to GcsArtifactService Add methods to GcsArtifactService for generating Google Cloud Storage authenticated browser URLs (https://storage.cloud.google.com/...) and time-limited signed URLs directly for client and frontend consumption. --- .../adk/artifacts/gcs_artifact_service.py | 145 ++++++++++++++++++ .../artifacts/test_artifact_service.py | 96 ++++++++++++ 2 files changed, 241 insertions(+) diff --git a/src/google/adk/artifacts/gcs_artifact_service.py b/src/google/adk/artifacts/gcs_artifact_service.py index c554024a67..ea1f1a67c5 100644 --- a/src/google/adk/artifacts/gcs_artifact_service.py +++ b/src/google/adk/artifacts/gcs_artifact_service.py @@ -24,6 +24,7 @@ from __future__ import annotations import asyncio +import datetime import logging from typing import Any from typing import Optional @@ -605,3 +606,147 @@ async def get_artifact_version( filename, version, ) + + def _get_authenticated_url_sync( + self, + app_name: str, + user_id: str, + session_id: Optional[str], + filename: str, + version: Optional[int] = None, + ) -> Optional[str]: + """Generates an authenticated browser URL for an artifact.""" + if version is None: + versions = self._list_versions( + app_name=app_name, + user_id=user_id, + session_id=session_id, + filename=filename, + ) + if not versions: + return None + version = max(versions) + + blob_name = self._get_blob_name( + app_name, user_id, filename, version, session_id + ) + blob = self.bucket.get_blob(blob_name) + if not blob: + return None + + return f"https://storage.cloud.google.com/{self.bucket_name}/{blob_name}" + + async def get_authenticated_url( + self, + *, + app_name: str, + user_id: str, + filename: str, + session_id: Optional[str] = None, + version: Optional[int] = None, + ) -> Optional[str]: + """Generates an authenticated browser URL for an artifact. + + The URL returned requires the user to be authenticated with a Google + Account that has read permission for the object. + + Args: + app_name: The name of the application. + user_id: The ID of the user who owns the artifact. + filename: The name of the artifact file. + session_id: The ID of the session (ignored for user-namespaced files). + version: The version of the artifact. If None, the latest version will be used. + + Returns: + The authenticated GCS URL (https://storage.cloud.google.com/...), or None if the artifact does not exist. + """ + return await asyncio.to_thread( + self._get_authenticated_url_sync, + app_name, + user_id, + session_id, + filename, + version, + ) + + def _get_signed_url_sync( + self, + app_name: str, + user_id: str, + session_id: Optional[str], + filename: str, + version: Optional[int] = None, + expiration: Optional[ + Union[datetime.datetime, datetime.timedelta, int] + ] = None, + method: str = "GET", + kwargs: Optional[dict[str, Any]] = None, + ) -> Optional[str]: + """Generates a time-limited signed URL for an artifact.""" + if version is None: + versions = self._list_versions( + app_name=app_name, + user_id=user_id, + session_id=session_id, + filename=filename, + ) + if not versions: + return None + version = max(versions) + + blob_name = self._get_blob_name( + app_name, user_id, filename, version, session_id + ) + blob = self.bucket.get_blob(blob_name) + if not blob: + return None + + if expiration is None: + expiration = datetime.timedelta(hours=1) + + return blob.generate_signed_url( + expiration=expiration, + method=method, + **(kwargs or {}), + ) + + async def get_signed_url( + self, + *, + app_name: str, + user_id: str, + filename: str, + session_id: Optional[str] = None, + version: Optional[int] = None, + expiration: Optional[ + Union[datetime.datetime, datetime.timedelta, int] + ] = None, + method: str = "GET", + **kwargs: Any, + ) -> Optional[str]: + """Generates a time-limited signed URL for an artifact. + + Args: + app_name: The name of the application. + user_id: The ID of the user who owns the artifact. + filename: The name of the artifact file. + session_id: The ID of the session (ignored for user-namespaced files). + version: The version of the artifact. If None, the latest version will be used. + expiration: Time when the signed URL expires (datetime, timedelta, or epoch seconds). Defaults to 1 hour if not specified. + method: HTTP method allowed for the signed URL (default: "GET"). + **kwargs: Additional keyword arguments forwarded to `Blob.generate_signed_url`. + + Returns: + The signed URL string, or None if the artifact does not exist. + """ + return await asyncio.to_thread( + self._get_signed_url_sync, + app_name, + user_id, + session_id, + filename, + version, + expiration, + method, + kwargs, + ) diff --git a/tests/unittests/artifacts/test_artifact_service.py b/tests/unittests/artifacts/test_artifact_service.py index de9dff6b0b..43e5c3c8a3 100644 --- a/tests/unittests/artifacts/test_artifact_service.py +++ b/tests/unittests/artifacts/test_artifact_service.py @@ -17,6 +17,7 @@ """Tests for the artifact service.""" from datetime import datetime +from datetime import timedelta import enum import json from pathlib import Path @@ -111,6 +112,15 @@ def delete(self) -> None: self.content = None self.content_type = None + def generate_signed_url( + self, + expiration: Any = None, + method: str = "GET", + **kwargs: Any, + ) -> str: + """Mocks generating a signed URL for the blob.""" + return f"https://storage.googleapis.com/test_bucket/{self.name}?signed=true&method={method}" + class MockBucket: """Mocks a GCS Bucket object.""" @@ -2126,6 +2136,92 @@ async def test_gcs_load_artifact_returns_none_for_missing_version() -> None: ) +@pytest.mark.asyncio # type: ignore[untyped-decorator] +async def test_gcs_get_authenticated_url_latest_version() -> None: + """GcsArtifactService generates an authenticated URL for latest version.""" + service = mock_gcs_artifact_service() # type: ignore[no-untyped-call] + scope = {"app_name": "app", "user_id": "user1", "session_id": "sess1"} + + await service.save_artifact( + **scope, filename="notes.txt", artifact=types.Part.from_text(text="v0") + ) + await service.save_artifact( + **scope, filename="notes.txt", artifact=types.Part.from_text(text="v1") + ) + + url = await service.get_authenticated_url(**scope, filename="notes.txt") + assert ( + url + == "https://storage.cloud.google.com/test_bucket/app/user1/sess1/notes.txt/1" + ) + + +@pytest.mark.asyncio # type: ignore[untyped-decorator] +async def test_gcs_get_authenticated_url_specific_version() -> None: + """GcsArtifactService generates an authenticated URL for a specific version.""" + service = mock_gcs_artifact_service() # type: ignore[no-untyped-call] + scope = {"app_name": "app", "user_id": "user1", "session_id": "sess1"} + + await service.save_artifact( + **scope, filename="notes.txt", artifact=types.Part.from_text(text="v0") + ) + await service.save_artifact( + **scope, filename="notes.txt", artifact=types.Part.from_text(text="v1") + ) + + url = await service.get_authenticated_url( + **scope, filename="notes.txt", version=0 + ) + assert ( + url + == "https://storage.cloud.google.com/test_bucket/app/user1/sess1/notes.txt/0" + ) + + +@pytest.mark.asyncio # type: ignore[untyped-decorator] +async def test_gcs_get_authenticated_url_returns_none_for_missing() -> None: + """GcsArtifactService returns None for authenticated URL of missing artifact.""" + service = mock_gcs_artifact_service() # type: ignore[no-untyped-call] + scope = {"app_name": "app", "user_id": "user1", "session_id": "sess1"} + + assert ( + await service.get_authenticated_url(**scope, filename="nonexistent.txt") + is None + ) + + +@pytest.mark.asyncio # type: ignore[untyped-decorator] +async def test_gcs_get_signed_url_latest_version() -> None: + """GcsArtifactService generates a signed URL for latest version.""" + service = mock_gcs_artifact_service() # type: ignore[no-untyped-call] + scope = {"app_name": "app", "user_id": "user1", "session_id": "sess1"} + + await service.save_artifact( + **scope, filename="notes.txt", artifact=types.Part.from_text(text="v0") + ) + + url = await service.get_signed_url( + **scope, + filename="notes.txt", + expiration=timedelta(hours=2), + ) + assert ( + url + == "https://storage.googleapis.com/test_bucket/app/user1/sess1/notes.txt/0?signed=true&method=GET" + ) + + +@pytest.mark.asyncio # type: ignore[untyped-decorator] +async def test_gcs_get_signed_url_returns_none_for_missing() -> None: + """GcsArtifactService returns None for signed URL of missing artifact.""" + service = mock_gcs_artifact_service() # type: ignore[no-untyped-call] + scope = {"app_name": "app", "user_id": "user1", "session_id": "sess1"} + + assert ( + await service.get_signed_url(**scope, filename="nonexistent.txt") is None + ) + + @pytest.mark.asyncio @pytest.mark.parametrize( "service_type",