From c2ff2855b727bf39d8f456dee08eda2a0c6c49f3 Mon Sep 17 00:00:00 2001 From: iomgaa Date: Mon, 6 Jul 2026 22:43:01 -0400 Subject: [PATCH] =?UTF-8?q?feat(adapters):=20=E5=AE=9E=E7=8E=B0=E4=B8=89?= =?UTF-8?q?=E5=B1=82=E6=B5=81=E5=BC=8F=E7=9C=8B=E9=97=A8=E7=8B=97?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit TTFT / inter_token / total 超时保护。 参考 CHSAnalyzer2 streaming.py。 Co-Authored-By: Claude Opus 4.6 (1M context) --- adapters/streaming.py | 131 ++++++++++++++++++++++++++++++ tests/unit/test_streaming.py | 152 +++++++++++++++++++++++++++++++++++ 2 files changed, 283 insertions(+) create mode 100644 adapters/streaming.py create mode 100644 tests/unit/test_streaming.py diff --git a/adapters/streaming.py b/adapters/streaming.py new file mode 100644 index 0000000..3bea48f --- /dev/null +++ b/adapters/streaming.py @@ -0,0 +1,131 @@ +"""三层流式活性超时看门狗:给任意异步迭代器施加首产出/产出间/总时长三层超时保护。 + +领域无关、不打日志、不触网(Humble Object):仅以 asyncio.timeout 包裹取值, +便于注入假异步迭代器做单测。超时抛 StreamLivenessTimeout(携 kind/elapsed/是否已出首项), +由 adapter 边界翻译为领域错误并埋点。 + +参考: CHSAnalyzer2 app/providers/streaming.py +""" + +from __future__ import annotations + +import asyncio +import contextlib +import time +from typing import TYPE_CHECKING, TypeVar + +if TYPE_CHECKING: + from collections.abc import AsyncIterator + +_T = TypeVar("_T") + + +class StreamLivenessTimeout(Exception): # noqa: N818 — 规格要求名称,非 Error 后缀 + """流活性超时异常。 + + 属性: + kind: 超时类型,取值 {"ttft", "inter_token", "total"}。 + elapsed_s: 自流开始到触发超时的秒数。 + first_token_seen: 触发时是否已产出过首项。 + """ + + def __init__(self, kind: str, elapsed_s: float, first_token_seen: bool) -> None: + self.kind = kind + self.elapsed_s = elapsed_s + self.first_token_seen = first_token_seen + super().__init__(f"流活性超时({kind}, elapsed={elapsed_s:.1f}s)") + + +async def _anext_within( + it: AsyncIterator[_T], + timeout_s: float, + *, + kind: str, + start: float, + first: bool, +) -> _T: + """限时取下一项;本层 deadline 触发抛 StreamLivenessTimeout(kind),上游自抛的 TimeoutError 原样上抛。 + + 参数: + it: 底层异步迭代器。 + timeout_s: 本次取项的超时预算(秒)。 + kind: 本层超时归类(ttft/inter_token/total),仅在本层 deadline 触发时使用。 + start: 流开始时刻(time.monotonic),用于计算 elapsed_s。 + first: 是否尚未产出过首项(用于 StreamLivenessTimeout.first_token_seen)。 + + 返回: + source 的下一项。 + + 异常: + StreamLivenessTimeout: 本层 asyncio.timeout deadline 触发(cm.expired())。 + StopAsyncIteration: 底层迭代器耗尽(由调用方处理收尾)。 + """ + try: + async with asyncio.timeout(timeout_s) as cm: + return await it.__anext__() + except TimeoutError: + if not cm.expired(): + raise # 上游自抛的 TimeoutError,非本层 deadline,原样上抛 + raise StreamLivenessTimeout(kind, time.monotonic() - start, not first) from None + + +async def stream_with_liveness_timeouts( + source: AsyncIterator[_T], + *, + ttft_s: float, + inter_token_s: float, + total_s: float, +) -> AsyncIterator[_T]: + """逐项产出 source,并施加三层活性超时。 + + 参数: + source: 待包裹的异步迭代器(如 VLM 流式分片)。 + ttft_s: 首项产出超时(time-to-first-token,秒)。 + inter_token_s: 相邻两项之间的产出超时(秒)。 + total_s: 整条流的总时长超时(秒)。 + + 产出: + 与 source 逐项一致的元素。 + + 异常: + StreamLivenessTimeout: 首项超 ttft_s、后续项间超 inter_token_s 或整体超 total_s。 + + 关键实现: + 超时**只包裹单次 __anext__**,绝不包裹 yield——否则总时长 deadline 会在生成器挂起 + (消费者处理已产出项)期间继续计时,可能取消消费者的 await 而非由本生成器稳定抛出 + StreamLivenessTimeout(取消泄漏)。总时长改由「每轮取 min(本层预算, 剩余总额) 钳制」 + 实现:钳到剩余总额(clamped)时若超时即判 total,否则判 ttft/inter_token。用 + asyncio.timeout(...) 的 cm.expired() 区分本层 deadline 与上游自抛的 TimeoutError。 + finally 中调用底层迭代器的 aclose(若有)以释放上游资源。 + """ + it = source.__aiter__() + start = time.monotonic() + deadline = start + total_s + first = True + try: + while True: + remaining_total = deadline - time.monotonic() + if remaining_total <= 0: + raise StreamLivenessTimeout("total", time.monotonic() - start, not first) + budget = ttft_s if first else inter_token_s + # 剩余总额比本层预算更紧 → 先到的是 total + clamped = remaining_total <= budget + kind = "total" if clamped else ("ttft" if first else "inter_token") + try: + item = await _anext_within( + it, + remaining_total if clamped else budget, + kind=kind, + start=start, + first=first, + ) + except StopAsyncIteration: + return + first = False + yield item + finally: + # 收尾关闭上游迭代器;清理失败不得掩盖主异常,故显式抑制 + aclose = getattr(it, "aclose", None) + if aclose is not None: + with contextlib.suppress(Exception): + await aclose() diff --git a/tests/unit/test_streaming.py b/tests/unit/test_streaming.py new file mode 100644 index 0000000..515b77e --- /dev/null +++ b/tests/unit/test_streaming.py @@ -0,0 +1,152 @@ +"""三层流式看门狗 (stream_with_liveness_timeouts) 单元测试。 + +验证 TTFT / inter_token / total 三层超时保护行为, +以及正常透传、空流等边界场景。 +""" + +from __future__ import annotations + +import asyncio +from typing import TYPE_CHECKING, TypeVar + +if TYPE_CHECKING: + from collections.abc import AsyncIterator + +import pytest + +from adapters.streaming import StreamLivenessTimeout, stream_with_liveness_timeouts + +_T = TypeVar("_T") + + +async def _items_with_delays( + items: list[_T], + delays: list[float], +) -> AsyncIterator[_T]: + """辅助异步生成器:按 delays 指定的延迟逐项产出 items。 + + 参数: + items: 待产出的元素列表。 + delays: 每个元素产出前的延迟秒数,长度须与 items 一致。 + """ + for item, delay in zip(items, delays, strict=True): + await asyncio.sleep(delay) + yield item + + +@pytest.mark.asyncio +async def test_all_items_yielded_within_budget() -> None: + """正常场景:所有 item 在超时预算内产出,全部透传。""" + items = ["a", "b", "c", "d"] + delays = [0.01, 0.01, 0.01, 0.01] + + result: list[str] = [] + async for chunk in stream_with_liveness_timeouts( + _items_with_delays(items, delays), + ttft_s=1.0, + inter_token_s=1.0, + total_s=5.0, + ): + result.append(chunk) + + assert result == items + + +@pytest.mark.asyncio +async def test_ttft_timeout_fires() -> None: + """首 token 超时:第一个 item 延迟超过 ttft_s,触发 kind='ttft'。""" + items = ["slow_first"] + delays = [0.5] + + with pytest.raises(StreamLivenessTimeout) as exc_info: + async for _ in stream_with_liveness_timeouts( + _items_with_delays(items, delays), + ttft_s=0.05, + inter_token_s=5.0, + total_s=10.0, + ): + pass # pragma: no cover — 不应到达 + + err = exc_info.value + assert err.kind == "ttft" + assert err.first_token_seen is False + + +@pytest.mark.asyncio +async def test_inter_token_timeout_fires() -> None: + """token 间隔超时:首项正常,第二项延迟超过 inter_token_s。""" + items = ["fast", "very_slow"] + delays = [0.01, 0.5] + + with pytest.raises(StreamLivenessTimeout) as exc_info: + async for _ in stream_with_liveness_timeouts( + _items_with_delays(items, delays), + ttft_s=5.0, + inter_token_s=0.05, + total_s=10.0, + ): + pass + + err = exc_info.value + assert err.kind == "inter_token" + assert err.first_token_seen is True + + +@pytest.mark.asyncio +async def test_total_timeout_fires() -> None: + """总时长超时:每个 item 都在单层预算内,但累计超过 total_s。""" + items = ["a", "b", "c", "d", "e", "f"] + # 每项 0.05s,累计 0.30s > total_s=0.15s + delays = [0.05, 0.05, 0.05, 0.05, 0.05, 0.05] + + with pytest.raises(StreamLivenessTimeout) as exc_info: + async for _ in stream_with_liveness_timeouts( + _items_with_delays(items, delays), + ttft_s=5.0, + inter_token_s=5.0, + total_s=0.15, + ): + pass + + err = exc_info.value + assert err.kind == "total" + + +@pytest.mark.asyncio +async def test_thinking_tokens_refresh_watchdog() -> None: + """(False, 'think') 元组正常透传,不影响看门狗计时。""" + items = [ + (False, "think_1"), + (False, "think_2"), + (True, "real_token"), + ] + delays = [0.01, 0.01, 0.01] + + result: list[tuple[bool, str]] = [] + async for chunk in stream_with_liveness_timeouts( + _items_with_delays(items, delays), + ttft_s=1.0, + inter_token_s=1.0, + total_s=5.0, + ): + result.append(chunk) + + assert result == items + + +@pytest.mark.asyncio +async def test_empty_stream_no_error() -> None: + """空流:source 立即耗尽,不抛异常,正常结束。""" + items: list[str] = [] + delays: list[float] = [] + + result: list[str] = [] + async for chunk in stream_with_liveness_timeouts( + _items_with_delays(items, delays), + ttft_s=1.0, + inter_token_s=1.0, + total_s=5.0, + ): + result.append(chunk) # pragma: no cover + + assert result == []