Skip to content

Commit affda8e

Browse files
its-amannits-amann
committed
Merge remote-tracking branch 'origin/main' into fix-thought-signature-pruning
2 parents bd33c6c + f863150 commit affda8e

5 files changed

Lines changed: 352 additions & 3 deletions

File tree

‎src/google/adk/artifacts/artifact_util.py‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,8 @@
2020

2121
from google.genai import types
2222

23+
from ..errors import input_validation_error
24+
2325

2426
class ParsedArtifactUri(NamedTuple):
2527
"""The result of parsing an artifact URI."""
@@ -113,3 +115,22 @@ def is_artifact_ref(artifact: types.Part) -> bool:
113115
and artifact.file_data.file_uri
114116
and artifact.file_data.file_uri.startswith("artifact://")
115117
)
118+
119+
120+
def validate_artifact_reference_scope(
121+
*,
122+
app_name: str,
123+
user_id: str,
124+
session_id: str | None,
125+
parsed_uri: ParsedArtifactUri,
126+
) -> None:
127+
"""Ensures artifact references cannot escape the caller's scope."""
128+
if parsed_uri.app_name != app_name or parsed_uri.user_id != user_id:
129+
raise input_validation_error.InputValidationError(
130+
"Artifact references must stay within the same app and user scope."
131+
)
132+
if parsed_uri.session_id is not None and parsed_uri.session_id != session_id:
133+
raise input_validation_error.InputValidationError(
134+
"Session-scoped artifact references must stay within the same"
135+
" session scope."
136+
)

‎src/google/adk/artifacts/gcs_artifact_service.py‎

Lines changed: 14 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -252,10 +252,17 @@ def _save_artifact(
252252
if not file_uri:
253253
raise InputValidationError("Artifact file_data must have a file_uri.")
254254
if artifact_util.is_artifact_ref(artifact):
255-
if not artifact_util.parse_artifact_uri(file_uri):
255+
parsed_uri = artifact_util.parse_artifact_uri(file_uri)
256+
if not parsed_uri:
256257
raise InputValidationError(
257258
f"Invalid artifact reference URI: {file_uri}"
258259
)
260+
artifact_util.validate_artifact_reference_scope(
261+
app_name=app_name,
262+
user_id=user_id,
263+
session_id=session_id,
264+
parsed_uri=parsed_uri,
265+
)
259266
# Store the URI and mime_type (if any) as blob metadata; no content to upload.
260267
metadata = {
261268
**(blob.metadata or {}),
@@ -315,6 +322,12 @@ def _load_artifact(
315322
raise InputValidationError(
316323
f"Invalid artifact reference URI: {file_uri}"
317324
)
325+
artifact_util.validate_artifact_reference_scope(
326+
app_name=app_name,
327+
user_id=user_id,
328+
session_id=session_id,
329+
parsed_uri=parsed_uri,
330+
)
318331
return self._load_artifact(
319332
app_name=parsed_uri.app_name,
320333
user_id=parsed_uri.user_id,

‎src/google/adk/artifacts/in_memory_artifact_service.py‎

Lines changed: 16 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -128,10 +128,19 @@ async def save_artifact(
128128
artifact_version.mime_type = "text/plain"
129129
elif artifact.file_data is not None:
130130
if artifact_util.is_artifact_ref(artifact):
131-
if not artifact_util.parse_artifact_uri(artifact.file_data.file_uri):
131+
parsed_uri = artifact_util.parse_artifact_uri(
132+
artifact.file_data.file_uri
133+
)
134+
if not parsed_uri:
132135
raise InputValidationError(
133136
f"Invalid artifact reference URI: {artifact.file_data.file_uri}"
134137
)
138+
artifact_util.validate_artifact_reference_scope(
139+
app_name=app_name,
140+
user_id=user_id,
141+
session_id=session_id,
142+
parsed_uri=parsed_uri,
143+
)
135144
# If it's a valid artifact URI, we store the artifact part as-is.
136145
# And we don't know the mime type until we load it.
137146
else:
@@ -180,6 +189,12 @@ async def load_artifact(
180189
"Invalid artifact reference URI:"
181190
f" {artifact_data.file_data.file_uri}"
182191
)
192+
artifact_util.validate_artifact_reference_scope(
193+
app_name=app_name,
194+
user_id=user_id,
195+
session_id=session_id,
196+
parsed_uri=parsed_uri,
197+
)
183198
return await self.load_artifact(
184199
app_name=parsed_uri.app_name,
185200
user_id=parsed_uri.user_id,

‎tests/unittests/artifacts/test_artifact_service.py‎

Lines changed: 250 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -905,6 +905,256 @@ async def test_file_save_artifact_rejects_absolute_path_within_scope(tmp_path):
905905
)
906906

907907

908+
@pytest.mark.asyncio
909+
@pytest.mark.parametrize(
910+
"service_type",
911+
[
912+
ArtifactServiceType.IN_MEMORY,
913+
ArtifactServiceType.GCS,
914+
],
915+
)
916+
async def test_artifact_reference_allows_same_session_scope(
917+
service_type, artifact_service_factory
918+
):
919+
"""ArtifactService allows references inside the same session scope."""
920+
artifact_service = artifact_service_factory(service_type)
921+
922+
await artifact_service.save_artifact(
923+
app_name="app0",
924+
user_id="user0",
925+
session_id="sess0",
926+
filename="source.txt",
927+
artifact=types.Part(text="hello"),
928+
)
929+
930+
ref = types.Part(
931+
file_data=types.FileData(
932+
file_uri=(
933+
"artifact://apps/app0/users/user0/sessions/sess0/"
934+
"artifacts/source.txt/versions/0"
935+
),
936+
mime_type="text/plain",
937+
)
938+
)
939+
await artifact_service.save_artifact(
940+
app_name="app0",
941+
user_id="user0",
942+
session_id="sess0",
943+
filename="ref.txt",
944+
artifact=ref,
945+
)
946+
947+
loaded = await artifact_service.load_artifact(
948+
app_name="app0",
949+
user_id="user0",
950+
session_id="sess0",
951+
filename="ref.txt",
952+
)
953+
assert loaded == types.Part(text="hello")
954+
955+
956+
@pytest.mark.asyncio
957+
@pytest.mark.parametrize(
958+
"service_type",
959+
[
960+
ArtifactServiceType.IN_MEMORY,
961+
ArtifactServiceType.GCS,
962+
],
963+
)
964+
async def test_artifact_reference_allows_same_user_user_scope(
965+
service_type, artifact_service_factory
966+
):
967+
"""ArtifactService allows references to user-scoped files from same user."""
968+
artifact_service = artifact_service_factory(service_type)
969+
970+
await artifact_service.save_artifact(
971+
app_name="app0",
972+
user_id="user0",
973+
session_id="sess0",
974+
filename="user:profile.txt",
975+
artifact=types.Part(text="profile"),
976+
)
977+
978+
ref = types.Part(
979+
file_data=types.FileData(
980+
file_uri=(
981+
"artifact://apps/app0/users/user0/artifacts/"
982+
"user:profile.txt/versions/0"
983+
),
984+
mime_type="text/plain",
985+
)
986+
)
987+
await artifact_service.save_artifact(
988+
app_name="app0",
989+
user_id="user0",
990+
session_id="sess1",
991+
filename="ref.txt",
992+
artifact=ref,
993+
)
994+
995+
loaded = await artifact_service.load_artifact(
996+
app_name="app0",
997+
user_id="user0",
998+
session_id="sess1",
999+
filename="ref.txt",
1000+
)
1001+
assert loaded == types.Part(text="profile")
1002+
1003+
1004+
@pytest.mark.asyncio
1005+
@pytest.mark.parametrize(
1006+
"service_type",
1007+
[
1008+
ArtifactServiceType.IN_MEMORY,
1009+
ArtifactServiceType.GCS,
1010+
],
1011+
)
1012+
async def test_artifact_reference_rejects_cross_user_on_save(
1013+
service_type, artifact_service_factory
1014+
):
1015+
"""ArtifactService rejects references to different users on save."""
1016+
artifact_service = artifact_service_factory(service_type)
1017+
1018+
await artifact_service.save_artifact(
1019+
app_name="app0",
1020+
user_id="victim",
1021+
session_id="victim-sess",
1022+
filename="user:secret.txt",
1023+
artifact=types.Part(text="secret"),
1024+
)
1025+
1026+
ref = types.Part(
1027+
file_data=types.FileData(
1028+
file_uri=(
1029+
"artifact://apps/app0/users/victim/artifacts/"
1030+
"user:secret.txt/versions/0"
1031+
),
1032+
mime_type="text/plain",
1033+
)
1034+
)
1035+
with pytest.raises(InputValidationError, match="same app and user scope"):
1036+
await artifact_service.save_artifact(
1037+
app_name="app0",
1038+
user_id="attacker",
1039+
session_id="attacker-sess",
1040+
filename="ref.txt",
1041+
artifact=ref,
1042+
)
1043+
1044+
1045+
@pytest.mark.asyncio
1046+
@pytest.mark.parametrize(
1047+
"service_type",
1048+
[
1049+
ArtifactServiceType.IN_MEMORY,
1050+
ArtifactServiceType.GCS,
1051+
],
1052+
)
1053+
async def test_artifact_reference_rejects_cross_app_on_save(
1054+
service_type, artifact_service_factory
1055+
):
1056+
"""ArtifactService rejects references to different apps on save."""
1057+
artifact_service = artifact_service_factory(service_type)
1058+
1059+
await artifact_service.save_artifact(
1060+
app_name="victim-app",
1061+
user_id="user0",
1062+
session_id="sess0",
1063+
filename="user:secret.txt",
1064+
artifact=types.Part(text="secret"),
1065+
)
1066+
1067+
ref = types.Part(
1068+
file_data=types.FileData(
1069+
file_uri=(
1070+
"artifact://apps/victim-app/users/user0/artifacts/"
1071+
"user:secret.txt/versions/0"
1072+
),
1073+
mime_type="text/plain",
1074+
)
1075+
)
1076+
with pytest.raises(InputValidationError, match="same app and user scope"):
1077+
await artifact_service.save_artifact(
1078+
app_name="attacker-app",
1079+
user_id="user0",
1080+
session_id="sess0",
1081+
filename="ref.txt",
1082+
artifact=ref,
1083+
)
1084+
1085+
1086+
@pytest.mark.asyncio
1087+
@pytest.mark.parametrize(
1088+
"service_type",
1089+
[
1090+
ArtifactServiceType.IN_MEMORY,
1091+
ArtifactServiceType.GCS,
1092+
],
1093+
)
1094+
async def test_artifact_reference_rejects_cross_session_on_load(
1095+
service_type, artifact_service_factory
1096+
):
1097+
"""ArtifactService rejects modified references to different sessions on load."""
1098+
artifact_service = artifact_service_factory(service_type)
1099+
1100+
await artifact_service.save_artifact(
1101+
app_name="app0",
1102+
user_id="user0",
1103+
session_id="sess0",
1104+
filename="source.txt",
1105+
artifact=types.Part(text="source"),
1106+
)
1107+
await artifact_service.save_artifact(
1108+
app_name="app0",
1109+
user_id="user0",
1110+
session_id="sess1",
1111+
filename="source.txt",
1112+
artifact=types.Part(text="other-session"),
1113+
)
1114+
1115+
ref = types.Part(
1116+
file_data=types.FileData(
1117+
file_uri=(
1118+
"artifact://apps/app0/users/user0/sessions/sess0/"
1119+
"artifacts/source.txt/versions/0"
1120+
),
1121+
mime_type="text/plain",
1122+
)
1123+
)
1124+
await artifact_service.save_artifact(
1125+
app_name="app0",
1126+
user_id="user0",
1127+
session_id="sess0",
1128+
filename="ref.txt",
1129+
artifact=ref,
1130+
)
1131+
1132+
new_uri = (
1133+
"artifact://apps/app0/users/user0/sessions/sess1/"
1134+
"artifacts/source.txt/versions/0"
1135+
)
1136+
# Manually modify the stored reference URI to point to a different session.
1137+
if service_type == ArtifactServiceType.GCS:
1138+
blob_name = artifact_service._get_blob_name(
1139+
"app0", "user0", "ref.txt", 0, "sess0"
1140+
)
1141+
blob = artifact_service.bucket.get_blob(blob_name)
1142+
blob.metadata["adkFileUri"] = new_uri
1143+
elif service_type == ArtifactServiceType.IN_MEMORY:
1144+
ref_path = artifact_service._artifact_path(
1145+
"app0", "user0", "ref.txt", "sess0"
1146+
)
1147+
artifact_service.artifacts[ref_path][0].data.file_data.file_uri = new_uri
1148+
1149+
with pytest.raises(InputValidationError, match="same session scope"):
1150+
await artifact_service.load_artifact(
1151+
app_name="app0",
1152+
user_id="user0",
1153+
session_id="sess0",
1154+
filename="ref.txt",
1155+
)
1156+
1157+
9081158
class TestEnsurePart:
9091159
"""Tests for the ensure_part normalization helper."""
9101160

0 commit comments

Comments
 (0)