本文介绍如何在 WebSocket 连接收到停止请求（stop_requested）时，立即取消正在执行的异步 AI 生成任务，避免等待 client.images.generate() 等长耗时协程完成，从而实现真正的即时响应与资源释放。
本文介绍如何在 websocket 连接收到停止请求（`stop_requested`）时，**立即取消正在执行的异步 ai 生成任务**，避免等待 `client.images.generate()` 等长耗时协程完成，从而实现真正的即时响应与资源释放。
在异步 Web 服务（如基于 FastAPI + WebSockets 的图像生成接口）中，一个常见痛点是：用户触发“中断”后，后端仍需等待当前 await client.images.generate(...) 完成才响应停止逻辑——这不仅造成延迟，还浪费 API 配额与计算资源。根本原因在于：原生 HTTP/WebSocket 协议不提供对运行中协程的强制取消能力，而 asyncio 的任务取消必须由协程主动协作或通过可等待对象显式监听中断信号。
✅ 正确方案：使用 asyncio.Event 实现协作式取消
推荐采用 asyncio.Event 作为轻量、高效、标准的同步原语。它天然支持 await event.wait()，且可在任意位置（如 WebSocket 断开处理逻辑）调用 event.set() 触发所有等待者退出。
以下是重构后的完整示例：
WebSocket 8.18.2
下载
WebSocket 8.18.2 是该协议规范的一个重要迭代版本，主要优化了连接稳定性与数据传输效率。它通过全双工通信机制，允许客户端与服务器在单一长连接上实时交换数据，大幅降低传统 HTTP 轮询的开销。该版本增强了心跳保活、自动重连及二进制帧传输能力，适用于即时通讯、在线游戏及金融行情推送等低延迟场景，为开发者提供更可靠的实时网络交互基础。
import asyncio
from typing import Optional, Any
class ImageGeneration(BaseAIGeneration):
async def process(self, websocket, prompt: str) -> Optional[str]:
# 创建与本次请求绑定的取消事件（建议在 process 开头创建）
stop_ev = asyncio.Event()
# 启动生成任务，并传入事件用于监听中断
task = asyncio.create_task(
self.generate_image_with_cancellation(prompt, self.model, "1024x1024", stop_ev)
)
try:
# 等待生成完成，或被 stop_ev 中断
image_url = await task
return image_url
except asyncio.CancelledError:
# 任务被显式取消（例如 stop_ev.set() + await stop_ev.wait() 中断）
print("Image generation cancelled by user.")
return None
except Exception as e:
print(f"Generation error: {e}")
return None
async def generate_image_with_cancellation(
self,
prompt: str,
model: str,
size: str = "1024x1024",
stop_event: Optional[asyncio.Event] = None
) -> Optional[str]:
# 第一步：检查初始状态（快速响应）
if stop_event and stop_event.is_set():
return None
# 第二步：发起 OpenAI 异步调用，但包装为可取消的等待
try:
# 使用 asyncio.wait_for + timeout（可选），但核心是监听 stop_event
done, pending = await asyncio.wait(
[self._raw_generate(prompt, model, size), stop_event.wait()],
return_when=asyncio.FIRST_COMPLETED
)
for t in pending:
t.cancel() # 取消未完成的等待（如 stop_event.wait()）
for res in done:
result = res.result()
if isinstance(result, Exception):
raise result
if stop_event and stop_event.is_set():
return None # 中断成功
# 否则 result 就是 generate 返回值
if isinstance(result, list) and len(result) > 0:
return result[0].url
return None
except asyncio.CancelledError:
return None
async def _raw_generate(self, prompt: str, model: str, size: str):
response = await client.images.generate(
model=model,
prompt=prompt,
size=size,
n=1
)
return response.data
? 关键改进点说明
-
无轮询、无
sleep：摒弃低效的while not task.done(): await asyncio.sleep(0.1)循环，改用asyncio.wait(..., FIRST_COMPLETED)让事件驱动代替主动轮询。 -
职责分离清晰：
stop_event由 WebSocket 生命周期管理（例如在websocket.disconnect()或自定义stop()方法中调用stop_event.set()），业务逻辑只负责监听和响应。 -
异常安全：捕获
CancelledError并优雅返回，避免未处理异常导致连接崩溃。 -
可扩展性强：同一
stop_event可同时用于多个嵌套await（如日志上传、后处理等），统一控制流。
⚠️ 注意事项
-
stop_event必须是每次请求独立创建的新实例，不可复用或全局共享，否则不同连接间会相互干扰； - 若底层
client.images.generate内部未正确传播取消信号（如使用了不支持取消的 HTTP 库），可能仍会短暂阻塞；建议确认其是否基于httpx.AsyncClient或aiohttp等原生支持asyncio取消的客户端； - 生产环境建议添加超时兜底：
await asyncio.wait_for(task, timeout=60.0)，防止因网络问题彻底卡死。
通过以上设计，你将获得真正响应迅速、资源可控、符合 asyncio 最佳实践的 WebSocket 生成中断能力。
