From 789b594370b50f6a9eeebb4a7a355f7fa2b45ac7 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Tue, 15 Sep 2026 00:00:53 -0700 Subject: [PATCH] feat(e2e): add bounded trusted provider capture and snapshot storage --- .circleci/config.yml | 9 +- .../test_provider_capture.py | 360 ++++++++++++++++++ .../test_provider_capture_publication.py | 168 ++++++++ .../test_provider_capture_state.py | 181 +++++++++ .../test_provider_capture_store.py | 97 +++++ .../test_provider_edge_control.py | 259 +++++++++++++ tests/e2e/CONTRIBUTING.md | 11 + tests/e2e/capture_policy.py | 207 ++++++++++ tests/e2e/capture_publication.py | 186 +++++++++ tests/e2e/capture_session.py | 123 ++++++ tests/e2e/capture_snapshot.py | 130 +++++++ tests/e2e/capture_store.py | 202 ++++++++++ tests/e2e/conftest.py | 28 +- tests/e2e/e2e_config.py | 6 + tests/e2e/fixture_mode.py | 11 +- tests/e2e/provider_edge.py | 124 +++++- tests/e2e/provider_edge_cli.py | 156 ++++++++ tests/e2e/provider_edge_control.py | 259 +++++++++++++ tests/e2e/provider_edge_remote.py | 85 +++++ .../e2e/router/test_reliability_cache_e2e.py | 2 + 20 files changed, 2578 insertions(+), 26 deletions(-) create mode 100644 tests/code_coverage_tests/test_provider_capture.py create mode 100644 tests/code_coverage_tests/test_provider_capture_publication.py create mode 100644 tests/code_coverage_tests/test_provider_capture_state.py create mode 100644 tests/code_coverage_tests/test_provider_capture_store.py create mode 100644 tests/code_coverage_tests/test_provider_edge_control.py create mode 100644 tests/e2e/capture_policy.py create mode 100644 tests/e2e/capture_publication.py create mode 100644 tests/e2e/capture_session.py create mode 100644 tests/e2e/capture_snapshot.py create mode 100644 tests/e2e/capture_store.py create mode 100644 tests/e2e/provider_edge_cli.py create mode 100644 tests/e2e/provider_edge_control.py create mode 100644 tests/e2e/provider_edge_remote.py diff --git a/.circleci/config.yml b/.circleci/config.yml index bb4ad0f4019..3ffdc676ee9 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -2927,10 +2927,15 @@ jobs: command: | mkdir -p test-results/provider-replay-harness uv run --no-sync pytest -q --noconftest -o addopts= -o pythonpath=tests/e2e -p no:rerunfailures \ - --junitxml=test-results/provider-replay-harness/junit.xml \ + --hypothesis-seed=20260915 --junitxml=test-results/provider-replay-harness/junit.xml \ tests/e2e/test_provider_edge.py tests/e2e/test_fixture_bundle.py \ tests/e2e/test_fixture_canonical.py tests/e2e/test_fixture_mode.py \ - tests/code_coverage_tests/test_provider_replay_harness.py + tests/code_coverage_tests/test_provider_replay_harness.py \ + tests/code_coverage_tests/test_provider_capture.py \ + tests/code_coverage_tests/test_provider_capture_store.py \ + tests/code_coverage_tests/test_provider_capture_publication.py \ + tests/code_coverage_tests/test_provider_capture_state.py \ + tests/code_coverage_tests/test_provider_edge_control.py - store_test_results: path: test-results/provider-replay-harness diff --git a/tests/code_coverage_tests/test_provider_capture.py b/tests/code_coverage_tests/test_provider_capture.py new file mode 100644 index 00000000000..36856b4c715 --- /dev/null +++ b/tests/code_coverage_tests/test_provider_capture.py @@ -0,0 +1,360 @@ +from __future__ import annotations + +import base64 +import hashlib +import threading +from dataclasses import dataclass, field +from datetime import datetime, timedelta, timezone +from pathlib import Path +from typing import Final + +import pytest +from capture_policy import ScenarioIdentity, ScenarioOutcome, canonical_scenario_id, publication_error +from capture_session import AttemptDenied, AttemptReserved, AttemptUncertain, CaptureSession, Reservation +from capture_snapshot import ( + CaptureProvenance, + ScenarioSnapshot, + SnapshotFailure, + build_snapshot, + materialize_snapshot, + refresh_due, + verify_snapshot, +) +from capture_store import StoreFailure +from fixture_bundle import ( + BundleRecorder, + LoadedBundle, + RecordedHttpResponse, + RecordedRequest, + RecordedStreamedResponse, + load_bundle, + prepare_bundle, +) +from fixture_mode import current_test_key +from provider_edge import REPLAY_MISS_STATUS, RecordEdge, ReplayEdge, ReplaySource, _persist +from test_provider_edge import CHAT_PATH, call_edge, fake_provider, provider_url, running_edge + + +@dataclass +class ScriptedReservations: + cap: int = 12 + uncertain: bool = False + calls: tuple[str, ...] = field(default=(), init=False) + + def reserve(self, *, scenario_key: str, owner: str, attempt_id: str) -> Reservation: + self.calls = (*self.calls, attempt_id) + if self.uncertain: + return AttemptUncertain("reservation outcome uncertain") + if len(self.calls) > self.cap: + return AttemptDenied("durable budget exhausted") + return AttemptReserved(attempt_id) + + def complete(self, *, scenario_key: str, owner: str, attempt_id: str, successful: bool) -> str | None: + return None + + def acquire(self, *, scenario_key: str, owner: str, expires_at: int) -> StoreFailure | None: + return None + + def release(self, *, scenario_key: str, owner: str) -> StoreFailure | None: + return None + + +def successful_response() -> RecordedHttpResponse: + return RecordedHttpResponse(status_code=200, headers={}, body_b64=base64.b64encode(b'{"answer":"blue"}').decode()) + + +class TestBoundedCapture: + def test_uncertain_reservation_stops_all_following_attempts(self) -> None: + store: Final = ScriptedReservations(uncertain=True) + session: Final = CaptureSession(ScenarioIdentity("test_example.py::test_one", "a" * 64, "test"), "owner", store) + assert session.before_attempt() == "reservation outcome uncertain" + assert session.before_attempt() == "reservation outcome uncertain" + assert len(store.calls) == 1 + assert not session.finish(ScenarioOutcome(True, True, True)).publishable + + def test_cap_plus_one_is_denied_before_http(self, tmp_path: Path) -> None: + store: Final = ScriptedReservations(cap=1) + scenario: Final = ScenarioIdentity("test_example.py::test_one", "a" * 64, "test") + session: Final = CaptureSession(scenario, "owner", store, max_attempts=1) + recorder: Final = prepare_bundle(tmp_path / "bounded", profile="stateless_v1") + assert isinstance(recorder, BundleRecorder) + backend: Final = RecordEdge( + recorder, + threading.Lock(), + test_key=lambda: scenario.node, + before_attempt=session.before_attempt, + response_finished=session.response_finished, + response_byte_limit=1024, + ) + with fake_provider() as provider: + with running_edge(backend, {"openai": provider_url(provider)}) as edge: + first: Final = call_edge( + edge, "POST", CHAT_PATH, body=b'{"prompt":"blue"}', headers={"content-type": "application/json"} + ) + assert first.status_code == 200 + second: Final = call_edge( + edge, "POST", CHAT_PATH, body=b'{"prompt":"blue"}', headers={"content-type": "application/json"} + ) + assert second.status_code == REPLAY_MISS_STATUS + assert b"attempt cap" in second.body + assert len(provider.hits) == 1 + assert len(store.calls) == 1 + assert not session.finish(ScenarioOutcome(True, True, True)).publishable + + def test_failure_is_not_recovered_by_a_second_attempt(self) -> None: + store: Final = ScriptedReservations() + session: Final = CaptureSession(ScenarioIdentity("test_example.py::test_one", "a" * 64, "test"), "owner", store) + assert session.before_attempt() is None + session.response_finished(RecordedHttpResponse(status_code=503, headers={}, body_b64="")) + assert session.before_attempt() == "provider response was not successful" + assert len(store.calls) == 1 + assert not session.finish(ScenarioOutcome(True, True, True)).publishable + + def test_finalizer_failure_and_uncertain_inflight_outcome_cannot_publish(self) -> None: + session: Final = CaptureSession( + ScenarioIdentity("test_example.py::test_one", "a" * 64, "test"), "owner", ScriptedReservations() + ) + assert session.before_attempt() is None + assert session.finish(ScenarioOutcome(True, True, True)).error == "outbound outcome is uncertain" + assert session.before_attempt() == "scenario is closed" + session.response_finished(successful_response()) + assert not session.finish(ScenarioOutcome(True, True, True)).publishable + + def test_success_requires_complete_trusted_outcome(self) -> None: + for teardown, expected in ((False, False), (True, True)): + session: Final = CaptureSession( + ScenarioIdentity("test_example.py::test_one", "a" * 64, "test"), "owner", ScriptedReservations() + ) + assert session.before_attempt() is None + session.response_finished(successful_response()) + result: Final = session.finish(ScenarioOutcome(True, True, teardown)) + assert result.publishable is expected + assert result.response_count == 1 + assert len(result.attempts) == 1 + + +class TestScenarioIdentity: + def test_test_roots_normalize_without_changing_parameters(self) -> None: + expected: Final = "tests/e2e/router/test_cache.py::TestCache::test_hit[prompt/2031-04-05]" + assert canonical_scenario_id(expected) == expected + assert canonical_scenario_id("/opt/runner/" + expected) == expected + assert canonical_scenario_id(expected.removeprefix("tests/e2e/")) == expected + assert canonical_scenario_id(expected.replace("2031-04-05", "2031-04-06")) != expected + + @pytest.mark.parametrize("node", ["", "test_cache.py", "../test_cache.py::test_hit", "x/../test.py::test_hit"]) + def test_invalid_scenario_ids_are_rejected(self, node: str) -> None: + with pytest.raises(ValueError): + canonical_scenario_id(node) + + def test_contract_and_profile_change_identity_but_candidate_revision_does_not_key_it(self) -> None: + node: Final = "router/test_cache.py::test_hit" + identity: Final = ScenarioIdentity(node, "a" * 64, "openai-test-v1") + assert identity.key == ScenarioIdentity("/opt/runner/tests/e2e/" + node, "a" * 64, "openai-test-v1").key + assert identity.key != ScenarioIdentity(node, "b" * 64, "openai-test-v1").key + assert identity.key != ScenarioIdentity(node, "a" * 64, "openai-test-v2").key + + def test_explicit_record_and_replay_identity_does_not_use_pytest_process(self, tmp_path: Path) -> None: + scenario: Final = "tests/e2e/router/test_cache.py::test_hit" + recorder: Final = prepare_bundle(tmp_path / "capture", profile="stateless_v1") + assert isinstance(recorder, BundleRecorder) + with fake_provider() as provider: + mounts: Final = {"openai": provider_url(provider)} + with running_edge(RecordEdge(recorder, threading.Lock(), test_key=lambda: scenario), mounts) as edge: + captured: Final = call_edge( + edge, + "POST", + CHAT_PATH, + body=b'{"prompt":"synthetic"}', + headers={"content-type": "application/json"}, + ) + assert captured.status_code == 200 + loaded: Final = load_bundle(recorder.root, profile="stateless_v1") + assert isinstance(loaded, LoadedBundle) + wrong: Final = ReplaySource(loaded, test_key=current_test_key) + with running_edge(ReplayEdge(wrong), mounts) as edge: + assert ( + call_edge( + edge, + "POST", + CHAT_PATH, + body=b'{"prompt":"synthetic"}', + headers={"content-type": "application/json"}, + ).status_code + == REPLAY_MISS_STATUS + ) + source: Final = ReplaySource(loaded, test_key=lambda: scenario) + with running_edge(ReplayEdge(source), mounts) as edge: + replayed: Final = call_edge( + edge, + "POST", + CHAT_PATH, + body=b'{"prompt":"synthetic"}', + headers={"content-type": "application/json"}, + ) + assert replayed.status_code == 200 + assert replayed.body == captured.body + assert source.leftover_error(scenario) is None + assert len(provider.hits) == 1 + + +class TestPublicationPolicy: + @pytest.mark.parametrize("phase", ["setup", "call", "teardown"]) + def test_every_trusted_phase_must_pass(self, phase: str) -> None: + outcome: Final = ScenarioOutcome(setup=phase != "setup", call=phase != "call", teardown=phase != "teardown") + assert publication_error(outcome, ()) is not None + + def test_empty_or_failed_capture_cannot_be_published(self) -> None: + outcome: Final = ScenarioOutcome(setup=True, call=True, teardown=True) + assert publication_error(outcome, ()) == "capture contains no interactions" + failed: Final = RecordedHttpResponse(status_code=503, headers={}, body_b64=base64.b64encode(b"failed").decode()) + assert publication_error(outcome, (failed,)) == "provider response was not successful" + + @pytest.mark.parametrize( + "bad", + [b'{"error":"failed"}', b'{"usage":{"prompt_tokens":2,"completion_tokens":3,"total_tokens":9}}', b"not-json"], + ) + def test_wrong_success_bodies_do_not_publish(self, bad: bytes) -> None: + response: Final = RecordedHttpResponse(status_code=200, headers={}, body_b64=base64.b64encode(bad).decode()) + assert publication_error(ScenarioOutcome(True, True, True), (response,)) is not None + + @pytest.mark.parametrize("tail", [b"", b'data: {"error":"failed"}\n\n', b"data: [DONE]\n"]) + def test_incomplete_or_error_stream_does_not_publish(self, tail: bytes) -> None: + data: Final = b'data: {"choices":[{"delta":{"content":"blue"}}]}\n\n' + tail + response: Final = RecordedStreamedResponse( + status_code=200, headers={}, chunks_b64=[base64.b64encode(data).decode()] + ) + assert publication_error(ScenarioOutcome(True, True, True), (response,)) is not None + + @pytest.mark.parametrize("provider", ["openai", "anthropic"]) + def test_complete_stream_across_arbitrary_chunk_boundaries_publishes(self, provider: str) -> None: + payload: Final = ( + b': keepalive\ndata: {"choices":[{"index":0,"delta":{"content":"blue"},"finish_reason":null}],"usage":null}\n\ndata: {"choices":[{"index":0,"delta":{},"finish_reason":"stop"}],"usage":null}\n\ndata: {"choices":[],"usage":{"prompt_tokens":2,"completion_tokens":1,"total_tokens":3}}\n\ndata: [DONE]\n\n' + if provider == "openai" + else b'event: message_start\ndata: {"type":"message_start","message":{"usage":{"input_tokens":2,"output_tokens":0}}}\n\nevent: content_block_delta\ndata: {"type":"content_block_delta","delta":{"text":"blue"}}\n\nevent: message_delta\ndata: {"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"output_tokens":1}}\n\nevent: message_stop\ndata: {"type":"message_stop"}\n\n' + ) + for split in range(1, len(payload)): + response: Final = RecordedStreamedResponse( + status_code=200, + headers={}, + chunks_b64=[base64.b64encode(part).decode() for part in (payload[:split], payload[split:])], + ) + assert publication_error(ScenarioOutcome(True, True, True), (response,)) is None + + @pytest.mark.parametrize( + "payload", + [ + b': heartbeat\ndata: {"error":{"message":"failed"}}\n\ndata: [DONE]\n\n', + b"data: [DONE]\n\n", + b': heartbeat\ndata: {"error":{"message":"failed"}}\n\ndata: {"choices":[{"delta":{"content":"blue"},"finish_reason":"stop"}]}\n\ndata: [DONE]\n\n', + b'data: {"choices":[{"index":0,"delta":{"content":"done"},"finish_reason":"stop"}]}\n\ndata: {"choices":[{"index":0,"delta":{"content":"late"},"finish_reason":null}]}\n\ndata: [DONE]\n\n', + b'data: {"choices":[{"delta":{"content":"partial"},"finish_reason":null}]}\n\ndata: [DONE]\n\n', + b'data: [DONE]\n\ndata: {"type":"message_stop"}\n\n', + b'data: {"choices":null}\n\ndata: [DONE]\n\n', + b'data: {"choices":[{"finish_reason":"length"}]}\n\ndata: [DONE]\n\n', + ], + ) + def test_malformed_completion_settles_attempt_as_failure(self, payload: bytes) -> None: + response: Final = RecordedStreamedResponse( + status_code=200, headers={}, chunks_b64=[base64.b64encode(payload).decode()] + ) + session: Final = CaptureSession( + ScenarioIdentity("test_example.py::test_one", "a" * 64, "synthetic"), "owner", ScriptedReservations() + ) + assert session.before_attempt() is None + session.response_finished(response) + result: Final = session.finish(ScenarioOutcome(True, True, True)) + assert not result.publishable + assert result.error != "outbound outcome is uncertain" + + +class TestSnapshotBoundary: + def test_capture_download_fresh_directory_replay_and_faults(self, tmp_path: Path) -> None: + identity: Final = ScenarioIdentity("test_example.py::test_one", "a" * 64, "synthetic-v1") + session: Final = CaptureSession(identity, "owner", ScriptedReservations()) + recorder: Final = prepare_bundle(tmp_path / "capture", profile="stateless_v1") + assert isinstance(recorder, BundleRecorder) + with fake_provider() as provider: + mounts: Final = {"openai": provider_url(provider)} + backend: Final = RecordEdge( + recorder, + threading.Lock(), + test_key=lambda: canonical_scenario_id(identity.node), + before_attempt=session.before_attempt, + response_finished=session.response_finished, + ) + with running_edge(backend, mounts) as edge: + response: Final = call_edge( + edge, "POST", CHAT_PATH, body=b'{"prompt":"blue"}', headers={"content-type": "application/json"} + ) + assert response.status_code == 200 + bundle: Final = load_bundle(recorder.root, profile="stateless_v1") + assert isinstance(bundle, LoadedBundle) + result: Final = session.finish(ScenarioOutcome(True, True, True)) + snapshot: Final = build_snapshot( + result, + bundle, + CaptureProvenance( + test_revision="b" * 40, candidate_revision="c" * 40, runner_digest="sha256:" + "d" * 64 + ), + ) + assert isinstance(snapshot, bytes) + digest: Final = hashlib.sha256(snapshot).hexdigest() + now: Final = datetime.now(timezone.utc) + verified: Final = verify_snapshot(snapshot, expected_sha256=digest, identity=identity, now=now) + assert isinstance(verified, ScenarioSnapshot) + assert not refresh_due(verified, now=now) + assert refresh_due(verified, now=now + timedelta(days=1)) + assert isinstance( + verify_snapshot(snapshot + b" ", expected_sha256=digest, identity=identity, now=now), SnapshotFailure + ) + assert isinstance( + verify_snapshot(snapshot, expected_sha256=digest, identity=identity, now=now + timedelta(days=7)), + SnapshotFailure, + ) + assert isinstance( + verify_snapshot( + snapshot, + expected_sha256=digest, + identity=ScenarioIdentity(identity.node, "f" * 64, "synthetic-v1"), + now=now, + ), + SnapshotFailure, + ) + materialize_snapshot(verified, tmp_path / "fresh") + with pytest.raises(FileExistsError): + materialize_snapshot(verified, tmp_path / "fresh") + loaded: Final = load_bundle(tmp_path / "fresh", profile="stateless_v1") + assert isinstance(loaded, LoadedBundle) + source: Final = ReplaySource(loaded, test_key=lambda: canonical_scenario_id(identity.node)) + with running_edge(ReplayEdge(source), mounts) as edge: + replayed: Final = call_edge( + edge, "POST", CHAT_PATH, body=b'{"prompt":"blue"}', headers={"content-type": "application/json"} + ) + assert replayed.status_code == 200 + assert replayed.body == response.body + assert source.leftover_error(canonical_scenario_id(identity.node)) is None + assert len(provider.hits) == 1 + + +@pytest.mark.parametrize("disk_failure", [True, False]) +def test_persist_must_reach_disk_before_success(tmp_path: Path, disk_failure: bool) -> None: + root: Final = tmp_path / "bundle" + if disk_failure: + root.write_text("occupied by a file") + session: Final = CaptureSession( + ScenarioIdentity("test_example.py::test_one", "a" * 64, "synthetic"), "owner", ScriptedReservations() + ) + backend: Final = RecordEdge( + BundleRecorder(root, profile="stateless_v1"), threading.Lock(), response_finished=session.response_finished + ) + request: Final = RecordedRequest(method="post", path="/openai/v1/chat/completions", headers={}) + assert session.before_attempt() is None + if disk_failure: + with pytest.raises(OSError): + _persist(backend, session.identity.node, request, successful_response()) + else: + _persist(backend, session.identity.node, request, successful_response()) + assert len(tuple(root.rglob("*.json"))) == 1 + result: Final = session.finish(ScenarioOutcome(True, True, True)) + assert result.publishable is not disk_failure + assert result.response_count == (0 if disk_failure else 1) diff --git a/tests/code_coverage_tests/test_provider_capture_publication.py b/tests/code_coverage_tests/test_provider_capture_publication.py new file mode 100644 index 00000000000..8c36219ff47 --- /dev/null +++ b/tests/code_coverage_tests/test_provider_capture_publication.py @@ -0,0 +1,168 @@ +from __future__ import annotations + +import base64 +import hashlib +from datetime import datetime, timezone +from io import BytesIO +from typing import Final, cast + +import pytest +from botocore.config import Config +from botocore.response import StreamingBody +from botocore.session import get_session +from botocore.stub import ANY, Stubber +from capture_policy import ScenarioIdentity, ScenarioOutcome +from capture_publication import ObjectClient, PointerClient, SnapshotPointer, SnapshotRepository +from capture_snapshot import CaptureProvenance, ScenarioSnapshot, SnapshotFailure +from fixture_bundle import Interaction, Manifest, RecordedHttpResponse, RecordedRequest +from fixture_profile import StrictIdentity, strict_identity + +NOW: Final = datetime(2031, 4, 5, tzinfo=timezone.utc) +IDENTITY: Final = ScenarioIdentity("test_example.py::test_one", "a" * 64, "synthetic") + + +def snapshot_content() -> bytes: + request: Final = strict_identity( + method="POST", + path="/openai/v1/chat/completions", + query="", + headers={"content-type": "application/json"}, + body=b'{"prompt":"blue"}', + mount="openai", + upstream_base="https://api.openai.com", + ) + assert isinstance(request, StrictIdentity) + return ( + ScenarioSnapshot( + identity=IDENTITY, + manifest=Manifest( + format_version=5, recorded_at=NOW, harness_version="synthetic", match_profile="stateless_v1" + ), + interactions=( + Interaction( + request=RecordedRequest( + method="post", path="/openai/v1/chat/completions", headers={}, strict_identity=request + ), + response=RecordedHttpResponse( + status_code=200, headers={}, body_b64=base64.b64encode(b'{"answer":"blue"}').decode() + ), + ), + ), + provenance=CaptureProvenance( + test_revision="b" * 40, candidate_revision="c" * 40, runner_digest="sha256:" + "d" * 64 + ), + owner="owner", + attempts=("attempt",), + outcome=ScenarioOutcome(True, True, True), + ) + .model_dump_json() + .encode() + ) + + +def clients() -> tuple[ObjectClient, PointerClient]: + session: Final = get_session() + config: Final = Config(retries={"total_max_attempts": 1}) + return ( + cast( + ObjectClient, + session.create_client( + "s3", + region_name="us-east-1", + aws_access_key_id="synthetic", + aws_secret_access_key="synthetic", + config=config, + ), + ), + cast( + PointerClient, + session.create_client( + "dynamodb", + region_name="us-east-1", + aws_access_key_id="synthetic", + aws_secret_access_key="synthetic", + config=config, + ), + ), + ) + + +class TestConditionalPublication: + @pytest.mark.parametrize("conflict", [False, True]) + def test_object_readback_precedes_conditional_pointer_promotion(self, conflict: bool) -> None: + objects, pointers = clients() + content: Final = snapshot_content() + sha: Final = hashlib.sha256(content).hexdigest() + key: Final = f"approved/{IDENTITY.key}/{sha}.json" + with Stubber(objects) as s3, Stubber(pointers) as dynamo: + s3.add_response( + "put_object", + {"VersionId": "version-one"}, + { + "Bucket": "synthetic-bucket", + "Key": key, + "Body": content, + "IfNoneMatch": "*", + "ContentType": "application/json", + "ChecksumSHA256": base64.b64encode(hashlib.sha256(content).digest()).decode(), + "ServerSideEncryption": "aws:kms", + "SSEKMSKeyId": "synthetic-key", + }, + ) + s3.add_response( + "get_object", + { + "VersionId": "version-one", + "ContentLength": len(content), + "Body": StreamingBody(BytesIO(content), len(content)), + }, + {"Bucket": "synthetic-bucket", "Key": key, "VersionId": "version-one"}, + ) + expected: Final = { + "TableName": "synthetic-table", + "Item": ANY, + "ConditionExpression": "pointer_revision = :expected", + "ExpressionAttributeValues": {":expected": {"S": "previous"}}, + } + if conflict: + dynamo.add_client_error( + "put_item", service_error_code="ConditionalCheckFailedException", expected_params=expected + ) + else: + dynamo.add_response("put_item", {}, expected) + result: Final = SnapshotRepository( + objects, pointers, "synthetic-bucket", "synthetic-table", "synthetic-key" + ).publish(content, IDENTITY, expected_revision="previous", now=NOW) + if conflict: + assert isinstance(result, SnapshotFailure) and "unreferenced" in result.reason + else: + assert ( + isinstance(result, SnapshotPointer) and result.version_id == "version-one" and result.sha256 == sha + ) + s3.assert_no_pending_responses() + dynamo.assert_no_pending_responses() + + @pytest.mark.parametrize("fault", ["upload", "version", "digest"]) + def test_interrupted_or_corrupt_blob_never_reaches_pointer_write(self, fault: str) -> None: + objects, pointers = clients() + content: Final = snapshot_content() + with Stubber(objects) as s3, Stubber(pointers) as dynamo: + if fault == "upload": + s3.add_client_error("put_object", service_error_code="InternalError") + else: + s3.add_response("put_object", {"VersionId": "version-one"}) + downloaded: Final = content.replace(b"blue", b"gold") if fault == "digest" else content + s3.add_response( + "get_object", + { + "VersionId": "wrong" if fault == "version" else "version-one", + "ContentLength": len(downloaded), + "Body": StreamingBody(BytesIO(downloaded), len(downloaded)), + }, + ) + result: Final = SnapshotRepository( + objects, pointers, "synthetic-bucket", "synthetic-table", "synthetic-key" + ).publish(content, IDENTITY, expected_revision=None, now=NOW) + assert isinstance(result, SnapshotFailure) + s3.assert_no_pending_responses() + dynamo.assert_no_pending_responses() diff --git a/tests/code_coverage_tests/test_provider_capture_state.py b/tests/code_coverage_tests/test_provider_capture_state.py new file mode 100644 index 00000000000..15df74f95c4 --- /dev/null +++ b/tests/code_coverage_tests/test_provider_capture_state.py @@ -0,0 +1,181 @@ +from __future__ import annotations + +import hashlib +from datetime import timedelta +from io import BytesIO +from typing import Final + +from botocore.response import StreamingBody +from botocore.stub import ANY, Stubber +from capture_policy import ScenarioIdentity, ScenarioOutcome +from capture_publication import SnapshotPointer, SnapshotRepository +from capture_session import CaptureSession +from capture_snapshot import ScenarioSnapshot, SnapshotFailure, refresh_due, verify_snapshot +from fixture_bundle import RecordedHttpResponse +from hypothesis import settings +from hypothesis import strategies as st +from hypothesis.stateful import RuleBasedStateMachine, invariant, precondition, rule +from test_provider_capture import ScriptedReservations, successful_response +from test_provider_capture_publication import IDENTITY, NOW, clients, snapshot_content + + +class CaptureLifecycleMachine(RuleBasedStateMachine): + def __init__(self) -> None: + super().__init__() + self.store = ScriptedReservations() + self.session = CaptureSession( + ScenarioIdentity("test_synthetic.py::test_capture", "a" * 64, "synthetic"), + "owner", + self.store, + max_attempts=3, + ) + self.reservations = 0 + self.responses = 0 + self.in_flight = False + self.failed = False + self.closed = False + + @rule(uncertain=st.booleans()) + def reserve(self, uncertain: bool) -> None: + self.store.uncertain = uncertain + expected_denial: Final = self.closed or self.failed or self.in_flight or self.reservations == 3 or uncertain + response: Final = self.session.before_attempt() + assert (response is not None) == expected_denial + if not expected_denial: + self.reservations += 1 + self.in_flight = True + elif not self.closed: + self.failed = True + + @precondition(lambda self: self.in_flight and not self.closed) + @rule(success=st.booleans()) + def receive(self, success: bool) -> None: + self.session.response_finished( + successful_response() if success else RecordedHttpResponse(status_code=503, headers={}, body_b64="") + ) + self.responses += 1 + self.in_flight = False + if not success: + self.failed = True + + @precondition(lambda self: not self.closed) + @rule(setup=st.booleans(), call=st.booleans(), teardown=st.booleans()) + def finish(self, setup: bool, call: bool, teardown: bool) -> None: + expected: Final = setup and call and teardown and not self.failed and not self.in_flight and self.responses > 0 + result: Final = self.session.finish(ScenarioOutcome(setup, call, teardown)) + assert result.publishable == expected + assert len(result.attempts) == self.reservations + assert result.response_count == self.responses + self.closed = True + + @invariant() + def budget_is_never_exceeded(self) -> None: + assert self.reservations <= 3 + assert len(self.store.calls) <= 4 + + +TestCaptureLifecycle = CaptureLifecycleMachine.TestCase +TestCaptureLifecycle.settings = settings(max_examples=30, stateful_step_count=25, derandomize=True, deadline=None) + + +class PublicationLifecycleMachine(RuleBasedStateMachine): + def __init__(self) -> None: + super().__init__() + objects, pointers = clients() + self.s3 = Stubber(objects) + self.ddb = Stubber(pointers) + self.s3.activate() + self.ddb.activate() + self.repository = SnapshotRepository(objects, pointers, "synthetic-bucket", "synthetic-table", "synthetic-key") + self.approved: SnapshotPointer | None = None + self.content: bytes | None = None + self.sequence = 0 + self.age = 0 + + @rule(interrupted=st.booleans(), stale_writer=st.booleans()) + def publish(self, interrupted: bool, stale_writer: bool) -> None: + self.sequence += 1 + candidate: Final = ( + ScenarioSnapshot.model_validate_json(snapshot_content()) + .model_copy(update={"owner": f"owner-{self.sequence}"}) + .model_dump_json() + .encode() + ) + digest: Final = hashlib.sha256(candidate).hexdigest() + key: Final = f"approved/{IDENTITY.key}/{digest}.json" + version: Final = f"version-{self.sequence}" + self.s3.add_response("put_object", {"VersionId": version}) + self.s3.add_response( + "get_object", + { + "VersionId": version, + "ContentLength": len(candidate), + "Body": StreamingBody(BytesIO(candidate), len(candidate)), + }, + {"Bucket": "synthetic-bucket", "Key": key, "VersionId": version}, + ) + expected: Final = "stale" if stale_writer else (self.approved.revision if self.approved else None) + params: Final = { + "TableName": "synthetic-table", + "Item": ANY, + "ConditionExpression": "attribute_not_exists(pk)" if expected is None else "pointer_revision = :expected", + } + if expected is not None: + params["ExpressionAttributeValues"] = {":expected": {"S": expected}} + if interrupted or stale_writer: + self.ddb.add_client_error( + "put_item", + service_error_code="InternalServerError" if interrupted else "ConditionalCheckFailedException", + expected_params=params, + ) + else: + self.ddb.add_response("put_item", {}, params) + result: Final = self.repository.publish(candidate, IDENTITY, expected_revision=expected, now=NOW) + assert isinstance(result, SnapshotFailure) == (interrupted or stale_writer) + if isinstance(result, SnapshotPointer): + self.approved = result + self.content = candidate + self.age = 0 + self.s3.assert_no_pending_responses() + self.ddb.assert_no_pending_responses() + + @rule(hours=st.sampled_from([0, 23, 24, 167, 168, 169])) + def expire(self, hours: int) -> None: + self.age = hours + + @precondition(lambda self: self.approved is not None) + @rule(outage=st.booleans(), corrupt=st.booleans()) + def download(self, outage: bool, corrupt: bool) -> None: + assert self.approved is not None and self.content is not None + pointer: Final = self.approved + if outage: + self.s3.add_client_error("get_object", service_error_code="ServiceUnavailable") + else: + data: Final = self.content.replace(b"blue", b"gold") if corrupt else self.content + self.s3.add_response( + "get_object", + { + "VersionId": pointer.version_id, + "ContentLength": len(data), + "Body": StreamingBody(BytesIO(data), len(data)), + }, + ) + result: Final = self.repository.download(pointer, IDENTITY, now=NOW + timedelta(hours=self.age)) + assert isinstance(result, bytes) == (not outage and not corrupt and self.age < 168) + local: Final = verify_snapshot( + self.content, expected_sha256=pointer.sha256, identity=IDENTITY, now=NOW + timedelta(hours=self.age) + ) + assert isinstance(local, ScenarioSnapshot) == (self.age < 168) + if isinstance(local, ScenarioSnapshot): + assert refresh_due(local, now=NOW + timedelta(hours=self.age)) == (self.age >= 24) + self.s3.assert_no_pending_responses() + + def teardown(self) -> None: + self.s3.assert_no_pending_responses() + self.ddb.assert_no_pending_responses() + self.s3.deactivate() + self.ddb.deactivate() + + +TestPublicationLifecycle = PublicationLifecycleMachine.TestCase +TestPublicationLifecycle.settings = settings(max_examples=20, stateful_step_count=15, derandomize=True, deadline=None) diff --git a/tests/code_coverage_tests/test_provider_capture_store.py b/tests/code_coverage_tests/test_provider_capture_store.py new file mode 100644 index 00000000000..079212486be --- /dev/null +++ b/tests/code_coverage_tests/test_provider_capture_store.py @@ -0,0 +1,97 @@ +from __future__ import annotations + +from typing import Final, cast + +import boto3 +import pytest +from botocore.config import Config +from botocore.stub import Stubber +from capture_session import AttemptDenied, AttemptReserved, AttemptUncertain +from capture_store import DynamoCaptureStore, DynamoClient + + +def dynamo_client() -> DynamoClient: + return cast( + DynamoClient, + boto3.client( + "dynamodb", + region_name="us-east-1", + aws_access_key_id="synthetic", + aws_secret_access_key="synthetic", + config=Config(retries={"max_attempts": 0}), + ), + ) + + +class TestDurableReservations: + def test_run_cannot_reset_an_existing_budget(self) -> None: + client: Final = dynamo_client() + with Stubber(client) as stub: + stub.add_response( + "put_item", + {}, + { + "TableName": "synthetic-table", + "Item": { + "pk": {"S": "run#owner"}, + "attempt_count": {"N": "0"}, + "attempt_cap": {"N": "2"}, + "lease_expires": {"N": "200"}, + }, + "ConditionExpression": "attribute_not_exists(pk)", + }, + ) + store: Final = DynamoCaptureStore(client, "synthetic-table", clock=lambda: 100) + assert store.create_run(owner="owner", cap=2, expires_at=200) is None + stub.add_client_error("put_item", service_error_code="ConditionalCheckFailedException") + failure: Final = store.create_run(owner="owner", cap=2, expires_at=200) + assert failure is not None and not failure.uncertain + stub.assert_no_pending_responses() + + def test_expired_but_unsettled_lease_cannot_be_stolen(self) -> None: + client: Final = dynamo_client() + with Stubber(client) as stub: + expected: Final = { + "TableName": "synthetic-table", + "Item": { + "pk": {"S": "lease#scenario"}, + "lease_owner": {"S": "owner"}, + "lease_expires": {"N": "200"}, + "lease_settled": {"BOOL": False}, + }, + "ConditionExpression": "attribute_not_exists(pk) OR (lease_settled = :true AND lease_expires <= :now)", + "ExpressionAttributeValues": {":true": {"BOOL": True}, ":now": {"N": "100"}}, + } + stub.add_client_error( + "put_item", service_error_code="ConditionalCheckFailedException", expected_params=expected + ) + failure: Final = DynamoCaptureStore(client, "synthetic-table", clock=lambda: 100).acquire( + scenario_key="scenario", owner="owner", expires_at=200 + ) + assert failure is not None and not failure.uncertain + stub.assert_no_pending_responses() + + @pytest.mark.parametrize( + "code,kind", [("TransactionCanceledException", AttemptDenied), ("InternalServerError", AttemptUncertain)] + ) + def test_reservation_denial_and_uncertainty_remain_distinct( + self, code: str, kind: type[AttemptDenied | AttemptUncertain] + ) -> None: + client: Final = dynamo_client() + with Stubber(client) as stub: + stub.add_client_error("transact_write_items", service_error_code=code) + result: Final = DynamoCaptureStore(client, "synthetic-table", clock=lambda: 100).reserve( + scenario_key="scenario", owner="owner", attempt_id="attempt-one" + ) + assert isinstance(result, kind) + stub.assert_no_pending_responses() + + def test_successful_atomic_reservation_keeps_attempt_identity(self) -> None: + client: Final = dynamo_client() + with Stubber(client) as stub: + stub.add_response("transact_write_items", {}) + result: Final = DynamoCaptureStore(client, "synthetic-table", clock=lambda: 100).reserve( + scenario_key="scenario", owner="owner", attempt_id="attempt-one" + ) + assert result == AttemptReserved("attempt-one") + stub.assert_no_pending_responses() diff --git a/tests/code_coverage_tests/test_provider_edge_control.py b/tests/code_coverage_tests/test_provider_edge_control.py new file mode 100644 index 00000000000..372b965efb4 --- /dev/null +++ b/tests/code_coverage_tests/test_provider_edge_control.py @@ -0,0 +1,259 @@ +from __future__ import annotations + +import base64 +import hashlib +import os +import subprocess +import sys +import threading +from collections.abc import Generator +from contextlib import contextmanager +from pathlib import Path +from typing import Final + +import pytest +from capture_policy import SCENARIO_BYTES, ScenarioIdentity, canonical_scenario_id +from capture_store import StoreFailure +from fixture_bundle import ( + BundleRecorder, + LoadedBundle, + RecordedHttpResponse, + RecordedRequest, + load_bundle, + prepare_bundle, +) +from provider_edge import REPLAY_MISS_STATUS, _persist, start_provider_edge +from provider_edge_control import ControlRequest, ControlServer, EdgeController +from provider_edge_remote import RemoteEdge +from test_provider_capture import ScriptedReservations, successful_response +from test_provider_edge import CHAT_PATH, call_edge, fake_provider, provider_url + + +@contextmanager +def remote_controller(controller: EdgeController, mounts: dict[str, str]) -> Generator[RemoteEdge, None, None]: + running: Final = start_provider_edge( + controller.backend({}), mounts=mounts, guard=controller, observation=controller + ) + control: Final = ControlServer(controller) + thread: Final = threading.Thread(target=control.serve_forever, daemon=True) + thread.start() + try: + yield RemoteEdge(f"http://127.0.0.1:{control.server_port}", f"http://127.0.0.1:{running.edge.port}") + finally: + control.shutdown() + control.server_close() + running.shutdown() + thread.join(timeout=5) + + +class TestTrustedRemoteEdge: + def test_capture_and_replay_keep_lifecycle_and_observed_counts(self, tmp_path: Path) -> None: + identity: Final = ScenarioIdentity("test_example.py::test_one", "a" * 64, "synthetic") + node: Final = canonical_scenario_id(identity.node) + recorder: Final = prepare_bundle(tmp_path / "capture", profile="stateless_v1") + assert isinstance(recorder, BundleRecorder) + controller: Final = EdgeController( + {node: identity}, "owner", 200, store=ScriptedReservations(), recorder=recorder + ) + with fake_provider() as provider: + mounts: Final = {"openai": provider_url(provider)} + with remote_controller(controller, mounts) as remote: + denied: Final = call_edge(remote.edge, "POST", CHAT_PATH, body=b"{}") + assert denied.status_code == REPLAY_MISS_STATUS + remote.begin(node) + remote.phase(node, "setup", True) + observer: Final = remote.command(ControlRequest(action="observe", node=node, marker="blue")) + assert observer.observation_id is not None + captured: Final = call_edge( + remote.edge, + "POST", + CHAT_PATH, + body=b'{"prompt":"blue"}', + headers={"content-type": "application/json"}, + ) + assert captured.status_code == 200 + observed: Final = remote.command( + ControlRequest(action="count", node=node, observation_id=observer.observation_id) + ) + assert observed.count == 1 + assert call_edge(remote.edge, "POST", "/control", body=b"{}").status_code == 404 + remote.phase(node, "call", True) + remote.phase(node, "teardown", True) + assert len(controller.results) == 1 and controller.results[0].publishable + with pytest.raises(RuntimeError, match="already started"): + remote.begin(node) + bundle: Final = load_bundle(recorder.root, profile="stateless_v1") + assert isinstance(bundle, LoadedBundle) + replay: Final = EdgeController({node: identity}, "replay", 200, replay_bundle=bundle) + with remote_controller(replay, mounts) as remote: + remote.begin("/opt/runner/" + node) + remote.phase(node, "setup", True) + response: Final = call_edge( + remote.edge, + "POST", + CHAT_PATH, + body=b'{"prompt":"blue"}', + headers={"content-type": "application/json"}, + ) + assert response.status_code == 200 and response.body == captured.body + remote.phase(node, "call", True) + remote.phase(node, "teardown", True) + assert len(provider.hits) == 1 + + def test_finalizer_failure_is_a_failed_remote_outcome(self, tmp_path: Path) -> None: + identity: Final = ScenarioIdentity("test_example.py::test_one", "a" * 64, "synthetic") + node: Final = canonical_scenario_id(identity.node) + recorder: Final = prepare_bundle(tmp_path / "capture", profile="stateless_v1") + assert isinstance(recorder, BundleRecorder) + controller: Final = EdgeController( + {node: identity}, "owner", 200, store=ScriptedReservations(), recorder=recorder + ) + with fake_provider() as provider, remote_controller(controller, {"openai": provider_url(provider)}) as remote: + remote.begin(node) + remote.phase(node, "setup", True) + assert ( + call_edge( + remote.edge, + "POST", + CHAT_PATH, + body=b'{"prompt":"blue"}', + headers={"content-type": "application/json"}, + ).status_code + == 200 + ) + remote.phase(node, "call", True) + with pytest.raises(RuntimeError, match="teardown must all pass"): + remote.phase(node, "teardown", False) + assert not controller.results[0].publishable + assert call_edge(remote.edge, "POST", CHAT_PATH, body=b"{}").status_code == REPLAY_MISS_STATUS + assert len(provider.hits) == 1 + + @pytest.mark.parametrize( + "control", + ["http://0.0.0.0:8081", "http://10.1.0.2:8081", "https://127.0.0.1:8081", "http://user@127.0.0.1:8081"], + ) + def test_driver_cannot_address_non_loopback_management(self, control: str) -> None: + with pytest.raises(ValueError, match="loopback"): + RemoteEdge(control, "http://10.1.0.2:8080") + + +@pytest.mark.parametrize("fault", ["duplicate", "release_denied", "release_uncertain"]) +def test_rejected_lifecycle_cannot_export_success(tmp_path: Path, fault: str) -> None: + class Store(ScriptedReservations): + def release(self, *, scenario_key: str, owner: str) -> StoreFailure | None: + return ( + None + if fault == "duplicate" + else StoreFailure("lease release failed", uncertain=fault == "release_uncertain") + ) + + identity: Final = ScenarioIdentity("test_example.py::test_one", "a" * 64, "synthetic") + node: Final = canonical_scenario_id(identity.node) + recorder: Final = prepare_bundle(tmp_path / "capture", profile="stateless_v1") + assert isinstance(recorder, BundleRecorder) + saved: Final = [] + controller: Final = EdgeController( + {node: identity}, "owner", 200, store=Store(), recorder=recorder, outcome_sink=saved.append + ) + assert controller.command(ControlRequest(action="begin", node=node)).ok + assert controller.command(ControlRequest(action="phase", node=node, phase="setup", passed=True)).ok + assert controller.before_attempt() is None + controller.response_finished(successful_response()) + if fault == "duplicate": + assert not controller.command(ControlRequest(action="phase", node=node, phase="setup", passed=True)).ok + controller.command(ControlRequest(action="phase", node=node, phase="call", passed=True)) + assert not controller.command(ControlRequest(action="phase", node=node, phase="teardown", passed=True)).ok + assert not controller.command(ControlRequest(action="status")).ok + assert saved and not saved[0][0].publishable + + +@pytest.mark.parametrize("fail_finalizer, fail_call", [(False, False), (True, False), (False, True)]) +def test_real_pytest_hooks_persist_finalizer_outcome(tmp_path: Path, fail_finalizer: bool, fail_call: bool) -> None: + harness: Final = Path(__file__).resolve().parents[1] / "e2e" + suite: Final = tmp_path / "suite" + suite.mkdir() + (suite / "conftest.py").write_bytes((harness / "conftest.py").read_bytes()) + (suite / "test_synthetic.py").write_text("""import os +import pytest +import requests +from e2e_config import unique_marker + +@pytest.fixture +def resource(): + yield + assert os.environ["FAIL_FINALIZER"] == "0", "synthetic finalizer failure" + +@pytest.mark.e2e +def test_one(resource): + assert unique_marker() == os.environ["EXPECTED_MARKER"] + response = requests.post(os.environ["EDGE_DATA"] + "/openai/v1/chat/completions", json={"prompt":"blue"}, timeout=5) + assert response.status_code == 200 + assert os.environ["FAIL_CALL"] == "0", "synthetic body failure" +""") + identity: Final = ScenarioIdentity("test_synthetic.py::test_one", "a" * 64, "synthetic") + node: Final = canonical_scenario_id(identity.node) + recorder: Final = prepare_bundle(tmp_path / "capture", profile="stateless_v1") + assert isinstance(recorder, BundleRecorder) + saved: Final = [] + controller: Final = EdgeController( + {node: identity}, "owner", 200, store=ScriptedReservations(), recorder=recorder, outcome_sink=saved.append + ) + with fake_provider() as provider, remote_controller(controller, {"openai": provider_url(provider)}) as remote: + environment: Final = { + "PATH": os.environ["PATH"], + "PYTHONPATH": str(harness), + "PYTEST_DISABLE_PLUGIN_AUTOLOAD": "1", + "E2E_FIXTURE_MODE": " Record ", + "E2E_PROVIDER_EDGE_CONTROL_URL": remote.control_url, + "E2E_PROVIDER_EDGE_DATA_URL": f"http://127.0.0.1:{remote.edge.port}", + "EDGE_DATA": f"http://127.0.0.1:{remote.edge.port}", + "LITELLM_PROXY_URL": provider_url(provider), + "FAIL_FINALIZER": "1" if fail_finalizer else "0", + "FAIL_CALL": "1" if fail_call else "0", + "EXPECTED_MARKER": hashlib.sha1(f"{node}#0".encode()).hexdigest()[:12], + } + result: Final = subprocess.run( + [sys.executable, "-m", "pytest", "-q", "-o", "addopts=", "test_synthetic.py"], + cwd=suite, + env=environment, + capture_output=True, + text=True, + timeout=30, + ) + assert result.returncode == (1 if fail_finalizer or fail_call else 0), result.stdout + result.stderr + assert len(saved) == 1 and len(saved[0]) == 1 + outcome: Final = saved[0][0] + assert outcome.outcome.setup + assert outcome.outcome.call is not fail_call + assert outcome.outcome.teardown is not fail_finalizer + assert outcome.publishable == (not fail_finalizer and not fail_call) + + +def test_aggregate_recording_limit_stops_before_second_disk_write(tmp_path: Path) -> None: + identity: Final = ScenarioIdentity("test_example.py::test_one", "a" * 64, "synthetic") + node: Final = canonical_scenario_id(identity.node) + recorder: Final = prepare_bundle(tmp_path / "capture", profile="stateless_v1") + assert isinstance(recorder, BundleRecorder) + controller: Final = EdgeController({node: identity}, "owner", 200, store=ScriptedReservations(), recorder=recorder) + backend: Final = controller.backend({}) + from provider_edge import RecordEdge + + assert isinstance(backend, RecordEdge) + request: Final = RecordedRequest(method="post", path="/openai/v1/chat/completions", headers={}) + response: Final = RecordedHttpResponse( + status_code=200, + headers={}, + body_b64=base64.b64encode(b'{"answer":"' + b"x" * (3 * 1024 * 1024) + b'"}').decode(), + ) + assert controller.command(ControlRequest(action="begin", node=node)).ok + assert controller.command(ControlRequest(action="phase", node=node, phase="setup", passed=True)).ok + assert controller.before_attempt() is None + _persist(backend, node, request, response) + assert 0 < controller.remaining_bytes() < SCENARIO_BYTES // 2 + assert controller.before_attempt() is None + with pytest.raises(OSError, match="capture byte limit"): + _persist(backend, node, request, response) + controller.command(ControlRequest(action="phase", node=node, phase="call", passed=True)) + assert not controller.command(ControlRequest(action="phase", node=node, phase="teardown", passed=True)).ok + assert not controller.results[0].publishable + assert sum(path.stat().st_size for path in recorder.root.rglob("*.json")) < SCENARIO_BYTES diff --git a/tests/e2e/CONTRIBUTING.md b/tests/e2e/CONTRIBUTING.md index 53a05931ca7..61a64087086 100644 --- a/tests/e2e/CONTRIBUTING.md +++ b/tests/e2e/CONTRIBUTING.md @@ -248,3 +248,14 @@ The semantic header set is `content-type`, `accept`, `anthropic-version`, `anthr Excluded transport and telemetry headers are `host`, `content-length`, `connection`, `accept-encoding`, `user-agent`, `traceparent`, `tracestate`, `x-request-id`, `x-client-request-id` and `x-stainless-*`. Inbound transfer-encoding is unsupported; send JSON with content-length framing. The destination represents host identity and the relay carries original body bytes. Replay does not verify credentials, SDK timeout/retry behavior, transport performance, model availability or stateful remote IDs. Live relay uses original request bytes and header values, never the stored identity Strict replay harness regression tests live in `tests/code_coverage_tests/test_provider_replay_harness.py`. The CircleCI `provider_replay_harness` job runs them alongside the existing legacy harness files with `--noconftest -o pythonpath=tests/e2e`; they need only synthetic HTTP providers and temporary fixture storage + + +### Trusted external provider edge + +An independently controlled harness can use `provider_edge_cli.py --config ` to serve the existing strict recorder or replay source. The JSON configuration declares the mode, owned bundle/output paths, canonical scenario roster, upstream mounts, ports and immutable scenario contract identities. Capture additionally requires a DynamoDB table/region, a lease deadline, an attempt cap (at most 12), and a `request_budgets` entry for every enrolled scenario containing the exact model and maximum output tokens. Provider header values can come from named environment variables through `credential_env`; those variables belong only to the trusted capture process. Replay configuration rejects capture credentials and a store. + +Set both `E2E_PROVIDER_EDGE_CONTROL_URL` and `E2E_PROVIDER_EDGE_DATA_URL` in the trusted pytest process, together with `E2E_FIXTURE_MODE=record` or `replay`. The control URL must address literal `127.0.0.1`; the data URL must be reachable by the proxy. Run pytest and the control server within the same isolated trust boundary. The pytest hooks send canonical scenario begin/setup/call/teardown outcomes and reset deterministic markers; the proxy receives only a provider mount URL. Calls before a scenario begins are rejected. Existing local edge configuration remains available when the two external URLs are absent. + +Capture reserves each outbound attempt before forwarding it. An uncertain reservation stops further attempts. Publication requires successful trusted setup, call and finalizers, complete successful responses, matching attempt/interaction counts and settled lease release. The CLI accepts at most four scenarios with a 5 MiB snapshot limit each, bounding an approved run to 20 MiB. A rejected capture remains diagnostic evidence; it never becomes a reusable success. + +The generic snapshot repository creates an immutable object, reads back its exact version/digest, then conditionally promotes the approved pointer. A failed pointer update leaves an unreferenced object. Readers reject snapshots at seven days and report refresh eligibility at 24 hours. Verified snapshots can be materialized to a new local bundle for replay without storage access. Deployment code must separately enforce trusted publisher inputs, IAM permissions, network isolation, full roster completion and cleanup; the harness does not establish those boundaries by itself. Keep configuration, recordings and run outcomes out of Git. diff --git a/tests/e2e/capture_policy.py b/tests/e2e/capture_policy.py new file mode 100644 index 00000000000..4cc8831f0d2 --- /dev/null +++ b/tests/e2e/capture_policy.py @@ -0,0 +1,207 @@ +from __future__ import annotations + +import base64 +import binascii +import hashlib +import json +from dataclasses import dataclass +from pathlib import PurePosixPath +from typing import Final + +from fixture_bundle import RecordedHttpResponse, RecordedResponse, RecordedStreamedResponse +from pydantic import BaseModel, ConfigDict, Field, JsonValue, TypeAdapter, ValidationError + +JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +SCENARIO_BYTES: Final = 5 * 1024 * 1024 +RUN_BYTES: Final = 20 * 1024 * 1024 +SOFT_AGE_SECONDS: Final = 24 * 60 * 60 +HARD_AGE_SECONDS: Final = 7 * SOFT_AGE_SECONDS + + +class RequestBudget(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + model: str = Field(min_length=1) + max_output_tokens: int = Field(gt=0, le=4096) + max_request_bytes: int = Field(default=65536, gt=0, le=65536) + + def error(self, body: bytes | None) -> str | None: + if body is None or len(body) > self.max_request_bytes: + return "capture request exceeds its byte budget" + try: + payload: Final = JSON_OBJECT.validate_json(body) + except ValidationError: + return "capture request is not JSON" + if payload.get("model") != self.model or payload.get("n", 1) != 1: + return "capture request changes its approved model or completion count" + limits: Final = tuple(payload[key] for key in ("max_tokens", "max_completion_tokens") if key in payload) + if not limits or any( + not isinstance(limit, int) or isinstance(limit, bool) or not 0 < limit <= self.max_output_tokens + for limit in limits + ): + return "capture request exceeds its output-token budget" + return None + + +def canonical_scenario_id(node: str) -> str: + path, separator, test = node.partition("::") + relative: Final = path.split("tests/e2e/", 1)[-1] + if ( + not separator + or not test + or not relative.endswith(".py") + or PurePosixPath(relative).is_absolute() + or ".." in PurePosixPath(relative).parts + or "\\" in path + or "\x00" in node + ): + raise ValueError("scenario must name an E2E test node under tests/e2e") + return f"tests/e2e/{relative}::{test}" + + +@dataclass(frozen=True, slots=True) +class ScenarioIdentity: + node: str + contract_sha256: str + credential_profile: str + matcher_version: str = "stateless_v1" + + def __post_init__(self) -> None: + canonical_scenario_id(self.node) + if len(self.contract_sha256) != 64 or any(c not in "0123456789abcdef" for c in self.contract_sha256): + raise ValueError("scenario contract must be a SHA-256 digest") + if not self.credential_profile or self.matcher_version != "stateless_v1": + raise ValueError("capture requires a credential profile and stateless_v1 matcher") + + @property + def key(self) -> str: + return hashlib.sha256( + json.dumps( + ( + canonical_scenario_id(self.node), + self.contract_sha256, + self.credential_profile, + self.matcher_version, + ), + separators=(",", ":"), + ).encode() + ).hexdigest() + + +@dataclass(frozen=True, slots=True) +class ScenarioOutcome: + setup: bool = False + call: bool = False + teardown: bool = False + + +def _invalid_count(value: JsonValue) -> bool: + return not isinstance(value, int) or isinstance(value, bool) or value < 0 + + +def _usage_error(value: JsonValue) -> str | None: + if not isinstance(value, dict): + return "invalid usage object" + counts: Final = tuple(value.get(key) for key in ("prompt_tokens", "completion_tokens", "total_tokens")) + if any(key in value for key in ("prompt_tokens", "completion_tokens", "total_tokens")): + if not all(type(count) is int and count >= 0 for count in counts): + return "invalid token usage" + prompt, completion, total = counts + if isinstance(prompt, int) and isinstance(completion, int) and total != prompt + completion: + return "inconsistent token usage" + if any(_invalid_count(value[key]) for key in ("input_tokens", "output_tokens") if key in value): + return "invalid token usage" + return None + + +def _json_error(body: bytes, *, streaming: bool = False) -> str | None: + try: + value: Final = JSON_OBJECT.validate_json(body) + except ValidationError: + return "invalid response JSON" + if not value or "error" in value or value.get("type") == "error": + return "provider response contains an error" + if value.get("status") in ("incomplete", "failed", "cancelled"): + return "provider response did not complete" + if streaming and value.get("usage") is None: + return None + return _usage_error(value["usage"]) if "usage" in value else None + + +def _stream_error(response: RecordedStreamedResponse) -> str | None: + if response.truncated is not None: + return "stream was truncated" + try: + body: Final = b"".join(base64.b64decode(chunk, validate=True) for chunk in response.chunks_b64).decode() + except (ValueError, UnicodeError, binascii.Error): + return "invalid stream encoding" + normalized: Final = body.replace("\r\n", "\n") + if not normalized.endswith("\n\n"): + return "stream lacks complete event framing" + events: Final = tuple( + data + for event in normalized.split("\n\n") + if (data := "\n".join(line[5:].lstrip(" ") for line in event.splitlines() if line.startswith("data:"))) + ) + if not events or any(not data for data in events): + return "stream contains no data event" + if "[DONE]" in events[:-1]: + return "stream has events after completion" + errors: Final = tuple(_json_error(data.encode(), streaming=True) for data in events if data != "[DONE]") + if any(errors): + return next(error for error in errors if error is not None) + parsed: Final = tuple(JSON_OBJECT.validate_json(data) for data in events if data != "[DONE]") + if not parsed: + return "stream contains no completion payload" + if events[-1] == "[DONE]": + if any( + not isinstance(entries := value.get("choices"), list) + or any(not isinstance(entry, dict) for entry in entries) + for value in parsed + ): + return "stream contains invalid choices" + choices: Final = tuple( + entry + for value in parsed + if isinstance(entries := value.get("choices"), list) + for entry in entries + if isinstance(entry, dict) + ) + if not choices or choices[-1].get("finish_reason") not in ("stop", "tool_calls", "function_call"): + return "stream lacks a successful completion reason" + if any(choice.get("index", 0) != 0 for choice in choices): + return "stream changes its approved completion count" + if any(choice.get("finish_reason") is not None for choice in choices[:-1]): + return "stream has choice events after completion" + return None + if parsed[-1].get("type") != "message_stop" or any(v.get("type") == "message_stop" for v in parsed[:-1]): + return "stream lacks a final success terminator" + if parsed[0].get("type") != "message_start": + return "stream lacks message start" + deltas: Final = tuple(value.get("delta") for value in parsed if value.get("type") == "message_delta") + if not any( + isinstance(delta, dict) and delta.get("stop_reason") in ("end_turn", "tool_use", "stop_sequence") + for delta in deltas + ): + return "stream lacks a successful completion reason" + return None + + +def _response_error(response: RecordedResponse) -> str | None: + if not 200 <= response.status_code < 300: + return "provider response was not successful" + if isinstance(response, RecordedHttpResponse): + try: + return _json_error(base64.b64decode(response.body_b64, validate=True)) + except (ValueError, binascii.Error): + return "invalid response encoding" + return _stream_error(response) + + +def publication_error(outcome: ScenarioOutcome, responses: tuple[RecordedResponse, ...]) -> str | None: + if not (outcome.setup and outcome.call and outcome.teardown): + return "trusted scenario setup, call and teardown must all pass" + if not responses: + return "capture contains no interactions" + if sum(len(response.model_dump_json().encode()) for response in responses) > SCENARIO_BYTES: + return "scenario exceeds capture byte limit" + return next((error for response in responses if (error := _response_error(response)) is not None), None) diff --git a/tests/e2e/capture_publication.py b/tests/e2e/capture_publication.py new file mode 100644 index 00000000000..49d721b2d45 --- /dev/null +++ b/tests/e2e/capture_publication.py @@ -0,0 +1,186 @@ +from __future__ import annotations + +import base64 +import hashlib +import uuid +from dataclasses import dataclass +from datetime import datetime +from typing import Final, Protocol + +from botocore.exceptions import BotoCoreError, ClientError +from botocore.response import StreamingBody +from capture_policy import SCENARIO_BYTES, ScenarioIdentity +from capture_snapshot import SnapshotFailure, verify_snapshot +from pydantic import BaseModel, ConfigDict, Field, JsonValue, ValidationError + + +class SnapshotPointer(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + revision: str + scenario_key: str = Field(pattern=r"^[a-f0-9]{64}$") + object_key: str + version_id: str = Field(min_length=1) + sha256: str = Field(pattern=r"^[a-f0-9]{64}$") + size: int = Field(gt=0, le=SCENARIO_BYTES) + + +class ObjectVersion(BaseModel): + VersionId: str = Field(min_length=1) + + +class ObjectDownload(ObjectVersion): + model_config = ConfigDict(arbitrary_types_allowed=True) + Body: StreamingBody + ContentLength: int + + +class PointerRead(BaseModel): + Item: dict[str, dict[str, str]] = {} + + +class ObjectClient(Protocol): + def put_object( + self, + *, + Bucket: str, + Key: str, + Body: bytes, + IfNoneMatch: str, + ContentType: str, + ChecksumSHA256: str, + ServerSideEncryption: str, + SSEKMSKeyId: str, + ) -> object: ... + + def get_object(self, *, Bucket: str, Key: str, VersionId: str) -> object: ... + + +class PointerClient(Protocol): + def get_item(self, *, TableName: str, Key: dict[str, JsonValue], ConsistentRead: bool) -> object: ... + + def put_item( + self, + *, + TableName: str, + Item: dict[str, JsonValue], + ConditionExpression: str, + ExpressionAttributeValues: dict[str, JsonValue] = ..., + ) -> object: ... + + +@dataclass(frozen=True, slots=True) +class SnapshotRepository: + objects: ObjectClient + pointers: PointerClient + bucket: str + table: str + kms_key: str + + def read_pointer(self, identity: ScenarioIdentity) -> SnapshotPointer | SnapshotFailure | None: + try: + response: Final = PointerRead.model_validate( + self.pointers.get_item( + TableName=self.table, + Key={"pk": {"S": f"pointer#{identity.key}"}}, + ConsistentRead=True, + ) + ) + if not response.Item: + return None + encoded: Final = response.Item.get("pointer_json", {}).get("S") + if encoded is None: + return SnapshotFailure("approved pointer is malformed") + pointer: Final = SnapshotPointer.model_validate_json(encoded) + return ( + pointer + if pointer.scenario_key == identity.key + else SnapshotFailure("approved pointer identity mismatch") + ) + except (ClientError, BotoCoreError, ValidationError): + return SnapshotFailure("approved pointer unavailable") + + def download( + self, pointer: SnapshotPointer, identity: ScenarioIdentity, *, now: datetime + ) -> bytes | SnapshotFailure: + if pointer.scenario_key != identity.key or pointer.version_id == "null": + return SnapshotFailure("approved pointer identity or version mismatch") + if pointer.object_key != f"approved/{identity.key}/{pointer.sha256}.json": + return SnapshotFailure("approved object key mismatch") + try: + response: Final = ObjectDownload.model_validate( + self.objects.get_object( + Bucket=self.bucket, + Key=pointer.object_key, + VersionId=pointer.version_id, + ) + ) + try: + if response.VersionId != pointer.version_id or response.ContentLength != pointer.size: + return SnapshotFailure("approved object version or size mismatch") + content: Final = response.Body.read(SCENARIO_BYTES + 1) + finally: + response.Body.close() + if len(content) != pointer.size: + return SnapshotFailure("approved object size mismatch") + verified: Final = verify_snapshot(content, expected_sha256=pointer.sha256, identity=identity, now=now) + return verified if isinstance(verified, SnapshotFailure) else content + except (ClientError, BotoCoreError, ValidationError, OSError): + return SnapshotFailure("approved snapshot unavailable") + + def publish( + self, + content: bytes, + identity: ScenarioIdentity, + *, + expected_revision: str | None, + now: datetime, + ) -> SnapshotPointer | SnapshotFailure: + digest: Final = hashlib.sha256(content).hexdigest() + verified: Final = verify_snapshot(content, expected_sha256=digest, identity=identity, now=now) + if isinstance(verified, SnapshotFailure): + return verified + key: Final = f"approved/{identity.key}/{digest}.json" + try: + version: Final = ObjectVersion.model_validate( + self.objects.put_object( + Bucket=self.bucket, + Key=key, + Body=content, + IfNoneMatch="*", + ContentType="application/json", + ChecksumSHA256=base64.b64encode(hashlib.sha256(content).digest()).decode(), + ServerSideEncryption="aws:kms", + SSEKMSKeyId=self.kms_key, + ) + ) + except (ClientError, BotoCoreError, ValidationError): + return SnapshotFailure("create-only snapshot publication failed or is uncertain") + pointer: Final = SnapshotPointer( + revision=uuid.uuid4().hex, + scenario_key=identity.key, + object_key=key, + version_id=version.VersionId, + sha256=digest, + size=len(content), + ) + readback: Final = self.download(pointer, identity, now=now) + if isinstance(readback, SnapshotFailure): + return readback + item: Final[dict[str, JsonValue]] = { + "pk": {"S": f"pointer#{identity.key}"}, + "pointer_revision": {"S": pointer.revision}, + "pointer_json": {"S": pointer.model_dump_json()}, + } + try: + if expected_revision is None: + self.pointers.put_item(TableName=self.table, Item=item, ConditionExpression="attribute_not_exists(pk)") + else: + self.pointers.put_item( + TableName=self.table, + Item=item, + ConditionExpression="pointer_revision = :expected", + ExpressionAttributeValues={":expected": {"S": expected_revision}}, + ) + except (ClientError, BotoCoreError): + return SnapshotFailure("pointer promotion failed or is uncertain; object remains unreferenced") + return pointer diff --git a/tests/e2e/capture_session.py b/tests/e2e/capture_session.py new file mode 100644 index 00000000000..5fd640d38b4 --- /dev/null +++ b/tests/e2e/capture_session.py @@ -0,0 +1,123 @@ +from __future__ import annotations + +import threading +import uuid +from collections.abc import Callable +from dataclasses import dataclass, field +from typing import Final, Protocol + +from capture_policy import ScenarioIdentity, ScenarioOutcome, publication_error +from fixture_bundle import RecordedResponse + + +@dataclass(frozen=True, slots=True) +class AttemptReserved: + attempt_id: str + + +@dataclass(frozen=True, slots=True) +class AttemptDenied: + reason: str + + +@dataclass(frozen=True, slots=True) +class AttemptUncertain: + reason: str + + +type Reservation = AttemptReserved | AttemptDenied | AttemptUncertain + + +class AttemptStore(Protocol): + def reserve(self, *, scenario_key: str, owner: str, attempt_id: str) -> Reservation: ... + + def complete(self, *, scenario_key: str, owner: str, attempt_id: str, successful: bool) -> str | None: ... + + +@dataclass(frozen=True, slots=True) +class CaptureResult: + identity: ScenarioIdentity + owner: str + attempts: tuple[str, ...] + outcome: ScenarioOutcome + response_count: int + error: str | None + + @property + def publishable(self) -> bool: + return self.error is None + + +@dataclass(slots=True) +class CaptureSession: + identity: ScenarioIdentity + owner: str + store: AttemptStore + new_attempt_id: Callable[[], str] = lambda: uuid.uuid4().hex + max_attempts: int = 12 + _lock: threading.Lock = field(default_factory=threading.Lock, init=False) + _attempts: tuple[str, ...] = field(default=(), init=False) + _responses: tuple[RecordedResponse, ...] = field(default=(), init=False) + _in_flight: bool = field(default=False, init=False) + _error: str | None = field(default=None, init=False) + _finished: CaptureResult | None = field(default=None, init=False) + + def __post_init__(self) -> None: + if not self.owner or not 1 <= self.max_attempts <= 12: + raise ValueError("capture needs an owner and an attempt cap between 1 and 12") + + def before_attempt(self) -> str | None: + with self._lock: + if self._finished is not None: + return "scenario is closed" + if self._error is not None: + return self._error + if self._in_flight: + self._error = "concurrent scenario request rejected" + return self._error + if len(self._attempts) >= self.max_attempts: + self._error = "capture attempt cap reached" + return self._error + attempt_id: Final = self.new_attempt_id() + reservation: Final = self.store.reserve( + scenario_key=self.identity.key, owner=self.owner, attempt_id=attempt_id + ) + if not isinstance(reservation, AttemptReserved): + self._error = reservation.reason + return self._error + if reservation.attempt_id != attempt_id or attempt_id in self._attempts: + self._error = "reservation identity mismatch" + return self._error + self._attempts = (*self._attempts, attempt_id) + self._in_flight = True + return None + + def response_finished(self, response: RecordedResponse) -> None: + with self._lock: + if not self._in_flight or self._finished is not None: + self._error = "response does not belong to an active attempt" + return + self._responses = (*self._responses, response) + self._error = self._error or publication_error(ScenarioOutcome(True, True, True), self._responses) + completion_error: Final = self.store.complete( + scenario_key=self.identity.key, + owner=self.owner, + attempt_id=self._attempts[-1], + successful=self._error is None, + ) + self._error = self._error or completion_error + self._in_flight = completion_error is not None + + def finish(self, outcome: ScenarioOutcome) -> CaptureResult: + with self._lock: + if self._finished is not None: + return self._finished + error: Final = ( + self._error + or ("outbound outcome is uncertain" if self._in_flight else None) + or publication_error(outcome, self._responses) + ) + self._finished = CaptureResult( + self.identity, self.owner, self._attempts, outcome, len(self._responses), error + ) + return self._finished diff --git a/tests/e2e/capture_snapshot.py b/tests/e2e/capture_snapshot.py new file mode 100644 index 00000000000..9a03288d216 --- /dev/null +++ b/tests/e2e/capture_snapshot.py @@ -0,0 +1,130 @@ +from __future__ import annotations + +import hashlib +from dataclasses import dataclass +from datetime import datetime +from pathlib import Path +from typing import Final, Literal + +from capture_policy import ( + HARD_AGE_SECONDS, + SCENARIO_BYTES, + SOFT_AGE_SECONDS, + ScenarioIdentity, + ScenarioOutcome, + canonical_scenario_id, + publication_error, +) +from capture_session import CaptureResult +from fixture_bundle import Interaction, LoadedBundle, Manifest, interaction_filename, slug_for_test +from pydantic import BaseModel, ConfigDict, Field, ValidationError + + +class CaptureProvenance(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + test_revision: str = Field(pattern=r"^[a-f0-9]{40}$") + candidate_revision: str = Field(pattern=r"^[a-f0-9]{40}$") + runner_digest: str = Field(pattern=r"^sha256:[a-f0-9]{64}$") + + +class ScenarioSnapshot(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + schema_version: Literal[1] = 1 + identity: ScenarioIdentity + manifest: Manifest + interactions: tuple[Interaction, ...] + provenance: CaptureProvenance + owner: str + attempts: tuple[str, ...] + outcome: ScenarioOutcome + + +@dataclass(frozen=True, slots=True) +class SnapshotFailure: + reason: str + + +def build_snapshot( + result: CaptureResult, bundle: LoadedBundle, provenance: CaptureProvenance +) -> bytes | SnapshotFailure: + if not result.publishable: + return SnapshotFailure(result.error or "capture was not approved") + node: Final = canonical_scenario_id(result.identity.node) + interactions: Final = bundle.interactions.get(slug_for_test(node), ()) + if result.response_count != len(interactions) or result.response_count != len(result.attempts): + return SnapshotFailure("capture interaction and attempt counts differ") + if bundle.manifest.match_profile != "stateless_v1" or bundle.manifest.format_version != 5: + return SnapshotFailure("capture requires a strict bundle") + if any(interaction.request.strict_identity is None for interaction in interactions): + return SnapshotFailure("capture contains a legacy request") + error: Final = publication_error(result.outcome, tuple(i.response for i in interactions)) + if error is not None: + return SnapshotFailure(error) + encoded: Final = ( + ScenarioSnapshot( + identity=result.identity, + manifest=bundle.manifest, + interactions=interactions, + provenance=provenance, + owner=result.owner, + attempts=result.attempts, + outcome=result.outcome, + ) + .model_dump_json() + .encode() + ) + return encoded if len(encoded) <= SCENARIO_BYTES else SnapshotFailure("scenario exceeds capture byte limit") + + +def verify_snapshot( + content: bytes, + *, + expected_sha256: str, + identity: ScenarioIdentity, + now: datetime, +) -> ScenarioSnapshot | SnapshotFailure: + if len(content) > SCENARIO_BYTES or hashlib.sha256(content).hexdigest() != expected_sha256: + return SnapshotFailure("snapshot size or digest mismatch") + try: + snapshot: Final = ScenarioSnapshot.model_validate_json(content) + except (ValueError, ValidationError): + return SnapshotFailure("snapshot schema invalid") + if ( + snapshot.identity.key != identity.key + or snapshot.manifest.match_profile != "stateless_v1" + or snapshot.manifest.format_version != 5 + ): + return SnapshotFailure("snapshot identity or matcher mismatch") + recorded_at: Final = snapshot.manifest.recorded_at + if recorded_at.tzinfo is None or now.tzinfo is None: + return SnapshotFailure("snapshot timestamps must include a timezone") + age: Final = (now - recorded_at).total_seconds() + if age < 0 or age >= HARD_AGE_SECONDS: + return SnapshotFailure("snapshot is outside its hard age limit") + if ( + not snapshot.interactions + or len(snapshot.interactions) != len(snapshot.attempts) + or len(set(snapshot.attempts)) != len(snapshot.attempts) + ): + return SnapshotFailure("snapshot interaction and attempt counts differ") + if any(i.request.strict_identity is None for i in snapshot.interactions): + return SnapshotFailure("snapshot contains a legacy request") + error: Final = publication_error(snapshot.outcome, tuple(i.response for i in snapshot.interactions)) + if error is not None: + return SnapshotFailure(error) + return snapshot + + +def refresh_due(snapshot: ScenarioSnapshot, *, now: datetime) -> bool: + return (now - snapshot.manifest.recorded_at).total_seconds() >= SOFT_AGE_SECONDS + + +def materialize_snapshot(snapshot: ScenarioSnapshot, destination: Path) -> None: + destination.mkdir(parents=True, exist_ok=False) + (destination / "manifest.json").write_text(snapshot.manifest.model_dump_json(), encoding="utf-8") + scenario_dir: Final = destination / slug_for_test(canonical_scenario_id(snapshot.identity.node)) + scenario_dir.mkdir() + for ordinal, interaction in enumerate(snapshot.interactions): + (scenario_dir / interaction_filename(ordinal, interaction.request)).write_text( + interaction.model_dump_json(), encoding="utf-8" + ) diff --git a/tests/e2e/capture_store.py b/tests/e2e/capture_store.py new file mode 100644 index 00000000000..146cf03b638 --- /dev/null +++ b/tests/e2e/capture_store.py @@ -0,0 +1,202 @@ +from __future__ import annotations + +import hashlib +import time +from collections.abc import Callable +from dataclasses import dataclass +from typing import Final, Protocol + +from botocore.exceptions import BotoCoreError, ClientError +from capture_session import AttemptDenied, AttemptReserved, AttemptUncertain, Reservation +from pydantic import JsonValue + + +@dataclass(frozen=True, slots=True) +class StoreFailure: + reason: str + uncertain: bool = False + + +class CaptureLeaseStore(Protocol): + def acquire(self, *, scenario_key: str, owner: str, expires_at: int) -> StoreFailure | None: ... + + def reserve(self, *, scenario_key: str, owner: str, attempt_id: str) -> Reservation: ... + + def complete(self, *, scenario_key: str, owner: str, attempt_id: str, successful: bool) -> str | None: ... + + def release(self, *, scenario_key: str, owner: str) -> StoreFailure | None: ... + + +class DynamoClient(Protocol): + def put_item( + self, + *, + TableName: str, + Item: dict[str, JsonValue], + ConditionExpression: str, + ExpressionAttributeValues: dict[str, JsonValue] = ..., + ) -> object: ... + + def update_item( + self, + *, + TableName: str, + Key: dict[str, JsonValue], + UpdateExpression: str, + ConditionExpression: str, + ExpressionAttributeValues: dict[str, JsonValue], + ) -> object: ... + + def transact_write_items(self, *, TransactItems: list[dict[str, JsonValue]], ClientRequestToken: str) -> object: ... + + +def _store_failure(error: ClientError | BotoCoreError) -> StoreFailure: + if isinstance(error, ClientError): + code: Final = str(error.response.get("Error", {}).get("Code", "")) + if code in ("ConditionalCheckFailedException", "TransactionCanceledException"): + return StoreFailure("store condition rejected") + return StoreFailure("store operation outcome uncertain", uncertain=True) + + +@dataclass(frozen=True, slots=True) +class DynamoCaptureStore: + client: DynamoClient + table: str + clock: Callable[[], float] = time.time + + def create_run(self, *, owner: str, cap: int, expires_at: int) -> StoreFailure | None: + if not owner or not 1 <= cap <= 12 or expires_at <= self.clock(): + return StoreFailure("invalid capture run bounds") + try: + self.client.put_item( + TableName=self.table, + Item={ + "pk": {"S": f"run#{owner}"}, + "attempt_count": {"N": "0"}, + "attempt_cap": {"N": str(cap)}, + "lease_expires": {"N": str(expires_at)}, + }, + ConditionExpression="attribute_not_exists(pk)", + ) + except (ClientError, BotoCoreError) as error: + return _store_failure(error) + return None + + def acquire(self, *, scenario_key: str, owner: str, expires_at: int) -> StoreFailure | None: + now: Final = int(self.clock()) + if not owner or not scenario_key or expires_at <= now: + return StoreFailure("invalid lease bounds") + try: + self.client.put_item( + TableName=self.table, + Item={ + "pk": {"S": f"lease#{scenario_key}"}, + "lease_owner": {"S": owner}, + "lease_expires": {"N": str(expires_at)}, + "lease_settled": {"BOOL": False}, + }, + ConditionExpression="attribute_not_exists(pk) OR (lease_settled = :true AND lease_expires <= :now)", + ExpressionAttributeValues={":true": {"BOOL": True}, ":now": {"N": str(now)}}, + ) + except (ClientError, BotoCoreError) as error: + return _store_failure(error) + return None + + def reserve(self, *, scenario_key: str, owner: str, attempt_id: str) -> Reservation: + if not owner or not scenario_key or not 1 <= len(attempt_id) <= 36: + return AttemptDenied("invalid attempt identity") + now: Final = str(int(self.clock())) + transaction: Final[list[dict[str, JsonValue]]] = [ + { + "Update": { + "TableName": self.table, + "Key": {"pk": {"S": f"run#{owner}"}}, + "UpdateExpression": "SET attempt_count = attempt_count + :one", + "ConditionExpression": "attempt_count < attempt_cap AND lease_expires > :now", + "ExpressionAttributeValues": {":one": {"N": "1"}, ":now": {"N": now}}, + } + }, + { + "Update": { + "TableName": self.table, + "Key": {"pk": {"S": f"lease#{scenario_key}"}}, + "UpdateExpression": "SET active_attempt = :attempt", + "ConditionExpression": "lease_owner = :owner AND lease_expires > :now AND lease_settled = :false AND attribute_not_exists(active_attempt)", + "ExpressionAttributeValues": { + ":owner": {"S": owner}, + ":now": {"N": now}, + ":false": {"BOOL": False}, + ":attempt": {"S": attempt_id}, + }, + } + }, + { + "Put": { + "TableName": self.table, + "Item": { + "pk": {"S": f"attempt#{owner}#{attempt_id}"}, + "recorded_scenario": {"S": scenario_key}, + "attempt_completed": {"BOOL": False}, + }, + "ConditionExpression": "attribute_not_exists(pk)", + } + }, + ] + try: + self.client.transact_write_items(TransactItems=transaction, ClientRequestToken=attempt_id) + except (ClientError, BotoCoreError) as error: + failure: Final = _store_failure(error) + return AttemptUncertain(failure.reason) if failure.uncertain else AttemptDenied(failure.reason) + return AttemptReserved(attempt_id) + + def complete(self, *, scenario_key: str, owner: str, attempt_id: str, successful: bool) -> str | None: + transaction: Final[list[dict[str, JsonValue]]] = [ + { + "Update": { + "TableName": self.table, + "Key": {"pk": {"S": f"lease#{scenario_key}"}}, + "UpdateExpression": "REMOVE active_attempt", + "ConditionExpression": "lease_owner = :owner AND active_attempt = :attempt", + "ExpressionAttributeValues": {":owner": {"S": owner}, ":attempt": {"S": attempt_id}}, + } + }, + { + "Update": { + "TableName": self.table, + "Key": {"pk": {"S": f"attempt#{owner}#{attempt_id}"}}, + "UpdateExpression": "SET attempt_completed = :true, attempt_successful = :attempt_successful", + "ConditionExpression": "recorded_scenario = :recorded_scenario AND attempt_completed = :false", + "ExpressionAttributeValues": { + ":recorded_scenario": {"S": scenario_key}, + ":false": {"BOOL": False}, + ":true": {"BOOL": True}, + ":attempt_successful": {"BOOL": successful}, + }, + } + }, + ] + try: + self.client.transact_write_items( + TransactItems=transaction, + ClientRequestToken=hashlib.sha256(f"done:{owner}:{attempt_id}".encode()).hexdigest()[:36], + ) + except (ClientError, BotoCoreError) as error: + return _store_failure(error).reason + return None + + def release(self, *, scenario_key: str, owner: str) -> StoreFailure | None: + try: + self.client.update_item( + TableName=self.table, + Key={"pk": {"S": f"lease#{scenario_key}"}}, + UpdateExpression="SET lease_settled = :true, lease_expires = :now", + ConditionExpression="lease_owner = :owner AND attribute_not_exists(active_attempt)", + ExpressionAttributeValues={ + ":owner": {"S": owner}, + ":true": {"BOOL": True}, + ":now": {"N": str(int(self.clock()))}, + }, + ) + except (ClientError, BotoCoreError) as error: + return _store_failure(error) + return None diff --git a/tests/e2e/conftest.py b/tests/e2e/conftest.py index 36569896125..01aeb6e8ae4 100644 --- a/tests/e2e/conftest.py +++ b/tests/e2e/conftest.py @@ -21,15 +21,21 @@ from typing import Final import pytest import requests -from e2e_config import CONTROL_PLANE_BASE_URL, FIXTURE_DIR, FIXTURE_MODE_RAW, PROXY_BASE_URL, unique_marker +from e2e_config import CONTROL_PLANE_BASE_URL, FIXTURE_DIR, FIXTURE_MODE_RAW, PROXY_BASE_URL, REMOTE_EDGE, unique_marker from e2e_db import RESET_OPT_IN_ENV, reset_spend_logs, run_spend_log_cleanup from e2e_http import unwrap -from fixture_mode import fixture_mode_collection_error, fixture_report_lines +from fixture_mode import ( + fixture_mode_collection_error, + fixture_report_lines, + parse_fixture_mode, + reset_deterministic_markers, +) from idp import Identity, Keycloak, keycloak_from_env from junit_properties import attach_result_properties from lifecycle import ProxyClientProvider, ResourceManager from models import TeamNewBody, UserNewBody, UserNewResponse from provider_edge import replay_leftover_error +from provider_edge_control import ControlRequest from proxy_client import ProxyClient, build_proxy_client _E2E_TEST_RAN = pytest.StashKey[bool]() @@ -100,6 +106,11 @@ def pytest_sessionstart(session: pytest.Session) -> None: """Abort before collection when E2E_FIXTURE_MODE can never work: an unknown mode value, or replay against a missing, unreadable, or stale bundle (the stale message names the bundle's age). Live and record modes pass through.""" + if REMOTE_EDGE is not None: + status: Final = REMOTE_EDGE.command(ControlRequest(action="status")) + if status.mode != parse_fixture_mode(FIXTURE_MODE_RAW): + raise pytest.UsageError("remote edge mode does not match requested fixture mode") + return reason = fixture_mode_collection_error( FIXTURE_MODE_RAW, FIXTURE_DIR, now=datetime.now(timezone.utc) ) @@ -108,6 +119,8 @@ def pytest_sessionstart(session: pytest.Session) -> None: def pytest_report_header(config: pytest.Config) -> list[str]: + if REMOTE_EDGE is not None: + return [f"e2e fixture mode: {FIXTURE_MODE_RAW} through trusted remote edge"] return fixture_report_lines(FIXTURE_MODE_RAW, FIXTURE_DIR, now=datetime.now(timezone.utc)) @@ -157,6 +170,9 @@ def pytest_runtest_setup(item: pytest.Item) -> None: the proxy too: only provider-bound traffic replays from the bundle.""" if item.get_closest_marker("e2e") is None: return + if REMOTE_EDGE is not None: + REMOTE_EDGE.begin(item.nodeid) + reset_deterministic_markers(item.nodeid) reason = _proxy_fail_reason() if reason is not None: pytest.fail(reason) @@ -182,6 +198,12 @@ def pytest_runtest_makereport( report = yield if report.when == "call": item.stash[_CALL_PASSED] = report.passed + if REMOTE_EDGE is not None and item.get_closest_marker("e2e") is not None: + try: + REMOTE_EDGE.phase(item.nodeid, report.when, report.passed) + except RuntimeError as error: + report.outcome = "failed" + report.longrepr = str(error) return report @@ -193,6 +215,8 @@ def pytest_runtest_teardown(item: pytest.Item) -> Generator[None, None, None]: yield so fixture finalizers replay their recorded calls first. Failed tests are left alone - their own failure already explains any unconsumed tail.""" result = yield + if REMOTE_EDGE is not None: + return result if not item.stash.get(_CALL_PASSED, False): return result reason = replay_leftover_error( diff --git a/tests/e2e/e2e_config.py b/tests/e2e/e2e_config.py index e15a0cd0f8f..5d3c5c7e74e 100644 --- a/tests/e2e/e2e_config.py +++ b/tests/e2e/e2e_config.py @@ -15,6 +15,7 @@ from typing import Final from dotenv import load_dotenv from fixture_mode import deterministic_marker, parse_fixture_mode from provider_edge import provider_edge_api_base +from provider_edge_remote import RemoteEdge # Local runs keep provider / DataDog keys in tests/e2e/.env (see CONTRIBUTING.md). # Compose injects them into the proxy container, but pytest on the host does not @@ -120,6 +121,10 @@ PROVIDER_EDGE_BIND_HOST = os.environ.get("E2E_PROVIDER_EDGE_BIND_HOST", "").stri PROVIDER_EDGE_ADVERTISE_HOST = ( os.environ.get("E2E_PROVIDER_EDGE_ADVERTISE_HOST", "").strip() or PROVIDER_EDGE_BIND_HOST ) +REMOTE_EDGE: Final = ( + RemoteEdge(os.environ.get("E2E_PROVIDER_EDGE_CONTROL_URL", ""), os.environ.get("E2E_PROVIDER_EDGE_DATA_URL", "")) + if os.environ.get("E2E_PROVIDER_EDGE_CONTROL_URL") or os.environ.get("E2E_PROVIDER_EDGE_DATA_URL") else None +) # Deliberately modest concurrency. The suite shares its proxy with every other # suite in the run, and 750 users at spawn rate 50 saturated the request path hard @@ -206,6 +211,7 @@ def provider_edge_base(mount: str) -> str | None: bind_host=PROVIDER_EDGE_BIND_HOST, advertise_host=PROVIDER_EDGE_ADVERTISE_HOST, forward_timeout=REQUEST_TIMEOUT, + remote=REMOTE_EDGE, ) diff --git a/tests/e2e/fixture_mode.py b/tests/e2e/fixture_mode.py index 9a7c1b6db12..bb96cb7bebf 100644 --- a/tests/e2e/fixture_mode.py +++ b/tests/e2e/fixture_mode.py @@ -19,6 +19,7 @@ from datetime import datetime from pathlib import Path from typing import Final, Literal, assert_never +from capture_policy import canonical_scenario_id from fixture_bundle import ( FreshBundle, StaleBundle, @@ -72,12 +73,20 @@ def deterministic_marker() -> str: the Nth marker of a test is a pure function of the test's node id and N, so a replay run regenerates exactly the model names, prompts, and tags the record run sent and every recorded provider interaction still matches its key.""" - test_key = current_test_key() + test_key: Final = ( + canonical_scenario_id(current_test_key()) + if os.environ.get("E2E_PROVIDER_EDGE_CONTROL_URL") + else current_test_key() + ) ordinal = _marker_ordinals.get(test_key, 0) _marker_ordinals[test_key] = ordinal + 1 return hashlib.sha1(f"{test_key}#{ordinal}".encode()).hexdigest()[:12] +def reset_deterministic_markers(node: str) -> None: + _marker_ordinals.pop(canonical_scenario_id(node), None) + + def fixture_mode_collection_error(mode_raw: str, bundle_dir: Path, *, now: datetime) -> str | None: """Session-abort reason for a fixture-mode setup that can never work, or None. Called at collection time (conftest pytest_sessionstart) so a stale or missing diff --git a/tests/e2e/provider_edge.py b/tests/e2e/provider_edge.py index de36895ebb6..0e83c1bd431 100644 --- a/tests/e2e/provider_edge.py +++ b/tests/e2e/provider_edge.py @@ -45,14 +45,14 @@ import hashlib import re import threading from collections import deque -from collections.abc import Generator, Mapping, Sequence +from collections.abc import Callable, Generator, Mapping, Sequence from contextlib import closing, contextmanager from dataclasses import dataclass, field from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from itertools import islice from pathlib import Path from types import MappingProxyType -from typing import Final, Literal, assert_never +from typing import TYPE_CHECKING, Final, Literal, Protocol, assert_never from urllib.parse import parse_qsl, urlsplit from e2e_http import ( @@ -95,6 +95,9 @@ from fixture_mode import ( from fixture_profile import IneligibleRequest, MatchProfile, match_profile, strict_identity from pydantic import JsonValue, TypeAdapter +if TYPE_CHECKING: + from provider_edge_remote import RemoteEdge + EDGE_MOUNTS: Final[Mapping[str, str]] = MappingProxyType( { "openai": "https://api.openai.com", @@ -130,13 +133,9 @@ _RESPONSE_DROPPED_HEADERS: Final[frozenset[str]] = _HOP_BY_HOP_HEADERS | { _JSON: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue) -_BOUNDARY_PATTERN: Final = re.compile( - r'(?:^|;)\s*boundary\s*=\s*(?:"([^"]*)"|([^;,\s]+))', re.IGNORECASE -) +_BOUNDARY_PATTERN: Final = re.compile(r'(?:^|;)\s*boundary\s*=\s*(?:"([^"]*)"|([^;,\s]+))', re.IGNORECASE) _DISPOSITION_NAME_PATTERN: Final = re.compile(r'(?:^|;)\s*name="([^"]*)"', re.IGNORECASE) -_DISPOSITION_FILENAME_PATTERN: Final = re.compile( - r'(?:^|;)\s*filename="([^"]*)"', re.IGNORECASE -) +_DISPOSITION_FILENAME_PATTERN: Final = re.compile(r'(?:^|;)\s*filename="([^"]*)"', re.IGNORECASE) _UNPARSED_MULTIPART: Final = "" _BOUNDARY_PLACEHOLDER: Final = b"--" _BINARY_FIELD_PREFIX: Final = " None: @@ -452,7 +452,7 @@ class ReplaySource: return self._pools.get(slug, {}) def next_interaction(self, request: RecordedRequest) -> Interaction: - test_key: Final = current_test_key() + test_key: Final = self.test_key() slug: Final = slug_for_test(test_key) pool: Final = self._pool(slug) canonical: Final = canonicalize(request) @@ -494,6 +494,13 @@ class RecordEdge: recorder: BundleRecorder lock: threading.Lock + test_key: Callable[[], str] = current_test_key + before_attempt: Callable[[], str | None] | None = None + response_finished: Callable[[RecordedResponse], None] | None = None + response_byte_limit: int | Callable[[], int] | None = None + recording_error: Callable[[RecordedRequest, RecordedResponse], str | None] | None = None + upstream_headers: Mapping[str, Mapping[str, str]] = field(default_factory=lambda: MappingProxyType({})) + request_error: Callable[[bytes | None], str | None] | None = None @dataclass(frozen=True, slots=True) @@ -514,6 +521,7 @@ class ProviderRequestObservation: marker: str _count: int = field(default=0, init=False) _lock: threading.Lock = field(default_factory=threading.Lock, init=False) + _remote_count: Callable[[], int] | None = field(default=None, init=False) def observe(self, body: bytes | None) -> None: if body is not None and self.marker.encode() in body: @@ -523,7 +531,17 @@ class ProviderRequestObservation: @property def count(self) -> int: with self._lock: - return self._count + return self._remote_count() if self._remote_count is not None else self._count + + def bind_remote_count(self, reader: Callable[[], int]) -> None: + with self._lock: + self._remote_count = reader + + def finish_remote_count(self) -> None: + with self._lock: + if self._remote_count is not None: + self._count = self._remote_count() + self._remote_count = None @dataclass(frozen=True, slots=True) @@ -646,11 +664,17 @@ def _upstream_url(upstream_base: str, upstream_path: str, query: str) -> str: return f"{url}?{query}" if query else url -def _persist( - backend: RecordEdge, test_key: str, request: RecordedRequest, response: RecordedResponse -) -> None: +def _persist(backend: RecordEdge, test_key: str, request: RecordedRequest, response: RecordedResponse) -> None: with backend.lock: + if backend.recording_error is not None and (error := backend.recording_error(request, response)) is not None: + raise OSError(error) backend.recorder.record(test_key=test_key, request=request, response=response) + if backend.response_finished is not None: + backend.response_finished(response) + + +def _response_byte_limit(backend: RecordEdge) -> int | None: + return backend.response_byte_limit() if callable(backend.response_byte_limit) else backend.response_byte_limit def _recording_steps( @@ -670,12 +694,16 @@ def _recording_steps( have gone red.""" collected: list[bytes] = [] truncated: str | None = None + byte_limit: Final = _response_byte_limit(backend) try: with closing(head.steps) as steps: for step in steps: match step: case StreamChunk(): - pass + if byte_limit is not None and sum(map(len, collected)) + len(step.data) > byte_limit: + truncated = "capture response byte limit reached" # rebind-ok: preserve truncation in the persisted stream + yield StreamTruncation(reason=truncated) + return case StreamTruncation(reason=reason): truncated = f"upstream: {reason}" case _: @@ -696,7 +724,7 @@ def _recording_steps( ) -def _drain_to_response(head: StreamHead) -> RecordedHttpResponse: +def _drain_to_response(head: StreamHead, byte_limit: int | None = None) -> RecordedHttpResponse: """A response the detection rule did not call streamed: drain the same step iterator, join the pieces, and store today's buffered shape byte for byte. A truncation part way through degrades to the synthetic 502 exactly as the eager @@ -707,6 +735,8 @@ def _drain_to_response(head: StreamHead) -> RecordedHttpResponse: for step in steps: match step: case StreamChunk(data=data): + if byte_limit is not None and sum(map(len, pieces)) + len(data) > byte_limit: + return _network_error_response("capture response byte limit reached") pieces.append(data) case StreamTruncation(reason=reason): return _network_error_response(reason) @@ -725,9 +755,18 @@ def _handle_record( body: bytes | None, timeout: float, ) -> EdgeOutcome: - test_key: Final = current_test_key() + test_key: Final = backend.test_key() + if backend.request_error is not None: + request_error: Final = backend.request_error(body) + if request_error is not None: + return _text_reply(REPLAY_MISS_STATUS, request_error) + if backend.before_attempt is not None: + denial: Final = backend.before_attempt() + if denial is not None: + return _text_reply(REPLAY_MISS_STATUS, denial) forwarded: Final = { - name: value for name, value in headers.items() if name.lower() not in _REQUEST_DROPPED_HEADERS + **{name: value for name, value in headers.items() if name.lower() not in _REQUEST_DROPPED_HEADERS}, + **backend.upstream_headers.get(request.path.lstrip("/").partition("/")[0], {}), } head: Final = forward_stream(method, url, headers=forwarded, body=body, timeout=timeout) match head: @@ -742,7 +781,7 @@ def _handle_record( steps=_recording_steps(backend, test_key, request, head), ) case StreamHead(): - buffered: Final = _drain_to_response(head) + buffered: Final = _drain_to_response(head, _response_byte_limit(backend)) _persist(backend, test_key, request, buffered) return _recorded_outcome(buffered) case _: @@ -841,6 +880,16 @@ def handle_edge_request( assert_never(backend) +class RequestGuard(Protocol): + def begin_request(self) -> str | None: ... + + def end_request(self) -> None: ... + + +class RequestObserver(Protocol): + def observe(self, body: bytes | None) -> None: ... + + class _EdgeHandler(BaseHTTPRequestHandler): protocol_version = "HTTP/1.1" @@ -862,6 +911,19 @@ class _EdgeHandler(BaseHTTPRequestHandler): def _handle(self) -> None: edge_server: Final = self.server assert isinstance(edge_server, _EdgeHTTPServer) + if edge_server.guard is not None: + denial: Final = edge_server.guard.begin_request() + if denial is not None: + self._write_reply(_text_reply(REPLAY_MISS_STATUS, denial)) + self.close_connection = True + return + try: + self._serve(edge_server) + finally: + if edge_server.guard is not None: + edge_server.guard.end_request() + + def _serve(self, edge_server: _EdgeHTTPServer) -> None: length: Final = int(self.headers.get("content-length") or "0") body: Final = self.rfile.read(length) if length else None if edge_server.observation is not None: @@ -939,13 +1001,15 @@ class _EdgeHTTPServer(ThreadingHTTPServer): backend: EdgeBackend, mounts: Mapping[str, str], forward_timeout: float, - observation: ProviderRequestObservation | None, + observation: RequestObserver | None, + guard: RequestGuard | None, ) -> None: super().__init__(bind, _EdgeHandler) self.backend: Final = backend self.mounts: Final = mounts self.forward_timeout: Final = forward_timeout self.observation: Final = observation + self.guard: Final = guard @dataclass(frozen=True, slots=True) @@ -974,14 +1038,16 @@ def start_provider_edge( bind_host: str = "127.0.0.1", advertise_host: str | None = None, forward_timeout: float = 60.0, - observation: ProviderRequestObservation | None = None, + observation: RequestObserver | None = None, + guard: RequestGuard | None = None, + bind_port: int = 0, ) -> RunningEdge: """Boot an edge server on an OS-assigned port in a daemon thread. ``advertise_host`` is what api_base URLs name (it differs from the bind host when the proxy runs in a container and reaches the host machine via a gateway address like host.docker.internal).""" server: Final = _EdgeHTTPServer( - (bind_host, 0), backend=backend, mounts=mounts, forward_timeout=forward_timeout, observation=observation + (bind_host, bind_port), backend=backend, mounts=mounts, forward_timeout=forward_timeout, observation=observation, guard=guard ) thread: Final = threading.Thread(target=server.serve_forever, name="e2e-provider-edge", daemon=True) thread.start() @@ -1047,10 +1113,15 @@ def provider_edge_api_base( bind_host: str, advertise_host: str, forward_timeout: float = 60.0, + remote: RemoteEdge | None = None, ) -> str | None: """The api_base a suite gives an edge-wired deployment: None in live mode (the deployment keeps its real provider api_base) and the process-wide edge server's mount URL in record and replay, booting the server on first use.""" + if remote is not None: + if parse_fixture_mode(mode_raw) not in ("record", "replay") or mount not in EDGE_MOUNTS: + raise ValueError("remote edge requires a supported record/replay mount") + return remote.edge.api_base(mount) mode: Final = parse_fixture_mode(mode_raw) match mode: case InvalidFixtureMode(value=value): @@ -1092,7 +1163,18 @@ def observed_provider_edge( advertise_host: str, forward_timeout: float = 60.0, mounts: Mapping[str, str] = EDGE_MOUNTS, + remote: RemoteEdge | None = None, ) -> Generator[ProviderEdge, None, None]: + if remote is not None: + if parse_fixture_mode(mode_raw) not in ("record", "replay"): + raise ValueError("remote observation requires record/replay mode") + observation_id: Final = remote.observe(observation.marker) + observation.bind_remote_count(lambda: remote.count(observation_id)) + try: + yield remote.edge + finally: + observation.finish_remote_count() + return running: Final = start_provider_edge( _observed_backend(mode_raw, bundle_dir), mounts=mounts, bind_host=bind_host, advertise_host=advertise_host, diff --git a/tests/e2e/provider_edge_cli.py b/tests/e2e/provider_edge_cli.py new file mode 100644 index 00000000000..7839b400040 --- /dev/null +++ b/tests/e2e/provider_edge_cli.py @@ -0,0 +1,156 @@ +from __future__ import annotations + +import argparse +import os +import signal +import threading +from datetime import datetime, timezone +from pathlib import Path +from types import FrameType +from typing import Final, Literal, cast + +from botocore.config import Config +from botocore.session import get_session +from capture_policy import RUN_BYTES, SCENARIO_BYTES, RequestBudget, ScenarioIdentity, canonical_scenario_id +from capture_session import CaptureResult +from capture_store import DynamoCaptureStore, DynamoClient +from fixture_bundle import BundleRecorder, FreshBundle, UnreadableBundle, check_freshness, load_bundle, prepare_bundle +from provider_edge import start_provider_edge +from provider_edge_control import ControlServer, EdgeController +from pydantic import BaseModel, ConfigDict, Field, TypeAdapter + + +class CredentialHeader(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + variable: str + prefix: str = "" + + +class EdgeConfiguration(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + mode: Literal["record", "replay"] + owner: str = Field(min_length=1) + bundle_dir: Path + outcome_file: Path + identities: tuple[ScenarioIdentity, ...] = Field(min_length=1, max_length=RUN_BYTES // SCENARIO_BYTES) + mounts: dict[str, str] + advertise_host: str + data_port: int = Field(default=8080, gt=0, le=65535) + control_port: int = Field(default=8081, gt=0, le=65535) + lease_deadline: int + attempt_cap: int = Field(default=12, gt=0, le=12) + table: str | None = None + region: str | None = None + credential_env: dict[str, dict[str, CredentialHeader]] = {} + request_budgets: dict[str, RequestBudget] = {} + + +def _write_outcomes(path: Path, results: tuple[CaptureResult, ...]) -> None: + temporary: Final = path.with_suffix(path.suffix + ".tmp") + temporary.write_bytes(TypeAdapter(tuple[CaptureResult, ...]).dump_json(results)) + os.replace(temporary, path) + + +def configured_controller(config: EdgeConfiguration) -> EdgeController: + identities: Final = {canonical_scenario_id(identity.node): identity for identity in config.identities} + if len(identities) != len(config.identities): + raise ValueError("duplicate configured scenario identity") + if config.mode == "replay": + if config.credential_env or config.table: + raise ValueError("replay cannot receive capture credentials or a store") + if not isinstance( + check_freshness(config.bundle_dir, now=datetime.now(timezone.utc), profile="stateless_v1"), FreshBundle + ): + raise ValueError("replay bundle is outside its age limit") + bundle: Final = load_bundle(config.bundle_dir, profile="stateless_v1") + if isinstance(bundle, UnreadableBundle): + raise ValueError("verified replay bundle is unavailable") + return EdgeController(identities, config.owner, config.lease_deadline, replay_bundle=bundle) + if not config.table or not config.region: + raise ValueError("capture requires a configured attempt store") + if set(config.request_budgets) != set(identities): + raise ValueError("capture requires a model and request budget for every enrolled scenario") + if config.bundle_dir.exists(): + raise ValueError("capture requires a new owned bundle directory") + store: Final = DynamoCaptureStore( + cast( + DynamoClient, + get_session().create_client( + "dynamodb", + region_name=config.region, + config=Config( + retries={"total_max_attempts": 1}, + connect_timeout=5, + read_timeout=5, + ), + ), + ), + config.table, + ) + run_error: Final = store.create_run(owner=config.owner, cap=config.attempt_cap, expires_at=config.lease_deadline) + if run_error is not None: + raise ValueError(run_error.reason) + recorder: Final = prepare_bundle(config.bundle_dir, profile="stateless_v1") + assert isinstance(recorder, BundleRecorder) + return EdgeController( + identities, + config.owner, + config.lease_deadline, + store=store, + recorder=recorder, + outcome_sink=lambda results: _write_outcomes(config.outcome_file, results), + request_budgets=config.request_budgets, + ) + + +def serve(config: EdgeConfiguration) -> None: + controller: Final = configured_controller(config) + upstream_headers: Final = { + mount: { + header.lower(): credential.prefix + os.environ[credential.variable] + for header, credential in headers.items() + } + for mount, headers in config.credential_env.items() + } + stopped: Final = threading.Event() + + def stop(signum: int, frame: FrameType | None) -> None: + stopped.set() + + signal.signal(signal.SIGTERM, stop) + signal.signal(signal.SIGINT, stop) + control: Final = ControlServer(controller, port=config.control_port) + running: Final = start_provider_edge( + controller.backend(upstream_headers), + mounts=config.mounts, + bind_host="0.0.0.0", + advertise_host=config.advertise_host, + bind_port=config.data_port, + guard=controller, + observation=controller, + ) + thread: Final = threading.Thread(target=control.serve_forever, daemon=True) + thread.start() + try: + stopped.wait() + finally: + control.shutdown() + control.server_close() + running.shutdown() + thread.join(timeout=5) + + +class CommandArguments(argparse.Namespace): + config: Path + + +def main() -> None: + parser: Final = argparse.ArgumentParser(description="Serve an explicitly configured trusted provider edge") + parser.add_argument("--config", required=True, type=Path) + arguments: Final = CommandArguments() + parser.parse_args(namespace=arguments) + serve(EdgeConfiguration.model_validate_json(arguments.config.read_bytes())) + + +if __name__ == "__main__": + main() diff --git a/tests/e2e/provider_edge_control.py b/tests/e2e/provider_edge_control.py new file mode 100644 index 00000000000..2d1810466ab --- /dev/null +++ b/tests/e2e/provider_edge_control.py @@ -0,0 +1,259 @@ +from __future__ import annotations + +import threading +import uuid +from collections.abc import Callable, Mapping +from dataclasses import dataclass, field, replace +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from typing import Final, Literal + +from capture_policy import SCENARIO_BYTES, RequestBudget, ScenarioIdentity, ScenarioOutcome, canonical_scenario_id +from capture_session import CaptureResult, CaptureSession +from capture_store import CaptureLeaseStore +from fixture_bundle import BundleRecorder, Interaction, LoadedBundle, RecordedRequest, RecordedResponse +from provider_edge import EdgeBackend, ProviderRequestObservation, RecordEdge, ReplayEdge, ReplaySource +from pydantic import BaseModel, ConfigDict, Field, ValidationError + + +class ControlRequest(BaseModel): + model_config = ConfigDict(extra="forbid") + action: Literal["begin", "phase", "observe", "count", "status"] + node: str = "" + phase: Literal["setup", "call", "teardown"] | None = None + passed: bool = False + marker: str = Field(default="", max_length=256) + observation_id: str = "" + + +class ControlReply(BaseModel): + ok: bool + error: str | None = None + observation_id: str | None = None + count: int | None = None + mode: Literal["record", "replay"] | None = None + + +@dataclass(slots=True) +class EdgeController: + identities: Mapping[str, ScenarioIdentity] + owner: str + lease_deadline: int + store: CaptureLeaseStore | None = None + recorder: BundleRecorder | None = None + replay_bundle: LoadedBundle | None = None + outcome_sink: Callable[[tuple[CaptureResult, ...]], None] | None = None + request_budgets: Mapping[str, RequestBudget] = field(default_factory=lambda: dict[str, RequestBudget]()) + _replay: ReplaySource | None = field(default=None, init=False) + _active: str | None = field(default=None, init=False) + _session: CaptureSession | None = field(default=None, init=False) + _phases: tuple[tuple[str, bool], ...] = field(default=(), init=False) + _completed: frozenset[str] = field(default_factory=frozenset, init=False) + _results: tuple[CaptureResult, ...] = field(default=(), init=False) + _observations: tuple[tuple[str, ProviderRequestObservation], ...] = field(default=(), init=False) + _in_flight: int = field(default=0, init=False) + _recorded_bytes: int = field(default=0, init=False) + _halted: bool = field(default=False, init=False) + _lock: threading.RLock = field(default_factory=threading.RLock, init=False) + + def __post_init__(self) -> None: + if (self.store is None) == (self.replay_bundle is None) or (self.store is not None and self.recorder is None): + raise ValueError("controller requires exactly one capture store or replay source") + if not self.identities or any( + key != canonical_scenario_id(identity.node) for key, identity in self.identities.items() + ): + raise ValueError("controller requires canonical enrolled identities") + if self.replay_bundle is not None: + self._replay = ReplaySource(self.replay_bundle, test_key=self.test_key) + + def backend(self, upstream_headers: Mapping[str, Mapping[str, str]]) -> EdgeBackend: + if self._replay is not None: + return ReplayEdge(self._replay) + assert self.recorder is not None + return RecordEdge( + self.recorder, + threading.Lock(), + test_key=self.test_key, + before_attempt=self.before_attempt, + response_finished=self.response_finished, + response_byte_limit=self.remaining_bytes, + recording_error=self.recording_error, + upstream_headers=upstream_headers, + request_error=self.request_error, + ) + + def request_error(self, body: bytes | None) -> str | None: + budget: Final = self.request_budgets.get(self.test_key()) + return budget.error(body) if budget is not None else None + + def remaining_bytes(self) -> int: + with self._lock: + return max(0, SCENARIO_BYTES - self._recorded_bytes) + + def recording_error(self, request: RecordedRequest, response: RecordedResponse) -> str | None: + with self._lock: + size: Final = len(Interaction(request=request, response=response).model_dump_json(indent=2).encode()) + if self._recorded_bytes + size > SCENARIO_BYTES: + return "scenario exceeds capture byte limit" + self._recorded_bytes += size + return None + + def test_key(self) -> str: + with self._lock: + if self._active is None: + raise RuntimeError("no active scenario") + return self._active + + def begin_request(self) -> str | None: + with self._lock: + if self._active is None or self._halted: + return "no active trusted scenario" + self._in_flight += 1 + return None + + def end_request(self) -> None: + with self._lock: + self._in_flight -= 1 + + def before_attempt(self) -> str | None: + with self._lock: + return self._session.before_attempt() if self._session is not None else "no active capture session" + + def response_finished(self, response: RecordedResponse) -> None: + with self._lock: + if self._session is not None: + self._session.response_finished(response) + + def observe(self, body: bytes | None) -> None: + with self._lock: + for _, observation in self._observations: + observation.observe(body) + + @property + def results(self) -> tuple[CaptureResult, ...]: + with self._lock: + return self._results + + def command(self, request: ControlRequest) -> ControlReply: + with self._lock: + if request.action == "status": + return ControlReply( + ok=not self._halted, + count=len(self._completed), + mode="record" if self.store is not None else "replay", + ) + if request.action == "begin": + return self._begin(request.node) + if self._active is None or canonical_scenario_id(request.node) != self._active: + return ControlReply(ok=False, error="scenario is not active") + if request.action == "phase": + return self._phase(request) + if request.action == "observe": + if not request.marker: + return ControlReply(ok=False, error="observation marker is empty") + identifier: Final = uuid.uuid4().hex + self._observations = (*self._observations, (identifier, ProviderRequestObservation(request.marker))) + return ControlReply(ok=True, observation_id=identifier) + observation: Final = next((obs for key, obs in self._observations if key == request.observation_id), None) + return ( + ControlReply(ok=True, count=observation.count) + if observation is not None + else ControlReply(ok=False, error="unknown observation") + ) + + def _begin(self, raw_node: str) -> ControlReply: + node: Final = canonical_scenario_id(raw_node) + if self._active is not None or self._halted or node in self._completed: + return ControlReply(ok=False, error="scenario already started or run halted") + identity: Final = self.identities.get(node) + if identity is None: + return ControlReply(ok=False, error="scenario is not enrolled") + if self.store is not None: + failure: Final = self.store.acquire( + scenario_key=identity.key, owner=self.owner, expires_at=self.lease_deadline + ) + if failure is not None: + self._halted = True + return ControlReply(ok=False, error=failure.reason) + self._session = CaptureSession(identity, self.owner, self.store) + self._active = node + self._phases = () + self._observations = () + self._recorded_bytes = 0 + return ControlReply(ok=True) + + def _phase(self, request: ControlRequest) -> ControlReply: + if request.phase is None or request.phase in dict(self._phases): + self._halted = True + return ControlReply(ok=False, error="invalid or duplicate scenario phase") + if request.phase != "setup" and "setup" not in dict(self._phases): + self._halted = True + return ControlReply(ok=False, error="scenario setup outcome missing") + self._phases = (*self._phases, (request.phase, request.passed)) + if not request.passed: + self._halted = True + if request.phase != "teardown": + return ControlReply(ok=True) + outcome: Final = ScenarioOutcome( + **{phase: dict(self._phases).get(phase, False) for phase in ("setup", "call", "teardown")} + ) + error: Final = self._finish(outcome) + self._completed = self._completed | {self.test_key()} + self._active = None + self._halted = self._halted or error is not None + if self.outcome_sink is not None: + self.outcome_sink(self._results) + return ControlReply(ok=error is None, error=error) + + def _finish(self, outcome: ScenarioOutcome) -> str | None: + if self._in_flight: + return "scenario ended with in-flight provider requests" + if self._session is not None: + result: Final = self._session.finish(outcome) + assert self.store is not None + release: Final = self.store.release(scenario_key=result.identity.key, owner=self.owner) + error: Final = ( + result.error + or (release.reason if release is not None else None) + or ("scenario lifecycle was rejected" if self._halted else None) + ) + self._results = (*self._results, replace(result, error=error)) + return error + assert self._replay is not None + if not (outcome.setup and outcome.call and outcome.teardown): + return "trusted scenario setup, call and teardown must all pass" + if self._halted: + return "scenario lifecycle was rejected" + return self._replay.leftover_error(self.test_key()) + + +class _ControlHandler(BaseHTTPRequestHandler): + def log_message(self, format: str, *args: object) -> None: + pass + + def do_POST(self) -> None: + server: Final = self.server + assert isinstance(server, ControlServer) + try: + length: Final = int(self.headers.get("content-length", "0")) + if self.path != "/control" or not 0 < length <= 4096: + self.send_error(400) + return + request: Final = ControlRequest.model_validate_json(self.rfile.read(length)) + reply: Final = server.controller.command(request) + except (ValueError, ValidationError): + self.send_error(400) + return + encoded: Final = reply.model_dump_json().encode() + self.send_response(200 if reply.ok else 409) + self.send_header("content-type", "application/json") + self.send_header("content-length", str(len(encoded))) + self.end_headers() + self.wfile.write(encoded) + + +class ControlServer(ThreadingHTTPServer): + daemon_threads = True + + def __init__(self, controller: EdgeController, *, port: int = 0) -> None: + self.controller: Final = controller + super().__init__(("127.0.0.1", port), _ControlHandler) diff --git a/tests/e2e/provider_edge_remote.py b/tests/e2e/provider_edge_remote.py new file mode 100644 index 00000000000..1f284406b06 --- /dev/null +++ b/tests/e2e/provider_edge_remote.py @@ -0,0 +1,85 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Final, Literal +from urllib.parse import urlsplit + +from capture_policy import canonical_scenario_id +from e2e_http import NetworkError, forward +from fixture_mode import current_test_key +from provider_edge import ProviderEdge +from provider_edge_control import ControlReply, ControlRequest + + +@dataclass(frozen=True, slots=True) +class RemoteEdge: + control_url: str + data_url: str + + def __post_init__(self) -> None: + control: Final = urlsplit(self.control_url) + data: Final = urlsplit(self.data_url) + if ( + control.scheme != "http" + or control.hostname != "127.0.0.1" + or control.path not in ("", "/") + or control.username + or control.query + or control.fragment + ): + raise ValueError("edge management URL must be loopback HTTP") + if ( + data.scheme != "http" + or not data.hostname + or data.port is None + or data.path not in ("", "/") + or data.username + or data.query + or data.fragment + ): + raise ValueError("edge data URL must name an HTTP host and port") + + @property + def edge(self) -> ProviderEdge: + parsed: Final = urlsplit(self.data_url) + assert parsed.hostname is not None and parsed.port is not None + return ProviderEdge(parsed.port, parsed.hostname) + + def command(self, request: ControlRequest) -> ControlReply: + response: Final = forward( + "POST", + self.control_url.rstrip("/") + "/control", + headers={"content-type": "application/json"}, + body=request.model_dump_json().encode(), + timeout=10, + ) + if isinstance(response, NetworkError): + raise RuntimeError("trusted edge control unavailable") + reply: Final = ControlReply.model_validate_json(response.body) + if response.status_code != 200 or not reply.ok: + raise RuntimeError(reply.error or "trusted edge control rejected operation") + return reply + + def begin(self, node: str) -> None: + self.command(ControlRequest(action="begin", node=canonical_scenario_id(node))) + + def phase(self, node: str, phase: Literal["setup", "call", "teardown"], passed: bool) -> None: + self.command(ControlRequest(action="phase", node=canonical_scenario_id(node), phase=phase, passed=passed)) + + def observe(self, marker: str) -> str: + reply: Final = self.command( + ControlRequest(action="observe", node=canonical_scenario_id(current_test_key()), marker=marker) + ) + if reply.observation_id is None: + raise RuntimeError("trusted edge did not register observation") + return reply.observation_id + + def count(self, observation_id: str) -> int: + reply: Final = self.command( + ControlRequest( + action="count", node=canonical_scenario_id(current_test_key()), observation_id=observation_id + ) + ) + if reply.count is None: + raise RuntimeError("trusted edge did not return observed count") + return reply.count diff --git a/tests/e2e/router/test_reliability_cache_e2e.py b/tests/e2e/router/test_reliability_cache_e2e.py index f7a2f2ffeb7..2ebf0cc68c5 100644 --- a/tests/e2e/router/test_reliability_cache_e2e.py +++ b/tests/e2e/router/test_reliability_cache_e2e.py @@ -16,6 +16,7 @@ from e2e_config import ( PROVIDER_EDGE_ADVERTISE_HOST, PROVIDER_EDGE_BIND_HOST, REQUEST_TIMEOUT, + REMOTE_EDGE, unique_marker, ) from lifecycle import ResourceManager @@ -53,6 +54,7 @@ class TestReliabilityCache: bind_host=PROVIDER_EDGE_BIND_HOST, advertise_host=PROVIDER_EDGE_ADVERTISE_HOST, forward_timeout=REQUEST_TIMEOUT, + remote=REMOTE_EDGE, ) as edge: model_id: Final = client.proxy.create_model( model,