fix(e2e): reject incomplete captures and drain teardown safely

tests/code_coverage_tests/test_provider_capture.py

@@ -2,7 +2,11 @@ from __future__ import annotations
 
 import base64
 import hashlib
+import json
 import threading
+from collections.abc import Generator
+from contextlib import contextmanager
+from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
 from dataclasses import dataclass, field
 from datetime import datetime, timedelta, timezone
 from pathlib import Path
@@ -32,7 +36,7 @@ from fixture_bundle import (
 )
 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
+from test_provider_edge import CHAT_PATH, call_edge, provider_url, running_edge
 
 
 @dataclass
@@ -59,8 +63,61 @@ class ScriptedReservations:
         return None
 
 
+def completion_payload(text: str = "blue") -> bytes:
+    return json.dumps(
+        {
+            "id": "chatcmpl-synthetic",
+            "object": "chat.completion",
+            "model": "synthetic",
+            "choices": [{"index": 0, "message": {"role": "assistant", "content": text}, "finish_reason": "stop"}],
+            "usage": {"prompt_tokens": 2, "completion_tokens": 1, "total_tokens": 3},
+        }
+    ).encode()
+
+
 def successful_response() -> RecordedHttpResponse:
-    return RecordedHttpResponse(status_code=200, headers={}, body_b64=base64.b64encode(b'{"answer":"blue"}').decode())
+    return RecordedHttpResponse(status_code=200, headers={}, body_b64=base64.b64encode(completion_payload()).decode())
+
+
+class CompletionProvider(ThreadingHTTPServer):
+    daemon_threads = True
+
+    def __init__(self) -> None:
+        super().__init__(("127.0.0.1", 0), CompletionHandler)
+        self.hits: list[str] = []
+
+
+class CompletionHandler(BaseHTTPRequestHandler):
+    def do_GET(self) -> None:
+        self.send_response(200)
+        self.send_header("content-length", "0")
+        self.end_headers()
+
+    def do_POST(self) -> None:
+        assert isinstance(self.server, CompletionProvider)
+        self.rfile.read(int(self.headers.get("content-length", "0")))
+        self.server.hits.append(self.path)
+        payload = completion_payload()
+        self.send_response(200)
+        self.send_header("content-type", "application/json")
+        self.send_header("content-length", str(len(payload)))
+        self.end_headers()
+        self.wfile.write(payload)
+
+    def log_message(self, format: str, *args: object) -> None:
+        pass
+
+
+@contextmanager
+def fake_provider() -> Generator[CompletionProvider]:
+    with CompletionProvider() as server:
+        thread = threading.Thread(target=server.serve_forever, daemon=True)
+        thread.start()
+        try:
+            yield server
+        finally:
+            server.shutdown()
+            thread.join()
 
 
 class TestBoundedCapture:
@@ -143,7 +200,7 @@ class TestScenarioIdentity:
 
     @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):
+        with pytest.raises(ValueError, match="scenario must name an E2E test node"):
             canonical_scenario_id(node)
 
     def test_contract_and_profile_change_identity_but_candidate_revision_does_not_key_it(self) -> None:
@@ -350,7 +407,7 @@ def test_persist_must_reach_disk_before_success(tmp_path: Path, disk_failure: bo
     request: Final = RecordedRequest(method="post", path="/openai/v1/chat/completions", headers={})
     assert session.before_attempt() is None
     if disk_failure:
-        with pytest.raises(OSError):
+        with pytest.raises(OSError, match=r"Not a directory|File exists"):
             _persist(backend, session.identity.node, request, successful_response())
     else:
         _persist(backend, session.identity.node, request, successful_response())
@@ -358,3 +415,61 @@ def test_persist_must_reach_disk_before_success(tmp_path: Path, disk_failure: bo
     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)
+
+
+@pytest.mark.parametrize(
+    "fault",
+    [
+        "empty_choices",
+        "missing_finish",
+        "partial_finish",
+        "missing_message",
+        "missing_usage",
+        "empty_usage",
+        "missing_identity",
+    ],
+)
+def test_incomplete_nonstream_completion_cannot_publish(fault: str) -> None:
+    payload = json.loads(completion_payload())
+    if fault == "empty_choices":
+        payload["choices"] = []
+    elif fault == "missing_finish":
+        del payload["choices"][0]["finish_reason"]
+    elif fault == "partial_finish":
+        payload["choices"][0]["finish_reason"] = "length"
+    elif fault == "missing_message":
+        del payload["choices"][0]["message"]
+    elif fault == "missing_usage":
+        del payload["usage"]
+    elif fault == "empty_usage":
+        payload["usage"] = {}
+    else:
+        del payload["id"]
+    response = RecordedHttpResponse(
+        status_code=200, headers={}, body_b64=base64.b64encode(json.dumps(payload).encode()).decode()
+    )
+    assert publication_error(ScenarioOutcome(True, True, True), (response,)) is not None
+    assert publication_error(ScenarioOutcome(True, True, True), (successful_response(),)) is None
+
+
+@pytest.mark.parametrize("fault", ["empty_content", "missing_stop", "partial_usage", "healthy"])
+def test_anthropic_nonstream_requires_finished_content_and_usage(fault: str) -> None:
+    payload = {
+        "id": "msg-synthetic",
+        "model": "synthetic",
+        "type": "message",
+        "role": "assistant",
+        "content": [{"type": "text", "text": "blue"}],
+        "stop_reason": "end_turn",
+        "usage": {"input_tokens": 2, "output_tokens": 1},
+    }
+    if fault == "empty_content":
+        payload["content"] = []
+    elif fault == "missing_stop":
+        del payload["stop_reason"]
+    elif fault == "partial_usage":
+        payload["usage"] = {"input_tokens": 2}
+    response = RecordedHttpResponse(
+        status_code=200, headers={}, body_b64=base64.b64encode(json.dumps(payload).encode()).decode()
+    )
+    assert (publication_error(ScenarioOutcome(True, True, True), (response,)) is None) is (fault == "healthy")

tests/code_coverage_tests/test_provider_capture_publication.py

@@ -4,6 +4,7 @@ import base64
 import hashlib
 from datetime import datetime, timezone
 from io import BytesIO
+from pathlib import Path
 from typing import Final, cast
 
 import pytest
@@ -13,9 +14,10 @@ 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 capture_snapshot import CaptureProvenance, ScenarioSnapshot, SnapshotFailure, materialize_snapshot
+from fixture_bundle import Interaction, Manifest, RecordedRequest
 from fixture_profile import StrictIdentity, strict_identity
+from test_provider_capture import successful_response
 
 NOW: Final = datetime(2031, 4, 5, tzinfo=timezone.utc)
 IDENTITY: Final = ScenarioIdentity("test_example.py::test_one", "a" * 64, "synthetic")
@@ -43,9 +45,7 @@ def snapshot_content() -> bytes:
                     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()
-                    ),
+                    response=successful_response(),
                 ),
             ),
             provenance=CaptureProvenance(
@@ -166,3 +166,20 @@ class TestConditionalPublication:
             assert isinstance(result, SnapshotFailure)
             s3.assert_no_pending_responses()
             dynamo.assert_no_pending_responses()
+
+
+def test_materialization_write_failure_removes_partial_bundle_and_can_retry(tmp_path: Path) -> None:
+    snapshot = ScenarioSnapshot.model_validate_json(snapshot_content())
+    first = snapshot.interactions[0]
+    # Force a real filesystem failure after the manifest and first interaction were written.
+    bad_request = first.request.model_copy(update={"method": "x" * 300})
+    broken = snapshot.model_copy(update={"interactions": (first, first.model_copy(update={"request": bad_request}))})
+    destination = tmp_path / "snapshot"
+    with pytest.raises(OSError, match="File name too long"):
+        materialize_snapshot(broken, destination)
+    assert not destination.exists()
+    materialize_snapshot(snapshot, destination)
+    assert len(list(destination.rglob("*.json"))) == 2
+    with pytest.raises(FileExistsError):
+        materialize_snapshot(snapshot, destination)
+    assert len(list(destination.rglob("*.json"))) == 2

tests/code_coverage_tests/test_provider_edge_control.py

@@ -13,6 +13,7 @@ from typing import Final
 
 import pytest
 from capture_policy import SCENARIO_BYTES, ScenarioIdentity, canonical_scenario_id
+from capture_session import CaptureResult
 from capture_store import StoreFailure
 from fixture_bundle import (
     BundleRecorder,
@@ -25,8 +26,8 @@ from fixture_bundle import (
 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
+from test_provider_capture import ScriptedReservations, completion_payload, fake_provider, successful_response
+from test_provider_edge import CHAT_PATH, call_edge, provider_url
 
 
 @contextmanager
@@ -243,7 +244,7 @@ def test_aggregate_recording_limit_stops_before_second_disk_write(tmp_path: Path
     response: Final = RecordedHttpResponse(
         status_code=200,
         headers={},
-        body_b64=base64.b64encode(b'{"answer":"' + b"x" * (3 * 1024 * 1024) + b'"}').decode(),
+        body_b64=base64.b64encode(completion_payload("x" * (3 * 1024 * 1024))).decode(),
     )
     assert controller.command(ControlRequest(action="begin", node=node)).ok
     assert controller.command(ControlRequest(action="phase", node=node, phase="setup", passed=True)).ok
@@ -257,3 +258,42 @@ def test_aggregate_recording_limit_stops_before_second_disk_write(tmp_path: Path
     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
+
+
+def test_inflight_teardown_drains_before_releasing_lease(tmp_path: Path) -> None:
+    events: list[str] = []
+
+    class Store(ScriptedReservations):
+        def complete(self, *, scenario_key: str, owner: str, attempt_id: str, successful: bool) -> str | None:
+            events.append("settled")
+            return None
+
+        def release(self, *, scenario_key: str, owner: str) -> StoreFailure | None:
+            events.append("released")
+            return None
+
+    identity = ScenarioIdentity("test_example.py::test_one", "a" * 64, "synthetic")
+    node = canonical_scenario_id(identity.node)
+    recorder = prepare_bundle(tmp_path / "capture", profile="stateless_v1")
+    assert isinstance(recorder, BundleRecorder)
+    saved: list[tuple[CaptureResult, ...]] = []
+    controller = 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.begin_request() is None
+    assert controller.before_attempt() is None
+    assert controller.command(ControlRequest(action="phase", node=node, phase="call", passed=True)).ok
+    ended = controller.command(ControlRequest(action="phase", node=node, phase="teardown", passed=True))
+    assert not ended.ok and ended.error == "scenario ended with in-flight provider requests"
+    assert events == [] and saved == []
+    assert controller.begin_request() == "no active trusted scenario"
+    assert controller.test_key() == node
+    controller.response_finished(successful_response())
+    controller.end_request()
+    assert events == ["settled", "released"]
+    assert len(saved) == 1 and not saved[0][0].publishable
+    assert saved[0][0].error == "scenario lifecycle was rejected"
+    assert controller.command(ControlRequest(action="status")).count == 1
+    assert not controller.command(ControlRequest(action="begin", node=node)).ok

tests/e2e/capture_policy.py

@@ -124,7 +124,53 @@ def _json_error(body: bytes, *, streaming: bool = False) -> str | None:
         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
+    usage_error: Final = _usage_error(value["usage"]) if "usage" in value else None
+    if usage_error is not None or streaming:
+        return usage_error
+    return _completion_error(value)
+
+
+def _completion_error(value: dict[str, JsonValue]) -> str | None:
+    if not all(isinstance(value.get(key), str) and value[key] for key in ("id", "model")):
+        return "completion lacks provider identity"
+    usage: Final = value.get("usage")
+    if not isinstance(usage, dict):
+        return "completion lacks token usage"
+    if value.get("object") == "chat.completion":
+        choices: Final = value.get("choices")
+        if not isinstance(choices, list) or len(choices) != 1 or not isinstance(choices[0], dict):
+            return "completion requires exactly one choice"
+        choice: Final = choices[0]
+        message: Final = choice.get("message")
+        if (
+            type(choice.get("index")) is not int
+            or choice["index"] != 0
+            or choice.get("finish_reason") != "stop"
+            or not isinstance(message, dict)
+            or message.get("role") != "assistant"
+            or not isinstance(message.get("content"), str)
+        ):
+            return "completion lacks a finished assistant message"
+        if not all(key in usage for key in ("prompt_tokens", "completion_tokens", "total_tokens")):
+            return "completion lacks token usage"
+        return None
+    if value.get("type") == "message":
+        content: Final = value.get("content")
+        if (
+            value.get("role") != "assistant"
+            or value.get("stop_reason") not in ("end_turn", "stop_sequence")
+            or not isinstance(content, list)
+            or not content
+            or any(
+                not isinstance(block, dict) or block.get("type") != "text" or not isinstance(block.get("text"), str)
+                for block in content
+            )
+        ):
+            return "completion lacks finished message content"
+        if not all(key in usage for key in ("input_tokens", "output_tokens")):
+            return "completion lacks token usage"
+        return None
+    return "unsupported completion response"
 
 
 def _stream_error(response: RecordedStreamedResponse) -> str | None:

tests/e2e/capture_snapshot.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import hashlib
+import shutil
 from dataclasses import dataclass
 from datetime import datetime
 from pathlib import Path
@@ -121,6 +122,14 @@ def refresh_due(snapshot: ScenarioSnapshot, *, now: datetime) -> bool:
 
 def materialize_snapshot(snapshot: ScenarioSnapshot, destination: Path) -> None:
     destination.mkdir(parents=True, exist_ok=False)
+    try:
+        _write_snapshot(snapshot, destination)
+    except BaseException:
+        shutil.rmtree(destination)
+        raise
+
+
+def _write_snapshot(snapshot: ScenarioSnapshot, destination: Path) -> None:
     (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()

tests/e2e/provider_edge_control.py

@@ -53,6 +53,7 @@ class EdgeController:
     _in_flight: int = field(default=0, init=False)
     _recorded_bytes: int = field(default=0, init=False)
     _halted: bool = field(default=False, init=False)
+    _pending_outcome: ScenarioOutcome | None = field(default=None, init=False)
     _lock: threading.RLock = field(default_factory=threading.RLock, init=False)
 
     def __post_init__(self) -> None:
@@ -113,6 +114,10 @@ class EdgeController:
     def end_request(self) -> None:
         with self._lock:
             self._in_flight -= 1
+            if self._in_flight == 0 and self._pending_outcome is not None:
+                outcome = self._pending_outcome
+                self._pending_outcome = None
+                self._complete(outcome)
 
     def before_attempt(self) -> str | None:
         with self._lock:
@@ -196,9 +201,17 @@ class EdgeController:
         outcome: Final = ScenarioOutcome(
             **{phase: dict(self._phases).get(phase, False) for phase in ("setup", "call", "teardown")}
         )
+        if self._in_flight:
+            self._halted = True
+            self._pending_outcome = outcome
+            return ControlReply(ok=False, error="scenario ended with in-flight provider requests")
+        return self._complete(outcome)
+
+    def _complete(self, outcome: ScenarioOutcome) -> ControlReply:
         error: Final = self._finish(outcome)
         self._completed = self._completed | {self.test_key()}
         self._active = None
+        self._session = None
         self._halted = self._halted or error is not None
         if self.outcome_sink is not None:
             self.outcome_sink(self._results)