fix: Store and retrieve usage_metadata in Vertex AI custom_metadata

The Vertex AI session service does not natively support persisting usage_metadata. This change serializes usage_metadata into the custom_metadata field under the key '_usage_metadata' when appending events and deserializes it back when retrieving events. This allows usage information to be round-tripped through the Vertex AI session service.

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

Author: GWeale

src/google/adk/sessions/vertex_ai_session_service.py

@@ -42,6 +42,20 @@
 
 logger = logging.getLogger('google_adk.' + __name__)
 
+_COMPACTION_CUSTOM_METADATA_KEY = '_compaction'
+_USAGE_METADATA_CUSTOM_METADATA_KEY = '_usage_metadata'
+
+
+def _set_internal_custom_metadata(
+    metadata_dict: dict[str, Any], *, key: str, value: dict[str, Any]
+) -> None:
+  """Stores internal metadata alongside user-provided custom metadata."""
+  existing_custom_metadata = metadata_dict.get('custom_metadata') or {}
+  metadata_dict['custom_metadata'] = {
+      **existing_custom_metadata,
+      key: value,
+  }
+
 
 class VertexAiSessionService(BaseSessionService):
   """Connects to the Vertex AI Agent Engine Session Service using Agent Engine SDK.
@@ -301,11 +315,22 @@ async def append_event(self, session: Session, event: Event) -> Event:
       compaction_dict = event.actions.compaction.model_dump(
           exclude_none=True, mode='json'
       )
-      existing_custom = metadata_dict.get('custom_metadata') or {}
-      metadata_dict['custom_metadata'] = {
-          **existing_custom,
-          '_compaction': compaction_dict,
-      }
+      _set_internal_custom_metadata(
+          metadata_dict,
+          key=_COMPACTION_CUSTOM_METADATA_KEY,
+          value=compaction_dict,
+      )
+    # Store usage_metadata in custom_metadata since the Vertex AI service
+    # does not persist it in EventMetadata.
+    if event.usage_metadata:
+      usage_dict = event.usage_metadata.model_dump(
+          exclude_none=True, mode='json'
+      )
+      _set_internal_custom_metadata(
+          metadata_dict,
+          key=_USAGE_METADATA_CUSTOM_METADATA_KEY,
+          value=usage_dict,
+      )
     config['event_metadata'] = metadata_dict
 
     async with self._get_api_client() as api_client:
@@ -378,11 +403,20 @@ def _from_api_event(api_event_obj: vertexai.types.SessionEvent) -> Event:
     # Extract compaction data stored in custom_metadata.
     # NOTE: This read path must be kept permanently because sessions
     # written before native compaction support store compaction data
-    # in custom_metadata under the '_compaction' key.
+    # in custom_metadata under the compaction metadata key.
     compaction_data = None
-    if custom_metadata and '_compaction' in custom_metadata:
+    usage_metadata_data = None
+    if custom_metadata and (
+        _COMPACTION_CUSTOM_METADATA_KEY in custom_metadata
+        or _USAGE_METADATA_CUSTOM_METADATA_KEY in custom_metadata
+    ):
       custom_metadata = dict(custom_metadata)  # avoid mutating the API response
-      compaction_data = custom_metadata.pop('_compaction')
+      compaction_data = custom_metadata.pop(
+          _COMPACTION_CUSTOM_METADATA_KEY, None
+      )
+      usage_metadata_data = custom_metadata.pop(
+          _USAGE_METADATA_CUSTOM_METADATA_KEY, None
+      )
       if not custom_metadata:
         custom_metadata = None
     grounding_metadata = _session_util.decode_model(
@@ -397,6 +431,7 @@ def _from_api_event(api_event_obj: vertexai.types.SessionEvent) -> Event:
     branch = None
     custom_metadata = None
     compaction_data = None
+    usage_metadata_data = None
     grounding_metadata = None
 
   if actions:
@@ -416,6 +451,12 @@ def _from_api_event(api_event_obj: vertexai.types.SessionEvent) -> Event:
     else:
       event_actions = EventActions()
 
+  usage_metadata = None
+  if usage_metadata_data:
+    usage_metadata = types.GenerateContentResponseUsageMetadata.model_validate(
+        usage_metadata_data
+    )
+
   return Event(
       id=api_event_obj.name.split('/')[-1],
       invocation_id=api_event_obj.invocation_id,
@@ -434,4 +475,5 @@ def _from_api_event(api_event_obj: vertexai.types.SessionEvent) -> Event:
       custom_metadata=custom_metadata,
       grounding_metadata=grounding_metadata,
       long_running_tool_ids=long_running_tool_ids,
+      usage_metadata=usage_metadata,
   )

tests/unittests/sessions/test_vertex_ai_session_service.py

@@ -911,3 +911,131 @@ async def test_append_event_with_compaction_and_custom_metadata():
   # User custom_metadata is preserved without the internal _compaction key
   assert appended_event.custom_metadata == {'user_key': 'user_value'}
   assert '_compaction' not in (appended_event.custom_metadata or {})
+
+
+@pytest.mark.asyncio
+@pytest.mark.usefixtures('mock_get_api_client')
+async def test_append_event_with_usage_metadata():
+  """usage_metadata round-trips through append_event and get_session."""
+  session_service = mock_vertex_ai_session_service()
+  session = await session_service.get_session(
+      app_name='123', user_id='user', session_id='1'
+  )
+  assert session is not None
+
+  event_to_append = Event(
+      invocation_id='usage_invocation',
+      author='model',
+      timestamp=1734005536.0,
+      usage_metadata=genai_types.GenerateContentResponseUsageMetadata(
+          prompt_token_count=150,
+          candidates_token_count=50,
+          total_token_count=200,
+      ),
+  )
+
+  await session_service.append_event(session, event_to_append)
+
+  retrieved_session = await session_service.get_session(
+      app_name='123', user_id='user', session_id='1'
+  )
+  assert retrieved_session is not None
+
+  appended_event = retrieved_session.events[-1]
+  assert appended_event.usage_metadata is not None
+  assert appended_event.usage_metadata.prompt_token_count == 150
+  assert appended_event.usage_metadata.candidates_token_count == 50
+  assert appended_event.usage_metadata.total_token_count == 200
+  # custom_metadata should remain None when only usage_metadata was stored
+  assert appended_event.custom_metadata is None
+
+
+@pytest.mark.asyncio
+@pytest.mark.usefixtures('mock_get_api_client')
+async def test_append_event_with_usage_metadata_and_custom_metadata():
+  """Both usage_metadata and user custom_metadata survive the round-trip."""
+  session_service = mock_vertex_ai_session_service()
+  session = await session_service.get_session(
+      app_name='123', user_id='user', session_id='1'
+  )
+  assert session is not None
+
+  event_to_append = Event(
+      invocation_id='usage_and_meta_invocation',
+      author='model',
+      timestamp=1734005537.0,
+      usage_metadata=genai_types.GenerateContentResponseUsageMetadata(
+          prompt_token_count=300,
+          total_token_count=400,
+      ),
+      custom_metadata={'my_key': 'my_value'},
+  )
+
+  await session_service.append_event(session, event_to_append)
+
+  retrieved_session = await session_service.get_session(
+      app_name='123', user_id='user', session_id='1'
+  )
+  assert retrieved_session is not None
+
+  appended_event = retrieved_session.events[-1]
+  # usage_metadata is restored
+  assert appended_event.usage_metadata is not None
+  assert appended_event.usage_metadata.prompt_token_count == 300
+  assert appended_event.usage_metadata.total_token_count == 400
+  # User custom_metadata is preserved without internal keys
+  assert appended_event.custom_metadata == {'my_key': 'my_value'}
+  assert '_usage_metadata' not in (appended_event.custom_metadata or {})
+
+
+@pytest.mark.asyncio
+@pytest.mark.usefixtures('mock_get_api_client')
+async def test_append_event_with_usage_metadata_and_compaction():
+  """usage_metadata, compaction, and user custom_metadata all coexist."""
+  session_service = mock_vertex_ai_session_service()
+  session = await session_service.get_session(
+      app_name='123', user_id='user', session_id='1'
+  )
+  assert session is not None
+
+  compaction = EventCompaction(
+      start_timestamp=500.0,
+      end_timestamp=600.0,
+      compacted_content=genai_types.Content(
+          parts=[genai_types.Part(text='compacted')]
+      ),
+  )
+  event_to_append = Event(
+      invocation_id='all_three_invocation',
+      author='model',
+      timestamp=1734005538.0,
+      actions=EventActions(compaction=compaction),
+      usage_metadata=genai_types.GenerateContentResponseUsageMetadata(
+          prompt_token_count=1000,
+          candidates_token_count=250,
+          total_token_count=1250,
+      ),
+      custom_metadata={'extra': 'info'},
+  )
+
+  await session_service.append_event(session, event_to_append)
+
+  retrieved_session = await session_service.get_session(
+      app_name='123', user_id='user', session_id='1'
+  )
+  assert retrieved_session is not None
+
+  appended_event = retrieved_session.events[-1]
+  # Compaction is restored
+  assert appended_event.actions.compaction is not None
+  assert appended_event.actions.compaction.start_timestamp == 500.0
+  assert appended_event.actions.compaction.end_timestamp == 600.0
+  # usage_metadata is restored
+  assert appended_event.usage_metadata is not None
+  assert appended_event.usage_metadata.prompt_token_count == 1000
+  assert appended_event.usage_metadata.candidates_token_count == 250
+  assert appended_event.usage_metadata.total_token_count == 1250
+  # User custom_metadata is preserved without internal keys
+  assert appended_event.custom_metadata == {'extra': 'info'}
+  assert '_compaction' not in (appended_event.custom_metadata or {})
+  assert '_usage_metadata' not in (appended_event.custom_metadata or {})