mirror of
https://github.com/BerriAI/litellm.git
synced 2026-08-28 05:25:59 +00:00
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:
parent
ec47bbaaaa
commit
b2f7216a6b
2 changed files with 56 additions and 15 deletions
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue