fix: preserve cache_metadata in VertexAiSessionService event round-trip#4714
fix: preserve cache_metadata in VertexAiSessionService event round-trip#4714OiPunk wants to merge 1 commit intogoogle:mainfrom
Conversation
…vice 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 google#4698
Summary of ChangesHello, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request fixes a critical bug in the Highlights
🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Changelog
Activity
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
adk-bot
added
the
services
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request correctly addresses the issue where cache_metadata and usage_metadata were not being preserved during the event round-trip in VertexAiSessionService. The changes to serialize and deserialize these fields are implemented correctly, following the existing pattern for other metadata. The addition of a new unit test, test_append_event_with_cache_and_usage_metadata, is great for verifying the fix. I have one suggestion to make the new test more concise and maintainable.
| # 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 |
There was a problem hiding this comment.
To improve the readability and maintainability of this test, you can directly compare the cache_metadata and usage_metadata objects with the originals (cache_meta and usage_meta) instead of asserting each field individually. Pydantic models, which are used here, support equality comparison out of the box. This makes the test more concise and robust against future changes to these models.
# cache_metadata is preserved
assert appended_event.cache_metadata == cache_meta
# usage_metadata is preserved
assert appended_event.usage_metadata == usage_meta2 participants
Summary
VertexAiSessionServicedropscache_metadataandusage_metadatafields when serializing/deserializingEventobjects. This causesContextCacheRequestProcessorto never find previous cache metadata from session events, so it creates a new context cache on every LLM call instead of reusing existing ones.Root cause
append_event):cache_metadataandusage_metadatawere not included in theevent_metadatadict sent to the Vertex AI API._from_api_event):cache_metadataandusage_metadatawere not extracted from the API response when reconstructing theEventobject.Other session service implementations (
InMemorySessionService, database-backedStorageEvent) preserve these fields correctly.Fix
cache_metadataandusage_metadataintoevent_metadataduringappend_event._from_api_eventusing_session_util.decode_model, consistent with howgrounding_metadatais already handled.Test plan
test_append_event_with_cache_and_usage_metadatathat verifies both fields survive a full round-trip (append -> get_session).Fixes #4698