From 5634216f9115e7effb1bcbc3cb0f9f1e1b70c5b0 Mon Sep 17 00:00:00 2001
From: iomgaa <iomgaaycz@gmail.com>
Date: Mon, 20 Jul 2026 06:37:35 -0400
Subject: [PATCH] feat: port three-layer stream liveness watchdog

---
 src/polygateway/streaming.py | 107 ++++++++++++++++++++++++++++++++++
 tests/unit/test_streaming.py | 108 +++++++++++++++++++++++++++++++++++
 2 files changed, 215 insertions(+)
 create mode 100644 src/polygateway/streaming.py
 create mode 100644 tests/unit/test_streaming.py

diff --git a/src/polygateway/streaming.py b/src/polygateway/streaming.py
new file mode 100644
index 0000000..6a64820
--- /dev/null
+++ b/src/polygateway/streaming.py
@@ -0,0 +1,107 @@
+"""三层流式活性超时看门狗: 给任意异步迭代器施加首产出/产出间/总时长三层超时保护。
+
+领域无关、不打日志、不触网: 仅以 asyncio.timeout 包裹取值,便于注入假异步
+迭代器做单测。超时抛 StreamLivenessTimeout(携 kind/elapsed/是否已出首项),
+由 transport 边界翻译为 TransientError 并埋点。
+
+移植自 Video-Tree-TRM5 adapters/streaming.py(三项目同款,语义零改动)。
+活性口径: 看门狗只感知"有无产出",不区分产出内容——transport 把
+reasoning_content 增量同样作为流元素产出,思考流因此天然刷新计时(CHS R1)。
+"""
+
+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 — 三项目冻结的公共名
+    """流活性超时异常。
+
+    属性:
+        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 用 cm.expired() 区分,原样上抛不误吞。
+    """
+    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,并施加三层活性超时。
+
+    关键实现: 超时**只包裹单次 __anext__**,绝不包裹 yield——否则总时长
+    deadline 会在生成器挂起(消费者处理已产出项)期间继续计时,可能取消
+    消费者的 await 而非由本生成器稳定抛出(取消泄漏)。总时长由「每轮取
+    min(本层预算, 剩余总额) 钳制」实现: 钳到剩余总额时若超时即判 total,
+    否则判 ttft/inter_token。finally 关闭底层迭代器释放上游资源。
+    """
+    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..f20fbab
--- /dev/null
+++ b/tests/unit/test_streaming.py
@@ -0,0 +1,108 @@
+"""streaming.py 三层活性看门狗测试(移植自 VT adapters/streaming.py,语义零改动)。"""
+
+import asyncio
+
+import pytest
+
+from polygateway.streaming import StreamLivenessTimeout, stream_with_liveness_timeouts
+
+
+async def _emit(items, *, delay_s=0.0, first_delay_s=None):
+    """可控节奏的假流;first_delay_s 单独控制首项延迟。"""
+    for i, item in enumerate(items):
+        wait = first_delay_s if (i == 0 and first_delay_s is not None) else delay_s
+        if wait:
+            await asyncio.sleep(wait)
+        yield item
+
+
+async def _collect(agen):
+    return [item async for item in agen]
+
+
+class TestPassthrough:
+    async def test_items_pass_through_in_order(self):
+        wrapped = stream_with_liveness_timeouts(
+            _emit(["a", "b", "c"]), ttft_s=1.0, inter_token_s=1.0, total_s=5.0
+        )
+        assert await _collect(wrapped) == ["a", "b", "c"]
+
+    async def test_empty_stream_ends_cleanly(self):
+        wrapped = stream_with_liveness_timeouts(
+            _emit([]), ttft_s=1.0, inter_token_s=1.0, total_s=5.0
+        )
+        assert await _collect(wrapped) == []
+
+
+class TestThreeLayers:
+    async def test_ttft_timeout(self):
+        wrapped = stream_with_liveness_timeouts(
+            _emit(["a"], first_delay_s=0.2), ttft_s=0.05, inter_token_s=1.0, total_s=5.0
+        )
+        with pytest.raises(StreamLivenessTimeout) as ei:
+            await _collect(wrapped)
+        assert ei.value.kind == "ttft"
+        assert ei.value.first_token_seen is False
+
+    async def test_inter_token_timeout(self):
+        wrapped = stream_with_liveness_timeouts(
+            _emit(["a", "b"], delay_s=0.2, first_delay_s=0.0),
+            ttft_s=1.0, inter_token_s=0.05, total_s=5.0,
+        )
+        with pytest.raises(StreamLivenessTimeout) as ei:
+            await _collect(wrapped)
+        assert ei.value.kind == "inter_token"
+        assert ei.value.first_token_seen is True
+
+    async def test_total_timeout(self):
+        # 每项都快,但项数多到总时长超限 → 判 total 而非 inter_token
+        wrapped = stream_with_liveness_timeouts(
+            _emit(list(range(100)), delay_s=0.02), ttft_s=1.0, inter_token_s=1.0, total_s=0.1
+        )
+        with pytest.raises(StreamLivenessTimeout) as ei:
+            await _collect(wrapped)
+        assert ei.value.kind == "total"
+
+
+class TestForeignTimeoutAndCleanup:
+    async def test_upstream_timeout_error_not_reclassified(self):
+        """上游自抛的 TimeoutError 不得被误判为本层活性超时。"""
+
+        async def bad_stream():
+            raise TimeoutError("upstream own timeout")
+            yield  # pragma: no cover
+
+        wrapped = stream_with_liveness_timeouts(
+            bad_stream(), ttft_s=1.0, inter_token_s=1.0, total_s=5.0
+        )
+        with pytest.raises(TimeoutError) as ei:
+            await _collect(wrapped)
+        assert not isinstance(ei.value, StreamLivenessTimeout)
+
+    async def test_underlying_iterator_closed_on_abandon(self):
+        closed = asyncio.Event()
+
+        async def tracked():
+            try:
+                for i in range(100):
+                    await asyncio.sleep(0.01)
+                    yield i
+            finally:
+                closed.set()
+
+        wrapped = stream_with_liveness_timeouts(
+            tracked(), ttft_s=1.0, inter_token_s=1.0, total_s=5.0
+        )
+        assert (await wrapped.__anext__()) == 0
+        await wrapped.aclose()  # 消费方提前放弃 → finally 关闭底层迭代器
+        assert closed.is_set()
+
+    async def test_cancellation_propagates(self):
+        wrapped = stream_with_liveness_timeouts(
+            _emit(["a"], first_delay_s=10.0), ttft_s=30.0, inter_token_s=30.0, total_s=60.0
+        )
+        task = asyncio.ensure_future(_collect(wrapped))
+        await asyncio.sleep(0.05)
+        task.cancel()
+        with pytest.raises(asyncio.CancelledError):
+            await task