Merge pull request #41366 from BerriAI/litellm_/back-002-litellm-e2e-replay-8de5b1

fix(e2e): record cookie-setting provider responses and keep prompt-caching tests live

tests/code_coverage_tests/test_provider_cache.py

@@ -38,6 +38,7 @@ class Provider(ThreadingHTTPServer):
     delay: float = 0
     stream: bool = False
     truncated: bool = False
+    cookie: str = ""
 
 
 class Handler(BaseHTTPRequestHandler):
@@ -62,6 +63,8 @@ class Handler(BaseHTTPRequestHandler):
             return
         self.send_header("content-type", "application/json")
         self.send_header("content-length", str(len(server.response)))
+        if server.cookie:
+            self.send_header("set-cookie", server.cookie)
         self.end_headers()
         self.wfile.write(server.response)
 
@@ -172,6 +175,14 @@ def test_failed_provider_responses_never_enter_cache(store: RedisResponseStore,
     assert len(provider.hits) == 2
 
 
+def test_cookie_setting_success_is_reused_without_the_cookie(store: RedisResponseStore, provider: Provider) -> None:
+    provider.cookie = "__cf_bm=synthetic-bot-management; Path=/; HttpOnly; Secure"
+    with edge(CacheEdge(store, SECRET), provider) as url:
+        replies: Final = tuple(call(url) for _ in range(2))
+    assert len(provider.hits) == 1
+    assert all(reply.body == SUCCESS and "set-cookie" not in reply.headers for reply in replies)
+
+
 def test_expiry_does_not_slide(store: RedisResponseStore, provider: Provider) -> None:
     short: Final = replace(store, lifetime_ms=250)
     with edge(CacheEdge(short, SECRET), provider) as url:

tests/e2e/PROVIDER_CACHE.md

@@ -2,7 +2,7 @@
 
 `E2E_PROVIDER_CACHE=1` enables automatic response reuse in the live E2E mode. Standard OpenAI and Anthropic model registrations use the provider edge. Existing custom API bases, named credentials, mocked models and realtime WebSocket deployments keep their existing routing. Other provider protocols remain live
 
-The edge caches complete successful POST responses for `/v1/chat/completions` and `/v1/messages`, including streams. Unsupported endpoints pass through. It matches the method, original URL, effective outbound headers (including authentication and HTTP-library defaults), body presence and exact body bytes using a full keyed digest. It sends the same prepared request used for matching. No prompts, random markers, JSON values or credentials are normalized away
+The edge caches complete successful POST responses for `/v1/chat/completions` and `/v1/messages`, including streams. Unsupported endpoints pass through. It matches the method, original URL, effective outbound headers (including authentication and HTTP-library defaults), body presence and exact body bytes using a full keyed digest. It sends the same prepared request used for matching. No prompts, random markers, JSON values or credentials are normalized away. Provider `Set-Cookie` headers are dropped before validation and never recorded: the edge already withholds them from the proxy, and OpenAI responses always carry Cloudflare bot-management cookies
 
 An eligible miss calls the provider. A complete successful response is stored immediately even if a later test assertion fails. Provider errors, malformed responses, truncated streams and cancelled captures are not stored. Cache reads, writes and lease failures fall through to normal provider behavior; they introduce no provider retry. An already-started response cannot be restarted after a delivery failure
 
@@ -20,7 +20,7 @@ The trusted runner receives:
 
 Do not give cache credentials to candidate deployments. Counter artifacts contain no recorded payloads or credentials. Hits count shared-cache responses; upstream attempts count actual forwards from the edge. Existing application-cache observations still count requests arriving at the edge, including shared-cache hits
 
-Tests that require real provider timing, limits or state use `@pytest.mark.provider_live`. The marker keeps newly registered models on live routes without weakening their assertions. Ordinary assertion failures still fail E2E. The shared cache does not modify provider response IDs or make the proxy aware of replay
+Tests that require real provider timing, limits or state use `@pytest.mark.provider_live`. The marker keeps newly registered models on live routes without weakening their assertions. The provider prompt-caching tests carry it because a replayed priming response reports cache creation rather than a cache read. Ordinary assertion failures still fail E2E. The shared cache does not modify provider response IDs or make the proxy aware of replay
 
 ## Recorded response semantics
 

tests/e2e/llm_translation/test_cache_control.py

@@ -44,7 +44,7 @@ from models import ChatBody, ChatMessage, ChatResponse, LiteLLMParamsBody, Usage
 from passthrough_client import PassthroughClient
 import os
 
-pytestmark = pytest.mark.e2e
+pytestmark = [pytest.mark.e2e, pytest.mark.provider_live]
 
 BEDROCK_MODEL = "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0"
 VERTEX_MODEL = "vertex_ai/gemini-2.5-flash"

tests/e2e/provider_cache.py

@@ -28,6 +28,7 @@ from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter, ValidationEr
 LIFETIME_SECONDS: Final = 86_400
 MAX_REQUEST_BYTES: Final = 256 * 1024
 MAX_RESPONSE_BYTES: Final = 8 * 1024 * 1024
+UNRECORDED_RESPONSE_HEADERS: Final = frozenset({"set-cookie"})
 JSON_VALUE: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue)
 
 
@@ -104,8 +105,6 @@ def cacheable_endpoint(method: str, url: str, body: bytes | None) -> bool:
 def successful_response(url: str, status: int, headers: Mapping[str, str], body: bytes) -> bool:
     if not 200 <= status < 300 or len(body) > MAX_RESPONSE_BYTES:
         return False
-    if any(name.lower() == "set-cookie" for name in headers):
-        return False
     streaming: Final = "text/event-stream" in headers.get("content-type", "").lower()
     if streaming:
         try:
@@ -283,11 +282,14 @@ class CacheEdge:
                     yield step
                     capture.observe(step)
             chunks: Final = capture.chunks() if capture.eligible else ()
-            if not capture.eligible or not successful_response(url, head.status_code, head.headers, b"".join(chunks)):
+            headers: Final = {
+                name: value for name, value in head.headers.items() if name.lower() not in UNRECORDED_RESPONSE_HEADERS
+            }
+            if not capture.eligible or not successful_response(url, head.status_code, headers, b"".join(chunks)):
                 self.counters.increment("rejected")
                 return
             response: Final = CachedResponse(
-                request_key=key, status_code=head.status_code, headers=head.headers,
+                request_key=key, status_code=head.status_code, headers=headers,
                 chunks=tuple(base64.b64encode(chunk).decode("ascii") for chunk in chunks),
             )
             published: Final = self.store.publish(key, lease, encode_response(self.secret, response))

tests/e2e/quota_management/ratelimit/test_tpm_excludes_cached_tokens_e2e.py

@@ -26,7 +26,7 @@ from models import (
 )
 from quota_client import QuotaClient
 
-pytestmark = pytest.mark.e2e
+pytestmark = [pytest.mark.e2e, pytest.mark.provider_live]
 
 # Anthropic prompt caching (host has ANTHROPIC_API_KEY; Bedrock was "Operation not allowed").
 ANTHROPIC_MODEL = "anthropic/claude-haiku-4-5-20251001"

tests/e2e/test_provider_edge.py

@@ -1254,7 +1254,8 @@ class TestHandleEdgeRequestPure:
 
 
 class TestApiBaseSeam:
-    def test_live_mode_returns_none(self, tmp_path: Path) -> None:
+    def test_live_mode_returns_none(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
+        monkeypatch.delenv("E2E_PROVIDER_CACHE", raising=False)
         for mode_raw in ("live", ""):
             assert (
                 provider_edge_api_base(