Skip to content

Commit c0655d0

Browse files
fix(artifacts): normalize session_id across all artifact service backends
InMemoryArtifactService, FileArtifactService, and GcsArtifactService all keyed session-scoped artifact storage on the raw, un-normalized session_id string. Session services (InMemorySessionService, SqliteSessionService) strip whitespace from session_id before using it (#6892, #6941/#6942), so a session created with a padded id is stored under the trimmed one -- but its artifacts, saved with the same padded id, landed in a sibling namespace the session itself is never keyed under: one logical session, two artifact namespaces. This also let a whitespace-only session_id slip past validate_path_segment (which only rejects an actually-empty string) and be used verbatim as a literal path/key segment. Adds artifact_util.normalize_session_id(), used at every entry point that builds a storage key/path or path from a caller-supplied session_id: save/load/list/delete/list_versions/list_artifact_versions/ get_artifact_version in the in-memory and GCS services, and the shared _session_artifacts_dir() choke point in the file-based service. Also applied inside parse_artifact_uri() and get_artifact_uri() so a padded id can never leak into (or be read back out of) an artifact reference URI, which is what the same-session artifact-reference scope check compares against. Does not touch session services -- that normalization already landed via
1 parent 6de43b0 commit c0655d0

6 files changed

Lines changed: 184 additions & 1 deletion

File tree

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

Lines changed: 20 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,24 @@ class ParsedArtifactUri(NamedTuple):
4343
)
4444

4545

46+
def normalize_session_id(session_id: str | None) -> str | None:
47+
"""Normalizes a caller-supplied session id.
48+
49+
Strips surrounding whitespace the same way the session services do, so an
50+
artifact saved against a padded session id lands in the same storage
51+
namespace as the (normalized) session it belongs to, instead of a sibling
52+
namespace no caller using the trimmed id can ever reach.
53+
54+
Args:
55+
session_id: The caller-supplied session id, or None for a user-scoped
56+
artifact.
57+
58+
Returns:
59+
The stripped session id, or None if `session_id` was None.
60+
"""
61+
return session_id.strip() if session_id is not None else None
62+
63+
4664
def parse_artifact_uri(uri: str) -> ParsedArtifactUri | None:
4765
"""Parses an artifact URI.
4866
@@ -60,7 +78,7 @@ def parse_artifact_uri(uri: str) -> ParsedArtifactUri | None:
6078
return ParsedArtifactUri(
6179
app_name=match.group(1),
6280
user_id=match.group(2),
63-
session_id=match.group(3),
81+
session_id=normalize_session_id(match.group(3)),
6482
filename=match.group(4),
6583
version=int(match.group(5)),
6684
)
@@ -97,6 +115,7 @@ def get_artifact_uri(
97115
Returns:
98116
The constructed artifact URI.
99117
"""
118+
session_id = normalize_session_id(session_id)
100119
if session_id:
101120
return f"artifact://apps/{app_name}/users/{user_id}/sessions/{session_id}/artifacts/{filename}/versions/{version}"
102121
else:

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -247,6 +247,7 @@ def _user_artifacts_dir(base_root: Path) -> Path:
247247

248248
def _session_artifacts_dir(base_root: Path, session_id: str) -> Path:
249249
"""Returns the path that stores session-scoped artifacts."""
250+
session_id = session_id.strip()
250251
artifact_util.validate_path_segment(session_id, "session_id")
251252
return base_root / "sessions" / session_id / "artifacts"
252253

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

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -214,6 +214,7 @@ def _get_blob_prefix(
214214
if self._file_has_user_namespace(filename):
215215
return f"{app_name}/{user_id}/user/{filename}"
216216

217+
session_id = artifact_util.normalize_session_id(session_id)
217218
if session_id is None:
218219
raise InputValidationError(
219220
"Session ID must be provided for session-scoped artifacts."
@@ -257,6 +258,7 @@ def _save_artifact(
257258
from google.cloud import exceptions # pylint: disable=g-import-not-at-top
258259

259260
artifact = ensure_part(artifact)
261+
session_id = artifact_util.normalize_session_id(session_id)
260262
blob_metadata = {k: str(v) for k, v in (custom_metadata or {}).items()}
261263
if artifact.inline_data and artifact.inline_data.display_name:
262264
blob_metadata[_GCS_DISPLAY_NAME_METADATA_KEY] = (
@@ -350,6 +352,7 @@ def _load_artifact(
350352
filename: str,
351353
version: Optional[int] = None,
352354
) -> Optional[types.Part]:
355+
session_id = artifact_util.normalize_session_id(session_id)
353356
if version is None:
354357
versions = self._list_versions(
355358
app_name=app_name,
@@ -430,6 +433,7 @@ def _list_artifact_keys(
430433
) -> list[str]:
431434
artifact_util.validate_path_segment(app_name, "app_name")
432435
artifact_util.validate_path_segment(user_id, "user_id")
436+
session_id = artifact_util.normalize_session_id(session_id)
433437
if session_id is not None:
434438
artifact_util.validate_path_segment(session_id, "session_id")
435439
filenames = set()
@@ -635,6 +639,7 @@ def _get_authenticated_url_sync(
635639
max_depth: int = _MAX_ARTIFACT_REFERENCE_DEPTH,
636640
) -> Optional[str]:
637641
"""Generates an authenticated browser URL for an artifact."""
642+
session_id = artifact_util.normalize_session_id(session_id)
638643
if version is None:
639644
versions = self._list_versions(
640645
app_name=app_name,
@@ -742,6 +747,7 @@ def _get_signed_url_sync(
742747
max_depth: int = _MAX_ARTIFACT_REFERENCE_DEPTH,
743748
) -> Optional[str]:
744749
"""Generates a time-limited signed URL for an artifact."""
750+
session_id = artifact_util.normalize_session_id(session_id)
745751
if version is None:
746752
versions = self._list_versions(
747753
app_name=app_name,

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

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -91,6 +91,7 @@ def _artifact_path(
9191
if self._file_has_user_namespace(filename):
9292
return f"{app_name}/{user_id}/user/{filename}"
9393

94+
session_id = artifact_util.normalize_session_id(session_id)
9495
if session_id is None:
9596
raise InputValidationError(
9697
"Session ID must be provided for session-scoped artifacts."
@@ -110,6 +111,7 @@ async def save_artifact(
110111
custom_metadata: Optional[dict[str, Any]] = None,
111112
) -> int:
112113
artifact = ensure_part(artifact)
114+
session_id = artifact_util.normalize_session_id(session_id)
113115
path = self._artifact_path(app_name, user_id, filename, session_id)
114116
if path not in self.artifacts:
115117
self.artifacts[path] = []
@@ -167,6 +169,7 @@ async def load_artifact(
167169
session_id: Optional[str] = None,
168170
version: Optional[int] = None,
169171
) -> Optional[types.Part]:
172+
session_id = artifact_util.normalize_session_id(session_id)
170173
path = self._artifact_path(app_name, user_id, filename, session_id)
171174
versions = self.artifacts.get(path)
172175
if not versions:
@@ -222,6 +225,7 @@ async def list_artifact_keys(
222225
) -> list[str]:
223226
artifact_util.validate_path_segment(app_name, "app_name")
224227
artifact_util.validate_path_segment(user_id, "user_id")
228+
session_id = artifact_util.normalize_session_id(session_id)
225229
if session_id is not None:
226230
artifact_util.validate_path_segment(session_id, "session_id")
227231
usernamespace_prefix = f"{app_name}/{user_id}/user/"

‎tests/unittests/artifacts/test_artifact_service.py‎

Lines changed: 119 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -305,6 +305,125 @@ async def test_save_load_delete(service_type, artifact_service_factory):
305305
)
306306

307307

308+
@pytest.mark.asyncio
309+
@pytest.mark.parametrize(
310+
"service_type",
311+
[
312+
ArtifactServiceType.IN_MEMORY,
313+
ArtifactServiceType.GCS,
314+
ArtifactServiceType.FILE,
315+
],
316+
)
317+
async def test_padded_session_id_lands_in_same_namespace_as_trimmed(
318+
service_type, artifact_service_factory
319+
):
320+
"""An artifact saved against a whitespace-padded session id must be
321+
reachable, listable, and deletable using the trimmed id, since that is the
322+
id the corresponding session is actually stored under (session services
323+
normalize session_id the same way)."""
324+
artifact_service = artifact_service_factory(service_type)
325+
artifact = types.Part(text="hello")
326+
327+
await artifact_service.save_artifact(
328+
app_name="app0",
329+
user_id="user0",
330+
session_id="sess0\n",
331+
filename="report.txt",
332+
artifact=artifact,
333+
)
334+
335+
# Reachable under the trimmed id, the one the session itself is keyed on.
336+
assert (
337+
await artifact_service.load_artifact(
338+
app_name="app0",
339+
user_id="user0",
340+
session_id="sess0",
341+
filename="report.txt",
342+
)
343+
== artifact
344+
)
345+
assert "report.txt" in await artifact_service.list_artifact_keys(
346+
app_name="app0", user_id="user0", session_id="sess0"
347+
)
348+
349+
# Also reachable under the original padded id: both forms must resolve to
350+
# the same stored artifact rather than the padded id silently shadowing it
351+
# in a namespace only the padded id itself could ever reach again.
352+
assert (
353+
await artifact_service.load_artifact(
354+
app_name="app0",
355+
user_id="user0",
356+
session_id="sess0\n",
357+
filename="report.txt",
358+
)
359+
== artifact
360+
)
361+
362+
await artifact_service.delete_artifact(
363+
app_name="app0",
364+
user_id="user0",
365+
session_id="sess0",
366+
filename="report.txt",
367+
)
368+
assert not await artifact_service.load_artifact(
369+
app_name="app0",
370+
user_id="user0",
371+
session_id="sess0\n",
372+
filename="report.txt",
373+
)
374+
375+
376+
@pytest.mark.asyncio
377+
@pytest.mark.parametrize(
378+
"service_type",
379+
[
380+
ArtifactServiceType.IN_MEMORY,
381+
ArtifactServiceType.GCS,
382+
],
383+
)
384+
async def test_artifact_reference_allows_padded_session_id_at_call_site(
385+
service_type, artifact_service_factory
386+
):
387+
"""A caller that consistently uses a padded session id must still be able
388+
to save and resolve an artifact reference within that session: the scope
389+
check must compare normalized ids on both sides, not the caller's raw
390+
string against a URI minted from the (already normalized) stored id."""
391+
artifact_service = artifact_service_factory(service_type)
392+
393+
await artifact_service.save_artifact(
394+
app_name="app0",
395+
user_id="user0",
396+
session_id="sess0\n",
397+
filename="source.txt",
398+
artifact=types.Part(text="hello"),
399+
)
400+
401+
ref = types.Part(
402+
file_data=types.FileData(
403+
file_uri=(
404+
"artifact://apps/app0/users/user0/sessions/sess0/"
405+
"artifacts/source.txt/versions/0"
406+
),
407+
mime_type="text/plain",
408+
)
409+
)
410+
await artifact_service.save_artifact(
411+
app_name="app0",
412+
user_id="user0",
413+
session_id="sess0\n",
414+
filename="ref.txt",
415+
artifact=ref,
416+
)
417+
418+
loaded = await artifact_service.load_artifact(
419+
app_name="app0",
420+
user_id="user0",
421+
session_id="sess0\n",
422+
filename="ref.txt",
423+
)
424+
assert loaded == types.Part(text="hello")
425+
426+
308427
@pytest.mark.asyncio
309428
async def test_in_memory_loads_nested_artifact_reference(
310429
artifact_service_factory,

‎tests/unittests/artifacts/test_artifact_util.py‎

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -112,6 +112,40 @@ def test_get_user_scoped_artifact_uri():
112112
assert uri == "artifact://apps/app2/users/user2/artifacts/file2/versions/456"
113113

114114

115+
def test_normalize_session_id_strips_whitespace():
116+
assert artifact_util.normalize_session_id(" sess0\n") == "sess0"
117+
118+
119+
def test_normalize_session_id_passes_none_through():
120+
assert artifact_util.normalize_session_id(None) is None
121+
122+
123+
def test_get_artifact_uri_normalizes_padded_session_id():
124+
"""A padded session id must not leak into the constructed URI, since the
125+
session it points at is stored under the trimmed id."""
126+
uri = artifact_util.get_artifact_uri(
127+
app_name="app1",
128+
user_id="user1",
129+
session_id=" session1\n",
130+
filename="file1",
131+
version=123,
132+
)
133+
assert (
134+
uri
135+
== "artifact://apps/app1/users/user1/sessions/session1/artifacts/file1/versions/123"
136+
)
137+
138+
139+
def test_parse_artifact_uri_normalizes_a_legacy_padded_session_id():
140+
"""A URI minted before this fix could carry a padded session id in its own
141+
path segment; parsing it must still yield the trimmed id so downstream
142+
scope checks compare like with like."""
143+
uri = "artifact://apps/app1/users/user1/sessions/session1\n/artifacts/file1/versions/123"
144+
parsed = artifact_util.parse_artifact_uri(uri)
145+
assert parsed is not None
146+
assert parsed.session_id == "session1"
147+
148+
115149
def test_is_artifact_ref_true():
116150
"""Tests is_artifact_ref with a valid artifact reference."""
117151
artifact = types.Part(

0 commit comments

Comments
 (0)