Skip to content

feat(models): add client_max_age and invalidate_clients() to Gemini - #7797

Open
inamdarmihir wants to merge 1 commit into
google:mainfrom
inamdarmihir:fix/7697-gemini-client-refresh
Open

inamdarmihir wants to merge 1 commit into
google:mainfrom
inamdarmihir:fix/7697-gemini-client-refresh

Conversation

@inamdarmihir

Copy link
Copy Markdown
Contributor

Link to Issue or Description of Change

1. Link to an existing issue (if applicable):

Problem:
Gemini caches its genai.Client per 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_client that is a silent no-op, and Gemini has 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 of api_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 client passed to the constructor is never replaced. A subclass that overrides api_client with its own caching still owns that client's lifetime (documented). _PerLoopCache gained max-age support and a generation counter so a build running across clear() is not cached; its hit path is a single dict lookup so a concurrent clear() 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:

  • I have added or updated unit tests for my change.
  • All unit tests pass locally.

New tests/unittests/models/test_google_llm_client_refresh.py and additions to tests/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(); invalid client_max_age values; provided client untouched; subclasses overriding api_client.

pytest tests/unittests/utils/test_event_loop_cache.py tests/unittests/models
1855 passed

Manual End-to-End (E2E) Tests:

Throwaway script against a real Gemini with 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.05 reuses 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

  • I have read the CONTRIBUTING.md document.
  • I have performed a self-review of my own code.
  • I have commented my code, particularly in hard-to-understand areas.
  • I have added tests that prove my fix is effective or that my feature works.
  • New and existing unit tests pass locally with my changes.
  • I have manually tested my changes end-to-end.
  • Any dependent changes have been merged and published in downstream modules.

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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Gemini model calls return 401 about 24h after start on Agent Runtime (Agent Identity)

2 participants