diff --git a/src/google/adk/artifacts/file_artifact_service.py b/src/google/adk/artifacts/file_artifact_service.py index 948a356d14..b84cf906bc 100644 --- a/src/google/adk/artifacts/file_artifact_service.py +++ b/src/google/adk/artifacts/file_artifact_service.py @@ -50,7 +50,10 @@ def _iter_artifact_dirs(root: Path) -> list[Path]: current = Path(dirpath) if (current / "versions").exists(): artifact_dirs.append(current) - dirnames.clear() + # An artifact directory doubles as the parent of anything nested under + # it ("doc" and "doc/nested"), so keep walking, skipping only the + # stored versions of this artifact. + dirnames[:] = [name for name in dirnames if name != "versions"] return artifact_dirs @@ -293,6 +296,28 @@ def _canonical_uri(artifact_dir: Path, version: int) -> str: return payload_path.resolve().as_uri() +def _prune_empty_dirs(leaf: Path, stop_at: Path) -> None: + """Removes `leaf` and any parents it leaves empty, stopping at `stop_at`. + + Filenames may contain "/", so the directory of an artifact doubles as the + parent directory of every artifact nested under it: "doc" is stored at + ``{scope}/doc`` and "doc/nested" at ``{scope}/doc/nested``. A directory may + therefore only be removed once it holds nothing, or deleting "doc" would + take "doc/nested" with it. + + Args: + leaf: Directory to remove, if it is empty. + stop_at: Scope root. It and everything above it are never removed. + """ + current = leaf + while current != stop_at and current.is_relative_to(stop_at): + try: + current.rmdir() # Only succeeds on an empty directory. + except OSError: + return + current = current.parent + + def _list_versions_on_disk(artifact_dir: Path) -> list[int]: """Returns sorted versions discovered under the artifact directory.""" versions_dir = _versions_dir(artifact_dir) @@ -703,9 +728,18 @@ def _delete_artifact_sync( session_id=session_id, filename=filename, ) - if artifact_dir.exists(): - shutil.rmtree(artifact_dir) - logger.debug("Deleted artifact %s at %s", filename, artifact_dir) + versions_dir = _versions_dir(artifact_dir) + if not versions_dir.exists(): + return + # Only this artifact's own versions go. Its directory may also be the + # parent of a nested artifact ("doc" vs "doc/nested"), so it is pruned + # separately and only if nothing is left under it. + shutil.rmtree(versions_dir) + scope_root = self._scope_root( + self._base_root(app_name, user_id), session_id, filename + ) + _prune_empty_dirs(artifact_dir, scope_root) + logger.debug("Deleted artifact %s at %s", filename, artifact_dir) @override async def list_versions( diff --git a/src/google/adk/artifacts/gcs_artifact_service.py b/src/google/adk/artifacts/gcs_artifact_service.py index ee4a6b0fa3..220106d6ea 100644 --- a/src/google/adk/artifacts/gcs_artifact_service.py +++ b/src/google/adk/artifacts/gcs_artifact_service.py @@ -41,6 +41,42 @@ logger = logging.getLogger("google_adk." + __name__) +def _parse_version(blob_name: str, prefix: str) -> Optional[int]: + """Extracts the version of an artifact from one of its blob names. + + GCS has a flat namespace, so listing by prefix is a plain string match with + no notion of nesting depth. Because filenames are allowed to contain "/", + the prefix of an artifact is also a prefix of every artifact nested under it: + scanning "a/" to find versions of "a" also returns "a/b/3", which is version + 3 of the distinct artifact "a/b". + + A blob only holds a version of the artifact denoted by ``prefix`` when its + name is exactly ``{prefix}{version}``, so anything with a further "/" in it + belongs to some other artifact and must be skipped. + + Args: + blob_name: The full name of the blob, which must start with ``prefix``. + prefix: The blob prefix of the artifact, including the trailing "/". + + Returns: + The version number, or None if the blob does not hold a version of this + artifact. + """ + suffix = blob_name[len(prefix) :] + if "/" in suffix: + # Belongs to a distinct artifact nested under this one. + return None + # int() also accepts surrounding whitespace, underscores and non-ASCII + # digits, none of which _get_blob_name can produce. + if not (suffix.isascii() and suffix.isdigit()): + logger.warning( + "Skipping blob %s because it does not end with a version number.", + blob_name, + ) + return None + return int(suffix) + + class GcsArtifactService(BaseArtifactService): """An artifact service implementation using Google Cloud Storage (GCS).""" @@ -357,12 +393,17 @@ def _list_versions( artifact. Returns an empty list if no versions are found. """ - prefix = self._get_blob_prefix(app_name, user_id, filename, session_id) - blobs = self.storage_client.list_blobs(self.bucket, prefix=f"{prefix}/") + prefix = ( + f"{self._get_blob_prefix(app_name, user_id, filename, session_id)}/" + ) + blobs = self.storage_client.list_blobs(self.bucket, prefix=prefix) versions = [] for blob in blobs: - *_, version = blob.name.split("/") - versions.append(int(version)) + version = _parse_version(blob.name, prefix) + if version is None: + continue + + versions.append(version) return versions def _get_artifact_version_sync( @@ -410,17 +451,14 @@ def _list_artifact_versions_sync( filename: str, ) -> list[ArtifactVersion]: """Lists all versions and their metadata of an artifact.""" - prefix = self._get_blob_prefix(app_name, user_id, filename, session_id) - blobs = self.storage_client.list_blobs(self.bucket, prefix=f"{prefix}/") + prefix = ( + f"{self._get_blob_prefix(app_name, user_id, filename, session_id)}/" + ) + blobs = self.storage_client.list_blobs(self.bucket, prefix=prefix) artifact_versions = [] for blob in blobs: - try: - version = int(blob.name.split("/")[-1]) - except ValueError: - logger.warning( - "Skipping blob %s because it does not end with a version number.", - blob.name, - ) + version = _parse_version(blob.name, prefix) + if version is None: continue canonical_uri = f"gs://{self.bucket_name}/{blob.name}" diff --git a/src/google/adk/integrations/skill_registry/gcp_skill_registry.py b/src/google/adk/integrations/skill_registry/gcp_skill_registry.py index d06b32d15e..4abd4eb9a5 100644 --- a/src/google/adk/integrations/skill_registry/gcp_skill_registry.py +++ b/src/google/adk/integrations/skill_registry/gcp_skill_registry.py @@ -64,8 +64,11 @@ async def get_skill(self, *, name: str) -> models.Skill: # be a single path segment before it is interpolated into the resource # path. Accept the same character set skill names are already held to; the # snake-or-kebab pattern is the superset of the two accepted spellings. + # fullmatch, not match: the pattern ends in `$`, which outside MULTILINE + # also matches just before a trailing newline, so `match` would let + # "my-skill\n" through into the resource name. # pylint: disable-next=protected-access - if not models._SNAKE_OR_KEBAB_NAME_PATTERN.match(name): + if not models._SNAKE_OR_KEBAB_NAME_PATTERN.fullmatch(name): raise ValueError( f"Invalid skill name {name!r}: name must be lowercase kebab-case" " (a-z, 0-9, hyphens) or snake_case (a-z, 0-9, underscores), with" diff --git a/src/google/adk/memory/vertex_ai_rag_memory_service.py b/src/google/adk/memory/vertex_ai_rag_memory_service.py index 1696fb3450..faa04d86af 100644 --- a/src/google/adk/memory/vertex_ai_rag_memory_service.py +++ b/src/google/adk/memory/vertex_ai_rag_memory_service.py @@ -20,6 +20,7 @@ import binascii from collections import OrderedDict import json +import logging import os import tempfile from typing import Optional @@ -37,6 +38,7 @@ from ..events.event import Event from ..sessions.session import Session +logger = logging.getLogger("google_adk." + __name__) _SOURCE_DISPLAY_NAME_PREFIX = "adk-memory-v1." @@ -178,8 +180,15 @@ async def add_session_to_memory(self, session: Session) -> None: if temp_file_path: try: os.remove(temp_file_path) - except FileNotFoundError: - pass + except OSError: + # Best effort: this runs in a finally, so raising here would + # displace the exception already propagating, and a cancelled + # upload can leave the worker thread still holding the file. + logger.warning( + "Could not remove the temporary transcript at %s", + temp_file_path, + exc_info=True, + ) @override async def search_memory( diff --git a/tests/unittests/artifacts/test_artifact_service.py b/tests/unittests/artifacts/test_artifact_service.py index 6913b20657..8278aaf6a9 100644 --- a/tests/unittests/artifacts/test_artifact_service.py +++ b/tests/unittests/artifacts/test_artifact_service.py @@ -51,6 +51,21 @@ class ArtifactServiceType(Enum): GCS = "GCS" +def _artifact_text(part: Optional[types.Part]) -> Optional[str]: + """Returns an artifact's text regardless of which Part field carries it. + + The GCS backend round-trips a text part as inline_data, so comparing Part + objects across backends fails on the field name rather than the content. + """ + if part is None: + return None + if part.text is not None: + return part.text + if part.inline_data is not None and part.inline_data.data is not None: + return part.inline_data.data.decode("utf-8") + return None + + class MockBlob: """Mocks a GCS Blob object. @@ -2006,3 +2021,216 @@ async def test_list_artifact_keys_survives_metadata_path_shadowed_by_dir( # The shadowed artifact has no readable metadata, so it is listed by its # scope-relative path rather than dropped or raised on. assert keys == ["user:a"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "service_type", + [ + ArtifactServiceType.IN_MEMORY, + ArtifactServiceType.GCS, + ArtifactServiceType.FILE, + ], +) +async def test_nested_artifact_does_not_leak_versions_into_parent( + service_type, artifact_service_factory +): + """A nested artifact must not contribute versions to its parent. + + Filenames may contain "/", so "doc" and "doc/nested" are two distinct + artifacts. On a flat keyspace the records of "doc/nested" live under the + prefix used to scan for versions of "doc", and must not be counted as + versions of "doc". + """ + artifact_service = artifact_service_factory(service_type) + app_name = "app0" + user_id = "user0" + session_id = "123" + parent = types.Part.from_text(text="parent v0") + + await artifact_service.save_artifact( + app_name=app_name, + user_id=user_id, + session_id=session_id, + filename="doc", + artifact=parent, + ) + # Give the nested artifact more versions than the parent has, so that a leak + # would push max(versions) past any version "doc" actually has. + for i in range(3): + await artifact_service.save_artifact( + app_name=app_name, + user_id=user_id, + session_id=session_id, + filename="doc/nested", + artifact=types.Part.from_text(text=f"nested v{i}"), + ) + + assert await artifact_service.list_versions( + app_name=app_name, + user_id=user_id, + session_id=session_id, + filename="doc", + ) == [0] + + # Loading without an explicit version resolves max(versions). A leaked + # version points at a record that does not exist, silently yielding None. + # Compared by content: the GCS backend round-trips a text part as + # inline_data, so the Part objects are not equal even when the bytes are. + assert _artifact_text( + await artifact_service.load_artifact( + app_name=app_name, + user_id=user_id, + session_id=session_id, + filename="doc", + ) + ) == _artifact_text(parent) + + # The next version of "doc" must be 1, not 3. + assert ( + await artifact_service.save_artifact( + app_name=app_name, + user_id=user_id, + session_id=session_id, + filename="doc", + artifact=types.Part.from_text(text="parent v1"), + ) + == 1 + ) + + # The nested artifact is unaffected. + assert await artifact_service.list_versions( + app_name=app_name, + user_id=user_id, + session_id=session_id, + filename="doc/nested", + ) == [0, 1, 2] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "service_type", + [ + ArtifactServiceType.IN_MEMORY, + ArtifactServiceType.GCS, + ArtifactServiceType.FILE, + ], +) +async def test_list_artifact_versions_excludes_nested_artifact( + service_type, artifact_service_factory +): + """Version metadata of a nested artifact must not surface under its parent.""" + artifact_service = artifact_service_factory(service_type) + app_name = "app0" + user_id = "user0" + session_id = "123" + + for filename in ("doc", "doc/nested"): + await artifact_service.save_artifact( + app_name=app_name, + user_id=user_id, + session_id=session_id, + filename=filename, + artifact=types.Part.from_text(text=filename), + ) + + versions = await artifact_service.list_artifact_versions( + app_name=app_name, + user_id=user_id, + session_id=session_id, + filename="doc", + ) + + assert [v.version for v in versions] == [0] + # The returned handle must address "doc", not the nested artifact. + if service_type == ArtifactServiceType.GCS: + assert versions[0].canonical_uri.endswith("/doc/0") + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "service_type", + [ + ArtifactServiceType.IN_MEMORY, + ArtifactServiceType.GCS, + ArtifactServiceType.FILE, + ], +) +async def test_delete_artifact_keeps_nested_artifact( + service_type, artifact_service_factory +): + """Deleting an artifact must not disturb artifacts nested under it.""" + artifact_service = artifact_service_factory(service_type) + app_name = "app0" + user_id = "user0" + session_id = "123" + nested = types.Part.from_text(text="nested v0") + + await artifact_service.save_artifact( + app_name=app_name, + user_id=user_id, + session_id=session_id, + filename="doc", + artifact=types.Part.from_text(text="parent v0"), + ) + await artifact_service.save_artifact( + app_name=app_name, + user_id=user_id, + session_id=session_id, + filename="doc/nested", + artifact=nested, + ) + + await artifact_service.delete_artifact( + app_name=app_name, + user_id=user_id, + session_id=session_id, + filename="doc", + ) + + assert not await artifact_service.list_versions( + app_name=app_name, + user_id=user_id, + session_id=session_id, + filename="doc", + ) + assert _artifact_text( + await artifact_service.load_artifact( + app_name=app_name, + user_id=user_id, + session_id=session_id, + filename="doc/nested", + ) + ) == _artifact_text(nested) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "service_type", + [ + ArtifactServiceType.IN_MEMORY, + ArtifactServiceType.GCS, + ArtifactServiceType.FILE, + ], +) +async def test_list_keys_includes_nested_artifact( + service_type, artifact_service_factory +): + """An artifact nested under another artifact must still be listed.""" + artifact_service = artifact_service_factory(service_type) + app_name = "app0" + user_id = "user0" + session_id = "123" + + for filename in ("doc", "doc/nested"): + await artifact_service.save_artifact( + app_name=app_name, + user_id=user_id, + session_id=session_id, + filename=filename, + artifact=types.Part.from_text(text=filename), + ) + + assert await artifact_service.list_artifact_keys( + app_name=app_name, user_id=user_id, session_id=session_id + ) == ["doc", "doc/nested"] diff --git a/tests/unittests/integrations/skill_registry/test_gcp_skill_registry.py b/tests/unittests/integrations/skill_registry/test_gcp_skill_registry.py index 7834c32013..7d75ce9bad 100644 --- a/tests/unittests/integrations/skill_registry/test_gcp_skill_registry.py +++ b/tests/unittests/integrations/skill_registry/test_gcp_skill_registry.py @@ -196,6 +196,10 @@ async def test_get_skill_raises_on_invalid_skill_name(mock_vertex_client): "my-skill/revisions/rev-123", "My-Skill", "", + # The pattern ends in `$`, which outside MULTILINE also matches just + # before a trailing newline, so `match` let this reach the resource + # name. + "my-skill\n", ], ) @pytest.mark.asyncio diff --git a/tests/unittests/memory/test_vertex_ai_rag_memory_service.py b/tests/unittests/memory/test_vertex_ai_rag_memory_service.py index 397776af66..214aa31d46 100644 --- a/tests/unittests/memory/test_vertex_ai_rag_memory_service.py +++ b/tests/unittests/memory/test_vertex_ai_rag_memory_service.py @@ -189,6 +189,34 @@ def upload_file(**_kwargs): assert not list(temp_dir.iterdir()) +@pytest.mark.asyncio +async def test_failed_temp_file_removal_does_not_mask_the_upload_error( + mocker, temp_dir +): + """Cleanup runs in a finally, so it must not displace the real exception. + + A cancelled upload can leave the worker thread still holding the file, and + on some platforms the remove then fails. Raising from the finally block + would replace the exception already propagating. + """ + + def upload_file(**_kwargs): + raise RuntimeError("upload failed") + + mocker.patch( + "google.adk.dependencies.vertexai.rag", + SimpleNamespace(upload_file=upload_file), + ) + mocker.patch( + "google.adk.memory.vertex_ai_rag_memory_service.os.remove", + side_effect=PermissionError("file still in use"), + ) + memory_service = VertexAiRagMemoryService(rag_corpus="corpus") + + with pytest.raises(RuntimeError, match="upload failed"): + await memory_service.add_session_to_memory(_session()) + + @pytest.mark.asyncio async def test_add_session_leaves_no_temp_file_when_corpus_missing(temp_dir): memory_service = VertexAiRagMemoryService(rag_corpus=None)