chore(session): refine raw_event storage with safety fixes and fallback test

- Remove outdated TODOs regarding requested_tool_confirmations and native compaction support
- Add a warning log when falling back to legacy storage
- Avoid overwriting valid timestamps from `raw_event` with `None` if the API returns `None`
- Provide a fallback timestamp (current time) to prevent Pydantic validation errors when the API returns no timestamp

Co-authored-by: Shangjie Chen <deanchen@google.com>
PiperOrigin-RevId: 899195056

Author: DeanChensj

src/google/adk/sessions/vertex_ai_session_service.py

@@ -270,6 +270,7 @@ async def append_event(self, session: Session, event: Event) -> Event:
 
     reasoning_engine_id = self._get_reasoning_engine_id(session.app_name)
 
+    # Build config (Monolithic approach)
     config = {}
     if event.content:
       config['content'] = event.content.model_dump(
@@ -286,9 +287,6 @@ async def append_event(self, session: Session, event: Event) -> Event:
               k: json.loads(v.model_dump_json(exclude_none=True, by_alias=True))
               for k, v in event.actions.requested_auth_configs.items()
           },
-          # TODO: add requested_tool_confirmations, agent_state once
-          # they are available in the API.
-          # Note: compaction is stored via event_metadata.custom_metadata.
       }
     if event.error_code:
       config['error_code'] = event.error_code
@@ -311,10 +309,8 @@ async def append_event(self, session: Session, event: Event) -> Event:
       metadata_dict['grounding_metadata'] = event.grounding_metadata.model_dump(
           exclude_none=True, mode='json'
       )
-    # Store compaction data in custom_metadata since the Vertex AI service
-    # does not yet support the compaction field.
-    # TODO: Stop writing to custom_metadata once the Vertex AI service
-    # supports the compaction field natively in EventActions.
+
+    # ALWAYS write to custom_metadata
     if event.actions and event.actions.compaction:
       compaction_dict = event.actions.compaction.model_dump(
           exclude_none=True, mode='json'
@@ -324,8 +320,6 @@ async def append_event(self, session: Session, event: Event) -> Event:
           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'
@@ -335,7 +329,12 @@ async def append_event(self, session: Session, event: Event) -> Event:
           key=_USAGE_METADATA_CUSTOM_METADATA_KEY,
           value=usage_dict,
       )
+
     config['event_metadata'] = metadata_dict
+
+    # Persist the full event state using raw_event. If the client-side SDK
+    # does not support this field, it will raise a ValidationError, and we
+    # will fall back to legacy field-based storage.
     config['raw_event'] = event.model_dump(
         exclude_none=True,
         mode='json',
@@ -345,7 +344,8 @@ async def append_event(self, session: Session, event: Event) -> Event:
     # Retry without raw_event if client side validation fails for older SDK
     # versions.
     async with self._get_api_client() as api_client:
-      try:
+
+      async def _do_append(cfg: dict[str, Any]):
         await api_client.agent_engines.sessions.events.append(
             name=(
                 f'reasoningEngines/{reasoning_engine_id}/sessions/{session.id}'
@@ -355,22 +355,16 @@ async def append_event(self, session: Session, event: Event) -> Event:
             timestamp=datetime.datetime.fromtimestamp(
                 event.timestamp, tz=datetime.timezone.utc
             ),
-            config=config,
+            config=cfg,
         )
+
+      try:
+        await _do_append(config)
       except pydantic.ValidationError:
+        logger.warning('Vertex SDK does not support raw_event, falling back.')
         if 'raw_event' in config:
           del config['raw_event']
-        await api_client.agent_engines.sessions.events.append(
-            name=(
-                f'reasoningEngines/{reasoning_engine_id}/sessions/{session.id}'
-            ),
-            author=event.author,
-            invocation_id=event.invocation_id,
-            timestamp=datetime.datetime.fromtimestamp(
-                event.timestamp, tz=datetime.timezone.utc
-            ),
-            config=config,
-        )
+        await _do_append(config)
     return event
 
   def _get_reasoning_engine_id(self, app_name: str):
@@ -429,8 +423,8 @@ def _get_raw_event(api_event_obj: Any) -> dict[str, Any] | None:
 
 def _from_api_event(api_event_obj: vertexai.types.SessionEvent) -> Event:
   """Converts an API event object to an Event object."""
-  # Read event data from raw_event first before falling back to top level
-  # fields.
+  # Prioritize reading from raw_event to restore full state. Fall back to
+  # top-level fields for older data that lacks raw_event.
   raw_event_dict = _get_raw_event(api_event_obj)
   if raw_event_dict:
     event_dict = copy.deepcopy(raw_event_dict)
@@ -439,8 +433,9 @@ def _from_api_event(api_event_obj: vertexai.types.SessionEvent) -> Event:
         'id': api_event_obj.name.split('/')[-1],
         'invocation_id': getattr(api_event_obj, 'invocation_id', None),
         'author': getattr(api_event_obj, 'author', None),
-        'timestamp': timestamp_obj.timestamp() if timestamp_obj else None,
     })
+    if timestamp_obj:
+      event_dict['timestamp'] = timestamp_obj.timestamp()
     return Event.model_validate(event_dict)
 
   actions = getattr(api_event_obj, 'actions', None)
@@ -514,6 +509,13 @@ def _from_api_event(api_event_obj: vertexai.types.SessionEvent) -> Event:
         usage_metadata_data
     )
 
+  timestamp_obj = getattr(api_event_obj, 'timestamp', None)
+  timestamp = (
+      timestamp_obj.timestamp()
+      if timestamp_obj
+      else datetime.datetime.now(datetime.timezone.utc).timestamp()
+  )
+
   return Event(
       id=api_event_obj.name.split('/')[-1],
       invocation_id=api_event_obj.invocation_id,
@@ -522,7 +524,7 @@ def _from_api_event(api_event_obj: vertexai.types.SessionEvent) -> Event:
       content=_session_util.decode_model(
           getattr(api_event_obj, 'content', None), types.Content
       ),
-      timestamp=api_event_obj.timestamp.timestamp(),
+      timestamp=timestamp,
       error_code=getattr(api_event_obj, 'error_code', None),
       error_message=getattr(api_event_obj, 'error_message', None),
       partial=partial,

tests/unittests/sessions/test_vertex_ai_session_service.py

@@ -35,6 +35,7 @@
 from google.api_core import exceptions as api_core_exceptions
 from google.genai import types as genai_types
 from google.genai.errors import ClientError
+import pydantic
 import pytest
 
 MOCK_SESSION_JSON_1 = {
@@ -1223,3 +1224,55 @@ async def test_append_event_with_usage_metadata_and_compaction():
   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 {})
+
+
+@pytest.mark.asyncio
+@pytest.mark.usefixtures('mock_get_api_client')
+async def test_append_event_fallback_for_older_sdk(mock_api_client_instance):
+  """Tests that append_event falls back to custom_metadata when SDK fails on raw_event."""
+  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=1000.0,
+      end_timestamp=2000.0,
+      compacted_content=genai_types.Content(
+          parts=[genai_types.Part(text='compacted summary')]
+      ),
+  )
+  event_to_append = Event(
+      invocation_id='fallback_invocation',
+      author='model',
+      timestamp=1734005534.0,
+      actions=EventActions(compaction=compaction),
+  )
+
+  mock_client = mock_api_client_instance
+
+  async def side_effect(name, author, invocation_id, timestamp, config):
+    if 'raw_event' in config:
+      # Trigger a real ValidationError since Pydantic V2 doesn't allow easy
+      # instantiation
+      class DummyModel(pydantic.BaseModel):
+        a: int
+
+      DummyModel(a='not an int')
+    return await mock_client._append_event(
+        name, author, invocation_id, timestamp, config
+    )
+
+  mock_client.agent_engines.sessions.events.append.side_effect = side_effect
+
+  await session_service.append_event(session, event_to_append)
+
+  # Verify that it was written and restored correctly via custom_metadata
+  retrieved_session = await session_service.get_session(
+      app_name='123', user_id='user', session_id='1'
+  )
+  appended_event = retrieved_session.events[-1]
+
+  assert appended_event.actions.compaction is not None
+  assert appended_event.actions.compaction.start_timestamp == 1000.0