From b2e0cb15ddfcef962df1405168b978dd1085cb7d Mon Sep 17 00:00:00 2001 From: weiguangli-io Date: Thu, 5 Mar 2026 11:46:01 +0800 Subject: [PATCH 1/4] fix: preserve cache_metadata and usage_metadata in VertexAiSessionService event round-trip VertexAiSessionService was dropping cache_metadata and usage_metadata fields during Event serialization/deserialization. This caused ContextCacheRequestProcessor to never find previous cache metadata, creating a new cache on every LLM call instead of reusing existing ones. The fix adds cache_metadata and usage_metadata to the event_metadata dict during append_event (write path) and restores them in _from_api_event (read path), matching the behavior of other session service implementations. Fixes #4698 --- .../adk/sessions/vertex_ai_session_service.py | 21 +++++++ .../test_vertex_ai_session_service.py | 62 +++++++++++++++++++ 2 files changed, 83 insertions(+) diff --git a/src/google/adk/sessions/vertex_ai_session_service.py b/src/google/adk/sessions/vertex_ai_session_service.py index 8c1fdc134e..068ccd93e2 100644 --- a/src/google/adk/sessions/vertex_ai_session_service.py +++ b/src/google/adk/sessions/vertex_ai_session_service.py @@ -36,6 +36,7 @@ from ..events.event import Event from ..events.event_actions import EventActions from ..events.event_actions import EventCompaction +from ..models.cache_metadata import CacheMetadata from ..utils.vertex_ai_utils import get_express_mode_api_key from .base_session_service import BaseSessionService from .base_session_service import GetSessionConfig @@ -311,6 +312,14 @@ async def append_event(self, session: Session, event: Event) -> Event: else None ), } + if event.usage_metadata: + metadata_dict['usage_metadata'] = event.usage_metadata.model_dump( + exclude_none=True, mode='json' + ) + if event.cache_metadata: + metadata_dict['cache_metadata'] = event.cache_metadata.model_dump( + exclude_none=True, mode='json' + ) if event.grounding_metadata: metadata_dict['grounding_metadata'] = event.grounding_metadata.model_dump( exclude_none=True, mode='json' @@ -481,6 +490,14 @@ def _from_api_event(api_event_obj: vertexai.types.SessionEvent) -> Event: getattr(event_metadata, 'grounding_metadata', None), types.GroundingMetadata, ) + usage_metadata = _session_util.decode_model( + getattr(event_metadata, 'usage_metadata', None), + types.GenerateContentResponseUsageMetadata, + ) + cache_metadata = _session_util.decode_model( + getattr(event_metadata, 'cache_metadata', None), + CacheMetadata, + ) else: long_running_tool_ids = None partial = None @@ -491,6 +508,8 @@ def _from_api_event(api_event_obj: vertexai.types.SessionEvent) -> Event: compaction_data = None usage_metadata_data = None grounding_metadata = None + usage_metadata = None + cache_metadata = None if actions: actions_dict = actions.model_dump(exclude_none=True, mode='python') @@ -539,6 +558,8 @@ def _from_api_event(api_event_obj: vertexai.types.SessionEvent) -> Event: branch=branch, custom_metadata=custom_metadata, grounding_metadata=grounding_metadata, + usage_metadata=usage_metadata, + cache_metadata=cache_metadata, long_running_tool_ids=long_running_tool_ids, usage_metadata=usage_metadata, ) diff --git a/tests/unittests/sessions/test_vertex_ai_session_service.py b/tests/unittests/sessions/test_vertex_ai_session_service.py index 0b265c3cfb..78841200d1 100644 --- a/tests/unittests/sessions/test_vertex_ai_session_service.py +++ b/tests/unittests/sessions/test_vertex_ai_session_service.py @@ -344,6 +344,8 @@ def _convert_to_object(data): 'requested_auth_configs', 'rawEvent', 'raw_event', + 'cache_metadata', + 'usage_metadata', ]: kwargs[key] = value else: @@ -1306,3 +1308,63 @@ class DummyModel(pydantic.BaseModel): assert appended_event.actions.compaction is not None assert appended_event.actions.compaction.start_timestamp == 1000.0 + + +@pytest.mark.asyncio +@pytest.mark.usefixtures('mock_get_api_client') +async def test_append_event_with_cache_and_usage_metadata(): + """cache_metadata and usage_metadata round-trip through append and get.""" + 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 + + cache_meta = CacheMetadata( + cache_name='projects/123/locations/us-central1/cachedContents/456', + expire_time=9999999999.0, + fingerprint='abc123hash', + invocations_used=3, + contents_count=10, + created_at=1700000000.0, + ) + usage_meta = genai_types.GenerateContentResponseUsageMetadata( + prompt_token_count=100, + candidates_token_count=50, + total_token_count=150, + cached_content_token_count=80, + ) + event_to_append = Event( + invocation_id='cache_test_invocation', + author='model', + timestamp=1734005536.0, + content=genai_types.Content( + parts=[genai_types.Part(text='cached response')] + ), + cache_metadata=cache_meta, + usage_metadata=usage_meta, + ) + + 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] + # cache_metadata is preserved + assert appended_event.cache_metadata is not None + assert appended_event.cache_metadata.cache_name == ( + 'projects/123/locations/us-central1/cachedContents/456' + ) + assert appended_event.cache_metadata.fingerprint == 'abc123hash' + assert appended_event.cache_metadata.invocations_used == 3 + assert appended_event.cache_metadata.contents_count == 10 + assert appended_event.cache_metadata.created_at == 1700000000.0 + # usage_metadata is preserved + assert appended_event.usage_metadata is not None + assert appended_event.usage_metadata.prompt_token_count == 100 + assert appended_event.usage_metadata.candidates_token_count == 50 + assert appended_event.usage_metadata.total_token_count == 150 + assert appended_event.usage_metadata.cached_content_token_count == 80 From a0e51e82232e93b802645b57ded4dfd3157f95e3 Mon Sep 17 00:00:00 2001 From: weiguangli-io Date: Mon, 9 Mar 2026 20:56:35 +0800 Subject: [PATCH 2/4] refactor(test): simplify metadata assertions with object equality Applied Gemini Code Assist suggestion to make test more concise and maintainable by using direct object comparison instead of field-by-field assertions. Pydantic models support equality comparison natively. --- .../sessions/test_vertex_ai_session_service.py | 15 ++------------- 1 file changed, 2 insertions(+), 13 deletions(-) diff --git a/tests/unittests/sessions/test_vertex_ai_session_service.py b/tests/unittests/sessions/test_vertex_ai_session_service.py index 78841200d1..1be9ac0638 100644 --- a/tests/unittests/sessions/test_vertex_ai_session_service.py +++ b/tests/unittests/sessions/test_vertex_ai_session_service.py @@ -1354,17 +1354,6 @@ async def test_append_event_with_cache_and_usage_metadata(): appended_event = retrieved_session.events[-1] # cache_metadata is preserved - assert appended_event.cache_metadata is not None - assert appended_event.cache_metadata.cache_name == ( - 'projects/123/locations/us-central1/cachedContents/456' - ) - assert appended_event.cache_metadata.fingerprint == 'abc123hash' - assert appended_event.cache_metadata.invocations_used == 3 - assert appended_event.cache_metadata.contents_count == 10 - assert appended_event.cache_metadata.created_at == 1700000000.0 + assert appended_event.cache_metadata == cache_meta # usage_metadata is preserved - assert appended_event.usage_metadata is not None - assert appended_event.usage_metadata.prompt_token_count == 100 - assert appended_event.usage_metadata.candidates_token_count == 50 - assert appended_event.usage_metadata.total_token_count == 150 - assert appended_event.usage_metadata.cached_content_token_count == 80 + assert appended_event.usage_metadata == usage_meta From 1d0c123103e48fd024f9640d9ecdc361139b3257 Mon Sep 17 00:00:00 2001 From: weiguangli-io Date: Tue, 17 Mar 2026 22:24:09 +0800 Subject: [PATCH 3/4] style: autoformat contributing samples --- contributing/samples/gepa/experiment.py | 1 + contributing/samples/gepa/run_experiment.py | 1 + 2 files changed, 2 insertions(+) diff --git a/contributing/samples/gepa/experiment.py b/contributing/samples/gepa/experiment.py index b03179c231..df2079b727 100644 --- a/contributing/samples/gepa/experiment.py +++ b/contributing/samples/gepa/experiment.py @@ -44,6 +44,7 @@ from tau_bench.types import EnvRunResult from tau_bench.types import RunConfig import tau_bench_agent as tau_bench_agent_lib +import utils def run_tau_bench_rollouts( diff --git a/contributing/samples/gepa/run_experiment.py b/contributing/samples/gepa/run_experiment.py index ca3f5e7696..92299bcc12 100644 --- a/contributing/samples/gepa/run_experiment.py +++ b/contributing/samples/gepa/run_experiment.py @@ -26,6 +26,7 @@ import experiment import gepa_utils from google.genai import types +import utils _OUTPUT_DIR = flags.DEFINE_string( 'output_dir', From 8c4ffcf73a3cfd525a6a046986f5e5ba32ced229 Mon Sep 17 00:00:00 2001 From: weiguangli-io Date: Sun, 22 Mar 2026 21:19:57 +0800 Subject: [PATCH 4/4] fix: remove duplicate usage_metadata handling after rebase Upstream main already handles usage_metadata round-trip via _USAGE_METADATA_CUSTOM_METADATA_KEY in custom_metadata. Remove the redundant event_metadata-based usage_metadata storage/retrieval that was added in the initial commit, which caused a SyntaxError (repeated keyword argument) when merged with upstream. --- src/google/adk/sessions/vertex_ai_session_service.py | 10 ---------- 1 file changed, 10 deletions(-) diff --git a/src/google/adk/sessions/vertex_ai_session_service.py b/src/google/adk/sessions/vertex_ai_session_service.py index 068ccd93e2..aa8f8306f1 100644 --- a/src/google/adk/sessions/vertex_ai_session_service.py +++ b/src/google/adk/sessions/vertex_ai_session_service.py @@ -312,10 +312,6 @@ async def append_event(self, session: Session, event: Event) -> Event: else None ), } - if event.usage_metadata: - metadata_dict['usage_metadata'] = event.usage_metadata.model_dump( - exclude_none=True, mode='json' - ) if event.cache_metadata: metadata_dict['cache_metadata'] = event.cache_metadata.model_dump( exclude_none=True, mode='json' @@ -490,10 +486,6 @@ def _from_api_event(api_event_obj: vertexai.types.SessionEvent) -> Event: getattr(event_metadata, 'grounding_metadata', None), types.GroundingMetadata, ) - usage_metadata = _session_util.decode_model( - getattr(event_metadata, 'usage_metadata', None), - types.GenerateContentResponseUsageMetadata, - ) cache_metadata = _session_util.decode_model( getattr(event_metadata, 'cache_metadata', None), CacheMetadata, @@ -508,7 +500,6 @@ def _from_api_event(api_event_obj: vertexai.types.SessionEvent) -> Event: compaction_data = None usage_metadata_data = None grounding_metadata = None - usage_metadata = None cache_metadata = None if actions: @@ -558,7 +549,6 @@ def _from_api_event(api_event_obj: vertexai.types.SessionEvent) -> Event: branch=branch, custom_metadata=custom_metadata, grounding_metadata=grounding_metadata, - usage_metadata=usage_metadata, cache_metadata=cache_metadata, long_running_tool_ids=long_running_tool_ids, usage_metadata=usage_metadata,