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.
This commit is contained in:
Yuneng Jiang 2026-09-16 02:34:09 -07:00
parent 2d40254b57
commit b68e60f706
No known key found for this signature in database
5 changed files with 186 additions and 24 deletions

View file

@ -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:

View file

@ -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

View file

@ -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,

View file

@ -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)

View file

@ -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"