diff --git a/src/google/adk/artifacts/artifact_util.py b/src/google/adk/artifacts/artifact_util.py index 8227bbb92c4..57a2a2c2f96 100644 --- a/src/google/adk/artifacts/artifact_util.py +++ b/src/google/adk/artifacts/artifact_util.py @@ -143,6 +143,31 @@ def _is_drive_qualified(value: str) -> bool: return _WINDOWS_DRIVE_RE.match(value) is not None +def validate_session_id_segment(session_id: str) -> None: + """Validates a session_id that will be used as a literal storage segment. + + In addition to the checks in `validate_path_segment`, rejects the literal + value "user". Backends that lay out session-scoped and user-scoped + artifacts in the same flat namespace (in-memory, GCS) use that exact + string as a reserved segment marking user-scoped artifacts, so a session + actually named "user" would silently write into -- and read out of -- that + reserved namespace instead of its own. + + Args: + session_id: The caller-supplied session id. + + Raises: + InputValidationError: If `session_id` fails `validate_path_segment`, or + is the reserved value "user". + """ + validate_path_segment(session_id, "session_id") + if session_id == "user": + raise input_validation_error.InputValidationError( + "session_id must not be the reserved value 'user', which this" + " backend uses internally to mark user-scoped artifacts." + ) + + def validate_path_segment(value: str, field_name: str) -> None: """Rejects values that could alter the constructed path. diff --git a/src/google/adk/artifacts/gcs_artifact_service.py b/src/google/adk/artifacts/gcs_artifact_service.py index 26a3ea21db3..57afaad1022 100644 --- a/src/google/adk/artifacts/gcs_artifact_service.py +++ b/src/google/adk/artifacts/gcs_artifact_service.py @@ -217,7 +217,7 @@ def _get_blob_prefix( raise InputValidationError( "Session ID must be provided for session-scoped artifacts." ) - artifact_util.validate_path_segment(session_id, "session_id") + artifact_util.validate_session_id_segment(session_id) return f"{app_name}/{user_id}/{session_id}/{filename}" def _get_blob_name( @@ -418,7 +418,7 @@ def _list_artifact_keys( artifact_util.validate_path_segment(app_name, "app_name") artifact_util.validate_path_segment(user_id, "user_id") if session_id is not None: - artifact_util.validate_path_segment(session_id, "session_id") + artifact_util.validate_session_id_segment(session_id) filenames = set() if session_id: diff --git a/src/google/adk/artifacts/in_memory_artifact_service.py b/src/google/adk/artifacts/in_memory_artifact_service.py index 2ed4e0a9ac6..49683e08955 100644 --- a/src/google/adk/artifacts/in_memory_artifact_service.py +++ b/src/google/adk/artifacts/in_memory_artifact_service.py @@ -95,7 +95,7 @@ def _artifact_path( raise InputValidationError( "Session ID must be provided for session-scoped artifacts." ) - artifact_util.validate_path_segment(session_id, "session_id") + artifact_util.validate_session_id_segment(session_id) return f"{app_name}/{user_id}/{session_id}/{filename}" @override @@ -223,7 +223,7 @@ async def list_artifact_keys( artifact_util.validate_path_segment(app_name, "app_name") artifact_util.validate_path_segment(user_id, "user_id") if session_id is not None: - artifact_util.validate_path_segment(session_id, "session_id") + artifact_util.validate_session_id_segment(session_id) usernamespace_prefix = f"{app_name}/{user_id}/user/" session_prefix = ( f"{app_name}/{user_id}/{session_id}/" if session_id else None diff --git a/tests/unittests/artifacts/test_artifact_service.py b/tests/unittests/artifacts/test_artifact_service.py index a6dd4b642c9..8297319ffd3 100644 --- a/tests/unittests/artifacts/test_artifact_service.py +++ b/tests/unittests/artifacts/test_artifact_service.py @@ -286,6 +286,82 @@ async def test_save_load_delete(service_type, artifact_service_factory): ) +@pytest.mark.asyncio +@pytest.mark.parametrize( + "service_type", + [ + ArtifactServiceType.IN_MEMORY, + ArtifactServiceType.GCS, + ], +) +async def test_save_artifact_rejects_reserved_user_as_session_id( + service_type, artifact_service_factory +): + """IN_MEMORY and GCS lay session-scoped and user-scoped artifacts out in + the same flat namespace, using the literal segment "user" to mark + user-scoped ones. A session actually named "user" must be rejected rather + than silently colliding with that reserved segment.""" + artifact_service = artifact_service_factory(service_type) + + with pytest.raises(InputValidationError, match="reserved value 'user'"): + await artifact_service.save_artifact( + app_name="app0", + user_id="user0", + session_id="user", + filename="report.txt", + artifact=types.Part(text="hello"), + ) + + +@pytest.mark.asyncio +async def test_file_allows_reserved_user_as_session_id( + artifact_service_factory, +): + """Unlike IN_MEMORY and GCS, FILE lays session-scoped artifacts out under + their own `sessions//` subtree, distinct from the user-scoped + `artifacts/` subtree, so a session literally named "user" cannot collide + with it and is not rejected.""" + artifact_service = artifact_service_factory(ArtifactServiceType.FILE) + + await artifact_service.save_artifact( + app_name="app0", + user_id="user0", + session_id="user", + filename="report.txt", + artifact=types.Part(text="hello"), + ) + loaded = await artifact_service.load_artifact( + app_name="app0", + user_id="user0", + session_id="user", + filename="report.txt", + ) + assert loaded == types.Part(text="hello") + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "service_type", + [ + ArtifactServiceType.IN_MEMORY, + ArtifactServiceType.GCS, + ], +) +async def test_list_artifact_keys_rejects_reserved_user_as_session_id( + service_type, artifact_service_factory +): + """A session literally named "user" must be rejected by + list_artifact_keys too, not just by save/load/delete -- otherwise a + caller's listing could silently return another session's (or the + user-scope's) filenames.""" + artifact_service = artifact_service_factory(service_type) + + with pytest.raises(InputValidationError, match="reserved value 'user'"): + await artifact_service.list_artifact_keys( + app_name="app0", user_id="user0", session_id="user" + ) + + @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 d5b79242d3c..6731182aa80 100644 --- a/tests/unittests/artifacts/test_artifact_util.py +++ b/tests/unittests/artifacts/test_artifact_util.py @@ -192,6 +192,20 @@ def test_validate_path_segment_invalid(value, field_name): artifact_util.validate_path_segment(value, field_name) +def test_validate_session_id_segment_rejects_reserved_user(): + with pytest.raises(InputValidationError, match="reserved value 'user'"): + artifact_util.validate_session_id_segment("user") + + +def test_validate_session_id_segment_allows_ordinary_values(): + artifact_util.validate_session_id_segment("session1") + + +def test_validate_session_id_segment_still_runs_path_segment_checks(): + with pytest.raises(InputValidationError, match="must not be empty"): + artifact_util.validate_session_id_segment("") + + @pytest.mark.parametrize( "caller_session_id, uri_session_id", [