JARVIS/backend/app/services/task_dispatch.py

import asyncio
from datetime import UTC, datetime
from uuid import uuid4

from sqlalchemy import select
from sqlalchemy.ext.asyncio import async_sessionmaker
from sqlalchemy.orm import object_session
from sqlalchemy.orm import selectinload

from app.models.task import (
    Task,
    TaskDispatchStatus,
    TaskHistory,
    TaskPriority,
    TaskStatus,
    TaskSubTask,
)


def _now() -> datetime:
    return datetime.now(UTC)


def _stringify(value: object | None) -> str | None:
    if value is None:
        return None
    return str(value)


def append_task_history(
    task: Task,
    *,
    action: str,
    old_value: object | None = None,
    new_value: object | None = None,
) -> None:
    entry = TaskHistory(
        task_id=task.id,
        action=action,
        old_value=_stringify(old_value),
        new_value=_stringify(new_value),
    )
    session = object_session(task)
    if session is not None:
        session.add(entry)
        return
    task.history.append(entry)


def build_dispatch_payload(task: Task, subtasks: list[TaskSubTask]) -> dict[str, object]:
    return {
        "business_task_id": task.id,
        "title": task.title,
        "description": task.description,
        "priority": task.priority.value if isinstance(task.priority, TaskPriority) else str(task.priority),
        "due_date": task.due_date.isoformat() if task.due_date else None,
        "conversation_id": task.conversation_id,
        "user_id": task.user_id,
        "subtasks": [
            {
                "id": item.id,
                "title": item.title,
                "description": item.description,
                "status": item.status.value if isinstance(item.status, TaskStatus) else str(item.status),
                "assignee_type": item.assignee_type.value if item.assignee_type else None,
                "assignee_id": item.assignee_id,
                "dispatch_status": (
                    item.dispatch_status.value
                    if isinstance(item.dispatch_status, TaskDispatchStatus)
                    else str(item.dispatch_status)
                ),
                "order_index": item.order_index,
            }
            for item in subtasks
        ],
    }


async def _run_dispatch_flow(
    task_id: str,
    run_id: str,
    *,
    session_factory,
    subtask_id: str | None = None,
) -> None:
    await asyncio.sleep(0.01)

    async with session_factory() as db:
        task = await db.get(Task, task_id)
        if task is None:
            return
        target = await db.get(TaskSubTask, subtask_id) if subtask_id else None
        if subtask_id and target is None:
            return

        if subtask_id:
            previous = target.dispatch_status
            target.dispatch_status = TaskDispatchStatus.RUNNING
            target.dispatch_run_id = run_id
            target.completed_at = None
            task.dispatch_status = TaskDispatchStatus.RUNNING
            task.dispatch_run_id = run_id
            task.started_at = task.started_at or _now()
            task.last_synced_at = _now()
            append_task_history(
                task,
                action="dispatch_status_changed",
                old_value=f"{subtask_id}:{previous.value}",
                new_value=f"{subtask_id}:{TaskDispatchStatus.RUNNING.value}",
            )
        else:
            previous = task.dispatch_status
            task.dispatch_status = TaskDispatchStatus.RUNNING
            task.dispatch_run_id = run_id
            task.started_at = task.started_at or _now()
            task.last_synced_at = _now()
            task.status = TaskStatus.IN_PROGRESS
            append_task_history(
                task,
                action="dispatch_status_changed",
                old_value=previous.value,
                new_value=TaskDispatchStatus.RUNNING.value,
            )
        await db.commit()

    await asyncio.sleep(0.01)

    async with session_factory() as db:
        task = await db.get(Task, task_id)
        if task is None:
            return
        target = await db.get(TaskSubTask, subtask_id) if subtask_id else None
        if subtask_id and target is None:
            return

        synced_at = _now()
        if subtask_id:
            previous = target.dispatch_status
            target.dispatch_status = TaskDispatchStatus.COMPLETED
            target.dispatch_run_id = run_id
            target.status = TaskStatus.DONE
            target.completed_at = synced_at
            task.dispatch_status = TaskDispatchStatus.COMPLETED
            task.dispatch_run_id = run_id
            task.result_summary = f"Commander completed subtask {target.title}"
            task.last_synced_at = synced_at
            append_task_history(
                task,
                action="dispatch_status_changed",
                old_value=f"{subtask_id}:{previous.value}",
                new_value=f"{subtask_id}:{TaskDispatchStatus.COMPLETED.value}",
            )
        else:
            previous = task.dispatch_status
            task.dispatch_status = TaskDispatchStatus.COMPLETED
            task.dispatch_run_id = run_id
            task.result_summary = f"Commander completed task {task.title}"
            task.last_synced_at = synced_at
            task.status = TaskStatus.DONE
            task.completed_at = synced_at
            append_task_history(
                task,
                action="dispatch_status_changed",
                old_value=previous.value,
                new_value=TaskDispatchStatus.COMPLETED.value,
            )
        await db.commit()


def schedule_dispatch(task_id: str, run_id: str, *, session_factory, subtask_id: str | None = None) -> None:
    asyncio.create_task(
        _run_dispatch_flow(
            task_id,
            run_id,
            session_factory=session_factory,
            subtask_id=subtask_id,
        )
    )


async def queue_task_dispatch(
    task: Task,
    *,
    db,
    subtask: TaskSubTask | None = None,
) -> tuple[str, dict[str, object]]:
    subtasks = list(task.subtasks)
    run_id = uuid4().hex[:12]
    synced_at = _now()

    if subtask is not None:
        previous = subtask.dispatch_status
        subtask.dispatch_status = TaskDispatchStatus.QUEUED
        subtask.dispatch_run_id = run_id
        task.dispatch_status = TaskDispatchStatus.QUEUED
        task.dispatch_run_id = run_id
        task.result_summary = None
        task.last_synced_at = synced_at
        append_task_history(
            task,
            action="dispatched_to_commander",
            old_value=f"{subtask.id}:{previous.value}",
            new_value=f"{subtask.id}:{TaskDispatchStatus.QUEUED.value}",
        )
    else:
        previous = task.dispatch_status
        task.dispatch_status = TaskDispatchStatus.QUEUED
        task.dispatch_run_id = run_id
        task.result_summary = None
        task.started_at = None
        task.last_synced_at = synced_at
        append_task_history(
            task,
            action="dispatched_to_commander",
            old_value=previous.value,
            new_value=TaskDispatchStatus.QUEUED.value,
        )

    await db.commit()
    await db.refresh(task)
    payload = build_dispatch_payload(task, subtasks)
    session_factory = async_sessionmaker(bind=db.bind, expire_on_commit=False)
    schedule_dispatch(
        task.id,
        run_id,
        session_factory=session_factory,
        subtask_id=subtask.id if subtask else None,
    )
    return run_id, payload


async def load_task_with_details(db, *, task_id: str, user_id: str) -> Task | None:
    result = await db.execute(
        select(Task)
        .options(selectinload(Task.subtasks), selectinload(Task.history))
        .where(Task.id == task_id, Task.user_id == user_id)
    )
    return result.scalar_one_or_none()