diff --git a/README.md b/README.md index 423ed21..866ac60 100644 --- a/README.md +++ b/README.md @@ -264,6 +264,93 @@ async def upload(file: UploadFile = File(...)): > **约定**：插件 `__init__.py` 中导入 `router` 实例（`from .routes import router`），网关通过 `getattr(module, "router", None)` 自动发现。 +### 流式端点（SSE） + +基于 Zenoh Query/Reply 原生"多回复"实现流式输出：Worker 的 `@rpc_method` 返回生成器时， +框架逐块 `query.reply()` 推送（见 `utils/worker` README「流式端点」），网关以 SSE +（`text/event-stream`）逐块转发给 HTTP 客户端。**无额外端口、无额外进程、无新依赖**； +未声明 `stream` 的端点行为与之前完全一致。 + +#### TOML 配置 + +在方法节中声明 `stream = true`（示例 `config/workers/example/stream_chat.toml`）： + +```toml +[POST] +summary = "模拟 AI 流式对话（SSE，逐 token 输出）" +auth_required = false +stream = true +timeout = 300.0 # 流总寿命上限（秒，zenoh query deadline 硬上限） +stream_idle_timeout = 10.0 # 网关侧连续无数据块的最大间隔（秒） +params = ["prompt: str", "tokens: int"] +``` + +| 字段 | 默认值 | 说明 | +| --- | --- | --- | +| `stream` | `false` | 流式端点：以 SSE（`text/event-stream`）输出 | +| `timeout` | `3.0` | 流总寿命上限（秒）。**流式端点必须配大**（如 300s），到期后 zenoh 丢弃后续回复 | +| `stream_idle_timeout` | `10.0` | 空闲超时（秒）：超过该间隔无新数据块则网关终止流 | + +> `timeout`（总寿命硬上限）与 `stream_idle_timeout`（空闲超时）是两回事，流式场景两个都要配。 + +#### 请求与响应 + +- 请求与普通端点一致（GET 走 query/path 参数，POST/PUT/PATCH 走 body 模型），鉴权依赖同样生效 +- 响应为 `text/event-stream`，每块一个 `data: {json}\n\n` 事件（SSE 标准分帧） +- 响应头：`Cache-Control: no-cache`、`X-Accel-Buffering: no`（配合反代关闭缓冲） +- OpenAPI 不声明流式端点的响应 schema（逐块 JSON dict，无法用单一模型描述） + +```bash +curl -N -X POST localhost:8000/api/v1/calc/stream_chat \ + -H 'Content-Type: application/json' -d '{"prompt":"你好","tokens":3}' +# data: {"token":"你好·1","index":0} +# data: {"token":"你好·2","index":1} +# data: {"token":"你好·3","index":2} +# data: {"done":true} +``` + +#### 错误语义 + +| 场景 | 表现 | +| --- | --- | +| 无可用 worker（首个数据块前） | HTTP 503（与普通端点一致） | +| Worker 中途异常 / 流式端点初始化失败 | SSE error 事件 `data: {"error":"..."}` 收尾（HTTP 200 已提交，无法再改状态码） | +| 空闲超时 | SSE error 事件 `data: {"error":"流式空闲超时 ..."}` 收尾 | +| 客户端断开 | 网关终止迭代，Worker 侧 reply 失败自动停流，不空转 | + +#### 部署注意 + +生产环境经 nginx 反代时需关闭缓冲，否则 SSE 会被整段缓冲、失去逐块即时性： + +```nginx +location /api/v1/ { + proxy_buffering off; +} +``` + +#### 本地联调前注意 + +流式能力在 `chongming-worker` 框架（`utils/worker/`，独立仓库）。**本地未 push 的框架改动 +不会自动生效**——venv 里安装的是 `uv sync` 从 git 源拉取的旧版副本，旧版 `@rpc_method` +不认 `stream` 参数，会报 `TypeError: rpc_method() got an unexpected keyword argument 'stream'` +并连锁导致契约校验失败（启动终止）。本地联调时先以可编辑方式装入 venv： + +```bash +uv pip install -e utils/worker +``` + +正式发布仍走「开发与发布检查清单」流程（工具包 push → `uv lock --upgrade-package +chongming-worker` → `uv sync`，可编辑安装会被锁文件版本替换）。 + +#### 测试 + +``` +uv run python test_stream_worker.py # Worker 框架流式回复单测 +uv run python test_stream_gateway.py # 网关 call_stream 单测 +uv run python test_stream_e2e_zenoh.py # 真实 Zenoh 会话端到端（peer 模式，无需 router） +uv run python test_stream_route.py # 路由层 SSE 输出与兼容性回归（非流式端点不变） +``` + --- ## 进程插件与混合部署 @@ -450,6 +537,9 @@ python publisher.py example add '{"a":3,"b":5}' # 调用 add 5. 重启服务 ``` +流式端点（SSE）在此基础上额外两步：TOML 方法节声明 `stream = true` 并配大 `timeout`； +`@rpc_method(..., stream=True)` 函数返回生成器（yield 逐块），详见上文「流式端点（SSE）」。 + ### 常用命令 ```bash @@ -566,6 +656,8 @@ uv pip install -e utils/ # 安装/更新工具包子包 | `internal` | `<操作>.toml` | `true` 则仅限内部 RPC，不注册 HTTP 路由 | | `cacheable` | `<操作>.toml` | 端点是否可缓存 | | `summary` | `<操作>.toml` | 端点描述（用于 OpenAPI / 日志） | +| `stream` | `<操作>.toml` | `true` 则流式输出（SSE），需同时配大 `timeout` | +| `stream_idle_timeout`| `<操作>.toml` | 流式空闲超时（秒，默认 10），超时无新数据块则终止流 | --- diff --git a/src/chongming/app/core/route_generator.py b/src/chongming/app/core/route_generator.py index b3ee88f..bc70de0 100644 --- a/src/chongming/app/core/route_generator.py +++ b/src/chongming/app/core/route_generator.py @@ -5,15 +5,17 @@ 为每个端点生成带有正确 Pydantic 类型注解的处理函数，确保 OpenAPI 文档 正确反映参数类型、响应结构，并注册到 components/schemas。 """ -from typing import Any, Dict, List, Optional +from typing import Any, AsyncGenerator, Dict, List, Optional from inspect import Parameter, Signature +import orjson from fastapi import FastAPI, HTTPException, APIRouter, Query, Path as FastAPIPath, Request, Depends +from fastapi.responses import StreamingResponse from pydantic import create_model, Field, BaseModel from chongming_worker.route_config import WORKER_CONFIG_DIR, load_worker_routes_from_dir -from .worker_manager import WorkerManager +from .worker_manager import WorkerManager, StreamIdleTimeoutError from .logger import get_logger from .auth import store_user_to_request @@ -92,6 +94,10 @@ class RouteConfig: self.cacheable: bool = data.get("cacheable", False) self.auth_required: bool = data.get("auth_required", True) self.internal: bool = data.get("internal", False) + # 流式端点：以 SSE（text/event-stream）逐块输出 + self.stream: bool = data.get("stream", False) + # 流式空闲超时（秒）：连续未收到数据块的最大间隔（timeout 为流总寿命上限） + self.stream_idle_timeout: float = data.get("stream_idle_timeout", 10.0) def parse_params(self) -> Dict[str, str]: """解析 params 列表为 {name: type} 字典""" @@ -353,6 +359,116 @@ def _make_query_endpoint(worker_manager: WorkerManager, rc: RouteConfig, worker_ return endpoint +# 流式端点响应头：禁止缓存；通知反向代理（nginx）关闭缓冲保证 SSE 即时性 +_STREAM_HEADERS = { + "Cache-Control": "no-cache", + "X-Accel-Buffering": "no", +} + + +def _make_stream_endpoint( + worker_manager: WorkerManager, + rc: RouteConfig, + worker_type: str, + path_params: List[str] = (), + query_params: Optional[Dict[str, str]] = None, + RequestModel: Optional[type[BaseModel]] = None, +): + """ + 创建流式端点（TOML 中 stream = true）：以 SSE（text/event-stream）逐块输出。 + + 请求参数形态与普通端点一致（GET 走 query/path 参数，POST/PUT/PATCH 走 body 模型）。 + 错误语义: + - 初始错误（无可用 worker 等）-> HTTP 503 + - 流中错误（Worker 返回 {"error": ...} 收尾块 / 空闲超时）-> SSE error 事件 + （此时 HTTP 200 已提交，无法改状态码） + 不校验响应模型：流式输出是逐块 JSON dict。 + """ + path = rc.path + subject = rc.subject + http_method = rc.method + timeout = rc.timeout + idle_timeout = rc.stream_idle_timeout + param_names = list((query_params or {}).keys()) + + def build_payload(request: Request, body=None, kwargs=None) -> Dict[str, Any]: + """构造信封协议 payload（与普通端点一致）""" + meta = {"method": http_method} + user_info = getattr(request.state, "user", None) + if user_info: + meta["user_info"] = user_info + if body is not None: + data = body.model_dump() + else: + data = {k: kwargs[k] for k in param_names} + return {"meta": meta, "data": data} + + def _sse_frame(data: Dict[str, Any]) -> str: + return f"data: {orjson.dumps(data).decode('utf-8')}\n\n" + + async def _sse_body(gen) -> AsyncGenerator[str, None]: + """SSE 分帧；流中异常转为 SSE error 事件（HTTP 200 已提交）""" + try: + async for chunk in gen: + yield f"data: {chunk.decode('utf-8')}\n\n" + except StreamIdleTimeoutError as e: + logger.warning("路由 %s 流式空闲超时: %s", path, e) + yield _sse_frame({"error": str(e)}) + except Exception as e: + logger.error("路由 %s 流式中断: %s", path, e) + yield _sse_frame({"error": "内部服务器错误"}) + + async def _start_stream(payload: Dict[str, Any]) -> StreamingResponse: + try: + gen = await worker_manager.call_stream( + worker_type=worker_type, + subject=subject, + payload=payload, + timeout=timeout, + idle_timeout=idle_timeout, + ) + except RuntimeError as e: + logger.warning("路由 %s [%s] 流式调用失败: %s", path, worker_type, e) + raise HTTPException(status_code=503, detail=str(e)) + return StreamingResponse( + _sse_body(gen), + media_type="text/event-stream", + headers=_STREAM_HEADERS, + ) + + if RequestModel is not None: + async def endpoint(request: Request, body: RequestModel) -> StreamingResponse: # type: ignore + return await _start_stream(build_payload(request, body=body)) + return endpoint + + # GET/DELETE 形态 + if not param_names: + async def _simple_endpoint(request: Request) -> StreamingResponse: + return await _start_stream(build_payload(request, kwargs=None)) + return _simple_endpoint + + async def endpoint(request: Request, **kwargs: Any) -> StreamingResponse: + return await _start_stream(build_payload(request, kwargs=kwargs)) + + # 附加类型注解与签名（与 _make_query_endpoint 一致），保证 OpenAPI 参数文档正确 + params: List[Parameter] = [ + Parameter("request", Parameter.POSITIONAL_OR_KEYWORD, annotation=Request) + ] + annotations: Dict[str, Any] = {} + for name, type_str in (query_params or {}).items(): + py_type = _parse_python_type(type_str) + annotations[name] = py_type + if name in path_params: + params.append(Parameter(name, Parameter.KEYWORD_ONLY, default=FastAPIPath(...), annotation=py_type)) + else: + params.append(Parameter(name, Parameter.KEYWORD_ONLY, default=Query(...), annotation=py_type)) + + endpoint.__annotations__ = annotations + endpoint.__signature__ = Signature(params, return_annotation=StreamingResponse) # type: ignore + + return endpoint + + def _register_single_route( router: APIRouter, worker_manager: WorkerManager, @@ -391,16 +507,41 @@ def _register_single_route( request_model_name = f"{worker_type}_{clean_path}_{method_tag}_Request" # 生成响应模型 - ResponseModel = generate_response_model( - rc.response_model, response_model_name - ) - - if rc.method in ("POST", "PUT", "PATCH"): + if rc.stream: + # 流式端点：逐块 SSE 输出，不校验响应模型（OpenAPI 不声明响应 schema） + ResponseModel = None + if rc.method in ("POST", "PUT", "PATCH"): + # body 形态：RequestModel 传给端点做请求体校验 + RequestModel = generate_request_model( + param_spec, request_model_name + ) + endpoint = _make_stream_endpoint( + worker_manager, rc, worker_type, path_params, query_params, RequestModel + ) + else: + # GET/DELETE 形态：query/path 参数走签名注解； + # RequestModel 仅注册到 extra_schemas 供 OpenAPI 展示（与普通 GET 一致） + if param_spec: + RequestModel = generate_request_model( + param_spec, request_model_name + ) + if extra_schemas is not None and request_model_name not in extra_schemas: + extra_schemas[request_model_name] = RequestModel + endpoint = _make_stream_endpoint( + worker_manager, rc, worker_type, path_params, query_params, None + ) + elif rc.method in ("POST", "PUT", "PATCH"): + ResponseModel = generate_response_model( + rc.response_model, response_model_name + ) RequestModel = generate_request_model( param_spec, request_model_name ) endpoint = _make_body_endpoint(worker_manager, rc, worker_type, RequestModel) else: + ResponseModel = generate_response_model( + rc.response_model, response_model_name + ) # 对 GET 端点也生成 Request 模型并注册到 extra_schemas if param_spec: RequestModel = generate_request_model( diff --git a/src/chongming/app/core/worker_manager.py b/src/chongming/app/core/worker_manager.py index 4a271ec..3001b09 100644 --- a/src/chongming/app/core/worker_manager.py +++ b/src/chongming/app/core/worker_manager.py @@ -2,7 +2,7 @@ import asyncio import random import json import orjson -from typing import Any, Dict, List, Optional +from typing import Any, AsyncGenerator, Dict, List, Optional import zenoh from chongming_cache import Cache from .logger import get_logger @@ -12,6 +12,22 @@ logger = get_logger("worker_manager") _SENTINEL = object() # 哨兵值，用于标记"No reply received" +class StreamIdleTimeoutError(TimeoutError): + """流式调用空闲超时：在 idle_timeout 秒内未收到任何数据块""" + + +def _next_reply(iterator): + """线程安全地获取迭代器的下一个元素。 + + 用 for 循环 + 哨兵代替直接 next()，避免 StopIteration 在 + asyncio.to_thread 中传播的问题（与 call() 的 _get_first 同模式）。 + 迭代耗尽时返回 _SENTINEL。 + """ + for item in iterator: + return item + return _SENTINEL + + class WorkerManager: """ 统一 Worker 管理器 @@ -130,6 +146,114 @@ class WorkerManager: return False return True + async def _raise_no_worker(self, worker_type: str) -> None: + """ + 无可用 Worker 时的统一诊断：列出当前已知的所有 worker 类型并抛 RuntimeError。 + + 供 call() / call_stream() 共用，保证两类调用错误信息一致。 + """ + async with self._lock: + all_known = list(self._workers_cache.keys()) + local_known = list(self._local_workers.keys()) + all_types = sorted(set(all_known + local_known)) + msg = ( + f"类型 '{worker_type}' 没有可用 Worker。" + f"当前已知的 Worker 类型: {all_types if all_types else '(无)'}。" + f"请检查: 1) config/workers/config.toml 是否配置了该类型 " + f"2) 插件初始化是否成功（查看启动日志中的 '加载 Worker 插件失败' 错误）" + f"3) Worker 进程是否正常运行" + ) + logger.warning(msg) + raise RuntimeError(msg) + + async def call_stream( + self, + worker_type: str, + subject: str, + payload: Dict[str, Any], + timeout: float = 300.0, + idle_timeout: float = 10.0, + ) -> AsyncGenerator[bytes, None]: + """ + 流式调用：消费 Worker 的多回复流，逐块产出 JSON 字节。 + + 与 call() 的区别: + - 不合并回复（ConsolidationMode.NONE），逐块 yield，实现流式传输 + - 不重试（流一旦开始输出无法安全重放，中途失败只能重建） + - 超时模型: timeout = 流总寿命上限（zenoh query deadline，硬上限， + 流式端点须配大如 300s）；idle_timeout = 连续无数据块的最大间隔 + - 错误约定: 块内含 "error" 键（Worker 收尾错误块）或 reply.err + （zenoh 层错误）均为流的终止信号，转发后结束 + + Args: + worker_type: Worker 类型标识（如 "example"） + subject: 调用的主题（如 "stream_chat"） + payload: 请求参数（信封格式，由路由层构造） + timeout: 流总寿命上限（秒），需大于业务最长生成时长 + idle_timeout: 空闲超时（秒），超时抛 StreamIdleTimeoutError + + Yields: + bytes: 每个回复块的 JSON 字节 + + Raises: + RuntimeError: 无可用 worker + StreamIdleTimeoutError: 空闲超时 + """ + # 急切解析 worker：无可用时立即抛错，让路由层映射 HTTP 503 + # （异步生成器函数体在首次迭代才执行，故这里用普通 async 函数返回内部生成器） + worker_expr = await self.get_worker(worker_type) + if not worker_expr: + await self._raise_no_worker(worker_type) + + full_key = f"{worker_expr}/{subject}" + reply_iter = self.session.get( + full_key, + payload=orjson.dumps(payload), + consolidation=zenoh.ConsolidationMode.NONE, + timeout=timeout, + ) + + async def _stream(): + # 逐块消费；zenoh reply 迭代器是同步阻塞的，必须放线程执行， + # 避免阻塞事件循环（与 call() 的 _get_first 同模式） + while True: + try: + reply = await asyncio.wait_for( + asyncio.to_thread(_next_reply, reply_iter), + timeout=idle_timeout, + ) + except asyncio.TimeoutError: + raise StreamIdleTimeoutError( + f"Worker {worker_expr} 流式空闲超时 ({idle_timeout}s 无数据)" + ) from None + + if reply is _SENTINEL: + break + + if reply.ok: # type: ignore + data = bytes(reply.ok.payload) if reply.ok.payload else b"" # type: ignore + if not data: + continue + decoded = orjson.loads(data) + if isinstance(decoded, dict) and decoded.get("error"): + # Worker 错误收尾块：转发后终止流 + yield data + return + yield data + else: + # zenoh 层错误回复（如 queryable 异常）：转发后终止流 + error_payload = reply.err.payload # type: ignore + error_msg = str(error_payload) if error_payload else "unknown error" + logger.warning( + f"Worker {worker_expr} 流式返回错误: {error_msg}" + ) + yield orjson.dumps( + {"error": f"Worker {worker_expr} 返回错误: {error_msg}"} + ) + return + + return _stream() + async def call( self, worker_type: str, @@ -166,20 +290,7 @@ class WorkerManager: # 2. 获取可用 worker worker_expr = await self.get_worker(worker_type) if not worker_expr: - # 给出诊断信息：列出当前已知的所有 worker 类型 - async with self._lock: - all_known = list(self._workers_cache.keys()) - local_known = list(self._local_workers.keys()) - all_types = sorted(set(all_known + local_known)) - msg = ( - f"类型 '{worker_type}' 没有可用 Worker。" - f"当前已知的 Worker 类型: {all_types if all_types else '(无)'}。" - f"请检查: 1) config/workers/config.toml 是否配置了该类型 " - f"2) 插件初始化是否成功（查看启动日志中的 '加载 Worker 插件失败' 错误）" - f"3) Worker 进程是否正常运行" - ) - logger.warning(msg) - raise RuntimeError(msg) + await self._raise_no_worker(worker_type) full_key = f"{worker_expr}/{subject}" payload_bytes = orjson.dumps(payload) diff --git a/test_stream_e2e_zenoh.py b/test_stream_e2e_zenoh.py new file mode 100644 index 0000000..f8014c4 --- /dev/null +++ b/test_stream_e2e_zenoh.py @@ -0,0 +1,180 @@ +""" +真实 Zenoh 会话下的流式端到端测试（peer 模式，无需 router） + +单进程内起两个 Zenoh 会话（listen + connect），跑通完整生产链路: + @rpc_method(stream=True) 注册 -> create_worker 声明 Queryable + -> WorkerManager.call_stream 消费多回复 -> 逐块到达 + +验证: + 1. 流式方法：3+1 块按序到达（真实 Zenoh 多回复 + ConsolidationMode.NONE） + 2. 流式中途异常：{"error": ...} 收尾块转发后终止 + 3. 非流式方法：仍是单回复（兼容性回归） + 4. generate_route_configs 输出 stream 标记、非流式不含 + 5. decorators 对生成器函数的调用不消费生成器 + +运行: uv run python test_stream_e2e_zenoh.py +""" +import asyncio +import importlib.util +import logging +import sys +import types +from pathlib import Path + +import orjson +import zenoh + +REPO_ROOT = Path(__file__).resolve().parent +sys.path.insert(0, str(REPO_ROOT / "utils" / "worker" / "src")) +sys.path.insert(0, str(REPO_ROOT / "src")) + +# ---- 网关依赖隔离（同 test_stream_gateway）：桩掉 chongming 包链与 chongming_cache ---- +_pkg = types.ModuleType("chongming") +_pkg.__path__ = [] +_app = types.ModuleType("chongming.app") +_app.__path__ = [] +_core = types.ModuleType("chongming.app.core") +_core.__path__ = [] +_logger_stub = types.ModuleType("chongming.app.core.logger") +_logger_stub.get_logger = lambda name: logging.getLogger(name) +_cache_stub = types.ModuleType("chongming_cache") +_cache_stub.Cache = object +for mod in (_pkg, _app, _core, _logger_stub, _cache_stub): + sys.modules[mod.__name__] = mod + +_wm_path = REPO_ROOT / "src" / "chongming" / "app" / "core" / "worker_manager.py" +_spec = importlib.util.spec_from_file_location( + "chongming.app.core.worker_manager", _wm_path +) +_wm_mod = importlib.util.module_from_spec(_spec) +sys.modules["chongming.app.core.worker_manager"] = _wm_mod +_spec.loader.exec_module(_wm_mod) # type: ignore[union-attr] + +WorkerManager = _wm_mod.WorkerManager +StreamIdleTimeoutError = _wm_mod.StreamIdleTimeoutError + +from chongming_worker import clear_registry, rpc_method # noqa: E402 +from chongming_worker.decorators import generate_route_configs # noqa: E402 +from chongming_worker.runtime import create_worker # noqa: E402 + +failures = [] + + +def check(name, cond, detail=""): + status = "PASS" if cond else "FAIL" + print(f"[{status}] {name}" + (f" — {detail}" if detail else "")) + if not cond: + failures.append(name) + + +LISTEN_PORT = 7447 +LISTEN = f"tcp/127.0.0.1:{LISTEN_PORT}" + + +def open_session(listen=None): + """peer 模式会话：worker 侧 listen，网关侧 connect（显式端点，禁用组播）""" + c = zenoh.Config() + c.insert_json5("scouting/multicast/enabled", "false") + if listen: + c.insert_json5("listen/endpoints", f'["{listen}"]') + else: + c.insert_json5("connect/endpoints", f'["{LISTEN}"]') + return zenoh.open(c) + + +async def run_checks(): + clear_registry() + + # ---- 流式端点（返回异步生成器）---- + @rpc_method(subject="streamtest.echo", method="POST", stream=True) + async def echo_stream(prompt: str = "hi", tokens: int = 3): + for i in range(tokens): + await asyncio.sleep(0.01) + yield {"token": f"{prompt}-{i}"} + yield {"done": True} + + # ---- 流式端点：中途异常（框架应发 {"error": ...} 收尾块）---- + @rpc_method(subject="streamtest.fail", method="POST", stream=True) + async def fail_stream(): + yield {"token": "a"} + raise ValueError("mid-boom") + + # ---- 非流式端点（兼容性对照）---- + @rpc_method(subject="streamtest.add", method="GET") + async def add(a: float, b: float): + return a + b + + # 1. 配置生成：流式标记 + configs = generate_route_configs("streamtest") + stream_cfg = next(c for c in configs if c["subject"] == "streamtest.echo") + check("generate_route_configs 输出 stream 标记", stream_cfg.get("stream") is True) + add_cfg = next(c for c in configs if c["subject"] == "streamtest.add") + check("非流式配置不含 stream 标记", "stream" not in add_cfg) + + # 2. 真实 Zenoh 会话（peer 模式，无 router） + session_worker = open_session(listen=LISTEN) + session_gw = open_session() + + # 3. 启动真实 Worker（create_worker 声明 Queryable） + worker_task = asyncio.create_task( + create_worker(session_worker, "streamtest", "w1") + ) + await asyncio.sleep(0.5) # 等待 queryable 在 peer 间可见 + + try: + # 4. 网关 WorkerManager.call_stream 真实消费流式端点 + wm = WorkerManager(session_gw) + wm.register_worker("streamtest", "w1") + + async def _consume(subject, data, idle_timeout=2.0): + gen = await wm.call_stream( + "streamtest", subject, + {"meta": {"method": "POST"}, "data": data}, + timeout=10.0, idle_timeout=idle_timeout, + ) + return [orjson.loads(b) async for b in gen] + + chunks = await _consume("streamtest.echo", {"prompt": "hi", "tokens": 3}) + check( + "流式 3+1 块按序到达（真实多回复）", + chunks == [ + {"token": "hi-0"}, {"token": "hi-1"}, {"token": "hi-2"}, {"done": True}, + ], + f"chunks={chunks}", + ) + + # 5. 流式中途异常：错误块收尾 + chunks = await _consume("streamtest.fail", {}) + check( + "流式中途异常 error 块收尾", + chunks == [{"token": "a"}, {"error": "mid-boom"}], + f"chunks={chunks}", + ) + + # 6. 非流式兼容：call() 仍单回复 + result = await wm.call( + "streamtest", "streamtest.add", + {"meta": {"method": "GET"}, "data": {"a": 2, "b": 3}}, + timeout=10.0, + ) + check("非流式 call() 正常单回复", result == 5.0, f"result={result}") + + # 7. 负载均衡路径兼容：liveliness 发现的 worker 也能流式（同 call） + worker_expr = await wm.get_worker("streamtest") + check("get_worker 仍可解析本地 worker", worker_expr is not None) + + finally: + worker_task.cancel() + await asyncio.gather(worker_task, return_exceptions=True) + session_gw.close() + session_worker.close() + + +def run(): + asyncio.run(run_checks()) + print(f"\n结果: {len(failures)} 个失败") + return 1 if failures else 0 + + +if __name__ == "__main__": + sys.exit(run()) diff --git a/test_stream_gateway.py b/test_stream_gateway.py new file mode 100644 index 0000000..f55da12 --- /dev/null +++ b/test_stream_gateway.py @@ -0,0 +1,246 @@ +""" +网关 call_stream 单测（依赖隔离：桩掉 chongming 包链与 chongming_cache） + +直接加载 src/chongming/app/core/worker_manager.py 模块（避开 chongming.app +整链导入与 config.toml 依赖），验证: + 1. 多块顺序产出（ConsolidationMode.NONE 逐块消费） + 2. Worker 错误收尾块 {"error": ...}：转发后终止流 + 3. zenoh 错误回复（reply.err）：转发错误块后终止流 + 4. 空闲超时：抛 StreamIdleTimeoutError + 5. 无可用 worker：抛 RuntimeError + 6. 兼容性回归：call()（非流式路径）仍只取第一个回复 + +运行: uv run python test_stream_gateway.py +""" +import asyncio +import importlib.util +import logging +import sys +import time +import types +from pathlib import Path +from types import SimpleNamespace + +import orjson +import zenoh + +REPO_ROOT = Path(__file__).resolve().parent + +# ---- 依赖隔离：桩掉网关包链（chongming.app 顶层导入 fastapi/uvicorn/config）---- +_pkg = types.ModuleType("chongming") +_pkg.__path__ = [] +_app = types.ModuleType("chongming.app") +_app.__path__ = [] +_core = types.ModuleType("chongming.app.core") +_core.__path__ = [] +_logger_stub = types.ModuleType("chongming.app.core.logger") +_logger_stub.get_logger = lambda name: logging.getLogger(name) +_cache_stub = types.ModuleType("chongming_cache") +_cache_stub.Cache = object # WorkerManager 仅作类型标注使用 +for mod in (_pkg, _app, _core, _logger_stub, _cache_stub): + sys.modules[mod.__name__] = mod + +_wm_path = REPO_ROOT / "src" / "chongming" / "app" / "core" / "worker_manager.py" +_spec = importlib.util.spec_from_file_location( + "chongming.app.core.worker_manager", _wm_path +) +_wm_mod = importlib.util.module_from_spec(_spec) +sys.modules["chongming.app.core.worker_manager"] = _wm_mod +_spec.loader.exec_module(_wm_mod) # type: ignore[union-attr] + +WorkerManager = _wm_mod.WorkerManager +StreamIdleTimeoutError = _wm_mod.StreamIdleTimeoutError + +failures = [] + + +def check(name, cond, detail=""): + status = "PASS" if cond else "FAIL" + print(f"[{status}] {name}" + (f" — {detail}" if detail else "")) + if not cond: + failures.append(name) + + +class FakeReply: + """模拟 zenoh.Reply：ok / err 二选一""" + + def __init__(self, payload=None, error=None): + if error is None: + self.ok = SimpleNamespace(payload=payload) + self.err = None + else: + self.ok = None + self.err = SimpleNamespace(payload=error) + + +class FakeSession: + """模拟 zenoh.Session.get()：返回 reply 迭代器（可迭代对象）""" + + def __init__(self, reply_source): + self._reply_source = reply_source + self.get_calls = [] + + def get(self, key_expr, payload, consolidation=None, timeout=None): + self.get_calls.append((str(key_expr), payload, consolidation, timeout)) + return iter(self._reply_source) + + +def make_manager(replies, worker_id="worker-0"): + wm = WorkerManager(FakeSession(replies)) + wm.register_worker("example", worker_id) + return wm + + +def test_idle_timeout(): + # 模拟阻塞式 reply 迭代器：每个 __next__ 阻塞 0.3s（首个块 0.3s 后才到达）， + # 而空闲超时仅 0.05s -> 应抛 StreamIdleTimeoutError + class BlockingReplyIterator: + def __init__(self): + self._done = False + + def __iter__(self): + return self + + def __next__(self): + time.sleep(0.3) + if self._done: + raise StopIteration + self._done = True + return FakeReply(payload=orjson.dumps({"token": "slow"})) + + wm = WorkerManager(FakeSession(BlockingReplyIterator())) + + +def collect_stream(wm, idle_timeout=1.0): + """消费 call_stream 全部块并返回解码后的列表""" + + async def _run(): + gen = await wm.call_stream( + "example", "stream_chat", {"meta": {"method": "POST"}, "data": {}}, + timeout=5.0, idle_timeout=idle_timeout, + ) + return [orjson.loads(b) async for b in gen] + + return asyncio.run(_run()) + + +def test_multiple_chunks(): + replies = [ + FakeReply(payload=orjson.dumps({"token": "a"})), + FakeReply(payload=orjson.dumps({"token": "b"})), + ] + wm = make_manager(replies) + chunks = collect_stream(wm) + check("多块顺序产出", chunks == [{"token": "a"}, {"token": "b"}]) + check( + "使用 ConsolidationMode.NONE", + wm.session.get_calls[0][2] is zenoh.ConsolidationMode.NONE, + ) + check("key_expr 正确", wm.session.get_calls[0][0] == "worker/example/worker-0/stream_chat") + + +def test_error_chunk_terminates(): + replies = [ + FakeReply(payload=orjson.dumps({"token": "a"})), + FakeReply(payload=orjson.dumps({"error": "boom"})), + FakeReply(payload=orjson.dumps({"token": "c"})), # 不应被消费 + ] + wm = make_manager(replies) + chunks = collect_stream(wm) + check("错误块转发后终止", chunks == [{"token": "a"}, {"error": "boom"}]) + + +def test_err_reply_terminates(): + replies = [ + FakeReply(payload=orjson.dumps({"token": "a"})), + FakeReply(error="oops"), + ] + wm = make_manager(replies) + chunks = collect_stream(wm) + check( + "reply.err 转发错误块后终止", + chunks == [{"token": "a"}, {"error": "Worker worker/example/worker-0 返回错误: oops"}], + ) + + +def test_idle_timeout(): + # 模拟阻塞式 reply 迭代器：每个 __next__ 阻塞 0.3s（首个块 0.3s 后才到达）， + # 而空闲超时仅 0.05s -> 应抛 StreamIdleTimeoutError + class BlockingReplyIterator: + def __init__(self): + self._done = False + + def __iter__(self): + return self + + def __next__(self): + time.sleep(0.3) + if self._done: + raise StopIteration + self._done = True + return FakeReply(payload=orjson.dumps({"token": "slow"})) + + wm = WorkerManager(FakeSession(BlockingReplyIterator())) + wm.register_worker("example", "worker-0") + + async def _run(): + gen = await wm.call_stream( + "example", "stream_chat", {}, timeout=5.0, idle_timeout=0.05 + ) + async for _ in gen: + pass + + try: + asyncio.run(_run()) + check("空闲超时抛 StreamIdleTimeoutError", False) + except StreamIdleTimeoutError as e: + check("空闲超时抛 StreamIdleTimeoutError", True, str(e)) + + +def test_no_worker(): + wm = WorkerManager(FakeSession([])) # 未注册任何 worker + + async def _run(): + gen = await wm.call_stream("missing", "x", {}) + async for _ in gen: + pass + + try: + asyncio.run(_run()) + check("无 worker 抛 RuntimeError", False) + except RuntimeError as e: + check("无 worker 抛 RuntimeError", True, str(e)[:30]) + + +def test_call_compat_first_reply_only(): + """兼容性回归：非流式 call() 仍只取第一个回复，行为不变""" + replies = [ + FakeReply(payload=orjson.dumps({"result": 1})), + FakeReply(payload=orjson.dumps({"result": 2})), + ] + wm = make_manager(replies) + + async def _run(): + return await wm.call( + "example", "stream_chat", {"meta": {"method": "POST"}, "data": {}}, + timeout=5.0, + ) + + result = asyncio.run(_run()) + check("call() 仍只取第一个回复（非流式兼容）", result == {"result": 1}) + + +def run(): + test_multiple_chunks() + test_error_chunk_terminates() + test_err_reply_terminates() + test_idle_timeout() + test_no_worker() + test_call_compat_first_reply_only() + + print(f"\n结果: {len(failures)} 个失败") + return 1 if failures else 0 + + +if __name__ == "__main__": + sys.exit(run()) diff --git a/test_stream_route.py b/test_stream_route.py new file mode 100644 index 0000000..719f3d0 --- /dev/null +++ b/test_stream_route.py @@ -0,0 +1,221 @@ +""" +网关路由层流式端点单测（依赖隔离：桩掉 chongming 包链、auth 与 chongming_cache） + +直接加载 src/chongming/app/core/route_generator.py，用 FastAPI TestClient 验证: + 1. 流式 POST 端点：SSE 输出（media_type / 逐块 data: 帧 / 响应头） + 2. 流式 GET 端点：query 参数正确透传到 payload + 3. 初始错误（无可用 worker）-> HTTP 503 + 4. 流中空闲超时 -> SSE error 事件（HTTP 200 已提交） + 5. 兼容性回归：非流式端点输出与改动前一致（JSON 响应） + 6. 非流式端点仍走原 call() 逻辑（不误入流式分支） + +运行: uv run python test_stream_route.py +""" +import importlib.util +import logging +import sys +import types +from pathlib import Path + +import orjson +from fastapi import FastAPI +from fastapi.testclient import TestClient + +REPO_ROOT = Path(__file__).resolve().parent +sys.path.insert(0, str(REPO_ROOT / "src")) +sys.path.insert(0, str(REPO_ROOT / "utils" / "worker" / "src")) + +# ---- 依赖隔离：桩掉 chongming 包链 / auth / chongming_cache ---- +_pkg = types.ModuleType("chongming") +_pkg.__path__ = [] +_app = types.ModuleType("chongming.app") +_app.__path__ = [] +_core = types.ModuleType("chongming.app.core") +_core.__path__ = [] +_logger_stub = types.ModuleType("chongming.app.core.logger") +_logger_stub.get_logger = lambda name: logging.getLogger(name) +_auth_stub = types.ModuleType("chongming.app.core.auth") +_auth_stub.store_user_to_request = lambda x: x # 仅用于 Depends，桩掉鉴权 +_cache_stub = types.ModuleType("chongming_cache") +_cache_stub.Cache = object +for mod in (_pkg, _app, _core, _logger_stub, _auth_stub, _cache_stub): + sys.modules[mod.__name__] = mod + +# 加载真实 worker_manager（route_generator 依赖其 StreamIdleTimeoutError） +_wm_path = REPO_ROOT / "src" / "chongming" / "app" / "core" / "worker_manager.py" +_wm_spec = importlib.util.spec_from_file_location( + "chongming.app.core.worker_manager", _wm_path +) +_wm_mod = importlib.util.module_from_spec(_wm_spec) +sys.modules["chongming.app.core.worker_manager"] = _wm_mod +_wm_spec.loader.exec_module(_wm_mod) # type: ignore[union-attr] + +# 加载真实 route_generator +_rg_path = REPO_ROOT / "src" / "chongming" / "app" / "core" / "route_generator.py" +_rg_spec = importlib.util.spec_from_file_location( + "chongming.app.core.route_generator", _rg_path +) +_rg_mod = importlib.util.module_from_spec(_rg_spec) +sys.modules["chongming.app.core.route_generator"] = _rg_mod +_rg_spec.loader.exec_module(_rg_mod) # type: ignore[union-attr] + +register_route_group = _rg_mod.register_route_group +StreamIdleTimeoutError = _wm_mod.StreamIdleTimeoutError + +failures = [] + + +def check(name, cond, detail=""): + status = "PASS" if cond else "FAIL" + print(f"[{status}] {name}" + (f" — {detail}" if detail else "")) + if not cond: + failures.append(name) + + +class FakeWorkerManager: + """模拟 WorkerManager：call() 同步返回；call_stream() 逐块产出""" + + def __init__(self, chunks=None, stream_error=None): + self.chunks = chunks if chunks is not None else [{"token": "a"}, {"token": "b"}, {"done": True}] + self.stream_error = stream_error # "no_worker" | "idle" | None + self.call_payloads = [] + self.stream_payloads = [] + + async def call(self, worker_type, subject, payload, timeout=3.0, retries=1, cacheable=False): + self.call_payloads.append((worker_type, subject, payload)) + data = payload["data"] + return {"result": data.get("a", 0) + data.get("b", 0)} + + async def call_stream(self, worker_type, subject, payload, timeout=300.0, idle_timeout=10.0): + if self.stream_error == "no_worker": + raise RuntimeError(f"类型 '{worker_type}' 没有可用 Worker。") + self.stream_payloads.append((worker_type, subject, payload)) + + async def _gen(): + if self.stream_error == "idle": + yield orjson.dumps({"token": "a"}) + raise StreamIdleTimeoutError("流式空闲超时 (0.1s 无数据)") + for c in self.chunks: + yield orjson.dumps(c) + + return _gen() + + +def build_app(fwm): + app = FastAPI() + routes = [ + # 流式 POST（body 参数） + { + "subject": "streamtest.chat", "path": "/stream/chat", "method": "POST", + "params": ["prompt: str", "tokens: int"], "summary": "sse chat", + "stream": True, "timeout": 300.0, "stream_idle_timeout": 1.0, + "auth_required": False, + }, + # 流式 GET（query 参数） + { + "subject": "streamtest.progress", "path": "/stream/progress", "method": "GET", + "params": ["job_id: str", "interval: float"], "summary": "sse progress", + "stream": True, "timeout": 60.0, "auth_required": False, + }, + # 非流式对照（兼容性回归） + { + "subject": "streamtest.add", "path": "/calc/add", "method": "GET", + "params": ["a: float", "b: float"], "summary": "add", + "timeout": 3.0, "auth_required": False, + }, + { + "subject": "streamtest.echo", "path": "/calc/echo", "method": "POST", + "params": ["msg: str"], "summary": "echo", + "timeout": 3.0, "auth_required": False, + }, + ] + register_route_group(app, fwm, "streamtest", routes, prefix="/api/v1") + return TestClient(app) + + +def test_stream_post_sse(): + fwm = FakeWorkerManager() + client = build_app(fwm) + r = client.post("/api/v1/stream/chat", json={"prompt": "hi", "tokens": 3}) + check("流式 POST 状态码 200", r.status_code == 200, str(r.status_code)) + check( + "SSE media_type", + r.headers["content-type"].startswith("text/event-stream"), + r.headers.get("content-type", ""), + ) + check( + "SSE 逐块 data: 帧", + r.text + == 'data: {"token":"a"}\n\ndata: {"token":"b"}\n\ndata: {"done":true}\n\n', + repr(r.text), + ) + check("防缓冲响应头", r.headers.get("cache-control") == "no-cache") + check( + "body 参数透传", + fwm.stream_payloads[0][2]["data"] == {"prompt": "hi", "tokens": 3}, + ) + + +def test_stream_get_query_params(): + fwm = FakeWorkerManager() + client = build_app(fwm) + r = client.get("/api/v1/stream/progress?job_id=job-1&interval=0.5") + check("流式 GET 状态码 200", r.status_code == 200) + check( + "GET query 参数透传", + fwm.stream_payloads[0][2]["data"] == {"job_id": "job-1", "interval": 0.5}, + str(fwm.stream_payloads[0][2]["data"]), + ) + check("GET 流式输出 SSE", r.text.startswith("data: "), repr(r.text[:40])) + + +def test_stream_no_worker_503(): + fwm = FakeWorkerManager(stream_error="no_worker") + client = build_app(fwm) + r = client.post("/api/v1/stream/chat", json={"prompt": "hi", "tokens": 1}) + check("初始错误（无 worker）-> 503", r.status_code == 503, str(r.status_code)) + + +def test_stream_idle_timeout_sse_event(): + fwm = FakeWorkerManager(stream_error="idle") + client = build_app(fwm) + r = client.post("/api/v1/stream/chat", json={"prompt": "hi", "tokens": 1}) + check("空闲超时保持 200（已提交）", r.status_code == 200) + check( + "空闲超时转为 SSE error 事件", + r.text == 'data: {"token":"a"}\n\ndata: {"error":"流式空闲超时 (0.1s 无数据)"}\n\n', + repr(r.text), + ) + + +def test_non_stream_compat(): + """兼容性回归：非流式端点走原 call() 逻辑，输出不变""" + fwm = FakeWorkerManager() + client = build_app(fwm) + + r = client.get("/api/v1/calc/add?a=2&b=3") + check("非流式 GET 仍为 JSON 响应", r.status_code == 200 and r.json() == {"result": 5.0}, + f"status={r.status_code} body={r.text}") + + r = client.post("/api/v1/calc/echo", json={"msg": "hello"}) + check("非流式 POST 仍为 JSON 响应", + r.status_code == 200 and r.json() == {"result": 0.0}, + f"status={r.status_code} body={r.text}") + + check("非流式调用未进入 call_stream", len(fwm.stream_payloads) == 0) + check("非流式调用仍走 call()", len(fwm.call_payloads) == 2) + + +def run(): + test_stream_post_sse() + test_stream_get_query_params() + test_stream_no_worker_503() + test_stream_idle_timeout_sse_event() + test_non_stream_compat() + + print(f"\n结果: {len(failures)} 个失败") + return 1 if failures else 0 + + +if __name__ == "__main__": + sys.exit(run()) diff --git a/test_stream_worker.py b/test_stream_worker.py new file mode 100644 index 0000000..3722a3b --- /dev/null +++ b/test_stream_worker.py @@ -0,0 +1,118 @@ +""" +Worker 框架流式回复单测（不依赖真实 zenoh 会话） + +直接测 runtime.stream_replies_to_query（fake query 仅需 key_expr 与 reply）: + 1. 异步生成器：多块顺序 reply，payload 为 orjson JSON + 2. 同步生成器：同样逐块 reply + 3. 生成器中途异常：{"error": ...} 收尾块后停流 + 4. reply 失败（查询已终止）：立即停流且不抛异常 + +运行: uv run python test_stream_worker.py +""" +import asyncio +import sys +from pathlib import Path + +import orjson + +REPO_ROOT = Path(__file__).resolve().parent +# 优先使用本地框架 checkout（即使 chongming_worker 已通过 uv 安装也以本地为准） +sys.path.insert(0, str(REPO_ROOT / "utils" / "worker" / "src")) + +from chongming_worker.runtime import stream_replies_to_query # noqa: E402 + +failures = [] + + +def check(name, cond, detail=""): + status = "PASS" if cond else "FAIL" + print(f"[{status}] {name}" + (f" — {detail}" if detail else "")) + if not cond: + failures.append(name) + + +class FakeQuery: + """仅实现 stream_replies_to_query 用到的两个接口（key_expr / reply）""" + + def __init__(self, key_expr="worker/example/worker-0/example.stream_chat"): + self.key_expr = key_expr + self.replies = [] + self.fail_reply = False + + def reply(self, key_expr, payload): + if self.fail_reply: + raise RuntimeError("query terminated") + self.replies.append((key_expr, payload)) + + +def test_async_generator(): + q = FakeQuery() + + async def gen(): + yield {"token": "a"} + yield {"token": "b"} + + asyncio.run(stream_replies_to_query(q, gen())) + check("异步生成器逐块回复", len(q.replies) == 2) + check("回复 key_expr 正确", all(k == q.key_expr for k, _ in q.replies)) + check( + "payload 为 JSON dict", + [orjson.loads(p) for _, p in q.replies] == [{"token": "a"}, {"token": "b"}], + ) + + +def test_sync_generator(): + q = FakeQuery() + + def gen(): + yield {"n": 1} + yield {"n": 2} + + asyncio.run(stream_replies_to_query(q, gen())) + check( + "同步生成器逐块回复", + [orjson.loads(p) for _, p in q.replies] == [{"n": 1}, {"n": 2}], + ) + + +def test_mid_stream_exception(): + q = FakeQuery() + + async def gen(): + yield {"token": "a"} + raise ValueError("boom") + + asyncio.run(stream_replies_to_query(q, gen())) + check("中途异常收尾 error 块", len(q.replies) == 2) + check("error 块内容", orjson.loads(q.replies[1][1]) == {"error": "boom"}) + + +def test_reply_failure_stops(): + q = FakeQuery() + q.fail_reply = True + yielded = [] + + async def gen(): + yielded.append("yield-1") + yield {"token": "a"} + yielded.append("yield-2") # 首个块 reply 失败后不应被执行到 + yield {"token": "b"} + + # 不应抛异常，且首个块 reply 失败后立即停流 + asyncio.run(stream_replies_to_query(q, gen())) + check("reply 失败不抛异常", True) + check("reply 失败后立即停流", yielded == ["yield-1"]) + + +def run(): + test_async_generator() + test_sync_generator() + test_mid_stream_exception() + test_reply_failure_stops() + + print(f"\n结果: {len(failures)} 个失败") + return 1 if failures else 0 + + +if __name__ == "__main__": + sys.exit(run()) diff --git a/uv.lock b/uv.lock index 4417525..53824e1 100644 --- a/uv.lock +++ b/uv.lock @@ -486,7 +486,7 @@ source = { git = "http://124.71.68.6:3000/chongming/chongming-stream-manager.git [[package]] name = "chongming-worker" version = "0.2.0" -source = { git = "http://124.71.68.6:3000/chongming/chongming-utils-worker.git#d8230bc0ca913a1253d9db9d9af231da9fa7c779" } +source = { git = "http://124.71.68.6:3000/chongming/chongming-utils-worker.git#6dbdce177b2e44f2d0fd84df51d49bd690a8926b" } dependencies = [ { name = "eclipse-zenoh" }, { name = "pydantic" },