server/process_manager.py

"""Управление потоками и отмена генерации"""

import threading
from typing import List, Dict
from concurrent.futures import ThreadPoolExecutor

class ProcessManager:
    def __init__(self):
        self.active_threads: List[threading.Thread] = []
        self.cancel_flags: Dict[int, bool] = {}
        self.lock = threading.Lock()
        self.executor = ThreadPoolExecutor(max_workers=10)
        self.futures = []

    def register_thread(self, thread: threading.Thread) -> None:
        with self.lock:
            self.active_threads.append(thread)
            self.cancel_flags[thread.ident] = False

    def register_future(self, future) -> None:
        with self.lock:
            self.futures.append(future)

    def cancel_all(self) -> str:
        with self.lock:
            for future in self.futures:
                if not future.done():
                    future.cancel()
            self.futures.clear()
            for thread_id in self.cancel_flags:
                self.cancel_flags[thread_id] = True
            for thread in self.active_threads:
                if thread.is_alive():
                    try:
                        thread.join(timeout=0.5)
                    except Exception:
                        pass
            self.active_threads.clear()
            self.cancel_flags.clear()

        try:
            from interpreter import interpreter
            if hasattr(interpreter, 'cancel'):
                interpreter.cancel()
        except Exception:
            pass

        from .state import STATE
        STATE.cancel_flag = True
        STATE.current_mode = "paused"
        return "✅ Все процессы генерации остановлены!"

    def is_cancelled(self, thread_id: int = None) -> bool:
        if thread_id is None:
            thread_id = threading.current_thread().ident
        with self.lock:
            return self.cancel_flags.get(thread_id, False)

    def clear(self) -> None:
        with self.lock:
            self.active_threads = [t for t in self.active_threads if t.is_alive()]

    def get_active_count(self) -> int:
        with self.lock:
            return len([t for t in self.active_threads if t.is_alive()])

PROCESS_MANAGER = ProcessManager()