feat: bound fusion stream lifecycle

Author: jiwenxu025-boop

scripts/bazi_engine/_streaming.py

@@ -0,0 +1,51 @@
+"""Small, bounded bridge for sending worker events to an async SSE generator."""
+
+import asyncio
+import logging
+import os
+import threading
+from typing import Any
+
+logger = logging.getLogger(__name__)
+
+
+class StreamEventQueue:
+    """Bound worker-to-event-loop events and stop accepting them after cancellation."""
+
+    def __init__(self, loop: asyncio.AbstractEventLoop, *, name: str) -> None:
+        self._loop = loop
+        self._name = name
+        self._queue: asyncio.Queue[tuple[str, Any]] = asyncio.Queue(
+            maxsize=max(1, int(os.getenv("BAZI_STREAM_EVENT_QUEUE_MAX", "128")))
+        )
+        self._closed = threading.Event()
+        self._overflowed = threading.Event()
+
+    @property
+    def closed(self) -> bool:
+        return self._closed.is_set()
+
+    @property
+    def overflowed(self) -> bool:
+        return self._overflowed.is_set()
+
+    def publish(self, event_type: str, payload: Any) -> None:
+        """Schedule a non-blocking event offer from a synchronous worker callback."""
+        if not self._closed.is_set():
+            self._loop.call_soon_threadsafe(self._offer, event_type, payload)
+
+    def _offer(self, event_type: str, payload: Any) -> None:
+        if self._closed.is_set():
+            return
+        try:
+            self._queue.put_nowait((event_type, payload))
+        except asyncio.QueueFull:
+            self._overflowed.set()
+            self._closed.set()
+            logger.warning("SSE event queue overflowed stream=%s", self._name)
+
+    async def get(self, timeout: float) -> tuple[str, Any]:
+        return await asyncio.wait_for(self._queue.get(), timeout=timeout)
+
+    def close(self) -> None:
+        self._closed.set()

scripts/bazi_engine/api.py

@@ -35,6 +35,7 @@
 
 from ._http import close_shared_clients
 from ._runtime import close_blocking_executor, submit_blocking
+from ._streaming import StreamEventQueue
 from ._version import __version__
 from .chart import build_chart
 
@@ -52,6 +53,9 @@
 }
 _MAX_AI_STREAMS = int(os.getenv("BAZI_MAX_AI_STREAMS", "3"))
 _AI_STREAM_SLOTS = threading.BoundedSemaphore(max(1, _MAX_AI_STREAMS))
+_FUSION_STREAM_TOTAL_TIMEOUT = float(os.getenv("BAZI_FUSION_STREAM_TOTAL_TIMEOUT", "150"))
+_FUSION_STREAM_IDLE_TIMEOUT = float(os.getenv("BAZI_FUSION_STREAM_IDLE_TIMEOUT", "45"))
+_STREAM_HEARTBEAT_INTERVAL = float(os.getenv("BAZI_STREAM_HEARTBEAT_INTERVAL", "15"))
 _CORS_ORIGINS = [
     origin.strip()
     for origin in os.getenv("BAZI_CORS_ORIGINS", "").split(",")
@@ -324,39 +328,47 @@ def run_build():
                         chart_data.get("life_stage", ""),
                         age_info=age_info,
                     )
-                    fusion_queue: asyncio.Queue = asyncio.Queue()
+                    fusion_events = StreamEventQueue(loop, name="chart-fusion")
 
-                    def on_token(tok: str, _fusion_queue=fusion_queue):
-                        loop.call_soon_threadsafe(
-                            _fusion_queue.put_nowait, ("token", tok))
+                    def on_token(tok: str, _fusion_events=fusion_events):
+                        _fusion_events.publish("token", tok)
 
                     fusion_done_flag = {"done": False}
 
                     def run_fusion(
                         _fusion_key=fusion_key,
-                        _fusion_queue=fusion_queue,
+                        _fusion_events=fusion_events,
                         _pkg=pkg,
                         _fusion_done_flag=fusion_done_flag,
                         _fusion_meta=fusion_meta,
                         _fusion_started_at=fusion_started_at,
                     ):
                         try:
                             if not _fusion_key:
-                                loop.call_soon_threadsafe(_fusion_queue.put_nowait, ("fusion_error", "DEEPSEEK_API_KEY未设置"))
+                                _fusion_events.publish("fusion_error", "DEEPSEEK_API_KEY未设置")
                                 return
                             full = generate_fusion_report(
                                 _pkg,
                                 on_chunk=on_token,
                                 result_metadata=_fusion_meta,
                             )
+                            if _fusion_events.closed:
+                                _record_fusion_generation(
+                                    _fusion_meta["generation_id"],
+                                    _fusion_started_at,
+                                    "cancelled",
+                                    _fusion_meta,
+                                    "stream_closed",
+                                )
+                                return
                             if full:
                                 _record_fusion_generation(
                                     _fusion_meta["generation_id"],
                                     _fusion_started_at,
                                     "success",
                                     _fusion_meta,
                                 )
-                                loop.call_soon_threadsafe(_fusion_queue.put_nowait, ("fusion_done", full))
+                                _fusion_events.publish("fusion_done", full)
                             else:
                                 _record_fusion_generation(
                                     _fusion_meta["generation_id"],
@@ -365,10 +377,7 @@ def run_fusion(
                                     _fusion_meta,
                                     "empty_response",
                                 )
-                                loop.call_soon_threadsafe(
-                                    _fusion_queue.put_nowait,
-                                    ("fusion_error", "融合报告暂时不可用，请稍后重试"),
-                                )
+                                _fusion_events.publish("fusion_error", "融合报告暂时不可用，请稍后重试")
                             _fusion_done_flag["done"] = True
                         except Exception as error:
                             error_class = _fusion_error_class(error)
@@ -383,36 +392,41 @@ def run_fusion(
                                 "fusion generation failed id=%s class=%s type=%s",
                                 _fusion_meta["generation_id"], error_class, type(error).__name__,
                             )
-                            loop.call_soon_threadsafe(
-                                _fusion_queue.put_nowait,
-                                ("fusion_error", "融合报告暂时不可用，请稍后重试"),
-                            )
+                            _fusion_events.publish("fusion_error", "融合报告暂时不可用，请稍后重试")
                             _fusion_done_flag["done"] = True
 
                     submit_blocking(loop, run_fusion)
 
-                    _fusion_start = loop.time()
+                    _fusion_last_activity_at = loop.time()
                     while True:
                         try:
-                            ft, fd = await asyncio.wait_for(fusion_queue.get(), timeout=15.0)
+                            ft, fd = await fusion_events.get(_STREAM_HEARTBEAT_INTERVAL)
                         except TimeoutError:
                             if fusion_done_flag["done"]:
                                 break  # thread finished but queue empty
-                            elapsed = loop.time() - _fusion_start
-                            if elapsed > 90:
-                                yield f"data: {json.dumps({'phase': 'personality_error', 'message': f'融合超时({elapsed:.0f}s)'})}\n\n"
+                            total_elapsed = loop.time() - fusion_started_at
+                            idle_elapsed = loop.time() - _fusion_last_activity_at
+                            if total_elapsed > _FUSION_STREAM_TOTAL_TIMEOUT:
+                                fusion_events.close()
+                                yield f"data: {json.dumps({'phase': 'personality_error', 'message': '融合报告超时，请稍后重试'})}\n\n"
+                                break
+                            if idle_elapsed > _FUSION_STREAM_IDLE_TIMEOUT or fusion_events.overflowed:
+                                fusion_events.close()
+                                yield f"data: {json.dumps({'phase': 'personality_error', 'message': '融合报告响应超时，请稍后重试'})}\n\n"
                                 break
+                            yield f"data: {json.dumps({'phase': 'heartbeat'})}\n\n"
                             continue
                         if ft == "token":
                             yield f"data: {json.dumps({'phase': 'personality_token', 'token': fd})}\n\n"
-                            _fusion_start = loop.time()  # reset timeout on activity
+                            _fusion_last_activity_at = loop.time()
                         elif ft == "fusion_done":
                             if fd:
                                 yield f"data: {json.dumps({'phase': 'personality_done', 'full': fd, 'meta': fusion_meta})}\n\n"
                             break
                         elif ft == "fusion_error":
                             yield f"data: {json.dumps({'phase': 'personality_error', 'message': fd})}\n\n"
                             break
+                    fusion_events.close()
                 except Exception as error:
                     error_class = _fusion_error_class(error)
                     _record_fusion_generation(
@@ -876,14 +890,25 @@ async def err_gen():
 
     # SSE 生成器 — 用 asyncio.Queue 桥接同步 LLM 流，实现真正的逐 token 推送
     async def stream_fusion():
-        queue: asyncio.Queue = asyncio.Queue()
         loop = asyncio.get_event_loop()
+        events = StreamEventQueue(loop, name="fusion-direct")
         fusion_meta: dict = {"generation_id": uuid.uuid4().hex}
         fusion_started_at = time.monotonic()
+        generation_recorded = threading.Event()
+
+        def record_once(
+            outcome: Literal["success", "failure", "cancelled"],
+            error_class: str | None = None,
+        ) -> None:
+            if not generation_recorded.is_set():
+                generation_recorded.set()
+                _record_fusion_generation(
+                    fusion_meta["generation_id"], fusion_started_at, outcome, fusion_meta, error_class,
+                )
 
         def on_token(token: str):
             """LLM 每吐一个 token，立刻推入 queue"""
-            loop.call_soon_threadsafe(queue.put_nowait, ("token", token))
+            events.publish("token", token)
 
         def run_llm():
             """在线程中跑同步流式 LLM 调用"""
@@ -893,45 +918,75 @@ def run_llm():
                     on_chunk=on_token,
                     result_metadata=fusion_meta,
                 )
+                if events.closed:
+                    record_once("cancelled", "client_disconnected")
+                    return
                 if full:
-                    _record_fusion_generation(
-                        fusion_meta["generation_id"], fusion_started_at, "success", fusion_meta,
-                    )
-                    loop.call_soon_threadsafe(queue.put_nowait, ("done", full))
+                    record_once("success")
+                    events.publish("done", full)
                 else:
-                    _record_fusion_generation(
-                        fusion_meta["generation_id"], fusion_started_at, "failure", fusion_meta,
-                        "empty_response",
-                    )
-                    loop.call_soon_threadsafe(queue.put_nowait, ("error", "融合报告暂时不可用，请稍后重试"))
+                    record_once("failure", "empty_response")
+                    events.publish("error", "融合报告暂时不可用，请稍后重试")
             except Exception as error:
                 error_class = _fusion_error_class(error)
-                _record_fusion_generation(
-                    fusion_meta["generation_id"], fusion_started_at, "failure", fusion_meta, error_class,
-                )
+                record_once("failure", error_class)
                 logger.warning(
                     "fusion generation failed id=%s class=%s type=%s",
                     fusion_meta["generation_id"], error_class, type(error).__name__,
                 )
-                loop.call_soon_threadsafe(queue.put_nowait, ("error", "融合报告暂时不可用，请稍后重试"))
+                events.publish("error", "融合报告暂时不可用，请稍后重试")
 
         submit_blocking(loop, run_llm)
 
-        while True:
-            msg_type, msg_data = await queue.get()
-            if msg_type == "token":
-                yield f"data: {json.dumps({'token': msg_data})}\n\n"
-            elif msg_type == "done":
-                if msg_data:
-                    yield f"data: {json.dumps({'done': True, 'length': len(msg_data), 'full': msg_data, 'meta': fusion_meta})}\n\n"
-                else:
-                    yield f"data: {json.dumps({'error': 'LLM 调用失败，请稍后重试'})}\n\n"
-                yield "data: [DONE]\n\n"
-                return
-            elif msg_type == "error":
-                yield f"data: {json.dumps({'error': msg_data})}\n\n"
-                yield "data: [DONE]\n\n"
-                return
+        last_event_at = loop.time()
+        try:
+            while True:
+                elapsed = loop.time() - fusion_started_at
+                idle = loop.time() - last_event_at
+                if elapsed >= _FUSION_STREAM_TOTAL_TIMEOUT:
+                    events.close()
+                    record_once("failure", "total_timeout")
+                    yield f"data: {json.dumps({'error': '融合报告超时，请稍后重试'})}\n\n"
+                    yield "data: [DONE]\n\n"
+                    return
+                if idle >= _FUSION_STREAM_IDLE_TIMEOUT:
+                    events.close()
+                    record_once("failure", "idle_timeout")
+                    yield f"data: {json.dumps({'error': '融合报告响应超时，请稍后重试'})}\n\n"
+                    yield "data: [DONE]\n\n"
+                    return
+                timeout = min(
+                    _STREAM_HEARTBEAT_INTERVAL,
+                    _FUSION_STREAM_TOTAL_TIMEOUT - elapsed,
+                    _FUSION_STREAM_IDLE_TIMEOUT - idle,
+                )
+                try:
+                    msg_type, msg_data = await events.get(max(0.01, timeout))
+                except TimeoutError:
+                    if events.overflowed:
+                        events.close()
+                        record_once("failure", "queue_overflow")
+                        yield f"data: {json.dumps({'error': '融合报告处理繁忙，请稍后重试'})}\n\n"
+                        yield "data: [DONE]\n\n"
+                        return
+                    yield f"data: {json.dumps({'heartbeat': True})}\n\n"
+                    continue
+                last_event_at = loop.time()
+                if msg_type == "token":
+                    yield f"data: {json.dumps({'token': msg_data})}\n\n"
+                elif msg_type == "done":
+                    if msg_data:
+                        yield f"data: {json.dumps({'done': True, 'length': len(msg_data), 'full': msg_data, 'meta': fusion_meta})}\n\n"
+                    else:
+                        yield f"data: {json.dumps({'error': 'LLM 调用失败，请稍后重试'})}\n\n"
+                    yield "data: [DONE]\n\n"
+                    return
+                elif msg_type == "error":
+                    yield f"data: {json.dumps({'error': msg_data})}\n\n"
+                    yield "data: [DONE]\n\n"
+                    return
+        finally:
+            events.close()
 
     return _limited_stream_response(stream_fusion())
 

scripts/tests/test_api.py

@@ -4,6 +4,7 @@
 import importlib
 import json
 import sys
+import time
 from concurrent.futures import ThreadPoolExecutor
 from contextlib import asynccontextmanager
 from datetime import date
@@ -288,6 +289,40 @@ def fake_generate(*_args, **_kwargs):
     assert generation["error_class"] == "provider_rejected"
 
 
+def test_fusion_stream_stops_after_idle_timeout(monkeypatch, tmp_path):
+    import bazi_engine.api as api_module
+    import bazi_engine.personality_fusion as fusion_module
+    from bazi_engine.api import app
+
+    monkeypatch.setenv("BAZI_FUSION_ENGINE", "1")
+    monkeypatch.setenv("DEEPSEEK_API_KEY", "test-key")
+    monkeypatch.setattr(api_module, "_GENERATION_DIR", tmp_path)
+    monkeypatch.setattr(api_module, "_FUSION_STREAM_IDLE_TIMEOUT", 0.01)
+    monkeypatch.setattr(api_module, "_FUSION_STREAM_TOTAL_TIMEOUT", 0.5)
+    monkeypatch.setattr(api_module, "_STREAM_HEARTBEAT_INTERVAL", 0.005)
+
+    def fake_generate(*_args, **_kwargs):
+        time.sleep(0.05)
+        return "late report"
+
+    monkeypatch.setattr(fusion_module, "generate_fusion_report", fake_generate)
+    with TestClient(app).stream(
+        "POST",
+        "/api/personality/fusion/stream",
+        json={"personality": {"traits": {"社交": "内敛"}}},
+    ) as response:
+        body = "".join(response.iter_text())
+
+    assert response.status_code == 200
+    events = [
+        json.loads(line[6:])
+        for line in body.splitlines()
+        if line.startswith("data: {")
+    ]
+    assert any(event.get("error") == "融合报告响应超时，请稍后重试" for event in events)
+    assert "late report" not in body
+
+
 def test_fusion_feedback_saves_metadata_without_report_or_birth_data(monkeypatch, tmp_path):
     """融合反馈只保存分析元数据和报告哈希，不落报告正文或出生资料。"""
     import bazi_engine.api as api_module