diff --git a/src/google/adk/artifacts/artifact_util.py b/src/google/adk/artifacts/artifact_util.py index 8227bbb92c..f7efcd55fb 100644 --- a/src/google/adk/artifacts/artifact_util.py +++ b/src/google/adk/artifacts/artifact_util.py @@ -43,6 +43,24 @@ class ParsedArtifactUri(NamedTuple): ) +def normalize_session_id(session_id: str | None) -> str | None: + """Normalizes a caller-supplied session id. + + Strips surrounding whitespace the same way the session services do, so an + artifact saved against a padded session id lands in the same storage + namespace as the (normalized) session it belongs to, instead of a sibling + namespace no caller using the trimmed id can ever reach. + + Args: + session_id: The caller-supplied session id, or None for a user-scoped + artifact. + + Returns: + The stripped session id, or None if `session_id` was None. + """ + return session_id.strip() if session_id is not None else None + + def parse_artifact_uri(uri: str) -> ParsedArtifactUri | None: """Parses an artifact URI. @@ -60,7 +78,7 @@ def parse_artifact_uri(uri: str) -> ParsedArtifactUri | None: return ParsedArtifactUri( app_name=match.group(1), user_id=match.group(2), - session_id=match.group(3), + session_id=normalize_session_id(match.group(3)), filename=match.group(4), version=int(match.group(5)), ) @@ -97,6 +115,7 @@ def get_artifact_uri( Returns: The constructed artifact URI. """ + session_id = normalize_session_id(session_id) if session_id: return f"artifact://apps/{app_name}/users/{user_id}/sessions/{session_id}/artifacts/{filename}/versions/{version}" else: diff --git a/src/google/adk/artifacts/file_artifact_service.py b/src/google/adk/artifacts/file_artifact_service.py index 9c53754bd5..e66c5bcae9 100644 --- a/src/google/adk/artifacts/file_artifact_service.py +++ b/src/google/adk/artifacts/file_artifact_service.py @@ -247,6 +247,7 @@ def _user_artifacts_dir(base_root: Path) -> Path: def _session_artifacts_dir(base_root: Path, session_id: str) -> Path: """Returns the path that stores session-scoped artifacts.""" + session_id = session_id.strip() artifact_util.validate_path_segment(session_id, "session_id") return base_root / "sessions" / session_id / "artifacts" diff --git a/src/google/adk/artifacts/gcs_artifact_service.py b/src/google/adk/artifacts/gcs_artifact_service.py index 26a3ea21db..a92fb22e4a 100644 --- a/src/google/adk/artifacts/gcs_artifact_service.py +++ b/src/google/adk/artifacts/gcs_artifact_service.py @@ -213,6 +213,7 @@ def _get_blob_prefix( if self._file_has_user_namespace(filename): return f"{app_name}/{user_id}/user/{filename}" + session_id = artifact_util.normalize_session_id(session_id) if session_id is None: raise InputValidationError( "Session ID must be provided for session-scoped artifacts." @@ -254,6 +255,7 @@ def _save_artifact( custom_metadata: Optional[dict[str, Any]] = None, ) -> int: artifact = ensure_part(artifact) + session_id = artifact_util.normalize_session_id(session_id) versions = self._list_versions( app_name=app_name, user_id=user_id, @@ -337,6 +339,7 @@ def _load_artifact( filename: str, version: Optional[int] = None, ) -> Optional[types.Part]: + session_id = artifact_util.normalize_session_id(session_id) if version is None: versions = self._list_versions( app_name=app_name, @@ -417,6 +420,7 @@ def _list_artifact_keys( ) -> list[str]: artifact_util.validate_path_segment(app_name, "app_name") artifact_util.validate_path_segment(user_id, "user_id") + session_id = artifact_util.normalize_session_id(session_id) if session_id is not None: artifact_util.validate_path_segment(session_id, "session_id") filenames = set() @@ -622,6 +626,7 @@ def _get_authenticated_url_sync( max_depth: int = _MAX_ARTIFACT_REFERENCE_DEPTH, ) -> Optional[str]: """Generates an authenticated browser URL for an artifact.""" + session_id = artifact_util.normalize_session_id(session_id) if version is None: versions = self._list_versions( app_name=app_name, @@ -729,6 +734,7 @@ def _get_signed_url_sync( max_depth: int = _MAX_ARTIFACT_REFERENCE_DEPTH, ) -> Optional[str]: """Generates a time-limited signed URL for an artifact.""" + session_id = artifact_util.normalize_session_id(session_id) if version is None: versions = self._list_versions( app_name=app_name, diff --git a/src/google/adk/artifacts/in_memory_artifact_service.py b/src/google/adk/artifacts/in_memory_artifact_service.py index 2ed4e0a9ac..259a536951 100644 --- a/src/google/adk/artifacts/in_memory_artifact_service.py +++ b/src/google/adk/artifacts/in_memory_artifact_service.py @@ -91,6 +91,7 @@ def _artifact_path( if self._file_has_user_namespace(filename): return f"{app_name}/{user_id}/user/{filename}" + session_id = artifact_util.normalize_session_id(session_id) if session_id is None: raise InputValidationError( "Session ID must be provided for session-scoped artifacts." @@ -110,6 +111,7 @@ async def save_artifact( custom_metadata: Optional[dict[str, Any]] = None, ) -> int: artifact = ensure_part(artifact) + session_id = artifact_util.normalize_session_id(session_id) path = self._artifact_path(app_name, user_id, filename, session_id) if path not in self.artifacts: self.artifacts[path] = [] @@ -167,6 +169,7 @@ async def load_artifact( session_id: Optional[str] = None, version: Optional[int] = None, ) -> Optional[types.Part]: + session_id = artifact_util.normalize_session_id(session_id) path = self._artifact_path(app_name, user_id, filename, session_id) versions = self.artifacts.get(path) if not versions: @@ -222,6 +225,7 @@ async def list_artifact_keys( ) -> list[str]: artifact_util.validate_path_segment(app_name, "app_name") artifact_util.validate_path_segment(user_id, "user_id") + session_id = artifact_util.normalize_session_id(session_id) if session_id is not None: artifact_util.validate_path_segment(session_id, "session_id") usernamespace_prefix = f"{app_name}/{user_id}/user/" diff --git a/tests/unittests/artifacts/test_artifact_service.py b/tests/unittests/artifacts/test_artifact_service.py index a6dd4b642c..a17e853004 100644 --- a/tests/unittests/artifacts/test_artifact_service.py +++ b/tests/unittests/artifacts/test_artifact_service.py @@ -286,6 +286,125 @@ async def test_save_load_delete(service_type, artifact_service_factory): ) +@pytest.mark.asyncio +@pytest.mark.parametrize( + "service_type", + [ + ArtifactServiceType.IN_MEMORY, + ArtifactServiceType.GCS, + ArtifactServiceType.FILE, + ], +) +async def test_padded_session_id_lands_in_same_namespace_as_trimmed( + service_type, artifact_service_factory +): + """An artifact saved against a whitespace-padded session id must be + reachable, listable, and deletable using the trimmed id, since that is the + id the corresponding session is actually stored under (session services + normalize session_id the same way).""" + artifact_service = artifact_service_factory(service_type) + artifact = types.Part(text="hello") + + await artifact_service.save_artifact( + app_name="app0", + user_id="user0", + session_id="sess0\n", + filename="report.txt", + artifact=artifact, + ) + + # Reachable under the trimmed id, the one the session itself is keyed on. + assert ( + await artifact_service.load_artifact( + app_name="app0", + user_id="user0", + session_id="sess0", + filename="report.txt", + ) + == artifact + ) + assert "report.txt" in await artifact_service.list_artifact_keys( + app_name="app0", user_id="user0", session_id="sess0" + ) + + # Also reachable under the original padded id: both forms must resolve to + # the same stored artifact rather than the padded id silently shadowing it + # in a namespace only the padded id itself could ever reach again. + assert ( + await artifact_service.load_artifact( + app_name="app0", + user_id="user0", + session_id="sess0\n", + filename="report.txt", + ) + == artifact + ) + + await artifact_service.delete_artifact( + app_name="app0", + user_id="user0", + session_id="sess0", + filename="report.txt", + ) + assert not await artifact_service.load_artifact( + app_name="app0", + user_id="user0", + session_id="sess0\n", + filename="report.txt", + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "service_type", + [ + ArtifactServiceType.IN_MEMORY, + ArtifactServiceType.GCS, + ], +) +async def test_artifact_reference_allows_padded_session_id_at_call_site( + service_type, artifact_service_factory +): + """A caller that consistently uses a padded session id must still be able + to save and resolve an artifact reference within that session: the scope + check must compare normalized ids on both sides, not the caller's raw + string against a URI minted from the (already normalized) stored id.""" + artifact_service = artifact_service_factory(service_type) + + await artifact_service.save_artifact( + app_name="app0", + user_id="user0", + session_id="sess0\n", + filename="source.txt", + artifact=types.Part(text="hello"), + ) + + ref = types.Part( + file_data=types.FileData( + file_uri=( + "artifact://apps/app0/users/user0/sessions/sess0/" + "artifacts/source.txt/versions/0" + ), + mime_type="text/plain", + ) + ) + await artifact_service.save_artifact( + app_name="app0", + user_id="user0", + session_id="sess0\n", + filename="ref.txt", + artifact=ref, + ) + + loaded = await artifact_service.load_artifact( + app_name="app0", + user_id="user0", + session_id="sess0\n", + filename="ref.txt", + ) + assert loaded == types.Part(text="hello") + + @pytest.mark.asyncio async def test_in_memory_loads_nested_artifact_reference( artifact_service_factory, diff --git a/tests/unittests/artifacts/test_artifact_util.py b/tests/unittests/artifacts/test_artifact_util.py index d5b79242d3..43986b1a62 100644 --- a/tests/unittests/artifacts/test_artifact_util.py +++ b/tests/unittests/artifacts/test_artifact_util.py @@ -112,6 +112,40 @@ def test_get_user_scoped_artifact_uri(): assert uri == "artifact://apps/app2/users/user2/artifacts/file2/versions/456" +def test_normalize_session_id_strips_whitespace(): + assert artifact_util.normalize_session_id(" sess0\n") == "sess0" + + +def test_normalize_session_id_passes_none_through(): + assert artifact_util.normalize_session_id(None) is None + + +def test_get_artifact_uri_normalizes_padded_session_id(): + """A padded session id must not leak into the constructed URI, since the + session it points at is stored under the trimmed id.""" + uri = artifact_util.get_artifact_uri( + app_name="app1", + user_id="user1", + session_id=" session1\n", + filename="file1", + version=123, + ) + assert ( + uri + == "artifact://apps/app1/users/user1/sessions/session1/artifacts/file1/versions/123" + ) + + +def test_parse_artifact_uri_normalizes_a_legacy_padded_session_id(): + """A URI minted before this fix could carry a padded session id in its own + path segment; parsing it must still yield the trimmed id so downstream + scope checks compare like with like.""" + uri = "artifact://apps/app1/users/user1/sessions/session1\n/artifacts/file1/versions/123" + parsed = artifact_util.parse_artifact_uri(uri) + assert parsed is not None + assert parsed.session_id == "session1" + + def test_is_artifact_ref_true(): """Tests is_artifact_ref with a valid artifact reference.""" artifact = types.Part(