microfish/backend/app/models/task.py

"""Durable task status management with a test-only in-memory fallback.

The Flask application configures a SQLAlchemy session factory at startup. Code
that uses TaskManager outside an application (small unit tests and legacy
adapters) keeps the old in-memory behavior, but production requests do not.
"""

from __future__ import annotations

import threading
import uuid
from dataclasses import dataclass, field
from datetime import datetime, timedelta, timezone
from enum import Enum
from typing import Any, Dict, Optional, cast

from flask import current_app, has_app_context
from sqlalchemy import delete, select

from ..models.operations import Job, JobStatus
from ..utils.locale import t


class TaskStatus(str, Enum):
    PENDING = "pending"
    PROCESSING = "processing"
    COMPLETED = "completed"
    FAILED = "failed"


@dataclass
class Task:
    task_id: str
    task_type: str
    status: TaskStatus
    created_at: datetime
    updated_at: datetime
    progress: int = 0
    message: str = ""
    result: Optional[Dict] = None
    error: Optional[str] = None
    metadata: Dict = field(default_factory=dict)
    progress_detail: Dict = field(default_factory=dict)

    def to_dict(self) -> Dict[str, Any]:
        return {
            "task_id": self.task_id,
            "task_type": self.task_type,
            "status": self.status.value,
            "created_at": self.created_at.isoformat(),
            "updated_at": self.updated_at.isoformat(),
            "progress": self.progress,
            "message": self.message,
            "progress_detail": self.progress_detail,
            "result": self.result,
            "error": self.error,
            "metadata": self.metadata,
        }


class TaskManager:
    """Thread-safe task facade bound to one app/session factory."""

    _configured_session_factory = None
    _config_lock = threading.Lock()
    _fallback_tasks: Dict[str, Task] = {}
    _fallback_lock = threading.Lock()

    def __init__(self, session_factory=None):
        """Bind this manager to an explicit or current-app session factory.

        A manager created during a request/app context keeps that app's factory
        for background work, but it cannot be reused inside a different Flask
        app. The explicit class configuration remains only for legacy tests and
        callers that run outside Flask.
        """
        self._bound_app = None
        if has_app_context():
            app = cast(Any, current_app)._get_current_object()
            current_factory = app.extensions.get("crowdsight_session_factory")
            if not callable(current_factory):
                raise RuntimeError("task_session_factory_required")
            if session_factory is not None and session_factory is not current_factory:
                raise RuntimeError("task_session_factory_mismatch")
            bound_factory = current_factory
            self._bound_app = app
        elif session_factory is not None:
            bound_factory = session_factory
        else:
            # No Flask app context: use the explicit legacy/test binding.
            bound_factory = type(self)._configured_session_factory

        self._session_factory = bound_factory
        self._tasks = type(self)._fallback_tasks
        self._task_lock = type(self)._fallback_lock

    @classmethod
    def configure(cls, session_factory) -> None:
        """Set an explicit outside-Flask binding for tests/legacy adapters."""
        with cls._config_lock:
            cls._configured_session_factory = session_factory
            with cls._fallback_lock:
                cls._fallback_tasks.clear()

    def _factory(self):
        if not has_app_context():
            return self._session_factory

        app = cast(Any, current_app)._get_current_object()
        current_factory = app.extensions.get("crowdsight_session_factory")
        if not callable(current_factory):
            raise RuntimeError("task_session_factory_required")
        if self._bound_app is not None and self._bound_app is not app:
            raise RuntimeError("task_app_context_mismatch")
        if self._session_factory is not None and self._session_factory is not current_factory:
            raise RuntimeError("task_session_factory_mismatch")
        if self._session_factory is None:
            self._session_factory = current_factory
        self._bound_app = app
        return self._session_factory

    @staticmethod
    def _bounded_text(value: Any, limit: int = 4000) -> str:
        if value is None:
            return ""
        return str(value)[:limit]

    @staticmethod
    def _job_status(status: TaskStatus | str | None) -> str | None:
        if status is None:
            return None
        value = status.value if isinstance(status, TaskStatus) else str(status)
        return {
            TaskStatus.PENDING.value: JobStatus.QUEUED.value,
            TaskStatus.PROCESSING.value: JobStatus.RUNNING.value,
            TaskStatus.COMPLETED.value: JobStatus.SUCCEEDED.value,
            TaskStatus.FAILED.value: JobStatus.FAILED.value,
            JobStatus.QUEUED.value: JobStatus.QUEUED.value,
            JobStatus.RUNNING.value: JobStatus.RUNNING.value,
            JobStatus.SUCCEEDED.value: JobStatus.SUCCEEDED.value,
            JobStatus.FAILED.value: JobStatus.FAILED.value,
            JobStatus.CANCELLED.value: JobStatus.CANCELLED.value,
        }.get(value)

    @staticmethod
    def _task_status(status: str | JobStatus) -> TaskStatus:
        value = status.value if isinstance(status, JobStatus) else str(status)
        return {
            JobStatus.QUEUED.value: TaskStatus.PENDING,
            JobStatus.RUNNING.value: TaskStatus.PROCESSING,
            JobStatus.SUCCEEDED.value: TaskStatus.COMPLETED,
            JobStatus.FAILED.value: TaskStatus.FAILED,
            JobStatus.CANCELLED.value: TaskStatus.FAILED,
        }.get(value, TaskStatus.FAILED)

    @classmethod
    def _from_job(cls, job: Job) -> Task:
        metadata = job.job_metadata if isinstance(job.job_metadata, dict) else {}
        result = job.result if isinstance(job.result, dict) else job.result
        detail = job.progress_detail if isinstance(job.progress_detail, dict) else {}
        return Task(
            task_id=job.id,
            task_type=job.operation,
            status=cls._task_status(job.status),
            created_at=job.created_at,
            updated_at=job.updated_at,
            progress=job.progress,
            message=job.message,
            result=result,
            error=job.error_code,
            metadata=metadata,
            progress_detail=detail,
        )

    def create_task(self, task_type: str, metadata: Optional[Dict] = None) -> str:
        metadata = metadata or {}
        factory = self._factory()
        if factory is not None:
            organization_id = metadata.get("organization_id")
            if not isinstance(organization_id, str) or not organization_id:
                raise ValueError("task_scope_required")
            with factory() as session:
                job = Job(
                    organization_id=organization_id,
                    owner_user_id=metadata.get("owner_user_id"),
                    project_id=metadata.get("project_id"),
                    graph_id=metadata.get("graph_id"),
                    operation=self._bounded_text(task_type, 120),
                    status=JobStatus.QUEUED.value,
                    job_metadata=metadata,
                    progress_detail={},
                )
                session.add(job)
                session.commit()
                return job.id

        task_id = str(uuid.uuid4())
        now = datetime.now(timezone.utc)
        task = Task(
            task_id=task_id,
            task_type=task_type,
            status=TaskStatus.PENDING,
            created_at=now,
            updated_at=now,
            metadata=metadata,
        )
        with self._task_lock:
            self._tasks[task_id] = task
        return task_id

    def get_task(
        self,
        task_id: str,
        *,
        organization_id: Optional[str] = None,
        owner_user_id: Optional[str] = None,
    ) -> Optional[Task]:
        factory = self._factory()
        if factory is not None:
            with factory() as session:
                statement = select(Job).where(Job.id == task_id)
                if organization_id is not None:
                    statement = statement.where(Job.organization_id == organization_id)
                if owner_user_id is not None:
                    statement = statement.where(Job.owner_user_id == owner_user_id)
                job = session.scalar(statement)
                return self._from_job(job) if job is not None else None
        with self._task_lock:
            task = self._tasks.get(task_id)
            if task is None:
                return None
            if organization_id is not None and task.metadata.get("organization_id") != organization_id:
                return None
            if owner_user_id is not None and task.metadata.get("owner_user_id") != owner_user_id:
                return None
            return task

    def update_task(
        self,
        task_id: str,
        status: Optional[TaskStatus] = None,
        progress: Optional[int] = None,
        message: Optional[str] = None,
        result: Optional[Dict] = None,
        error: Optional[str] = None,
        progress_detail: Optional[Dict] = None,
    ):
        factory = self._factory()
        if factory is not None:
            with factory() as session:
                job = session.get(Job, task_id)
                if job is None:
                    return
                mapped_status = self._job_status(status)
                if mapped_status is not None:
                    job.status = mapped_status
                    if mapped_status in {
                        JobStatus.SUCCEEDED.value,
                        JobStatus.FAILED.value,
                        JobStatus.CANCELLED.value,
                    }:
                        job.finished_at = datetime.now(timezone.utc)
                if progress is not None:
                    job.progress = min(max(int(progress), 0), 100)
                if message is not None:
                    job.message = self._bounded_text(message)
                if result is not None:
                    job.result = result
                if error is not None:
                    job.error_code = self._bounded_text(error, 120)
                if progress_detail is not None:
                    job.progress_detail = progress_detail
                job.updated_at = datetime.now(timezone.utc)
                session.commit()
            return

        with self._task_lock:
            task = self._tasks.get(task_id)
            if task:
                task.updated_at = datetime.now(timezone.utc)
                if status is not None:
                    task.status = status
                if progress is not None:
                    task.progress = min(max(int(progress), 0), 100)
                if message is not None:
                    task.message = message
                if result is not None:
                    task.result = result
                if error is not None:
                    task.error = error
                if progress_detail is not None:
                    task.progress_detail = progress_detail

    def complete_task(self, task_id: str, result: Dict):
        self.update_task(
            task_id,
            status=TaskStatus.COMPLETED,
            progress=100,
            message=t("progress.taskComplete"),
            result=result,
        )

    def fail_task(self, task_id: str, error: str):
        self.update_task(
            task_id,
            status=TaskStatus.FAILED,
            message=t("progress.taskFailed"),
            error=error,
        )

    def list_tasks(
        self,
        task_type: Optional[str] = None,
        *,
        organization_id: Optional[str] = None,
        owner_user_id: Optional[str] = None,
    ) -> list:
        factory = self._factory()
        if factory is not None:
            with factory() as session:
                statement = select(Job).order_by(Job.created_at.desc())
                if task_type:
                    statement = statement.where(Job.operation == task_type)
                if organization_id is not None:
                    statement = statement.where(Job.organization_id == organization_id)
                if owner_user_id is not None:
                    statement = statement.where(Job.owner_user_id == owner_user_id)
                jobs = session.scalars(statement.limit(100)).all()
                return [self._from_job(job) for job in jobs]
        with self._task_lock:
            tasks = list(self._tasks.values())
            if task_type:
                tasks = [task for task in tasks if task.task_type == task_type]
            if organization_id is not None:
                tasks = [task for task in tasks if task.metadata.get("organization_id") == organization_id]
            if owner_user_id is not None:
                tasks = [task for task in tasks if task.metadata.get("owner_user_id") == owner_user_id]
            return [task for task in sorted(tasks, key=lambda item: item.created_at, reverse=True)]

    def cleanup_old_tasks(self, max_age_hours: int = 24):
        cutoff = datetime.now(timezone.utc) - timedelta(hours=max_age_hours)
        factory = self._factory()
        if factory is not None:
            with factory() as session:
                session.execute(
                    delete(Job).where(
                        Job.created_at < cutoff,
                        Job.status.in_([JobStatus.SUCCEEDED.value, JobStatus.FAILED.value]),
                    )
                )
                session.commit()
            return
        with self._task_lock:
            old_ids = [
                task_id
                for task_id, task in self._tasks.items()
                if task.created_at < cutoff and task.status in [TaskStatus.COMPLETED, TaskStatus.FAILED]
            ]
            for task_id in old_ids:
                del self._tasks[task_id]