From b17b9c97b84f60a69d71fa03063b3d15419afc88 Mon Sep 17 00:00:00 2001
From: Yatsuiii <battyrises@gmail.com>
Date: Sun, 30 Aug 2026 02:18:36 +0530
Subject: [PATCH] fix(memory): return custom_metadata from
 VertexAiMemoryBankService.search_memory

search_memory constructed every returned MemoryEntry with only author,
content, and timestamp, silently dropping custom_metadata even though
this service's own write path already treats it as first-class:
_add_memories_via_create merges MemoryEntry.custom_metadata via
_merge_custom_metadata_for_memory before writing, and add_memory's own
docstring documents metadata as meaningful, round-tripped data.

Adds _from_vertex_metadata / _from_vertex_metadata_value as the
symmetric read-side counterpart to the existing write-side
_to_vertex_metadata_value, handling the same four value kinds
(bool/double/string/timestamp) rather than only strings.

Verified with a live create -> retrieve round trip against a real
Vertex AI Memory Bank instance before implementing, confirming the
underlying SDK genuinely returns populated metadata that this class
was never reading.

Fixes #6946
---
 .../memory/vertex_ai_memory_bank_service.py   | 36 +++++++++++
 .../test_vertex_ai_memory_bank_service.py     | 60 +++++++++++++++++++
 2 files changed, 96 insertions(+)

diff --git a/src/google/adk/memory/vertex_ai_memory_bank_service.py b/src/google/adk/memory/vertex_ai_memory_bank_service.py
index c5ae3f36a95..689f2ef838a 100644
--- a/src/google/adk/memory/vertex_ai_memory_bank_service.py
+++ b/src/google/adk/memory/vertex_ai_memory_bank_service.py
@@ -579,6 +579,7 @@ async def search_memory(
                       role='user',
                   ),
                   timestamp=update_time.isoformat() if update_time else None,
+                  custom_metadata=_from_vertex_metadata(memory.metadata),
               )
           )
         except AttributeError:
@@ -1033,3 +1034,38 @@ def _to_vertex_metadata_value(
     )
     return None
   return {'string_value': str(value)}
+
+
+def _from_vertex_metadata(
+    vertex_metadata: Mapping[str, object] | None,
+) -> dict[str, object]:
+  """Converts Vertex MemoryMetadataValue objects back to plain Python values.
+
+  The read-side counterpart to `_to_vertex_metadata_value`: a `Memory`
+  retrieved from Vertex AI Memory Bank carries whatever `custom_metadata` was
+  written at creation time in this same oneof shape, and this reconstructs
+  the values a caller passed in, rather than leaving them wrapped.
+  """
+  if not vertex_metadata:
+    return {}
+  return {
+      key: _from_vertex_metadata_value(value)
+      for key, value in vertex_metadata.items()
+  }
+
+
+def _from_vertex_metadata_value(value: object) -> object:
+  """Converts one Vertex MemoryMetadataValue back to a plain Python value."""
+  bool_value = getattr(value, 'bool_value', None)
+  if bool_value is not None:
+    return bool_value
+  double_value = getattr(value, 'double_value', None)
+  if double_value is not None:
+    return double_value
+  string_value = getattr(value, 'string_value', None)
+  if string_value is not None:
+    return string_value
+  timestamp_value = getattr(value, 'timestamp_value', None)
+  if timestamp_value is not None:
+    return timestamp_value
+  return value
diff --git a/tests/unittests/memory/test_vertex_ai_memory_bank_service.py b/tests/unittests/memory/test_vertex_ai_memory_bank_service.py
index df923f808c5..4e33cdebee5 100644
--- a/tests/unittests/memory/test_vertex_ai_memory_bank_service.py
+++ b/tests/unittests/memory/test_vertex_ai_memory_bank_service.py
@@ -1174,6 +1174,66 @@ async def test_search_memory(mock_vertexai_client):
 
   assert len(result.memories) == 1
   assert result.memories[0].content.parts[0].text == 'test_content'
+  assert result.memories[0].custom_metadata == {}
+
+
+@pytest.mark.asyncio
+async def test_search_memory_returns_custom_metadata(mock_vertexai_client):
+  """`search_memory` must round-trip `custom_metadata`: the write path
+  (`_merge_custom_metadata_for_memory`, `_add_memories_via_create`) already
+  treats it as first-class, so the read path silently dropping it is a
+  correctness gap, not a missing feature. Covers every value kind the write
+  side supports (bool/double/string/timestamp), not just strings."""
+  retrieved_memory = mock.MagicMock()
+  retrieved_memory.memory.fact = 'test_content'
+  retrieved_memory.memory.update_time = datetime.datetime(
+      2024, 12, 12, 12, 12, 12, 123456
+  )
+  retrieved_memory.memory.metadata = {
+      'a_bool': vertex_types.MemoryMetadataValue(bool_value=True),
+      'a_double': vertex_types.MemoryMetadataValue(double_value=1.5),
+      'a_string': vertex_types.MemoryMetadataValue(string_value='record-123'),
+  }
+
+  mock_vertexai_client.agent_engines.memories.retrieve.return_value = (
+      _AsyncListIterator([retrieved_memory])
+  )
+  memory_service = mock_vertex_ai_memory_bank_service()
+
+  result = await memory_service.search_memory(
+      app_name=MOCK_APP_NAME, user_id=MOCK_USER_ID, query='query'
+  )
+
+  assert len(result.memories) == 1
+  assert result.memories[0].custom_metadata == {
+      'a_bool': True,
+      'a_double': 1.5,
+      'a_string': 'record-123',
+  }
+
+
+@pytest.mark.asyncio
+async def test_search_memory_with_no_metadata_returns_empty_dict(
+    mock_vertexai_client,
+):
+  """A memory with no metadata (or written before this field existed) must
+  round-trip to an empty dict, never `None` and never a placeholder."""
+  retrieved_memory = mock.MagicMock()
+  retrieved_memory.memory.fact = 'test_content'
+  retrieved_memory.memory.update_time = None
+  retrieved_memory.memory.metadata = None
+
+  mock_vertexai_client.agent_engines.memories.retrieve.return_value = (
+      _AsyncListIterator([retrieved_memory])
+  )
+  memory_service = mock_vertex_ai_memory_bank_service()
+
+  result = await memory_service.search_memory(
+      app_name=MOCK_APP_NAME, user_id=MOCK_USER_ID, query='query'
+  )
+
+  assert len(result.memories) == 1
+  assert result.memories[0].custom_metadata == {}
 
 
 @pytest.mark.asyncio