JARVIS/backend/app/agents/schemas/task.py

from __future__ import annotations

from datetime import datetime, timezone
from typing import Any, Literal

from pydantic import BaseModel, Field


TaskLifecycleStatus = Literal["pending", "in_progress", "completed", "failed", "blocked"]
VerificationStatus = Literal["passed", "failed", "skipped"]
TaskResultStatus = Literal["completed", "failed", "blocked", "passed", "skipped"]
InterruptStatus = Literal["requested", "acknowledged", "resolved"]
BudgetMode = Literal["direct", "collaboration"]


class InterruptRecord(BaseModel):
    interrupt_id: str
    reason: str
    status: InterruptStatus = "requested"
    requested_by: str | None = None
    source_event_id: str | None = None
    requested_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))
    payload: dict[str, Any] = Field(default_factory=dict)


class RecoveryRecord(BaseModel):
    recovery_id: str
    source_interrupt_id: str | None = None
    strategy: str | None = None
    resumed_from_task_id: str | None = None
    resumed_from_thread_id: str | None = None
    recovered_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))
    payload: dict[str, Any] = Field(default_factory=dict)


class CollaborationBudget(BaseModel):
    mode: BudgetMode = "direct"
    max_parallel_tasks: int | None = None
    remaining_parallel_tasks: int | None = None
    max_tool_calls: int | None = None
    remaining_tool_calls: int | None = None
    max_iterations: int | None = None
    remaining_iterations: int | None = None
    escalation_threshold: int | None = None
    metadata: dict[str, Any] = Field(default_factory=dict)


class AgentTask(BaseModel):
    task_id: str
    title: str
    status: TaskLifecycleStatus = "pending"
    owner_agent_id: str | None = None
    role: str | None = None
    goal: str | None = None
    parent_task_id: str | None = None
    child_task_ids: list[str] = Field(default_factory=list)
    thread_id: str | None = None
    message_id: str | None = None
    message_index: int | None = None
    expected_evidence: list[dict[str, Any]] = Field(default_factory=list)
    evidence: list[dict[str, Any]] = Field(default_factory=list)
    interrupt_records: list[InterruptRecord | dict[str, Any]] = Field(default_factory=list)
    recovery_records: list[RecoveryRecord | dict[str, Any]] = Field(default_factory=list)
    collaboration_budget: CollaborationBudget | dict[str, Any] | None = None
    result_summary: str | None = None
    created_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))
    updated_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))


class TaskResult(BaseModel):
    task_id: str
    status: TaskResultStatus
    summary: str | None = None
    evidence: list[dict[str, Any]] = Field(default_factory=list)
    owner_agent_id: str | None = None
    parent_task_id: str | None = None
    child_task_ids: list[str] = Field(default_factory=list)
    thread_id: str | None = None
    message_id: str | None = None
    message_index: int | None = None
    interrupt_records: list[InterruptRecord | dict[str, Any]] = Field(default_factory=list)
    recovery_records: list[RecoveryRecord | dict[str, Any]] = Field(default_factory=list)
    budget_snapshot: CollaborationBudget | dict[str, Any] | None = None
    next_action: str | None = None
    output_data: dict[str, Any] | None = None