From 9deeaa7ee7de4c27a97db7135a1bfdaf08dfd7bd Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 20 Oct 2026 01:29:40 -0700 Subject: [PATCH] fix(vertex_ai): keep the remembered cache on request-shape 400s, and stop a waiter at its own timeout Vertex and Gemini answer a request that sends system_instruction, tools or tool_config beside a cachedContent with a 400 that names the cache. That is not a missing cache, so it no longer triggers a rebuild or drops the remembered cache. A request that joins another request's lookup-or-create now waits at most its own timeout (the read limit of an httpx.Timeout) and then fails with a 408, without cancelling the first request. --- .../context_caching/known_cached_contents.py | 18 ++++++-- .../vertex_ai_context_caching.py | 16 ++++++- .../test_known_cached_contents.py | 42 +++++++++++++++++-- 3 files changed, 67 insertions(+), 9 deletions(-) diff --git a/litellm/llms/vertex_ai/context_caching/known_cached_contents.py b/litellm/llms/vertex_ai/context_caching/known_cached_contents.py index eefff5d5beb..a4292495014 100644 --- a/litellm/llms/vertex_ai/context_caching/known_cached_contents.py +++ b/litellm/llms/vertex_ai/context_caching/known_cached_contents.py @@ -11,12 +11,17 @@ KNOWN_CACHE_EXPIRY_MARGIN_SECONDS: Final = 10.0 MAX_KNOWN_CACHED_CONTENTS: Final = 10_000 _CACHE_REJECTED_STATUS_CODES: Final = frozenset({400, 403, 404}) _NAMES_CACHED_CONTENT: Final = re.compile(r"cached?\s*content", re.IGNORECASE) +_FORBIDS_FIELDS_BESIDE_A_CACHE: Final = re.compile(r"system[_\s]instruction", re.IGNORECASE) _RFC3339_FRACTION: Final = re.compile(r"\.(\d{1,6})\d*") def names_rejected_cache(status_code: int, message: str) -> bool: - return status_code in _CACHE_REJECTED_STATUS_CODES and _NAMES_CACHED_CONTENT.search(message) is not None + return ( + status_code in _CACHE_REJECTED_STATUS_CODES + and _NAMES_CACHED_CONTENT.search(message) is not None + and _FORBIDS_FIELDS_BESIDE_A_CACHE.search(message) is None + ) def parse_expire_time(expire_time: object) -> float | None: @@ -103,13 +108,18 @@ class KnownCachedContents: self._entries = {key: entry for key, entry in self._entries.items() if entry[0] != cached_content} async def single_flight( - self, scope: tuple[str, str], cache_key: str, resolve: Callable[[], Awaitable[str | None]] + self, + scope: tuple[str, str], + cache_key: str, + resolve: Callable[[], Awaitable[str | None]], + wait_seconds: float | None, ) -> str | None: """Run ``resolve`` once per (scope, cache_key) at a time within this event loop. Concurrent requests for the same prefix wait for the first one's lookup-or-create instead of each listing (and creating a duplicate cache). If that first attempt - fails, each waiter resolves on its own. + fails, each waiter resolves on its own. A waiter waits at most ``wait_seconds`` and + then raises ``asyncio.TimeoutError`` without cancelling the first request. """ loop: Final = asyncio.get_running_loop() own: Final[asyncio.Future[str | None]] = loop.create_future() @@ -117,7 +127,7 @@ class KnownCachedContents: inflight: Final = self._inflight.setdefault(loop, {}) leader: Final = inflight.setdefault((scope, cache_key), own) if leader is not own: - name: Final = await asyncio.shield(leader) + name: Final = await asyncio.wait_for(asyncio.shield(leader), wait_seconds) if name: return name return await resolve() diff --git a/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py b/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py index f3343eb5597..fa934568f18 100644 --- a/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py +++ b/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py @@ -1,3 +1,4 @@ +import asyncio from collections.abc import Sequence from typing import Final, Literal @@ -70,6 +71,14 @@ def _record_cache_creation_tokens(logging_obj: Logging, created: _CreatedCachedC logging_obj.vertex_ai_cache_creation_input_tokens = max(tokens, 0) # rebind-ok: set provider count +def _wait_seconds(timeout: float | httpx.Timeout | None) -> float | None: + match timeout: + case httpx.Timeout(): + return timeout.read + case _: + return timeout + + class ContextCachingEndpoints(VertexBase): """ Covers context caching endpoints for Vertex AI + Google AI Studio @@ -643,7 +652,12 @@ class ContextCachingEndpoints(VertexBase): ) return raw_response_cached.get("name") - cache_name: Final = await self._known_caches.single_flight(scope, generated_cache_key, lookup_or_create) + try: + cache_name: Final = await self._known_caches.single_flight( + scope, generated_cache_key, lookup_or_create, _wait_seconds(timeout) + ) + except asyncio.TimeoutError: + raise VertexAIError(status_code=408, message="Timeout error occurred.") return non_cached_messages, optional_params, cache_name def get_cache(self): diff --git a/tests/unit/llms/vertex_ai/context_caching/test_known_cached_contents.py b/tests/unit/llms/vertex_ai/context_caching/test_known_cached_contents.py index 14216c443b5..37a8a47edc7 100644 --- a/tests/unit/llms/vertex_ai/context_caching/test_known_cached_contents.py +++ b/tests/unit/llms/vertex_ai/context_caching/test_known_cached_contents.py @@ -84,7 +84,9 @@ def _sync_call( ) -async def _async_call(endpoints: ContextCachingEndpoints, client: AsyncHTTPHandler): +async def _async_call( + endpoints: ContextCachingEndpoints, client: AsyncHTTPHandler, timeout: float | httpx.Timeout = 30.0 +): return await endpoints.async_check_and_create_cache( messages=_messages(), optional_params={}, @@ -92,7 +94,7 @@ async def _async_call(endpoints: ContextCachingEndpoints, client: AsyncHTTPHandl api_base=None, model="gemini-3.1-pro", client=client, - timeout=30.0, + timeout=timeout, logging_obj=MagicMock(model_call_details={}), custom_llm_provider="vertex_ai", vertex_project="p", @@ -303,6 +305,29 @@ async def test_waiters_resolve_on_their_own_when_the_first_request_fails(endpoin assert results[1][2] == "cachedContents/2" +@pytest.mark.parametrize("waiter_timeout", [0.05, httpx.Timeout(30.0, read=0.05)]) +@pytest.mark.asyncio +async def test_waiter_stops_waiting_for_the_first_request_at_its_own_timeout( + endpoints: ContextCachingEndpoints, waiter_timeout: float | httpx.Timeout +): + client: Final = MagicMock(spec=AsyncHTTPHandler) + + async def slow_list(**_): + await asyncio.sleep(0.5) + return _response({}) + + client.get = AsyncMock(side_effect=slow_list) + client.post = AsyncMock(return_value=_response(_created("cachedContents/1", _expire_time(60)))) + + results: Final = await asyncio.gather( + _async_call(endpoints, client), _async_call(endpoints, client, waiter_timeout), return_exceptions=True + ) + + assert results[0][2] == "cachedContents/1" + assert isinstance(results[1], VertexAIError) + assert results[1].status_code == 408 + + def test_parse_expire_time_accepts_nanosecond_precision(): assert parse_expire_time("2026-10-14T18:06:21.123456789Z") == pytest.approx( datetime(2026, 10, 14, 18, 6, 21, 123456, tzinfo=timezone.utc).timestamp() @@ -417,13 +442,22 @@ async def test_remembered_cache_rejected_by_generate_content_is_replaced_within_ @pytest.mark.parametrize("sync_mode", [True, False]) @pytest.mark.parametrize("stream", [True, False]) +@pytest.mark.parametrize( + "message", + [ + "Invalid value at 'generation_config.temperature'", + "Tool config, tools and system instruction should not be set in therequest when using cached content.", + "CachedContent can not be used with GenerateContent request setting system_instruction, tools or tool_config.", + ], +) @pytest.mark.asyncio -async def test_error_that_does_not_name_the_cache_keeps_the_remembered_cache( +async def test_error_that_is_not_about_a_missing_cache_keeps_the_remembered_cache( respx_mock: respx.MockRouter, fake_provider_credentials: None, monkeypatch: pytest.MonkeyPatch, sync_mode: bool, stream: bool, + message: str, ) -> None: monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) list_route: Final = respx_mock.get(url__regex=r".*/cachedContents.*").mock( @@ -435,7 +469,7 @@ async def test_error_that_does_not_name_the_cache_keeps_the_remembered_cache( generate_route: Final = respx_mock.post(url__regex=r".*:(stream)?[gG]enerateContent.*").mock( side_effect=[ _generated(stream), - _rejected(400, "Invalid value at 'generation_config.temperature'"), + _rejected(400, message), _generated(stream), ] )