From 3d020c8ceb935b726aa252f5bbe9adf525a6f9fd Mon Sep 17 00:00:00 2001 From: jxxghp Date: Sun, 26 Jul 2026 08:55:34 +0800 Subject: [PATCH] fix(search): keep large SSE searches connected (#6186) --- app/api/endpoints/search.py | 108 ++++++++++++++++++++++--- docs/mcp-api.md | 2 + skills/moviepilot-api/SKILL.md | 4 +- tests/test_builtin_skill_boundaries.py | 2 +- tests/test_search_stream.py | 103 +++++++++++++++++++++++ 5 files changed, 208 insertions(+), 11 deletions(-) create mode 100644 tests/test_search_stream.py diff --git a/app/api/endpoints/search.py b/app/api/endpoints/search.py index bfd24d8f1..d4477d1df 100644 --- a/app/api/endpoints/search.py +++ b/app/api/endpoints/search.py @@ -1,6 +1,8 @@ import asyncio import json -from typing import List, Any, Optional, AsyncIterator +import time +from typing import Any, AsyncIterator, Iterator, List, Optional +from uuid import uuid4 from fastapi import APIRouter, Depends, Body, Request from fastapi.responses import StreamingResponse @@ -23,6 +25,12 @@ router = APIRouter() _SSE_APPEND_FLUSH_INTERVAL = 1 _SSE_APPEND_MAX_ITEMS = 48 +_SSE_HEARTBEAT_INTERVAL = 15 +_SSE_REPLACE_MAX_ITEMS = 48 +_SSE_RESPONSE_HEADERS = { + "Cache-Control": "no-cache", + "X-Accel-Buffering": "no", +} def _parse_site_list(sites: Optional[str]) -> Optional[List[int]]: @@ -180,11 +188,40 @@ def _merge_append_event(pending_event: Optional[dict], event: dict) -> dict: return merged_event +def _iter_replace_event_batches(event: dict) -> Iterator[dict]: + """ + 将超大的最终替换事件拆成有序批次，避免单个 SSE 消息承载全部完整对象。 + """ + items = event.get("items") + if ( + event.get("type") != "replace" + or not isinstance(items, list) + or len(items) <= _SSE_REPLACE_MAX_ITEMS + ): + yield event + return + + batch_count = (len(items) + _SSE_REPLACE_MAX_ITEMS - 1) // _SSE_REPLACE_MAX_ITEMS + for batch_index in range(batch_count): + start = batch_index * _SSE_REPLACE_MAX_ITEMS + batch_event = dict(event) + batch_event.update( + { + "type": "replace" if batch_index == 0 else "append", + "items": items[start:start + _SSE_REPLACE_MAX_ITEMS], + "replace_batch": True, + "batch_index": batch_index, + "batch_count": batch_count, + } + ) + yield batch_event + + async def _iter_batched_search_events( event_source: AsyncIterator[dict], ) -> AsyncIterator[dict]: """ - 对搜索流事件做轻量批处理，避免站点结果集中返回时产生过密 SSE。 + 对搜索流事件做轻量批处理，并在上游长时间静默时发送心跳。 """ iterator = event_source.__aiter__() pending_append_event: Optional[dict] = None @@ -195,13 +232,19 @@ async def _iter_batched_search_events( if next_event_task is None: next_event_task = asyncio.create_task(anext(iterator)) - timeout = _SSE_APPEND_FLUSH_INTERVAL if pending_append_event else None + timeout = ( + _SSE_APPEND_FLUSH_INTERVAL + if pending_append_event + else _SSE_HEARTBEAT_INTERVAL + ) done, _ = await asyncio.wait({next_event_task}, timeout=timeout) if not done: if pending_append_event: yield pending_append_event pending_append_event = None + else: + yield {"type": "heartbeat"} continue try: @@ -227,7 +270,8 @@ async def _iter_batched_search_events( yield pending_append_event pending_append_event = None - yield event + for batched_event in _iter_replace_event_batches(event): + yield batched_event finally: if next_event_task and not next_event_task.done(): next_event_task.cancel() @@ -239,13 +283,29 @@ async def _iter_batched_search_events( async def _stream_search_events(request: Request, event_source: AsyncIterator[dict]): """ - 输出搜索SSE事件 + 输出搜索 SSE 事件，并记录连接生命周期与传输规模。 """ locale = LocaleHelper.get_locale_from_request(request) + search_id = uuid4().hex[:12] + request_path = getattr(getattr(request, "url", None), "path", "unknown") + started_at = time.monotonic() + event_count = 0 + transmitted_bytes = 0 + last_event_type = "none" + last_stage = "none" + termination_reason = "source_exhausted" + logger.info(f"渐进式搜索流已建立，搜索ID：{search_id}，路径：{request_path}") try: has_sent_final_replace = False async for event in _iter_batched_search_events(event_source): + last_event_type = event.get("type") or "unknown" + last_stage = event.get("stage") or last_stage if await request.is_disconnected(): + termination_reason = "client_disconnected" + logger.warning( + f"渐进式搜索客户端已断开，搜索ID：{search_id}，路径：{request_path}，" + f"事件：{last_event_type}，阶段：{last_stage}" + ) break # 精确搜索会先发送 replace，再发送 done。done 再带整包 items 只会重复占用带宽和前端内存。 if event.get("type") == "replace" and event.get("items"): @@ -257,13 +317,36 @@ async def _stream_search_events(request: Request, event_source: AsyncIterator[di and event.get("items") ): event = {key: value for key, value in event.items() if key != "items"} - yield _sse_event(event, locale=locale) + payload = _sse_event(event, locale=locale) + event_count += 1 + transmitted_bytes += len(payload.encode("utf-8")) + if event.get("type") == "done": + termination_reason = "completed" + yield payload + except asyncio.CancelledError: + termination_reason = "cancelled" + logger.warning( + f"渐进式搜索流已取消，搜索ID：{search_id}，路径：{request_path}，" + f"事件：{last_event_type}，阶段：{last_stage}" + ) + raise except Exception as err: + termination_reason = "error" logger.error(f"渐进式搜索出错：{err}", exc_info=True) - yield _sse_event( + payload = _sse_event( {"type": "error", "success": False, "message": str(err)}, locale=locale, ) + event_count += 1 + transmitted_bytes += len(payload.encode("utf-8")) + yield payload + finally: + elapsed = time.monotonic() - started_at + logger.info( + f"渐进式搜索流结束，搜索ID：{search_id}，路径：{request_path}，" + f"状态：{termination_reason}，事件数：{event_count}，" + f"发送字节：{transmitted_bytes}，耗时：{elapsed:.2f}秒" + ) @router.get("/last", summary="查询搜索结果", response_model=List[schemas.Context]) @@ -319,6 +402,7 @@ async def search_by_id_stream( search_chain = SearchChain() async def event_source(): + """解析媒体身份并输出精确搜索流事件。""" search_params, message = await _resolve_media_search_params( mediaid=mediaid, media_type=media_type, @@ -341,7 +425,9 @@ async def search_by_id_stream( yield event return StreamingResponse( - _stream_search_events(request, event_source()), media_type="text/event-stream" + _stream_search_events(request, event_source()), + media_type="text/event-stream", + headers=_SSE_RESPONSE_HEADERS, ) @@ -401,7 +487,9 @@ async def search_by_title_stream( title=keyword, page=page, sites=_parse_site_list(sites), cache_local=True ) return StreamingResponse( - _stream_search_events(request, event_source), media_type="text/event-stream" + _stream_search_events(request, event_source), + media_type="text/event-stream", + headers=_SSE_RESPONSE_HEADERS, ) @@ -446,6 +534,7 @@ async def search_subtitle_by_title_stream( _iter_signed_subtitle_search_events(event_source), ), media_type="text/event-stream", + headers=_SSE_RESPONSE_HEADERS, ) @@ -558,6 +647,7 @@ async def search_subtitle_by_id_stream( _iter_signed_subtitle_search_events(event_source()), ), media_type="text/event-stream", + headers=_SSE_RESPONSE_HEADERS, ) diff --git a/docs/mcp-api.md b/docs/mcp-api.md index 135755ff7..a2020083b 100644 --- a/docs/mcp-api.md +++ b/docs/mcp-api.md @@ -142,6 +142,8 @@ FastAPI 异常响应保留 `detail` 字段，并在错误详情为文本时返 | GET | `/api/v1/search/last/context` | 获取上一次搜索结果及可复用搜索参数，`params.result_type` 为 `torrent` 或 `subtitle` | | POST | `/api/v1/search/recommend` | 获取 AI 推荐资源，请求体：`filtered_indices`、`check_only`、`force` | +渐进式搜索在无业务事件时每 15 秒发送 `{"type":"heartbeat"}`，客户端应将其仅用于连接保活。超过 48 条的最终 `replace` 会分批发送：首批 `type=replace`，后续批次 `type=append`，所有批次均带 `replace_batch=true`、从 0 开始的 `batch_index`、`batch_count` 和最终 `total_items`；客户端必须按顺序收齐后再原子替换结果。最终 `done` 在已发送 `replace` 后不重复携带 `items`。 + #### AniList 榜单 / 探索 AniList 榜单、探索、详情、人物和推荐接口优先通过 `anilist-chinese` 代理查询。代理不可用时自动回退 AniList 官方 GraphQL，并合并 `anilist-chinese` 每日数据集；媒体标题优先使用项目提供的中文标题，未提供中文标题时回退 AniList 原语言标题。 diff --git a/skills/moviepilot-api/SKILL.md b/skills/moviepilot-api/SKILL.md index 0e6f61568..66d83e017 100644 --- a/skills/moviepilot-api/SKILL.md +++ b/skills/moviepilot-api/SKILL.md @@ -1,6 +1,6 @@ --- name: moviepilot-api -version: 4 +version: 5 description: >- Use this skill when you need to call MoviePilot REST API endpoints directly with the bundled Python client. Covers MoviePilot HTTP endpoints across media @@ -169,6 +169,8 @@ AniList endpoints prefer the `anilist-chinese` proxy and fall back to official A | GET | `/api/v1/search/last/context` | Get latest search results with replayable params. `params.result_type` is `torrent` or `subtitle` | | POST | `/api/v1/search/recommend` | AI recommended resources. Body: `filtered_indices`, `check_only`, `force` | +Streaming search sends `{"type":"heartbeat"}` every 15 seconds without business events; use it only to keep the connection alive. Final `replace` payloads above 48 items are batched: the first event uses `type=replace`, later events use `type=append`, and every batch includes `replace_batch=true`, zero-based `batch_index`, `batch_count`, and final `total_items`. Collect all batches in order and replace the visible result atomically. After a `replace`, the final `done` event omits duplicate `items`. + ### Download (8 endpoints) | Method | Path | Description | diff --git a/tests/test_builtin_skill_boundaries.py b/tests/test_builtin_skill_boundaries.py index c724e5b22..ccbab9c0c 100644 --- a/tests/test_builtin_skill_boundaries.py +++ b/tests/test_builtin_skill_boundaries.py @@ -22,7 +22,7 @@ def test_modified_builtin_skills_have_incremented_versions() -> None: """本次修改过的内置技能必须递增版本，确保用户端同步更新。""" expected_versions = { "database-operation": "3", - "moviepilot-api": "4", + "moviepilot-api": "5", "moviepilot-cli": "6", "moviepilot-update": "3", } diff --git a/tests/test_search_stream.py b/tests/test_search_stream.py new file mode 100644 index 000000000..aaf22f6bc --- /dev/null +++ b/tests/test_search_stream.py @@ -0,0 +1,103 @@ +import asyncio +from types import SimpleNamespace + +import app.api.endpoints.search as search_endpoint + + +def test_large_replace_event_is_split_into_ordered_batches(monkeypatch): + """超大最终结果应拆成首个 replace 和后续 append 批次。""" + monkeypatch.setattr(search_endpoint, "_SSE_REPLACE_MAX_ITEMS", 2) + source_event = { + "type": "replace", + "stage": "filtered", + "items": [1, 2, 3, 4, 5], + "total_items": 5, + } + + async def _collect_events(): + """通过完整批处理适配器收集拆分后的最终结果。""" + + async def _source(): + """输出一个超大最终替换事件。""" + yield source_event + + return [ + event + async for event in search_endpoint._iter_batched_search_events(_source()) + ] + + events = asyncio.run(_collect_events()) + + assert [event["type"] for event in events] == ["replace", "append", "append"] + assert [event["batch_index"] for event in events] == [0, 1, 2] + assert all(event["batch_count"] == 3 for event in events) + assert all(event["replace_batch"] for event in events) + assert [item for event in events for item in event["items"]] == source_event["items"] + + +def test_small_replace_event_keeps_original_protocol(monkeypatch): + """小型最终结果应保持单个 replace 事件，兼容现有客户端。""" + monkeypatch.setattr(search_endpoint, "_SSE_REPLACE_MAX_ITEMS", 2) + source_event = { + "type": "replace", + "items": [1, 2], + "total_items": 2, + } + + assert list(search_endpoint._iter_replace_event_batches(source_event)) == [source_event] + + +def test_batched_search_events_emit_heartbeat_while_source_is_idle(monkeypatch): + """上游长时间无业务事件时应持续输出心跳，避免连接被空闲超时关闭。""" + monkeypatch.setattr(search_endpoint, "_SSE_HEARTBEAT_INTERVAL", 0.01) + + async def _read_heartbeat(): + """读取首个心跳并关闭仍在等待的上游迭代器。""" + blocker = asyncio.Event() + + async def _source(): + """模拟长时间处于过滤匹配阶段的事件源。""" + await blocker.wait() + yield {"type": "done"} + + events = search_endpoint._iter_batched_search_events(_source()) + try: + return await asyncio.wait_for(anext(events), timeout=0.5) + finally: + await events.aclose() + + assert asyncio.run(_read_heartbeat()) == {"type": "heartbeat"} + + +def test_search_stream_response_disables_proxy_buffering(monkeypatch): + """搜索 SSE 响应应显式禁用缓存和 Nginx 代理缓冲。""" + + class FakeSearchChain: + """提供无需外部依赖的空搜索流。""" + + def async_search_by_title_stream(self, **_kwargs): + """返回立即完成的搜索流。""" + + async def _source(): + """输出一个完成事件。""" + yield {"type": "done", "stage": "done", "total_items": 0} + + return _source() + + monkeypatch.setattr(search_endpoint, "SearchChain", FakeSearchChain) + + async def _never_disconnected(): + """模拟始终在线的 SSE 客户端。""" + return False + + request = SimpleNamespace( + url=SimpleNamespace(path="/api/v1/search/title/stream"), + is_disconnected=_never_disconnected, + ) + + response = asyncio.run( + search_endpoint.search_by_title_stream(request=request, keyword="Demo", _=None) + ) + + assert response.headers["cache-control"] == "no-cache" + assert response.headers["x-accel-buffering"] == "no"