fix(memory): scope Vertex RAG retrieval to the requesting app and user

VertexAiRagMemoryService ran top-k retrieval across the whole configured
corpus and dropped the other tenants' contexts afterwards. The response filter
made the result correct, but ranking still competed against every app and user
in the corpus, so a busy corpus could crowd a caller's own memories out of the
top-k entirely, and foreign context was transferred only to be discarded. This
is a recall and data transfer problem, not a disclosure one: no memory
belonging to another app or user was ever returned to the caller.

search_memory now lists the corpus, keeps the files whose display name names
the requesting app and user, and passes those file ids to VertexRagStore so
ranking happens inside that set. When the caller owns no files, retrieval is
skipped and an empty response is returned. Callers can now see memories that
the previous ranking crowded out, so result counts can go up.

Scoping is best effort and adds no permission requirement, but it does add up
to 10 list calls to each search. Those calls run on the SDK async surface, so
they are awaited rather than blocking the event loop. Scoping is abandoned,
rather than applied to the files listed so far, whenever the listing cannot be
completed: either a listing failure, such as a deployment whose credentials
can retrieve but not list, or a corpus larger than the roughly 1000 files that
page budget covers. Retrieval then runs unscoped exactly as it did before,
which the response filter still makes correct, and the reason is logged.
Applying a partial listing would instead hide the caller's own memories.
Corpora past that size therefore keep the old ranking behavior permanently.

Server-side metadata filtering was considered and rejected: it matches on the
RagFile user_metadata field, which is output only on uploaded files and which
this service has never populated, so it would exclude every memory already
stored.

No public interface changes.

Co-authored-by: George Weale <gweale@google.com>
PiperOrigin-RevId: 963641027

src/google/adk/memory/vertex_ai_rag_memory_service.py

@@ -20,6 +20,7 @@ import base64
 import binascii
 from collections import OrderedDict
 import json
+import logging
 import os
 import tempfile
 from typing import Optional
@@ -34,12 +35,22 @@ from .base_memory_service import SearchMemoryResponse
 from .memory_entry import MemoryEntry
 
 if TYPE_CHECKING:
+  import agentplatform
+
   from ..events.event import Event
   from ..sessions.session import Session
 
 
+logger = logging.getLogger("google_adk." + __name__)
+
 _SOURCE_DISPLAY_NAME_PREFIX = "adk-memory-v1."
 
+_RAG_FILE_PAGE_SIZE = 100
+
+# Scoping walks the corpus file list, so it is capped to keep the cost of a
+# search independent of how large a shared corpus grows.
+_MAX_RAG_FILE_PAGES = 10
+
 
 def _encode_source_display_name_part(value: str) -> str:
   return (
@@ -90,6 +101,69 @@ def _parse_source_display_name(
   return parts[0], parts[1], parts[2]
 
 
+async def _scoped_rag_resources(
+    client: agentplatform.AsyncClient,
+    rag_resources: list[types.VertexRagStoreRagResource],
+    app_name: str,
+    user_id: str,
+) -> list[types.VertexRagStoreRagResource] | None:
+  """Returns resources naming only the files owned by one app and user.
+
+  Returns None when a corpus cannot be listed within the page budget, in which
+  case the caller retrieves without narrowing the resources.
+  """
+  from agentplatform import types as agentplatform_types
+
+  scoped_resources: list[types.VertexRagStoreRagResource] = []
+  for rag_resource in rag_resources:
+    rag_corpus = rag_resource.rag_corpus
+    if not rag_corpus:
+      return None
+
+    rag_file_ids: list[str] = []
+    page_token: str | None = None
+    for _ in range(_MAX_RAG_FILE_PAGES):
+      response = await client.rag.list_files(
+          name=rag_corpus,
+          config=agentplatform_types.ListRagFilesConfig(
+              page_size=_RAG_FILE_PAGE_SIZE, page_token=page_token
+          ),
+      )
+      for rag_file in response.rag_files or []:
+        session_info = _parse_source_display_name(rag_file.display_name or "")
+        if (
+            not session_info
+            or session_info[0] != app_name
+            or session_info[1] != user_id
+            or not rag_file.name
+        ):
+          continue
+        # rag_file_ids takes the bare file id, not the full resource name.
+        rag_file_ids.append(rag_file.name.rsplit("/", 1)[-1])
+      page_token = response.next_page_token
+      if not page_token:
+        break
+
+    if page_token:
+      # Scoping to the files seen so far would hide the caller's own memories,
+      # so an incomplete listing is abandoned instead.
+      logger.warning(
+          "Listing %s did not finish within %d pages, so retrieval is not"
+          " scoped to the requesting app and user.",
+          rag_corpus,
+          _MAX_RAG_FILE_PAGES,
+      )
+      return None
+    if rag_file_ids:
+      scoped_resources.append(
+          types.VertexRagStoreRagResource(
+              rag_corpus=rag_corpus, rag_file_ids=rag_file_ids
+          )
+      )
+
+  return scoped_resources
+
+
 class VertexAiRagMemoryService(BaseMemoryService):
   """A memory service that uses Agent Platform RAG for storage and retrieval."""
 
@@ -223,12 +297,38 @@ class VertexAiRagMemoryService(BaseMemoryService):
 
     from ..events.event import Event
 
+    rag_resources = self._vertex_rag_store.rag_resources
+    if not rag_resources:
+      raise ValueError("Rag resources must be set.")
+
     client = agentplatform.Client(
         project=self._project, location=self._location
     ).aio
     try:
+      try:
+        scoped_resources = await _scoped_rag_resources(
+            client, rag_resources, app_name, user_id
+        )
+      except Exception:  # pylint: disable=broad-except
+        # Narrowing the resources only improves ranking and transfer; the
+        # response filter below keeps an unnarrowed retrieval correct.
+        logger.warning(
+            "Listing the corpus failed, so retrieval is not scoped to the"
+            " requesting app and user.",
+            exc_info=True,
+        )
+        scoped_resources = None
+
+      vertex_rag_store = self._vertex_rag_store
+      if scoped_resources is not None:
+        if not scoped_resources:
+          return SearchMemoryResponse()
+        vertex_rag_store = self._vertex_rag_store.model_copy(
+            update={"rag_resources": scoped_resources}
+        )
+
       response = await client.rag.retrieve_contexts(
-          vertex_rag_store=self._vertex_rag_store,
+          vertex_rag_store=vertex_rag_store,
           query=agentplatform_types.RagQuery(
               text=query,
               similarity_top_k=self._similarity_top_k,
@@ -241,8 +341,7 @@ class VertexAiRagMemoryService(BaseMemoryService):
     memory_results = []
     session_events_map: OrderedDict[str, list[list[Event]]] = OrderedDict()
     for context in response.contexts.contexts:
-      # filter out context that is not related
-      # TODO: Add server side filtering by app_name and user_id.
+      # Still required: retrieval is unscoped whenever listing did not finish.
       source_display_name = getattr(context, "source_display_name", "")
       if not isinstance(source_display_name, str):
         continue

tests/unittests/memory/test_vertex_ai_rag_memory_service.py

@@ -14,12 +14,14 @@
 
 import asyncio
 import json
+import logging
 import os
 import tempfile
 from types import SimpleNamespace
 
 from google.adk.events.event import Event
 from google.adk.memory.vertex_ai_rag_memory_service import _build_source_display_name
+from google.adk.memory.vertex_ai_rag_memory_service import _MAX_RAG_FILE_PAGES
 from google.adk.memory.vertex_ai_rag_memory_service import _SOURCE_DISPLAY_NAME_PREFIX
 from google.adk.memory.vertex_ai_rag_memory_service import VertexAiRagMemoryService
 from google.adk.sessions.session import Session
@@ -35,6 +37,42 @@ def _rag_context(source_display_name: str, text: str) -> SimpleNamespace:
   )
 
 
+def _rag_file(rag_file_id: str, source_display_name: str) -> SimpleNamespace:
+  # A listing entry reports the full resource name, not the bare file id.
+  return SimpleNamespace(
+      name=(
+          "projects/test-project/locations/us-central1/ragCorpora/1/ragFiles/"
+          + rag_file_id
+      ),
+      display_name=source_display_name,
+  )
+
+
+def _memory_texts(response) -> list[str]:
+  return [memory.content.parts[0].text for memory in response.memories]
+
+
+def _retrieved_store(fake_client) -> types.VertexRagStore:
+  return fake_client.rag.retrieve_contexts.call_args.kwargs["vertex_rag_store"]
+
+
+def _async_client(mocker):
+  """A client on the SDK async surface, where every RAG call is awaited."""
+  fake_client = mocker.Mock()
+  fake_client.aclose = mocker.AsyncMock()
+  fake_client.rag.list_files = mocker.AsyncMock()
+  fake_client.rag.retrieve_contexts = mocker.AsyncMock()
+  fake_client.rag.upload_file = mocker.AsyncMock()
+  return fake_client
+
+
+def _unlistable_client(mocker):
+  """A client for the tests that exercise the unscoped retrieval path."""
+  fake_client = _async_client(mocker)
+  fake_client.rag.list_files.side_effect = PermissionError("cannot list files")
+  return fake_client
+
+
 def _session() -> Session:
   return Session(
       app_name="demo.app",
@@ -96,10 +134,9 @@ async def test_search_memory_forwards_similarity_top_k(
       rag_corpus="unused",
       similarity_top_k=configured_top_k,
   )
-  fake_client = mocker.Mock()
-  fake_client.aclose = mocker.AsyncMock()
-  fake_client.rag.retrieve_contexts = mocker.AsyncMock(
-      return_value=SimpleNamespace(contexts=SimpleNamespace(contexts=[]))
+  fake_client = _unlistable_client(mocker)
+  fake_client.rag.retrieve_contexts.return_value = SimpleNamespace(
+      contexts=SimpleNamespace(contexts=[])
   )
   mocker.patch(
       "agentplatform.Client", return_value=mocker.Mock(aio=fake_client)
@@ -117,32 +154,191 @@ async def test_search_memory_forwards_similarity_top_k(
   assert kwargs["vertex_rag_store"].similarity_top_k is None
 
 
+@pytest.mark.asyncio
+async def test_search_memory_scopes_retrieval_to_tenant_files(mocker):
+  """Ranking happens over only the requesting app and user's files."""
+  memory_service = VertexAiRagMemoryService(
+      rag_corpus="corpus", similarity_top_k=5
+  )
+
+  fake_client = _async_client(mocker)
+  fake_client.rag.list_files.side_effect = [
+      SimpleNamespace(
+          rag_files=[
+              _rag_file(
+                  "alice-1",
+                  _build_source_display_name("demo", "alice", "session-1"),
+              ),
+              _rag_file(
+                  "bob-1",
+                  _build_source_display_name("demo", "bob", "session-2"),
+              ),
+          ],
+          next_page_token="page-2",
+      ),
+      SimpleNamespace(
+          rag_files=[
+              _rag_file("alice-2", "demo.alice.legacy-session"),
+              _rag_file(
+                  "other-app-1",
+                  _build_source_display_name("other", "alice", "session-3"),
+              ),
+          ],
+          next_page_token=None,
+      ),
+  ]
+  fake_client.rag.retrieve_contexts.return_value = SimpleNamespace(
+      contexts=SimpleNamespace(
+          contexts=[
+              _rag_context(
+                  _build_source_display_name("demo", "alice", "session-1"),
+                  "ALICE_MEMORY",
+              )
+          ]
+      )
+  )
+  mocker.patch(
+      "agentplatform.Client", return_value=mocker.Mock(aio=fake_client)
+  )
+
+  response = await memory_service.search_memory(
+      app_name="demo", user_id="alice", query="memory"
+  )
+
+  assert _memory_texts(response) == ["ALICE_MEMORY"]
+  retrieve_kwargs = fake_client.rag.retrieve_contexts.call_args.kwargs
+  scoped_store = retrieve_kwargs["vertex_rag_store"]
+  scoped_resources = scoped_store.rag_resources
+  assert [resource.rag_corpus for resource in scoped_resources] == ["corpus"]
+  assert scoped_resources[0].rag_file_ids == ["alice-1", "alice-2"]
+  # Rebuilding the store keeps top-k on the query and off the store.
+  assert retrieve_kwargs["query"].similarity_top_k == 5
+  assert scoped_store.similarity_top_k is None
+  assert fake_client.rag.list_files.await_count == 2
+  assert (
+      fake_client.rag.list_files.call_args_list[1].kwargs["config"].page_token
+      == "page-2"
+  )
+  fake_client.aclose.assert_awaited_once()
+
+
+@pytest.mark.asyncio
+async def test_search_memory_skips_retrieval_without_tenant_files(mocker):
+  memory_service = VertexAiRagMemoryService(rag_corpus="corpus")
+
+  fake_client = _async_client(mocker)
+  fake_client.rag.list_files.return_value = SimpleNamespace(
+      rag_files=[
+          _rag_file(
+              "bob-1",
+              _build_source_display_name("demo", "bob", "session-2"),
+          )
+      ],
+      next_page_token=None,
+  )
+  mocker.patch(
+      "agentplatform.Client", return_value=mocker.Mock(aio=fake_client)
+  )
+
+  response = await memory_service.search_memory(
+      app_name="demo", user_id="alice", query="memory"
+  )
+
+  assert response.memories == []
+  fake_client.rag.retrieve_contexts.assert_not_awaited()
+  # The early return still leaves through the finally that closes the client.
+  fake_client.aclose.assert_awaited_once()
+
+
+@pytest.mark.asyncio
+async def test_search_memory_retrieves_unscoped_when_listing_fails(mocker):
+  """A deployment that cannot list files still retrieves its own memories."""
+  memory_service = VertexAiRagMemoryService(rag_corpus="corpus")
+
+  fake_client = _unlistable_client(mocker)
+  fake_client.rag.retrieve_contexts.return_value = SimpleNamespace(
+      contexts=SimpleNamespace(
+          contexts=[
+              _rag_context(
+                  _build_source_display_name("demo", "alice", "session-1"),
+                  "ALICE_MEMORY",
+              ),
+              _rag_context(
+                  _build_source_display_name("demo", "bob", "session-2"),
+                  "BOB_MEMORY",
+              ),
+          ]
+      )
+  )
+  mocker.patch(
+      "agentplatform.Client", return_value=mocker.Mock(aio=fake_client)
+  )
+
+  response = await memory_service.search_memory(
+      app_name="demo", user_id="alice", query="memory"
+  )
+
+  assert _memory_texts(response) == ["ALICE_MEMORY"]
+  assert _retrieved_store(fake_client).rag_resources[0].rag_file_ids is None
+
+
+@pytest.mark.asyncio
+async def test_search_memory_retrieves_unscoped_when_corpus_is_too_large(
+    mocker, caplog
+):
+  """Listing is capped so search cost stays independent of corpus size."""
+  memory_service = VertexAiRagMemoryService(rag_corpus="corpus")
+
+  fake_client = _async_client(mocker)
+  fake_client.rag.list_files.return_value = SimpleNamespace(
+      rag_files=[
+          _rag_file(
+              "alice-1",
+              _build_source_display_name("demo", "alice", "session-1"),
+          )
+      ],
+      next_page_token="another-page",
+  )
+  fake_client.rag.retrieve_contexts.return_value = SimpleNamespace(
+      contexts=SimpleNamespace(contexts=[])
+  )
+  mocker.patch(
+      "agentplatform.Client", return_value=mocker.Mock(aio=fake_client)
+  )
+
+  with caplog.at_level(logging.WARNING):
+    await memory_service.search_memory(
+        app_name="demo", user_id="alice", query="memory"
+    )
+
+  assert fake_client.rag.list_files.await_count == _MAX_RAG_FILE_PAGES
+  assert _retrieved_store(fake_client).rag_resources[0].rag_file_ids is None
+  assert "not scoped to the requesting app and user" in caplog.text
+
+
 @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."""
   memory_service = VertexAiRagMemoryService(rag_corpus="unused")
 
-  fake_client = mocker.Mock()
-  fake_client.aclose = mocker.AsyncMock()
-  fake_client.rag.retrieve_contexts = mocker.AsyncMock(
-      return_value=SimpleNamespace(
-          contexts=SimpleNamespace(
-              contexts=[
-                  _rag_context(
-                      "demo.alice.smith.session_secret",
-                      "SECRET_FROM_ALICE_SMITH",
-                  ),
-                  _rag_context(
-                      _build_source_display_name("demo", "alice", "session_ok"),
-                      "NORMAL_ALICE_MEMORY",
-                  ),
-                  _rag_context(
-                      "demo.alice.legacy_session",
-                      "LEGACY_ALICE_MEMORY",
-                  ),
-                  _rag_context("demo.bob.session_other", "BOB_MEMORY"),
-              ]
-          )
+  fake_client = _unlistable_client(mocker)
+  fake_client.rag.retrieve_contexts.return_value = SimpleNamespace(
+      contexts=SimpleNamespace(
+          contexts=[
+              _rag_context(
+                  "demo.alice.smith.session_secret",
+                  "SECRET_FROM_ALICE_SMITH",
+              ),
+              _rag_context(
+                  _build_source_display_name("demo", "alice", "session_ok"),
+                  "NORMAL_ALICE_MEMORY",
+              ),
+              _rag_context(
+                  "demo.alice.legacy_session",
+                  "LEGACY_ALICE_MEMORY",
+              ),
+              _rag_context("demo.bob.session_other", "BOB_MEMORY"),
+          ]
       )
   )
 
@@ -165,10 +361,7 @@ async def test_add_and_search_memory_uses_unambiguous_display_names(
 ):
   memory_service = VertexAiRagMemoryService(rag_corpus="unused")
 
-  fake_client = mocker.Mock()
-  fake_client.aclose = mocker.AsyncMock()
-  fake_client.rag.upload_file = mocker.AsyncMock()
-  fake_client.rag.retrieve_contexts = mocker.AsyncMock()
+  fake_client = _unlistable_client(mocker)
   mocker.patch(
       "agentplatform.Client", return_value=mocker.Mock(aio=fake_client)
   )