core/task_queue.py

"""
core/task_queue.py — lightweight persistent task queue.

This is intentionally small and dependency-free: it gives the app an async
execution path without requiring Celery/Redis just to package a deployable
ZIP. Queue state is persisted to JSON so the UI/API can poll task progress.
"""

from __future__ import annotations

import copy
import json
import os
import queue
import tempfile
import threading
import time
import uuid
from typing import Any, Callable, Dict, Optional


class InMemoryTaskQueue:
    def __init__(self, workspace: str, max_workers: int = 1) -> None:
        self.workspace = os.path.abspath(workspace)
        self.max_workers = max(1, int(max_workers))
        self.tasks_dir = os.path.join(self.workspace, "queue", "tasks")
        os.makedirs(self.tasks_dir, exist_ok=True)
        self._queue: "queue.Queue[tuple[str, Callable[..., Any], Dict[str, Any]]]" = queue.Queue()
        self._lock = threading.RLock()
        self._started = False
        self._workers: list[threading.Thread] = []

    def _task_path(self, task_id: str) -> str:
        return os.path.join(self.tasks_dir, f"{task_id}.json")

    def _atomic_write(self, path: str, payload: Dict[str, Any]) -> None:
        fd, tmp = tempfile.mkstemp(prefix=".task_", suffix=".tmp", dir=os.path.dirname(path))
        try:
            with os.fdopen(fd, "w", encoding="utf-8") as fh:
                json.dump(payload, fh, indent=2, ensure_ascii=False, default=str)
                fh.flush()
                try:
                    os.fsync(fh.fileno())
                except OSError:
                    pass
            os.replace(tmp, path)
        finally:
            try:
                if os.path.exists(tmp):
                    os.unlink(tmp)
            except OSError:
                pass

    def _load_task(self, task_id: str) -> Optional[Dict[str, Any]]:
        path = self._task_path(task_id)
        if not os.path.exists(path):
            return None
        try:
            with open(path, "r", encoding="utf-8") as fh:
                payload = json.load(fh)
            return payload if isinstance(payload, dict) else None
        except Exception:
            return None

    def _save_task(self, task_id: str, payload: Dict[str, Any]) -> None:
        self._atomic_write(self._task_path(task_id), payload)

    def _ensure_started(self) -> None:
        if self._started:
            return
        with self._lock:
            if self._started:
                return
            for idx in range(self.max_workers):
                t = threading.Thread(target=self._worker, name=f"devai-queue-{idx}", daemon=True)
                t.start()
                self._workers.append(t)
            self._started = True

    def enqueue(
        self,
        func: Callable[..., Any],
        *,
        kwargs: Optional[Dict[str, Any]] = None,
        tenant_id: str = "default",
        label: str = "pipeline",
    ) -> str:
        self._ensure_started()
        task_id = f"task-{uuid.uuid4().hex[:16]}"
        now = time.time()
        payload = {
            "task_id": task_id,
            "tenant_id": tenant_id,
            "label": label,
            "status": "queued",
            "created_at": now,
            "started_at": None,
            "finished_at": None,
            "progress_events": [],
            "result": None,
            "error": None,
        }
        self._save_task(task_id, payload)
        self._queue.put((task_id, func, dict(kwargs or {})))
        return task_id

    def _worker(self) -> None:
        while True:
            task_id, func, kwargs = self._queue.get()
            try:
                payload = self._load_task(task_id) or {"task_id": task_id}
                payload.update({"status": "running", "started_at": time.time(), "error": None})
                self._save_task(task_id, payload)

                def _progress(event: Dict[str, Any]) -> None:
                    current = self._load_task(task_id) or payload
                    events = list(current.get("progress_events") or [])
                    events.append(copy.deepcopy(event))
                    current["progress_events"] = events[-500:]
                    self._save_task(task_id, current)

                call_kwargs = dict(kwargs)
                call_kwargs.setdefault("_progress_callback", _progress)
                result = func(**call_kwargs)

                current = self._load_task(task_id) or payload
                current.update({
                    "status": "completed",
                    "finished_at": time.time(),
                    "result": result,
                })
                self._save_task(task_id, current)
            except Exception as exc:
                current = self._load_task(task_id) or {"task_id": task_id}
                current.update({
                    "status": "failed",
                    "finished_at": time.time(),
                    "error": str(exc),
                })
                self._save_task(task_id, current)
            finally:
                self._queue.task_done()

    def get_status(self, task_id: str) -> Optional[Dict[str, Any]]:
        payload = self._load_task(task_id)
        if payload is None:
            return None
        return {
            "task_id": payload.get("task_id"),
            "tenant_id": payload.get("tenant_id"),
            "label": payload.get("label"),
            "status": payload.get("status"),
            "created_at": payload.get("created_at"),
            "started_at": payload.get("started_at"),
            "finished_at": payload.get("finished_at"),
            "progress_events": list(payload.get("progress_events") or []),
            "error": payload.get("error"),
        }

    def get_result(self, task_id: str) -> Optional[Dict[str, Any]]:
        payload = self._load_task(task_id)
        if payload is None:
            return None
        return copy.deepcopy(payload)


__all__ = ["InMemoryTaskQueue"]