feat(e2e): add bounded trusted provider capture and snapshot storage

This commit is contained in:
Yuneng Jiang 2026-09-15 00:00:53 -07:00
parent 4805c6d51f
commit 789b594370
No known key found for this signature in database
20 changed files with 2578 additions and 26 deletions

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

@ -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 <path>` 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.

207
tests/e2e/capture_policy.py Normal file
View file

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

View file

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

View file

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

View file

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

202
tests/e2e/capture_store.py Normal file
View file

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

View file

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

View file

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

View file

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

View file

@ -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 = "<unparsed-multipart>"
_BOUNDARY_PLACEHOLDER: Final = b"--<boundary>"
_BINARY_FIELD_PREFIX: Final = "<binary:sha256:"
@ -441,6 +440,7 @@ class ReplaySource:
or poll loop replays its recorded responses in recorded order)."""
bundle: LoadedBundle
test_key: Callable[[], str] = current_test_key
_pools: dict[str, dict[str, deque[Interaction]]] = field(init=False)
def __post_init__(self) -> 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,

View file

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

View file

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

View file

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

View file

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