long_articles_agentic_supply/supply/domain/task/base.py

"""Task 领域 —— 任务抽象、处理器基类、注册表

所有任务概念集中于此，替代旧的分散在 jobs/task_*.py 的全局注册表模式。
"""

from __future__ import annotations

from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from typing import Any, Callable, Dict, List, Optional

from supply.domain.task.config import TaskConfig, TaskStatus, get_default_task_config
from supply.domain.task.exceptions import (
    TaskError,
    TaskValidationError,
    TaskTimeoutError,
    TaskConcurrencyError,
    TaskLockError,
)


# ==================== Task / TaskResult ====================


@dataclass
class Task:
    """任务数据模型"""

    task_name: str
    params: Dict[str, Any] = field(default_factory=dict)
    trace_id: str = ""
    config: Optional[TaskConfig] = None


@dataclass
class TaskResult:
    """任务执行结果"""

    success: bool
    data: Any = None
    error: Optional[str] = None
    trace_id: str = ""

    @classmethod
    def ok(cls, data: Any = None, trace_id: str = "") -> "TaskResult":
        return cls(success=True, data=data, trace_id=trace_id)

    @classmethod
    def fail(cls, error: str, trace_id: str = "") -> "TaskResult":
        return cls(success=False, error=error, trace_id=trace_id)


# ==================== TaskRegistry（实例级） ====================


class TaskRegistry:
    """任务注册表 —— 实例级，替代全局 _TASK_HANDLER_REGISTRY

    使用方式：
        registry = TaskRegistry()

        @registry.register("content_generation")
        class ContentGenHandler(TaskHandler):
            ...

        handler = registry.get("content_generation")
        result = await handler.execute(task)
    """

    def __init__(self):
        self._handlers: Dict[str, type] = {}

    def register(self, task_name: str) -> Callable:
        """装饰器：将 TaskHandler 子类注册到注册表

        使用方式：
            @registry.register("my_task")
            class MyHandler(TaskHandler):
                ...
        """

        def decorator(handler_cls: type) -> type:
            if not issubclass(handler_cls, TaskHandler):
                raise TypeError(
                    f"@register 只能用于 TaskHandler 子类，收到 {handler_cls}"
                )
            if task_name in self._handlers:
                raise ValueError(f"任务 '{task_name}' 已注册")
            self._handlers[task_name] = handler_cls
            handler_cls.task_name = task_name
            return handler_cls

        return decorator

    def get(self, task_name: str) -> Optional[type]:
        return self._handlers.get(task_name)

    def list_names(self) -> List[str]:
        return list(self._handlers.keys())

    def __len__(self) -> int:
        return len(self._handlers)

    def __contains__(self, task_name: str) -> bool:
        return task_name in self._handlers


# ==================== TaskHandler 抽象基类 ====================


class TaskHandler(ABC):
    """任务处理器抽象基类

    业务方继承此类实现具体任务逻辑：

        class ContentGenHandler(TaskHandler):
            task_name = "content_generation"

            async def validate(self, task: Task) -> bool:
                return "topic" in task.params

            async def execute(self, task: Task) -> TaskResult:
                ...

    task_name 类属性由 @TaskRegistry.register 装饰器自动设置。
    """

    task_name: str = ""

    def __init__(
        self,
        log_service: Any = None,
        db_client: Any = None,
    ):
        self.log_service = log_service
        self.db_client = db_client

    @abstractmethod
    async def execute(self, task: Task) -> TaskResult:
        """执行任务主逻辑"""
        ...

    async def validate(self, task: Task) -> bool:
        """校验任务参数，默认通过。子类可按需覆写"""
        return True

    async def on_success(self, task: Task) -> None:
        """任务成功回调（可选覆写）"""
        pass

    async def on_failure(self, task: Task, error: Exception) -> None:
        """任务失败回调（可选覆写）"""
        pass

    def __repr__(self) -> str:
        return f"<TaskHandler name={self.task_name}>"