From 326279ab868530663b1fba3151e73c8267de74b9 Mon Sep 17 00:00:00 2001 From: GWeale Date: Mon, 24 Aug 2026 22:59:33 +0000 Subject: [PATCH] fix: close the residual gaps from the v1 backport review Three findings from the review of the v1 security backport batch that were not caused by those PRs. Deleting an artifact removed its whole directory, and a nested artifact lives beneath its parent, so deleting doc destroyed doc/nested; listing never reported it either. Deletion now takes only the artifact's own versions directory and prunes what it leaves empty. The GCS backend had the same conflation in its version scan, where a nested artifact's versions were reported as the parent's. A skill name could carry a trailing newline into the resource name, because the pattern ends in $ and was applied with match rather than fullmatch. The RAG temp-file cleanup caught only FileNotFoundError. It runs in a finally, so any other OSError displaced the exception already propagating, including CancelledError. --- .../adk/artifacts/file_artifact_service.py | 42 +++- .../adk/artifacts/gcs_artifact_service.py | 64 ++++- .../skill_registry/gcp_skill_registry.py | 5 +- .../memory/vertex_ai_rag_memory_service.py | 13 +- .../artifacts/test_artifact_service.py | 228 ++++++++++++++++++ .../skill_registry/test_gcp_skill_registry.py | 4 + .../test_vertex_ai_rag_memory_service.py | 28 +++ 7 files changed, 364 insertions(+), 20 deletions(-) diff --git a/src/google/adk/artifacts/file_artifact_service.py b/src/google/adk/artifacts/file_artifact_service.py index 948a356d14e..b84cf906bc0 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 ee4a6b0fa3e..220106d6eab 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 d06b32d15ef..4abd4eb9a5f 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 1696fb34508..faa04d86af8 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 6913b20657d..8278aaf6a96 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 7834c320130..7d75ce9bad4 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 397776af66f..214aa31d463 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)