refactor(channels): 重构Agent运行事件流式消费逻辑，移除SSE协议转换
主要变更： 1. 移除`stream_agent_run_events`相关SSE文本解析逻辑，改为直接读取Redis Stream结构化事件 2. 新增`_stream_structured_events`异步生成器，实现Redis Stream增量轮询与事件产出 3. 重构执行适配器与端口契约，更新文档注释与权限校验逻辑 4. 移除冗余的verbose参数与SSE协议相关代码 5. 新增超时兜底、终态run检测等安全机制 6. 更新单元测试用例以适配新的流式消费逻辑
This commit is contained in:
parent
1f47786dc8
commit
c216894a3e
@ -57,7 +57,7 @@ if TYPE_CHECKING:
|
||||
|
||||
__all__ = ["AgentRunAdapter"]
|
||||
|
||||
# 单次分页读取的事件上限，与 stream_agent_run_events 保持一致
|
||||
# 单次分页读取的事件上限，与 _stream_structured_events 保持一致
|
||||
_EVENTS_PAGE_LIMIT = 200
|
||||
|
||||
|
||||
@ -468,7 +468,6 @@ class AgentRunAdapter(AgentRunPort):
|
||||
run_id=run_id.value,
|
||||
after_seq="",
|
||||
current_uid=current_uid,
|
||||
verbose=True,
|
||||
):
|
||||
yield StreamEvent(
|
||||
event_type=event.get("event_type", ""),
|
||||
|
||||
@ -12,84 +12,140 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import asyncio
|
||||
import os
|
||||
from collections.abc import AsyncIterator
|
||||
from typing import Any
|
||||
|
||||
from yuxi.channels.contract import DependencyError, Error, NotFoundError
|
||||
from yuxi.channels.contract.ports.driven.shared.agent_run_execution_port import (
|
||||
AgentRunExecutionPort,
|
||||
)
|
||||
from yuxi.repositories.agent_run_repository import (
|
||||
TERMINAL_RUN_STATUSES,
|
||||
AgentRunRepository,
|
||||
)
|
||||
from yuxi.services.agent_run_service import (
|
||||
create_agent_run_view,
|
||||
get_run_terminal_info,
|
||||
get_run_uid,
|
||||
stream_agent_run_events,
|
||||
)
|
||||
from yuxi.services.run_queue_service import list_run_stream_events
|
||||
from yuxi.services.run_queue_service import (
|
||||
list_run_stream_events,
|
||||
normalize_after_seq,
|
||||
)
|
||||
from yuxi.storage.postgres.manager import pg_manager
|
||||
from yuxi.utils.datetime_utils import utc_now_naive
|
||||
|
||||
__all__ = ["AgentRunExecutionAdapter"]
|
||||
|
||||
|
||||
async def _stream_parsed_events(
|
||||
# 与 services 层 stream_agent_run_events 保持一致，避免引入新配置项。
|
||||
# 注意：env 变量名沿用 SSE 命名以保持运维配置兼容，但实际语义为
|
||||
# 「单次运行事件流最大消费时长（分钟）」，覆盖 SSE 与结构化流两条路径。
|
||||
_STREAM_POLL_INTERVAL_SECONDS = float(os.getenv("RUN_SSE_POLL_INTERVAL_SECONDS", "1.0"))
|
||||
_STREAM_MAX_CONNECTION_MINUTES = float(os.getenv("RUN_SSE_MAX_CONNECTION_MINUTES", "30"))
|
||||
_EVENTS_PAGE_LIMIT = 200
|
||||
|
||||
|
||||
async def _stream_structured_events(
|
||||
*,
|
||||
run_id: str,
|
||||
after_seq: str,
|
||||
current_uid: str,
|
||||
verbose: bool = True,
|
||||
) -> AsyncIterator[dict[str, Any]]:
|
||||
"""包装 ``stream_agent_run_events`` 并将 SSE 文本解析为结构化字典。
|
||||
"""基于 Redis Stream 增量轮询产出结构化事件字典。
|
||||
|
||||
作为 ``AgentRunExecutionAdapter.streamAgentRunEvents`` 的底层异步生成器，
|
||||
逐项将 SSE 文本解析为 ``{event_type, payload, seq}`` 字典后 yield。
|
||||
直接读取 Redis Stream 结构化事件（不经过 SSE 文本协议），避免 SSE
|
||||
heartbeat 等协议行触发 ``StreamEvent.seq`` 非空契约校验失败。
|
||||
|
||||
循环逻辑：
|
||||
1. 查 DB 获取 run 状态（不存在抛 ``NotFoundError``）
|
||||
2. ``list_run_stream_events`` 增量读取 Redis Stream
|
||||
3. 逐事件 yield ``{event_type, payload, seq}`` 字典（seq 为
|
||||
Redis Stream ID，永远非空，满足 StreamEvent DTO 契约）
|
||||
4. 遇 ``end`` 事件退出
|
||||
5. run 已终态且 Redis Stream 无新事件时退出（worker 崩溃场景，
|
||||
由 ``LoadBuildStage`` 通过 ``getAgentRunFinalOutput`` +
|
||||
``_raiseOnTerminalStatus`` 处理）
|
||||
6. 累计消费时长超 ``_STREAM_MAX_CONNECTION_MINUTES`` 退出
|
||||
（兜底安全阀，对齐 SSE 路径的 ``SSE_MAX_CONNECTION_MINUTES``，
|
||||
防 worker 卡死在 pending/running 永不终态导致无限轮询）
|
||||
7. 休眠 ``_STREAM_POLL_INTERVAL_SECONDS`` 后进入下一轮
|
||||
|
||||
参数：
|
||||
run_id: Agent 运行 ID。
|
||||
after_seq: 起始 Redis Stream seq（``"0-0"`` 或 ``""`` 表示从头）。
|
||||
current_uid: 调用方 UID，用于 run 权限校验。
|
||||
|
||||
生成：
|
||||
``{event_type, payload, seq}`` 字典。
|
||||
|
||||
抛出：
|
||||
NotFoundError: run 不存在或 current_uid 无权访问。
|
||||
DependencyError: DB / Redis 故障。
|
||||
"""
|
||||
last_seq = normalize_after_seq(after_seq)
|
||||
started_at = utc_now_naive()
|
||||
|
||||
def _parse_sse(sse_text: str) -> dict[str, Any]:
|
||||
"""将单个 SSE 文本块解析为 ``{event_type, payload, seq}`` 字典。
|
||||
try:
|
||||
while True:
|
||||
# 1. 查 DB 获取 run 状态（每轮循环一次，用于终态兜底判断）
|
||||
try:
|
||||
async with pg_manager.get_async_session_context() as db:
|
||||
repo = AgentRunRepository(db)
|
||||
run = await repo.get_run_for_user(run_id, str(current_uid))
|
||||
if not run:
|
||||
raise NotFoundError("agent_run", run_id)
|
||||
except (NotFoundError, DependencyError):
|
||||
raise
|
||||
except Exception as exc:
|
||||
raise DependencyError("postgres", cause=Error(str(exc))) from exc
|
||||
|
||||
底层 ``stream_agent_run_events`` 产出 SSE 格式文本（``event:`` 与
|
||||
``data:`` 行），本函数解析为结构化字典。解析失败时返回
|
||||
``{"event_type": "raw", "payload": {"text": sse_text}, "seq": ""}``。
|
||||
"""
|
||||
event_type = "raw"
|
||||
data_str = ""
|
||||
seq = ""
|
||||
for line in sse_text.splitlines():
|
||||
if line.startswith("event: "):
|
||||
event_type = line[len("event: ") :]
|
||||
elif line.startswith("data: "):
|
||||
data_str = line[len("data: ") :]
|
||||
elif line.startswith("id: "):
|
||||
seq = line[len("id: ") :]
|
||||
if not data_str:
|
||||
return {
|
||||
"event_type": event_type,
|
||||
"payload": {"text": sse_text},
|
||||
"seq": seq,
|
||||
}
|
||||
try:
|
||||
payload = json.loads(data_str)
|
||||
except json.JSONDecodeError:
|
||||
return {
|
||||
"event_type": event_type,
|
||||
"payload": {"text": sse_text},
|
||||
"seq": seq,
|
||||
}
|
||||
# SSE data 中可能嵌套 seq 字段，优先使用外层 id:，回退 data 内 seq
|
||||
if isinstance(payload, dict) and not seq:
|
||||
seq = payload.get("seq", "")
|
||||
return {
|
||||
"event_type": event_type,
|
||||
"payload": payload,
|
||||
"seq": seq,
|
||||
}
|
||||
# 2. 拉 Redis Stream 增量事件
|
||||
try:
|
||||
events = await list_run_stream_events(
|
||||
run_id,
|
||||
after_seq=last_seq,
|
||||
limit=_EVENTS_PAGE_LIMIT,
|
||||
)
|
||||
except (NotFoundError, DependencyError):
|
||||
raise
|
||||
except Exception as exc:
|
||||
raise DependencyError("redis_stream", cause=Error(str(exc))) from exc
|
||||
|
||||
async for sse_text in stream_agent_run_events(
|
||||
run_id=run_id,
|
||||
after_seq=after_seq,
|
||||
current_uid=current_uid,
|
||||
verbose=verbose,
|
||||
):
|
||||
yield _parse_sse(sse_text)
|
||||
# 3. yield 事件 + 检查 end
|
||||
for event in events:
|
||||
seq = str(event.get("seq") or "0-0")
|
||||
last_seq = seq
|
||||
event_type = event.get("event_type") or "message"
|
||||
envelope = event.get("payload") or {}
|
||||
yield {
|
||||
"event_type": event_type,
|
||||
"payload": envelope,
|
||||
"seq": seq,
|
||||
}
|
||||
if event_type == "end":
|
||||
return
|
||||
|
||||
# 4. run 已终态且 Redis Stream 无 end 事件（worker 崩溃）：
|
||||
# 直接退出，由 LoadBuildStage 通过 getAgentRunFinalOutput +
|
||||
# _raiseOnTerminalStatus 按终态抛领域错误
|
||||
if run.status in TERMINAL_RUN_STATUSES:
|
||||
return
|
||||
|
||||
# 5. 兜底安全阀：累计消费时长超阈值退出（worker 卡死在
|
||||
# pending/running 永不终态时防无限轮询，由调用方按
|
||||
# 流式提前结束处理，LoadBuildStage 仍可重放事件流兜底）
|
||||
elapsed_seconds = (utc_now_naive() - started_at).total_seconds()
|
||||
if elapsed_seconds >= _STREAM_MAX_CONNECTION_MINUTES * 60:
|
||||
return
|
||||
|
||||
# 6. 轮询休眠
|
||||
await asyncio.sleep(_STREAM_POLL_INTERVAL_SECONDS)
|
||||
except asyncio.CancelledError:
|
||||
return
|
||||
|
||||
|
||||
class AgentRunExecutionAdapter(AgentRunExecutionPort):
|
||||
@ -100,8 +156,9 @@ class AgentRunExecutionAdapter(AgentRunExecutionPort):
|
||||
函数的转发，``db`` 等技术资源由调用方通过方法参数传入或由 service 内部
|
||||
自管理。
|
||||
|
||||
``streamAgentRunEvents`` 额外承担 SSE 文本到结构化字典的解析，使
|
||||
``AgentRunAdapter`` 仅处理结构化数据而非线协议格式。
|
||||
``streamAgentRunEvents`` 通过 ``_stream_structured_events`` 直接读取
|
||||
Redis Stream 结构化事件，不经过 SSE 文本协议转换，避免 heartbeat 等协议行
|
||||
触发 ``StreamEvent.seq`` 非空契约校验失败。
|
||||
"""
|
||||
|
||||
async def createAgentRunView(
|
||||
@ -156,19 +213,31 @@ class AgentRunExecutionAdapter(AgentRunExecutionPort):
|
||||
run_id: str,
|
||||
after_seq: str,
|
||||
current_uid: str,
|
||||
verbose: bool = True,
|
||||
) -> AsyncIterator[dict[str, Any]]:
|
||||
"""委托 ``stream_agent_run_events`` 并解析 SSE 为结构化事件字典。
|
||||
"""基于 Redis Stream 增量轮询产出结构化事件字典。
|
||||
|
||||
底层 ``stream_agent_run_events`` 产出 SSE 文本行，本方法返回的异步
|
||||
迭代器逐项解析为 ``{event_type, payload, seq}`` 字典，供
|
||||
``AgentRunAdapter`` 转换为 ``StreamEvent`` DTO。
|
||||
委托 ``_stream_structured_events`` 异步生成器，直接读取
|
||||
``run_queue_service.list_run_stream_events`` 返回的结构化事件
|
||||
（含 ``seq`` 字段，源自 Redis Stream ID，永远非空），不经过 SSE
|
||||
文本协议转换，避免 heartbeat 等协议行触发 ``StreamEvent.seq``
|
||||
非空契约校验失败。
|
||||
|
||||
@pre
|
||||
- run_id 非空且为已创建的运行 ID
|
||||
- current_uid 与 run 记录的 uid 一致（``get_run_for_user`` 权限校验）
|
||||
|
||||
@post
|
||||
- 返回 AsyncIterator[dict[str, Any]]，逐项产出结构化事件字典
|
||||
（含 event_type / payload / seq 等键，seq 永远非空）
|
||||
|
||||
@failure
|
||||
- NotFoundError: run 不存在或 current_uid 无权访问
|
||||
- DependencyError: DB / Redis 故障
|
||||
"""
|
||||
return _stream_parsed_events(
|
||||
return _stream_structured_events(
|
||||
run_id=run_id,
|
||||
after_seq=after_seq,
|
||||
current_uid=current_uid,
|
||||
verbose=verbose,
|
||||
)
|
||||
|
||||
async def listRunStreamEvents(
|
||||
|
||||
@ -9,7 +9,8 @@ DB 技术访问逻辑，供 ``AgentRunAdapter`` 通过本端口调用，消除
|
||||
- ``createAgentRunView`` ← ``agent_run_service.create_agent_run_view``
|
||||
- ``getRunUid`` ← ``agent_run_service.get_run_uid``
|
||||
- ``getRunTerminalInfo`` ← ``agent_run_service.get_run_terminal_info``
|
||||
- ``streamAgentRunEvents`` ← ``agent_run_service.stream_agent_run_events``
|
||||
- ``streamAgentRunEvents`` ← 基于 ``run_queue_service.list_run_stream_events``
|
||||
的 Redis Stream 轮询（不依赖 SSE 文本协议）
|
||||
- ``listRunStreamEvents`` ← ``run_queue_service.list_run_stream_events``
|
||||
|
||||
关联约束：ADP-001（适配器不得直接依赖 ``yuxi.services``）。
|
||||
@ -117,25 +118,30 @@ class AgentRunExecutionPort(Protocol):
|
||||
run_id: str,
|
||||
after_seq: str,
|
||||
current_uid: str,
|
||||
verbose: bool = True,
|
||||
) -> AsyncIterator[dict[str, Any]]:
|
||||
"""流式消费 Agent 运行事件。
|
||||
"""基于 Redis Stream 增量轮询产出结构化事件字典。
|
||||
|
||||
委托 ``agent_run_service.stream_agent_run_events`` 实现（异步
|
||||
生成器），将底层 SSE 文本行解析为结构化事件字典后产出，供
|
||||
``AgentRunAdapter`` 转换为 ``StreamEvent`` DTO。
|
||||
委托 ``run_queue_service.list_run_stream_events`` 直接读取 Redis
|
||||
Stream 结构化事件（含 ``seq`` 字段，源自 Redis Stream ID，永远
|
||||
非空），不经过 SSE 文本协议转换，避免 heartbeat 等协议行触发
|
||||
``StreamEvent.seq`` 非空契约校验失败。
|
||||
|
||||
与 ``agent_run_service.stream_agent_run_events`` 的区别：后者服务
|
||||
HTTP SSE 端点（前端订阅场景），产出 SSE 文本协议；本方法服务
|
||||
channels 内部消费场景（``AgentRunAdapter.streamAgentRun``），直接
|
||||
产出结构化字典，消除 SSE 文本→解析→结构化的冗余转换链路。
|
||||
|
||||
@pre
|
||||
- run_id 非空且为已创建的运行 ID
|
||||
- current_uid 与 run 记录的 uid 一致
|
||||
- current_uid 与 run 记录的 uid 一致（``get_run_for_user`` 权限校验）
|
||||
|
||||
@post
|
||||
- 返回 AsyncIterator[dict[str, Any]]，逐项产出结构化事件字典
|
||||
（含 event_type / payload / seq 等键）
|
||||
（含 event_type / payload / seq 等键，seq 永远非空）
|
||||
|
||||
@failure
|
||||
- DependencyError：DB / Redis 故障
|
||||
- OperationTimeoutError：流式消费超时
|
||||
- NotFoundError: run 不存在或 current_uid 无权访问
|
||||
- DependencyError: DB / Redis 故障
|
||||
"""
|
||||
...
|
||||
|
||||
|
||||
@ -2,7 +2,8 @@
|
||||
|
||||
定义核心层对 Agent 运行能力的依赖契约，由 application 层调用，framework
|
||||
层实现。端口覆盖 Agent 运行的创建与流式消费，复用现有 agent_run_service
|
||||
与 chat_service，不另起执行体系。
|
||||
（创建运行视图）与 run_queue_service（增量读取 Redis Stream 事件），
|
||||
不另起执行体系。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
@ -34,8 +35,10 @@ class AgentRunPort(Protocol):
|
||||
复用约束（INV-I3 / FF-REUSE-02）：
|
||||
- ``createAgentRun`` **必须** 调用现有 ``agent_run_service.create_agent_run_view``，
|
||||
**不得** 另起执行体系。
|
||||
- ``streamAgentRun`` **必须** 复用现有 ``chat_service.stream_agent_chat``，
|
||||
**不得** 另起流式实现（FR-31）。
|
||||
- ``streamAgentRun`` **必须** 复用现有 ``run_queue_service.list_run_stream_events``
|
||||
增量读取 Redis Stream 事件（FR-31），**不得** 另起流式实现或绕过事件存储
|
||||
直接订阅底层 LLM 流。``current_uid`` 由 ``getRunUid`` 从 ``agent_runs`` 表
|
||||
查询获取（创建时由 ``service_account.uid`` 写入），不依赖外部订阅凭证。
|
||||
|
||||
关键约束：
|
||||
- **必须** 支持 FR-10 Agent 渠道上下文注入（渠道类型、账户 ID、会话键等
|
||||
@ -86,7 +89,6 @@ class AgentRunPort(Protocol):
|
||||
@failure
|
||||
- ValidationError：cmd 字段校验失败
|
||||
- DependencyError：依赖服务（DB/Redis/ARQ）故障
|
||||
- OperationTimeoutError：Agent 运行创建超时
|
||||
|
||||
@consistency
|
||||
- tx 非空时 AgentRun 创建与主事务原子提交（C-I1）
|
||||
@ -98,20 +100,25 @@ class AgentRunPort(Protocol):
|
||||
async def streamAgentRun(self, run_id: AgentRunId) -> AsyncIterator[StreamEvent]:
|
||||
"""流式消费 Agent 运行事件。
|
||||
|
||||
复用现有 ``chat_service.stream_agent_chat``（FR-31），不另起流式
|
||||
实现。覆盖 FR-13（流式）、FR-31（流式能力）。
|
||||
增量轮询 Redis Stream 事件存储（``run_queue_service.list_run_stream_events``）
|
||||
并转换为 ``StreamEvent`` DTO（FR-31），不订阅底层 LLM 流。覆盖 FR-13
|
||||
（流式）、FR-31（流式能力）。``current_uid`` 通过 ``getRunUid`` 从
|
||||
``agent_runs`` 表查询获取，用于 ``get_run_for_user`` 权限校验。
|
||||
|
||||
@pre
|
||||
- run_id 非空且为已创建的 Agent 运行 ID
|
||||
- 调用方持有有效的 Agent 运行订阅凭证
|
||||
- run 记录的 uid 字段已写入（创建时由 ``service_account.uid`` 注入）
|
||||
|
||||
@post
|
||||
- 返回 AsyncIterator[StreamEvent]，迭代产出流式事件
|
||||
- 流式结束时代理运行状态已回写
|
||||
- 遇 ``end`` 事件、run 转入终态、或累计消费时长超最大连接
|
||||
时长（``RUN_SSE_MAX_CONNECTION_MINUTES``，默认 30 分钟）时
|
||||
迭代结束
|
||||
- 仅读取 Redis Stream 与 ``agent_runs`` 表，不修改持久化状态
|
||||
|
||||
@failure
|
||||
- NotFoundError：run 不存在或 current_uid 无权访问
|
||||
- DependencyError：依赖服务（DB/Redis）故障
|
||||
- OperationTimeoutError：流式消费超时
|
||||
|
||||
@consistency
|
||||
- 最终一致性：流式分片可能滞后于 Agent 实际执行状态
|
||||
|
||||
@ -549,7 +549,7 @@ class TestAgentRunAdapterStream:
|
||||
execution_port = MagicMock()
|
||||
execution_port.getRunUid = AsyncMock(return_value="sender-1")
|
||||
|
||||
async def _events(*, run_id, after_seq, current_uid, verbose):
|
||||
async def _events(*, run_id, after_seq, current_uid):
|
||||
yield {"event_type": "messages", "payload": {"k": 1}, "seq": "1-0"}
|
||||
yield {"event_type": "tool_call", "payload": {"k": 2}, "seq": "2-0"}
|
||||
|
||||
@ -577,7 +577,7 @@ class TestAgentRunAdapterStream:
|
||||
execution_port = MagicMock()
|
||||
execution_port.getRunUid = AsyncMock(return_value="uid-1")
|
||||
|
||||
async def _events(*, run_id, after_seq, current_uid, verbose):
|
||||
async def _events(*, run_id, after_seq, current_uid):
|
||||
# 仅提供 event_type 与 seq，payload 缺失时默认 {}
|
||||
yield {"event_type": "messages", "seq": "1-0"}
|
||||
|
||||
|
||||
@ -2,22 +2,41 @@
|
||||
|
||||
覆盖 ``AgentRunExecutionAdapter`` 的 ``createAgentRunView`` / ``getRunUid`` /
|
||||
``streamAgentRunEvents`` / ``listRunStreamEvents`` 方法，patch
|
||||
``agent_run_service`` 与 ``run_queue_service`` 模块级函数，不连接真实 Redis / DB。
|
||||
``agent_run_service`` / ``run_queue_service`` / ``pg_manager`` /
|
||||
``AgentRunRepository`` 模块级函数与对象，不连接真实 Redis / DB。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from contextlib import asynccontextmanager
|
||||
from datetime import datetime, timedelta
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from yuxi.channels.adapters.agent_run_execution_adapter import (
|
||||
AgentRunExecutionAdapter,
|
||||
_stream_parsed_events,
|
||||
_stream_structured_events,
|
||||
)
|
||||
from yuxi.channels.contract import DependencyError, NotFoundError
|
||||
|
||||
pytestmark = pytest.mark.unit
|
||||
|
||||
|
||||
def _make_pg_session_mock():
|
||||
"""构造 ``pg_manager.get_async_session_context`` 的 mock。
|
||||
|
||||
``get_async_session_context`` 是 ``@asynccontextmanager`` 装饰的，
|
||||
调用时返回 async context manager（不需 await），通过 ``async with``
|
||||
进入。返回的 session 用 ``MagicMock`` 占位（不实际使用）。
|
||||
"""
|
||||
|
||||
@asynccontextmanager
|
||||
async def _ctx():
|
||||
yield MagicMock()
|
||||
|
||||
return _ctx
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestAgentRunExecutionAdapterCreateView:
|
||||
@pytest.mark.asyncio
|
||||
@ -126,23 +145,53 @@ class TestAgentRunExecutionAdapterGetRunUid:
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestAgentRunExecutionAdapterStreamEvents:
|
||||
"""``streamAgentRunEvents`` 基于 Redis Stream 增量轮询产出结构化事件。
|
||||
|
||||
覆盖：
|
||||
- 正常事件 yield（含 messages / tool_call / end 等事件类型）
|
||||
- 遇 ``end`` 事件立即退出
|
||||
- run 不存在抛 ``NotFoundError``
|
||||
- run 已终态且无 end 事件时退出（worker 崩溃兜底）
|
||||
- 累计消费时长超 ``_STREAM_MAX_CONNECTION_MINUTES`` 时退出（兜底安全阀）
|
||||
- DB / Redis 异常翻译为 ``DependencyError``
|
||||
- seq 字段直接取自 Redis Stream ID（永远非空）
|
||||
"""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stream_agent_run_events_parses_sse_with_data(self):
|
||||
# Arrange
|
||||
adapter = AgentRunExecutionAdapter()
|
||||
sse_text = 'event: messages\ndata: {"k": "v"}\nid: 1-0\n\n'
|
||||
async def test_stream_events_yields_redis_stream_events(self):
|
||||
# Arrange — list_run_stream_events 返回 2 个事件，第 2 个为 end
|
||||
events_batch = [
|
||||
{"event_type": "messages", "payload": {"k": 1}, "seq": "1-0"},
|
||||
{
|
||||
"event_type": "end",
|
||||
"payload": {"status": "completed"},
|
||||
"seq": "2-0",
|
||||
},
|
||||
]
|
||||
run_mock = MagicMock()
|
||||
run_mock.status = "completed"
|
||||
repo_mock = MagicMock()
|
||||
repo_mock.get_run_for_user = AsyncMock(return_value=run_mock)
|
||||
|
||||
async def _fake_stream(*, run_id, after_seq, current_uid, verbose):
|
||||
yield sse_text
|
||||
|
||||
with patch(
|
||||
"yuxi.channels.adapters.agent_run_execution_adapter.stream_agent_run_events",
|
||||
new=_fake_stream,
|
||||
with (
|
||||
patch(
|
||||
"yuxi.channels.adapters.agent_run_execution_adapter.pg_manager"
|
||||
) as mock_pg,
|
||||
patch(
|
||||
"yuxi.channels.adapters.agent_run_execution_adapter.AgentRunRepository"
|
||||
) as mock_repo_cls,
|
||||
patch(
|
||||
"yuxi.channels.adapters.agent_run_execution_adapter.list_run_stream_events",
|
||||
new=AsyncMock(return_value=events_batch),
|
||||
),
|
||||
):
|
||||
mock_pg.get_async_session_context = _make_pg_session_mock()
|
||||
mock_repo_cls.return_value = repo_mock
|
||||
|
||||
# Act
|
||||
results = [
|
||||
ev
|
||||
async for ev in adapter.streamAgentRunEvents(
|
||||
async for ev in _stream_structured_events(
|
||||
run_id="run-1",
|
||||
after_seq="",
|
||||
current_uid="uid-1",
|
||||
@ -150,88 +199,267 @@ class TestAgentRunExecutionAdapterStreamEvents:
|
||||
]
|
||||
|
||||
# Assert
|
||||
assert len(results) == 1
|
||||
assert len(results) == 2
|
||||
assert results[0]["event_type"] == "messages"
|
||||
assert results[0]["payload"] == {"k": "v"}
|
||||
assert results[0]["payload"] == {"k": 1}
|
||||
assert results[0]["seq"] == "1-0"
|
||||
assert results[1]["event_type"] == "end"
|
||||
assert results[1]["seq"] == "2-0"
|
||||
repo_mock.get_run_for_user.assert_awaited_once_with("run-1", "uid-1")
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stream_agent_run_events_returns_raw_when_no_data(self):
|
||||
# Arrange
|
||||
adapter = AgentRunExecutionAdapter()
|
||||
sse_text = "event: ping\n\n"
|
||||
async def test_stream_events_exits_on_end_event(self):
|
||||
# Arrange — 第一批含 end 事件，list_run_stream_events 不应被再次调用
|
||||
events_batch = [
|
||||
{"event_type": "messages", "payload": {}, "seq": "1-0"},
|
||||
{"event_type": "end", "payload": {}, "seq": "2-0"},
|
||||
]
|
||||
run_mock = MagicMock()
|
||||
run_mock.status = "running"
|
||||
repo_mock = MagicMock()
|
||||
repo_mock.get_run_for_user = AsyncMock(return_value=run_mock)
|
||||
|
||||
async def _fake_stream(*, run_id, after_seq, current_uid, verbose):
|
||||
yield sse_text
|
||||
|
||||
with patch(
|
||||
"yuxi.channels.adapters.agent_run_execution_adapter.stream_agent_run_events",
|
||||
new=_fake_stream,
|
||||
with (
|
||||
patch(
|
||||
"yuxi.channels.adapters.agent_run_execution_adapter.pg_manager"
|
||||
) as mock_pg,
|
||||
patch(
|
||||
"yuxi.channels.adapters.agent_run_execution_adapter.AgentRunRepository"
|
||||
) as mock_repo_cls,
|
||||
patch(
|
||||
"yuxi.channels.adapters.agent_run_execution_adapter.list_run_stream_events",
|
||||
new=AsyncMock(return_value=events_batch),
|
||||
) as mock_list,
|
||||
):
|
||||
mock_pg.get_async_session_context = _make_pg_session_mock()
|
||||
mock_repo_cls.return_value = repo_mock
|
||||
|
||||
# Act
|
||||
results = [
|
||||
ev
|
||||
async for ev in adapter.streamAgentRunEvents(
|
||||
async for ev in _stream_structured_events(
|
||||
run_id="run-1",
|
||||
after_seq="",
|
||||
current_uid="uid-1",
|
||||
)
|
||||
]
|
||||
|
||||
# Assert
|
||||
assert results[0]["event_type"] == "ping"
|
||||
assert results[0]["payload"] == {"text": sse_text}
|
||||
# Assert — 遇 end 立即退出，list_run_stream_events 仅被调用一次
|
||||
assert len(results) == 2
|
||||
assert mock_list.await_count == 1
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stream_agent_run_events_falls_back_to_payload_seq(self):
|
||||
# Arrange
|
||||
adapter = AgentRunExecutionAdapter()
|
||||
sse_text = 'data: {"seq": "9-9", "x": 1}\n\n'
|
||||
async def test_stream_events_raises_not_found_when_run_missing(self):
|
||||
# Arrange — get_run_for_user 返回 None
|
||||
repo_mock = MagicMock()
|
||||
repo_mock.get_run_for_user = AsyncMock(return_value=None)
|
||||
|
||||
async def _fake_stream(*, run_id, after_seq, current_uid, verbose):
|
||||
yield sse_text
|
||||
|
||||
with patch(
|
||||
"yuxi.channels.adapters.agent_run_execution_adapter.stream_agent_run_events",
|
||||
new=_fake_stream,
|
||||
with (
|
||||
patch(
|
||||
"yuxi.channels.adapters.agent_run_execution_adapter.pg_manager"
|
||||
) as mock_pg,
|
||||
patch(
|
||||
"yuxi.channels.adapters.agent_run_execution_adapter.AgentRunRepository"
|
||||
) as mock_repo_cls,
|
||||
):
|
||||
mock_pg.get_async_session_context = _make_pg_session_mock()
|
||||
mock_repo_cls.return_value = repo_mock
|
||||
|
||||
# Act / Assert
|
||||
with pytest.raises(NotFoundError):
|
||||
async for _ in _stream_structured_events(
|
||||
run_id="run-1",
|
||||
after_seq="",
|
||||
current_uid="uid-1",
|
||||
):
|
||||
pass
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stream_events_exits_when_run_terminal_without_end_event(self):
|
||||
# Arrange — Redis Stream 空，run 已终态（worker 崩溃场景）
|
||||
run_mock = MagicMock()
|
||||
run_mock.status = "failed" # 终态
|
||||
repo_mock = MagicMock()
|
||||
repo_mock.get_run_for_user = AsyncMock(return_value=run_mock)
|
||||
|
||||
with (
|
||||
patch(
|
||||
"yuxi.channels.adapters.agent_run_execution_adapter.pg_manager"
|
||||
) as mock_pg,
|
||||
patch(
|
||||
"yuxi.channels.adapters.agent_run_execution_adapter.AgentRunRepository"
|
||||
) as mock_repo_cls,
|
||||
patch(
|
||||
"yuxi.channels.adapters.agent_run_execution_adapter.list_run_stream_events",
|
||||
new=AsyncMock(return_value=[]), # Redis Stream 空
|
||||
),
|
||||
patch(
|
||||
"yuxi.channels.adapters.agent_run_execution_adapter.asyncio.sleep",
|
||||
new=AsyncMock(),
|
||||
),
|
||||
):
|
||||
mock_pg.get_async_session_context = _make_pg_session_mock()
|
||||
mock_repo_cls.return_value = repo_mock
|
||||
|
||||
# Act
|
||||
results = [
|
||||
ev
|
||||
async for ev in adapter.streamAgentRunEvents(
|
||||
async for ev in _stream_structured_events(
|
||||
run_id="run-1",
|
||||
after_seq="",
|
||||
current_uid="uid-1",
|
||||
)
|
||||
]
|
||||
|
||||
# Assert
|
||||
assert results[0]["seq"] == "9-9"
|
||||
assert results[0]["payload"] == {"seq": "9-9", "x": 1}
|
||||
# Assert — run 终态 + 无新事件，直接退出，不 yield 任何事件
|
||||
assert results == []
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stream_parsed_events_returns_raw_when_json_invalid(self):
|
||||
# Arrange
|
||||
async def _fake_stream(*, run_id, after_seq, current_uid, verbose):
|
||||
yield "data: not-json\n\n"
|
||||
async def test_stream_events_exits_when_elapsed_exceeds_max_connection_minutes(self):
|
||||
# Arrange — run 卡在 running 永不终态，模拟 utc_now_naive 推进超阈值
|
||||
run_mock = MagicMock()
|
||||
run_mock.status = "running" # 非终态，触发 elapsed 检查
|
||||
repo_mock = MagicMock()
|
||||
repo_mock.get_run_for_user = AsyncMock(return_value=run_mock)
|
||||
|
||||
# Act
|
||||
with patch(
|
||||
"yuxi.channels.adapters.agent_run_execution_adapter.stream_agent_run_events",
|
||||
new=_fake_stream,
|
||||
# utc_now_naive 调用序列：
|
||||
# call 1: 函数入口设置 started_at（T0）
|
||||
# call 2: 第一轮循环 elapsed 检查（T0 + 31min > 30min 阈值）
|
||||
base_time = datetime(2026, 1, 1, 12, 0, 0)
|
||||
expired_time = base_time + timedelta(minutes=31)
|
||||
|
||||
with (
|
||||
patch(
|
||||
"yuxi.channels.adapters.agent_run_execution_adapter.pg_manager"
|
||||
) as mock_pg,
|
||||
patch(
|
||||
"yuxi.channels.adapters.agent_run_execution_adapter.AgentRunRepository"
|
||||
) as mock_repo_cls,
|
||||
patch(
|
||||
"yuxi.channels.adapters.agent_run_execution_adapter.list_run_stream_events",
|
||||
new=AsyncMock(return_value=[]),
|
||||
),
|
||||
patch(
|
||||
"yuxi.channels.adapters.agent_run_execution_adapter.utc_now_naive",
|
||||
side_effect=[base_time, expired_time],
|
||||
),
|
||||
patch(
|
||||
"yuxi.channels.adapters.agent_run_execution_adapter.asyncio.sleep",
|
||||
new=AsyncMock(),
|
||||
),
|
||||
):
|
||||
mock_pg.get_async_session_context = _make_pg_session_mock()
|
||||
mock_repo_cls.return_value = repo_mock
|
||||
|
||||
# Act
|
||||
results = [
|
||||
ev
|
||||
async for ev in _stream_parsed_events(
|
||||
run_id="r",
|
||||
async for ev in _stream_structured_events(
|
||||
run_id="run-1",
|
||||
after_seq="",
|
||||
current_uid="u",
|
||||
current_uid="uid-1",
|
||||
)
|
||||
]
|
||||
|
||||
# Assert — elapsed 超阈值，安全阀触发，直接退出不 yield
|
||||
assert results == []
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stream_events_propagates_redis_error_as_dependency_error(self):
|
||||
# Arrange — list_run_stream_events 抛异常
|
||||
run_mock = MagicMock()
|
||||
run_mock.status = "running"
|
||||
repo_mock = MagicMock()
|
||||
repo_mock.get_run_for_user = AsyncMock(return_value=run_mock)
|
||||
|
||||
with (
|
||||
patch(
|
||||
"yuxi.channels.adapters.agent_run_execution_adapter.pg_manager"
|
||||
) as mock_pg,
|
||||
patch(
|
||||
"yuxi.channels.adapters.agent_run_execution_adapter.AgentRunRepository"
|
||||
) as mock_repo_cls,
|
||||
patch(
|
||||
"yuxi.channels.adapters.agent_run_execution_adapter.list_run_stream_events",
|
||||
new=AsyncMock(side_effect=RuntimeError("redis down")),
|
||||
),
|
||||
):
|
||||
mock_pg.get_async_session_context = _make_pg_session_mock()
|
||||
mock_repo_cls.return_value = repo_mock
|
||||
|
||||
# Act / Assert
|
||||
with pytest.raises(DependencyError):
|
||||
async for _ in _stream_structured_events(
|
||||
run_id="run-1",
|
||||
after_seq="",
|
||||
current_uid="uid-1",
|
||||
):
|
||||
pass
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stream_events_propagates_db_error_as_dependency_error(self):
|
||||
# Arrange — pg_manager.get_async_session_context 抛异常（非正常路径）
|
||||
@asynccontextmanager
|
||||
async def _raising_ctx():
|
||||
raise RuntimeError("db down")
|
||||
yield # pragma: no cover
|
||||
|
||||
with patch(
|
||||
"yuxi.channels.adapters.agent_run_execution_adapter.pg_manager"
|
||||
) as mock_pg:
|
||||
mock_pg.get_async_session_context = _raising_ctx
|
||||
|
||||
# Act / Assert
|
||||
with pytest.raises(DependencyError):
|
||||
async for _ in _stream_structured_events(
|
||||
run_id="run-1",
|
||||
after_seq="",
|
||||
current_uid="uid-1",
|
||||
):
|
||||
pass
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stream_agent_run_events_delegates_to_stream_structured_events(
|
||||
self,
|
||||
):
|
||||
# Arrange — 验证 streamAgentRunEvents 是同步方法返回异步迭代器
|
||||
adapter = AgentRunExecutionAdapter()
|
||||
events_batch = [
|
||||
{"event_type": "messages", "payload": {"x": 1}, "seq": "1-0"},
|
||||
{"event_type": "end", "payload": {}, "seq": "2-0"},
|
||||
]
|
||||
run_mock = MagicMock()
|
||||
run_mock.status = "completed"
|
||||
repo_mock = MagicMock()
|
||||
repo_mock.get_run_for_user = AsyncMock(return_value=run_mock)
|
||||
|
||||
with (
|
||||
patch(
|
||||
"yuxi.channels.adapters.agent_run_execution_adapter.pg_manager"
|
||||
) as mock_pg,
|
||||
patch(
|
||||
"yuxi.channels.adapters.agent_run_execution_adapter.AgentRunRepository"
|
||||
) as mock_repo_cls,
|
||||
patch(
|
||||
"yuxi.channels.adapters.agent_run_execution_adapter.list_run_stream_events",
|
||||
new=AsyncMock(return_value=events_batch),
|
||||
),
|
||||
):
|
||||
mock_pg.get_async_session_context = _make_pg_session_mock()
|
||||
mock_repo_cls.return_value = repo_mock
|
||||
|
||||
# Act — streamAgentRunEvents 同步调用返回异步迭代器
|
||||
iterator = adapter.streamAgentRunEvents(
|
||||
run_id="run-1",
|
||||
after_seq="",
|
||||
current_uid="uid-1",
|
||||
)
|
||||
results = [ev async for ev in iterator]
|
||||
|
||||
# Assert
|
||||
assert results[0]["event_type"] == "raw"
|
||||
assert "text" in results[0]["payload"]
|
||||
assert len(results) == 2
|
||||
assert results[0]["seq"] == "1-0"
|
||||
assert results[1]["event_type"] == "end"
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
|
||||
Loading…
Reference in New Issue
Block a user