JARVIS/backend/app/agents/state.py

from dataclasses import dataclass
from enum import Enum
from typing import Annotated, Any, Literal, TypedDict

from app.agents.schemas.event import AgentEvent
from app.agents.schemas.message import AgentMessage
from app.agents.schemas.task import (
    AgentTask,
    CollaborationBudget,
    InterruptRecord,
    RecoveryRecord,
    TaskResult,
    VerificationStatus,
)
from langchain_core.messages import AIMessage, BaseMessage, HumanMessage
from langgraph.graph.message import add_messages

AgentPhase = Literal[
    "phase_0_bootstrap",
    "phase_1_routing",
    "phase_2_controlled_collaboration",
    "phase_3_dynamic_collaboration",
    "phase_4_visibility_and_verification",
]


class AgentRole(str, Enum):
    MASTER = "master"
    SCHEDULE_PLANNER = "schedule_planner"
    EXECUTOR = "executor"
    LIBRARIAN = "librarian"
    ANALYST = "analyst"
    CODE_COMMANDER = "code_commander"


@dataclass
class ConversationTurn:
    role: str  # "user" | "assistant"
    content: str
    agent: AgentRole | None = None
    model: str | None = None


def turn_to_message(turn: ConversationTurn) -> BaseMessage:
    if turn.role == "user":
        return HumanMessage(content=turn.content)
    return AIMessage(content=turn.content)


class AgentState(TypedDict):
    messages: Annotated[list[BaseMessage], add_messages]
    user_id: str
    conversation_id: str
    parent_conversation_id: str | None
    thread_id: str | None
    last_message_id: str | None
    message_sequence: int
    agent_id: str | None
    parent_agent_id: str | None
    root_agent_id: str | None
    collaboration_depth: int
    spawned_agent_ids: list[str]

    execution_mode: Literal["direct", "collaboration", "delegated", "verified"]
    current_agent: str | None
    next_step: str | None
    active_agents: list[AgentRole]
    current_sub_commander: str | None
    active_sub_commanders: list[str]
    sub_commander_trace: list[dict[str, Any]]
    agent_trace: list[str]
    event_trace: list[AgentEvent | dict[str, Any]]
    message_trace: list[AgentMessage | dict[str, Any]]

    pending_tasks: list[dict[str, Any]]
    completed_tasks: list[dict[str, Any]]
    active_tasks: list[AgentTask | dict[str, Any]]
    task_results: list[TaskResult | dict[str, Any]]
    task_hierarchy: dict[str, list[str]]
    interrupted_tasks: list[InterruptRecord | dict[str, Any]]
    recovery_trace: list[RecoveryRecord | dict[str, Any]]
    recovery_points: list[dict[str, Any]]
    tool_calls: list[dict[str, Any]]
    last_tool_result: str | None
    action_results: list[dict[str, Any]]
    created_entities: list[dict[str, Any]]
    tool_outcomes: list[dict[str, Any]]
    task_result_summary: dict[str, Any] | None
    verifier_hints: dict[str, Any] | None

    verification_status: VerificationStatus | None
    verification_summary: str | None
    verification_evidence: list[dict[str, Any]]
    isolation_mode: str
    isolation_id: str | None
    isolation_workspace_path: str | None
    isolation_parent_conversation_id: str | None
    isolation_metadata: dict[str, Any]
    input_tokens: int
    output_tokens: int
    estimated_cost: float | None
    budget_warning: bool
    cost_by_agent: dict[str, dict[str, Any]]
    cost_thresholds: dict[str, Any]
    budget_state: CollaborationBudget | dict[str, Any] | None
    collaboration_budget_history: list[CollaborationBudget | dict[str, Any]]
    current_phase: AgentPhase
    phase_history: list[dict[str, Any]]
    current_checkpoint: str | None
    checkpoint_history: list[dict[str, Any]]

    tool_strategy_used: str | None
    tool_round_count: int
    max_tool_rounds: int
    retry_count: int
    max_retries: int
    iteration_count: int
    max_iterations: int
    routing_hops: int
    max_routing_hops: int
    terminated_due_to_loop_guard: bool
    retrieval_trace: list[dict[str, Any]]
    stop_reason: str | None

    clarification_needed: bool
    clarification_question: str | None
    fallback_parse_error: str | None
    should_respond: bool

    knowledge_context: str | None
    graph_context: str | None
    schedule_context_summary: str | None
    plan: str | None
    plan_steps: list[dict[str, Any]]
    analysis_report: str | None
    final_response: str | None

    memory_context: str | None
    current_datetime_context: str | None
    current_datetime_reference: dict[str, str] | None
    runtime_request_context: dict[str, Any] | None
    task_graph: dict[str, Any] | None
    scheduled_subtasks: list[dict[str, Any]]
    recalled_retrospectives: list[dict[str, Any]]
    retrospective_shortlist: list[dict[str, Any]]
    skill_shortlist: list[dict[str, Any]]
    skill_activation_records: list[dict[str, Any]]
    execution_decision: dict[str, Any] | None
    merge_report: dict[str, Any] | None
    verification_report: dict[str, Any] | None
    feature_flags: dict[str, bool]
    observability_report: dict[str, Any] | None

    turn_context: dict[str, Any] | None
    routing_decision: dict[str, Any] | None
    continuity_state: dict[str, Any] | None
    pending_action: dict[str, Any] | None
    last_completed_action: dict[str, Any] | None
    clarification_context: dict[str, Any] | None

    user_llm_config: dict[str, Any] | None
    provider_capabilities: dict[str, Any] | None

    # Code Commander state
    code_task_type: Literal["demo", "project", "modification"] | None
    code_ai_provider: Literal["claude", "gemini", "codex", "opencode"] | None
    code_sandbox_mode: bool | None
    code_workspace_path: str | None
    code_execution_session_id: str | None
    code_execution_result: dict[str, Any] | None


def initial_state(user_id: str, conversation_id: str) -> AgentState:
    return AgentState(
        messages=[],
        user_id=user_id,
        conversation_id=conversation_id,
        parent_conversation_id=None,
        thread_id=None,
        last_message_id=None,
        message_sequence=0,
        agent_id=AgentRole.MASTER.value,
        parent_agent_id=None,
        root_agent_id=AgentRole.MASTER.value,
        collaboration_depth=0,
        spawned_agent_ids=[],
        execution_mode="direct",
        current_agent=AgentRole.MASTER.value,
        next_step=None,
        active_agents=[AgentRole.MASTER],
        current_sub_commander=None,
        active_sub_commanders=[],
        sub_commander_trace=[],
        agent_trace=[AgentRole.MASTER.value],
        event_trace=[],
        message_trace=[],
        pending_tasks=[],
        completed_tasks=[],
        active_tasks=[],
        task_results=[],
        task_hierarchy={},
        interrupted_tasks=[],
        recovery_trace=[],
        recovery_points=[],
        tool_calls=[],
        last_tool_result=None,
        action_results=[],
        created_entities=[],
        tool_outcomes=[],
        task_result_summary=None,
        verifier_hints=None,
        verification_status=None,
        verification_summary=None,
        verification_evidence=[],
        isolation_mode="none",
        isolation_id=None,
        isolation_workspace_path=None,
        isolation_parent_conversation_id=None,
        isolation_metadata={},
        input_tokens=0,
        output_tokens=0,
        estimated_cost=None,
        budget_warning=False,
        cost_by_agent={},
        cost_thresholds={},
        budget_state=None,
        collaboration_budget_history=[],
        current_phase="phase_0_bootstrap",
        phase_history=[
            {
                "phase": "phase_0_bootstrap",
                "reason": "initial_state_created",
            }
        ],
        current_checkpoint="bootstrap.initialized",
        checkpoint_history=[
            {
                "checkpoint": "bootstrap.initialized",
                "phase": "phase_0_bootstrap",
                "reason": "initial_state_created",
            }
        ],
        tool_strategy_used=None,
        tool_round_count=0,
        max_tool_rounds=2,
        retry_count=0,
        max_retries=1,
        iteration_count=0,
        max_iterations=3,
        routing_hops=0,
        max_routing_hops=2,
        terminated_due_to_loop_guard=False,
        retrieval_trace=[],
        stop_reason=None,
        clarification_needed=False,
        clarification_question=None,
        fallback_parse_error=None,
        should_respond=True,
        knowledge_context=None,
        graph_context=None,
        schedule_context_summary=None,
        plan=None,
        plan_steps=[],
        analysis_report=None,
        final_response=None,
        memory_context=None,
        current_datetime_context=None,
        current_datetime_reference=None,
        runtime_request_context=None,
        task_graph=None,
        scheduled_subtasks=[],
        recalled_retrospectives=[],
        retrospective_shortlist=[],
        skill_shortlist=[],
        skill_activation_records=[],
        execution_decision=None,
        merge_report=None,
        verification_report=None,
        feature_flags={},
        observability_report=None,
        turn_context=None,
        routing_decision=None,
        continuity_state=None,
        pending_action=None,
        last_completed_action=None,
        clarification_context=None,
        user_llm_config=None,
        provider_capabilities=None,
    )