Repository navigation
feat(models): add client_max_age and invalidate_clients() to Gemini - #7797
Open
inamdarmihir wants to merge 1 commit into
Open
inamdarmihir wants to merge 1 commit into
inamdarmihir wants to merge 1 commit into
Conversation
Since the cached genai clients moved to a per-event-loop cache, evicting api_client from the instance __dict__ no longer rebuilds the client, so a long-running process that built its client from short-lived credentials has no supported way to refresh it and starts failing with 401s. Add an optional Gemini.client_max_age (seconds, default None keeps today's behavior) that rebuilds the api and live clients once they reach that age, and Gemini.invalidate_clients() to drop them on demand. Old clients are dropped, never closed, so in-flight requests keep working. The per-loop cache gains max-age support and a clear operation to back both. Fixes google#7697
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Link to Issue or Description of Change
1. Link to an existing issue (if applicable):
Problem:
Geminicaches itsgenai.Clientper event loop (PerLoopCachedProperty) for as long as the loop lives. On Agent Runtime with Agent Identity the access tokens are bound to a rotating client certificate, so the cached client starts returning 401 about 24h after start and keeps failing until the process restarts. Before 2.10 a subclass could evict__dict__["api_client"]; since the cache moved to_per_loop_api_clientthat is a silent no-op, andGeminihas no supported way to bound the client's age or rebuild it.Solution:
Two optional, additive controls on
Gemini; defaults are unchanged.client_max_age: Optional[float] = None(must be > 0, not serialized): once a cached client has existed for this many seconds, the next read ofapi_client/ the live client builds a new one. Age is counted from when the build started, so credentials captured during a slow build still count as old.invalidate_clients(): drops the cached api and live clients for every event loop so the next use rebuilds them. Safe from any thread.Old clients are dropped, never closed, so an in-flight request or live session keeps working on its client. A
clientpassed to the constructor is never replaced. A subclass that overridesapi_clientwith its own caching still owns that client's lifetime (documented)._PerLoopCachegained max-age support and a generation counter so a build running acrossclear()is not cached; its hit path is a single dict lookup so a concurrentclear()cannot raise.I deliberately did not add a rebuild-and-retry on 401: the error can occur mid-stream and a retry could replay non-idempotent calls. Callers can combine
invalidate_clients()with their own error handling.Testing Plan
Unit Tests:
New
tests/unittests/models/test_google_llm_client_refresh.pyand additions totests/unittests/utils/test_event_loop_cache.py, covering: default never expires; reuse just before and rebuild exactly at the max age; invalidate before first use, after use and across loops; api and live clients; per-loop isolation; an in-flight coroutine keeps its old client; invalidation during a build; concurrent readers vs.clear(); invalidclient_max_agevalues; providedclientuntouched; subclasses overridingapi_client.Manual End-to-End (E2E) Tests:
Throwaway script against a real
Geminiwith a stubbed client factory: the old__dict__.pop("api_client")workaround still does not refresh;invalidate_clients()yields a new client;client_max_age=0.05reuses the client before expiry and rebuilds api and live clients after it; the default never refreshes. I could not reproduce the 24h Agent Identity 401 itself (needs Agent Runtime).Checklist