From 34c2abf3003a457822c44d80d760b7938b37f677 Mon Sep 17 00:00:00 2001 From: George Weale Date: Mon, 17 Aug 2026 23:06:40 +0000 Subject: [PATCH 1/4] fix(sessions): validate session_id and enforce ownership in delete_session (v1) Port of upstream commit b2916c71 (PR #5580) to the v1 branch. VertexAiSessionService.delete_session accepted a user_id argument and never read it, so the caller's identity had no bearing on which session was deleted; get_session has checked ownership since it was written. delete_session now fetches the session first, returns silently if it is already gone, and raises ValueError when the session belongs to a different user. Session ids were also interpolated straight into the resource path. A new _validate_session_id rejects anything outside ^[A-Za-z0-9_-]+$ in create_session, get_session, delete_session and append_event, so an id can no longer contribute extra path segments. Behaviour changes for existing users: a caller-chosen session_id containing a dot, a colon or any non-ASCII character is now rejected with ValueError, where before it was passed through. Server-generated Vertex session ids are numeric and unaffected. delete_session now costs one extra get round trip, and a missing session returns without raising instead of surfacing the delete's 404. --- .../adk/sessions/vertex_ai_session_service.py | 38 ++++++++++++++++-- .../test_vertex_ai_session_service.py | 39 +++++++++++++++++++ 2 files changed, 74 insertions(+), 3 deletions(-) diff --git a/src/google/adk/sessions/vertex_ai_session_service.py b/src/google/adk/sessions/vertex_ai_session_service.py index 8c1fdc134e2..1f1edd3d5d3 100644 --- a/src/google/adk/sessions/vertex_ai_session_service.py +++ b/src/google/adk/sessions/vertex_ai_session_service.py @@ -47,6 +47,19 @@ _COMPACTION_CUSTOM_METADATA_KEY = '_compaction' _USAGE_METADATA_CUSTOM_METADATA_KEY = '_usage_metadata' +_SESSION_ID_PATTERN = re.compile(r'^[A-Za-z0-9_-]+$') + + +def _validate_session_id(session_id: str) -> None: + """Rejects session IDs that could escape the URL path segment.""" + if not isinstance(session_id, str) or not _SESSION_ID_PATTERN.fullmatch( + session_id + ): + raise ValueError( + f'Invalid session_id {session_id!r}: must match' + f' {_SESSION_ID_PATTERN.pattern}.' + ) + def _quote_filter_literal(value: str) -> str: """Quotes filter values so embedded metacharacters stay inside the literal.""" @@ -127,6 +140,7 @@ async def create_session( config = {'session_state': state} if state else {} if session_id: + _validate_session_id(session_id) config['session_id'] = session_id config.update(kwargs) async with self._get_api_client() as api_client: @@ -157,6 +171,7 @@ async def get_session( session_id: str, config: Optional[GetSessionConfig] = None, ) -> Optional[Session]: + _validate_session_id(session_id) reasoning_engine_id = self._get_reasoning_engine_id(app_name) session_resource_name = ( f'reasoningEngines/{reasoning_engine_id}/sessions/{session_id}' @@ -256,14 +271,30 @@ async def list_sessions( async def delete_session( self, *, app_name: str, user_id: str, session_id: str ) -> None: + _validate_session_id(session_id) reasoning_engine_id = self._get_reasoning_engine_id(app_name) + session_resource_name = ( + f'reasoningEngines/{reasoning_engine_id}/sessions/{session_id}' + ) async with self._get_api_client() as api_client: + # Enforce ownership: delete_session otherwise ignores user_id entirely. + try: + existing = await api_client.agent_engines.sessions.get( + name=session_resource_name + ) + except ClientError as e: + if e.code == 404: + return + raise + if existing.user_id != user_id: + raise ValueError( + f'Session {session_id} does not belong to user {user_id}.' + ) + try: await api_client.agent_engines.sessions.delete( - name=( - f'reasoningEngines/{reasoning_engine_id}/sessions/{session_id}' - ), + name=session_resource_name, ) except Exception as e: logger.error('Error deleting session %s: %s', session_id, e) @@ -274,6 +305,7 @@ async def append_event(self, session: Session, event: Event) -> Event: # Update the in-memory session. await super().append_event(session=session, event=event) + _validate_session_id(session.id) reasoning_engine_id = self._get_reasoning_engine_id(session.app_name) # Build config (Monolithic approach) diff --git a/tests/unittests/sessions/test_vertex_ai_session_service.py b/tests/unittests/sessions/test_vertex_ai_session_service.py index c5c9996ef5e..b8c71701dc8 100644 --- a/tests/unittests/sessions/test_vertex_ai_session_service.py +++ b/tests/unittests/sessions/test_vertex_ai_session_service.py @@ -725,6 +725,45 @@ async def test_get_and_delete_session(): assert str(excinfo.value) == '404 Session not found: 1' +@pytest.mark.asyncio +@pytest.mark.usefixtures('mock_get_api_client') +async def test_delete_session_rejects_other_users_session(): + """delete_session must not delete a session owned by a different user.""" + session_service = mock_vertex_ai_session_service() + + # session '1' belongs to 'user'; 'user2' must not be allowed to delete it. + with pytest.raises(ValueError) as excinfo: + await session_service.delete_session( + app_name='123', user_id='user2', session_id='1' + ) + assert 'does not belong to user user2' in str(excinfo.value) + + # Session must still exist. + assert ( + await session_service.get_session( + app_name='123', user_id='user', session_id='1' + ) + == MOCK_SESSION + ) + + +@pytest.mark.asyncio +@pytest.mark.usefixtures('mock_get_api_client') +async def test_session_id_path_traversal_rejected(): + """Session IDs containing path-traversal characters must be rejected.""" + session_service = mock_vertex_ai_session_service() + + for bad_id in ['..', '../foo', '..?force=true', 'a/b', '']: + with pytest.raises(ValueError): + await session_service.delete_session( + app_name='123', user_id='user', session_id=bad_id + ) + with pytest.raises(ValueError): + await session_service.get_session( + app_name='123', user_id='user', session_id=bad_id + ) + + @pytest.mark.asyncio @pytest.mark.usefixtures('mock_get_api_client') async def test_get_session_with_page_token(): From fc98ceec12a15d21e61241bbca271ff1b434ebee Mon Sep 17 00:00:00 2001 From: George Weale Date: Mon, 17 Aug 2026 23:08:34 +0000 Subject: [PATCH 2/4] fix(agents): scope LangGraph checkpointer thread id to app and user (v1) Port of upstream commit 19e2a728 to the v1 branch. LangGraphAgent passed the bare session id to the LangGraph checkpointer as its thread id. Session ids are caller-chosen and only unique within an (app name, user id) pair, so two users who picked the same session id read and wrote each other's graph state, and so did two apps. The thread id is now the SHA-256 digest of the app name, user id and session id together. Each component is length-prefixed before hashing, so a component containing the separator cannot stand in for a different triple, and hashing keeps the user id out of checkpointer storage. Behaviour change for existing users: checkpoints written by earlier releases are keyed on the old thread id and are not reused. With a persistent checkpointer, the first turn after upgrading resumes from empty graph state, and because the adapter sends only the last human message when a checkpointer is configured, the history that lived in the checkpointer is gone rather than re-sent. There is no migration path. Reading the old thread id as a fallback would mean reading the ambiguous key this change exists to remove, so it would reintroduce the cross-user collision. Affected users should finish in-flight conversations before upgrading. --- src/google/adk/agents/langgraph_agent.py | 45 ++++++- .../unittests/agents/test_langgraph_agent.py | 111 +++++++++++++++++- 2 files changed, 153 insertions(+), 3 deletions(-) diff --git a/src/google/adk/agents/langgraph_agent.py b/src/google/adk/agents/langgraph_agent.py index 2de1c2b8a97..5d527c88dad 100644 --- a/src/google/adk/agents/langgraph_agent.py +++ b/src/google/adk/agents/langgraph_agent.py @@ -14,6 +14,7 @@ from __future__ import annotations +import hashlib from typing import AsyncGenerator from typing import Union @@ -31,6 +32,33 @@ from .invocation_context import InvocationContext +def _get_thread_id(app_name: str, user_id: str, session_id: str) -> str: + """Derives the LangGraph checkpointer thread id for a session. + + Session ids are caller-chosen and are only unique within an + (app_name, user_id) pair, so all three components have to take part in the + thread id. Each component is length-prefixed before hashing so that a + component containing the separator cannot stand in for a different triple. + The composite is hashed rather than used verbatim so that the thread id is a + fixed-length token no checkpointer backend has to escape, and so the user id + is not written into checkpointer storage; the cost is that a stored row can + only be tied back to a session by recomputing the digest. + + Args: + app_name: the app the session belongs to + user_id: the user the session belongs to + session_id: the session id + + Returns: + a deterministic thread id for the session + """ + key = '|'.join( + f'{len(component)}:{component}' + for component in (app_name, user_id, session_id) + ) + return hashlib.sha256(key.encode('utf-8')).hexdigest() + + def _get_last_human_messages(events: list[Event]) -> list[HumanMessage]: """Extracts last human messages from given list of events. @@ -50,7 +78,14 @@ def _get_last_human_messages(events: list[Event]) -> list[HumanMessage]: class LangGraphAgent(BaseAgent): - """Currently a concept implementation, supports single and multi-turn.""" + """Currently a concept implementation, supports single and multi-turn. + + The checkpointer thread id is derived from the session's app name, user id + and id together, because session ids are only unique within an + (app name, user id) pair. Checkpoints written by earlier releases, which + keyed the thread on the session id alone, are not reused: with a persistent + checkpointer the first turn after upgrading resumes from empty graph state. + """ model_config = ConfigDict( arbitrary_types_allowed=True, @@ -68,7 +103,13 @@ async def _run_async_impl( ) -> AsyncGenerator[Event, None]: # Needed for langgraph checkpointer (for subsequent invocations; multi-turn) - config: RunnableConfig = {'configurable': {'thread_id': ctx.session.id}} + config: RunnableConfig = { + 'configurable': { + 'thread_id': _get_thread_id( + ctx.session.app_name, ctx.session.user_id, ctx.session.id + ) + } + } # Add instruction as SystemMessage if graph state is empty current_graph_state = self.graph.get_state(config) diff --git a/tests/unittests/agents/test_langgraph_agent.py b/tests/unittests/agents/test_langgraph_agent.py index 19ba55a6873..0174a561a21 100644 --- a/tests/unittests/agents/test_langgraph_agent.py +++ b/tests/unittests/agents/test_langgraph_agent.py @@ -20,6 +20,7 @@ LANGGRAPH_AVAILABLE = True try: from google.adk.agents.invocation_context import InvocationContext + from google.adk.agents.langgraph_agent import _get_thread_id from google.adk.agents.langgraph_agent import LangGraphAgent from google.adk.events.event import Event from google.adk.plugins.plugin_manager import PluginManager @@ -77,6 +78,7 @@ def __call__(self, *args, **kwargs): return DummyTypes() InvocationContext = DummyTypes() + _get_thread_id = DummyTypes() LangGraphAgent = DummyTypes() Event = DummyTypes() PluginManager = DummyTypes() @@ -231,6 +233,9 @@ async def test_langgraph_agent( mock_parent_context = MagicMock(spec=InvocationContext) mock_session = MagicMock() + mock_session.app_name = "test_app" + mock_session.user_id = "test_user" + mock_session.id = "test_session_id" mock_parent_context.session = mock_session mock_parent_context.user_content = types.Content( role="user", parts=[types.Part.from_text(text="test prompt")] @@ -256,8 +261,112 @@ async def test_langgraph_agent( assert result_event.author == "weather_agent" assert result_event.content.parts[0].text == "test response" + expected_thread_id = _get_thread_id( + mock_session.app_name, mock_session.user_id, mock_session.id + ) mock_graph.invoke.assert_called_once() mock_graph.invoke.assert_called_with( {"messages": expected_messages}, - {"configurable": {"thread_id": mock_session.id}}, + {"configurable": {"thread_id": expected_thread_id}}, + ) + + +def test_get_thread_id_is_stable_across_processes(): + """A literal digest also rejects a swap to a per-process-salted hash.""" + assert ( + _get_thread_id("app", "alice", "session-id") + == "8c95b75b65efd3d1ddd363cfdcb7d1d4bdd9a747aa8f505a887b1dc217fca0e2" + ) + + +def test_get_thread_id_separates_users_and_apps(): + assert _get_thread_id("app", "alice", "shared-id") != _get_thread_id( + "app", "bob", "shared-id" + ) + assert _get_thread_id("app_one", "alice", "shared-id") != _get_thread_id( + "app_two", "alice", "shared-id" + ) + + +def test_get_thread_id_cannot_be_forged_with_a_separator(): + assert _get_thread_id("app", "alice", "bob|s1") != _get_thread_id( + "app", "alice|bob", "s1" + ) + assert _get_thread_id("app|alice", "bob", "s1") != _get_thread_id( + "app", "alice|bob", "s1" ) + assert _get_thread_id("app", "alice", "1:x") != _get_thread_id( + "app", "alice|1:x", "" + ) + + +def _make_parent_context(app_name, user_id, session_id): + """Builds a mock invocation context for the given session triple.""" + parent_context = MagicMock(spec=InvocationContext) + mock_session = MagicMock() + mock_session.app_name = app_name + mock_session.user_id = user_id + mock_session.id = session_id + mock_session.events = [] + parent_context.session = mock_session + parent_context.user_content = types.Content( + role="user", parts=[types.Part.from_text(text="test prompt")] + ) + parent_context.branch = "parent_agent" + parent_context.end_invocation = False + parent_context.invocation_id = "test_invocation_id" + parent_context.model_copy.return_value = parent_context + parent_context.plugin_manager = PluginManager(plugins=[]) + return parent_context + + +async def _run_and_get_thread_id(app_name, user_id, session_id): + """Runs the agent once and returns the checkpointer thread id it used.""" + mock_graph = MagicMock(spec=CompiledGraph) + mock_graph_state = MagicMock() + mock_graph_state.values = {} + mock_graph.get_state.return_value = mock_graph_state + mock_graph.checkpointer = MagicMock() + mock_graph.invoke.return_value = { + "messages": [AIMessage(content="test response")] + } + agent = LangGraphAgent( + name="weather_agent", + instruction="test system prompt", + graph=mock_graph, + ) + + async for _ in agent.run_async( + _make_parent_context(app_name, user_id, session_id) + ): + pass + + read_config = mock_graph.get_state.call_args.args[0] + write_config = mock_graph.invoke.call_args.args[1] + assert read_config == write_config + return write_config["configurable"]["thread_id"] + + +@pytest.mark.asyncio +async def test_same_session_id_across_users_does_not_share_a_thread(): + """Session ids are caller-chosen and only unique within a user.""" + alice_thread_id = await _run_and_get_thread_id("app", "alice", "shared-id") + bob_thread_id = await _run_and_get_thread_id("app", "bob", "shared-id") + + assert alice_thread_id != bob_thread_id + + +@pytest.mark.asyncio +async def test_same_session_id_across_apps_does_not_share_a_thread(): + first_thread_id = await _run_and_get_thread_id("app_one", "a", "shared-id") + second_thread_id = await _run_and_get_thread_id("app_two", "a", "shared-id") + + assert first_thread_id != second_thread_id + + +@pytest.mark.asyncio +async def test_same_session_resolves_to_the_same_thread(): + first_thread_id = await _run_and_get_thread_id("app", "alice", "session-id") + second_thread_id = await _run_and_get_thread_id("app", "alice", "session-id") + + assert first_thread_id == second_thread_id From 0d1804ac98454cf73d8c5344d1637d6cd462fe0f Mon Sep 17 00:00:00 2001 From: George Weale Date: Mon, 17 Aug 2026 23:12:43 +0000 Subject: [PATCH 3/4] fix(memory): clean up Vertex RAG transcript and stop blocking the loop (v1) Adapted from upstream commit 80a05b7f. That commit rewrote this service onto a client SDK the v1 branch does not have, so the code here is written against v1's module-level vertexai RAG surface rather than cherry-picked. VertexAiRagMemoryService.add_session_to_memory writes the full session transcript to a plaintext file in the system temp directory before uploading it. The removal was the last statement of the happy path, so any upload that raised, or a cancelled task, left the transcript on disk indefinitely. The write and the uploads now sit in a try whose finally removes the file, tolerating FileNotFoundError, and the path is recorded before the write so a failed write is cleaned up too. The file is opened with an explicit utf-8 encoding instead of the locale default. Corpus names are also validated before the file is created rather than after, so a misconfigured service writes nothing at all. Both methods are declared async but called the synchronous RAG SDK, so every upload and retrieval blocked the event loop for the whole HTTP round trip. Both calls now run through asyncio.to_thread. This is the same two functions the transcript cleanup rewrites, which is why it is one change. to_thread copies the current contextvars context, so tracing spans still propagate into the worker thread, and both SDK entry points read the module globals set by vertexai.init and build a fresh client per call, so neither is thread-affine. Behaviour change for existing users: VertexAiRagMemoryService() constructed with no rag_corpus now raises ValueError("rag_corpus must be set on every RAG resource.") before doing any work. Previously the guard it replaces could never fire, because __init__ always builds a one-element rag_resources list even when rag_corpus is None, so that configuration wrote the transcript to disk and then failed inside the SDK with corpus_name=None. The CLI factory always passes a corpus, so it is unaffected. Separately, the SDK call now runs on a worker thread, which anyone substituting a thread-hostile fake for rag.upload_file or rag.retrieval_query would notice. --- .../memory/vertex_ai_rag_memory_service.py | 103 ++++++----- .../test_vertex_ai_rag_memory_service.py | 160 +++++++++++++++--- 2 files changed, 204 insertions(+), 59 deletions(-) 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 fb81508ddda..1696fb34508 100644 --- a/src/google/adk/memory/vertex_ai_rag_memory_service.py +++ b/src/google/adk/memory/vertex_ai_rag_memory_service.py @@ -15,6 +15,7 @@ from __future__ import annotations +import asyncio import base64 import binascii from collections import OrderedDict @@ -118,48 +119,67 @@ def __init__( @override async def add_session_to_memory(self, session: Session) -> None: - with tempfile.NamedTemporaryFile( - mode="w", delete=False, suffix=".txt" - ) as temp_file: - - output_lines = [] - for event in session.events: - if not event.content or not event.content.parts: - continue - text_parts = [ - part.text.replace("\n", " ") - for part in event.content.parts - if part.text - ] - if text_parts: - output_lines.append( - json.dumps({ - "author": event.author, - "timestamp": event.timestamp, - "text": ".".join(text_parts), - }) - ) - output_string = "\n".join(output_lines) - temp_file.write(output_string) - temp_file_path = temp_file.name - - if not self._vertex_rag_store.rag_resources: - raise ValueError("Rag resources must be set.") + rag_resources = self._vertex_rag_store.rag_resources or () + corpus_names = tuple( + resource.rag_corpus for resource in rag_resources if resource.rag_corpus + ) + if not corpus_names or len(corpus_names) != len(rag_resources): + raise ValueError("rag_corpus must be set on every RAG resource.") - from ..dependencies.vertexai import rag + output_lines = [] + for event in session.events: + if not event.content or not event.content.parts: + continue + text_parts = [ + part.text.replace("\n", " ") + for part in event.content.parts + if part.text + ] + if text_parts: + output_lines.append( + json.dumps({ + "author": event.author, + "timestamp": event.timestamp, + "text": ".".join(text_parts), + }) + ) + output_string = "\n".join(output_lines) - for rag_resource in self._vertex_rag_store.rag_resources: - rag.upload_file( - corpus_name=rag_resource.rag_corpus, - path=temp_file_path, - # this is the temp workaround as upload file does not support - # adding metadata, thus use display_name to store the session info. - display_name=_build_source_display_name( - session.app_name, session.user_id, session.id - ), - ) + from ..dependencies.vertexai import rag - os.remove(temp_file_path) + temp_file_path: str | None = None + try: + with tempfile.NamedTemporaryFile( + mode="w", + delete=False, + encoding="utf-8", + suffix=".txt", + ) as temp_file: + temp_file_path = temp_file.name + temp_file.write(output_string) + + # Fails fast: the RAG API cannot roll back corpora already written. + for corpus_name in corpus_names: + # The SDK call is synchronous, so it runs on a worker thread rather + # than blocking the event loop for the whole HTTP round trip. + await asyncio.to_thread( + rag.upload_file, + corpus_name=corpus_name, + path=temp_file_path, + # this is the temp workaround as upload file does not support + # adding metadata, thus use display_name to store the session info. + display_name=_build_source_display_name( + session.app_name, session.user_id, session.id + ), + ) + finally: + # The transcript is plaintext, so it is removed on failure and + # cancellation as well as on success. + if temp_file_path: + try: + os.remove(temp_file_path) + except FileNotFoundError: + pass @override async def search_memory( @@ -169,7 +189,10 @@ async def search_memory( from ..dependencies.vertexai import rag from ..events.event import Event - response = rag.retrieval_query( + # The SDK call is synchronous, so it runs on a worker thread rather than + # blocking the event loop for the whole HTTP round trip. + response = await asyncio.to_thread( + rag.retrieval_query, text=query, rag_resources=self._vertex_rag_store.rag_resources, rag_corpora=self._vertex_rag_store.rag_corpora, 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 7c20de87f5a..397776af66f 100644 --- a/tests/unittests/memory/test_vertex_ai_rag_memory_service.py +++ b/tests/unittests/memory/test_vertex_ai_rag_memory_service.py @@ -12,7 +12,11 @@ # See the License for the specific language governing permissions and # limitations under the License. +import asyncio import json +import os +import tempfile +import threading from types import SimpleNamespace from google.adk.events.event import Event @@ -31,6 +35,32 @@ def _rag_context(source_display_name: str, text: str) -> SimpleNamespace: ) +def _session() -> Session: + return Session( + app_name="demo.app", + user_id="alice.smith", + id="session.secret", + last_update_time=1, + events=[ + Event( + id="event-1", + author="user", + timestamp=1, + content=types.Content( + parts=[types.Part(text="sensitive memory")] + ), + ) + ], + ) + + +@pytest.fixture(name="temp_dir") +def _temp_dir(tmp_path, monkeypatch): + """Redirects NamedTemporaryFile so a leaked transcript is observable.""" + monkeypatch.setattr(tempfile, "tempdir", str(tmp_path)) + return tmp_path + + @pytest.mark.asyncio async def test_search_memory_rejects_ambiguous_legacy_display_names(mocker): """Ensures dotted user IDs cannot match another user's legacy memory.""" @@ -71,30 +101,15 @@ async def test_search_memory_rejects_ambiguous_legacy_display_names(mocker): @pytest.mark.asyncio -async def test_add_and_search_memory_uses_unambiguous_display_names(mocker): +async def test_add_and_search_memory_uses_unambiguous_display_names( + mocker, temp_dir +): memory_service = VertexAiRagMemoryService(rag_corpus="unused") upload_file = mocker.Mock() fake_rag = SimpleNamespace(upload_file=upload_file) mocker.patch("google.adk.dependencies.vertexai.rag", fake_rag) - await memory_service.add_session_to_memory( - Session( - app_name="demo.app", - user_id="alice.smith", - id="session.secret", - last_update_time=1, - events=[ - Event( - id="event-1", - author="user", - timestamp=1, - content=types.Content( - parts=[types.Part(text="sensitive memory")] - ), - ) - ], - ) - ) + await memory_service.add_session_to_memory(_session()) display_name = upload_file.call_args.kwargs["display_name"] assert display_name.startswith(_SOURCE_DISPLAY_NAME_PREFIX) @@ -115,3 +130,110 @@ async def test_add_and_search_memory_uses_unambiguous_display_names(mocker): assert [memory.content.parts[0].text for memory in response.memories] == [ "sensitive memory" ] + assert not list(temp_dir.iterdir()) + + +@pytest.mark.asyncio +async def test_add_session_cleans_temp_file_after_partial_upload_failure( + mocker, temp_dir +): + attempted_corpora = [] + + def upload_file(*, corpus_name, **_kwargs): + attempted_corpora.append(corpus_name) + if corpus_name == "second": + raise RuntimeError("upload failed") + + fake_rag = SimpleNamespace(upload_file=upload_file) + mocker.patch("google.adk.dependencies.vertexai.rag", fake_rag) + memory_service = VertexAiRagMemoryService(rag_corpus="first") + memory_service._vertex_rag_store.rag_resources = [ + types.VertexRagStoreRagResource(rag_corpus="first"), + types.VertexRagStoreRagResource(rag_corpus="second"), + types.VertexRagStoreRagResource(rag_corpus="third"), + ] + + with pytest.raises(RuntimeError, match="upload failed"): + await memory_service.add_session_to_memory(_session()) + + assert attempted_corpora == ["first", "second"] + assert not list(temp_dir.iterdir()) + + +@pytest.mark.asyncio +async def test_add_session_cleans_temp_file_when_cancelled(mocker, temp_dir): + upload_started = threading.Event() + allow_upload_to_finish = threading.Event() + + def upload_file(**_kwargs): + upload_started.set() + # Bounded so a regression that runs this on the event loop, where nothing + # can set the event, fails the test instead of hanging the suite. + allow_upload_to_finish.wait(timeout=30) + + fake_rag = SimpleNamespace(upload_file=upload_file) + mocker.patch("google.adk.dependencies.vertexai.rag", fake_rag) + memory_service = VertexAiRagMemoryService(rag_corpus="corpus") + + add_session = asyncio.create_task( + memory_service.add_session_to_memory(_session()) + ) + try: + await asyncio.to_thread(upload_started.wait) + add_session.cancel() + with pytest.raises(asyncio.CancelledError): + await add_session + finally: + allow_upload_to_finish.set() + + assert not list(temp_dir.iterdir()) + + +@pytest.mark.asyncio +async def test_add_session_leaves_no_temp_file_when_corpus_missing(temp_dir): + memory_service = VertexAiRagMemoryService(rag_corpus=None) + + with pytest.raises(ValueError, match="rag_corpus must be set"): + await memory_service.add_session_to_memory(_session()) + + assert not list(temp_dir.iterdir()) + + +@pytest.mark.asyncio +async def test_add_session_uploads_off_the_event_loop(mocker, temp_dir): + upload_thread_ids = [] + + def upload_file(*, path, **_kwargs): + upload_thread_ids.append(threading.get_ident()) + # The transcript is still on disk while the upload is in flight. + assert os.path.exists(path) + + fake_rag = SimpleNamespace(upload_file=upload_file) + mocker.patch("google.adk.dependencies.vertexai.rag", fake_rag) + memory_service = VertexAiRagMemoryService(rag_corpus="corpus") + + await memory_service.add_session_to_memory(_session()) + + assert upload_thread_ids == [upload_thread_ids[0]] + assert threading.get_ident() not in upload_thread_ids + assert not list(temp_dir.iterdir()) + + +@pytest.mark.asyncio +async def test_search_memory_queries_off_the_event_loop(mocker): + query_thread_ids = [] + + def retrieval_query(**_kwargs): + query_thread_ids.append(threading.get_ident()) + return SimpleNamespace(contexts=SimpleNamespace(contexts=[])) + + fake_rag = SimpleNamespace(retrieval_query=retrieval_query) + mocker.patch("google.adk.dependencies.vertexai.rag", fake_rag) + memory_service = VertexAiRagMemoryService(rag_corpus="corpus") + + await memory_service.search_memory( + app_name="demo", user_id="alice", query="memory" + ) + + assert query_thread_ids + assert threading.get_ident() not in query_thread_ids From 39e6b192b4d47e909ecef3ef7e51b59b5a342c1e Mon Sep 17 00:00:00 2001 From: George Weale Date: Mon, 17 Aug 2026 23:13:36 +0000 Subject: [PATCH 4/4] fix(memory): key in-memory memory store by app_name and user_id tuple (v1) Port of upstream commit fd8f7eb2 to the v1 branch. InMemoryMemoryService kept its per-user event store under the string f'{app_name}/{user_id}'. That flattening is not injective: an app named 'acme/alice' with user 'bob' produced the same key as app 'acme' with user 'alice/bob', so the two shared a store and each could read the other's sessions. The key is now the (app_name, user_id) tuple, which cannot be aliased. The store is private and in-memory, so nothing persists across a restart and there is no stored format to migrate. The only observable change is for a caller or test reaching into service._session_events with a string key. --- .../adk/memory/in_memory_memory_service.py | 8 ++-- .../memory/test_in_memory_memory_service.py | 45 ++++++++++++++++--- 2 files changed, 42 insertions(+), 11 deletions(-) diff --git a/src/google/adk/memory/in_memory_memory_service.py b/src/google/adk/memory/in_memory_memory_service.py index 02276598cb9..5ff78a83733 100644 --- a/src/google/adk/memory/in_memory_memory_service.py +++ b/src/google/adk/memory/in_memory_memory_service.py @@ -33,8 +33,8 @@ _UNKNOWN_SESSION_ID = '__unknown_session_id__' -def _user_key(app_name: str, user_id: str) -> str: - return f'{app_name}/{user_id}' +def _user_key(app_name: str, user_id: str) -> tuple[str, str]: + return (app_name, user_id) def _extract_words_lower(text: str) -> set[str]: @@ -54,8 +54,8 @@ class InMemoryMemoryService(BaseMemoryService): def __init__(self): self._lock = threading.Lock() - self._session_events: dict[str, dict[str, list[Event]]] = {} - """Keys are "{app_name}/{user_id}". Values are dicts of session_id to + self._session_events: dict[tuple[str, str], dict[str, list[Event]]] = {} + """Keys are (app_name, user_id). Values are dicts of session_id to session event lists. """ diff --git a/tests/unittests/memory/test_in_memory_memory_service.py b/tests/unittests/memory/test_in_memory_memory_service.py index d50692f0bcd..bb03b985529 100644 --- a/tests/unittests/memory/test_in_memory_memory_service.py +++ b/tests/unittests/memory/test_in_memory_memory_service.py @@ -108,7 +108,7 @@ async def test_add_session_to_memory(): memory_service = InMemoryMemoryService() await memory_service.add_session_to_memory(MOCK_SESSION_1) - user_key = f'{MOCK_APP_NAME}/{MOCK_USER_ID}' + user_key = (MOCK_APP_NAME, MOCK_USER_ID) assert user_key in memory_service._session_events session_memory = memory_service._session_events[user_key] assert MOCK_SESSION_1.id in session_memory @@ -129,7 +129,7 @@ async def test_add_events_to_memory_with_explicit_events(): events=[MOCK_SESSION_1.events[0]], ) - user_key = f'{MOCK_APP_NAME}/{MOCK_USER_ID}' + user_key = (MOCK_APP_NAME, MOCK_USER_ID) session_memory = memory_service._session_events[user_key] assert len(session_memory[MOCK_SESSION_1.id]) == 1 assert session_memory[MOCK_SESSION_1.id][0].id == 'event-1a' @@ -145,7 +145,7 @@ async def test_add_events_to_memory_without_session_id_uses_default_bucket(): events=[MOCK_SESSION_1.events[0]], ) - user_key = f'{MOCK_APP_NAME}/{MOCK_USER_ID}' + user_key = (MOCK_APP_NAME, MOCK_USER_ID) session_memory = memory_service._session_events[user_key] assert len(session_memory) == 1 unknown_session_events = next(iter(session_memory.values())) @@ -164,7 +164,7 @@ async def test_add_events_to_memory_alias_is_supported(): events=[MOCK_SESSION_1.events[0]], ) - user_key = f'{MOCK_APP_NAME}/{MOCK_USER_ID}' + user_key = (MOCK_APP_NAME, MOCK_USER_ID) session_memory = memory_service._session_events[user_key] assert [event.id for event in session_memory[MOCK_SESSION_1.id]] == [ 'event-1a' @@ -191,7 +191,7 @@ async def test_add_events_to_memory_appends_without_replacing(): events=[new_event], ) - user_key = f'{MOCK_APP_NAME}/{MOCK_USER_ID}' + user_key = (MOCK_APP_NAME, MOCK_USER_ID) session_memory = memory_service._session_events[user_key] assert [event.id for event in session_memory[MOCK_SESSION_1.id]] == [ 'event-1a', @@ -220,7 +220,7 @@ async def test_add_events_to_memory_deduplicates_event_ids(): events=[duplicate_event], ) - user_key = f'{MOCK_APP_NAME}/{MOCK_USER_ID}' + user_key = (MOCK_APP_NAME, MOCK_USER_ID) session_memory = memory_service._session_events[user_key] assert [event.id for event in session_memory[MOCK_SESSION_1.id]] == [ 'event-1a', @@ -234,7 +234,7 @@ async def test_add_session_with_no_events_to_memory(): memory_service = InMemoryMemoryService() await memory_service.add_session_to_memory(MOCK_SESSION_WITH_NO_EVENTS) - user_key = f'{MOCK_APP_NAME}/{MOCK_USER_ID}' + user_key = (MOCK_APP_NAME, MOCK_USER_ID) assert user_key in memory_service._session_events session_memory = memory_service._session_events[user_key] assert MOCK_SESSION_WITH_NO_EVENTS.id in session_memory @@ -327,3 +327,34 @@ async def test_search_memory_is_scoped_by_user(): assert ( result_other_user.memories[0].content.parts[0].text == 'This is a secret.' ) + + +@pytest.mark.asyncio +async def test_search_memory_does_not_collide_on_slash_in_identifiers(): + """Tests that a slash in app_name cannot alias another app/user pair.""" + memory_service = InMemoryMemoryService() + await memory_service.add_session_to_memory( + Session( + app_name='app/other-user', + user_id='user', + id='session-slashed-app', + last_update_time=1000, + events=[ + Event( + id='event-slashed-app', + invocation_id='inv-slashed-app', + author='user', + timestamp=12345, + content=types.Content( + parts=[types.Part(text='This is a secret.')] + ), + ), + ], + ) + ) + + result = await memory_service.search_memory( + app_name='app', user_id='other-user/user', query='secret' + ) + + assert not result.memories