diff --git a/apps/application/workflow/workflow_run_registry.py b/apps/application/workflow/workflow_run_registry.py index 5e240eb84c4..03a70184f77 100644 --- a/apps/application/workflow/workflow_run_registry.py +++ b/apps/application/workflow/workflow_run_registry.py @@ -1,17 +1,20 @@ # coding=utf-8 """ - @project: MaxKB - @file： workflow_run_registry.py - @desc: 工作流运行注册表，用于管理和取消正在运行的工作流实例 +@project: MaxKB +@file： workflow_run_registry.py +@desc: 工作流运行注册表，用于管理和取消正在运行的工作流实例 """ + import threading from enum import Enum from common.utils.logger import maxkb_logger +from common.bus import on class CancelResult(Enum): """取消操作结果""" + CANCELLED = "CANCELLED" NOT_FOUND = "NOT_FOUND" FAILED = "FAILED" @@ -166,3 +169,6 @@ def running_ids(cls) -> list: """ with cls._lock: return list(cls._running.keys()) + + +on("cancel_chat", lambda args: WorkflowRunRegistry.cancel_by_chat_id(args.get("chat_id"))) diff --git a/apps/chat/views/v3/chat.py b/apps/chat/views/v3/chat.py index 7bbbd7af6ad..83298c54ab7 100644 --- a/apps/chat/views/v3/chat.py +++ b/apps/chat/views/v3/chat.py @@ -19,7 +19,7 @@ from rest_framework.parsers import MultiPartParser from rest_framework.request import Request from rest_framework.views import APIView - +from common.bus import emit from application.api.application_api import SpeechToTextAPI, TextToSpeechAPI from application.models import ChatUserType, ChatSourceChoices, ApplicationAccessToken, ChatRecord from chat.api.chat_api import ChatAPI @@ -285,15 +285,8 @@ class CancelWorkflowView(APIView): tags=[_("V3 Chat")], # type: ignore ) def post(self, request: Request, chat_id: str): - from application.workflow.workflow_run_registry import WorkflowRunRegistry, CancelResult - - result_enum = WorkflowRunRegistry.cancel_by_chat_id(chat_id) - if result_enum == CancelResult.CANCELLED: - return result.success({"status": "cancelled", "chat_id": chat_id}) - elif result_enum == CancelResult.NOT_FOUND: - return result.success({"status": "not_found", "chat_id": chat_id}) - else: - return result.fail(500, _("Failed to cancel workflow")) + emit("cancel_chat", {"chat_id": str(chat_id)}) + return result.success({"status": "cancelled", "chat_id": chat_id}) class ResumeStreamView(APIView): diff --git a/apps/common/bus/__init__.py b/apps/common/bus/__init__.py new file mode 100644 index 00000000000..b13cf4c889b --- /dev/null +++ b/apps/common/bus/__init__.py @@ -0,0 +1 @@ +from common.bus.message_bus import on, emit diff --git a/apps/common/bus/message_bus.py b/apps/common/bus/message_bus.py new file mode 100644 index 00000000000..fbd751e8b0e --- /dev/null +++ b/apps/common/bus/message_bus.py @@ -0,0 +1,109 @@ +# coding=utf-8 +""" +@project: MaxKB +@file： message_bus.py +@desc: 基于 Redis Pub/Sub 的跨进程、跨系统消息总线。提供 emit / on 两个接口。 +""" + +import json +import threading +from typing import Callable, Dict, List + +from common.utils.logger import maxkb_logger + +CHANNEL = "maxkb:bus" + +_handlers: Dict[str, List[Callable]] = {} +_handlers_lock = threading.Lock() + +_subscriber_started = False +_subscriber_thread = None +_subscriber_lock = threading.Lock() + + +def _redis_client(): + """取底层 redis-py 客户端（含 pubsub 能力），已由 Django settings 配置支持 sentinel/普通 Redis。""" + from django_redis import get_redis_connection + + return get_redis_connection("default") + + +def _start_subscriber() -> None: + """懒启动：每个进程只起一条 daemon 订阅线程；线程因建连失败退出后，下次调用会重新拉起。""" + global _subscriber_started, _subscriber_thread + if _subscriber_started and _subscriber_thread is not None and _subscriber_thread.is_alive(): + return + with _subscriber_lock: + if _subscriber_started and _subscriber_thread is not None and _subscriber_thread.is_alive(): + return + _subscriber_started = True + _subscriber_thread = threading.Thread(target=_listen, name="maxkb-bus", daemon=True) + _subscriber_thread.start() + + +def _listen() -> None: + # 建连失败(Redis 不可用)则记日志退出,不自动重连;由 emit 检测线程存活后重新拉起。 + # 建连成功后的消息循环由 redis-py 内部维持,无需自行重连。 + ps = None + try: + ps = _redis_client().pubsub(ignore_subscribe_messages=True) + ps.subscribe(CHANNEL) + for message in ps.listen(): + _dispatch(message) + except Exception as e: + maxkb_logger.error(f"Message bus subscriber failed: {e}") + finally: + if ps is not None: + try: + ps.close() + except Exception: + pass + + +def _dispatch(message) -> None: + try: + data = message.get("data") + if isinstance(data, (bytes, bytearray)): + data = data.decode("utf-8") + body = json.loads(data) + event = body.get("event") + if not event: + return + args = body.get("args") + for handler in list(_handlers.get(event, ())): + try: + handler(args) + except Exception as e: + maxkb_logger.error(f"Message bus handler error for event[{event}]: {e}") + except Exception as e: + maxkb_logger.error("Failed to dispatch message bus event: %s", e) + + +def on(event: str, handler: Callable) -> None: + """ + 注册一个事件监听，本进程内生效；收到该事件时异步调用 handler(args)。 + :param event: 事件名称 + :param handler: 回调函数 def handler(args): + pass + """ + with _handlers_lock: + _handlers.setdefault(event, []).append(handler) + _start_subscriber() + + +def emit(event: str, args=None) -> bool: + """ + 向所有进程广播一个事件。Redis 不可用时返回 False。 + :param event: 事件名称 + :param args: 事件参数 + """ + _start_subscriber() + try: + _redis_client().publish( + CHANNEL, + json.dumps({"event": event, "args": args}, ensure_ascii=False), + ) + return True + except Exception as e: + maxkb_logger.warning(f"Message bus emit[{event}] failed: {e}") + return False