6a2632da70
将 Clover 从上层产品包旧仓库中独立出来,建立专属版本控制。 当前状态=纵切片端到端已打通(登录→选品→出文出图→审核→下载包), M1文案质量去套路化已验收。此提交作为后续按核销清单逐条修复的基线。 Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
123 lines
3.9 KiB
Python
123 lines
3.9 KiB
Python
"""
|
||
app/utils/sse_utils.py — SSE 工具函数（按 Lead 规范路径）
|
||
补发历史事件、event_seq 去重、Redis pub/sub 推送。
|
||
供 api/v1/stream.py 和 AIE worker 共用。
|
||
"""
|
||
|
||
import asyncio
|
||
import json
|
||
import logging
|
||
from typing import Any, AsyncGenerator
|
||
|
||
from app.core.config import get_settings
|
||
|
||
logger = logging.getLogger(__name__)
|
||
settings = get_settings()
|
||
|
||
# SSE 事件类型（契约§2全列）
|
||
SSE_EVENT_TYPES = frozenset({
|
||
"task_started", "analyze_done", "text_progress", "text_candidate",
|
||
"image_progress", "image_candidate", "flywheel_injected",
|
||
"batch_failed", "task_done", "error", "heartbeat",
|
||
})
|
||
|
||
_HISTORY_KEY_TPL = "sse:task:{task_id}:events"
|
||
_CHANNEL_TPL = "sse:task:{task_id}"
|
||
_HISTORY_TTL_SECONDS = 3600 # 历史事件保留 1h
|
||
|
||
|
||
def format_sse(event: str, data: dict, seq: int | None = None) -> str:
|
||
"""格式化单条 SSE 消息（text/event-stream 格式）。"""
|
||
payload = json.dumps(data, ensure_ascii=False)
|
||
lines = [f"event: {event}", f"data: {payload}"]
|
||
if seq is not None:
|
||
lines.append(f"id: {seq}")
|
||
return "\n".join(lines) + "\n\n"
|
||
|
||
|
||
async def push_event(
|
||
redis_client,
|
||
task_id: int,
|
||
workspace_id: int,
|
||
event: str,
|
||
data: dict[str, Any],
|
||
seq: int,
|
||
) -> None:
|
||
"""
|
||
推送事件到 Redis：
|
||
1. 追加到历史 list（供断线重连补发）
|
||
2. Publish 到 channel（供在线客户端实时收）
|
||
"""
|
||
if event not in SSE_EVENT_TYPES:
|
||
logger.warning("Unknown SSE event type: %s", event)
|
||
|
||
record = {"event": event, "data": data, "seq": seq, "workspace_id": workspace_id}
|
||
payload = json.dumps(record, ensure_ascii=False)
|
||
|
||
hist_key = _HISTORY_KEY_TPL.format(task_id=task_id)
|
||
channel = _CHANNEL_TPL.format(task_id=task_id)
|
||
|
||
await redis_client.rpush(hist_key, payload)
|
||
await redis_client.expire(hist_key, _HISTORY_TTL_SECONDS)
|
||
await redis_client.publish(channel, payload)
|
||
|
||
|
||
async def get_history_events(
|
||
redis_client, task_id: int, after_seq: int
|
||
) -> list[dict]:
|
||
"""取 task 历史事件中 seq > after_seq 的部分（断线重连补发）。"""
|
||
hist_key = _HISTORY_KEY_TPL.format(task_id=task_id)
|
||
try:
|
||
raw_list = await redis_client.lrange(hist_key, 0, -1)
|
||
result = []
|
||
for raw in raw_list:
|
||
ev = json.loads(raw)
|
||
if ev.get("seq", 0) > after_seq:
|
||
result.append(ev)
|
||
return result
|
||
except Exception as exc:
|
||
logger.warning("get_history_events failed task=%s: %s", task_id, exc)
|
||
return []
|
||
|
||
|
||
async def stream_events(
|
||
redis_client,
|
||
task_id: int,
|
||
workspace_id: int,
|
||
last_seq: int = 0,
|
||
heartbeat_interval: int = 25,
|
||
) -> AsyncGenerator[str, None]:
|
||
"""
|
||
主 SSE 生成器：
|
||
1. 先补发 last_seq 之后的历史事件
|
||
2. 再订阅 Redis channel 实时推新事件
|
||
3. 每 heartbeat_interval 秒推一次保活 heartbeat
|
||
"""
|
||
# 1. 补发历史
|
||
history = await get_history_events(redis_client, task_id, last_seq)
|
||
for ev in history:
|
||
yield format_sse(ev["event"], ev["data"], ev.get("seq"))
|
||
|
||
# 2. 实时订阅
|
||
pubsub = redis_client.pubsub()
|
||
channel = _CHANNEL_TPL.format(task_id=task_id)
|
||
await pubsub.subscribe(channel)
|
||
elapsed = 0
|
||
try:
|
||
while True:
|
||
msg = await pubsub.get_message(ignore_subscribe_messages=True, timeout=1.0)
|
||
if msg and msg["type"] == "message":
|
||
ev = json.loads(msg["data"])
|
||
if ev.get("workspace_id") != workspace_id:
|
||
continue # 防越权订阅他人任务
|
||
yield format_sse(ev["event"], ev["data"], ev.get("seq"))
|
||
if ev["event"] in ("task_done", "error"):
|
||
break
|
||
else:
|
||
elapsed += 1
|
||
if elapsed >= heartbeat_interval:
|
||
elapsed = 0
|
||
yield format_sse("heartbeat", {"ts": asyncio.get_event_loop().time()})
|
||
finally:
|
||
await pubsub.unsubscribe(channel)
|