mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-19 00:01:29 +00:00
Merge pull request #41402 from BerriAI/litellm_/buildkite-litellm-e2e-setup-ff714d
feat(e2e): make the provider cache reusable across builds and mount Bedrock behind it
This commit is contained in:
commit
765e6e498d
16 changed files with 1794 additions and 149 deletions
File diff suppressed because it is too large
Load diff
|
|
@ -1,11 +1,41 @@
|
|||
# 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, 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
|
||||
|
||||
## 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
|
||||
|
||||
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
|
||||
|
||||
## 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
|
||||
|
||||
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
|
||||
|
||||
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, 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
|
||||
|
||||
## Configuration
|
||||
|
|
@ -18,16 +48,18 @@ 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. 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. 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
|
||||
|
||||
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
|
||||
|
|
|
|||
|
|
@ -0,0 +1,162 @@
|
|||
"""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 _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",
|
||||
"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}
|
||||
|
||||
|
||||
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"]
|
||||
|
|
@ -132,6 +132,62 @@ 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.
|
||||
|
||||
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:
|
||||
if json.load(handle).get("userID") == _FIXED_CLI_USER_ID:
|
||||
return
|
||||
except (OSError, ValueError):
|
||||
pass
|
||||
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)
|
||||
|
||||
|
||||
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/<cwd slug>/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 +278,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 +303,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 +323,7 @@ def run_claude(
|
|||
completed = run_fn(
|
||||
cmd,
|
||||
env=env,
|
||||
cwd=workspace,
|
||||
input=stdin_input,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
}
|
||||
)
|
||||
|
||||
|
|
@ -85,7 +87,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",
|
||||
|
|
@ -116,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 "
|
||||
|
|
|
|||
|
|
@ -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"))
|
||||
|
|
|
|||
|
|
@ -51,6 +51,9 @@ SECRET_FIELD_SUFFIXES: Final[tuple[str, ...]] = (
|
|||
)
|
||||
SECRET_PLACEHOLDER: Final = "<secret>"
|
||||
|
||||
MARKER_PATTERN: Final = re.compile(r"(?<![0-9a-fA-F])[0-9a-f]{12}(?![0-9a-fA-F])")
|
||||
MARKER_PLACEHOLDER: Final = "<marker>"
|
||||
|
||||
PLACEHOLDER_RULES: Final[tuple[tuple[re.Pattern[str], str], ...]] = (
|
||||
(re.compile(r"(?<![0-9a-fA-F])[0-9a-f]{64}(?![0-9a-fA-F])"), "<sha256>"),
|
||||
(
|
||||
|
|
@ -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"),
|
||||
"<id>",
|
||||
),
|
||||
(re.compile(r"(?<![0-9a-fA-F])[0-9a-f]{12}(?![0-9a-fA-F])"), "<marker>"),
|
||||
(MARKER_PATTERN, MARKER_PLACEHOLDER),
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -4,14 +4,18 @@ import base64
|
|||
import hashlib
|
||||
import hmac
|
||||
import io
|
||||
import json
|
||||
import os
|
||||
import threading
|
||||
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
|
||||
|
||||
from botocore.eventstream import EventStreamBuffer, ParserError
|
||||
from e2e_http import (
|
||||
NetworkError,
|
||||
StreamChunk,
|
||||
|
|
@ -23,12 +27,36 @@ 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"
|
||||
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
|
||||
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"})
|
||||
JSON_VALUE: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue)
|
||||
|
||||
|
||||
|
|
@ -56,6 +84,24 @@ class CacheUnavailable:
|
|||
|
||||
|
||||
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):
|
||||
|
|
@ -83,28 +129,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(BEDROCK_SUFFIXES)
|
||||
return path in OPENAI_JSON_PATHS
|
||||
|
||||
|
||||
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:
|
||||
|
|
@ -118,28 +187,33 @@ def successful_response(url: str, status: int, headers: Mapping[str, str], body:
|
|||
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)
|
||||
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:
|
||||
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
|
||||
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)
|
||||
|
|
@ -147,6 +221,144 @@ 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."""
|
||||
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 path.endswith(BEDROCK_CONVERSE_SUFFIX):
|
||||
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 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
|
||||
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
|
||||
|
|
@ -172,7 +384,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,11 +395,58 @@ 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
|
||||
|
||||
|
||||
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], ...] = ()
|
||||
|
|
@ -199,6 +458,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)
|
||||
|
|
@ -226,14 +503,21 @@ 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)
|
||||
probe: KeyProbe = field(default_factory=KeyProbe)
|
||||
slots: SlotCounter = field(default_factory=SlotCounter)
|
||||
policies: Mapping[str, MountPolicy] = NO_POLICIES
|
||||
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,59 +525,122 @@ 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 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
|
||||
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]:
|
||||
"""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 unkeyed}
|
||||
|
||||
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.reject(mount, UNREACHABLE)
|
||||
return prepared
|
||||
key: Final = exact_key(self.secret, method, url, 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, 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")
|
||||
self.record_key(mount, "hit", test_key, method, url, keyed_headers, body)
|
||||
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")
|
||||
self.record_key(mount, "miss", test_key, method, url, keyed_headers, body)
|
||||
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.reject(mount, UNREACHABLE)
|
||||
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()
|
||||
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(url, head.status_code, headers, b"".join(chunks)):
|
||||
self.counters.increment("rejected")
|
||||
return
|
||||
response: Final = CachedResponse(
|
||||
request_key=key, status_code=head.status_code, headers=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")
|
||||
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
|
||||
}
|
||||
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
|
||||
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}")
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -8,14 +8,80 @@ 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_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/"
|
||||
|
||||
|
||||
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 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
|
||||
return declared
|
||||
|
||||
|
||||
def bedrock_mount(params: LiteLLMParamsBody) -> str | None:
|
||||
"""The edge mount a Bedrock deployment belongs to, or None.
|
||||
|
||||
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 model not in BEDROCK_EDGE_MODELS:
|
||||
return None
|
||||
return f"bedrock/{bedrock_region(params.aws_region_name)}"
|
||||
|
||||
|
||||
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:
|
||||
|
|
|
|||
|
|
@ -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 SIGNATURE_HEADERS, CacheEdge, MountPolicy, 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,22 @@ 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, policies=bedrock_policies())
|
||||
|
||||
|
||||
@functools.lru_cache(maxsize=1)
|
||||
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}": MountPolicy(sign=bedrock_signer(region), unkeyed_headers=SIGNATURE_HEADERS)
|
||||
for region in BEDROCK_REGIONS
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@functools.lru_cache(maxsize=8)
|
||||
|
|
|
|||
72
tests/e2e/provider_edge_bedrock.py
Normal file
72
tests/e2e/provider_edge_bedrock.py
Normal file
|
|
@ -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: # 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"
|
||||
)
|
||||
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)
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue