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/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/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/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/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 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 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 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():