feat: 添加流式请求客户端断连检测

在等待上游首字节期间检测客户端是否已断开连接,若检测到断连则及时取消请求并清理资源,避免资源浪费。

- 新增 ClientDisconnectedException 和 wait_for_with_disconnect_detection 工具函数
- ChatHandlerBase 和 CliMessageHandlerBase 均支持 http_request 参数用于断连检测
- 流式传输期间使用后台任务主动检测断连状态
- 断连时设置状态码 499 并记录日志
This commit is contained in:
fawney19
2026-01-28 17:10:21 +08:00
parent bae278a0c1
commit 91986baa57
4 changed files with 222 additions and 15 deletions

View File

@@ -27,8 +27,20 @@
from __future__ import annotations
import asyncio
import time
from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Protocol, runtime_checkable
from typing import (
TYPE_CHECKING,
Any,
Awaitable,
Callable,
Coroutine,
Dict,
Optional,
Protocol,
TypeVar,
runtime_checkable,
)
from fastapi import Request
from fastapi.responses import JSONResponse, StreamingResponse
@@ -322,8 +334,7 @@ class MessageHandlerProtocol(Protocol):
"""
消息处理器协议 - 定义标准接口
ChatHandlerBase 使用完整签名(含 request, http_request
CliMessageHandlerBase 使用简化签名(仅 original_request_body, original_headers
ChatHandlerBase 和 CliMessageHandlerBase 均支持 http_request 参数用于客户端断连检测
"""
async def process_stream(
@@ -562,3 +573,79 @@ class BaseMessageHandler:
else:
# 未知异常:完整堆栈
logger.exception(f"{message}: {error}")
# ============================================================================
# 客户端断连检测
# ============================================================================
class ClientDisconnectedException(Exception):
"""客户端在等待首字节时断开连接"""
pass
_T = TypeVar("_T")
async def wait_for_with_disconnect_detection(
coro: Coroutine[Any, Any, _T],
timeout: float,
is_disconnected: Callable[[], Awaitable[bool]],
request_id: str,
check_interval: float = 0.5,
) -> _T:
"""
等待协程完成,同时检测客户端断连
在等待上游响应(如首字节)时,定期检测客户端是否已断连。
若检测到断连,取消任务并抛出 ClientDisconnectedException。
Args:
coro: 要等待的协程
timeout: 超时时间(秒)
is_disconnected: 异步断连检测函数(如 http_request.is_disconnected
request_id: 请求 ID用于日志
check_interval: 断连检测间隔(秒),默认 0.5s
Returns:
协程的返回值
Raises:
ClientDisconnectedException: 客户端断连
asyncio.TimeoutError: 超时
asyncio.CancelledError: 任务被外部取消
"""
task = asyncio.create_task(coro)
client_disconnected = False
async def check_client_disconnect() -> None:
nonlocal client_disconnected
while not task.done():
await asyncio.sleep(check_interval)
try:
if await is_disconnected():
client_disconnected = True
logger.debug(f" [{request_id}] 检测到客户端断连,取消预取任务")
task.cancel()
break
except Exception as e:
logger.debug(f" [{request_id}] 断连检测异常: {e}")
disconnect_task = asyncio.create_task(check_client_disconnect())
try:
return await asyncio.wait_for(task, timeout=timeout)
except asyncio.CancelledError:
if client_disconnected:
raise ClientDisconnectedException("Client disconnected during prefetch")
raise
finally:
disconnect_task.cancel()
try:
await disconnect_task
except asyncio.CancelledError:
pass

View File

@@ -22,14 +22,18 @@ Chat Handler Base - Chat API 格式的通用基类
import asyncio
import json
from abc import ABC, abstractmethod
from typing import Any, AsyncGenerator, Callable, Dict, Optional, Union
from typing import Any, AsyncGenerator, Awaitable, Callable, Dict, Optional, Union
import httpx
from fastapi import BackgroundTasks, Request
from fastapi.responses import JSONResponse, StreamingResponse
from sqlalchemy.orm import Session
from src.api.handlers.base.base_handler import BaseMessageHandler
from src.api.handlers.base.base_handler import (
BaseMessageHandler,
ClientDisconnectedException,
wait_for_with_disconnect_detection,
)
from src.api.handlers.base.parsers import get_parser_for_format
from src.api.handlers.base.request_builder import PassthroughRequestBuilder
from src.api.handlers.base.response_parser import ResponseParser
@@ -499,6 +503,7 @@ class ChatHandlerBase(BaseMessageHandler, ABC):
original_headers,
query_params,
candidate,
is_disconnected=http_request.is_disconnected,
)
try:
@@ -608,6 +613,7 @@ class ChatHandlerBase(BaseMessageHandler, ABC):
original_headers: Dict[str, str],
query_params: Optional[Dict[str, str]] = None,
candidate: Optional[ProviderCandidate] = None,
is_disconnected: Optional[Callable[[], Awaitable[bool]]] = None,
) -> AsyncGenerator[bytes, None]:
"""执行流式请求并返回流生成器"""
# 重置上下文状态(重试时清除之前的数据)
@@ -757,7 +763,29 @@ class ChatHandlerBase(BaseMessageHandler, ABC):
try:
# 使用 asyncio.wait_for 包裹整个"建立连接 + 获取首字节"阶段
# stream_first_byte_timeout 控制首字节超时,避免上游长时间无响应
await asyncio.wait_for(_connect_and_prefetch(), timeout=request_timeout)
# 同时检测客户端断连,避免客户端已断开但服务端仍在等待上游响应
if is_disconnected is not None:
await wait_for_with_disconnect_detection(
_connect_and_prefetch(),
timeout=request_timeout,
is_disconnected=is_disconnected,
request_id=self.request_id,
)
else:
await asyncio.wait_for(_connect_and_prefetch(), timeout=request_timeout)
except ClientDisconnectedException:
# 客户端断开连接,清理资源
if response_ctx is not None:
try:
await response_ctx.__aexit__(None, None, None)
except Exception:
pass
await http_client.aclose()
logger.warning(f" [{self.request_id}] 客户端在等待首字节时断开连接")
ctx.status_code = 499
ctx.error_message = "client_disconnected_during_prefetch"
raise
except asyncio.TimeoutError:
# 整体请求超时(建立连接 + 获取首字节)

View File

@@ -197,6 +197,7 @@ class CliAdapterBase(ApiAdapter):
original_headers=original_headers,
query_params=query_params,
path_params=context.path_params,
http_request=http_request,
)
return await handler.process_sync(
original_request_body=original_request_body,

View File

@@ -26,7 +26,7 @@ from typing import (
)
import httpx
from fastapi import BackgroundTasks
from fastapi import BackgroundTasks, Request
from fastapi.responses import JSONResponse, StreamingResponse
from sqlalchemy.orm import Session
@@ -35,7 +35,9 @@ if TYPE_CHECKING:
from src.api.handlers.base.base_handler import (
BaseMessageHandler,
ClientDisconnectedException,
MessageTelemetry,
wait_for_with_disconnect_detection,
)
from src.api.handlers.base.parsers import get_parser_for_format
from src.api.handlers.base.request_builder import PassthroughRequestBuilder
@@ -505,6 +507,7 @@ class CliMessageHandlerBase(BaseMessageHandler):
original_headers: Dict[str, str],
query_params: Optional[Dict[str, str]] = None,
path_params: Optional[Dict[str, Any]] = None,
http_request: Optional[Request] = None,
) -> StreamingResponse:
"""
处理流式请求
@@ -514,6 +517,13 @@ class CliMessageHandlerBase(BaseMessageHandler):
2. 定义请求函数(供 FallbackOrchestrator 调用)
3. 执行请求并返回 StreamingResponse
4. 后台任务记录统计信息
Args:
original_request_body: 原始请求体
original_headers: 原始请求头
query_params: 查询参数
path_params: 路径参数
http_request: FastAPI Request 对象,用于检测客户端断连
"""
logger.debug(f"开始流式响应处理 ({self.FORMAT_ID})")
@@ -550,6 +560,7 @@ class CliMessageHandlerBase(BaseMessageHandler):
original_headers,
query_params,
candidate,
http_request, # 传递 http_request 用于断连检测
)
try:
@@ -600,8 +611,8 @@ class CliMessageHandlerBase(BaseMessageHandler):
original_request_body,
)
# 创建监控流
monitored_stream = self._create_monitored_stream(ctx, stream_generator)
# 创建监控流(传递 http_request 用于断连检测)
monitored_stream = self._create_monitored_stream(ctx, stream_generator, http_request)
# 透传提供商的响应头给客户端
# 同时添加必要的 SSE 头以确保流式传输正常工作
@@ -640,6 +651,7 @@ class CliMessageHandlerBase(BaseMessageHandler):
original_headers: Dict[str, str],
query_params: Optional[Dict[str, str]] = None,
candidate: Optional[ProviderCandidate] = None,
http_request: Optional[Request] = None,
) -> AsyncGenerator[bytes, None]:
"""执行流式请求并返回流生成器"""
# 重置上下文状态(重试时清除之前的数据,避免累积)
@@ -789,7 +801,16 @@ class CliMessageHandlerBase(BaseMessageHandler):
try:
# 使用 asyncio.wait_for 包裹整个"建立连接 + 获取首字节"阶段
# stream_first_byte_timeout 控制首字节超时,避免上游长时间无响应
await asyncio.wait_for(_connect_and_prefetch(), timeout=request_timeout)
# 同时检测客户端断连,避免客户端已断开但服务端仍在等待上游响应
if http_request is not None:
await wait_for_with_disconnect_detection(
_connect_and_prefetch(),
timeout=request_timeout,
is_disconnected=http_request.is_disconnected,
request_id=self.request_id,
)
else:
await asyncio.wait_for(_connect_and_prefetch(), timeout=request_timeout)
except asyncio.TimeoutError:
# 整体请求超时(建立连接 + 获取首字节)
@@ -808,6 +829,19 @@ class CliMessageHandlerBase(BaseMessageHandler):
timeout=int(request_timeout),
)
except ClientDisconnectedException:
# 客户端断开连接,清理资源
if response_ctx is not None:
try:
await response_ctx.__aexit__(None, None, None)
except Exception:
pass
await http_client.aclose()
logger.warning(f" [{self.request_id}] 客户端在等待首字节时断开连接")
ctx.status_code = 499
ctx.error_message = "client_disconnected_during_prefetch"
raise
except httpx.HTTPStatusError as e:
error_text = await self._extract_error_text(e)
logger.error(
@@ -1692,17 +1726,74 @@ class CliMessageHandlerBase(BaseMessageHandler):
self,
ctx: StreamContext,
stream_generator: AsyncGenerator[bytes, None],
http_request: Optional[Request] = None,
) -> AsyncGenerator[bytes, None]:
"""创建带监控的流生成器"""
"""
创建带监控的流生成器
支持两种断连检测方式:
1. 如果提供了 http_request使用后台任务主动检测客户端断连
2. 如果未提供,仅依赖 asyncio.CancelledError 被动检测
Args:
ctx: 流上下文
stream_generator: 底层流生成器
http_request: FastAPI Request 对象,用于检测客户端断连
"""
import time as time_module
last_chunk_time = time_module.time()
chunk_count = 0
try:
async for chunk in stream_generator:
last_chunk_time = time_module.time()
chunk_count += 1
yield chunk
if http_request is not None:
# 使用后台任务检测断连,完全不阻塞流式传输
disconnected = False
async def check_disconnect_background() -> None:
nonlocal disconnected
while not disconnected and not ctx.has_completion:
await asyncio.sleep(0.5)
try:
if await http_request.is_disconnected():
disconnected = True
break
except Exception as e:
# 检测失败时不中断流,继续传输
logger.debug(f"ID:{ctx.request_id} | 断连检测异常: {e}")
# 启动后台检查任务
check_task = asyncio.create_task(check_disconnect_background())
try:
async for chunk in stream_generator:
if disconnected:
# 如果响应已完成,客户端断开不算失败
if ctx.has_completion:
logger.info(
f"ID:{ctx.request_id} | Client disconnected after completion"
)
else:
logger.warning(f"ID:{ctx.request_id} | Client disconnected")
ctx.status_code = 499
ctx.error_message = "client_disconnected"
break
last_chunk_time = time_module.time()
chunk_count += 1
yield chunk
finally:
check_task.cancel()
try:
await check_task
except asyncio.CancelledError:
pass
else:
# 无 http_request仅被动监控
async for chunk in stream_generator:
last_chunk_time = time_module.time()
chunk_count += 1
yield chunk
except asyncio.CancelledError:
# 计算距离上次收到 chunk 的时间
time_since_last_chunk = time_module.time() - last_chunk_time