From 2d40254b57a62be9bce95586cb902b90c8505c4c Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 16 Sep 2026 02:15:46 -0700 Subject: [PATCH 01/21] feat(e2e): key the provider cache per test and mount Bedrock behind it The exact-request cache reused 5% of routed traffic (build 218: 19 hits, 350 misses) because every test salts its prompt with a fresh unique_marker(), so the same test could never match itself across builds. It also routed only openai and anthropic, while the week's flakiness was Bedrock. Key is now HMAC(test id + method + URL + headers + body, with every unique_marker() token replaced by a placeholder, + FIFO slot index). The slot index is what keeps two marker-only-different calls in one test on two recordings and therefore two provider response ids, so spend rows still reconcile one per invocation. A call outside any test is not cacheable. Bedrock gets a region-qualified mount and SigV4 re-signing, since the edge rewrites the Host the proxy signed. Signature headers are excluded from the key for signing mounts only, because x-amz-date would otherwise make every Bedrock request a permanent miss; every other mount still keys on its credentials whole. Only Anthropic-on-Bedrock chat deployments route: embeddings, image generation, rerank and realtime keep their direct path, and so do deployments carrying their own aws_role_name or static keys, whose whole point is to prove the product's assume-role chain rather than the runner's. The two eventstream actions bypass the cache and go live, still signed. Counters are now attributed per mount as well as in total, so a build can report a per-provider hit rate instead of one number. --- .../test_provider_cache.py | 509 +++++++++++++++--- tests/e2e/fixture_canonical.py | 5 +- tests/e2e/models.py | 1 + tests/e2e/provider_cache.py | 171 ++++-- tests/e2e/provider_cache_routing.py | 47 +- tests/e2e/provider_edge.py | 62 ++- tests/e2e/provider_edge_bedrock.py | 72 +++ tests/e2e/test_provider_edge.py | 19 +- 8 files changed, 768 insertions(+), 118 deletions(-) create mode 100644 tests/e2e/provider_edge_bedrock.py diff --git a/tests/code_coverage_tests/test_provider_cache.py b/tests/code_coverage_tests/test_provider_cache.py index 828227ed239..5d35981344d 100644 --- a/tests/code_coverage_tests/test_provider_cache.py +++ b/tests/code_coverage_tests/test_provider_cache.py @@ -7,7 +7,7 @@ import subprocess import threading import time import uuid -from collections.abc import Generator +from collections.abc import Generator, Mapping from concurrent.futures import ThreadPoolExecutor from contextlib import contextmanager from dataclasses import dataclass, replace @@ -18,21 +18,58 @@ from urllib.parse import urlsplit import pytest from e2e_http import NetworkError, PreparedForward, RawResponse, StreamChunk, StreamHead, forward, prepare_forward -from models import LiteLLMParamsBody -from provider_cache import CacheEdge, CacheHit, CaptureLease, exact_key, successful_response +from models import LiteLLMParamsBody, ModelMode +from botocore.credentials import Credentials +from provider_cache import ( + SIGNATURE_HEADERS, + CacheEdge, + CacheHit, + CaptureLease, + ResponseStore, + cacheable_endpoint, + request_identity, + slotted_key, + successful_response, +) from provider_cache_redis import PUBLISH, RedisCommands, RedisResponseStore, configured_cache, redis_store from provider_cache_routing import LIVE_PROVIDER_REQUIRED, route_cache_model -from provider_edge import configured_cache_backend, start_provider_edge +from fixture_mode import SESSION_TEST_KEY +from provider_edge import EDGE_MOUNTS, configured_cache_backend, resolve_mount, start_provider_edge +from provider_edge_bedrock import bedrock_signer from redis.exceptions import ConnectionError as RedisConnectionError SECRET: Final = b"synthetic-cache-hmac-key-for-tests" BODY: Final = b'{"model":"test","messages":[{"role":"user","content":"hello"}]}' SUCCESS: Final = b'{"id":"provider-fixed-id","choices":[{"message":{"content":"hello"},"finish_reason":"stop"}],"usage":{"prompt_tokens":1,"completion_tokens":1,"total_tokens":2}}' HEADERS: Final = {"content-type": "application/json", "authorization": "Bearer synthetic-account-one"} +TEST_KEY: Final = "tests/e2e/synthetic_suite.py::TestCase::test_case" +OTHER_TEST_KEY: Final = "tests/e2e/synthetic_suite.py::TestCase::test_other_case" + + +def marked(marker: str) -> bytes: + """One request body shaped like the suite's own: a fixed prompt salted with a + 12-lowercase-hex ``unique_marker()`` token, fresh on every run.""" + return b'{"model":"test","messages":[{"role":"user","content":"hello %s"}]}' % marker.encode() + + +MARKED: Final = marked("0a1b2c3d4e5f") +BEDROCK_MOUNT: Final = "bedrock/us-east-1" +BEDROCK_MODEL: Final = "us.anthropic.claude-haiku-4-5-20251001-v1%3A0" +BEDROCK_BODY: Final = b'{"messages":[{"role":"user","content":[{"text":"hello 0a1b2c3d4e5f"}]}]}' +CONVERSE_SUCCESS: Final = ( + b'{"output":{"message":{"role":"assistant","content":[{"text":"hi"}]}},' + b'"stopReason":"end_turn","usage":{"inputTokens":1,"outputTokens":1,"totalTokens":2}}' +) +INVOKE_SUCCESS: Final = ( + b'{"id":"msg_synthetic","type":"message","role":"assistant",' + b'"content":[{"type":"text","text":"hi"}],"stop_reason":"end_turn"}' +) +STATIC_CREDENTIALS: Final = Credentials("AKIAIOSFODNN7EXAMPLE", "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY") class Provider(ThreadingHTTPServer): hits: tuple[tuple[str, bytes], ...] = () + authorizations: tuple[str, ...] = () response: bytes = SUCCESS status: int = 200 delay: float = 0 @@ -49,6 +86,7 @@ class Handler(BaseHTTPRequestHandler): assert isinstance(server, Provider) body: Final = self.rfile.read(int(self.headers.get("content-length", "0"))) server.hits += ((self.path, body),) + server.authorizations += (self.headers.get("authorization", ""),) time.sleep(server.delay) self.send_response(server.status) if server.stream: @@ -122,6 +160,29 @@ def store(redis_url: str) -> RedisResponseStore: return redis_store(redis_url, "test-" + uuid.uuid4().hex) +def cache_edge(store: ResponseStore, test_key: str = TEST_KEY) -> CacheEdge: + """A cache edge standing in for one pytest process. A fresh instance over the + same store is the next build running the same test: the recordings survive, + the per-test FIFO slot counters start over.""" + return CacheEdge(store, SECRET, test_key=lambda: test_key) + + +def slot_key( + url: str, slot: int = 0, body: bytes | None = BODY, + headers: dict[str, str] = HEADERS, test_key: str = TEST_KEY, +) -> str: + prepared: Final = prepare_forward("POST", url, headers, body) + assert isinstance(prepared, PreparedForward) + return slotted_key(SECRET, request_identity(SECRET, test_key, "POST", url, prepared.headers, body), slot) + + +def bedrock_cache_edge(store: ResponseStore, test_key: str = TEST_KEY) -> CacheEdge: + return CacheEdge( + store, SECRET, test_key=lambda: test_key, + signers={BEDROCK_MOUNT: bedrock_signer("us-east-1", lambda: STATIC_CREDENTIALS)}, + ) + + @contextmanager def edge(cache: CacheEdge, provider: Provider) -> Generator[str, None, None]: upstream: Final = f"http://127.0.0.1:{provider.server_port}" @@ -132,34 +193,54 @@ def edge(cache: CacheEdge, provider: Provider) -> Generator[str, None, None]: running.shutdown() +@contextmanager +def bedrock_edge(cache: CacheEdge, provider: Provider, action: str = "converse") -> Generator[str, None, None]: + upstream: Final = f"http://127.0.0.1:{provider.server_port}" + running: Final = start_provider_edge(cache, mounts={BEDROCK_MOUNT: upstream}) + try: + yield f"{running.edge.api_base(BEDROCK_MOUNT)}/model/{BEDROCK_MODEL}/{action}" + finally: + running.shutdown() + + def call(url: str, body: bytes = BODY, headers: dict[str, str] = HEADERS) -> RawResponse: result: Final = forward("POST", url, headers=headers, body=body, timeout=5) assert isinstance(result, RawResponse), result return result -def test_success_is_reusable_across_fresh_edges(store: RedisResponseStore, provider: Provider) -> None: - with edge(CacheEdge(store, SECRET), provider) as url: +def test_repeated_call_takes_its_own_slot_and_both_replay_next_run( + store: RedisResponseStore, provider: Provider, +) -> None: + with edge(cache_edge(store), provider) as url: assert call(url).body == SUCCESS assert call(url).body == SUCCESS - with edge(CacheEdge(store, SECRET), provider) as other: + assert len(provider.hits) == 2 + with edge(cache_edge(store), provider) as other: assert call(other).body == SUCCESS - assert len(provider.hits) == 1 + assert call(other).body == SUCCESS + assert len(provider.hits) == 2 @pytest.mark.parametrize("body", [BODY + b" ", BODY.replace(b"hello", b"Hello"), BODY.replace(b"test", b"test2")]) def test_any_body_change_calls_live(store: RedisResponseStore, provider: Provider, body: bytes) -> None: - with edge(CacheEdge(store, SECRET), provider) as url: + with edge(cache_edge(store), provider) as url: call(url) + assert len(provider.hits) == 1 + with edge(cache_edge(store), provider) as url: call(url, body) + assert len(provider.hits) == 2 + with edge(cache_edge(store), provider) as url: call(url, body) assert len(provider.hits) == 2 @pytest.mark.parametrize("name,value", [("authorization", "Bearer another-account"), ("x-request-id", "one"), ("anthropic-version", "new")]) def test_changed_header_cannot_reuse(store: RedisResponseStore, provider: Provider, name: str, value: str) -> None: - with edge(CacheEdge(store, SECRET), provider) as url: + with edge(cache_edge(store), provider) as url: call(url) + assert len(provider.hits) == 1 + with edge(cache_edge(store), provider) as url: call(url, headers=HEADERS | {name: value}) call(url + "?x=1") assert len(provider.hits) == 3 @@ -169,36 +250,59 @@ def test_changed_header_cannot_reuse(store: RedisResponseStore, provider: Provid def test_failed_provider_responses_never_enter_cache(store: RedisResponseStore, provider: Provider, status: int, response: bytes) -> None: provider.status = status provider.response = response - with edge(CacheEdge(store, SECRET), provider) as url: + with edge(cache_edge(store), provider) as url: assert call(url).status_code == status + assert len(provider.hits) == 1 + with edge(cache_edge(store), provider) as url: assert call(url).body == response 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)) + with edge(cache_edge(store), provider) as url: + live: Final = call(url) + with edge(cache_edge(store), provider) as url: + replayed: Final = call(url) assert len(provider.hits) == 1 - assert all(reply.body == SUCCESS and "set-cookie" not in reply.headers for reply in replies) + assert all(reply.body == SUCCESS and "set-cookie" not in reply.headers for reply in (live, replayed)) 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: - call(url) - call(url) - time.sleep(0.3) - call(url) - call(url) + url: Final = f"http://127.0.0.1:{provider.server_port}/v1/chat/completions" + + def drain() -> None: + head = cache_edge(short).forward("openai", "POST", url, dict(HEADERS), BODY, 5) + assert isinstance(head, StreamHead) + assert b"".join(step.data for step in head.steps if isinstance(step, StreamChunk)) == SUCCESS + + drain() + assert len(provider.hits) == 1 + drain() + assert len(provider.hits) == 1 + time.sleep(0.3) + drain() assert len(provider.hits) == 2 -def test_concurrent_requests_publish_atomically(store: RedisResponseStore, provider: Provider) -> None: +def test_concurrent_builds_publish_one_recording_atomically( + store: RedisResponseStore, provider: Provider, +) -> None: + """Five processes running the same test at the same time all reach slot 0 of + one key, which is the only way the capture lease is contended now that a + repeat inside a single test takes its own slot.""" provider.delay = 0.15 - with edge(CacheEdge(store, SECRET), provider) as url: - with ThreadPoolExecutor(max_workers=5) as executor: - replies: Final = tuple(executor.map(lambda _: call(url).body, range(5))) + url: Final = f"http://127.0.0.1:{provider.server_port}/v1/chat/completions" + edges: Final = tuple(cache_edge(store) for _ in range(5)) + + def drain(cache: CacheEdge) -> bytes: + head = cache.forward("openai", "POST", url, dict(HEADERS), BODY, 5) + assert isinstance(head, StreamHead) + return b"".join(step.data for step in head.steps if isinstance(step, StreamChunk)) + + with ThreadPoolExecutor(max_workers=5) as executor: + replies: Final = tuple(executor.map(drain, edges)) assert replies == (SUCCESS,) * 5 assert len(provider.hits) == 1 @@ -231,9 +335,9 @@ def test_stream_completion_controls_publication(store: RedisResponseStore, provi provider.stream = True provider.truncated = truncated provider.response = b'data: {"choices":[{"index":0,"delta":{"content":"hello"},"finish_reason":"stop"}]}\n\ndata: [DONE]\n\n' - with edge(CacheEdge(store, SECRET), provider) as url: - for _ in range(2): - result: Final = forward("POST", url, headers=HEADERS, body=BODY, timeout=5) + for _ in range(2): + with edge(cache_edge(store), provider) as url: + result = forward("POST", url, headers=HEADERS, body=BODY, timeout=5) if truncated: assert isinstance(result, NetworkError) else: @@ -246,9 +350,9 @@ def test_store_outage_preserves_provider_success(provider: Provider) -> None: probe.bind(("127.0.0.1", 0)) port: Final = probe.getsockname()[1] unavailable: Final = redis_store(f"redis://127.0.0.1:{port}/0", "unavailable") - with edge(CacheEdge(unavailable, SECRET), provider) as url: - assert call(url).body == SUCCESS - assert call(url).body == SUCCESS + for _ in range(2): + with edge(cache_edge(unavailable), provider) as url: + assert call(url).body == SUCCESS assert len(provider.hits) == 2 @@ -267,7 +371,8 @@ def test_old_lease_cannot_overwrite_new_owner(store: RedisResponseStore) -> None def test_identity_preserves_values_and_never_contains_credentials() -> None: variants: Final = (b'{}', b'{"a":null}', b'{"a":false}', b'{"a":0}', b'{"a":0.0}', b'{"a":"0"}', b' { }', None, b'') - keys: Final = tuple(exact_key(SECRET, "POST", "https://example.invalid/v1/chat/completions", HEADERS, body) for body in variants) + url: Final = "https://example.invalid/v1/chat/completions" + keys: Final = tuple(request_identity(SECRET, TEST_KEY, "POST", url, HEADERS, body) for body in variants) assert len(set(keys)) == len(variants) assert all(len(key) == 64 and "synthetic-account" not in key for key in keys) @@ -275,21 +380,22 @@ def test_identity_preserves_values_and_never_contains_credentials() -> None: @pytest.mark.parametrize("payload", [b"corrupt response", '{"response":"{}","signature":"é"}'.encode()]) def test_corrupt_entry_is_replaced_by_same_successful_request(store: RedisResponseStore, provider: Provider, payload: bytes) -> None: upstream: Final = f"http://127.0.0.1:{provider.server_port}/v1/chat/completions" - prepared: Final = prepare_forward("POST", upstream, HEADERS, BODY) - assert isinstance(prepared, PreparedForward) - key: Final = exact_key(SECRET, "POST", upstream, prepared.headers, BODY) + key: Final = slot_key(upstream) lease: Final = store.lookup(key) assert isinstance(lease, CaptureLease) assert store.publish(key, lease, payload) - cache: Final = CacheEdge(store, SECRET) - for _ in range(2): - head = cache.forward("POST", upstream, HEADERS, BODY, 5) + caches: Final = tuple(cache_edge(store) for _ in range(2)) + for cache in caches: + head = cache.forward("openai", "POST", upstream, dict(HEADERS), BODY, 5) assert isinstance(head, StreamHead) assert b"".join(step.data for step in head.steps if isinstance(step, StreamChunk)) == SUCCESS assert len(provider.hits) == 1 - assert dict(cache.counters.counts) == { - "corrupt": 1, "misses": 1, "upstream_attempts": 1, "writes": 1, "hits": 1, + assert dict(caches[0].counters.counts) == { + "corrupt": 1, "mount:openai:corrupt": 1, "misses": 1, "mount:openai:misses": 1, + "upstream_attempts": 1, "mount:openai:upstream_attempts": 1, + "writes": 1, "mount:openai:writes": 1, } + assert dict(caches[1].counters.counts) == {"hits": 1, "mount:openai:hits": 1} @pytest.mark.parametrize("payload", [ @@ -301,22 +407,242 @@ def test_corrupt_entry_is_replaced_by_same_successful_request(store: RedisRespon def test_malformed_success_stream_is_never_cached(store: RedisResponseStore, provider: Provider, payload: bytes) -> None: provider.stream = True provider.response = payload - with edge(CacheEdge(store, SECRET), provider) as url: - assert call(url).body == payload - assert call(url).body == payload + for _ in range(2): + with edge(cache_edge(store), provider) as url: + assert call(url).body == payload assert len(provider.hits) == 2 +def test_requests_differing_only_by_marker_share_one_recording_per_slot( + store: RedisResponseStore, provider: Provider, +) -> None: + """The whole point of the canonical key. Every e2e test salts its prompt with + a fresh ``unique_marker()``, so before this the same test could never reuse + anything across builds. The second run mints markers it has never sent, which + is what a later build actually does, and must still serve both from the two + slots the first run recorded.""" + with edge(cache_edge(store), provider) as url: + assert call(url, MARKED).body == SUCCESS + assert call(url, marked("f5e4d3c2b1a0")).body == SUCCESS + assert len(provider.hits) == 2 + with edge(cache_edge(store), provider) as url: + assert call(url, marked("7c6b5a493827")).body == SUCCESS + assert call(url, marked("1122334455ff")).body == SUCCESS + assert len(provider.hits) == 2 + + +@pytest.mark.parametrize("body", [ + b'{"model":"test","messages":[{"role":"user","content":"hello 0a1b2c3d4e5"}]}', + b'{"model":"test","messages":[{"role":"user","content":"hello 0a1b2c3d4e5f0"}]}', + b'{"model":"test","messages":[{"role":"user","content":"hello 0A1B2C3D4E5F"}]}', + b'{"model":"0a1b2c3d4e5f","messages":[{"role":"user","content":"hello"}]}', +]) +def test_a_token_that_is_not_a_marker_keeps_its_own_key( + store: RedisResponseStore, provider: Provider, body: bytes, +) -> None: + """Too short, too long, upper case, or in another field: none of these is the + 12-lowercase-hex token ``unique_marker`` mints, so none may fold onto it.""" + with edge(cache_edge(store), provider) as url: + call(url, MARKED) + assert len(provider.hits) == 1 + with edge(cache_edge(store), provider) as url: + call(url, body) + assert len(provider.hits) == 2 + + +def test_another_test_never_reuses_this_tests_recording( + store: RedisResponseStore, provider: Provider, +) -> None: + with edge(cache_edge(store), provider) as url: + call(url) + assert len(provider.hits) == 1 + with edge(cache_edge(store, OTHER_TEST_KEY), provider) as url: + call(url) + assert len(provider.hits) == 2 + with edge(cache_edge(store, OTHER_TEST_KEY), provider) as url: + call(url) + assert len(provider.hits) == 2 + + +def test_calls_outside_any_test_are_never_cached( + store: RedisResponseStore, provider: Provider, +) -> None: + url: Final = f"http://127.0.0.1:{provider.server_port}/v1/chat/completions" + cache: Final = CacheEdge(store, SECRET, test_key=lambda: SESSION_TEST_KEY) + for _ in range(2): + head = cache.forward("openai", "POST", url, dict(HEADERS), BODY, 5) + assert isinstance(head, StreamHead) + assert b"".join(step.data for step in head.steps if isinstance(step, StreamChunk)) == SUCCESS + assert len(provider.hits) == 2 + assert dict(cache.counters.counts) == { + "bypass": 2, "mount:openai:bypass": 2, + "upstream_attempts": 2, "mount:openai:upstream_attempts": 2, + } + + +def test_counters_attribute_every_outcome_to_its_mount( + store: RedisResponseStore, provider: Provider, +) -> None: + """The build report needs per-provider hit counts, and the flat totals cannot + supply them. Anthropic is served a chat-shaped body here, which its validator + rejects, so one mount writes and the other does not.""" + upstream: Final = f"http://127.0.0.1:{provider.server_port}" + cache: Final = cache_edge(store) + running: Final = start_provider_edge(cache, mounts={"openai": upstream, "anthropic": upstream}) + try: + call(running.edge.api_base("openai") + "/v1/chat/completions") + call(running.edge.api_base("anthropic") + "/v1/messages") + finally: + running.shutdown() + counts: Final = dict(cache.counters.counts) + assert counts["misses"] == 2 + assert counts["mount:openai:misses"] == 1 and counts["mount:anthropic:misses"] == 1 + assert counts["mount:openai:writes"] == 1 and "mount:anthropic:writes" not in counts + assert counts["mount:anthropic:rejected"] == 1 and "mount:openai:rejected" not in counts + + +class TestBedrockSigning: + """Bedrock is the reason the edge could not mount it before: SigV4 covers the + Host header, so forwarding through a rewritten api_base invalidates the + proxy's signature. The edge mints its own over the upstream URL instead.""" + + def test_the_proxys_signature_is_replaced_not_forwarded(self) -> None: + signer: Final = bedrock_signer("us-east-1", lambda: STATIC_CREDENTIALS) + signed: Final = signer( + "POST", + f"https://bedrock-runtime.us-east-1.amazonaws.com/model/{BEDROCK_MODEL}/converse", + {"content-type": "application/json", "Authorization": "AWS4-HMAC-SHA256 Credential=PROXY/...", + "X-Amz-Date": "19700101T000000Z", "X-Amz-Security-Token": "proxy-session-token"}, + BEDROCK_BODY, + ) + assert "PROXY" not in str(signed) and "proxy-session-token" not in str(signed) + assert signed["Authorization"].startswith("AWS4-HMAC-SHA256 Credential=AKIAIOSFODNN7EXAMPLE/") + assert "/us-east-1/bedrock/aws4_request" in signed["Authorization"] + assert signed["X-Amz-Date"] != "19700101T000000Z" + assert signed["content-type"] == "application/json" + + def test_the_signed_url_reaches_the_wire_byte_for_byte(self) -> None: + """SigV4 hashes the canonical URI, so if the HTTP layer re-encoded the + colon in an inference-profile id after signing, every call would fail + with a signature mismatch rather than anything that names the cause.""" + url: Final = f"https://bedrock-runtime.us-east-1.amazonaws.com/model/{BEDROCK_MODEL}/converse" + signer: Final = bedrock_signer("us-east-1", lambda: STATIC_CREDENTIALS) + prepared: Final = prepare_forward("POST", url, signer("POST", url, dict(HEADERS), BEDROCK_BODY), BEDROCK_BODY) + assert isinstance(prepared, PreparedForward) + assert urlsplit(prepared.url).path == urlsplit(url).path + + def test_signature_headers_are_excluded_from_the_key( + self, store: RedisResponseStore, provider: Provider, + ) -> None: + """A real signature is fresh on every call, so keying on it would make + every Bedrock request a permanent miss. The stub signer here varies its + stamp per call on purpose: the real one only varies once a second, which + would let this pass by luck when it should fail.""" + provider.response = CONVERSE_SUCCESS + stamps: Final = iter(("20260101T000000Z", "20260102T111111Z")) + + def varying(method: str, url: str, headers: Mapping[str, str], body: bytes | None) -> dict[str, str]: + return dict(headers) | {"authorization": f"AWS4-HMAC-SHA256 {url}", "x-amz-date": next(stamps)} + + def signing_edge() -> CacheEdge: + return CacheEdge(store, SECRET, test_key=lambda: TEST_KEY, signers={BEDROCK_MOUNT: varying}) + + for _ in range(2): + with bedrock_edge(signing_edge(), provider) as url: + assert call(url, BEDROCK_BODY).body == CONVERSE_SUCCESS + assert len(provider.hits) == 1 + assert provider.authorizations[0] == ( + f"AWS4-HMAC-SHA256 http://127.0.0.1:{provider.server_port}/model/{BEDROCK_MODEL}/converse" + ), "the signature must cover the upstream URL the edge calls, not the edge URL the proxy called" + + def test_a_mount_without_a_signer_still_keys_on_its_credentials( + self, store: RedisResponseStore, provider: Provider, + ) -> None: + """The exclusion is per mount. Dropping authorization globally would let + one OpenAI account read another's recording.""" + cache: Final = bedrock_cache_edge(store) + assert "authorization" in SIGNATURE_HEADERS + assert "authorization" in cache.keyed("openai", HEADERS) + assert "authorization" not in cache.keyed(BEDROCK_MOUNT, HEADERS) + with edge(cache, provider) as url: + call(url) + with edge(bedrock_cache_edge(store), provider) as url: + call(url, headers=HEADERS | {"authorization": "Bearer synthetic-account-two"}) + assert len(provider.hits) == 2 + + @pytest.mark.parametrize("action,response", [("converse", CONVERSE_SUCCESS), ("invoke", INVOKE_SUCCESS)]) + def test_complete_responses_replay_on_the_next_run( + self, store: RedisResponseStore, provider: Provider, action: str, response: bytes, + ) -> None: + provider.response = response + for _ in range(2): + with bedrock_edge(bedrock_cache_edge(store), provider, action) as url: + assert call(url, BEDROCK_BODY).body == response + assert len(provider.hits) == 1 + + @pytest.mark.parametrize("action,response", [ + ("converse", b'{"output":{"message":{}}}'), + ("converse", b'{"stopReason":"end_turn"}'), + ("converse", b'{"message":"The provided model identifier is invalid."}'), + ("converse", CONVERSE_SUCCESS[:-20]), + ("invoke", b'{"id":"msg_x","type":"message","content":[{"type":"text","text":"hi"}]}'), + ("invoke", b'{"id":"msg_x","type":"message","stop_reason":"end_turn"}'), + ("invoke", b'{"message":"Too many requests, please wait before trying again."}'), + ]) + def test_incomplete_or_error_bodies_never_enter_the_cache( + self, store: RedisResponseStore, provider: Provider, action: str, response: bytes, + ) -> None: + provider.response = response + for _ in range(2): + with bedrock_edge(bedrock_cache_edge(store), provider, action) as url: + assert call(url, BEDROCK_BODY).body == response + assert len(provider.hits) == 2 + + @pytest.mark.parametrize("action", ["converse-stream", "invoke-with-response-stream"]) + def test_streaming_endpoints_go_live_every_time( + self, store: RedisResponseStore, provider: Provider, action: str, + ) -> None: + """An eventstream's completeness cannot be proven without parsing its + frames, so these bypass rather than risk recording a truncated answer. + They are still signed: a bypass is a forward, not a passthrough.""" + provider.response = CONVERSE_SUCCESS + cache: Final = bedrock_cache_edge(store) + for _ in range(2): + with bedrock_edge(cache, provider, action) as url: + assert call(url, BEDROCK_BODY).body == CONVERSE_SUCCESS + assert len(provider.hits) == 2 + assert dict(cache.counters.counts)[f"mount:{BEDROCK_MOUNT}:bypass"] == 2 + assert all( + sent.startswith("AWS4-HMAC-SHA256 Credential=AKIAIOSFODNN7EXAMPLE/") + for sent in provider.authorizations + ), provider.authorizations + + @pytest.mark.parametrize("action,cacheable", [ + ("converse", True), ("invoke", True), + ("converse-stream", False), ("invoke-with-response-stream", False), + ]) + def test_only_the_unary_bedrock_actions_are_cacheable(self, action: str, cacheable: bool) -> None: + url: Final = f"https://bedrock-runtime.us-east-1.amazonaws.com/model/{BEDROCK_MODEL}/{action}" + assert cacheable_endpoint(BEDROCK_MOUNT, "POST", url, BEDROCK_BODY) is cacheable + + def test_a_region_mount_resolves_whole(self) -> None: + resolved: Final = resolve_mount(f"/{BEDROCK_MOUNT}/model/{BEDROCK_MODEL}/converse", EDGE_MOUNTS) + assert resolved is not None + assert resolved.mount == BEDROCK_MOUNT + assert resolved.upstream_base == "https://bedrock-runtime.us-east-1.amazonaws.com" + assert resolved.upstream_path == f"model/{BEDROCK_MODEL}/converse" + + def test_anthropic_stream_requires_start_finish_and_stop() -> None: start: Final = b'data: {"type":"message_start","message":{}}\n\n' finish: Final = b'data: {"type":"message_delta","delta":{"stop_reason":"end_turn"}}\n\n' stop: Final = b'data: {"type":"message_stop"}\n\n' url: Final = "https://example.invalid/v1/messages" headers: Final = {"content-type": "text/event-stream"} - assert successful_response(url, 200, headers, start + finish + stop) - assert not successful_response(url, 200, headers, start + stop) - assert not successful_response(url, 200, headers, finish + stop) - assert not successful_response(url, 200, headers, start + finish) + assert successful_response("anthropic", url, 200, headers, start + finish + stop) + assert not successful_response("anthropic", url, 200, headers, start + stop) + assert not successful_response("anthropic", url, 200, headers, finish + stop) + assert not successful_response("anthropic", url, 200, headers, start + finish) @pytest.mark.parametrize("provider,suffix", [("openai", "/v1"), ("anthropic", "")]) @@ -342,6 +668,53 @@ def test_registration_preserves_unsupported_or_explicit_routes(params: LiteLLMPa assert route_cache_model(params, unexpected_edge, enabled=True) is params +@pytest.mark.parametrize("model", [ + "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", + "bedrock/converse/us.anthropic.claude-sonnet-5", + "bedrock/invoke/us.anthropic.claude-haiku-4-5-20251001-v1:0", +]) +def test_anthropic_on_bedrock_registers_the_edge_as_its_runtime_endpoint(model: str) -> None: + params: Final = LiteLLMParamsBody(model=model) + routed: Final = route_cache_model(params, lambda mount: f"http://edge.invalid/{mount}", enabled=True) + assert routed.aws_bedrock_runtime_endpoint == "http://edge.invalid/bedrock/us-east-1" + assert routed.api_base is None + assert routed.model_dump(exclude={"aws_bedrock_runtime_endpoint"}) == params.model_dump( + exclude={"aws_bedrock_runtime_endpoint"} + ) + + +@pytest.mark.parametrize("params", [ + LiteLLMParamsBody(model="bedrock/amazon.titan-embed-text-v2:0"), + LiteLLMParamsBody(model="bedrock/amazon.nova-canvas-v1:0"), + LiteLLMParamsBody(model="bedrock/amazon.nova-sonic-v1:0"), + LiteLLMParamsBody(model="bedrock/arn:aws:bedrock:us-east-1::foundation-model/cohere.rerank-v3-5:0"), + LiteLLMParamsBody(model="bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", aws_role_name="arn:aws:iam::1:role/x"), + LiteLLMParamsBody(model="bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", aws_access_key_id="AKIA"), + LiteLLMParamsBody(model="bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", api_base="https://custom.invalid"), + LiteLLMParamsBody( + model="bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", + aws_bedrock_runtime_endpoint="https://custom.invalid", + ), + LiteLLMParamsBody(model="bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", aws_region_name="eu-west-1"), +]) +def test_bedrock_deployments_the_edge_must_not_touch_keep_their_direct_route(params: LiteLLMParamsBody) -> None: + """Non-Anthropic models the runner role cannot invoke, deployments carrying + their own AWS identity (routing those would replace the assume-role chain the + batch suite exists to prove), explicit endpoints, and unmounted regions.""" + routed: Final = route_cache_model( + params, lambda mount: None if mount not in EDGE_MOUNTS else f"http://edge.invalid/{mount}", enabled=True, + ) + assert routed is params or routed.aws_bedrock_runtime_endpoint == params.aws_bedrock_runtime_endpoint + + +@pytest.mark.parametrize("mode", ["batch", "realtime", "image_generation"]) +def test_a_bedrock_deployment_with_a_mode_keeps_its_direct_route(mode: ModelMode) -> None: + params: Final = LiteLLMParamsBody(model="bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0") + assert route_cache_model( + params, lambda mount: f"http://edge.invalid/{mount}", enabled=True, mode=mode, + ) is params + + def test_rollback_and_live_only_policy_keep_direct_provider_route() -> None: params: Final = LiteLLMParamsBody(model="openai/test") assert route_cache_model(params, lambda _: "http://edge.invalid", enabled=False) is params @@ -366,14 +739,16 @@ class PublishOutage: def test_write_outage_preserves_success_without_hidden_retry(store: RedisResponseStore, provider: Provider) -> None: unavailable: Final = replace(store, client=PublishOutage(store.client)) - cache: Final = CacheEdge(unavailable, SECRET) + cache: Final = cache_edge(unavailable) with edge(cache, provider) as url: assert call(url).body == SUCCESS assert call(url).body == SUCCESS assert len(provider.hits) == 2 assert dict(cache.counters.counts)["write_failures"] == 2 - with edge(CacheEdge(store, SECRET), provider) as url: + with edge(cache_edge(store), provider) as url: assert call(url).body == SUCCESS + assert len(provider.hits) == 3 + with edge(cache_edge(store), provider) as url: assert call(url).body == SUCCESS assert len(provider.hits) == 3 @@ -382,45 +757,41 @@ def test_connection_failure_releases_capture_lease(store: RedisResponseStore) -> with socket.socket() as unavailable: unavailable.bind(("127.0.0.1", 0)) url: Final = f"http://127.0.0.1:{unavailable.getsockname()[1]}/v1/chat/completions" - cache: Final = CacheEdge(store, SECRET) - assert isinstance(cache.forward("POST", url, HEADERS, BODY, 0.2), NetworkError) - prepared: Final = prepare_forward("POST", url, HEADERS, BODY) - assert isinstance(prepared, PreparedForward) - key: Final = exact_key(SECRET, "POST", url, prepared.headers, BODY) - slot: Final = store.lookup(key) - assert isinstance(slot, CaptureLease) - assert store.release(key, slot) + cache: Final = cache_edge(store) + assert isinstance(cache.forward("openai", "POST", url, dict(HEADERS), BODY, 0.2), NetworkError) + key: Final = slot_key(url) + lease: Final = store.lookup(key) + assert isinstance(lease, CaptureLease) + assert store.release(key, lease) assert dict(cache.counters.counts)["rejected"] == 1 def test_close_before_first_chunk_releases_lease(store: RedisResponseStore, provider: Provider) -> None: url: Final = f"http://127.0.0.1:{provider.server_port}/v1/chat/completions" - cache: Final = CacheEdge(store, SECRET) - head: Final = cache.forward("POST", url, HEADERS, BODY, 5) + cache: Final = cache_edge(store) + head: Final = cache.forward("openai", "POST", url, dict(HEADERS), BODY, 5) assert isinstance(head, StreamHead) head.steps.close() - prepared: Final = prepare_forward("POST", url, HEADERS, BODY) - assert isinstance(prepared, PreparedForward) - key: Final = exact_key(SECRET, "POST", url, prepared.headers, BODY) - slot: Final = store.lookup(key) - assert isinstance(slot, CaptureLease) - assert store.release(key, slot) + key: Final = slot_key(url) + lease: Final = store.lookup(key) + assert isinstance(lease, CaptureLease) + assert store.release(key, lease) def test_effective_account_change_cannot_reuse_cache( store: RedisResponseStore, provider: Provider, monkeypatch: pytest.MonkeyPatch, tmp_path, ) -> None: url: Final = f"http://127.0.0.1:{provider.server_port}/v1/chat/completions" - cache: Final = CacheEdge(store, SECRET) - for account in ("account-a", "account-b", "account-b"): + caches: Final = tuple(cache_edge(store) for _ in range(3)) + for account, cache in zip(("account-a", "account-b", "account-b"), caches, strict=True): netrc = tmp_path / account netrc.write_text(f"machine 127.0.0.1 login {account} password synthetic\n") monkeypatch.setenv("NETRC", str(netrc)) - head = cache.forward("POST", url, HEADERS, BODY, 5) + head = cache.forward("openai", "POST", url, dict(HEADERS), BODY, 5) assert isinstance(head, StreamHead) assert b"".join(step.data for step in head.steps if isinstance(step, StreamChunk)) == SUCCESS assert len(provider.hits) == 2 - assert dict(cache.counters.counts)["hits"] == 1 + assert dict(caches[2].counters.counts)["hits"] == 1 def test_enabled_environment_reuses_store_across_fresh_backends( @@ -449,7 +820,7 @@ def test_enabled_environment_reuses_store_across_fresh_backends( def test_duplicate_headers_bypass_cache_and_count_live_calls( store: RedisResponseStore, provider: Provider, known_mount: bool, ) -> None: - cache: Final = CacheEdge(store, SECRET) + cache: Final = cache_edge(store) with edge(cache, provider) as url: parsed: Final = urlsplit(url) for _ in range(2): diff --git a/tests/e2e/fixture_canonical.py b/tests/e2e/fixture_canonical.py index e76d63ca33b..019c011aa67 100644 --- a/tests/e2e/fixture_canonical.py +++ b/tests/e2e/fixture_canonical.py @@ -51,6 +51,9 @@ SECRET_FIELD_SUFFIXES: Final[tuple[str, ...]] = ( ) SECRET_PLACEHOLDER: Final = "" +MARKER_PATTERN: Final = re.compile(r"(?" + PLACEHOLDER_RULES: Final[tuple[tuple[re.Pattern[str], str], ...]] = ( (re.compile(r"(?"), ( @@ -67,7 +70,7 @@ PLACEHOLDER_RULES: Final[tuple[tuple[re.Pattern[str], str], ...]] = ( re.compile(r"\b(?:chatcmpl|msgbatch|msg|resp|batch|call|req|ftjob|gen|file)[-_][A-Za-z0-9]{8,}\b"), "", ), - (re.compile(r"(?"), + (MARKER_PATTERN, MARKER_PLACEHOLDER), ) diff --git a/tests/e2e/models.py b/tests/e2e/models.py index 7101438c5f8..7550bfdc150 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -951,6 +951,7 @@ class LiteLLMParamsBody(BaseModel): aws_access_key_id: str | None = None aws_secret_access_key: str | None = None aws_region_name: str | None = None + aws_bedrock_runtime_endpoint: str | None = None vertex_project: str | None = None vertex_location: str | None = None vertex_credentials: str | None = None diff --git a/tests/e2e/provider_cache.py b/tests/e2e/provider_cache.py index 0c6eac75a43..1dc2f99abe5 100644 --- a/tests/e2e/provider_cache.py +++ b/tests/e2e/provider_cache.py @@ -23,12 +23,18 @@ from e2e_http import ( prepare_forward, primed_steps, ) +from fixture_canonical import MARKER_PATTERN, MARKER_PLACEHOLDER +from fixture_mode import SESSION_TEST_KEY, current_test_key from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter, ValidationError 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"}) +SIGNATURE_HEADERS: Final = frozenset( + {"authorization", "x-amz-date", "x-amz-security-token", "x-amz-content-sha256"} +) +BEDROCK_MOUNT_PREFIX: Final = "bedrock" JSON_VALUE: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue) @@ -56,6 +62,7 @@ class CacheUnavailable: type CacheLookup = CacheHit | CaptureLease | CacheBusy | CacheUnavailable +type RequestSigner = Callable[[str, str, Mapping[str, str], bytes | None], dict[str, str]] class ResponseStore(Protocol): @@ -83,28 +90,51 @@ class SignedResponse(BaseModel): signature: str -def exact_key(secret: bytes, method: str, url: str, headers: Mapping[str, str], body: bytes | None) -> str: +def canonical_text(value: str) -> str: + return MARKER_PATTERN.sub(MARKER_PLACEHOLDER, value) + + +def canonical_body(body: bytes) -> bytes: + try: + return canonical_text(body.decode("utf-8")).encode("utf-8") + except UnicodeDecodeError: + return body + + +def request_identity( + secret: bytes, test_key: str, method: str, url: str, headers: Mapping[str, str], body: bytes | None, +) -> str: fields: Final = ( - b"provider-cache-exact-v1", method.encode(), url.encode(), + b"provider-cache-canonical-v2", test_key.encode(), method.encode(), canonical_text(url).encode(), *(part.encode() for pair in sorted(headers.items()) for part in pair), - b"no-body" if body is None else b"body", b"" if body is None else body, + b"no-body" if body is None else b"body", b"" if body is None else canonical_body(body), ) encoded: Final = b"".join(len(part).to_bytes(8, "big") + part for part in fields) return hmac.new(secret, encoded, hashlib.sha256).hexdigest() -def cacheable_endpoint(method: str, url: str, body: bytes | None) -> bool: - return ( - method == "POST" - and urlsplit(url).path in {"/v1/chat/completions", "/v1/messages"} - and body is not None - and len(body) <= MAX_REQUEST_BYTES - ) +def slotted_key(secret: bytes, identity: str, slot: int) -> str: + return hmac.new(secret, f"{identity}:{slot}".encode(), hashlib.sha256).hexdigest() -def successful_response(url: str, status: int, headers: Mapping[str, str], body: bytes) -> bool: +def is_bedrock(mount: str) -> bool: + return mount.partition("/")[0] == BEDROCK_MOUNT_PREFIX + + +def cacheable_endpoint(mount: str, method: str, url: str, body: bytes | None) -> bool: + if method != "POST" or body is None or len(body) > MAX_REQUEST_BYTES: + return False + path: Final = urlsplit(url).path + if is_bedrock(mount): + return path.startswith("/model/") and path.endswith(("/converse", "/invoke")) + return path in {"/v1/chat/completions", "/v1/messages"} + + +def successful_response(mount: str, 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 is_bedrock(mount): + return complete_bedrock_response(url, body) streaming: Final = "text/event-stream" in headers.get("content-type", "").lower() if streaming: try: @@ -147,6 +177,26 @@ def successful_response(url: str, status: int, headers: Mapping[str, str], body: ) +def complete_bedrock_response(url: str, body: bytes) -> bool: + """Converse answers with ``output`` plus a ``stopReason``; InvokeModel on an + Anthropic model answers the Anthropic message shape. Either way a truncated + or error body is missing the terminator field, which is what makes it safe to + record. The streaming variants never reach here: they are not cacheable.""" + try: + value: Final = JSON_VALUE.validate_json(body) + except ValidationError: + return False + if not isinstance(value, dict) or "message" in value: + return False + if urlsplit(url).path.endswith("/converse"): + return isinstance(value.get("output"), dict) and isinstance(value.get("stopReason"), str) + return ( + value.get("type") == "message" + and isinstance(value.get("content"), list) + and isinstance(value.get("stop_reason"), str) + ) + + def complete_chat_stream(values: tuple[JsonValue, ...]) -> bool: if any(not isinstance(value, dict) or not isinstance(value.get("choices"), list) for value in values): return False @@ -172,7 +222,7 @@ def encode_response(secret: bytes, response: CachedResponse) -> bytes: return SignedResponse(response=raw, signature=hmac.new(secret, raw.encode(), hashlib.sha256).hexdigest()).model_dump_json().encode() -def decode_response(secret: bytes, key: str, payload: bytes, url: str) -> CachedResponse | None: +def decode_response(secret: bytes, key: str, payload: bytes, mount: str, url: str) -> CachedResponse | None: if len(payload) > 2 * MAX_RESPONSE_BYTES: return None try: @@ -183,7 +233,9 @@ def decode_response(secret: bytes, key: str, payload: bytes, url: str) -> Cached chunks: Final = tuple(base64.b64decode(chunk, validate=True) for chunk in response.chunks) except (ValidationError, ValueError): return None - if response.request_key != key or not successful_response(url, response.status_code, response.headers, b"".join(chunks)): + if response.request_key != key or not successful_response( + mount, url, response.status_code, response.headers, b"".join(chunks) + ): return None return response @@ -199,6 +251,24 @@ class CacheCounters: self.counts = tuple((current | {name: current.get(name, 0) + 1}).items()) +@dataclass(slots=True) +class SlotCounter: + """FIFO position of a request among the canonically identical ones its test + has already sent. Two calls in one test that differ only by ``unique_marker`` + canonicalize the same, so without this they would share one recording and the + second would replay the first's provider response id.""" + + counts: tuple[tuple[str, int], ...] = () + lock: threading.Lock = field(default_factory=threading.Lock) + + def take(self, identity: str) -> int: + with self.lock: + current: Final = dict(self.counts) + taken: Final = current.get(identity, 0) + self.counts = tuple((current | {identity: taken + 1}).items()) + return taken + + @dataclass(slots=True) class ResponseCapture: buffer: io.BytesIO = field(default_factory=io.BytesIO) @@ -231,9 +301,12 @@ class CacheEdge: store: ResponseStore secret: bytes = field(repr=False) counters: CacheCounters = field(default_factory=CacheCounters) + slots: SlotCounter = field(default_factory=SlotCounter) + signers: Mapping[str, RequestSigner] = field(default_factory=dict) wait_seconds: float = 2.0 clock: Callable[[], float] = time.monotonic sleep: Callable[[float], None] = time.sleep + test_key: Callable[[], str] = current_test_key def lookup(self, key: str) -> CacheLookup: deadline: Final = self.clock() + self.wait_seconds @@ -241,39 +314,71 @@ class CacheEdge: self.sleep(min(0.05, max(0, deadline - self.clock()))) return result - def forward(self, method: str, url: str, headers: dict[str, str], body: bytes | None, timeout: float) -> StreamHead | NetworkError: - if not cacheable_endpoint(method, url, body): - self.counters.increment("bypass") - self.counters.increment("upstream_attempts") - return forward_stream(method, url, headers=headers, body=body, timeout=timeout) - prepared: Final = prepare_forward(method, url, headers, body) + def count(self, mount: str, name: str) -> None: + self.counters.increment(name) + self.counters.increment(f"mount:{mount}:{name}") + + def outbound(self, mount: str, method: str, url: str, headers: dict[str, str], body: bytes | None) -> dict[str, str]: + """The headers actually sent upstream. A signing mount gets a signature + minted over the upstream URL, because the edge rewrote the Host the proxy + signed and Bedrock verifies it.""" + signer: Final = self.signers.get(mount) + return headers if signer is None else signer(method, url, headers, body) + + def keyed(self, mount: str, headers: Mapping[str, str]) -> Mapping[str, str]: + """A signing mount's signature headers are the edge's own and carry a + timestamp, so keying on them would make every request a permanent miss. + Every other mount keys on its headers whole, credentials included, so a + different account can never read another's recording.""" + if mount not in self.signers: + return headers + return {name: value for name, value in headers.items() if name.lower() not in SIGNATURE_HEADERS} + + def forward( + self, mount: str, method: str, url: str, headers: dict[str, str], body: bytes | None, timeout: float, + ) -> StreamHead | NetworkError: + test_key: Final = self.test_key() + if test_key == SESSION_TEST_KEY or not cacheable_endpoint(mount, method, url, body): + self.count(mount, "bypass") + self.count(mount, "upstream_attempts") + return forward_stream( + method, url, headers=self.outbound(mount, method, url, headers, body), body=body, timeout=timeout, + ) + prepared: Final = prepare_forward(method, url, self.outbound(mount, method, url, headers, body), body) if isinstance(prepared, NetworkError): - self.counters.increment("rejected") + self.count(mount, "rejected") return prepared - key: Final = exact_key(self.secret, method, url, prepared.headers, body) + identity: Final = request_identity( + self.secret, test_key, method, url, self.keyed(mount, prepared.headers), body, + ) + key: Final = slotted_key(self.secret, identity, self.slots.take(identity)) found: Final = self.lookup(key) if isinstance(found, CacheHit): - response: Final = decode_response(self.secret, key, found.payload, url) + response: Final = decode_response(self.secret, key, found.payload, mount, url) if response is not None and self.clock() < found.valid_until: - self.counters.increment("hits") + self.count(mount, "hits") return StreamHead(response.status_code, response.headers, response_steps(response)) - self.counters.increment("corrupt" if response is None else "expired") + self.count(mount, "corrupt" if response is None else "expired") self.store.discard(key, found.payload) capture_slot: Final = self.lookup(key) if isinstance(found, CacheHit) else found - self.counters.increment("misses") + self.count(mount, "misses") if isinstance(capture_slot, CacheUnavailable): - self.counters.increment("cache_errors") - self.counters.increment("upstream_attempts") + self.count(mount, "cache_errors") + self.count(mount, "upstream_attempts") head: Final = forward_prepared_stream(prepared, timeout) if not isinstance(capture_slot, CaptureLease): return head if isinstance(head, NetworkError): self.store.release(key, capture_slot) - self.counters.increment("rejected") + self.count(mount, "rejected") return head - return StreamHead(head.status_code, head.headers, primed_steps(self.capture(key, capture_slot, url, head))) + return StreamHead( + head.status_code, head.headers, primed_steps(self.capture(mount, key, capture_slot, url, head)), + ) - def capture(self, key: str, lease: CaptureLease, url: str, head: StreamHead) -> Generator[StreamStep, None, None]: + def capture( + self, mount: str, key: str, lease: CaptureLease, url: str, head: StreamHead, + ) -> Generator[StreamStep, None, None]: capture: Final = ResponseCapture() try: with closing(head.steps): @@ -285,15 +390,15 @@ class CacheEdge: 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") + if not capture.eligible or not successful_response(mount, url, head.status_code, headers, b"".join(chunks)): + self.count(mount, "rejected") return response: Final = CachedResponse( 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)) - self.counters.increment("writes" if published else "write_failures") + self.count(mount, "writes" if published else "write_failures") finally: self.store.release(key, lease) capture.buffer.close() diff --git a/tests/e2e/provider_cache_routing.py b/tests/e2e/provider_cache_routing.py index 24599b5a313..e7e4899eb71 100644 --- a/tests/e2e/provider_cache_routing.py +++ b/tests/e2e/provider_cache_routing.py @@ -8,14 +8,57 @@ from models import LiteLLMParamsBody, ModelMode LIVE_PROVIDER_REQUIRED: Final[ContextVar[bool]] = ContextVar("live_provider_required", default=False) +DEFAULT_BEDROCK_REGION: Final = "us-east-1" +BEDROCK_ANTHROPIC_INFIX: Final = "anthropic." + + +def bedrock_mount(params: LiteLLMParamsBody) -> str | None: + """The edge mount an Anthropic-on-Bedrock deployment belongs to, or None. + + Only the Anthropic models route. The edge validates converse and invoke + bodies by their Anthropic and Converse terminator fields, and the runner role + is allowed to invoke exactly those models, so Bedrock embeddings, image + generation, rerank and realtime keep their existing direct path rather than + reaching an edge that could neither sign nor validate for them.""" + route: Final = params.model.partition("/")[2] + model: Final = route.partition("/")[2] or route + if BEDROCK_ANTHROPIC_INFIX not in model: + return None + return f"bedrock/{params.aws_region_name or DEFAULT_BEDROCK_REGION}" + + +def route_bedrock( + params: LiteLLMParamsBody, base_for: Callable[[str], str | None], mode: ModelMode | None, +) -> LiteLLMParamsBody: + """Deployments that carry their own AWS identity stay off the edge. The edge + re-signs with the run pod's role, so routing an `aws_role_name` deployment + would quietly replace the very assume-role chain that test exists to prove.""" + if mode is not None or params.aws_role_name is not None or params.aws_access_key_id is not None: + return params + if params.api_base is not None or params.aws_bedrock_runtime_endpoint is not None: + return params + mount: Final = bedrock_mount(params) + if mount is None: + return params + base: Final = base_for(mount) + if base is None: + return params + return params.model_copy(update={"aws_bedrock_runtime_endpoint": base}) + def route_cache_model( params: LiteLLMParamsBody, base_for: Callable[[str], str | None], *, enabled: bool, mode: ModelMode | None = None, ) -> LiteLLMParamsBody: - if not enabled or mode == "realtime" or LIVE_PROVIDER_REQUIRED.get() or params.api_base is not None or params.mock_response is not None: + if not enabled or LIVE_PROVIDER_REQUIRED.get() or params.mock_response is not None: + return params + if params.litellm_credential_name is not None: return params provider: Final = params.model.partition("/")[0] - if provider not in {"openai", "anthropic"} or params.litellm_credential_name is not None: + if provider == "bedrock": + return route_bedrock(params, base_for, mode) + if mode == "realtime" or params.api_base is not None: + return params + if provider not in {"openai", "anthropic"}: return params base: Final = base_for(provider) if base is None: diff --git a/tests/e2e/provider_edge.py b/tests/e2e/provider_edge.py index dda9e6f8e4f..8f718ad1967 100644 --- a/tests/e2e/provider_edge.py +++ b/tests/e2e/provider_edge.py @@ -48,7 +48,7 @@ import threading from collections import deque from collections.abc import Generator, Mapping, Sequence from contextlib import closing, contextmanager -from dataclasses import dataclass, field +from dataclasses import dataclass, field, replace from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from itertools import islice from pathlib import Path @@ -94,17 +94,41 @@ from fixture_mode import ( parse_fixture_mode, ) from fixture_profile import IneligibleRequest, MatchProfile, match_profile, strict_identity -from provider_cache import CacheEdge +from provider_cache import CacheEdge, RequestSigner, is_bedrock from provider_cache_routing import LIVE_PROVIDER_REQUIRED from pydantic import JsonValue, TypeAdapter +BEDROCK_REGIONS: Final[tuple[str, ...]] = ("us-east-1",) + EDGE_MOUNTS: Final[Mapping[str, str]] = MappingProxyType( { "openai": "https://api.openai.com", "anthropic": "https://api.anthropic.com", + **{ + f"bedrock/{region}": f"https://bedrock-runtime.{region}.amazonaws.com" + for region in BEDROCK_REGIONS + }, } ) + +@dataclass(frozen=True, slots=True) +class ResolvedMount: + mount: str + upstream_base: str + upstream_path: str + + +def resolve_mount(path: str, mounts: Mapping[str, str]) -> ResolvedMount | None: + """Longest mount prefix wins, so a region-qualified mount such as + ``bedrock/us-east-1`` resolves whole instead of leaving the region as the + first segment of the upstream path.""" + trimmed: Final = path.lstrip("/") + for mount in sorted(mounts, key=len, reverse=True): + if trimmed == mount or trimmed.startswith(f"{mount}/"): + return ResolvedMount(mount, mounts[mount], trimmed[len(mount):].lstrip("/")) + return None + REPLAY_MISS_STATUS: Final = 599 _HOP_BY_HOP_HEADERS: Final[frozenset[str]] = frozenset( @@ -754,14 +778,14 @@ def _handle_record( def _handle_live( method: str, url: str, headers: Mapping[str, str], body: bytes | None, timeout: float, - cache: CacheEdge | None = None, + cache: CacheEdge | None = None, mount: str = "", ) -> EdgeOutcome: forwarded: Final = { name: value for name, value in headers.items() if name.lower() not in _REQUEST_DROPPED_HEADERS } head: Final = ( forward_stream(method, url, headers=forwarded, body=body, timeout=timeout) - if cache is None else cache.forward(method, url, forwarded, body, timeout) + if cache is None else cache.forward(mount, method, url, forwarded, body, timeout) ) match head: case NetworkError(message=message): @@ -796,10 +820,13 @@ def handle_edge_request( prefix, then record (forward + persist) or replay (serve from the bundle). Socket-free so unit tests exercise every branch without a server.""" split: Final = urlsplit(raw_path) - mount, _, upstream_path = split.path.lstrip("/").partition("/") - upstream_base: Final = mounts.get(mount) - if upstream_base is None: - return _text_reply(404, f"unknown provider mount {mount!r}; known mounts: {', '.join(sorted(mounts))}") + resolved: Final = resolve_mount(split.path, mounts) + if resolved is None: + unknown: Final = split.path.lstrip("/").partition("/")[0] + return _text_reply(404, f"unknown provider mount {unknown!r}; known mounts: {', '.join(sorted(mounts))}") + mount: Final = resolved.mount + upstream_base: Final = resolved.upstream_base + upstream_path: Final = resolved.upstream_path profile: Final = ( backend.recorder.profile if isinstance(backend, RecordEdge) @@ -830,7 +857,8 @@ def handle_edge_request( match backend: case CacheEdge(): return _handle_live( - method, _upstream_url(upstream_base, upstream_path, split.query), headers, body, timeout, backend, + method, _upstream_url(upstream_base, upstream_path, split.query), headers, body, timeout, + backend, mount, ) case LiveEdge(): return _handle_live( @@ -891,7 +919,7 @@ class _EdgeHandler(BaseHTTPRequestHandler): ) if isinstance(edge_server.backend, CacheEdge) and duplicate_headers: edge_server.backend.counters.increment("duplicate_header_bypass") - if urlsplit(self.path).path.lstrip("/").partition("/")[0] in edge_server.mounts: + if resolve_mount(urlsplit(self.path).path, edge_server.mounts) is not None: edge_server.backend.counters.increment("upstream_attempts") outcome: Final = handle_edge_request( selected_backend, @@ -1079,6 +1107,8 @@ def provider_edge_api_base( return _shared_cache_edge(bind_host, advertise_host, forward_timeout).api_base(mount) return None case "record" | "replay": + if is_bedrock(mount): + return None if mount not in EDGE_MOUNTS: raise ValueError(f"unknown provider mount {mount!r}; known mounts: {', '.join(sorted(EDGE_MOUNTS))}") return _shared_edge(mode, bundle_dir, bind_host, advertise_host, forward_timeout, match_profile()).api_base( @@ -1108,7 +1138,17 @@ def configured_cache_backend() -> CacheEdge | None: return None from provider_cache_redis import configured_cache - return configured_cache() + cache: Final = configured_cache() + return None if cache is None else replace(cache, signers=bedrock_signers()) + + +@functools.lru_cache(maxsize=1) +def bedrock_signers() -> Mapping[str, RequestSigner]: + """One signer per mounted Bedrock region, built lazily so a run that never + mounts Bedrock neither imports botocore nor resolves an AWS identity.""" + from provider_edge_bedrock import bedrock_signer + + return MappingProxyType({f"bedrock/{region}": bedrock_signer(region) for region in BEDROCK_REGIONS}) @functools.lru_cache(maxsize=8) diff --git a/tests/e2e/provider_edge_bedrock.py b/tests/e2e/provider_edge_bedrock.py new file mode 100644 index 00000000000..5d8148482d7 --- /dev/null +++ b/tests/e2e/provider_edge_bedrock.py @@ -0,0 +1,72 @@ +"""SigV4 re-signing for Bedrock traffic routed through the provider edge. + +Bedrock is the one provider the edge could never mount. SigV4 signs the Host +header, so rewriting ``api_base`` to point at the edge invalidates the proxy's +signature and Bedrock rejects the call before it reaches a model. The edge +therefore has to drop the proxy's signature and mint its own over the upstream +URL it is actually about to call. + +The identity it signs with is the run pod's own, from the EKS Pod Identity +association on ServiceAccount ``buildkite-e2e-run``. That role carries Bedrock +invoke and converse on an allowlist of the Anthropic models the suite registers +and nothing else, so a re-signed call can reach exactly the models the suite +already uses. The proxy's own Bedrock credentials are not involved in a routed +deployment, which is why ``aws_role_name`` deployments stay off the edge: their +whole point is to prove the product's assume-role chain. + +Signature headers are excluded from the cache key by the caller, and they have +to be: ``x-amz-date`` is a timestamp, so keying on it would make every Bedrock +request a permanent miss. +""" + +from __future__ import annotations + +import functools +from collections.abc import Callable, Mapping +from dataclasses import dataclass +from typing import Final + +from botocore.auth import SigV4Auth +from botocore.awsrequest import AWSRequest +from botocore.credentials import Credentials +from botocore.session import Session +from provider_cache import SIGNATURE_HEADERS + +BEDROCK_SERVICE: Final = "bedrock" + + +class MissingAwsCredentials(RuntimeError): + """No AWS identity is resolvable, so the edge cannot sign for Bedrock.""" + + +@dataclass(frozen=True, slots=True) +class BedrockSigner: + region: str + credentials: Callable[[], Credentials] + + def __call__(self, method: str, url: str, headers: Mapping[str, str], body: bytes | None) -> dict[str, str]: + unsigned: Final = { + name: value for name, value in headers.items() if name.lower() not in SIGNATURE_HEADERS + } + request: Final = AWSRequest(method=method, url=url, headers=unsigned, data=body or b"") + SigV4Auth(self.credentials(), BEDROCK_SERVICE, self.region).add_auth(request) + return dict(request.headers) + + +@functools.lru_cache(maxsize=1) +def pod_credentials() -> Credentials: + """The run pod's own identity, resolved once per process through botocore's + ordinary chain, which reaches Pod Identity at the ``container-role`` link.""" + resolved: Final = Session().get_credentials() + if resolved is None: + raise MissingAwsCredentials( + "the provider edge is mounted for Bedrock but no AWS credentials resolve; " + "the run pod gets them from the Pod Identity association on buildkite-e2e-run" + ) + return resolved + + +def bedrock_signer(region: str, credentials: Callable[[], Credentials] = pod_credentials) -> BedrockSigner: + """Credentials are resolved on the first signed request, not here, so a run + that mounts Bedrock but never calls it needs no AWS identity at all.""" + return BedrockSigner(region, credentials) diff --git a/tests/e2e/test_provider_edge.py b/tests/e2e/test_provider_edge.py index 5d0c79f26f6..c8d70697182 100644 --- a/tests/e2e/test_provider_edge.py +++ b/tests/e2e/test_provider_edge.py @@ -1279,15 +1279,30 @@ class TestApiBaseSeam: ) def test_unknown_mount_raises_naming_the_known_mounts(self, tmp_path: Path) -> None: - with pytest.raises(ValueError, match="unknown provider mount 'bedrock'"): + with pytest.raises(ValueError, match="unknown provider mount 'cohere'"): provider_edge_api_base( - "bedrock", + "cohere", mode_raw="record", bundle_dir=tmp_path / "bundle", bind_host="127.0.0.1", advertise_host="127.0.0.1", ) + @pytest.mark.parametrize("mode_raw", ["record", "replay"]) + def test_bedrock_never_wires_a_bundle_because_the_edge_cannot_sign_into_one( + self, tmp_path: Path, mode_raw: str, + ) -> None: + """Record and replay serve from a bundle without re-signing, so a Bedrock + deployment pointed at that edge would send the proxy's signature over a + rewritten Host. It keeps its direct route in both modes.""" + assert provider_edge_api_base( + "bedrock/us-east-1", + mode_raw=mode_raw, + bundle_dir=tmp_path / "bundle", + bind_host="127.0.0.1", + advertise_host="127.0.0.1", + ) is None + def test_record_mode_boots_one_shared_edge_and_prepares_the_bundle(self, tmp_path: Path) -> None: root = tmp_path / "bundle" first = provider_edge_api_base( From b68e60f7061a2102fa076ff7ca6368f819d13770 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 16 Sep 2026 02:34:09 -0700 Subject: [PATCH 02/21] feat(e2e): cache the responses and embeddings endpoints behind the edge Chat completions and messages were the only cacheable paths. The suite also drives /v1/embeddings and /v1/responses through the same OpenAI mount, so both now cache, each with its own completeness rule: a chat response's `choices` check would reject a perfectly good embedding, and a Responses run that never reached `response.completed` must stay out of the cache the same way a truncated stream does. Vertex and Gemini stay off the edge. litellm's `_check_custom_proxy` rewrites a path-prefixed vertex api_base into `{api_base}:{endpoint}`, dropping project, location and model, so a mount under a path prefix cannot work without a root-mounted edge on its own port or a change in litellm. Shipping an unvalidated URL guess would have been worse than saying so in PROVIDER_CACHE.md. Also finishes the MountPolicy move: a mount now carries its signer and its unkeyed headers together instead of a bare signer map. --- .../test_provider_cache.py | 102 +++++++++++++++++- tests/e2e/PROVIDER_CACHE.md | 28 ++++- tests/e2e/provider_cache.py | 63 +++++++++-- tests/e2e/provider_edge.py | 15 ++- tests/e2e/provider_edge_bedrock.py | 2 +- 5 files changed, 186 insertions(+), 24 deletions(-) diff --git a/tests/code_coverage_tests/test_provider_cache.py b/tests/code_coverage_tests/test_provider_cache.py index 5d35981344d..5c491e1cdcd 100644 --- a/tests/code_coverage_tests/test_provider_cache.py +++ b/tests/code_coverage_tests/test_provider_cache.py @@ -25,6 +25,7 @@ from provider_cache import ( CacheEdge, CacheHit, CaptureLease, + MountPolicy, ResponseStore, cacheable_endpoint, request_identity, @@ -179,7 +180,9 @@ def slot_key( def bedrock_cache_edge(store: ResponseStore, test_key: str = TEST_KEY) -> CacheEdge: return CacheEdge( store, SECRET, test_key=lambda: test_key, - signers={BEDROCK_MOUNT: bedrock_signer("us-east-1", lambda: STATIC_CREDENTIALS)}, + policies={BEDROCK_MOUNT: MountPolicy( + sign=bedrock_signer("us-east-1", lambda: STATIC_CREDENTIALS), unkeyed_headers=SIGNATURE_HEADERS, + )}, ) @@ -501,6 +504,98 @@ def test_counters_attribute_every_outcome_to_its_mount( assert counts["mount:anthropic:rejected"] == 1 and "mount:openai:rejected" not in counts +EMBEDDING_SUCCESS: Final = ( + b'{"object":"list","data":[{"object":"embedding","index":0,"embedding":[0.1,0.2]}],' + b'"model":"text-embedding-3-small","usage":{"prompt_tokens":2,"total_tokens":2}}' +) +RESPONSE_SUCCESS: Final = b'{"id":"resp_synthetic","object":"response","status":"completed","output":[]}' +RESPONSE_STREAM_SUCCESS: Final = ( + b'data: {"type":"response.created","response":{"id":"resp_synthetic"}}\n\n' + b'data: {"type":"response.completed","response":{"id":"resp_synthetic","status":"completed"}}\n\n' +) + + +@contextmanager +def openai_edge(cache: CacheEdge, provider: Provider, path: str) -> Generator[str, None, None]: + upstream: Final = f"http://127.0.0.1:{provider.server_port}" + running: Final = start_provider_edge(cache, mounts={"openai": upstream}) + try: + yield running.edge.api_base("openai") + path + finally: + running.shutdown() + + +class TestNonChatOpenAiEndpoints: + """Chat and messages were the only cacheable paths. Embeddings and responses + are the other two JSON endpoints the suite drives through the same mount, and + each needs its own completeness rule: a chat response's ``choices`` check + would reject a perfectly good embedding.""" + + @pytest.mark.parametrize("path,response", [ + ("/v1/embeddings", EMBEDDING_SUCCESS), + ("/v1/responses", RESPONSE_SUCCESS), + ]) + def test_complete_responses_replay_on_the_next_run( + self, store: RedisResponseStore, provider: Provider, path: str, response: bytes, + ) -> None: + provider.response = response + for _ in range(2): + with openai_edge(cache_edge(store), provider, path) as url: + assert call(url, MARKED).body == response + assert len(provider.hits) == 1 + + def test_a_completed_response_stream_replays( + self, store: RedisResponseStore, provider: Provider, + ) -> None: + provider.stream = True + provider.response = RESPONSE_STREAM_SUCCESS + for _ in range(2): + with openai_edge(cache_edge(store), provider, "/v1/responses") as url: + assert call(url, MARKED).body == RESPONSE_STREAM_SUCCESS + assert len(provider.hits) == 1 + + @pytest.mark.parametrize("path,response", [ + ("/v1/embeddings", b'{"object":"list","data":[],"usage":{"prompt_tokens":0}}'), + ("/v1/embeddings", b'{"object":"list","data":[{"object":"embedding","index":0,"embedding":[]}],"usage":{}}'), + ("/v1/embeddings", b'{"object":"list","data":[{"object":"embedding","index":0,"embedding":[0.1]}]}'), + ("/v1/responses", b'{"id":"resp_x","object":"response","status":"incomplete","output":[]}'), + ("/v1/responses", b'{"id":"resp_x","object":"response","status":"in_progress","output":[]}'), + ("/v1/responses", b'{"id":"resp_x","object":"response","output":[]}'), + ]) + def test_incomplete_bodies_never_enter_the_cache( + self, store: RedisResponseStore, provider: Provider, path: str, response: bytes, + ) -> None: + provider.response = response + for _ in range(2): + with openai_edge(cache_edge(store), provider, path) as url: + assert call(url, MARKED).body == response + assert len(provider.hits) == 2 + + @pytest.mark.parametrize("payload", [ + b'data: {"type":"response.created","response":{"id":"resp_x"}}\n\n', + b'data: {"type":"response.created","response":{"id":"resp_x"}}\n\ndata: {"type":"response.failed"}\n\n', + b'data: {"type":"response.completed","response":{"id":"resp_x"}}\n\ndata: {"type":"response.created"}\n\n', + ]) + def test_a_response_stream_that_never_completed_is_never_cached( + self, store: RedisResponseStore, provider: Provider, payload: bytes, + ) -> None: + provider.stream = True + provider.response = payload + for _ in range(2): + with openai_edge(cache_edge(store), provider, "/v1/responses") as url: + assert call(url, MARKED).body == payload + assert len(provider.hits) == 2 + + @pytest.mark.parametrize("path,cacheable", [ + ("/v1/chat/completions", True), ("/v1/messages", True), + ("/v1/embeddings", True), ("/v1/responses", True), + ("/v1/audio/speech", False), ("/v1/images/generations", False), + ("/v1/files", False), ("/v1/batches", False), + ]) + def test_only_the_json_endpoints_are_cacheable(self, path: str, cacheable: bool) -> None: + assert cacheable_endpoint("openai", "POST", f"https://api.openai.com{path}", MARKED) is cacheable + + class TestBedrockSigning: """Bedrock is the reason the edge could not mount it before: SigV4 covers the Host header, so forwarding through a rewritten api_base invalidates the @@ -545,7 +640,10 @@ class TestBedrockSigning: return dict(headers) | {"authorization": f"AWS4-HMAC-SHA256 {url}", "x-amz-date": next(stamps)} def signing_edge() -> CacheEdge: - return CacheEdge(store, SECRET, test_key=lambda: TEST_KEY, signers={BEDROCK_MOUNT: varying}) + return CacheEdge( + store, SECRET, test_key=lambda: TEST_KEY, + policies={BEDROCK_MOUNT: MountPolicy(sign=varying, unkeyed_headers=SIGNATURE_HEADERS)}, + ) for _ in range(2): with bedrock_edge(signing_edge(), provider) as url: diff --git a/tests/e2e/PROVIDER_CACHE.md b/tests/e2e/PROVIDER_CACHE.md index 8635c9ed9ae..393a96e8c16 100644 --- a/tests/e2e/PROVIDER_CACHE.md +++ b/tests/e2e/PROVIDER_CACHE.md @@ -1,11 +1,29 @@ # Shared provider-response cache -`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 +`E2E_PROVIDER_CACHE=1` enables automatic response reuse in the live E2E mode. Standard OpenAI and Anthropic model registrations use the provider edge, as do Anthropic-on-Bedrock registrations that carry no AWS identity of their own. 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. 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 +The edge caches complete successful POST responses for `/v1/chat/completions`, `/v1/messages`, `/v1/embeddings` and `/v1/responses` on the OpenAI and Anthropic mounts, and for `/model/{id}/converse` and `/model/{id}/invoke` on a Bedrock mount, including streams. Unsupported endpoints pass through. Each endpoint family has its own completeness rule, so a truncated embedding or a Responses run that never reached `response.completed` is not stored + +## Request identity + +A recording belongs to one test. The key is a keyed digest over the test's node id, the method, the URL, the effective outbound headers (including authentication and HTTP-library defaults), body presence and the body bytes, with one normalization: a 12-hex-digit run, the shape `unique_marker()` mints, is replaced by a placeholder in both the URL and a UTF-8 body. Nothing else is normalized away. No prompts, JSON values or credentials are rewritten, and the rule is the one `fixture_canonical.py` already applies for record/replay, so there is a single definition of what a marker is + +Requests that differ only by their markers therefore share a canonical identity, which is what makes the cache reusable across builds: every e2e test salts its prompt afresh, so an exact-byte key would miss on every call. Within one test, calls that share a canonical identity are still recorded and replayed separately, by a FIFO slot index appended to the key. That matters because a replayed response carries the recorded provider response id, `LiteLLM_SpendLogs.request_id` is that id, and one shared recording answering two calls would collapse two spend rows into one + +Two different tests never share a recording, and a provider call made outside any test (fixtures, session setup) is never cached, because the identity has no test node id to bind to + +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 +## Bedrock + +Bedrock could not be mounted before because SigV4 signs the `Host` header, so a rewritten `api_base` failed signature verification at the provider. The edge now re-signs: it drops the proxy's signature headers, signs the upstream request with the run pod's own AWS identity from its EKS Pod Identity association, and forwards that. The signature headers are excluded from the key, since `x-amz-date` is a timestamp and keying on it would make every Bedrock call a permanent miss + +Only deployments that carry no AWS identity of their own route to the edge. A deployment with `aws_role_name`, `aws_access_key_id`, an `api_base` or an `aws_bedrock_runtime_endpoint` keeps its direct path, because re-signing it would quietly replace the very credential chain that test exists to prove. Only Anthropic models route, matching what the runner role is allowed to invoke and what the edge knows how to validate + +Vertex and Gemini are not mounted. litellm's `_check_custom_proxy` rewrites a path-prefixed Vertex `api_base` into `{api_base}:{endpoint}`, dropping project, location and model, so a mount under a path prefix cannot work without either a root-mounted edge on its own port or a change in litellm + Recordings are shared across workers and builds through dedicated Redis, separate from the candidate's own cache. They expire 86,400 seconds after capture starts, based on Redis time. Reads never extend expiry. There is no scheduled recapture: the next miss calls the provider again. Bounded coordination reduces duplicate concurrent calls, but slow or failed captures may lead to extra live calls after the wait expires ## Configuration @@ -18,16 +36,16 @@ The trusted runner receives: - `E2E_PROVIDER_CACHE_NAMESPACE`: shared environment namespace, independent of build and candidate revision - `E2E_PROVIDER_CACHE_METRICS_DIR`: optional per-process counter artifact directory -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 +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. Every counter is emitted twice, once as a flat total and once under `mount:{mount}:`, so a hit rate can be read per provider rather than only in aggregate. 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. 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 -Replay preserves the original response ID, usage and end-to-end headers. The proxy can therefore deduplicate repeated provider IDs when storing spend-log rows, just as it does when a live upstream returns the same ID twice. One spend-log row per invocation is not guaranteed for identical recorded responses. Existing spend reconciliation requests use distinct prompt markers and retain their distinct-ID and row-count assertions; accounting tests are not automatically excluded from caching +Replay preserves the original response ID, usage and end-to-end headers. The proxy can therefore deduplicate repeated provider IDs when storing spend-log rows, just as it does when a live upstream returns the same ID twice. One spend-log row per invocation is not guaranteed for identical recorded responses. Spend reconciliation keeps its distinct-ID and row-count assertions: its prompts differ by an index as well as a marker, so they stay distinct once markers are normalized, and calls that are canonically equal within one test take separate FIFO slots and separate recordings anyway. Accounting tests are not automatically excluded from caching Provider remaining-quota headers describe the captured response. Metrics derived from them are historical on a cache hit, not a measurement of current provider capacity. Gateway-generated API-key quota headers are a separate contract. A test of fresh provider quota or timing must use the live-provider policy; replay can still exercise how the proxy processes the recorded headers ## Qualification -`tests/code_coverage_tests/test_provider_cache.py` exercises local HTTP providers and disposable real Redis. CI runs these checks with the existing provider-edge and replay harness tests. These component checks do not establish Buildkite deployment, full-suite cross-build reuse or a genuine 24-hour expiry observation; those require separate runtime evidence +`tests/code_coverage_tests/test_provider_cache.py` exercises local HTTP providers and disposable real Redis, including the marker-canonical key, the FIFO slot index, per-test isolation, SigV4 re-signing against a local upstream, and each endpoint's completeness rule. CI runs these checks with the existing provider-edge and replay harness tests. These component checks do not establish Buildkite deployment, full-suite cross-build reuse or a genuine 24-hour expiry observation; those require separate runtime evidence diff --git a/tests/e2e/provider_cache.py b/tests/e2e/provider_cache.py index 1dc2f99abe5..a2b18a3a466 100644 --- a/tests/e2e/provider_cache.py +++ b/tests/e2e/provider_cache.py @@ -9,6 +9,7 @@ import time from collections.abc import Callable, Generator, Mapping from contextlib import closing from dataclasses import dataclass, field +from types import MappingProxyType from typing import Final, Literal, Protocol from urllib.parse import urlsplit @@ -35,6 +36,7 @@ SIGNATURE_HEADERS: Final = frozenset( {"authorization", "x-amz-date", "x-amz-security-token", "x-amz-content-sha256"} ) BEDROCK_MOUNT_PREFIX: Final = "bedrock" +OPENAI_JSON_PATHS: Final = frozenset({"/v1/chat/completions", "/v1/messages", "/v1/embeddings", "/v1/responses"}) JSON_VALUE: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue) @@ -65,6 +67,23 @@ type CacheLookup = CacheHit | CaptureLease | CacheBusy | CacheUnavailable type RequestSigner = Callable[[str, str, Mapping[str, str], bytes | None], dict[str, str]] +@dataclass(frozen=True, slots=True) +class MountPolicy: + """What a mount needs beyond plain forwarding. + + ``sign`` mints a fresh credential over the upstream URL, for providers whose + auth covers the Host the edge rewrote. ``unkeyed_headers`` names headers that + must stay out of the cache key because they change on every call and would + otherwise make the mount a permanent miss: a minted signature, or an OAuth + token the provider rotates. Naming one costs the guarantee that a recording + can never cross credentials, so a mount with a rotating token relies on the + environment holding one identity for that provider. Mounts with a static API + key name nothing here and keep the guarantee whole.""" + + sign: RequestSigner | None = None + unkeyed_headers: frozenset[str] = frozenset() + + class ResponseStore(Protocol): def lookup(self, key: str) -> CacheLookup: ... @@ -127,7 +146,7 @@ def cacheable_endpoint(mount: str, method: str, url: str, body: bytes | None) -> path: Final = urlsplit(url).path if is_bedrock(mount): return path.startswith("/model/") and path.endswith(("/converse", "/invoke")) - return path in {"/v1/chat/completions", "/v1/messages"} + return path in OPENAI_JSON_PATHS def successful_response(mount: str, url: str, status: int, headers: Mapping[str, str], body: bytes) -> bool: @@ -150,6 +169,8 @@ def successful_response(mount: str, url: str, status: int, headers: Mapping[str, return False if not values or any(not isinstance(value, dict) or "error" in value or value.get("type") == "error" for value in values): return False + if urlsplit(url).path == "/v1/responses": + return complete_responses_stream(values) if urlsplit(url).path == "/v1/chat/completions": return events[-1] == "[DONE]" and "[DONE]" not in events[:-1] and complete_chat_stream(values) return ( @@ -168,8 +189,17 @@ def successful_response(mount: str, url: str, status: int, headers: Mapping[str, return False if not isinstance(value, dict) or "error" in value: return False - if urlsplit(url).path == "/v1/messages": + path: Final = urlsplit(url).path + if path == "/v1/messages": return value.get("type") == "message" and isinstance(value.get("content"), list) and isinstance(value.get("stop_reason"), str) + if path == "/v1/embeddings": + data: Final = value.get("data") + return isinstance(data, list) and bool(data) and isinstance(value.get("usage"), dict) and all( + isinstance(item, dict) and isinstance(item.get("embedding"), list) and bool(item["embedding"]) + for item in data + ) + if path == "/v1/responses": + return value.get("object") == "response" and value.get("status") == "completed" choices: Final = value.get("choices") return isinstance(choices, list) and bool(choices) and all( isinstance(choice, dict) and isinstance(choice.get("message"), dict) and isinstance(choice.get("finish_reason"), str) @@ -197,6 +227,14 @@ def complete_bedrock_response(url: str, body: bytes) -> bool: ) +def complete_responses_stream(values: tuple[JsonValue, ...]) -> bool: + """The Responses API streams typed events and ends with ``response.completed``. + A run that failed, was cancelled, or ran out of tokens ends with a different + terminal event, so requiring that one keeps a half-finished response out.""" + last: Final = values[-1] + return isinstance(last, dict) and last.get("type") == "response.completed" + + def complete_chat_stream(values: tuple[JsonValue, ...]) -> bool: if any(not isinstance(value, dict) or not isinstance(value.get("choices"), list) for value in values): return False @@ -296,13 +334,16 @@ def response_steps(response: CachedResponse) -> Generator[StreamStep, None, None yield StreamChunk(base64.b64decode(chunk, validate=True)) +NO_POLICIES: Final[Mapping[str, MountPolicy]] = MappingProxyType({}) + + @dataclass(frozen=True, slots=True) class CacheEdge: store: ResponseStore secret: bytes = field(repr=False) counters: CacheCounters = field(default_factory=CacheCounters) slots: SlotCounter = field(default_factory=SlotCounter) - signers: Mapping[str, RequestSigner] = field(default_factory=dict) + policies: Mapping[str, MountPolicy] = NO_POLICIES wait_seconds: float = 2.0 clock: Callable[[], float] = time.monotonic sleep: Callable[[float], None] = time.sleep @@ -321,18 +362,18 @@ class CacheEdge: def outbound(self, mount: str, method: str, url: str, headers: dict[str, str], body: bytes | None) -> dict[str, str]: """The headers actually sent upstream. A signing mount gets a signature minted over the upstream URL, because the edge rewrote the Host the proxy - signed and Bedrock verifies it.""" - signer: Final = self.signers.get(mount) + signed and the provider verifies it.""" + signer: Final = self.policies.get(mount, MountPolicy()).sign return headers if signer is None else signer(method, url, headers, body) def keyed(self, mount: str, headers: Mapping[str, str]) -> Mapping[str, str]: - """A signing mount's signature headers are the edge's own and carry a - timestamp, so keying on them would make every request a permanent miss. - Every other mount keys on its headers whole, credentials included, so a - different account can never read another's recording.""" - if mount not in self.signers: + """Headers the cache key is built from. A mount keeps its credentials in + the key unless its policy names them unkeyed, so by default one account + can never read another's recording.""" + unkeyed: Final = self.policies.get(mount, MountPolicy()).unkeyed_headers + if not unkeyed: return headers - return {name: value for name, value in headers.items() if name.lower() not in SIGNATURE_HEADERS} + return {name: value for name, value in headers.items() if name.lower() not in unkeyed} def forward( self, mount: str, method: str, url: str, headers: dict[str, str], body: bytes | None, timeout: float, diff --git a/tests/e2e/provider_edge.py b/tests/e2e/provider_edge.py index 8f718ad1967..2606b26fe99 100644 --- a/tests/e2e/provider_edge.py +++ b/tests/e2e/provider_edge.py @@ -94,7 +94,7 @@ from fixture_mode import ( parse_fixture_mode, ) from fixture_profile import IneligibleRequest, MatchProfile, match_profile, strict_identity -from provider_cache import CacheEdge, RequestSigner, is_bedrock +from provider_cache import SIGNATURE_HEADERS, CacheEdge, MountPolicy, is_bedrock from provider_cache_routing import LIVE_PROVIDER_REQUIRED from pydantic import JsonValue, TypeAdapter @@ -1139,16 +1139,21 @@ def configured_cache_backend() -> CacheEdge | None: from provider_cache_redis import configured_cache cache: Final = configured_cache() - return None if cache is None else replace(cache, signers=bedrock_signers()) + return None if cache is None else replace(cache, policies=bedrock_policies()) @functools.lru_cache(maxsize=1) -def bedrock_signers() -> Mapping[str, RequestSigner]: - """One signer per mounted Bedrock region, built lazily so a run that never +def bedrock_policies() -> Mapping[str, MountPolicy]: + """One policy per mounted Bedrock region, built lazily so a run that never mounts Bedrock neither imports botocore nor resolves an AWS identity.""" from provider_edge_bedrock import bedrock_signer - return MappingProxyType({f"bedrock/{region}": bedrock_signer(region) for region in BEDROCK_REGIONS}) + return MappingProxyType( + { + f"bedrock/{region}": MountPolicy(sign=bedrock_signer(region), unkeyed_headers=SIGNATURE_HEADERS) + for region in BEDROCK_REGIONS + } + ) @functools.lru_cache(maxsize=8) diff --git a/tests/e2e/provider_edge_bedrock.py b/tests/e2e/provider_edge_bedrock.py index 5d8148482d7..73e4a16d272 100644 --- a/tests/e2e/provider_edge_bedrock.py +++ b/tests/e2e/provider_edge_bedrock.py @@ -58,7 +58,7 @@ def pod_credentials() -> Credentials: """The run pod's own identity, resolved once per process through botocore's ordinary chain, which reaches Pod Identity at the ``container-role`` link.""" resolved: Final = Session().get_credentials() - if resolved is None: + if resolved is None: # pyright: ignore[reportUnnecessaryComparison] # stubs miss the empty-chain None raise MissingAwsCredentials( "the provider edge is mounted for Bedrock but no AWS credentials resolve; " "the run pod gets them from the Pod Identity association on buildkite-e2e-run" From aebfcf7da3b5b13591232f4b559978d4f92aafb5 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 16 Sep 2026 02:45:06 -0700 Subject: [PATCH 03/21] fix(e2e): route Bedrock deployments whose region only the proxy can resolve Almost every Bedrock deployment in the suite declares aws_region_name="os.environ/AWS_REGION". The mount resolver treated that string as a region name, produced a mount nothing serves, and left the whole Anthropic-on-Bedrock surface on its direct path, which is the one thing mounting Bedrock was for. The run pod does not share the proxy's environment, so the harness genuinely cannot resolve that reference. A `us.` inference profile fans out across the US regions and is reachable from any of them, so those route to the default mount whatever the proxy resolved. A model that is not cross-region and declares its region that way keeps its direct path rather than being sent to a region it may not exist in. --- .../test_provider_cache.py | 27 +++++++++++++------ tests/e2e/PROVIDER_CACHE.md | 2 ++ tests/e2e/provider_cache_routing.py | 23 +++++++++++++++- 3 files changed, 43 insertions(+), 9 deletions(-) diff --git a/tests/code_coverage_tests/test_provider_cache.py b/tests/code_coverage_tests/test_provider_cache.py index 5c491e1cdcd..0278bf8fd48 100644 --- a/tests/code_coverage_tests/test_provider_cache.py +++ b/tests/code_coverage_tests/test_provider_cache.py @@ -766,13 +766,20 @@ def test_registration_preserves_unsupported_or_explicit_routes(params: LiteLLMPa assert route_cache_model(params, unexpected_edge, enabled=True) is params -@pytest.mark.parametrize("model", [ - "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", - "bedrock/converse/us.anthropic.claude-sonnet-5", - "bedrock/invoke/us.anthropic.claude-haiku-4-5-20251001-v1:0", +@pytest.mark.parametrize("model,region", [ + ("bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", None), + ("bedrock/converse/us.anthropic.claude-sonnet-5", None), + ("bedrock/invoke/us.anthropic.claude-haiku-4-5-20251001-v1:0", None), + ("bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", "us-east-1"), + ("bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", "os.environ/AWS_REGION"), + ("bedrock/invoke/us.anthropic.claude-sonnet-5", "os.environ/AWS_REGION"), ]) -def test_anthropic_on_bedrock_registers_the_edge_as_its_runtime_endpoint(model: str) -> None: - params: Final = LiteLLMParamsBody(model=model) +def test_anthropic_on_bedrock_registers_the_edge_as_its_runtime_endpoint(model: str, region: str | None) -> None: + """Almost every Bedrock deployment in the suite declares its region as + `os.environ/AWS_REGION`, which only the proxy can resolve. Treating that + string as a region name would leave the whole Anthropic-on-Bedrock surface + off the edge, which is the point of mounting it at all.""" + params: Final = LiteLLMParamsBody(model=model, aws_region_name=region) routed: Final = route_cache_model(params, lambda mount: f"http://edge.invalid/{mount}", enabled=True) assert routed.aws_bedrock_runtime_endpoint == "http://edge.invalid/bedrock/us-east-1" assert routed.api_base is None @@ -794,15 +801,19 @@ def test_anthropic_on_bedrock_registers_the_edge_as_its_runtime_endpoint(model: aws_bedrock_runtime_endpoint="https://custom.invalid", ), LiteLLMParamsBody(model="bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", aws_region_name="eu-west-1"), + LiteLLMParamsBody(model="bedrock/anthropic.claude-sonnet-5", aws_region_name="os.environ/AWS_REGION"), + LiteLLMParamsBody(model="bedrock/invoke/eu.anthropic.claude-sonnet-5", aws_region_name="os.environ/AWS_REGION"), ]) def test_bedrock_deployments_the_edge_must_not_touch_keep_their_direct_route(params: LiteLLMParamsBody) -> None: """Non-Anthropic models the runner role cannot invoke, deployments carrying their own AWS identity (routing those would replace the assume-role chain the - batch suite exists to prove), explicit endpoints, and unmounted regions.""" + batch suite exists to prove), explicit endpoints, unmounted regions, and a + region only the proxy can resolve on a model that is not cross-region, whose + real region the harness cannot know.""" routed: Final = route_cache_model( params, lambda mount: None if mount not in EDGE_MOUNTS else f"http://edge.invalid/{mount}", enabled=True, ) - assert routed is params or routed.aws_bedrock_runtime_endpoint == params.aws_bedrock_runtime_endpoint + assert routed is params @pytest.mark.parametrize("mode", ["batch", "realtime", "image_generation"]) diff --git a/tests/e2e/PROVIDER_CACHE.md b/tests/e2e/PROVIDER_CACHE.md index 393a96e8c16..c530574cda2 100644 --- a/tests/e2e/PROVIDER_CACHE.md +++ b/tests/e2e/PROVIDER_CACHE.md @@ -20,6 +20,8 @@ An eligible miss calls the provider. A complete successful response is stored im Bedrock could not be mounted before because SigV4 signs the `Host` header, so a rewritten `api_base` failed signature verification at the provider. The edge now re-signs: it drops the proxy's signature headers, signs the upstream request with the run pod's own AWS identity from its EKS Pod Identity association, and forwards that. The signature headers are excluded from the key, since `x-amz-date` is a timestamp and keying on it would make every Bedrock call a permanent miss +Almost every Bedrock deployment in the suite declares its region as `os.environ/AWS_REGION`, which only the proxy can resolve, and the run pod does not share that environment. A `us.` inference profile fans out across the US regions and is reachable from any of them, so those route to the default mount whatever the proxy resolved. A model that is not cross-region and declares its region that way keeps its direct path rather than being sent to a region it may not exist in. + Only deployments that carry no AWS identity of their own route to the edge. A deployment with `aws_role_name`, `aws_access_key_id`, an `api_base` or an `aws_bedrock_runtime_endpoint` keeps its direct path, because re-signing it would quietly replace the very credential chain that test exists to prove. Only Anthropic models route, matching what the runner role is allowed to invoke and what the edge knows how to validate Vertex and Gemini are not mounted. litellm's `_check_custom_proxy` rewrites a path-prefixed Vertex `api_base` into `{api_base}:{endpoint}`, dropping project, location and model, so a mount under a path prefix cannot work without either a root-mounted edge on its own port or a change in litellm diff --git a/tests/e2e/provider_cache_routing.py b/tests/e2e/provider_cache_routing.py index e7e4899eb71..d2237bb49bc 100644 --- a/tests/e2e/provider_cache_routing.py +++ b/tests/e2e/provider_cache_routing.py @@ -10,6 +10,26 @@ LIVE_PROVIDER_REQUIRED: Final[ContextVar[bool]] = ContextVar("live_provider_requ DEFAULT_BEDROCK_REGION: Final = "us-east-1" BEDROCK_ANTHROPIC_INFIX: Final = "anthropic." +BEDROCK_CROSS_REGION_PREFIX: Final = "us." +ENV_REFERENCE_PREFIX: Final = "os.environ/" + + +def bedrock_region(declared: str | None, model: str) -> str | None: + """The region whose edge mount a deployment belongs to, or None when the + harness cannot know it. + + Most Bedrock deployments declare `os.environ/AWS_REGION`, which the proxy + resolves from its own environment. The run pod does not share that + environment, so the harness genuinely does not know the region. A `us.` + inference profile fans out across the US regions and is reachable from any + of them, so the default entry point is correct for those whatever the proxy + resolved; anything else keeps its direct path rather than being sent to a + region the model may not exist in.""" + if declared is None: + return DEFAULT_BEDROCK_REGION + if not declared.startswith(ENV_REFERENCE_PREFIX): + return declared + return DEFAULT_BEDROCK_REGION if model.startswith(BEDROCK_CROSS_REGION_PREFIX) else None def bedrock_mount(params: LiteLLMParamsBody) -> str | None: @@ -24,7 +44,8 @@ def bedrock_mount(params: LiteLLMParamsBody) -> str | None: model: Final = route.partition("/")[2] or route if BEDROCK_ANTHROPIC_INFIX not in model: return None - return f"bedrock/{params.aws_region_name or DEFAULT_BEDROCK_REGION}" + region: Final = bedrock_region(params.aws_region_name, model) + return None if region is None else f"bedrock/{region}" def route_bedrock( From 30c6241e3a5f534037dc57a6ea544d5ff9d8bdeb Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 16 Sep 2026 03:01:07 -0700 Subject: [PATCH 04/21] fix(e2e): a null error field is not an error Every OpenAI Responses body carries `error: null` at the top level, and the completeness check tested the key's presence rather than its value, so it rejected every single one. The cost was silent: nothing failed, the endpoint simply never cached, which is exactly the outcome the endpoint was added for. Found by driving the edge against the real providers rather than the synthetic fixtures, which carried no error key at all. Reading the value instead of the key is also more accurate for chat completions and messages, where a real error body carries a populated error object. --- .../test_provider_cache.py | 42 +++++++++++++++++-- tests/e2e/provider_cache.py | 7 +++- 2 files changed, 44 insertions(+), 5 deletions(-) diff --git a/tests/code_coverage_tests/test_provider_cache.py b/tests/code_coverage_tests/test_provider_cache.py index 0278bf8fd48..3dafbadf508 100644 --- a/tests/code_coverage_tests/test_provider_cache.py +++ b/tests/code_coverage_tests/test_provider_cache.py @@ -508,10 +508,13 @@ EMBEDDING_SUCCESS: Final = ( b'{"object":"list","data":[{"object":"embedding","index":0,"embedding":[0.1,0.2]}],' b'"model":"text-embedding-3-small","usage":{"prompt_tokens":2,"total_tokens":2}}' ) -RESPONSE_SUCCESS: Final = b'{"id":"resp_synthetic","object":"response","status":"completed","output":[]}' +RESPONSE_SUCCESS: Final = ( + b'{"id":"resp_synthetic","object":"response","status":"completed","error":null,' + b'"incomplete_details":null,"output":[]}' +) RESPONSE_STREAM_SUCCESS: Final = ( - b'data: {"type":"response.created","response":{"id":"resp_synthetic"}}\n\n' - b'data: {"type":"response.completed","response":{"id":"resp_synthetic","status":"completed"}}\n\n' + b'data: {"type":"response.created","response":{"id":"resp_synthetic","error":null}}\n\n' + b'data: {"type":"response.completed","response":{"id":"resp_synthetic","status":"completed"},"error":null}\n\n' ) @@ -586,6 +589,39 @@ class TestNonChatOpenAiEndpoints: assert call(url, MARKED).body == payload assert len(provider.hits) == 2 + @pytest.mark.parametrize("path,response", [ + ("/v1/chat/completions", b'{"id":"x","error":null,"choices":[{"message":{"content":"hi"},' + b'"finish_reason":"stop"}]}'), + ("/v1/messages", b'{"id":"msg_x","type":"message","role":"assistant","error":null,' + b'"content":[{"type":"text","text":"hi"}],"stop_reason":"end_turn"}'), + ("/v1/responses", RESPONSE_SUCCESS), + ]) + def test_a_null_error_field_is_not_an_error( + self, store: RedisResponseStore, provider: Provider, path: str, response: bytes, + ) -> None: + """Every OpenAI Responses body carries `error: null`, and testing the key's + presence rather than its value rejected all of them. The cost was silent: + nothing failed, the endpoint simply never cached.""" + assert b'"error":null' in response + provider.response = response + for _ in range(2): + with openai_edge(cache_edge(store), provider, path) as url: + assert call(url, MARKED).body == response + assert len(provider.hits) == 1 + + @pytest.mark.parametrize("path,response", [ + ("/v1/chat/completions", b'{"error":{"message":"rate limited","type":"rate_limit_error"}}'), + ("/v1/responses", b'{"object":"response","status":"completed","error":{"message":"bad"},"output":[]}'), + ]) + def test_a_populated_error_field_still_rejects( + self, store: RedisResponseStore, provider: Provider, path: str, response: bytes, + ) -> None: + provider.response = response + for _ in range(2): + with openai_edge(cache_edge(store), provider, path) as url: + assert call(url, MARKED).body == response + assert len(provider.hits) == 2 + @pytest.mark.parametrize("path,cacheable", [ ("/v1/chat/completions", True), ("/v1/messages", True), ("/v1/embeddings", True), ("/v1/responses", True), diff --git a/tests/e2e/provider_cache.py b/tests/e2e/provider_cache.py index a2b18a3a466..0dee33f33c6 100644 --- a/tests/e2e/provider_cache.py +++ b/tests/e2e/provider_cache.py @@ -167,7 +167,10 @@ def successful_response(mount: str, url: str, status: int, headers: Mapping[str, values: Final = tuple(JSON_VALUE.validate_json(event) for event in events if event != "[DONE]") except (UnicodeDecodeError, ValidationError): return False - if not values or any(not isinstance(value, dict) or "error" in value or value.get("type") == "error" for value in values): + if not values or any( + not isinstance(value, dict) or value.get("error") is not None or value.get("type") == "error" + for value in values + ): return False if urlsplit(url).path == "/v1/responses": return complete_responses_stream(values) @@ -187,7 +190,7 @@ def successful_response(mount: str, url: str, status: int, headers: Mapping[str, value: Final = JSON_VALUE.validate_json(body) except ValidationError: return False - if not isinstance(value, dict) or "error" in value: + if not isinstance(value, dict) or value.get("error") is not None: return False path: Final = urlsplit(url).path if path == "/v1/messages": From c7246adc1dd571069a2cb3c77b3a40b720eb0da6 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 16 Sep 2026 03:29:02 -0700 Subject: [PATCH 05/21] fix(e2e): route only the Bedrock models the runner role can invoke The edge re-signs with the run pod's identity, whose IAM policy is an explicit per-model allowlist. Matching on the `anthropic.` infix instead routed every Anthropic-on-Bedrock model, so a model outside the policy came back 403 from Bedrock with no fallback, taking the whole claude_code Bedrock matrix red. An unlisted model now keeps its direct path and loses only caching. --- .../test_provider_cache.py | 4 ++++ tests/e2e/PROVIDER_CACHE.md | 4 +++- tests/e2e/provider_cache_routing.py | 23 ++++++++++++------- 3 files changed, 22 insertions(+), 9 deletions(-) diff --git a/tests/code_coverage_tests/test_provider_cache.py b/tests/code_coverage_tests/test_provider_cache.py index 3dafbadf508..ac1ccd5692b 100644 --- a/tests/code_coverage_tests/test_provider_cache.py +++ b/tests/code_coverage_tests/test_provider_cache.py @@ -809,6 +809,8 @@ def test_registration_preserves_unsupported_or_explicit_routes(params: LiteLLMPa ("bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", "us-east-1"), ("bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", "os.environ/AWS_REGION"), ("bedrock/invoke/us.anthropic.claude-sonnet-5", "os.environ/AWS_REGION"), + ("bedrock/us.anthropic.claude-opus-4-7", "us-east-1"), + ("bedrock/converse/us.anthropic.claude-opus-4-7", "us-east-1"), ]) def test_anthropic_on_bedrock_registers_the_edge_as_its_runtime_endpoint(model: str, region: str | None) -> None: """Almost every Bedrock deployment in the suite declares its region as @@ -839,6 +841,8 @@ def test_anthropic_on_bedrock_registers_the_edge_as_its_runtime_endpoint(model: LiteLLMParamsBody(model="bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", aws_region_name="eu-west-1"), LiteLLMParamsBody(model="bedrock/anthropic.claude-sonnet-5", aws_region_name="os.environ/AWS_REGION"), LiteLLMParamsBody(model="bedrock/invoke/eu.anthropic.claude-sonnet-5", aws_region_name="os.environ/AWS_REGION"), + LiteLLMParamsBody(model="bedrock/us.anthropic.claude-opus-4-5", aws_region_name="us-east-1"), + LiteLLMParamsBody(model="bedrock/converse/us.anthropic.claude-haiku-9-9", aws_region_name="us-east-1"), ]) def test_bedrock_deployments_the_edge_must_not_touch_keep_their_direct_route(params: LiteLLMParamsBody) -> None: """Non-Anthropic models the runner role cannot invoke, deployments carrying diff --git a/tests/e2e/PROVIDER_CACHE.md b/tests/e2e/PROVIDER_CACHE.md index c530574cda2..3a29c7b6fe4 100644 --- a/tests/e2e/PROVIDER_CACHE.md +++ b/tests/e2e/PROVIDER_CACHE.md @@ -22,7 +22,9 @@ Bedrock could not be mounted before because SigV4 signs the `Host` header, so a Almost every Bedrock deployment in the suite declares its region as `os.environ/AWS_REGION`, which only the proxy can resolve, and the run pod does not share that environment. A `us.` inference profile fans out across the US regions and is reachable from any of them, so those route to the default mount whatever the proxy resolved. A model that is not cross-region and declares its region that way keeps its direct path rather than being sent to a region it may not exist in. -Only deployments that carry no AWS identity of their own route to the edge. A deployment with `aws_role_name`, `aws_access_key_id`, an `api_base` or an `aws_bedrock_runtime_endpoint` keeps its direct path, because re-signing it would quietly replace the very credential chain that test exists to prove. Only Anthropic models route, matching what the runner role is allowed to invoke and what the edge knows how to validate +Only deployments that carry no AWS identity of their own route to the edge. A deployment with `aws_role_name`, `aws_access_key_id`, an `api_base` or an `aws_bedrock_runtime_endpoint` keeps its direct path, because re-signing it would quietly replace the very credential chain that test exists to prove + +Which models route is an explicit allowlist in `provider_cache_routing.py`, mirroring the runner role's IAM policy, which names its models one by one. That coupling is deliberate: the edge re-signs with the run pod's identity, so a model the role cannot invoke comes back 403 from Bedrock rather than falling back. An unlisted model keeps its direct path and loses only caching, so adding a Bedrock model to the suite can never turn it red. Adding one to the edge is a policy edit in litellm-ops plus a line here Vertex and Gemini are not mounted. litellm's `_check_custom_proxy` rewrites a path-prefixed Vertex `api_base` into `{api_base}:{endpoint}`, dropping project, location and model, so a mount under a path prefix cannot work without either a root-mounted edge on its own port or a change in litellm diff --git a/tests/e2e/provider_cache_routing.py b/tests/e2e/provider_cache_routing.py index d2237bb49bc..97e05344423 100644 --- a/tests/e2e/provider_cache_routing.py +++ b/tests/e2e/provider_cache_routing.py @@ -9,8 +9,15 @@ from models import LiteLLMParamsBody, ModelMode LIVE_PROVIDER_REQUIRED: Final[ContextVar[bool]] = ContextVar("live_provider_required", default=False) DEFAULT_BEDROCK_REGION: Final = "us-east-1" -BEDROCK_ANTHROPIC_INFIX: Final = "anthropic." BEDROCK_CROSS_REGION_PREFIX: Final = "us." +BEDROCK_EDGE_MODELS: Final = frozenset( + { + "us.anthropic.claude-haiku-4-5-20251001-v1:0", + "us.anthropic.claude-sonnet-4-5-20250929-v1:0", + "us.anthropic.claude-sonnet-5", + "us.anthropic.claude-opus-4-7", + } +) ENV_REFERENCE_PREFIX: Final = "os.environ/" @@ -33,16 +40,16 @@ def bedrock_region(declared: str | None, model: str) -> str | None: def bedrock_mount(params: LiteLLMParamsBody) -> str | None: - """The edge mount an Anthropic-on-Bedrock deployment belongs to, or None. + """The edge mount a Bedrock deployment belongs to, or None. - Only the Anthropic models route. The edge validates converse and invoke - bodies by their Anthropic and Converse terminator fields, and the runner role - is allowed to invoke exactly those models, so Bedrock embeddings, image - generation, rerank and realtime keep their existing direct path rather than - reaching an edge that could neither sign nor validate for them.""" + The allowlist mirrors the runner role's IAM policy, which names its models + one by one. A model outside it would be re-signed with an identity that + cannot invoke it and come back 403 from Bedrock, so an unlisted model keeps + its direct path and loses only caching. Adding a model is a policy edit in + litellm-ops and a line here.""" route: Final = params.model.partition("/")[2] model: Final = route.partition("/")[2] or route - if BEDROCK_ANTHROPIC_INFIX not in model: + if model not in BEDROCK_EDGE_MODELS: return None region: Final = bedrock_region(params.aws_region_name, model) return None if region is None else f"bedrock/{region}" From ebf34cd88006110c78a50648578e3dc7e0f115cd Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 16 Sep 2026 03:35:00 -0700 Subject: [PATCH 06/21] docs(e2e): say plainly that Bedrock streaming is not cached --- tests/e2e/PROVIDER_CACHE.md | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/e2e/PROVIDER_CACHE.md b/tests/e2e/PROVIDER_CACHE.md index 3a29c7b6fe4..fe289e406aa 100644 --- a/tests/e2e/PROVIDER_CACHE.md +++ b/tests/e2e/PROVIDER_CACHE.md @@ -2,7 +2,9 @@ `E2E_PROVIDER_CACHE=1` enables automatic response reuse in the live E2E mode. Standard OpenAI and Anthropic model registrations use the provider edge, as do Anthropic-on-Bedrock registrations that carry no AWS identity of their own. 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`, `/v1/messages`, `/v1/embeddings` and `/v1/responses` on the OpenAI and Anthropic mounts, and for `/model/{id}/converse` and `/model/{id}/invoke` on a Bedrock mount, including streams. Unsupported endpoints pass through. Each endpoint family has its own completeness rule, so a truncated embedding or a Responses run that never reached `response.completed` is not stored +The edge caches complete successful POST responses for `/v1/chat/completions`, `/v1/messages`, `/v1/embeddings` and `/v1/responses` on the OpenAI and Anthropic mounts, SSE streams included, and for `/model/{id}/converse` and `/model/{id}/invoke` on a Bedrock mount. Unsupported endpoints pass through. Each endpoint family has its own completeness rule, so a truncated embedding or a Responses run that never reached `response.completed` is not stored + +Bedrock's streaming endpoints, `converse-stream` and `invoke-with-response-stream`, are not cacheable. They still cross the edge and are still re-signed, so they need the same IAM, but they always call the provider. AWS frames them as binary `vnd.amazon.eventstream` rather than SSE, and reading a terminal event out of that is what a completeness rule for them would need. That matters more than the endpoint count suggests: the Claude Code compat cells drive the real CLI, which always streams, so most Bedrock traffic in the suite is not cached today ## Request identity From 7c2234be3a91a600c02f416e678b163dd53746ac Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 16 Sep 2026 04:41:29 -0700 Subject: [PATCH 07/21] test(e2e): enforce the cross-region invariant on the Bedrock allowlist The allowlist rejects an unlisted model before the region resolver runs, so the two negative cases that used to cover the resolver were passing for the wrong reason and two mutations of it survived. Answering an env-referenced region with the default mount is only sound because every allowlisted model is a `us.` profile that fans out across the US regions, so assert that on the list itself and drop the per-call branch it made unreachable. --- .../test_provider_cache.py | 33 ++++++++++++++++++- tests/e2e/provider_cache_routing.py | 27 +++++++-------- 2 files changed, 43 insertions(+), 17 deletions(-) diff --git a/tests/code_coverage_tests/test_provider_cache.py b/tests/code_coverage_tests/test_provider_cache.py index ac1ccd5692b..c24ba8d6221 100644 --- a/tests/code_coverage_tests/test_provider_cache.py +++ b/tests/code_coverage_tests/test_provider_cache.py @@ -33,7 +33,13 @@ from provider_cache import ( successful_response, ) from provider_cache_redis import PUBLISH, RedisCommands, RedisResponseStore, configured_cache, redis_store -from provider_cache_routing import LIVE_PROVIDER_REQUIRED, route_cache_model +from provider_cache_routing import ( + BEDROCK_CROSS_REGION_PREFIX, + BEDROCK_EDGE_MODELS, + LIVE_PROVIDER_REQUIRED, + bedrock_region, + route_cache_model, +) from fixture_mode import SESSION_TEST_KEY from provider_edge import EDGE_MOUNTS, configured_cache_backend, resolve_mount, start_provider_edge from provider_edge_bedrock import bedrock_signer @@ -856,6 +862,31 @@ def test_bedrock_deployments_the_edge_must_not_touch_keep_their_direct_route(par assert routed is params +@pytest.mark.parametrize("declared,expected", [ + (None, "us-east-1"), + ("us-west-2", "us-west-2"), + ("eu-west-1", "eu-west-1"), + ("os.environ/AWS_REGION", "us-east-1"), + ("os.environ/ANY_OTHER_NAME", "us-east-1"), +]) +def test_a_region_only_the_proxy_can_resolve_falls_back_to_the_default_mount( + declared: str | None, expected: str, +) -> None: + """A declared literal region is the one the deployment meant. A region the + proxy resolves from its own environment is one the run pod cannot see, and + the default mount answers it.""" + assert bedrock_region(declared) == expected + + +def test_every_model_on_the_edge_allowlist_is_a_cross_region_profile() -> None: + """Answering an env-referenced region with the default mount is only correct + for a profile that fans out across the US regions and is reachable from any + of them. A single-region model on this list would be sent to a region it may + not exist in, so the list is where that is caught.""" + assert BEDROCK_EDGE_MODELS + assert all(model.startswith(BEDROCK_CROSS_REGION_PREFIX) for model in BEDROCK_EDGE_MODELS) + + @pytest.mark.parametrize("mode", ["batch", "realtime", "image_generation"]) def test_a_bedrock_deployment_with_a_mode_keeps_its_direct_route(mode: ModelMode) -> None: params: Final = LiteLLMParamsBody(model="bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0") diff --git a/tests/e2e/provider_cache_routing.py b/tests/e2e/provider_cache_routing.py index 97e05344423..f9775a2b152 100644 --- a/tests/e2e/provider_cache_routing.py +++ b/tests/e2e/provider_cache_routing.py @@ -21,22 +21,18 @@ BEDROCK_EDGE_MODELS: Final = frozenset( ENV_REFERENCE_PREFIX: Final = "os.environ/" -def bedrock_region(declared: str | None, model: str) -> str | None: - """The region whose edge mount a deployment belongs to, or None when the - harness cannot know it. +def bedrock_region(declared: str | None) -> str: + """The region whose edge mount a deployment belongs to. - Most Bedrock deployments declare `os.environ/AWS_REGION`, which the proxy - resolves from its own environment. The run pod does not share that - environment, so the harness genuinely does not know the region. A `us.` - inference profile fans out across the US regions and is reachable from any - of them, so the default entry point is correct for those whatever the proxy - resolved; anything else keeps its direct path rather than being sent to a - region the model may not exist in.""" - if declared is None: + Most Bedrock deployments declare `os.environ/AWS_REGION`, which only the + proxy can resolve from its own environment; the run pod does not share it. + Answering those with the default mount is correct because every model on the + edge allowlist is a `us.` inference profile, which fans out across the US + regions and is reachable from any of them. That invariant is enforced on the + allowlist itself rather than re-checked per call.""" + if declared is None or declared.startswith(ENV_REFERENCE_PREFIX): return DEFAULT_BEDROCK_REGION - if not declared.startswith(ENV_REFERENCE_PREFIX): - return declared - return DEFAULT_BEDROCK_REGION if model.startswith(BEDROCK_CROSS_REGION_PREFIX) else None + return declared def bedrock_mount(params: LiteLLMParamsBody) -> str | None: @@ -51,8 +47,7 @@ def bedrock_mount(params: LiteLLMParamsBody) -> str | None: model: Final = route.partition("/")[2] or route if model not in BEDROCK_EDGE_MODELS: return None - region: Final = bedrock_region(params.aws_region_name, model) - return None if region is None else f"bedrock/{region}" + return f"bedrock/{bedrock_region(params.aws_region_name)}" def route_bedrock( From bd1c2d6f07d7b9edd46bb11d95ad22cb7e029ea9 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 16 Sep 2026 05:09:04 -0700 Subject: [PATCH 08/21] fix(e2e): keep the tool-continuation echo-back test on the live path The key normalizes a unique marker so two builds match, which is the whole point, but it makes this test's identity collide with an earlier run's: it mints a fresh receipt, sends it through a tool result, and asserts the model echoes it back verbatim, so a stale recording matched and answered with the old receipt. Build 223 is where that surfaced, once the corpus was full enough for the first call to hit. A test that asserts a provider echoed this run's own unique value belongs on the live path. --- tests/e2e/PROVIDER_CACHE.md | 4 +++- tests/e2e/conftest.py | 6 +++++- tests/e2e/llm_translation/test_messages_e2e.py | 1 + 3 files changed, 9 insertions(+), 2 deletions(-) diff --git a/tests/e2e/PROVIDER_CACHE.md b/tests/e2e/PROVIDER_CACHE.md index fe289e406aa..698e78a5aa1 100644 --- a/tests/e2e/PROVIDER_CACHE.md +++ b/tests/e2e/PROVIDER_CACHE.md @@ -44,7 +44,9 @@ 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. Every counter is emitted twice, once as a flat total and once under `mount:{mount}:`, so a hit rate can be read per provider rather than only in aggregate. 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. 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 +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. + +One more class needs it, and it is the cost of normalizing the marker. A test that mints a fresh marker, sends it, and then asserts the provider's answer contains that exact value is asserting on the marker rather than using it as a salt. The key treats two such requests as the same identity, so a stale recording matches and answers with the marker from the run that recorded it. `TestOpenAIMessagesToolContinuation` is the one in the suite today: it sends a freshly minted receipt through a tool result and asserts the model echoes it back verbatim. If you add a test that asserts a provider echoed your own unique value, it belongs on the live path. 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 diff --git a/tests/e2e/conftest.py b/tests/e2e/conftest.py index 829c84910a9..430e16525d5 100644 --- a/tests/e2e/conftest.py +++ b/tests/e2e/conftest.py @@ -85,7 +85,11 @@ def jwt_identity(idp: Keycloak, resources: ResourceManager, proxy: ProxyClient) def pytest_configure(config: pytest.Config) -> None: - config.addinivalue_line("markers", "provider_live: requires actual provider timing, limits or state; bypass shared cache") + config.addinivalue_line( + "markers", + "provider_live: requires actual provider timing, limits, state, or a response that echoes this" + " run's own unique value; bypass shared cache", + ) config.addinivalue_line( "markers", "e2e: live test that requires a running proxy and real provider keys", diff --git a/tests/e2e/llm_translation/test_messages_e2e.py b/tests/e2e/llm_translation/test_messages_e2e.py index 44c416a3e78..09ec48daa2f 100644 --- a/tests/e2e/llm_translation/test_messages_e2e.py +++ b/tests/e2e/llm_translation/test_messages_e2e.py @@ -372,6 +372,7 @@ def _request_tool( class TestOpenAIMessagesToolContinuation: + @pytest.mark.provider_live @pytest.mark.parametrize("stream", [True, False], ids=["stream", "nonstream"]) def test_required_tool_arguments_and_correlated_result( self, endpoints_client: EndpointsClient, resources: ResourceManager, stream: bool From 30a691ed55c59d8fd19c28e5356ba196bf5ae45a Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 16 Sep 2026 06:46:19 -0700 Subject: [PATCH 09/21] feat(e2e): cache Bedrock streaming responses The Claude Code compat cells drive the real CLI, which always streams, so converse-stream and invoke-with-response-stream were most of the suite's Bedrock traffic and all of it bypassed the edge. AWS frames those as binary vnd.amazon.eventstream rather than SSE, so botocore's own parser reads the frames and validates both CRCs, and each endpoint is then held to its terminal grammar. Two details drove the rule. A ConverseStream ends with metadata, not with messageStop, and metadata is what carries the token usage litellm prices the call from, so a stream cut between the two names a stop reason but would replay as a free call. And a dropped connection is invisible to the parser: it yields the frames it did receive and silently discards a trailing partial one, so a stream cut one byte short parses clean. The body is checked against the frame lengths it declares to catch that. The invoke stream carries the ordinary Anthropic event grammar inside its chunk frames, so it shares the completeness rule with the SSE mounts. Validated against three real Bedrock eventstream captures, and the tests build their own frames rather than pasting a capture, with one test holding that framing to botocore's parser. --- .../test_provider_cache.py | 258 ++++++++++++++++-- tests/e2e/PROVIDER_CACHE.md | 4 +- tests/e2e/provider_cache.py | 141 +++++++++- 3 files changed, 370 insertions(+), 33 deletions(-) diff --git a/tests/code_coverage_tests/test_provider_cache.py b/tests/code_coverage_tests/test_provider_cache.py index c24ba8d6221..4c131434ecd 100644 --- a/tests/code_coverage_tests/test_provider_cache.py +++ b/tests/code_coverage_tests/test_provider_cache.py @@ -1,9 +1,13 @@ from __future__ import annotations +import base64 +import binascii +import json import os import shutil import socket import subprocess +import struct import threading import time import uuid @@ -17,9 +21,11 @@ from typing import Final from urllib.parse import urlsplit import pytest +from pydantic import JsonValue from e2e_http import NetworkError, PreparedForward, RawResponse, StreamChunk, StreamHead, forward, prepare_forward from models import LiteLLMParamsBody, ModelMode from botocore.credentials import Credentials +from botocore.eventstream import EventStreamBuffer from provider_cache import ( SIGNATURE_HEADERS, CacheEdge, @@ -638,6 +644,53 @@ class TestNonChatOpenAiEndpoints: assert cacheable_endpoint("openai", "POST", f"https://api.openai.com{path}", MARKED) is cacheable +BEDROCK_STREAM_MODEL: Final = "us.anthropic.claude-haiku-4-5-20251001-v1:0" +CONVERSE_STREAM_URL: Final = f"https://bedrock.invalid/model/{BEDROCK_STREAM_MODEL}/converse-stream" +INVOKE_STREAM_URL: Final = f"https://bedrock.invalid/model/{BEDROCK_STREAM_MODEL}/invoke-with-response-stream" + + +def eventstream_frame(headers: Mapping[str, str], payload: bytes) -> bytes: + """AWS eventstream wire framing, the shape `vnd.amazon.eventstream` bodies + arrive in. Built here rather than pasted from a capture so a test can express + the stream it means; `test_the_frames_these_tests_build_are_real_aws_framing` + holds it to botocore's own parser.""" + encoded: Final = b"".join( + bytes([len(name)]) + name.encode() + b"\x07" + struct.pack(">H", len(value)) + value.encode() + for name, value in headers.items() + ) + prelude: Final = struct.pack(">II", 16 + len(encoded) + len(payload), len(encoded)) + framed: Final = prelude + struct.pack(">I", binascii.crc32(prelude)) + encoded + payload + return framed + struct.pack(">I", binascii.crc32(framed)) + + +def eventstream_event(event_type: str, payload: JsonValue, message_type: str = "event") -> bytes: + return eventstream_frame( + {":event-type": event_type, ":message-type": message_type, ":content-type": "application/json"}, + json.dumps(payload).encode(), + ) + + +def invoke_chunk(inner: JsonValue) -> bytes: + return eventstream_event("chunk", {"bytes": base64.b64encode(json.dumps(inner).encode()).decode("ascii")}) + + +CONVERSE_STREAM_OK: Final = ( + eventstream_event("messageStart", {"role": "assistant"}) + + eventstream_event("contentBlockDelta", {"contentBlockIndex": 0, "delta": {"text": "hi"}}) + + eventstream_event("contentBlockStop", {"contentBlockIndex": 0}) + + eventstream_event("messageStop", {"stopReason": "end_turn"}) + + eventstream_event("metadata", {"usage": {"inputTokens": 12, "outputTokens": 6, "totalTokens": 18}}) +) +INVOKE_STREAM_OK: Final = ( + invoke_chunk({"type": "message_start", "message": {"id": "msg_bdrk_x", "role": "assistant"}}) + + invoke_chunk({"type": "content_block_start", "index": 0}) + + invoke_chunk({"type": "content_block_delta", "index": 0, "delta": {"text": "hi"}}) + + invoke_chunk({"type": "content_block_stop", "index": 0}) + + invoke_chunk({"type": "message_delta", "delta": {"stop_reason": "end_turn"}}) + + invoke_chunk({"type": "message_stop"}) +) + + class TestBedrockSigning: """Bedrock is the reason the edge could not mount it before: SigV4 covers the Host header, so forwarding through a rewritten api_base invalidates the @@ -738,32 +791,55 @@ class TestBedrockSigning: assert call(url, BEDROCK_BODY).body == response assert len(provider.hits) == 2 - @pytest.mark.parametrize("action", ["converse-stream", "invoke-with-response-stream"]) - def test_streaming_endpoints_go_live_every_time( - self, store: RedisResponseStore, provider: Provider, action: str, + @pytest.mark.parametrize("action,response", [ + ("converse-stream", CONVERSE_STREAM_OK), + ("invoke-with-response-stream", INVOKE_STREAM_OK), + ], ids=["converse-stream", "invoke-stream"]) + def test_a_finished_stream_is_served_from_the_cache_the_second_time( + self, store: RedisResponseStore, provider: Provider, action: str, response: bytes, ) -> None: - """An eventstream's completeness cannot be proven without parsing its - frames, so these bypass rather than risk recording a truncated answer. - They are still signed: a bypass is a forward, not a passthrough.""" - provider.response = CONVERSE_SUCCESS - cache: Final = bedrock_cache_edge(store) - for _ in range(2): - with bedrock_edge(cache, provider, action) as url: - assert call(url, BEDROCK_BODY).body == CONVERSE_SUCCESS - assert len(provider.hits) == 2 - assert dict(cache.counters.counts)[f"mount:{BEDROCK_MOUNT}:bypass"] == 2 + provider.response = response + with bedrock_edge(bedrock_cache_edge(store), provider, action) as url: + assert call(url, BEDROCK_BODY).body == response + assert len(provider.hits) == 1 + replay: Final = bedrock_cache_edge(store) + with bedrock_edge(replay, provider, action) as url: + assert call(url, BEDROCK_BODY).body == response + assert len(provider.hits) == 1 + assert dict(replay.counters.counts)[f"mount:{BEDROCK_MOUNT}:hits"] == 1 assert all( sent.startswith("AWS4-HMAC-SHA256 Credential=AKIAIOSFODNN7EXAMPLE/") for sent in provider.authorizations ), provider.authorizations - @pytest.mark.parametrize("action,cacheable", [ - ("converse", True), ("invoke", True), - ("converse-stream", False), ("invoke-with-response-stream", False), - ]) - def test_only_the_unary_bedrock_actions_are_cacheable(self, action: str, cacheable: bool) -> None: + @pytest.mark.parametrize("action,response", [ + ("converse-stream", CONVERSE_STREAM_OK[:-1]), + ("invoke-with-response-stream", INVOKE_STREAM_OK[:-1]), + ], ids=["converse-stream", "invoke-stream"]) + def test_a_stream_the_connection_cut_short_calls_the_provider_every_time( + self, store: RedisResponseStore, provider: Provider, action: str, response: bytes, + ) -> None: + """The whole risk of caching an eventstream is recording a half-finished + one, so a truncated body has to be rejected rather than stored.""" + provider.response = response + with bedrock_edge(bedrock_cache_edge(store), provider, action) as url: + assert call(url, BEDROCK_BODY).body == response + replay: Final = bedrock_cache_edge(store) + with bedrock_edge(replay, provider, action) as url: + assert call(url, BEDROCK_BODY).body == response + assert len(provider.hits) == 2 + assert dict(replay.counters.counts)[f"mount:{BEDROCK_MOUNT}:rejected"] == 1 + assert f"mount:{BEDROCK_MOUNT}:hits" not in dict(replay.counters.counts) + + @pytest.mark.parametrize("action", ["converse", "invoke", "converse-stream", "invoke-with-response-stream"]) + def test_every_anthropic_bedrock_action_is_cacheable(self, action: str) -> None: url: Final = f"https://bedrock-runtime.us-east-1.amazonaws.com/model/{BEDROCK_MODEL}/{action}" - assert cacheable_endpoint(BEDROCK_MOUNT, "POST", url, BEDROCK_BODY) is cacheable + assert cacheable_endpoint(BEDROCK_MOUNT, "POST", url, BEDROCK_BODY) + + @pytest.mark.parametrize("action", ["count-tokens", "invoke-async", "converse-stream-x"]) + def test_an_unknown_bedrock_action_is_not_cacheable(self, action: str) -> None: + url: Final = f"https://bedrock-runtime.us-east-1.amazonaws.com/model/{BEDROCK_MODEL}/{action}" + assert not cacheable_endpoint(BEDROCK_MOUNT, "POST", url, BEDROCK_BODY) def test_a_region_mount_resolves_whole(self) -> None: resolved: Final = resolve_mount(f"/{BEDROCK_MOUNT}/model/{BEDROCK_MODEL}/converse", EDGE_MOUNTS) @@ -1021,3 +1097,147 @@ def test_duplicate_headers_bypass_cache_and_count_live_calls( assert len(provider.hits) == (2 if known_mount else 0) assert dict(cache.counters.counts)["duplicate_header_bypass"] == 2 assert dict(cache.counters.counts).get("upstream_attempts", 0) == (2 if known_mount else 0) + + +class TestBedrockStreams: + def test_the_frames_these_tests_build_are_real_aws_framing(self) -> None: + buffer: Final = EventStreamBuffer() + buffer.add_data(CONVERSE_STREAM_OK) + assert [event.headers[":event-type"] for event in buffer] == [ + "messageStart", "contentBlockDelta", "contentBlockStop", "messageStop", "metadata", + ] + + @pytest.mark.parametrize("url,body", [ + (CONVERSE_STREAM_URL, CONVERSE_STREAM_OK), + (INVOKE_STREAM_URL, INVOKE_STREAM_OK), + ]) + def test_a_finished_stream_is_recordable(self, url: str, body: bytes) -> None: + assert cacheable_endpoint(BEDROCK_MOUNT, "POST", url, b"{}") + assert successful_response(BEDROCK_MOUNT, url, 200, {}, body) + + @pytest.mark.parametrize("url,body", [ + (CONVERSE_STREAM_URL, CONVERSE_STREAM_OK), + (INVOKE_STREAM_URL, INVOKE_STREAM_OK), + ]) + @pytest.mark.parametrize("keep", [1, -1, -4]) + def test_a_stream_the_connection_cut_short_is_not_recordable( + self, url: str, body: bytes, keep: int, + ) -> None: + """botocore yields the frames it did receive and silently drops a trailing + partial one, so a stream cut a single byte short parses clean and only the + byte accounting and the terminator rule catch it.""" + assert not successful_response(BEDROCK_MOUNT, url, 200, {}, body[:keep]) + + @pytest.mark.parametrize("url,body", [ + (CONVERSE_STREAM_URL, CONVERSE_STREAM_OK), + (INVOKE_STREAM_URL, INVOKE_STREAM_OK), + ]) + def test_a_corrupted_frame_is_not_recordable(self, url: str, body: bytes) -> None: + flipped: Final = bytearray(body) + flipped[len(body) // 2] ^= 0xFF + assert not successful_response(BEDROCK_MOUNT, url, 200, {}, bytes(flipped)) + + def test_a_converse_stream_that_lost_its_usage_is_not_recordable(self) -> None: + """ConverseStream names its stop reason a frame before it reports usage, + and litellm prices the call from that usage, so a stream cut between the + two would replay as a free call.""" + without_metadata: Final = ( + eventstream_event("messageStart", {"role": "assistant"}) + + eventstream_event("messageStop", {"stopReason": "end_turn"}) + ) + assert not successful_response(BEDROCK_MOUNT, CONVERSE_STREAM_URL, 200, {}, without_metadata) + + def test_a_converse_stream_that_never_stopped_is_not_recordable(self) -> None: + assert not successful_response( + BEDROCK_MOUNT, CONVERSE_STREAM_URL, 200, {}, + eventstream_event("messageStart", {"role": "assistant"}) + + eventstream_event("metadata", {"usage": {"totalTokens": 18}}), + ) + + def test_a_stream_that_failed_after_answering_200_is_not_recordable(self) -> None: + """Bedrock reports a fault that began after the headers went out as an + exception frame in place of the terminator it never got to send.""" + assert not successful_response( + BEDROCK_MOUNT, CONVERSE_STREAM_URL, 200, {}, + eventstream_event("messageStart", {"role": "assistant"}) + + eventstream_event("contentBlockDelta", {"contentBlockIndex": 0, "delta": {"text": "hi"}}) + + eventstream_event("modelStreamErrorException", {"message": "boom"}, message_type="exception"), + ) + + @pytest.mark.parametrize("url,body", [ + (CONVERSE_STREAM_URL, CONVERSE_STREAM_OK), + (INVOKE_STREAM_URL, INVOKE_STREAM_OK), + ], ids=["converse-stream", "invoke-stream"]) + def test_a_stream_cut_after_its_terminator_is_not_recordable(self, url: str, body: bytes) -> None: + """The terminator rules cannot see this one. Every frame the stream owes + has arrived and the partial frame after them is the one botocore drops + without a word, so only counting the bytes against the frame lengths + tells this from a stream that ended where it meant to.""" + assert successful_response(BEDROCK_MOUNT, url, 200, {}, body) + assert not successful_response(BEDROCK_MOUNT, url, 200, {}, body + b"\x00\x00\x02") + + def test_a_converse_stream_whose_stop_frame_names_no_reason_is_not_recordable(self) -> None: + assert not successful_response( + BEDROCK_MOUNT, CONVERSE_STREAM_URL, 200, {}, + eventstream_event("messageStart", {"role": "assistant"}) + + eventstream_event("messageStop", {}) + + eventstream_event("metadata", {"usage": {"totalTokens": 18}}), + ) + + def test_an_invoke_stream_carrying_a_frame_that_is_not_a_chunk_is_not_recordable(self) -> None: + """Every frame of an invoke stream is a `chunk` holding one base64 event. + A frame that is not one carries an event this rule cannot read, so the + stream can no longer be judged complete.""" + assert not successful_response( + BEDROCK_MOUNT, INVOKE_STREAM_URL, 200, {}, + invoke_chunk({"type": "message_start", "message": {"id": "msg_bdrk_x"}}) + + eventstream_event("metadata", {"usage": {"totalTokens": 18}}) + + invoke_chunk({"type": "message_delta", "delta": {"stop_reason": "end_turn"}}) + + invoke_chunk({"type": "message_stop"}), + ) + + def test_a_frame_claiming_no_length_is_rejected_rather_than_walked_forever(self) -> None: + """A frame length of zero never advances the cursor. Rejecting it is what + keeps a corrupt body from spinning the edge instead of answering.""" + assert not successful_response(BEDROCK_MOUNT, CONVERSE_STREAM_URL, 200, {}, b"\x00\x00\x00\x00" * 4) + + @pytest.mark.parametrize("url,terminator", [ + (INVOKE_STREAM_URL, invoke_chunk({"type": "message_stop"})), + (CONVERSE_STREAM_URL, eventstream_event("metadata", {"usage": {"totalTokens": 18}})), + ], ids=["invoke-stream", "converse-stream"]) + def test_a_delta_that_names_no_stop_reason_does_not_finish_a_stream( + self, url: str, terminator: bytes, + ) -> None: + """A `message_delta` arriving without its stop reason is the shape of a + turn the connection cut short partway through the delta itself.""" + head: Final = ( + invoke_chunk({"type": "message_start", "message": {"id": "msg_bdrk_x"}}) + + invoke_chunk({"type": "message_delta", "delta": {}}) + ) + assert not successful_response(BEDROCK_MOUNT, url, 200, {}, head + terminator) + + def test_an_invoke_chunk_that_is_not_base64_is_not_recordable(self) -> None: + assert not successful_response( + BEDROCK_MOUNT, INVOKE_STREAM_URL, 200, {}, + invoke_chunk({"type": "message_start", "message": {"id": "msg_bdrk_x"}}) + + eventstream_event("chunk", {"bytes": "not base64 at all !!"}) + + invoke_chunk({"type": "message_stop"}), + ) + + def test_an_invoke_stream_missing_its_stop_reason_is_not_recordable(self) -> None: + assert not successful_response( + BEDROCK_MOUNT, INVOKE_STREAM_URL, 200, {}, + invoke_chunk({"type": "message_start", "message": {"id": "msg_bdrk_x"}}) + + invoke_chunk({"type": "message_stop"}), + ) + + def test_an_empty_stream_is_not_recordable(self) -> None: + for url in (CONVERSE_STREAM_URL, INVOKE_STREAM_URL): + assert not successful_response(BEDROCK_MOUNT, url, 200, {}, b"") + + def test_each_streaming_endpoint_is_held_to_its_own_grammar(self) -> None: + assert not successful_response(BEDROCK_MOUNT, CONVERSE_STREAM_URL, 200, {}, INVOKE_STREAM_OK) + assert not successful_response(BEDROCK_MOUNT, INVOKE_STREAM_URL, 200, {}, CONVERSE_STREAM_OK) + + def test_a_stream_that_errored_before_it_started_is_not_recordable(self) -> None: + assert not successful_response(BEDROCK_MOUNT, CONVERSE_STREAM_URL, 503, {}, CONVERSE_STREAM_OK) diff --git a/tests/e2e/PROVIDER_CACHE.md b/tests/e2e/PROVIDER_CACHE.md index 698e78a5aa1..29d6e5ab4f2 100644 --- a/tests/e2e/PROVIDER_CACHE.md +++ b/tests/e2e/PROVIDER_CACHE.md @@ -4,7 +4,9 @@ The edge caches complete successful POST responses for `/v1/chat/completions`, `/v1/messages`, `/v1/embeddings` and `/v1/responses` on the OpenAI and Anthropic mounts, SSE streams included, and for `/model/{id}/converse` and `/model/{id}/invoke` on a Bedrock mount. Unsupported endpoints pass through. Each endpoint family has its own completeness rule, so a truncated embedding or a Responses run that never reached `response.completed` is not stored -Bedrock's streaming endpoints, `converse-stream` and `invoke-with-response-stream`, are not cacheable. They still cross the edge and are still re-signed, so they need the same IAM, but they always call the provider. AWS frames them as binary `vnd.amazon.eventstream` rather than SSE, and reading a terminal event out of that is what a completeness rule for them would need. That matters more than the endpoint count suggests: the Claude Code compat cells drive the real CLI, which always streams, so most Bedrock traffic in the suite is not cached today +Bedrock's streaming endpoints, `converse-stream` and `invoke-with-response-stream`, cache too. AWS frames those as binary `vnd.amazon.eventstream` rather than SSE, so botocore's own parser reads the frames and validates both CRCs, and each endpoint is then held to its terminal grammar. That matters more than the endpoint count suggests: the Claude Code compat cells drive the real CLI, which always streams, so streaming is most of the suite's Bedrock traffic + +Two details of that rule are worth knowing before changing it. A ConverseStream ends with `metadata`, not with `messageStop`, and the `metadata` frame is what carries the token usage litellm prices the call from, so the rule requires it: a stream cut between the two still names a stop reason but would replay as a free call. And a dropped connection is invisible to the parser, which yields the frames it did receive and silently discards a trailing partial one, so the body is also checked against the frame lengths it declares. A stream cut one byte short parses clean and has to be caught that way ## Request identity diff --git a/tests/e2e/provider_cache.py b/tests/e2e/provider_cache.py index 0dee33f33c6..9ae89861f63 100644 --- a/tests/e2e/provider_cache.py +++ b/tests/e2e/provider_cache.py @@ -13,6 +13,7 @@ from types import MappingProxyType from typing import Final, Literal, Protocol from urllib.parse import urlsplit +from botocore.eventstream import EventStreamBuffer, ParserError from e2e_http import ( NetworkError, StreamChunk, @@ -36,6 +37,19 @@ SIGNATURE_HEADERS: Final = frozenset( {"authorization", "x-amz-date", "x-amz-security-token", "x-amz-content-sha256"} ) BEDROCK_MOUNT_PREFIX: Final = "bedrock" +BEDROCK_CONVERSE_SUFFIX: Final = "/converse" +BEDROCK_INVOKE_SUFFIX: Final = "/invoke" +BEDROCK_CONVERSE_STREAM_SUFFIX: Final = "/converse-stream" +BEDROCK_INVOKE_STREAM_SUFFIX: Final = "/invoke-with-response-stream" +BEDROCK_SUFFIXES: Final = ( + BEDROCK_CONVERSE_SUFFIX, + BEDROCK_INVOKE_SUFFIX, + BEDROCK_CONVERSE_STREAM_SUFFIX, + BEDROCK_INVOKE_STREAM_SUFFIX, +) +EVENTSTREAM_PRELUDE_BYTES: Final = 4 +EVENT_TYPE_HEADER: Final = ":event-type" +EVENTSTREAM_HEADERS: Final[TypeAdapter[dict[str, str]]] = TypeAdapter(dict[str, str]) OPENAI_JSON_PATHS: Final = frozenset({"/v1/chat/completions", "/v1/messages", "/v1/embeddings", "/v1/responses"}) JSON_VALUE: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue) @@ -145,7 +159,7 @@ def cacheable_endpoint(mount: str, method: str, url: str, body: bytes | None) -> return False path: Final = urlsplit(url).path if is_bedrock(mount): - return path.startswith("/model/") and path.endswith(("/converse", "/invoke")) + return path.startswith("/model/") and path.endswith(BEDROCK_SUFFIXES) return path in OPENAI_JSON_PATHS @@ -176,16 +190,7 @@ def successful_response(mount: str, url: str, status: int, headers: Mapping[str, return complete_responses_stream(values) if urlsplit(url).path == "/v1/chat/completions": return events[-1] == "[DONE]" and "[DONE]" not in events[:-1] and complete_chat_stream(values) - return ( - "[DONE]" not in events - and isinstance(values[0], dict) and values[0].get("type") == "message_start" - and isinstance(values[-1], dict) and values[-1].get("type") == "message_stop" - and any( - isinstance(value, dict) and value.get("type") == "message_delta" - and isinstance(delta := value.get("delta"), dict) and isinstance(delta.get("stop_reason"), str) - for value in values - ) - ) + return "[DONE]" not in events and complete_anthropic_stream(values) try: value: Final = JSON_VALUE.validate_json(body) except ValidationError: @@ -214,14 +219,19 @@ def complete_bedrock_response(url: str, body: bytes) -> bool: """Converse answers with ``output`` plus a ``stopReason``; InvokeModel on an Anthropic model answers the Anthropic message shape. Either way a truncated or error body is missing the terminator field, which is what makes it safe to - record. The streaming variants never reach here: they are not cacheable.""" + record.""" + path: Final = urlsplit(url).path + if path.endswith(BEDROCK_CONVERSE_STREAM_SUFFIX): + return complete_converse_stream(body) + if path.endswith(BEDROCK_INVOKE_STREAM_SUFFIX): + return complete_invoke_stream(body) try: value: Final = JSON_VALUE.validate_json(body) except ValidationError: return False if not isinstance(value, dict) or "message" in value: return False - if urlsplit(url).path.endswith("/converse"): + if path.endswith(BEDROCK_CONVERSE_SUFFIX): return isinstance(value.get("output"), dict) and isinstance(value.get("stopReason"), str) return ( value.get("type") == "message" @@ -230,6 +240,111 @@ def complete_bedrock_response(url: str, body: bytes) -> bool: ) +def whole_eventstream_messages(body: bytes) -> bool: + """Whether the body is exactly a whole number of eventstream messages. + + A dropped connection is the failure this catches, and it has to be caught + here: botocore yields the messages it did receive and silently discards a + trailing partial one, so a stream cut a single byte short parses clean. Each + message declares its own total length in its first four bytes, so walking + those is enough to tell a complete body from a cut one.""" + offset = 0 # rebind-ok: a cursor walking the declared frame lengths + while offset + EVENTSTREAM_PRELUDE_BYTES <= len(body): + total: int = int.from_bytes(body[offset : offset + EVENTSTREAM_PRELUDE_BYTES], "big") + if total <= 0 or offset + total > len(body): + return False + offset += total + return offset == len(body) + + +def eventstream_events(body: bytes) -> tuple[tuple[str, JsonValue], ...] | None: + """The stream's (event type, decoded payload) pairs, or None if it is not a + complete, uncorrupted stream. + + botocore validates both CRCs and raises ``ParserError`` rather than decoding + corruption into something plausible. A failure that began after Bedrock had + already answered 200 arrives as an ``exception`` frame in place of the + terminator, so it is the terminator rules below that reject it and this does + not need to inspect ``:message-type`` as well.""" + if not body or not whole_eventstream_messages(body): + return None + buffer: Final = EventStreamBuffer() + buffer.add_data(body) + try: + return tuple( + (event_type(event.headers), JSON_VALUE.validate_json(event.payload)) + for event in buffer + ) + except (ParserError, ValidationError, ValueError): + return None + + +def event_type(headers: object) -> str: + """botocore's eventstream headers come back untyped, so the one header this + reads is validated into a string rather than trusted.""" + parsed: Final = EVENTSTREAM_HEADERS.validate_python(headers) + return parsed.get(EVENT_TYPE_HEADER, "") + + +def complete_converse_stream(body: bytes) -> bool: + """ConverseStream ends with ``metadata``, not with ``messageStop``. + + Requiring the metadata frame rather than the stop frame is deliberate: it + carries the token usage litellm prices the call from, so a stream cut between + the two still names a stop reason but would replay as a free call.""" + events: Final = eventstream_events(body) + if not events or events[-1][0] != "metadata": + return False + return any( + event_type == "messageStop" and isinstance(payload, dict) and isinstance(payload.get("stopReason"), str) + for event_type, payload in events + ) + + +def complete_invoke_stream(body: bytes) -> bool: + """InvokeModelWithResponseStream wraps the ordinary Anthropic event grammar + in ``chunk`` frames, one base64 payload each, so it is held to the same + terminator rule as the Anthropic SSE path. A frame Bedrock sends instead of a + chunk, an exception among them, carries no such payload and fails the rule + without the frame type needing to be read.""" + events: Final = eventstream_events(body) + if not events: + return False + values: Final = tuple(invoke_chunk_value(payload) for _, payload in events) + return all(value is not None for value in values) and complete_anthropic_stream(values) + + +def invoke_chunk_value(payload: JsonValue) -> JsonValue | None: + """The Anthropic event inside one ``chunk`` frame, or None for a frame that + carries no readable one.""" + if not isinstance(payload, dict) or not isinstance(encoded := payload.get("bytes"), str): + return None + try: + return JSON_VALUE.validate_json(base64.b64decode(encoded, validate=True)) + except (ValidationError, ValueError): + return None + + +def complete_anthropic_stream(values: tuple[JsonValue, ...]) -> bool: + """The Anthropic event grammar, shared by the SSE mounts and by Bedrock's + invoke stream, which carries the same events inside eventstream frames. A + ``message_delta`` naming a stop reason is what separates a finished turn from + one the connection cut short.""" + if not values: + return False + first: Final = values[0] + last: Final = values[-1] + return ( + isinstance(first, dict) and first.get("type") == "message_start" + and isinstance(last, dict) and last.get("type") == "message_stop" + and any( + isinstance(value, dict) and value.get("type") == "message_delta" + and isinstance(delta := value.get("delta"), dict) and isinstance(delta.get("stop_reason"), str) + for value in values + ) + ) + + def complete_responses_stream(values: tuple[JsonValue, ...]) -> bool: """The Responses API streams typed events and ends with ``response.completed``. A run that failed, was cancelled, or ran out of tokens ends with a different From 8a553ceb58887c8aa2aa24c7cfb75ce662d2e321 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 16 Sep 2026 07:38:44 -0700 Subject: [PATCH 10/21] feat(e2e): mount Gemini on the provider cache Gemini needs none of the machinery Bedrock needed. litellm composes {api_base}/models/{model}:{endpoint} from a custom api_base, so a plain path-prefixed mount reaches it, and the credential travels as a static x-goog-api-key header that no host rewrite invalidates. Nothing is re-signed and nothing leaves the cache key, so a recording still cannot cross credentials. A finished turn names a finishReason on every candidate and reports usageMetadata. The reason is read as a string rather than compared to STOP: MAX_TOKENS and the safety reasons end a turn just as finally, and rejecting them would send every one of them upstream forever. Streaming is the half worth care. Gemini repeats usageMetadata on every chunk and names a finishReason only on the last, so the terminator is the final event rather than any event, and a stream the connection cut short ends on a chunk carrying usage and no reason. The mount's upstream base carries the API version, so the path the rules see is /v1beta/models/..., not the one the proxy sent. The first version of this anchored the rule at the start of that path, which passed every test against a stub with no version prefix and would have cached nothing at all in a real run. Caught by replaying the rules over responses captured from live gemini-2.5-flash, which is also why the tests now mount their stub under the version prefix. Vertex stays unmounted and is a separate provider here: litellm grafts the default Vertex path onto an api_base only when that api_base has no path of its own, so Vertex needs a root-mounted edge on its own port. --- .../test_provider_cache.py | 154 +++++++++++++++++- tests/e2e/PROVIDER_CACHE.md | 16 +- tests/e2e/provider_cache.py | 46 ++++++ tests/e2e/provider_cache_routing.py | 3 +- tests/e2e/provider_edge.py | 1 + 5 files changed, 215 insertions(+), 5 deletions(-) diff --git a/tests/code_coverage_tests/test_provider_cache.py b/tests/code_coverage_tests/test_provider_cache.py index 4c131434ecd..520317a0678 100644 --- a/tests/code_coverage_tests/test_provider_cache.py +++ b/tests/code_coverage_tests/test_provider_cache.py @@ -861,7 +861,7 @@ def test_anthropic_stream_requires_start_finish_and_stop() -> None: assert not successful_response("anthropic", url, 200, headers, start + finish) -@pytest.mark.parametrize("provider,suffix", [("openai", "/v1"), ("anthropic", "")]) +@pytest.mark.parametrize("provider,suffix", [("openai", "/v1"), ("anthropic", ""), ("gemini", "")]) def test_normal_registration_routes_supported_providers(provider: str, suffix: str) -> None: params: Final = LiteLLMParamsBody(model=f"{provider}/test", api_key="os.environ/SYNTHETIC_KEY", timeout=12) routed: Final = route_cache_model(params, lambda mount: f"http://edge.invalid/{mount}", enabled=True) @@ -873,6 +873,8 @@ def test_normal_registration_routes_supported_providers(provider: str, suffix: s @pytest.mark.parametrize("params", [ LiteLLMParamsBody(model="bedrock/test"), LiteLLMParamsBody(model="azure/test"), + LiteLLMParamsBody(model="vertex_ai/gemini-2.5-flash"), + LiteLLMParamsBody(model="gemini/gemini-2.5-flash", api_base="https://custom.invalid"), LiteLLMParamsBody(model="openai/test", api_base="https://custom.invalid/v1"), LiteLLMParamsBody(model="openai/test", api_base=""), LiteLLMParamsBody(model="openai/test", litellm_credential_name="named-credential"), @@ -1241,3 +1243,153 @@ class TestBedrockStreams: def test_a_stream_that_errored_before_it_started_is_not_recordable(self) -> None: assert not successful_response(BEDROCK_MOUNT, CONVERSE_STREAM_URL, 503, {}, CONVERSE_STREAM_OK) + + +GEMINI_MODEL: Final = "gemini-2.5-flash" +GEMINI_API_VERSION: Final = "/v1beta" +GEMINI_GENERATE_PATH: Final = f"/models/{GEMINI_MODEL}:generateContent" +GEMINI_STREAM_PATH: Final = f"/models/{GEMINI_MODEL}:streamGenerateContent" +GEMINI_USAGE: Final = {"promptTokenCount": 7, "candidatesTokenCount": 1, "totalTokenCount": 25} + + +def gemini_body(finish_reason: str | None, usage: bool = True, candidates: bool = True) -> JsonValue: + candidate: Final[dict[str, JsonValue]] = {"content": {"parts": [{"text": "OK"}], "role": "model"}, "index": 0} + return { + "candidates": [{**candidate, "finishReason": finish_reason} if finish_reason else candidate] + if candidates else [], + **({"usageMetadata": GEMINI_USAGE} if usage else {}), + "modelVersion": GEMINI_MODEL, + } + + +def gemini_unary(finish_reason: str | None = "STOP", usage: bool = True, candidates: bool = True) -> bytes: + return json.dumps(gemini_body(finish_reason, usage, candidates)).encode() + + +def gemini_stream(*finish_reasons: str | None) -> bytes: + return b"".join( + b"data: " + json.dumps(gemini_body(reason)).encode() + b"\r\n\r\n" for reason in finish_reasons + ) + + +@contextmanager +def gemini_edge(cache: CacheEdge, provider: Provider, path: str) -> Generator[str, None, None]: + upstream: Final = f"http://127.0.0.1:{provider.server_port}{GEMINI_API_VERSION}" + running: Final = start_provider_edge(cache, mounts={"gemini": upstream}) + try: + yield running.edge.api_base("gemini") + path + finally: + running.shutdown() + + +class TestGemini: + """Gemini reaches the edge by path prefix alone: litellm composes + `{api_base}/models/{model}:{endpoint}` and sends a static `x-goog-api-key`, + so nothing has to be re-signed and nothing leaves the cache key. The response + grammar is its own though, and the streaming one is the interesting half: every + chunk repeats `usageMetadata`, so only `finishReason` on the last chunk + separates a finished turn from a dropped connection.""" + + @pytest.mark.parametrize("path,response", [ + (GEMINI_GENERATE_PATH, gemini_unary()), + (GEMINI_STREAM_PATH, gemini_stream(None, None, "STOP")), + ], ids=["generate", "stream"]) + def test_a_finished_turn_replays_on_the_next_run( + self, store: RedisResponseStore, provider: Provider, path: str, response: bytes, + ) -> None: + provider.stream = path == GEMINI_STREAM_PATH + provider.response = response + for _ in range(2): + with gemini_edge(cache_edge(store), provider, path) as url: + assert call(url, MARKED).body == response + assert len(provider.hits) == 1 + + @pytest.mark.parametrize("reason", ["MAX_TOKENS", "SAFETY", "RECITATION"]) + def test_a_turn_the_provider_ended_for_its_own_reasons_is_still_finished( + self, store: RedisResponseStore, provider: Provider, reason: str, + ) -> None: + """Reading `finishReason` as a string rather than comparing it to STOP is + deliberate. A turn cut off by the token limit or a safety filter is over, + and rejecting those would send every one of them upstream forever.""" + provider.response = gemini_unary(reason) + for _ in range(2): + with gemini_edge(cache_edge(store), provider, GEMINI_GENERATE_PATH) as url: + assert call(url, MARKED).body == provider.response + assert len(provider.hits) == 1 + + @pytest.mark.parametrize("response", [ + gemini_unary(None), + gemini_unary("STOP", usage=False), + gemini_unary("STOP", candidates=False), + b'{"error":{"code":400,"message":"API key not valid","status":"INVALID_ARGUMENT"}}', + ], ids=["no-finish-reason", "no-usage", "no-candidates", "error-body"]) + def test_an_unfinished_or_failed_turn_never_enters_the_cache( + self, store: RedisResponseStore, provider: Provider, response: bytes, + ) -> None: + provider.response = response + for _ in range(2): + with gemini_edge(cache_edge(store), provider, GEMINI_GENERATE_PATH) as url: + assert call(url, MARKED).body == response + assert len(provider.hits) == 2 + + @pytest.mark.parametrize("response", [ + gemini_stream(None, None), + gemini_stream("STOP", None), + gemini_stream(), + ], ids=["cut-before-the-reason", "reason-then-another-chunk", "empty"]) + def test_a_stream_that_never_named_a_reason_calls_the_provider_every_time( + self, store: RedisResponseStore, provider: Provider, response: bytes, + ) -> None: + provider.stream = True + provider.response = response + for _ in range(2): + with gemini_edge(cache_edge(store), provider, GEMINI_STREAM_PATH) as url: + assert call(url, MARKED).body == response + assert len(provider.hits) == 2 + + def test_a_response_whose_candidates_did_not_all_finish_is_not_recordable( + self, store: RedisResponseStore, provider: Provider, + ) -> None: + """A request for more than one candidate is answered by more than one, and + the turn is over only when every one of them names a reason. Holding the + whole list to that rule rather than its first entry is what keeps a + half-finished answer from being stored and replayed as a finished one.""" + finished: Final = json.loads(gemini_unary("STOP"))["candidates"][0] + unfinished: Final = json.loads(gemini_unary(None))["candidates"][0] + provider.response = json.dumps( + {"candidates": [finished, {**unfinished, "index": 1}], "usageMetadata": GEMINI_USAGE} + ).encode() + for _ in range(2): + with gemini_edge(cache_edge(store), provider, GEMINI_GENERATE_PATH) as url: + assert call(url, MARKED).body == provider.response + assert len(provider.hits) == 2 + + @pytest.mark.parametrize("path,cacheable", [ + (GEMINI_GENERATE_PATH, True), + (GEMINI_STREAM_PATH, True), + (f"/models/{GEMINI_MODEL}:countTokens", False), + (f"/models/{GEMINI_MODEL}:embedContent", False), + ("/v1/chat/completions", False), + (f"/files/{GEMINI_MODEL}:generateContent", False), + ]) + @pytest.mark.parametrize("version", ["", GEMINI_API_VERSION], ids=["bare", "versioned"]) + def test_only_the_generate_endpoints_are_cacheable(self, version: str, path: str, cacheable: bool) -> None: + """The mount's upstream base carries the API version, so the path the cache + sees is the upstream one and starts `/v1beta`. A rule anchored at the start + of the path would pass every test against a stub with no version prefix and + then cache nothing at all in a real run.""" + assert cacheable_endpoint("gemini", "POST", f"https://gemini.invalid{version}{path}", MARKED) is cacheable + + def test_the_bodies_these_tests_build_match_a_real_gemini_response(self) -> None: + """The shapes above are hand-built so a test can express the turn it means. + This holds them to the fields a live `generativelanguage.googleapis.com` + answer carries, captured 2026-09-16 against gemini-2.5-flash.""" + captured: Final = json.loads( + '{"candidates":[{"content":{"parts":[{"text":"OK"}],"role":"model"},"finishReason":"STOP",' + '"index":0}],"usageMetadata":{"promptTokenCount":7,"candidatesTokenCount":1,' + '"totalTokenCount":25},"modelVersion":"gemini-2.5-flash","responseId":"1J6qauKFI8ut1MkPgNjI4AI"}' + ) + built: Final = json.loads(gemini_unary()) + assert captured.keys() >= built.keys() + assert captured["candidates"][0].keys() >= built["candidates"][0].keys() + assert successful_response("gemini", GEMINI_GENERATE_PATH, 200, {}, json.dumps(captured).encode()) diff --git a/tests/e2e/PROVIDER_CACHE.md b/tests/e2e/PROVIDER_CACHE.md index 29d6e5ab4f2..e2f8030074d 100644 --- a/tests/e2e/PROVIDER_CACHE.md +++ b/tests/e2e/PROVIDER_CACHE.md @@ -1,13 +1,23 @@ # Shared provider-response cache -`E2E_PROVIDER_CACHE=1` enables automatic response reuse in the live E2E mode. Standard OpenAI and Anthropic model registrations use the provider edge, as do Anthropic-on-Bedrock registrations that carry no AWS identity of their own. Existing custom API bases, named credentials, mocked models and realtime WebSocket deployments keep their existing routing. Other provider protocols remain live +`E2E_PROVIDER_CACHE=1` enables automatic response reuse in the live E2E mode. Standard OpenAI, Anthropic and Gemini model registrations use the provider edge, as do Anthropic-on-Bedrock registrations that carry no AWS identity of their own. 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`, `/v1/messages`, `/v1/embeddings` and `/v1/responses` on the OpenAI and Anthropic mounts, SSE streams included, and for `/model/{id}/converse` and `/model/{id}/invoke` on a Bedrock mount. Unsupported endpoints pass through. Each endpoint family has its own completeness rule, so a truncated embedding or a Responses run that never reached `response.completed` is not stored +The edge caches complete successful POST responses for `/v1/chat/completions`, `/v1/messages`, `/v1/embeddings` and `/v1/responses` on the OpenAI and Anthropic mounts, SSE streams included, for `/model/{id}/converse` and `/model/{id}/invoke` on a Bedrock mount, and for `models/{model}:generateContent` and `:streamGenerateContent` on the Gemini mount. Unsupported endpoints pass through. Each endpoint family has its own completeness rule, so a truncated embedding or a Responses run that never reached `response.completed` is not stored Bedrock's streaming endpoints, `converse-stream` and `invoke-with-response-stream`, cache too. AWS frames those as binary `vnd.amazon.eventstream` rather than SSE, so botocore's own parser reads the frames and validates both CRCs, and each endpoint is then held to its terminal grammar. That matters more than the endpoint count suggests: the Claude Code compat cells drive the real CLI, which always streams, so streaming is most of the suite's Bedrock traffic Two details of that rule are worth knowing before changing it. A ConverseStream ends with `metadata`, not with `messageStop`, and the `metadata` frame is what carries the token usage litellm prices the call from, so the rule requires it: a stream cut between the two still names a stop reason but would replay as a free call. And a dropped connection is invisible to the parser, which yields the frames it did receive and silently discards a trailing partial one, so the body is also checked against the frame lengths it declares. A stream cut one byte short parses clean and has to be caught that way +## Gemini + +Gemini needs nothing that Bedrock needed. litellm composes `{api_base}/models/{model}:{endpoint}` from a custom api_base, so a path-prefixed mount reaches it, and the credential travels as a static `x-goog-api-key` header that no host rewrite invalidates. Nothing is re-signed and nothing is excluded from the key, so a recording still cannot cross credentials + +The mount's upstream base carries the API version, which is the one detail worth remembering: the path the cache rules see is the upstream one, `/v1beta/models/...`, not the one the proxy sent. A rule anchored at the start of that path would look right against a local stub and then cache nothing at all in a real run + +A finished turn names a `finishReason` on every candidate and reports `usageMetadata`. The reason is read as a string rather than compared to `STOP`, because `MAX_TOKENS` and the safety reasons end a turn just as finally and rejecting them would send every one of them upstream forever. Streaming is the more interesting half: Gemini repeats `usageMetadata` on every chunk and names a `finishReason` only on the last one, so the terminator is the final event rather than any event, and a stream the connection cut short ends on a chunk with usage and no reason + +Vertex is not mounted. litellm grafts the default Vertex path onto an api_base only when that api_base has no path of its own, so a Vertex mount needs a root-mounted edge on its own port rather than a path prefix. Gemini and Vertex are separate providers in litellm and the Gemini mount does not cover Vertex deployments + ## Request identity A recording belongs to one test. The key is a keyed digest over the test's node id, the method, the URL, the effective outbound headers (including authentication and HTTP-library defaults), body presence and the body bytes, with one normalization: a 12-hex-digit run, the shape `unique_marker()` mints, is replaced by a placeholder in both the URL and a UTF-8 body. Nothing else is normalized away. No prompts, JSON values or credentials are rewritten, and the rule is the one `fixture_canonical.py` already applies for record/replay, so there is a single definition of what a marker is @@ -30,7 +40,7 @@ Only deployments that carry no AWS identity of their own route to the edge. A de Which models route is an explicit allowlist in `provider_cache_routing.py`, mirroring the runner role's IAM policy, which names its models one by one. That coupling is deliberate: the edge re-signs with the run pod's identity, so a model the role cannot invoke comes back 403 from Bedrock rather than falling back. An unlisted model keeps its direct path and loses only caching, so adding a Bedrock model to the suite can never turn it red. Adding one to the edge is a policy edit in litellm-ops plus a line here -Vertex and Gemini are not mounted. litellm's `_check_custom_proxy` rewrites a path-prefixed Vertex `api_base` into `{api_base}:{endpoint}`, dropping project, location and model, so a mount under a path prefix cannot work without either a root-mounted edge on its own port or a change in litellm +Vertex is not mounted. litellm's `_check_custom_proxy` rewrites a path-prefixed Vertex `api_base` into `{api_base}:{endpoint}`, dropping project, location and model, so a mount under a path prefix cannot work without either a root-mounted edge on its own port or a change in litellm. Gemini is a separate provider there and does have a working path-prefixed form, so it is mounted; see the Gemini section Recordings are shared across workers and builds through dedicated Redis, separate from the candidate's own cache. They expire 86,400 seconds after capture starts, based on Redis time. Reads never extend expiry. There is no scheduled recapture: the next miss calls the provider again. Bounded coordination reduces duplicate concurrent calls, but slow or failed captures may lead to extra live calls after the wait expires diff --git a/tests/e2e/provider_cache.py b/tests/e2e/provider_cache.py index 9ae89861f63..7bba93321b5 100644 --- a/tests/e2e/provider_cache.py +++ b/tests/e2e/provider_cache.py @@ -37,6 +37,10 @@ SIGNATURE_HEADERS: Final = frozenset( {"authorization", "x-amz-date", "x-amz-security-token", "x-amz-content-sha256"} ) BEDROCK_MOUNT_PREFIX: Final = "bedrock" +GEMINI_MOUNT: Final = "gemini" +GEMINI_MODELS_SEGMENT: Final = "/models" +GEMINI_GENERATE_SUFFIX: Final = ":generateContent" +GEMINI_STREAM_SUFFIX: Final = ":streamGenerateContent" BEDROCK_CONVERSE_SUFFIX: Final = "/converse" BEDROCK_INVOKE_SUFFIX: Final = "/invoke" BEDROCK_CONVERSE_STREAM_SUFFIX: Final = "/converse-stream" @@ -154,12 +158,21 @@ def is_bedrock(mount: str) -> bool: return mount.partition("/")[0] == BEDROCK_MOUNT_PREFIX +def is_gemini(mount: str) -> bool: + return mount == GEMINI_MOUNT + + def cacheable_endpoint(mount: str, method: str, url: str, body: bytes | None) -> bool: if method != "POST" or body is None or len(body) > MAX_REQUEST_BYTES: return False path: Final = urlsplit(url).path if is_bedrock(mount): return path.startswith("/model/") and path.endswith(BEDROCK_SUFFIXES) + if is_gemini(mount): + collection, _, resource = path.rpartition("/") + return collection.endswith(GEMINI_MODELS_SEGMENT) and resource.endswith( + (GEMINI_GENERATE_SUFFIX, GEMINI_STREAM_SUFFIX) + ) return path in OPENAI_JSON_PATHS @@ -186,6 +199,8 @@ def successful_response(mount: str, url: str, status: int, headers: Mapping[str, for value in values ): return False + if is_gemini(mount): + return complete_gemini_stream(values) if urlsplit(url).path == "/v1/responses": return complete_responses_stream(values) if urlsplit(url).path == "/v1/chat/completions": @@ -197,6 +212,8 @@ def successful_response(mount: str, url: str, status: int, headers: Mapping[str, return False if not isinstance(value, dict) or value.get("error") is not None: return False + if is_gemini(mount): + return complete_gemini_candidates(value) path: Final = urlsplit(url).path if path == "/v1/messages": return value.get("type") == "message" and isinstance(value.get("content"), list) and isinstance(value.get("stop_reason"), str) @@ -215,6 +232,35 @@ def successful_response(mount: str, url: str, status: int, headers: Mapping[str, ) +def complete_gemini_candidates(value: Mapping[str, JsonValue]) -> bool: + """A finished Gemini turn names a ``finishReason`` on every candidate and + reports the usage litellm prices the call from. ``finishReason`` is read as a + string rather than compared to ``STOP`` because ``MAX_TOKENS`` and the safety + reasons end a turn just as finally, and a cache that rejected them would send + every one of them upstream forever.""" + candidates: Final = value.get("candidates") + return ( + isinstance(value.get("usageMetadata"), dict) + and isinstance(candidates, list) + and bool(candidates) + and all( + isinstance(candidate, dict) and isinstance(candidate.get("finishReason"), str) + for candidate in candidates + ) + ) + + +def complete_gemini_stream(values: tuple[JsonValue, ...]) -> bool: + """Gemini repeats ``usageMetadata`` on every chunk but names a + ``finishReason`` only on the last one, so the terminator is the final event + rather than any event. A stream the connection cut short ends on a chunk that + carries usage and no reason, which is exactly what this rejects.""" + if not values: + return False + last: Final = values[-1] + return isinstance(last, dict) and complete_gemini_candidates(last) + + def complete_bedrock_response(url: str, body: bytes) -> bool: """Converse answers with ``output`` plus a ``stopReason``; InvokeModel on an Anthropic model answers the Anthropic message shape. Either way a truncated diff --git a/tests/e2e/provider_cache_routing.py b/tests/e2e/provider_cache_routing.py index f9775a2b152..c4b02beac2d 100644 --- a/tests/e2e/provider_cache_routing.py +++ b/tests/e2e/provider_cache_routing.py @@ -19,6 +19,7 @@ BEDROCK_EDGE_MODELS: Final = frozenset( } ) ENV_REFERENCE_PREFIX: Final = "os.environ/" +EDGE_PROVIDERS: Final = frozenset({"openai", "anthropic", "gemini"}) def bedrock_region(declared: str | None) -> str: @@ -81,7 +82,7 @@ def route_cache_model( return route_bedrock(params, base_for, mode) if mode == "realtime" or params.api_base is not None: return params - if provider not in {"openai", "anthropic"}: + if provider not in EDGE_PROVIDERS: return params base: Final = base_for(provider) if base is None: diff --git a/tests/e2e/provider_edge.py b/tests/e2e/provider_edge.py index 2606b26fe99..219df55233a 100644 --- a/tests/e2e/provider_edge.py +++ b/tests/e2e/provider_edge.py @@ -104,6 +104,7 @@ EDGE_MOUNTS: Final[Mapping[str, str]] = MappingProxyType( { "openai": "https://api.openai.com", "anthropic": "https://api.anthropic.com", + "gemini": "https://generativelanguage.googleapis.com/v1beta", **{ f"bedrock/{region}": f"https://bedrock-runtime.{region}.amazonaws.com" for region in BEDROCK_REGIONS From 7006da9cde950981379214f3d2dd0f645d68995e Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 16 Sep 2026 07:47:31 -0700 Subject: [PATCH 11/21] feat(e2e): say why a response was not recorded Build 226 routed Bedrock streaming for the first time and rejected 62 of 220 misses on that mount, and the counters could not say why. A flat rejected count covers three unrelated things with opposite fixes: the consumer walking away mid-capture, a body that arrived whole and failed its endpoint's rule, and a provider that could not be reached. Each now also counts its own reason. A consumer that walks away was counting nothing at all. Abandoning the capture generator raises GeneratorExit at its yield, so neither branch of the old accounting ran and the miss simply vanished from the report, which is also why misses could exceed writes plus rejected with nothing to explain the gap. The decision moves into settle() so the generator's finally owns the accounting and an abandoned capture is counted like any other rejection. --- .../test_provider_cache.py | 39 ++++++++++++++ tests/e2e/PROVIDER_CACHE.md | 2 +- tests/e2e/provider_cache.py | 53 +++++++++++++------ 3 files changed, 78 insertions(+), 16 deletions(-) diff --git a/tests/code_coverage_tests/test_provider_cache.py b/tests/code_coverage_tests/test_provider_cache.py index 520317a0678..1271b20438a 100644 --- a/tests/code_coverage_tests/test_provider_cache.py +++ b/tests/code_coverage_tests/test_provider_cache.py @@ -516,6 +516,45 @@ def test_counters_attribute_every_outcome_to_its_mount( assert counts["mount:anthropic:rejected"] == 1 and "mount:openai:rejected" not in counts +def test_a_rejection_says_whether_the_body_was_cut_short_or_simply_unfinished( + store: RedisResponseStore, provider: Provider, +) -> None: + """One `rejected` count cannot tell a connection that dropped from a body the + provider finished sending and the rules turned down, and those have opposite + fixes: the first is the client going away mid-capture, the second is a grammar + the cache does not accept. A mount whose rejections are mostly one or the other + is a different problem, so the report has to be able to say which.""" + upstream: Final = f"http://127.0.0.1:{provider.server_port}" + cut_short: Final = cache_edge(store) + provider.stream = True + provider.truncated = True + provider.response = b'data: {"choices":[{"index":0,"delta":{"content":"hi"},"finish_reason":"stop"}]}\n\ndata: [DONE]\n\n' + running: Final = start_provider_edge(cut_short, mounts={"openai": upstream}) + try: + forward("POST", running.edge.api_base("openai") + "/v1/chat/completions", + headers=HEADERS, body=MARKED, timeout=5) + finally: + running.shutdown() + + unfinished: Final = cache_edge(store) + provider.stream = False + provider.truncated = False + provider.response = b'{"choices":[{"index":0,"message":{"content":"hi"}}]}' + second: Final = start_provider_edge(unfinished, mounts={"openai": upstream}) + try: + call(second.edge.api_base("openai") + "/v1/chat/completions", MARKED) + finally: + second.shutdown() + + cut: Final = dict(cut_short.counters.counts) + turned_down: Final = dict(unfinished.counters.counts) + assert cut["mount:openai:rejected"] == 1 and turned_down["mount:openai:rejected"] == 1 + assert cut["mount:openai:rejected_cut_short"] == 1 + assert "mount:openai:rejected_incomplete" not in cut + assert turned_down["mount:openai:rejected_incomplete"] == 1 + assert "mount:openai:rejected_cut_short" not in turned_down + + EMBEDDING_SUCCESS: Final = ( b'{"object":"list","data":[{"object":"embedding","index":0,"embedding":[0.1,0.2]}],' b'"model":"text-embedding-3-small","usage":{"prompt_tokens":2,"total_tokens":2}}' diff --git a/tests/e2e/PROVIDER_CACHE.md b/tests/e2e/PROVIDER_CACHE.md index e2f8030074d..894d9be4efa 100644 --- a/tests/e2e/PROVIDER_CACHE.md +++ b/tests/e2e/PROVIDER_CACHE.md @@ -54,7 +54,7 @@ The trusted runner receives: - `E2E_PROVIDER_CACHE_NAMESPACE`: shared environment namespace, independent of build and candidate revision - `E2E_PROVIDER_CACHE_METRICS_DIR`: optional per-process counter artifact directory -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. Every counter is emitted twice, once as a flat total and once under `mount:{mount}:`, so a hit rate can be read per provider rather than only in aggregate. Existing application-cache observations still count requests arriving at the edge, including shared-cache hits +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. A rejection also counts its reason, one of `rejected_cut_short` (the consumer walked away mid-capture), `rejected_incomplete` (the body arrived whole and failed its endpoint's rule) or `rejected_unreachable` (the provider could not be reached). A mount whose rejections are nearly all one or the other is a different problem, and the flat count cannot tell them apart. Every counter is emitted twice, once as a flat total and once under `mount:{mount}:`, so a hit rate can be read per provider rather than only in aggregate. 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. The provider prompt-caching tests carry it because a replayed priming response reports cache creation rather than a cache read. diff --git a/tests/e2e/provider_cache.py b/tests/e2e/provider_cache.py index 7bba93321b5..f528d08b720 100644 --- a/tests/e2e/provider_cache.py +++ b/tests/e2e/provider_cache.py @@ -52,6 +52,9 @@ BEDROCK_SUFFIXES: Final = ( BEDROCK_INVOKE_STREAM_SUFFIX, ) EVENTSTREAM_PRELUDE_BYTES: Final = 4 +CUT_SHORT: Final = "cut_short" +INCOMPLETE: Final = "incomplete" +UNREACHABLE: Final = "unreachable" EVENT_TYPE_HEADER: Final = ":event-type" EVENTSTREAM_HEADERS: Final[TypeAdapter[dict[str, str]]] = TypeAdapter(dict[str, str]) OPENAI_JSON_PATHS: Final = frozenset({"/v1/chat/completions", "/v1/messages", "/v1/embeddings", "/v1/responses"}) @@ -551,7 +554,7 @@ class CacheEdge: ) prepared: Final = prepare_forward(method, url, self.outbound(mount, method, url, headers, body), body) if isinstance(prepared, NetworkError): - self.count(mount, "rejected") + self.reject(mount, UNREACHABLE) return prepared identity: Final = request_identity( self.secret, test_key, method, url, self.keyed(mount, prepared.headers), body, @@ -575,7 +578,7 @@ class CacheEdge: return head if isinstance(head, NetworkError): self.store.release(key, capture_slot) - self.count(mount, "rejected") + self.reject(mount, UNREACHABLE) return head return StreamHead( head.status_code, head.headers, primed_steps(self.capture(mount, key, capture_slot, url, head)), @@ -585,25 +588,45 @@ class CacheEdge: self, mount: str, key: str, lease: CaptureLease, url: str, head: StreamHead, ) -> Generator[StreamStep, None, None]: capture: Final = ResponseCapture() + reason = CUT_SHORT # rebind-ok: a consumer that walks away never reaches the settle call below try: with closing(head.steps): yield StreamChunk(b"") for step in head.steps: yield step capture.observe(step) - chunks: Final = capture.chunks() if capture.eligible else () - 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(mount, url, head.status_code, headers, b"".join(chunks)): - self.count(mount, "rejected") - return - response: Final = CachedResponse( - 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)) - self.count(mount, "writes" if published else "write_failures") + reason = self.settle(mount, key, lease, url, head, capture) finally: + self.reject(mount, reason) self.store.release(key, lease) capture.buffer.close() + + def settle( + self, mount: str, key: str, lease: CaptureLease, url: str, head: StreamHead, capture: ResponseCapture, + ) -> str | None: + """None once the response is stored, otherwise the reason it was not.""" + if not capture.eligible: + return CUT_SHORT + headers: Final = { + name: value for name, value in head.headers.items() if name.lower() not in UNRECORDED_RESPONSE_HEADERS + } + chunks: Final = capture.chunks() + if not successful_response(mount, url, head.status_code, headers, b"".join(chunks)): + return INCOMPLETE + response: Final = CachedResponse( + 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)) + self.count(mount, "writes" if published else "write_failures") + return None + + def reject(self, mount: str, reason: str | None) -> None: + """A flat rejection count cannot separate a connection that went away from + a body the provider finished sending and the rules turned down, and the two + have opposite fixes. A mount whose rejections are nearly all one or the + other is a different problem, so the report has to be able to say which.""" + if reason is None: + return + self.count(mount, "rejected") + self.count(mount, f"rejected_{reason}") From c447c3312db1b99e058062f8cd61f2904fa1e9ef Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 16 Sep 2026 08:13:04 -0700 Subject: [PATCH 12/21] feat(e2e): separate a provider error from a body that failed its rule Build 226's 62 Bedrock rejections are the question this is trying to answer, and "incomplete" would have covered both candidate causes at once. Replaying the completeness rules over eight streams captured from live Bedrock, covering tool use, extended thinking and a max-tokens stop on both streaming endpoints, accepts every one of them, so a rule that is too strict is the less likely half. A provider that answered 429 or 5xx and was retried out of sight is the other, and it now counts as rejected_error_status rather than being folded in with a grammar failure. --- .../code_coverage_tests/test_provider_cache.py | 18 +++++++++++++++--- tests/e2e/PROVIDER_CACHE.md | 2 +- tests/e2e/provider_cache.py | 3 +++ 3 files changed, 19 insertions(+), 4 deletions(-) diff --git a/tests/code_coverage_tests/test_provider_cache.py b/tests/code_coverage_tests/test_provider_cache.py index 1271b20438a..9668de35a02 100644 --- a/tests/code_coverage_tests/test_provider_cache.py +++ b/tests/code_coverage_tests/test_provider_cache.py @@ -546,13 +546,25 @@ def test_a_rejection_says_whether_the_body_was_cut_short_or_simply_unfinished( finally: second.shutdown() + refused: Final = cache_edge(store) + provider.status = 429 + provider.response = b'{"message":"Too many requests"}' + third: Final = start_provider_edge(refused, mounts={"openai": upstream}) + try: + call(third.edge.api_base("openai") + "/v1/chat/completions", MARKED) + finally: + third.shutdown() + cut: Final = dict(cut_short.counters.counts) turned_down: Final = dict(unfinished.counters.counts) - assert cut["mount:openai:rejected"] == 1 and turned_down["mount:openai:rejected"] == 1 + errored: Final = dict(refused.counters.counts) + assert cut["mount:openai:rejected"] == turned_down["mount:openai:rejected"] == errored["mount:openai:rejected"] == 1 assert cut["mount:openai:rejected_cut_short"] == 1 - assert "mount:openai:rejected_incomplete" not in cut assert turned_down["mount:openai:rejected_incomplete"] == 1 - assert "mount:openai:rejected_cut_short" not in turned_down + assert errored["mount:openai:rejected_error_status"] == 1 + assert not {"mount:openai:rejected_incomplete", "mount:openai:rejected_error_status"} & set(cut) + assert not {"mount:openai:rejected_cut_short", "mount:openai:rejected_error_status"} & set(turned_down) + assert not {"mount:openai:rejected_cut_short", "mount:openai:rejected_incomplete"} & set(errored) EMBEDDING_SUCCESS: Final = ( diff --git a/tests/e2e/PROVIDER_CACHE.md b/tests/e2e/PROVIDER_CACHE.md index 894d9be4efa..ed1e7d517b0 100644 --- a/tests/e2e/PROVIDER_CACHE.md +++ b/tests/e2e/PROVIDER_CACHE.md @@ -54,7 +54,7 @@ The trusted runner receives: - `E2E_PROVIDER_CACHE_NAMESPACE`: shared environment namespace, independent of build and candidate revision - `E2E_PROVIDER_CACHE_METRICS_DIR`: optional per-process counter artifact directory -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. A rejection also counts its reason, one of `rejected_cut_short` (the consumer walked away mid-capture), `rejected_incomplete` (the body arrived whole and failed its endpoint's rule) or `rejected_unreachable` (the provider could not be reached). A mount whose rejections are nearly all one or the other is a different problem, and the flat count cannot tell them apart. Every counter is emitted twice, once as a flat total and once under `mount:{mount}:`, so a hit rate can be read per provider rather than only in aggregate. Existing application-cache observations still count requests arriving at the edge, including shared-cache hits +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. A rejection also counts its reason, one of `rejected_cut_short` (the consumer walked away mid-capture), `rejected_error_status` (the provider answered, with an error), `rejected_incomplete` (the body arrived whole with a success status and failed its endpoint's rule) or `rejected_unreachable` (the provider could not be reached at all). A mount whose rejections are nearly all one or the other is a different problem, and the flat count cannot tell them apart. Every counter is emitted twice, once as a flat total and once under `mount:{mount}:`, so a hit rate can be read per provider rather than only in aggregate. 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. The provider prompt-caching tests carry it because a replayed priming response reports cache creation rather than a cache read. diff --git a/tests/e2e/provider_cache.py b/tests/e2e/provider_cache.py index f528d08b720..399a0379889 100644 --- a/tests/e2e/provider_cache.py +++ b/tests/e2e/provider_cache.py @@ -55,6 +55,7 @@ EVENTSTREAM_PRELUDE_BYTES: Final = 4 CUT_SHORT: Final = "cut_short" INCOMPLETE: Final = "incomplete" UNREACHABLE: Final = "unreachable" +ERROR_STATUS: Final = "error_status" EVENT_TYPE_HEADER: Final = ":event-type" EVENTSTREAM_HEADERS: Final[TypeAdapter[dict[str, str]]] = TypeAdapter(dict[str, str]) OPENAI_JSON_PATHS: Final = frozenset({"/v1/chat/completions", "/v1/messages", "/v1/embeddings", "/v1/responses"}) @@ -610,6 +611,8 @@ class CacheEdge: headers: Final = { name: value for name, value in head.headers.items() if name.lower() not in UNRECORDED_RESPONSE_HEADERS } + if not 200 <= head.status_code < 300: + return ERROR_STATUS chunks: Final = capture.chunks() if not successful_response(mount, url, head.status_code, headers, b"".join(chunks)): return INCOMPLETE From 1972a30defcea23d170ba431af99f2e82e652b24 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 16 Sep 2026 08:37:25 -0700 Subject: [PATCH 13/21] revert(e2e): unmount Gemini, its api_base means two things Build 227 mounted Gemini and turned TestGeminiFiles::test_gemini_file_upload red. litellm's two Gemini endpoints disagree about what api_base means. Chat composes {api_base}/models/{model}:{endpoint} and defaults api_base to https://generativelanguage.googleapis.com/v1beta, so the version lives inside it. File upload composes {api_base}/upload/v1beta/files and defaults to the host root, so the version lives outside it. A single api_base cannot satisfy both, and a registration carries no signal about which endpoint the deployment will be used for, so the edge cannot route one and not the other. Backing it out rather than working around it. The cache must never turn a passing test red, which is the same rule the Bedrock model allowlist follows, and Gemini was 7 of roughly 1030 edge calls in that build. Anyone pointing litellm's Gemini provider at an AI gateway or a corporate proxy hits this too, so the fix belongs in litellm; mounting Gemini is one line once it lands. This reverts commit 8a553ceb58887c8aa2aa24c7cfb75ce662d2e321. --- .../test_provider_cache.py | 154 +----------------- tests/e2e/PROVIDER_CACHE.md | 18 +- tests/e2e/provider_cache.py | 46 ------ tests/e2e/provider_cache_routing.py | 3 +- tests/e2e/provider_edge.py | 1 - 5 files changed, 7 insertions(+), 215 deletions(-) diff --git a/tests/code_coverage_tests/test_provider_cache.py b/tests/code_coverage_tests/test_provider_cache.py index 9668de35a02..e57d33a0406 100644 --- a/tests/code_coverage_tests/test_provider_cache.py +++ b/tests/code_coverage_tests/test_provider_cache.py @@ -912,7 +912,7 @@ def test_anthropic_stream_requires_start_finish_and_stop() -> None: assert not successful_response("anthropic", url, 200, headers, start + finish) -@pytest.mark.parametrize("provider,suffix", [("openai", "/v1"), ("anthropic", ""), ("gemini", "")]) +@pytest.mark.parametrize("provider,suffix", [("openai", "/v1"), ("anthropic", "")]) def test_normal_registration_routes_supported_providers(provider: str, suffix: str) -> None: params: Final = LiteLLMParamsBody(model=f"{provider}/test", api_key="os.environ/SYNTHETIC_KEY", timeout=12) routed: Final = route_cache_model(params, lambda mount: f"http://edge.invalid/{mount}", enabled=True) @@ -924,8 +924,6 @@ def test_normal_registration_routes_supported_providers(provider: str, suffix: s @pytest.mark.parametrize("params", [ LiteLLMParamsBody(model="bedrock/test"), LiteLLMParamsBody(model="azure/test"), - LiteLLMParamsBody(model="vertex_ai/gemini-2.5-flash"), - LiteLLMParamsBody(model="gemini/gemini-2.5-flash", api_base="https://custom.invalid"), LiteLLMParamsBody(model="openai/test", api_base="https://custom.invalid/v1"), LiteLLMParamsBody(model="openai/test", api_base=""), LiteLLMParamsBody(model="openai/test", litellm_credential_name="named-credential"), @@ -1294,153 +1292,3 @@ class TestBedrockStreams: def test_a_stream_that_errored_before_it_started_is_not_recordable(self) -> None: assert not successful_response(BEDROCK_MOUNT, CONVERSE_STREAM_URL, 503, {}, CONVERSE_STREAM_OK) - - -GEMINI_MODEL: Final = "gemini-2.5-flash" -GEMINI_API_VERSION: Final = "/v1beta" -GEMINI_GENERATE_PATH: Final = f"/models/{GEMINI_MODEL}:generateContent" -GEMINI_STREAM_PATH: Final = f"/models/{GEMINI_MODEL}:streamGenerateContent" -GEMINI_USAGE: Final = {"promptTokenCount": 7, "candidatesTokenCount": 1, "totalTokenCount": 25} - - -def gemini_body(finish_reason: str | None, usage: bool = True, candidates: bool = True) -> JsonValue: - candidate: Final[dict[str, JsonValue]] = {"content": {"parts": [{"text": "OK"}], "role": "model"}, "index": 0} - return { - "candidates": [{**candidate, "finishReason": finish_reason} if finish_reason else candidate] - if candidates else [], - **({"usageMetadata": GEMINI_USAGE} if usage else {}), - "modelVersion": GEMINI_MODEL, - } - - -def gemini_unary(finish_reason: str | None = "STOP", usage: bool = True, candidates: bool = True) -> bytes: - return json.dumps(gemini_body(finish_reason, usage, candidates)).encode() - - -def gemini_stream(*finish_reasons: str | None) -> bytes: - return b"".join( - b"data: " + json.dumps(gemini_body(reason)).encode() + b"\r\n\r\n" for reason in finish_reasons - ) - - -@contextmanager -def gemini_edge(cache: CacheEdge, provider: Provider, path: str) -> Generator[str, None, None]: - upstream: Final = f"http://127.0.0.1:{provider.server_port}{GEMINI_API_VERSION}" - running: Final = start_provider_edge(cache, mounts={"gemini": upstream}) - try: - yield running.edge.api_base("gemini") + path - finally: - running.shutdown() - - -class TestGemini: - """Gemini reaches the edge by path prefix alone: litellm composes - `{api_base}/models/{model}:{endpoint}` and sends a static `x-goog-api-key`, - so nothing has to be re-signed and nothing leaves the cache key. The response - grammar is its own though, and the streaming one is the interesting half: every - chunk repeats `usageMetadata`, so only `finishReason` on the last chunk - separates a finished turn from a dropped connection.""" - - @pytest.mark.parametrize("path,response", [ - (GEMINI_GENERATE_PATH, gemini_unary()), - (GEMINI_STREAM_PATH, gemini_stream(None, None, "STOP")), - ], ids=["generate", "stream"]) - def test_a_finished_turn_replays_on_the_next_run( - self, store: RedisResponseStore, provider: Provider, path: str, response: bytes, - ) -> None: - provider.stream = path == GEMINI_STREAM_PATH - provider.response = response - for _ in range(2): - with gemini_edge(cache_edge(store), provider, path) as url: - assert call(url, MARKED).body == response - assert len(provider.hits) == 1 - - @pytest.mark.parametrize("reason", ["MAX_TOKENS", "SAFETY", "RECITATION"]) - def test_a_turn_the_provider_ended_for_its_own_reasons_is_still_finished( - self, store: RedisResponseStore, provider: Provider, reason: str, - ) -> None: - """Reading `finishReason` as a string rather than comparing it to STOP is - deliberate. A turn cut off by the token limit or a safety filter is over, - and rejecting those would send every one of them upstream forever.""" - provider.response = gemini_unary(reason) - for _ in range(2): - with gemini_edge(cache_edge(store), provider, GEMINI_GENERATE_PATH) as url: - assert call(url, MARKED).body == provider.response - assert len(provider.hits) == 1 - - @pytest.mark.parametrize("response", [ - gemini_unary(None), - gemini_unary("STOP", usage=False), - gemini_unary("STOP", candidates=False), - b'{"error":{"code":400,"message":"API key not valid","status":"INVALID_ARGUMENT"}}', - ], ids=["no-finish-reason", "no-usage", "no-candidates", "error-body"]) - def test_an_unfinished_or_failed_turn_never_enters_the_cache( - self, store: RedisResponseStore, provider: Provider, response: bytes, - ) -> None: - provider.response = response - for _ in range(2): - with gemini_edge(cache_edge(store), provider, GEMINI_GENERATE_PATH) as url: - assert call(url, MARKED).body == response - assert len(provider.hits) == 2 - - @pytest.mark.parametrize("response", [ - gemini_stream(None, None), - gemini_stream("STOP", None), - gemini_stream(), - ], ids=["cut-before-the-reason", "reason-then-another-chunk", "empty"]) - def test_a_stream_that_never_named_a_reason_calls_the_provider_every_time( - self, store: RedisResponseStore, provider: Provider, response: bytes, - ) -> None: - provider.stream = True - provider.response = response - for _ in range(2): - with gemini_edge(cache_edge(store), provider, GEMINI_STREAM_PATH) as url: - assert call(url, MARKED).body == response - assert len(provider.hits) == 2 - - def test_a_response_whose_candidates_did_not_all_finish_is_not_recordable( - self, store: RedisResponseStore, provider: Provider, - ) -> None: - """A request for more than one candidate is answered by more than one, and - the turn is over only when every one of them names a reason. Holding the - whole list to that rule rather than its first entry is what keeps a - half-finished answer from being stored and replayed as a finished one.""" - finished: Final = json.loads(gemini_unary("STOP"))["candidates"][0] - unfinished: Final = json.loads(gemini_unary(None))["candidates"][0] - provider.response = json.dumps( - {"candidates": [finished, {**unfinished, "index": 1}], "usageMetadata": GEMINI_USAGE} - ).encode() - for _ in range(2): - with gemini_edge(cache_edge(store), provider, GEMINI_GENERATE_PATH) as url: - assert call(url, MARKED).body == provider.response - assert len(provider.hits) == 2 - - @pytest.mark.parametrize("path,cacheable", [ - (GEMINI_GENERATE_PATH, True), - (GEMINI_STREAM_PATH, True), - (f"/models/{GEMINI_MODEL}:countTokens", False), - (f"/models/{GEMINI_MODEL}:embedContent", False), - ("/v1/chat/completions", False), - (f"/files/{GEMINI_MODEL}:generateContent", False), - ]) - @pytest.mark.parametrize("version", ["", GEMINI_API_VERSION], ids=["bare", "versioned"]) - def test_only_the_generate_endpoints_are_cacheable(self, version: str, path: str, cacheable: bool) -> None: - """The mount's upstream base carries the API version, so the path the cache - sees is the upstream one and starts `/v1beta`. A rule anchored at the start - of the path would pass every test against a stub with no version prefix and - then cache nothing at all in a real run.""" - assert cacheable_endpoint("gemini", "POST", f"https://gemini.invalid{version}{path}", MARKED) is cacheable - - def test_the_bodies_these_tests_build_match_a_real_gemini_response(self) -> None: - """The shapes above are hand-built so a test can express the turn it means. - This holds them to the fields a live `generativelanguage.googleapis.com` - answer carries, captured 2026-09-16 against gemini-2.5-flash.""" - captured: Final = json.loads( - '{"candidates":[{"content":{"parts":[{"text":"OK"}],"role":"model"},"finishReason":"STOP",' - '"index":0}],"usageMetadata":{"promptTokenCount":7,"candidatesTokenCount":1,' - '"totalTokenCount":25},"modelVersion":"gemini-2.5-flash","responseId":"1J6qauKFI8ut1MkPgNjI4AI"}' - ) - built: Final = json.loads(gemini_unary()) - assert captured.keys() >= built.keys() - assert captured["candidates"][0].keys() >= built["candidates"][0].keys() - assert successful_response("gemini", GEMINI_GENERATE_PATH, 200, {}, json.dumps(captured).encode()) diff --git a/tests/e2e/PROVIDER_CACHE.md b/tests/e2e/PROVIDER_CACHE.md index ed1e7d517b0..aca81f26c5a 100644 --- a/tests/e2e/PROVIDER_CACHE.md +++ b/tests/e2e/PROVIDER_CACHE.md @@ -1,23 +1,13 @@ # Shared provider-response cache -`E2E_PROVIDER_CACHE=1` enables automatic response reuse in the live E2E mode. Standard OpenAI, Anthropic and Gemini model registrations use the provider edge, as do Anthropic-on-Bedrock registrations that carry no AWS identity of their own. Existing custom API bases, named credentials, mocked models and realtime WebSocket deployments keep their existing routing. Other provider protocols remain live +`E2E_PROVIDER_CACHE=1` enables automatic response reuse in the live E2E mode. Standard OpenAI and Anthropic model registrations use the provider edge, as do Anthropic-on-Bedrock registrations that carry no AWS identity of their own. 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`, `/v1/messages`, `/v1/embeddings` and `/v1/responses` on the OpenAI and Anthropic mounts, SSE streams included, for `/model/{id}/converse` and `/model/{id}/invoke` on a Bedrock mount, and for `models/{model}:generateContent` and `:streamGenerateContent` on the Gemini mount. Unsupported endpoints pass through. Each endpoint family has its own completeness rule, so a truncated embedding or a Responses run that never reached `response.completed` is not stored +The edge caches complete successful POST responses for `/v1/chat/completions`, `/v1/messages`, `/v1/embeddings` and `/v1/responses` on the OpenAI and Anthropic mounts, SSE streams included, and for `/model/{id}/converse` and `/model/{id}/invoke` on a Bedrock mount. Unsupported endpoints pass through. Each endpoint family has its own completeness rule, so a truncated embedding or a Responses run that never reached `response.completed` is not stored Bedrock's streaming endpoints, `converse-stream` and `invoke-with-response-stream`, cache too. AWS frames those as binary `vnd.amazon.eventstream` rather than SSE, so botocore's own parser reads the frames and validates both CRCs, and each endpoint is then held to its terminal grammar. That matters more than the endpoint count suggests: the Claude Code compat cells drive the real CLI, which always streams, so streaming is most of the suite's Bedrock traffic Two details of that rule are worth knowing before changing it. A ConverseStream ends with `metadata`, not with `messageStop`, and the `metadata` frame is what carries the token usage litellm prices the call from, so the rule requires it: a stream cut between the two still names a stop reason but would replay as a free call. And a dropped connection is invisible to the parser, which yields the frames it did receive and silently discards a trailing partial one, so the body is also checked against the frame lengths it declares. A stream cut one byte short parses clean and has to be caught that way -## Gemini - -Gemini needs nothing that Bedrock needed. litellm composes `{api_base}/models/{model}:{endpoint}` from a custom api_base, so a path-prefixed mount reaches it, and the credential travels as a static `x-goog-api-key` header that no host rewrite invalidates. Nothing is re-signed and nothing is excluded from the key, so a recording still cannot cross credentials - -The mount's upstream base carries the API version, which is the one detail worth remembering: the path the cache rules see is the upstream one, `/v1beta/models/...`, not the one the proxy sent. A rule anchored at the start of that path would look right against a local stub and then cache nothing at all in a real run - -A finished turn names a `finishReason` on every candidate and reports `usageMetadata`. The reason is read as a string rather than compared to `STOP`, because `MAX_TOKENS` and the safety reasons end a turn just as finally and rejecting them would send every one of them upstream forever. Streaming is the more interesting half: Gemini repeats `usageMetadata` on every chunk and names a `finishReason` only on the last one, so the terminator is the final event rather than any event, and a stream the connection cut short ends on a chunk with usage and no reason - -Vertex is not mounted. litellm grafts the default Vertex path onto an api_base only when that api_base has no path of its own, so a Vertex mount needs a root-mounted edge on its own port rather than a path prefix. Gemini and Vertex are separate providers in litellm and the Gemini mount does not cover Vertex deployments - ## Request identity A recording belongs to one test. The key is a keyed digest over the test's node id, the method, the URL, the effective outbound headers (including authentication and HTTP-library defaults), body presence and the body bytes, with one normalization: a 12-hex-digit run, the shape `unique_marker()` mints, is replaced by a placeholder in both the URL and a UTF-8 body. Nothing else is normalized away. No prompts, JSON values or credentials are rewritten, and the rule is the one `fixture_canonical.py` already applies for record/replay, so there is a single definition of what a marker is @@ -40,7 +30,9 @@ Only deployments that carry no AWS identity of their own route to the edge. A de Which models route is an explicit allowlist in `provider_cache_routing.py`, mirroring the runner role's IAM policy, which names its models one by one. That coupling is deliberate: the edge re-signs with the run pod's identity, so a model the role cannot invoke comes back 403 from Bedrock rather than falling back. An unlisted model keeps its direct path and loses only caching, so adding a Bedrock model to the suite can never turn it red. Adding one to the edge is a policy edit in litellm-ops plus a line here -Vertex is not mounted. litellm's `_check_custom_proxy` rewrites a path-prefixed Vertex `api_base` into `{api_base}:{endpoint}`, dropping project, location and model, so a mount under a path prefix cannot work without either a root-mounted edge on its own port or a change in litellm. Gemini is a separate provider there and does have a working path-prefixed form, so it is mounted; see the Gemini section +Vertex and Gemini are not mounted, for different reasons. litellm grafts the default Vertex path onto an `api_base` only when that `api_base` has no path of its own, so a path-prefixed Vertex mount instead becomes `{api_base}:{endpoint}`, dropping project, location and model. Vertex needs a root-mounted edge on its own port, or a change in litellm + +Gemini reaches a path-prefixed mount perfectly well and was mounted for one build, then backed out, because litellm's two Gemini endpoints disagree about what `api_base` means. Chat composes `{api_base}/models/{model}:{endpoint}` and defaults `api_base` to `https://generativelanguage.googleapis.com/v1beta`, so the version has to be inside it. File upload composes `{api_base}/upload/v1beta/files` and defaults to the host root, so the version has to be outside it. One `api_base` cannot satisfy both, and a deployment gives no signal at registration time about which it will be used for, so mounting Gemini turned `TestGeminiFiles::test_gemini_file_upload` red in build 227. Anyone pointing litellm's Gemini provider at an AI gateway or a corporate proxy hits the same thing; it is a litellm bug rather than a cache limitation, and mounting Gemini is one line once it is fixed Recordings are shared across workers and builds through dedicated Redis, separate from the candidate's own cache. They expire 86,400 seconds after capture starts, based on Redis time. Reads never extend expiry. There is no scheduled recapture: the next miss calls the provider again. Bounded coordination reduces duplicate concurrent calls, but slow or failed captures may lead to extra live calls after the wait expires diff --git a/tests/e2e/provider_cache.py b/tests/e2e/provider_cache.py index 399a0379889..444972ffce5 100644 --- a/tests/e2e/provider_cache.py +++ b/tests/e2e/provider_cache.py @@ -37,10 +37,6 @@ SIGNATURE_HEADERS: Final = frozenset( {"authorization", "x-amz-date", "x-amz-security-token", "x-amz-content-sha256"} ) BEDROCK_MOUNT_PREFIX: Final = "bedrock" -GEMINI_MOUNT: Final = "gemini" -GEMINI_MODELS_SEGMENT: Final = "/models" -GEMINI_GENERATE_SUFFIX: Final = ":generateContent" -GEMINI_STREAM_SUFFIX: Final = ":streamGenerateContent" BEDROCK_CONVERSE_SUFFIX: Final = "/converse" BEDROCK_INVOKE_SUFFIX: Final = "/invoke" BEDROCK_CONVERSE_STREAM_SUFFIX: Final = "/converse-stream" @@ -162,21 +158,12 @@ def is_bedrock(mount: str) -> bool: return mount.partition("/")[0] == BEDROCK_MOUNT_PREFIX -def is_gemini(mount: str) -> bool: - return mount == GEMINI_MOUNT - - def cacheable_endpoint(mount: str, method: str, url: str, body: bytes | None) -> bool: if method != "POST" or body is None or len(body) > MAX_REQUEST_BYTES: return False path: Final = urlsplit(url).path if is_bedrock(mount): return path.startswith("/model/") and path.endswith(BEDROCK_SUFFIXES) - if is_gemini(mount): - collection, _, resource = path.rpartition("/") - return collection.endswith(GEMINI_MODELS_SEGMENT) and resource.endswith( - (GEMINI_GENERATE_SUFFIX, GEMINI_STREAM_SUFFIX) - ) return path in OPENAI_JSON_PATHS @@ -203,8 +190,6 @@ def successful_response(mount: str, url: str, status: int, headers: Mapping[str, for value in values ): return False - if is_gemini(mount): - return complete_gemini_stream(values) if urlsplit(url).path == "/v1/responses": return complete_responses_stream(values) if urlsplit(url).path == "/v1/chat/completions": @@ -216,8 +201,6 @@ def successful_response(mount: str, url: str, status: int, headers: Mapping[str, return False if not isinstance(value, dict) or value.get("error") is not None: return False - if is_gemini(mount): - return complete_gemini_candidates(value) path: Final = urlsplit(url).path if path == "/v1/messages": return value.get("type") == "message" and isinstance(value.get("content"), list) and isinstance(value.get("stop_reason"), str) @@ -236,35 +219,6 @@ def successful_response(mount: str, url: str, status: int, headers: Mapping[str, ) -def complete_gemini_candidates(value: Mapping[str, JsonValue]) -> bool: - """A finished Gemini turn names a ``finishReason`` on every candidate and - reports the usage litellm prices the call from. ``finishReason`` is read as a - string rather than compared to ``STOP`` because ``MAX_TOKENS`` and the safety - reasons end a turn just as finally, and a cache that rejected them would send - every one of them upstream forever.""" - candidates: Final = value.get("candidates") - return ( - isinstance(value.get("usageMetadata"), dict) - and isinstance(candidates, list) - and bool(candidates) - and all( - isinstance(candidate, dict) and isinstance(candidate.get("finishReason"), str) - for candidate in candidates - ) - ) - - -def complete_gemini_stream(values: tuple[JsonValue, ...]) -> bool: - """Gemini repeats ``usageMetadata`` on every chunk but names a - ``finishReason`` only on the last one, so the terminator is the final event - rather than any event. A stream the connection cut short ends on a chunk that - carries usage and no reason, which is exactly what this rejects.""" - if not values: - return False - last: Final = values[-1] - return isinstance(last, dict) and complete_gemini_candidates(last) - - def complete_bedrock_response(url: str, body: bytes) -> bool: """Converse answers with ``output`` plus a ``stopReason``; InvokeModel on an Anthropic model answers the Anthropic message shape. Either way a truncated diff --git a/tests/e2e/provider_cache_routing.py b/tests/e2e/provider_cache_routing.py index c4b02beac2d..f9775a2b152 100644 --- a/tests/e2e/provider_cache_routing.py +++ b/tests/e2e/provider_cache_routing.py @@ -19,7 +19,6 @@ BEDROCK_EDGE_MODELS: Final = frozenset( } ) ENV_REFERENCE_PREFIX: Final = "os.environ/" -EDGE_PROVIDERS: Final = frozenset({"openai", "anthropic", "gemini"}) def bedrock_region(declared: str | None) -> str: @@ -82,7 +81,7 @@ def route_cache_model( return route_bedrock(params, base_for, mode) if mode == "realtime" or params.api_base is not None: return params - if provider not in EDGE_PROVIDERS: + if provider not in {"openai", "anthropic"}: return params base: Final = base_for(provider) if base is None: diff --git a/tests/e2e/provider_edge.py b/tests/e2e/provider_edge.py index 219df55233a..2606b26fe99 100644 --- a/tests/e2e/provider_edge.py +++ b/tests/e2e/provider_edge.py @@ -104,7 +104,6 @@ EDGE_MOUNTS: Final[Mapping[str, str]] = MappingProxyType( { "openai": "https://api.openai.com", "anthropic": "https://api.anthropic.com", - "gemini": "https://generativelanguage.googleapis.com/v1beta", **{ f"bedrock/{region}": f"https://bedrock-runtime.{region}.amazonaws.com" for region in BEDROCK_REGIONS From 7d42bc751debd3ffa7eaa6166d20f627cee0f067 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 16 Sep 2026 10:05:13 -0700 Subject: [PATCH 14/21] docs(e2e): name all four rejection reasons in the counter note --- tests/e2e/PROVIDER_CACHE.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/e2e/PROVIDER_CACHE.md b/tests/e2e/PROVIDER_CACHE.md index aca81f26c5a..e06b3c01653 100644 --- a/tests/e2e/PROVIDER_CACHE.md +++ b/tests/e2e/PROVIDER_CACHE.md @@ -46,7 +46,7 @@ The trusted runner receives: - `E2E_PROVIDER_CACHE_NAMESPACE`: shared environment namespace, independent of build and candidate revision - `E2E_PROVIDER_CACHE_METRICS_DIR`: optional per-process counter artifact directory -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. A rejection also counts its reason, one of `rejected_cut_short` (the consumer walked away mid-capture), `rejected_error_status` (the provider answered, with an error), `rejected_incomplete` (the body arrived whole with a success status and failed its endpoint's rule) or `rejected_unreachable` (the provider could not be reached at all). A mount whose rejections are nearly all one or the other is a different problem, and the flat count cannot tell them apart. Every counter is emitted twice, once as a flat total and once under `mount:{mount}:`, so a hit rate can be read per provider rather than only in aggregate. Existing application-cache observations still count requests arriving at the edge, including shared-cache hits +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. A rejection also counts its reason, one of `rejected_cut_short` (the consumer walked away mid-capture), `rejected_error_status` (the provider answered, with an error), `rejected_incomplete` (the body arrived whole with a success status and failed its endpoint's rule) or `rejected_unreachable` (the provider could not be reached at all). A mount whose rejections are nearly all of one kind is a different problem from one whose rejections are nearly all of another, and the flat count cannot tell them apart. Every counter is emitted twice, once as a flat total and once under `mount:{mount}:`, so a hit rate can be read per provider rather than only in aggregate. 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. The provider prompt-caching tests carry it because a replayed priming response reports cache creation rather than a cache read. From 39acea0754e0dd5f291e8f99126b1e0b3f5505b3 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 16 Sep 2026 12:42:15 -0700 Subject: [PATCH 15/21] feat(e2e): make Claude Code send the same bytes every build The compat cells drove the CLI with a fresh HOME per invocation and the pytest process's own working directory, and both reach the request body. The system prompt names a memory directory built from $CLAUDE_CONFIG_DIR/projects/, so a per-invocation config directory rewrote every body, and the CLI adds a git block for its working directory, so inheriting the checkout rewrote every body once per candidate. The device id churned for the same reason: the CLI mints it once and persists it in .claude.json, which we threw away each call. Nothing here was load-bearing. All three ride in metadata.user_id, whose job is abuse detection, not quota, caching or continuity. So pin the config directory and the working directory at fixed paths, seed the device id, and pin the session id. HOME stays fresh and empty per invocation, so the isolation is no weaker than before, and the CLI's own state no longer outlives the pod either. The working directory is deliberately not the checkout, so a model-directed Read now sees an empty directory rather than the repository. A pinned session id needs --no-session-persistence beside it: the CLI refuses a session id another live process holds, and the matrix runs its cells across xdist workers. Without the flag, six of eight concurrent invocations die on "Session ID is already in use". --- .../test_request_determinism.py | 143 ++++++++++++++++++ tests/e2e/claude_code/cli_driver.py | 57 +++++++ 2 files changed, 200 insertions(+) create mode 100644 tests/e2e/claude_code/_driver_unit_tests/test_request_determinism.py diff --git a/tests/e2e/claude_code/_driver_unit_tests/test_request_determinism.py b/tests/e2e/claude_code/_driver_unit_tests/test_request_determinism.py new file mode 100644 index 00000000000..09a181162d2 --- /dev/null +++ b/tests/e2e/claude_code/_driver_unit_tests/test_request_determinism.py @@ -0,0 +1,143 @@ +"""The CLI must send the same request bytes from one build to the next. + +Markerless harness test: it drives the real `claude` binary against a local +stub instead of a proxy, so it carries no `e2e` marker. The binary is a +prerequisite of this whole suite, so a missing one is a failure rather than a +skip. + +Two builds differ in ways the driver does not control: a fresh pod, so no CLI +state survives, and a different candidate checked out at a different commit. +Both used to reach the request body, through the memory path the system prompt +names and through the git block the CLI adds for its working directory, so the +shared provider cache missed on every Claude Code cell. This replays those two +differences across a pair of invocations and holds the bytes equal. + +A pinned session id is what makes the second test necessary. The matrix runs +its cells across xdist workers, and the CLI refuses to start a session id that +another live process already holds, so pinning one without also opting out of +session persistence turns most of a parallel run red. +""" + +from __future__ import annotations + +import json +import os +import shutil +import subprocess +import threading +from collections import Counter +from concurrent.futures import ThreadPoolExecutor +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from pathlib import Path +from typing import List, Tuple + +import pytest + +from claude_code.cli_driver import _stable_cli_state, run_claude +from claude_code.rate_limiter import RateLimiter + +_STUB_REPLY = { + "id": "msg_stub", + "type": "message", + "role": "assistant", + "model": "claude-haiku-4-5", + "content": [{"type": "text", "text": "ok"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 10, "output_tokens": 2}, +} + + +def _make_repo(root: Path, subject: str) -> Path: + root.mkdir(parents=True, exist_ok=True) + identity = {"NAME": "t", "EMAIL": "t@e2e"} + env = dict( + os.environ, + **{f"GIT_{role}_{key}": value for role in ("AUTHOR", "COMMITTER") for key, value in identity.items()}, + ) + (root / "file.txt").write_text(subject, encoding="utf-8") + for args in (["init", "-q"], ["add", "."], ["commit", "-q", "-m", subject]): + subprocess.run(["git", *args], cwd=root, env=env, check=True, capture_output=True) + return root + + +@pytest.fixture(name="captured") +def _captured() -> Tuple[str, List[bytes]]: + bodies: List[bytes] = [] + lock = threading.Lock() + + class Handler(BaseHTTPRequestHandler): + protocol_version = "HTTP/1.1" + + def do_POST(self) -> None: + raw = self.rfile.read(int(self.headers.get("content-length") or 0)) + if "count_tokens" not in self.path: + with lock: + bodies.append(raw) + payload = json.dumps({"input_tokens": 10} if "count_tokens" in self.path else _STUB_REPLY).encode() + self.send_response(200) + self.send_header("content-type", "application/json") + self.send_header("content-length", str(len(payload))) + self.end_headers() + self.wfile.write(payload) + + def log_message(self, *_args: object) -> None: + return + + server = ThreadingHTTPServer(("127.0.0.1", 0), Handler) + threading.Thread(target=server.serve_forever, daemon=True).start() + try: + yield f"http://127.0.0.1:{server.server_address[1]}", bodies + finally: + server.shutdown() + + +def test_two_builds_send_the_same_request_bytes(captured: Tuple[str, List[bytes]], tmp_path: Path) -> None: + base_url, bodies = captured + limiter = RateLimiter(state_dir=tmp_path / "limiter") + checkouts = (_make_repo(tmp_path / "build-1", "first"), _make_repo(tmp_path / "build-2", "second")) + origin = Path.cwd() + + sent = [] + for checkout in checkouts: + shutil.rmtree(Path(_stable_cli_state()[0]).parent, ignore_errors=True) + os.chdir(checkout) + try: + before = len(bodies) + run_claude( + prompt="say ok", + model="claude-haiku-4-5", + base_url=base_url, + api_key="stub", + extra_env={"CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC": "1"}, + rate_limiter=limiter, + ) + sent.append(bodies[before:]) + finally: + os.chdir(origin) + + assert sent[0], "the CLI sent no request to the stub, so there is nothing to compare" + assert sent[0] == sent[1] + + +def test_concurrent_cells_do_not_collide_on_the_pinned_session( + captured: Tuple[str, List[bytes]], tmp_path: Path +) -> None: + base_url, bodies = captured + limiter = RateLimiter(state_dir=tmp_path / "limiter") + + def one(_index: int) -> int: + return run_claude( + prompt="say ok", + model="claude-haiku-4-5", + base_url=base_url, + api_key="stub", + extra_env={"CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC": "1"}, + rate_limiter=limiter, + ).exit_code + + with ThreadPoolExecutor(max_workers=4) as pool: + codes = list(pool.map(one, range(4))) + + assert codes == [0, 0, 0, 0] + assert bodies, "the CLI sent no request to the stub, so there is nothing to compare" + assert set(Counter(bodies).values()) == {4} diff --git a/tests/e2e/claude_code/cli_driver.py b/tests/e2e/claude_code/cli_driver.py index 447e8cc0bbb..3fad87c7479 100644 --- a/tests/e2e/claude_code/cli_driver.py +++ b/tests/e2e/claude_code/cli_driver.py @@ -132,6 +132,57 @@ def _make_isolated_home() -> str: return tempfile.mkdtemp(prefix="claude-cli-home-") +_FIXED_CLI_USER_ID = "0" * 64 +_FIXED_CLI_SESSION_ID = "00000000-0000-4000-8000-000000000000" + + +def _seed_cli_identity(config_dir: str) -> None: + """Pin the device id the CLI would otherwise mint per config directory. + + It mints 32 random bytes on first run, writes them to `.claude.json` as + `userID`, and sends them in `metadata.user_id` forever after, so the value + is stable for exactly as long as that file lives. Pinning it, and the + session id passed beside it, costs nothing: both feed abuse detection + rather than quota, caching or continuity.""" + path = os.path.join(config_dir, ".claude.json") + try: + with open(path, encoding="utf-8") as handle: + if json.load(handle).get("userID") == _FIXED_CLI_USER_ID: + return + except (OSError, ValueError): + pass + staged = f"{path}.{os.getpid()}" + with open(staged, "w", encoding="utf-8") as handle: + json.dump({"userID": _FIXED_CLI_USER_ID}, handle) + os.replace(staged, path) + + +def _stable_cli_state() -> Tuple[str, str]: + """Config directory and working directory for the CLI, at fixed paths. + + Both reach the request body. The memory directory the system prompt + names is `$CLAUDE_CONFIG_DIR/projects//memory`, and a working + directory inside a git repository also contributes its branch and recent + commits. So a per-invocation config directory rewrites every body, and + inheriting the checkout rewrites every body once per candidate, which is + why the shared provider cache could never serve a Claude Code cell. + Pinning both makes the bodies repeatable across builds. + + This narrows what survives rather than widening it: HOME stays fresh and + empty per invocation, so the isolation `_make_isolated_home` describes is + unchanged, and the CLI's own state no longer outlives the pod either. The + working directory is deliberately not the checkout, so a model-directed + `Read` sees an empty directory instead of the repository. + """ + root = os.path.join(tempfile.gettempdir(), f"litellm-e2e-claude-{os.getuid()}") + config_dir = os.path.join(root, "config") + workspace = os.path.join(root, "workspace") + for path in (root, config_dir, workspace): + os.makedirs(path, mode=0o700, exist_ok=True) + _seed_cli_identity(config_dir) + return config_dir, workspace + + class ClaudeCLIError(RuntimeError): """Raised when the `claude` CLI cannot be invoked or returns a fatal error.""" @@ -222,6 +273,9 @@ def run_claude( "--verbose", "--model", model, + "--session-id", + _FIXED_CLI_SESSION_ID, + "--no-session-persistence", ] if extra_args: cmd.extend(extra_args) @@ -244,6 +298,8 @@ def run_claude( # regardless of how the subprocess exits. isolated_home = _make_isolated_home() env["HOME"] = isolated_home + config_dir, workspace = _stable_cli_state() + env["CLAUDE_CONFIG_DIR"] = config_dir if extra_env: env.update(extra_env) @@ -262,6 +318,7 @@ def run_claude( completed = run_fn( cmd, env=env, + cwd=workspace, input=stdin_input, capture_output=True, text=True, From 2481146727fe3b2613df96ccf8f568f5b0a3cc72 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 16 Sep 2026 12:44:53 -0700 Subject: [PATCH 16/21] docs(e2e): say why the CLI-driving cells needed a driver fix, not a rule --- tests/e2e/PROVIDER_CACHE.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/e2e/PROVIDER_CACHE.md b/tests/e2e/PROVIDER_CACHE.md index e06b3c01653..d37b37eeba7 100644 --- a/tests/e2e/PROVIDER_CACHE.md +++ b/tests/e2e/PROVIDER_CACHE.md @@ -16,6 +16,8 @@ Requests that differ only by their markers therefore share a canonical identity, Two different tests never share a recording, and a provider call made outside any test (fixtures, session setup) is never cached, because the identity has no test node id to bind to +A client that varies its own request between runs defeats that identity without breaking any rule, and the Claude Code compat cells did. The CLI sends a device id and a session id in `metadata.user_id`, and its system prompt names both its memory directory and its working directory, adding the branch and recent commits when that directory is a git repository. Driven with a fresh HOME and the checkout as its working directory, every cell sent different bytes every build. The fix belongs in the driver rather than here: `claude_code/cli_driver.py` pins the config directory, the working directory and both identifiers, which is why the cache needs no rule for any of it. Normalizing them instead would have hidden a real defect class, since a rule cannot tell a client's own churn from a value a test means to assert on + 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 From 9421b26bf6dcc98bee10e2cbbd45bf6ff1c23166 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 16 Sep 2026 13:47:07 -0700 Subject: [PATCH 17/21] fix(e2e): stage the seeded device id per thread, not per process Build 232 took two compat cells red with a FileNotFoundError renaming `.claude.json.197` onto `.claude.json`. `run_claude_models_parallel` drives several models from one process, so a pid-suffixed staged name is shared between threads: one thread renamed the file the other was still writing, and the loser died on a path that no longer existed. mkstemp in the same directory gives a name that is unique per thread as well as per process, and the rename stays atomic. --- .../test_request_determinism.py | 19 ++++++++++++++++++- tests/e2e/claude_code/cli_driver.py | 11 ++++++++--- 2 files changed, 26 insertions(+), 4 deletions(-) diff --git a/tests/e2e/claude_code/_driver_unit_tests/test_request_determinism.py b/tests/e2e/claude_code/_driver_unit_tests/test_request_determinism.py index 09a181162d2..5046f35c73b 100644 --- a/tests/e2e/claude_code/_driver_unit_tests/test_request_determinism.py +++ b/tests/e2e/claude_code/_driver_unit_tests/test_request_determinism.py @@ -33,7 +33,7 @@ from typing import List, Tuple import pytest -from claude_code.cli_driver import _stable_cli_state, run_claude +from claude_code.cli_driver import _FIXED_CLI_USER_ID, _seed_cli_identity, _stable_cli_state, run_claude from claude_code.rate_limiter import RateLimiter _STUB_REPLY = { @@ -141,3 +141,20 @@ def test_concurrent_cells_do_not_collide_on_the_pinned_session( assert codes == [0, 0, 0, 0] assert bodies, "the CLI sent no request to the stub, so there is nothing to compare" assert set(Counter(bodies).values()) == {4} + + +def test_seeding_the_device_id_survives_threads_racing_on_the_same_directory(tmp_path: Path) -> None: + """`run_claude_models_parallel` drives several models from one process, so the + seed's staged file has to be unique per thread and not merely per process.""" + config_dir = tmp_path / "config" + config_dir.mkdir() + seeded = config_dir / ".claude.json" + + for _round in range(20): + seeded.unlink(missing_ok=True) + with ThreadPoolExecutor(max_workers=16) as pool: + for outcome in [pool.submit(_seed_cli_identity, str(config_dir)) for _ in range(16)]: + outcome.result() + + assert json.loads(seeded.read_text(encoding="utf-8"))["userID"] == _FIXED_CLI_USER_ID + assert sorted(entry.name for entry in config_dir.iterdir()) == [".claude.json"] diff --git a/tests/e2e/claude_code/cli_driver.py b/tests/e2e/claude_code/cli_driver.py index 3fad87c7479..a01d8ab3e7c 100644 --- a/tests/e2e/claude_code/cli_driver.py +++ b/tests/e2e/claude_code/cli_driver.py @@ -143,7 +143,12 @@ def _seed_cli_identity(config_dir: str) -> None: `userID`, and sends them in `metadata.user_id` forever after, so the value is stable for exactly as long as that file lives. Pinning it, and the session id passed beside it, costs nothing: both feed abuse detection - rather than quota, caching or continuity.""" + rather than quota, caching or continuity. + + The staged name has to be unique per *thread*, not per process: + `run_claude_models_parallel` drives several models from one process, so a + pid-suffixed name lets one thread rename the file another is still + writing, and the loser dies on a missing path.""" path = os.path.join(config_dir, ".claude.json") try: with open(path, encoding="utf-8") as handle: @@ -151,8 +156,8 @@ def _seed_cli_identity(config_dir: str) -> None: return except (OSError, ValueError): pass - staged = f"{path}.{os.getpid()}" - with open(staged, "w", encoding="utf-8") as handle: + handle_fd, staged = tempfile.mkstemp(dir=config_dir, prefix=".claude.json.") + with os.fdopen(handle_fd, "w", encoding="utf-8") as handle: json.dump({"userID": _FIXED_CLI_USER_ID}, handle) os.replace(staged, path) From 76c0f8db1d415a0307f4188a0ca992688ec3b44b Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 16 Sep 2026 15:01:08 -0700 Subject: [PATCH 18/21] chore(e2e): report the key components behind a mount that never converges Builds 232 and 233 held the Bedrock hit rate at 9% with the Claude Code driver already sending byte-identical requests and headers, so something between the proxy's ingress and the upstream still moves per build and the flat key cannot say what. Emit a digest per key component next to the counters: the test id, the method, the URL, each keyed header, the whole body, and one digest per top-level JSON body field. Values are digested, so no payload or credential reaches the artifact. Diffing two builds' artifacts names the field that moved. Diagnostic, to be removed once it has answered. --- tests/e2e/provider_cache.py | 63 +++++++++++++++++++++++++++++-- tests/e2e/provider_cache_redis.py | 4 ++ 2 files changed, 64 insertions(+), 3 deletions(-) diff --git a/tests/e2e/provider_cache.py b/tests/e2e/provider_cache.py index 444972ffce5..22967869c78 100644 --- a/tests/e2e/provider_cache.py +++ b/tests/e2e/provider_cache.py @@ -4,6 +4,8 @@ import base64 import hashlib import hmac import io +import json +import os import threading import time from collections.abc import Callable, Generator, Mapping @@ -400,6 +402,51 @@ def decode_response(secret: bytes, key: str, payload: bytes, mount: str, url: st return response +def component_digests( + test_key: str, method: str, url: str, headers: Mapping[str, str], body: bytes | None, +) -> dict[str, str]: + """Per-component digests of everything the key covers. + + A mount whose corpus never converges is a mount where one of these moves + between builds, and the flat key cannot say which. Values are digested, so + no payload or credential is written, and a JSON body contributes one digest + per top-level field so the field that moved can be named.""" + parts: dict[str, str] = { # rebind-ok: a report assembled from three differently shaped sources + "test_key": test_key, + "method": method, + "url": short_digest(canonical_text(url).encode()), + } + for name, value in sorted(headers.items()): + parts[f"header:{name.lower()}"] = short_digest(value.encode()) + canonical: Final = b"" if body is None else canonical_body(body) + parts["body"] = short_digest(canonical) + try: + parsed: Final = JSON_VALUE.validate_json(canonical) + except ValidationError: + return parts + if isinstance(parsed, dict): + for name, value in sorted(parsed.items()): + parts[f"body:{name}"] = short_digest(json.dumps(value, sort_keys=True).encode()) + return parts + + +def short_digest(value: bytes) -> str: + return hashlib.sha256(value).hexdigest()[:16] + + +@dataclass(slots=True) +class KeyProbe: + """Every keyed request's components, when a metrics directory is configured.""" + + rows: tuple[tuple[tuple[str, str], ...], ...] = () + lock: threading.Lock = field(default_factory=threading.Lock) + + def observe(self, mount: str, outcome: str, parts: Mapping[str, str]) -> None: + row: Final = tuple({"mount": mount, "outcome": outcome, **parts}.items()) + with self.lock: + self.rows = (*self.rows, row) + + @dataclass(slots=True) class CacheCounters: counts: tuple[tuple[str, int], ...] = () @@ -464,6 +511,7 @@ class CacheEdge: store: ResponseStore secret: bytes = field(repr=False) counters: CacheCounters = field(default_factory=CacheCounters) + probe: KeyProbe = field(default_factory=KeyProbe) slots: SlotCounter = field(default_factory=SlotCounter) policies: Mapping[str, MountPolicy] = NO_POLICIES wait_seconds: float = 2.0 @@ -481,6 +529,14 @@ class CacheEdge: self.counters.increment(name) self.counters.increment(f"mount:{mount}:{name}") + def record_key( + self, mount: str, outcome: str, test_key: str, method: str, url: str, + headers: Mapping[str, str], body: bytes | None, + ) -> None: + if not os.environ.get("E2E_PROVIDER_CACHE_METRICS_DIR"): + return + self.probe.observe(mount, outcome, component_digests(test_key, method, url, headers, body)) + def outbound(self, mount: str, method: str, url: str, headers: dict[str, str], body: bytes | None) -> dict[str, str]: """The headers actually sent upstream. A signing mount gets a signature minted over the upstream URL, because the edge rewrote the Host the proxy @@ -511,20 +567,21 @@ class CacheEdge: if isinstance(prepared, NetworkError): self.reject(mount, UNREACHABLE) return prepared - identity: Final = request_identity( - self.secret, test_key, method, url, self.keyed(mount, prepared.headers), body, - ) + keyed_headers: Final = self.keyed(mount, prepared.headers) + identity: Final = request_identity(self.secret, test_key, method, url, keyed_headers, body) key: Final = slotted_key(self.secret, identity, self.slots.take(identity)) found: Final = self.lookup(key) if isinstance(found, CacheHit): response: Final = decode_response(self.secret, key, found.payload, mount, url) if response is not None and self.clock() < found.valid_until: self.count(mount, "hits") + self.record_key(mount, "hit", test_key, method, url, keyed_headers, body) return StreamHead(response.status_code, response.headers, response_steps(response)) self.count(mount, "corrupt" if response is None else "expired") self.store.discard(key, found.payload) capture_slot: Final = self.lookup(key) if isinstance(found, CacheHit) else found self.count(mount, "misses") + self.record_key(mount, "miss", test_key, method, url, keyed_headers, body) if isinstance(capture_slot, CacheUnavailable): self.count(mount, "cache_errors") self.count(mount, "upstream_attempts") diff --git a/tests/e2e/provider_cache_redis.py b/tests/e2e/provider_cache_redis.py index be4e31b2c49..2c7419cfc0f 100644 --- a/tests/e2e/provider_cache_redis.py +++ b/tests/e2e/provider_cache_redis.py @@ -134,6 +134,10 @@ def write_metrics(cache: CacheEdge) -> None: root: Final = Path(directory) root.mkdir(parents=True, exist_ok=True) (root / f"{os.getpid()}.json").write_text(report + "\n") + if cache.probe.rows: + (root / f"keys-{os.getpid()}.json").write_text( + json.dumps([dict(row) for row in cache.probe.rows]) + "\n" + ) except OSError: logging.getLogger(__name__).warning("provider cache metrics artifact unavailable") logging.getLogger(__name__).info("%s", report) From a9fc6d255b5bd82220f3c2265eaf568dcf179423 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 16 Sep 2026 15:15:42 -0700 Subject: [PATCH 19/21] fix(prices): dedupe Nova cache_read_input_token_cost keys left by a text merge Six Amazon Nova entries define cache_read_input_token_cost twice, which is what a clean text merge of two branches that both added the field looks like. JSON parsers keep the last occurrence, so this turned test_price_map_has_no_duplicate_keys red on every open PR's merge commit, including this one, which touches neither file. Both occurrences in all six entries carry the same value, so dropping the later one leaves every parsed price identical. Same change as #41496, carried here so this branch is not blocked on it. Identical deletions, so the two merge cleanly in either order. --- ...model_prices_and_context_window_backup.json | 18 ++++++------------ model_prices_and_context_window.json | 18 ++++++------------ 2 files changed, 12 insertions(+), 24 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index b1c38d0350a..92e5b1c4ff7 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -377,8 +377,7 @@ "supports_prompt_caching": true, "supports_response_schema": true, "supports_vision": true, - "supports_tool_choice": true, - "cache_read_input_token_cost": 1.5e-08 + "supports_tool_choice": true }, "amazon.nova-2-lite-v1:0": { "cache_read_input_token_cost": 7.5e-08, @@ -561,8 +560,7 @@ "supports_function_calling": true, "supports_prompt_caching": true, "supports_response_schema": true, - "supports_tool_choice": true, - "cache_read_input_token_cost": 8.75e-09 + "supports_tool_choice": true }, "amazon.nova-pro-v1:0": { "cache_read_input_token_cost": 2e-07, @@ -578,8 +576,7 @@ "supports_prompt_caching": true, "supports_response_schema": true, "supports_vision": true, - "supports_tool_choice": true, - "cache_read_input_token_cost": 2e-07 + "supports_tool_choice": true }, "amazon.nova-sonic-v1:0": { "deprecation_date": "2026-09-14", @@ -45794,8 +45791,7 @@ "supports_prompt_caching": true, "supports_response_schema": true, "supports_vision": true, - "supports_tool_choice": true, - "cache_read_input_token_cost": 1.5e-08 + "supports_tool_choice": true }, "us.amazon.nova-micro-v1:0": { "cache_read_input_token_cost": 8.75e-09, @@ -45809,8 +45805,7 @@ "supports_function_calling": true, "supports_prompt_caching": true, "supports_response_schema": true, - "supports_tool_choice": true, - "cache_read_input_token_cost": 8.75e-09 + "supports_tool_choice": true }, "us.amazon.nova-premier-v1:0": { "deprecation_date": "2026-09-14", @@ -45842,8 +45837,7 @@ "supports_prompt_caching": true, "supports_response_schema": true, "supports_vision": true, - "supports_tool_choice": true, - "cache_read_input_token_cost": 2e-07 + "supports_tool_choice": true }, "us.anthropic.claude-3-5-haiku-20241022-v1:0": { "cache_creation_input_token_cost": 1e-06, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index b1c38d0350a..92e5b1c4ff7 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -377,8 +377,7 @@ "supports_prompt_caching": true, "supports_response_schema": true, "supports_vision": true, - "supports_tool_choice": true, - "cache_read_input_token_cost": 1.5e-08 + "supports_tool_choice": true }, "amazon.nova-2-lite-v1:0": { "cache_read_input_token_cost": 7.5e-08, @@ -561,8 +560,7 @@ "supports_function_calling": true, "supports_prompt_caching": true, "supports_response_schema": true, - "supports_tool_choice": true, - "cache_read_input_token_cost": 8.75e-09 + "supports_tool_choice": true }, "amazon.nova-pro-v1:0": { "cache_read_input_token_cost": 2e-07, @@ -578,8 +576,7 @@ "supports_prompt_caching": true, "supports_response_schema": true, "supports_vision": true, - "supports_tool_choice": true, - "cache_read_input_token_cost": 2e-07 + "supports_tool_choice": true }, "amazon.nova-sonic-v1:0": { "deprecation_date": "2026-09-14", @@ -45794,8 +45791,7 @@ "supports_prompt_caching": true, "supports_response_schema": true, "supports_vision": true, - "supports_tool_choice": true, - "cache_read_input_token_cost": 1.5e-08 + "supports_tool_choice": true }, "us.amazon.nova-micro-v1:0": { "cache_read_input_token_cost": 8.75e-09, @@ -45809,8 +45805,7 @@ "supports_function_calling": true, "supports_prompt_caching": true, "supports_response_schema": true, - "supports_tool_choice": true, - "cache_read_input_token_cost": 8.75e-09 + "supports_tool_choice": true }, "us.amazon.nova-premier-v1:0": { "deprecation_date": "2026-09-14", @@ -45842,8 +45837,7 @@ "supports_prompt_caching": true, "supports_response_schema": true, "supports_vision": true, - "supports_tool_choice": true, - "cache_read_input_token_cost": 2e-07 + "supports_tool_choice": true }, "us.anthropic.claude-3-5-haiku-20241022-v1:0": { "cache_creation_input_token_cost": 1e-06, From b9751a38ab0a21f56703569b18b3db084a6fa1ae Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 16 Sep 2026 15:36:54 -0700 Subject: [PATCH 20/21] Revert "fix(prices): dedupe Nova cache_read_input_token_cost keys left by a text merge" This reverts commit a9fc6d255b5bd82220f3c2265eaf568dcf179423. --- ...model_prices_and_context_window_backup.json | 18 ++++++++++++------ model_prices_and_context_window.json | 18 ++++++++++++------ 2 files changed, 24 insertions(+), 12 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 92e5b1c4ff7..b1c38d0350a 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -377,7 +377,8 @@ "supports_prompt_caching": true, "supports_response_schema": true, "supports_vision": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "cache_read_input_token_cost": 1.5e-08 }, "amazon.nova-2-lite-v1:0": { "cache_read_input_token_cost": 7.5e-08, @@ -560,7 +561,8 @@ "supports_function_calling": true, "supports_prompt_caching": true, "supports_response_schema": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "cache_read_input_token_cost": 8.75e-09 }, "amazon.nova-pro-v1:0": { "cache_read_input_token_cost": 2e-07, @@ -576,7 +578,8 @@ "supports_prompt_caching": true, "supports_response_schema": true, "supports_vision": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "cache_read_input_token_cost": 2e-07 }, "amazon.nova-sonic-v1:0": { "deprecation_date": "2026-09-14", @@ -45791,7 +45794,8 @@ "supports_prompt_caching": true, "supports_response_schema": true, "supports_vision": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "cache_read_input_token_cost": 1.5e-08 }, "us.amazon.nova-micro-v1:0": { "cache_read_input_token_cost": 8.75e-09, @@ -45805,7 +45809,8 @@ "supports_function_calling": true, "supports_prompt_caching": true, "supports_response_schema": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "cache_read_input_token_cost": 8.75e-09 }, "us.amazon.nova-premier-v1:0": { "deprecation_date": "2026-09-14", @@ -45837,7 +45842,8 @@ "supports_prompt_caching": true, "supports_response_schema": true, "supports_vision": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "cache_read_input_token_cost": 2e-07 }, "us.anthropic.claude-3-5-haiku-20241022-v1:0": { "cache_creation_input_token_cost": 1e-06, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 92e5b1c4ff7..b1c38d0350a 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -377,7 +377,8 @@ "supports_prompt_caching": true, "supports_response_schema": true, "supports_vision": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "cache_read_input_token_cost": 1.5e-08 }, "amazon.nova-2-lite-v1:0": { "cache_read_input_token_cost": 7.5e-08, @@ -560,7 +561,8 @@ "supports_function_calling": true, "supports_prompt_caching": true, "supports_response_schema": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "cache_read_input_token_cost": 8.75e-09 }, "amazon.nova-pro-v1:0": { "cache_read_input_token_cost": 2e-07, @@ -576,7 +578,8 @@ "supports_prompt_caching": true, "supports_response_schema": true, "supports_vision": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "cache_read_input_token_cost": 2e-07 }, "amazon.nova-sonic-v1:0": { "deprecation_date": "2026-09-14", @@ -45791,7 +45794,8 @@ "supports_prompt_caching": true, "supports_response_schema": true, "supports_vision": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "cache_read_input_token_cost": 1.5e-08 }, "us.amazon.nova-micro-v1:0": { "cache_read_input_token_cost": 8.75e-09, @@ -45805,7 +45809,8 @@ "supports_function_calling": true, "supports_prompt_caching": true, "supports_response_schema": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "cache_read_input_token_cost": 8.75e-09 }, "us.amazon.nova-premier-v1:0": { "deprecation_date": "2026-09-14", @@ -45837,7 +45842,8 @@ "supports_prompt_caching": true, "supports_response_schema": true, "supports_vision": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "cache_read_input_token_cost": 2e-07 }, "us.anthropic.claude-3-5-haiku-20241022-v1:0": { "cache_creation_input_token_cost": 1e-06, From afb28540bbd3868dcebd83e8e7c7347c7611abfa Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 16 Sep 2026 16:01:32 -0700 Subject: [PATCH 21/21] fix(e2e): keep the CLI determinism test out of the in-cluster suite It drives the real CLI for several seconds. The edge stamps every upstream call with PYTEST_CURRENT_TEST, a process-global that names whichever test the worker is in when the call arrives rather than the one that made it, so a test that holds a worker that long collects other tests' in-flight calls. Build 234's key report credits this test with 20 Bedrock and 7 Anthropic misses, and it makes no provider call at all. Those misattributed calls take the wrong test id into the cache key and write recordings under it, so the test was polluting the shared corpus it exists to protect. Deselected unless E2E_CLI_DETERMINISM is set, the same opt-in shape the managed-files, prompt-caching and redis-chaos markers already use. The attribution bug itself is older than this branch and is reported, not fixed here. --- .../_driver_unit_tests/test_request_determinism.py | 2 ++ tests/e2e/conftest.py | 7 +++++++ tests/e2e/e2e_config.py | 1 + tests/e2e/pytest.ini | 1 + 4 files changed, 11 insertions(+) diff --git a/tests/e2e/claude_code/_driver_unit_tests/test_request_determinism.py b/tests/e2e/claude_code/_driver_unit_tests/test_request_determinism.py index 5046f35c73b..b7d330b7da6 100644 --- a/tests/e2e/claude_code/_driver_unit_tests/test_request_determinism.py +++ b/tests/e2e/claude_code/_driver_unit_tests/test_request_determinism.py @@ -36,6 +36,8 @@ import pytest from claude_code.cli_driver import _FIXED_CLI_USER_ID, _seed_cli_identity, _stable_cli_state, run_claude from claude_code.rate_limiter import RateLimiter +pytestmark = pytest.mark.cli_determinism + _STUB_REPLY = { "id": "msg_stub", "type": "message", diff --git a/tests/e2e/conftest.py b/tests/e2e/conftest.py index 430e16525d5..ac4cfb71407 100644 --- a/tests/e2e/conftest.py +++ b/tests/e2e/conftest.py @@ -23,6 +23,7 @@ from typing import Final import pytest import requests from e2e_config import ( + CLI_DETERMINISM_OPT_IN_ENV, CONTROL_PLANE_BASE_URL, FIXTURE_DIR, FIXTURE_MODE_RAW, @@ -53,6 +54,7 @@ OPT_IN_MARKERS: Final = MappingProxyType( "managed_files": MANAGED_FILES_OPT_IN_ENV, "prompt_caching_stack": PROMPT_CACHING_OPT_IN_ENV, "redis_chaos": REDIS_CHAOS_OPT_IN_ENV, + "cli_determinism": CLI_DETERMINISM_OPT_IN_ENV, } ) @@ -120,6 +122,11 @@ def pytest_configure(config: pytest.Config) -> None: "prompt_caching_stack: needs a proxy running with router_settings.optional_pre_call_checks including " "prompt_caching; deselected unless E2E_PROMPT_CACHING_STACK is set", ) + config.addinivalue_line( + "markers", + "cli_determinism: drives the real claude CLI for several seconds, which widens the window in which " + "another test's in-flight upstream call is attributed to it; deselected unless E2E_CLI_DETERMINISM is set", + ) config.addinivalue_line( "markers", "redis_chaos: load test that pauses the proxy's Redis outright mid-run; needs a proxy booted from " diff --git a/tests/e2e/e2e_config.py b/tests/e2e/e2e_config.py index 896cb3e7efe..82ddb09f7f5 100644 --- a/tests/e2e/e2e_config.py +++ b/tests/e2e/e2e_config.py @@ -145,6 +145,7 @@ WEEKLY_ANOMALY_OPT_IN_ENV = "E2E_WEEKLY_ANOMALY" MANAGED_FILES_OPT_IN_ENV = "E2E_MANAGED_FILES_STACK" PROMPT_CACHING_OPT_IN_ENV = "E2E_PROMPT_CACHING_STACK" REDIS_CHAOS_OPT_IN_ENV = "E2E_REDIS_CHAOS" +CLI_DETERMINISM_OPT_IN_ENV = "E2E_CLI_DETERMINISM" ANOMALY_SESSIONS = int(os.environ.get("E2E_ANOMALY_SESSIONS", "6")) ANOMALY_TURNS_PER_SESSION = int(os.environ.get("E2E_ANOMALY_TURNS_PER_SESSION", "6")) ANOMALY_TURN_ATTEMPTS = int(os.environ.get("E2E_ANOMALY_TURN_ATTEMPTS", "3")) diff --git a/tests/e2e/pytest.ini b/tests/e2e/pytest.ini index 1fdd3bd28ad..f6d23a3ec12 100644 --- a/tests/e2e/pytest.ini +++ b/tests/e2e/pytest.ini @@ -10,4 +10,5 @@ markers = weekly: real-provider anomaly load test that spends real money; deselected unless E2E_WEEKLY_ANOMALY is set managed_files: needs a proxy running with require_managed_files enabled; deselected unless E2E_MANAGED_FILES_STACK is set prompt_caching_stack: needs a proxy running with router_settings.optional_pre_call_checks including prompt_caching; deselected unless E2E_PROMPT_CACHING_STACK is set + cli_determinism: drives the real claude CLI for several seconds, which widens the window in which another test's in-flight upstream call is attributed to it; deselected unless E2E_CLI_DETERMINISM is set redis_chaos: load test that pauses the proxy's Redis outright mid-run; needs a proxy booted from gateway/redis_chaos_ci_config.yml on the same host, and is deselected unless E2E_REDIS_CHAOS is set