From 0ecb815b4ae6558f1cea5a92cd9940b1a9727dbf Mon Sep 17 00:00:00 2001 From: rookit Date: Thu, 14 May 2026 09:16:23 +0800 Subject: [PATCH] fix playbook framework bug --- Lib/baseplaybook.py | 8 ++--- Lib/montior.py | 71 +++++++++++++++++++++++++++++++++----- Lib/threadmodulemanager.py | 31 +++++++++++++---- PLAYBOOKS/Investigation.py | 3 +- 4 files changed, 91 insertions(+), 22 deletions(-) diff --git a/Lib/baseplaybook.py b/Lib/baseplaybook.py index 8493e43..74bc615 100644 --- a/Lib/baseplaybook.py +++ b/Lib/baseplaybook.py @@ -10,7 +10,6 @@ from PLUGINS.SIRP.sirpextramodel import PlaybookJobStatus, PlaybookModel class BasePlaybook(BaseAPI): - RUN_AS_JOB = True # 是否作为后台任务运行 NAME = None def __init__(self): @@ -40,11 +39,8 @@ class BasePlaybook(BaseAPI): return result def execute(self): - try: - self.run() - except Exception as e: - self.logger.exception(e) - self.update_playbook_status(PlaybookJobStatus.FAILED, str(e)) + result = self.run() + return result class LanggraphPlaybook(BasePlaybook): diff --git a/Lib/montior.py b/Lib/montior.py index 5a3f529..2a5f8e1 100644 --- a/Lib/montior.py +++ b/Lib/montior.py @@ -3,7 +3,8 @@ import importlib import threading import time -from typing import Callable +import uuid +from typing import Callable, Dict, Optional from apscheduler.schedulers.background import BackgroundScheduler @@ -20,6 +21,48 @@ from PLUGINS.SIRP.sirpextramodel import PlaybookJobStatus, PlaybookModel class MainMonitor(object): + @staticmethod + def on_playbook_task_finished(thread_id: str, task_obj: BasePlaybook, result: object, + exception: Optional[Exception], context: Optional[Dict[str, str]] = None): + context = context or {} + playbook_row_id = context.get("playbook_row_id") + if not isinstance(playbook_row_id, str) or playbook_row_id == "": + logger.error(f"[Thread {thread_id}] Missing playbook_row_id in callback context.") + return + + try: + playbook_current = Playbook.get(playbook_row_id, lazy_load=True) + except Exception as e: + logger.error(f"[Thread {thread_id}] Failed to load playbook for fallback status update.") + logger.exception(e) + return + + # Business code has already written a terminal state; do not overwrite. + if playbook_current.job_status != PlaybookJobStatus.RUNNING: + logger.info( + f"[Thread {thread_id}] Skip fallback status update, current status: {playbook_current.job_status}, " + f"row_id: {playbook_row_id}" + ) + return + + fallback_status = PlaybookJobStatus.SUCCESS if exception is None else PlaybookJobStatus.FAILED + if exception is None: + remark = f"source=thread_fallback; thread_id={thread_id}; message=Task finished without explicit business status update." + else: + remark = ( + f"source=thread_fallback; thread_id={thread_id}; " + f"exception={type(exception).__name__}; message={exception}" + ) + + model_tmp = PlaybookModel(row_id=playbook_row_id) + model_tmp.job_status = fallback_status + model_tmp.remark = remark + Playbook.update(model_tmp) + + logger.info( + f"[Thread {thread_id}] Applied fallback status update: {fallback_status}, row_id: {playbook_row_id}" + ) + MainScheduler: BackgroundScheduler _background_threads = {} @@ -119,17 +162,27 @@ class MainMonitor(object): Playbook.update(model_tmp) continue - job_id = thread_module_manager.start_task(playbook_intent) - if not job_id: + job_id = str(uuid.uuid1()) + model_tmp.job_status = PlaybookJobStatus.RUNNING + model_tmp.job_id = job_id + Playbook.update(model_tmp) + + try: + thread_module_manager.start_task( + playbook_intent, + thread_id=job_id, + on_finished=MainMonitor.on_playbook_task_finished, + callback_context={ + "playbook_row_id": model.row_id, + }, + ) + except Exception as e: model_tmp.job_status = PlaybookJobStatus.FAILED - model_tmp.remark = "Failed to create playbook job." + model_tmp.remark = f"Failed to create playbook job. exception={type(e).__name__}; message={e}" Playbook.update(model_tmp) continue - else: - logger.info(f"Create playbook job success: {job_id}") - model_tmp.job_status = PlaybookJobStatus.RUNNING - model_tmp.job_id = job_id - Playbook.update(model_tmp) + + logger.info(f"Create playbook job success: {job_id}") @staticmethod def subscribe_case_analysis_scheduler(): diff --git a/Lib/threadmodulemanager.py b/Lib/threadmodulemanager.py index c318638..5291e37 100644 --- a/Lib/threadmodulemanager.py +++ b/Lib/threadmodulemanager.py @@ -6,11 +6,16 @@ import threading import time import uuid from enum import Enum -from typing import Optional, Callable, Dict, Any +from typing import Optional, Callable, Dict, Protocol from Lib.log import logger +class TaskExecutable(Protocol): + def execute(self) -> object: + ... + + class ThreadStatus(Enum): PENDING = "pending" RUNNING = "running" @@ -25,7 +30,7 @@ class ThreadInfo: self.status = ThreadStatus.PENDING self.start_time: Optional[float] = None self.end_time: Optional[float] = None - self.result: Any = None + self.result: object = None self.exception: Optional[Exception] = None def get_duration(self) -> Optional[float]: @@ -54,7 +59,8 @@ class ThreadModuleManager: self._next_auto_id += 1 return f"thread_{self._next_auto_id}" - def _run_task(self, thread_id: str, task_obj) -> None: + def _run_task(self, thread_id: str, task_obj: TaskExecutable, on_finished: Optional[Callable] = None, + callback_context: Optional[Dict[str, str]] = None) -> None: thread_info = self._threads[thread_id] thread_info.status = ThreadStatus.RUNNING thread_info.start_time = time.time() @@ -71,14 +77,27 @@ class ThreadModuleManager: logger.error(f"[Thread {thread_id}] Task failed with exception.") finally: thread_info.end_time = time.time() + if on_finished: + try: + on_finished( + thread_id=thread_id, + task_obj=task_obj, + result=thread_info.result, + exception=thread_info.exception, + context=callback_context, + ) + except Exception as callback_exception: + logger.error(f"[Thread {thread_id}] on_finished callback failed: {type(callback_exception).__name__}") + logger.exception(callback_exception) - def start_task(self, task_obj, thread_id: Optional[str] = None) -> str: + def start_task(self, task_obj: TaskExecutable, thread_id: Optional[str] = None, + on_finished: Optional[Callable] = None, callback_context: Optional[Dict[str, str]] = None) -> str: if thread_id is None: thread_id = str(uuid.uuid1()) thread = threading.Thread( target=self._run_task, - args=(thread_id, task_obj), + args=(thread_id, task_obj, on_finished, callback_context), name=thread_id, daemon=False ) @@ -101,7 +120,7 @@ class ThreadModuleManager: with self._lock: return self._threads.get(thread_id) - def get_result(self, thread_id: str) -> Any: + def get_result(self, thread_id: str) -> object: with self._lock: thread_info = self._threads.get(thread_id) if thread_info and thread_info.status == ThreadStatus.COMPLETED: diff --git a/PLAYBOOKS/Investigation.py b/PLAYBOOKS/Investigation.py index 8aeff7c..2a0a7b3 100644 --- a/PLAYBOOKS/Investigation.py +++ b/PLAYBOOKS/Investigation.py @@ -9,7 +9,7 @@ from PLUGINS.SIRP.analysis import ( ) from PLUGINS.SIRP.sirpapi import Case from PLUGINS.SIRP.sirpbasemodel import AI_PROFILE_INVESTIGATION -from PLUGINS.SIRP.sirpextramodel import PlaybookModel +from PLUGINS.SIRP.sirpextramodel import PlaybookModel, PlaybookJobStatus class Playbook(BasePlaybook): @@ -46,6 +46,7 @@ class Playbook(BasePlaybook): Case.update(case_patch) self.logger.info(f"Case analysis completed. row_id: {case_row_id}, trigger: {trigger}") + self.update_playbook_status(PlaybookJobStatus.SUCCESS, "Case Investigation Success.") if __name__ == "__main__":