fix(e2e): record a streamed chunk only after its downstream write lands

A downstream disconnect mid-relay was recording the chunk whose write never
landed, so replay would hand back a byte the record run never delivered. Append
each chunk after its yield returns, and label the truncation from the generator
close, so the recording holds exactly what the proxy received.
This commit is contained in:
mateo-berri 2026-08-24 13:26:58 -07:00
parent ec47bbaaaa
commit b2f7216a6b
2 changed files with 56 additions and 15 deletions

View file

@ -628,33 +628,38 @@ def _persist(
def _recording_steps(
backend: RecordEdge, test_key: str, request: RecordedRequest, head: StreamHead
) -> Generator[StreamStep, None, None]:
"""Record mode's step source: hand each upstream chunk downstream as it arrives
while collecting it, then persist the whole sequence once, under the lock.
Relaying incrementally keeps record exercising the proxy's incremental parser
the way a live run does.
"""Record mode's step source: hand each upstream chunk downstream and record it
only once that write has returned, then persist the whole sequence once, under
the lock. Relaying incrementally keeps record exercising the proxy's incremental
parser the way a live run does.
The ``finally`` also covers the proxy hanging up mid-stream, which closes this
generator: what arrived is still recorded, marked truncated, because recording a
cut-short stream as a clean one would let a later replay serve a well-terminated
fraction of the response and pass a test that should have gone red."""
A chunk is appended after its ``yield`` returns, so a downstream that hangs up
mid-relay records exactly the chunks it took and never the one whose write
raised. The ``except`` covers that downstream close and the proxy hanging up
mid-stream; either way the ``finally`` persists what arrived, marked truncated,
because recording a cut-short stream as a clean one would let a later replay
serve a well-terminated fraction of the response and pass a test that should
have gone red."""
collected: list[bytes] = []
truncated: str | None = None
delivered = False
try:
with closing(head.steps) as steps:
for step in steps:
match step:
case StreamChunk(data=data):
collected.append(data)
case StreamChunk():
pass
case StreamTruncation(reason=reason):
truncated = f"upstream: {reason}"
case _:
assert_never(step)
yield step
delivered = True
finally:
if not delivered and truncated is None:
if isinstance(step, StreamChunk):
collected.append(step.data)
except GeneratorExit:
if truncated is None:
truncated = f"downstream: relay closed after {len(collected)} chunks"
raise
finally:
_persist(
backend,
test_key,

View file

@ -36,7 +36,7 @@ from typing import Final
import pytest
from pydantic import TypeAdapter
from e2e_http import RawResponse, forward
from e2e_http import RawResponse, StreamChunk, forward
from fixture_canonical import canonicalize
from fixture_bundle import (
BundleRecorder,
@ -54,6 +54,7 @@ from provider_edge import (
REPLAY_MISS_STATUS,
EdgeBackend,
EdgeReply,
EdgeStream,
ProviderEdge,
RecordEdge,
ReplayEdge,
@ -1077,6 +1078,41 @@ class TestStreamingFidelity:
assert recorded.truncated is not None
assert recorded.truncated.startswith("upstream: ")
def test_a_downstream_disconnect_mid_relay_records_only_the_delivered_chunks(
self, tmp_path: Path
) -> None:
"""The provider keeps sending, but the proxy the edge relays to hangs up after
two chunks. The chunk whose downstream write never landed must stay out of the
recording, or replay would hand back a byte the record run never delivered.
Driven through the pure ``handle_edge_request`` core because a socket client
cannot force these tiny chunks to block mid-write, so closing the relay
generator is the faithful stand-in for the downstream write raising: it lands
the generator on the same suspended yield a broken pipe would."""
root = tmp_path / "bundle"
with chunked_provider() as provider:
outcome = handle_edge_request(
record_backend(root),
{"openai": provider_url(provider)},
"POST",
STREAM_PATH,
{"content-type": "application/json"},
STREAM_BODY,
timeout=10.0,
)
assert isinstance(outcome, EdgeStream)
steps = outcome.steps
first = next(steps)
second = next(steps)
assert isinstance(first, StreamChunk) and isinstance(second, StreamChunk)
assert (first.data, second.data) == (SSE_CHUNKS[0], SSE_CHUNKS[1])
steps.close()
recorded = recorded_stream(root)
assert recorded.status_code == 200
assert stream_chunks(recorded) == [SSE_CHUNKS[0]]
assert recorded.truncated == "downstream: relay closed after 1 chunks"
def test_a_truncated_recording_replays_as_a_truncated_stream(self, tmp_path: Path) -> None:
root = tmp_path / "bundle"
record_stream(root, abort_after=2)