mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-01 17:00:21 +08:00
fix: 流式响应上游连接异常时优雅降级,避免丢失 usage 和 telemetry
上游 httpx.StreamClosed/HTTPError 发生时,若已向客户端输出数据则 best-effort flush 残留 SSE 并标记 502 结束流,保证 StreamingResponse 背景任务正常执行;若尚未输出则 re-raise 以触发 failover。
This commit is contained in:
@@ -426,6 +426,7 @@ class StreamProcessor:
|
|||||||
try:
|
try:
|
||||||
sse_parser = SSEEventParser()
|
sse_parser = SSEEventParser()
|
||||||
streaming_started = False
|
streaming_started = False
|
||||||
|
yielded_any = False
|
||||||
buffer = b""
|
buffer = b""
|
||||||
# 使用增量解码器处理跨 chunk 的 UTF-8 字符
|
# 使用增量解码器处理跨 chunk 的 UTF-8 字符
|
||||||
decoder = codecs.getincrementaldecoder("utf-8")(errors="replace")
|
decoder = codecs.getincrementaldecoder("utf-8")(errors="replace")
|
||||||
@@ -464,7 +465,8 @@ class StreamProcessor:
|
|||||||
ctx.needs_conversion = False
|
ctx.needs_conversion = False
|
||||||
|
|
||||||
def _mark_stream_started() -> None:
|
def _mark_stream_started() -> None:
|
||||||
nonlocal start_time, streaming_started
|
nonlocal start_time, streaming_started, yielded_any
|
||||||
|
yielded_any = True
|
||||||
# 记录首字时间 (TTFB) - 在 yield 之前记录
|
# 记录首字时间 (TTFB) - 在 yield 之前记录
|
||||||
if start_time is not None:
|
if start_time is not None:
|
||||||
ctx.record_first_byte_time(start_time)
|
ctx.record_first_byte_time(start_time)
|
||||||
@@ -808,6 +810,37 @@ class StreamProcessor:
|
|||||||
|
|
||||||
except GeneratorExit:
|
except GeneratorExit:
|
||||||
raise
|
raise
|
||||||
|
except (httpx.StreamClosed, httpx.HTTPError) as exc:
|
||||||
|
# 连接关闭/协议错误:best-effort flush 残留 SSE,避免丢失尾部 usage。
|
||||||
|
try:
|
||||||
|
if buffer:
|
||||||
|
remaining = decoder.decode(buffer, True)
|
||||||
|
buffer = b""
|
||||||
|
for line in remaining.split("\n"):
|
||||||
|
self._process_line(ctx, sse_parser, line, skip_record=needs_conversion)
|
||||||
|
|
||||||
|
# flush SSE parser 内部累积的未完成事件
|
||||||
|
for event in sse_parser.flush():
|
||||||
|
self.handle_sse_event(
|
||||||
|
ctx,
|
||||||
|
event.get("event"),
|
||||||
|
event.get("data") or "",
|
||||||
|
skip_record=needs_conversion,
|
||||||
|
)
|
||||||
|
except Exception:
|
||||||
|
# best-effort: 不应因 flush 失败影响后续流程
|
||||||
|
pass
|
||||||
|
|
||||||
|
# 若尚未向客户端输出任何数据,抛出异常以触发上层 failover。
|
||||||
|
if not yielded_any:
|
||||||
|
raise
|
||||||
|
|
||||||
|
# 已输出过数据:不要继续抛异常(否则 StreamingResponse 背景任务不会执行,
|
||||||
|
# usage/telemetry 可能无法落库)。标记为上游错误并结束流。
|
||||||
|
if not ctx.has_completion:
|
||||||
|
ctx.status_code = 502
|
||||||
|
ctx.error_message = f"upstream_stream_error:{type(exc).__name__}"
|
||||||
|
return
|
||||||
finally:
|
finally:
|
||||||
if metrics_enabled:
|
if metrics_enabled:
|
||||||
labels = {
|
labels = {
|
||||||
|
|||||||
@@ -1,5 +1,8 @@
|
|||||||
import json
|
import json
|
||||||
from typing import Any
|
from typing import Any, AsyncIterator
|
||||||
|
|
||||||
|
import httpx
|
||||||
|
import pytest
|
||||||
|
|
||||||
from src.api.handlers.base.response_parser import (
|
from src.api.handlers.base.response_parser import (
|
||||||
ParsedChunk,
|
ParsedChunk,
|
||||||
@@ -78,3 +81,49 @@ def test_process_line_handles_openai_usage_chunk_followed_by_done_without_blank_
|
|||||||
assert ctx.input_tokens == 7
|
assert ctx.input_tokens == 7
|
||||||
assert ctx.output_tokens == 3
|
assert ctx.output_tokens == 3
|
||||||
assert ctx.has_completion is True
|
assert ctx.has_completion is True
|
||||||
|
|
||||||
|
|
||||||
|
class _DummyResponseCtx:
|
||||||
|
async def __aexit__(self, exc_type, exc, tb) -> None: # noqa: ANN001
|
||||||
|
return None
|
||||||
|
|
||||||
|
|
||||||
|
class _DummyHTTPClient:
|
||||||
|
async def aclose(self) -> None:
|
||||||
|
return None
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_create_response_stream_flushes_usage_on_remote_protocol_error() -> None:
|
||||||
|
ctx = StreamContext(model="test-model", api_format="openai:chat")
|
||||||
|
ctx.provider_api_format = "openai:chat"
|
||||||
|
processor = StreamProcessor(request_id="test-request", default_parser=DummyParser())
|
||||||
|
|
||||||
|
usage_chunk = {
|
||||||
|
"id": "chatcmpl_test",
|
||||||
|
"object": "chat.completion.chunk",
|
||||||
|
"choices": [],
|
||||||
|
"usage": {"prompt_tokens": 11, "completion_tokens": 4, "total_tokens": 15},
|
||||||
|
}
|
||||||
|
|
||||||
|
async def _iter_bytes_then_remote_protocol_error() -> AsyncIterator[bytes]:
|
||||||
|
yield f"data: {json.dumps(usage_chunk)}\n".encode("utf-8")
|
||||||
|
raise httpx.RemoteProtocolError("boom")
|
||||||
|
|
||||||
|
out = b""
|
||||||
|
async for b in processor.create_response_stream(
|
||||||
|
ctx=ctx,
|
||||||
|
byte_iterator=_iter_bytes_then_remote_protocol_error(),
|
||||||
|
response_ctx=_DummyResponseCtx(),
|
||||||
|
http_client=_DummyHTTPClient(), # type: ignore[arg-type]
|
||||||
|
prefetched_chunks=[],
|
||||||
|
start_time=None,
|
||||||
|
):
|
||||||
|
out += b
|
||||||
|
|
||||||
|
# Stream ends gracefully (no exception), but usage is best-effort captured and request is marked failed.
|
||||||
|
assert b"data:" in out
|
||||||
|
assert ctx.input_tokens == 11
|
||||||
|
assert ctx.output_tokens == 4
|
||||||
|
assert ctx.status_code == 502
|
||||||
|
assert (ctx.error_message or "").startswith("upstream_stream_error:")
|
||||||
|
|||||||
Reference in New Issue
Block a user