beaver_project/app-instance/backend/beaver/tasks/service.py

"""Internal service for automatic Task mode."""

from __future__ import annotations

from datetime import datetime, timezone
from pathlib import Path
from typing import Any
from uuid import uuid4

from .models import TaskEvent, TaskRecord, ValidationResult
from .store import TaskStore


class TaskService:
    def __init__(self, root: str | Path) -> None:
        self.store = TaskStore(root)

    def create_task(
        self,
        *,
        session_id: str,
        description: str,
        creator: str = "main-agent",
        metadata: dict[str, Any] | None = None,
    ) -> TaskRecord:
        now = self._now()
        task = TaskRecord(
            task_id=uuid4().hex,
            session_id=session_id,
            description=description,
            goal=description,
            constraints=[],
            priority=0,
            status="open",
            creator=creator,
            created_at=now,
            updated_at=now,
            metadata=dict(metadata or {}),
        )
        self.store.upsert_task(task)
        self._event(task, "created", payload={"description": description})
        return task

    def get_task(self, task_id: str) -> TaskRecord | None:
        return self.store.get_task(task_id)

    def get_task_by_run_id(self, run_id: str) -> TaskRecord | None:
        return self.store.get_task_by_run_id(run_id)

    def get_latest_open_task(self, session_id: str) -> TaskRecord | None:
        return self.store.get_latest_open_task(session_id)

    def start_run(self, task_id: str, *, user_message: str, attempt_index: int) -> TaskRecord:
        task = self._require(task_id)
        task.status = "running"
        task.updated_at = self._now()
        task.metadata["latest_user_message"] = user_message
        task.metadata["latest_attempt_index"] = attempt_index
        self.store.upsert_task(task)
        self._event(task, "run_started", payload={"user_message": user_message, "attempt_index": attempt_index})
        return task

    def append_run(self, task_id: str, run_id: str, *, skill_names: list[str] | None = None) -> TaskRecord:
        task = self._require(task_id)
        if run_id not in task.run_ids:
            task.run_ids.append(run_id)
        for name in skill_names or []:
            if name not in task.skill_names:
                task.skill_names.append(name)
        task.updated_at = self._now()
        self.store.upsert_task(task)
        self._event(task, "run_completed", run_id=run_id, payload={"skill_names": skill_names or []})
        return task

    def record_validation(self, task_id: str, run_id: str, validation: ValidationResult) -> TaskRecord:
        task = self._require(task_id)
        task.status = "awaiting_feedback"
        task.updated_at = self._now()
        task.validation_result = validation.to_dict()
        self.store.upsert_task(task)
        self._event(task, "validated", run_id=run_id, payload=validation.to_dict())
        return task

    def add_feedback(
        self,
        task_id: str,
        *,
        feedback_type: str,
        comment: str | None = None,
        run_id: str | None = None,
    ) -> TaskRecord:
        task = self._require(task_id)
        now = self._now()
        matching_feedback = any(
            item.get("run_id") == run_id and item.get("feedback_type") == feedback_type
            for item in task.feedback
        )
        conflicting_feedback = next(
            (
                item
                for item in task.feedback
                if item.get("run_id") == run_id and item.get("feedback_type") != feedback_type
            ),
            None,
        )
        if conflicting_feedback is not None:
            raise ValueError(
                f"Feedback for run_id={run_id!r} was already recorded as "
                f"{conflicting_feedback.get('feedback_type')!r}"
            )
        if task.status in {"closed", "abandoned"} and not matching_feedback:
            raise ValueError(f"Task {task.task_id} is already finalized as {task.status!r}")
        if matching_feedback:
            return task

        entry = {
            "feedback_type": feedback_type,
            "comment": comment or "",
            "run_id": run_id,
            "created_at": now,
        }
        task.feedback.append(entry)
        if feedback_type == "revise":
            task.status = "needs_revision"
        elif feedback_type == "abandon":
            task.status = "abandoned"
            task.closed_at = now
            task.close_reason = comment or "abandoned"
        elif feedback_type == "satisfied":
            task.status = "closed"
            task.closed_at = now
            task.close_reason = "satisfied"
            task.satisfaction = 1.0
        task.updated_at = now
        self.store.upsert_task(task)
        self._event(task, f"feedback_{feedback_type}", run_id=run_id, payload=entry)
        return task

    def _require(self, task_id: str) -> TaskRecord:
        task = self.store.get_task(task_id)
        if task is None:
            raise ValueError(f"Unknown task_id: {task_id}")
        return task

    def _event(
        self,
        task: TaskRecord,
        event_type: str,
        *,
        run_id: str | None = None,
        payload: dict[str, Any] | None = None,
    ) -> None:
        self.store.append_event(
            TaskEvent(
                event_id=uuid4().hex,
                task_id=task.task_id,
                session_id=task.session_id,
                run_id=run_id,
                event_type=event_type,
                created_at=self._now(),
                payload=dict(payload or {}),
            )
        )

    @staticmethod
    def _now() -> str:
        return datetime.now(timezone.utc).isoformat()