Timmy-time-dashboard/src/timmy/dispatch/routing.py

"""Routing logic — enums, agent registry, and task-to-agent mapping.

Defines the core types (:class:`AgentType`, :class:`TaskType`,
:class:`DispatchStatus`), the :data:`AGENT_REGISTRY`, and the functions
that decide which agent handles a given task.
"""

from __future__ import annotations

import logging
from dataclasses import dataclass
from enum import StrEnum

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Enumerations
# ---------------------------------------------------------------------------


class AgentType(StrEnum):
    """Known agents in the swarm."""

    CLAUDE_CODE = "claude_code"
    KIMI_CODE = "kimi_code"
    AGENT_API = "agent_api"
    TIMMY = "timmy"


class TaskType(StrEnum):
    """Categories of engineering work."""

    # Claude Code strengths
    ARCHITECTURE = "architecture"
    REFACTORING = "refactoring"
    COMPLEX_REASONING = "complex_reasoning"
    CODE_REVIEW = "code_review"

    # Kimi Code strengths
    PARALLEL_IMPLEMENTATION = "parallel_implementation"
    ROUTINE_CODING = "routine_coding"
    FAST_ITERATION = "fast_iteration"

    # Agent API strengths
    RESEARCH = "research"
    ANALYSIS = "analysis"
    SPECIALIZED = "specialized"

    # Timmy strengths
    TRIAGE = "triage"
    PLANNING = "planning"
    CREATIVE = "creative"
    ORCHESTRATION = "orchestration"


class DispatchStatus(StrEnum):
    """Lifecycle state of a dispatched task."""

    PENDING = "pending"
    ASSIGNED = "assigned"
    IN_PROGRESS = "in_progress"
    COMPLETED = "completed"
    FAILED = "failed"
    ESCALATED = "escalated"
    TIMED_OUT = "timed_out"


# ---------------------------------------------------------------------------
# Agent registry
# ---------------------------------------------------------------------------


@dataclass
class AgentSpec:
    """Capabilities and limits for a single agent."""

    name: AgentType
    display_name: str
    strengths: frozenset[TaskType]
    gitea_label: str | None  # label to apply when dispatching
    max_concurrent: int = 1
    interface: str = "gitea"  # "gitea" | "api" | "local"
    api_endpoint: str | None = None  # for interface="api"


#: Authoritative agent registry — all known agents and their capabilities.
AGENT_REGISTRY: dict[AgentType, AgentSpec] = {
    AgentType.CLAUDE_CODE: AgentSpec(
        name=AgentType.CLAUDE_CODE,
        display_name="Claude Code",
        strengths=frozenset(
            {
                TaskType.ARCHITECTURE,
                TaskType.REFACTORING,
                TaskType.COMPLEX_REASONING,
                TaskType.CODE_REVIEW,
            }
        ),
        gitea_label="claude-ready",
        max_concurrent=1,
        interface="gitea",
    ),
    AgentType.KIMI_CODE: AgentSpec(
        name=AgentType.KIMI_CODE,
        display_name="Kimi Code",
        strengths=frozenset(
            {
                TaskType.PARALLEL_IMPLEMENTATION,
                TaskType.ROUTINE_CODING,
                TaskType.FAST_ITERATION,
            }
        ),
        gitea_label="kimi-ready",
        max_concurrent=1,
        interface="gitea",
    ),
    AgentType.AGENT_API: AgentSpec(
        name=AgentType.AGENT_API,
        display_name="Agent API",
        strengths=frozenset(
            {
                TaskType.RESEARCH,
                TaskType.ANALYSIS,
                TaskType.SPECIALIZED,
            }
        ),
        gitea_label=None,
        max_concurrent=5,
        interface="api",
    ),
    AgentType.TIMMY: AgentSpec(
        name=AgentType.TIMMY,
        display_name="Timmy",
        strengths=frozenset(
            {
                TaskType.TRIAGE,
                TaskType.PLANNING,
                TaskType.CREATIVE,
                TaskType.ORCHESTRATION,
            }
        ),
        gitea_label=None,
        max_concurrent=1,
        interface="local",
    ),
}

#: Map from task type to preferred agent (primary routing table).
_TASK_ROUTING: dict[TaskType, AgentType] = {
    TaskType.ARCHITECTURE: AgentType.CLAUDE_CODE,
    TaskType.REFACTORING: AgentType.CLAUDE_CODE,
    TaskType.COMPLEX_REASONING: AgentType.CLAUDE_CODE,
    TaskType.CODE_REVIEW: AgentType.CLAUDE_CODE,
    TaskType.PARALLEL_IMPLEMENTATION: AgentType.KIMI_CODE,
    TaskType.ROUTINE_CODING: AgentType.KIMI_CODE,
    TaskType.FAST_ITERATION: AgentType.KIMI_CODE,
    TaskType.RESEARCH: AgentType.AGENT_API,
    TaskType.ANALYSIS: AgentType.AGENT_API,
    TaskType.SPECIALIZED: AgentType.AGENT_API,
    TaskType.TRIAGE: AgentType.TIMMY,
    TaskType.PLANNING: AgentType.TIMMY,
    TaskType.CREATIVE: AgentType.TIMMY,
    TaskType.ORCHESTRATION: AgentType.TIMMY,
}


# ---------------------------------------------------------------------------
# Routing logic
# ---------------------------------------------------------------------------


def select_agent(task_type: TaskType) -> AgentType:
    """Return the best agent for *task_type* based on the routing table.

    Args:
        task_type: The category of engineering work to be done.

    Returns:
        The :class:`AgentType` best suited to handle this task.
    """
    return _TASK_ROUTING.get(task_type, AgentType.TIMMY)


def infer_task_type(title: str, description: str = "") -> TaskType:
    """Heuristic: guess the most appropriate :class:`TaskType` from text.

    Scans *title* and *description* for keyword signals and returns the
    strongest match.  Falls back to :attr:`TaskType.ROUTINE_CODING`.

    Args:
        title: Short task title.
        description: Longer task description (optional).

    Returns:
        The inferred :class:`TaskType`.
    """
    text = (title + " " + description).lower()

    _SIGNALS: list[tuple[TaskType, frozenset[str]]] = [
        (
            TaskType.ARCHITECTURE,
            frozenset({"architect", "design", "adr", "system design", "schema"}),
        ),
        (
            TaskType.REFACTORING,
            frozenset({"refactor", "clean up", "cleanup", "reorganise", "reorganize"}),
        ),
        (TaskType.CODE_REVIEW, frozenset({"review", "pr review", "pull request review", "audit"})),
        (
            TaskType.COMPLEX_REASONING,
            frozenset({"complex", "hard problem", "debug", "investigate", "diagnose"}),
        ),
        (
            TaskType.RESEARCH,
            frozenset({"research", "survey", "literature", "benchmark", "analyse", "analyze"}),
        ),
        (TaskType.ANALYSIS, frozenset({"analysis", "profil", "trace", "metric", "performance"})),
        (TaskType.TRIAGE, frozenset({"triage", "classify", "prioritise", "prioritize"})),
        (TaskType.PLANNING, frozenset({"plan", "roadmap", "milestone", "epic", "spike"})),
        (TaskType.CREATIVE, frozenset({"creative", "persona", "story", "write", "draft"})),
        (TaskType.ORCHESTRATION, frozenset({"orchestrat", "coordinat", "swarm", "dispatch"})),
        (TaskType.PARALLEL_IMPLEMENTATION, frozenset({"parallel", "concurrent", "batch"})),
        (TaskType.FAST_ITERATION, frozenset({"quick", "fast", "iterate", "prototype", "poc"})),
    ]

    for task_type, keywords in _SIGNALS:
        if any(kw in text for kw in keywords):
            return task_type

    return TaskType.ROUTINE_CODING