smart-org/apps/api/app/interfaces/http/routers/engine.py

"""The AI workforce engine's API: status, task board, task detail and recent turns.

Every route resolves its tenant through ``get_org_context``. The reads are open to
every member of the organization, so the credit notice and the board work for
every role. The four actions on a task (pause, resume, cancel, answer) need the
``tasks.run`` capability and each writes one audit row. The team (view, hire by role
and tier, fire) is in this router too; hire and fire need ``employees.manage``.

Three things are load-bearing here:

* **Model identity never leaves.** No schema in this module has a field that can
  hold a model id or name, a provider, an account, a container or a port. A turn
  shows the employee's *name*, the role, the mode and the stage, never the model
  that ran it. A future refactor cannot leak one without adding such a field.
* **Every query is tenant-scoped.** Tasks, states, turns, gates and runs are read
  through repositories that filter by organization, and a task id that belongs to
  another organization is a 404 (the same answer a missing id gets).
* **Bounded reads.** The board scans at most the newest 500 tasks of the
  organization when it filters by stage; a turn excerpt is capped at 1,000 characters.
"""

from __future__ import annotations

from datetime import datetime
from typing import Literal

from fastapi import APIRouter, Depends, Query, Request
from pydantic import BaseModel, Field
from sqlalchemy.ext.asyncio import AsyncSession

from app.application.container_org import build_model_grants
from app.application.ports.model_grants import GrantedModel
from app.application.services.audit_service import record_audit
from app.core.database import get_session
from app.core.errors import NotFoundError
from app.domain.audit.models import AuditAction
from app.domain.identity.models import CAN_MANAGE_EMPLOYEES, CAN_RUN_TASKS
from app.domain.workforce.engine import EngineTurn
from app.domain.workforce.models import Employee, EmployeeStatus, RoleKey, Run, Task
from app.domain.workforce.pipeline import Stage, TaskPipelineState, pipeline_stages
from app.infrastructure.repositories.billing import CreditWalletRepository
from app.infrastructure.repositories.engine import (
    EmployeeProfileRepository,
    EngineTurnRepository,
    GateResultRepository,
    PipelineStateRepository,
)
from app.infrastructure.repositories.workforce import (
    EmployeeRepository,
    RunRepository,
    TaskRepository,
)
from app.interfaces.http.deps import OrgContext, client_meta, get_org_context

router = APIRouter(prefix="/organizations/{org_id}/engine", tags=["engine"])

#: The most tasks a board reads, newest first (a stage filter looks only among these).
_BOARD_SCAN_LIMIT = 500
#: The most pipeline states read in one listing; the rest are looked up one by one.
_STATE_SCAN_LIMIT = 1000
#: The longest reply excerpt a turn carries. The full reply is not shown here.
_EXCERPT_CHARS = 1000
#: The most turns one task detail carries, and the most gate verdicts.
_DETAIL_LIMIT = 200


# ---------------------------------------------------------------------------
# Schemas
# ---------------------------------------------------------------------------


class EngineStatusOut(BaseModel):
    """Whether the organization's engine can work, and how much of the team is in use."""

    balance: float
    currency: str
    credit_ok: bool
    automation_enabled: bool
    max_employees: int
    team_size: int
    stage_counts: dict[str, int]


class EngineTaskOut(BaseModel):
    """One task on the board.

    ``stage``, ``pipeline`` and ``kind`` are null for a task the engine does not
    own (a legacy task). ``cycle`` and ``max_cycles`` are 0 for those tasks too.
    """

    id: str
    title: str
    status: str
    stage: str | None
    pipeline: str | None
    kind: str | None
    cycle: int
    max_cycles: int
    thread_id: str | None
    assigned_employee_id: str | None
    assignee_name: str | None
    created_at: datetime
    updated_at: datetime


class EngineTurnOut(BaseModel):
    """One employee turn: who ran, in which role and stage, and how it ended."""

    id: str
    run_id: str
    role: str
    mode: str
    stage: str
    employee_name: str | None
    ok: bool
    error_kind: str
    tool_calls: int
    reply_excerpt: str
    started_at: datetime
    finished_at: datetime | None


class EngineOrgTurnOut(EngineTurnOut):
    """A turn in the organization's recent list, with the task it belongs to."""

    task_id: str
    task_title: str


class EngineGateOut(BaseModel):
    id: str
    gate: str
    ok: bool
    detail: str
    created_at: datetime


class EngineRunOut(BaseModel):
    id: str
    status: str
    started_at: datetime | None
    finished_at: datetime | None


class EngineTaskDetailOut(EngineTaskOut):
    """One task with its pipeline, its turns, its gate verdicts and its last diff."""

    description: str
    #: The stages of the task's pipeline in order; empty for a legacy task.
    stages: list[str]
    waiver_reason: str
    branch: str
    depends_on: list[str]
    #: The sha the merge landed on, once the task is completed.
    merged_sha: str
    #: The last accepted diff (capped by the writer at 20,000 characters).
    diff: str
    turns: list[EngineTurnOut]
    gates: list[EngineGateOut]
    runs: list[EngineRunOut]


class EngineAnswerRequest(BaseModel):
    """A person's answer to the question a task is waiting on."""

    content: str = Field(min_length=1, max_length=4000)


class TierOut(BaseModel):
    tier: str
    label: str


class TeamMemberOut(BaseModel):
    """One team member. ``tier`` is neutral: no model, provider or model id is shown."""

    id: str
    name: str
    role_key: str
    status: str
    #: ``fast``, ``balanced``, ``strong``, or ``unknown`` when the model is not granted.
    tier: str
    busy_task_id: str | None
    busy_task_title: str | None
    hired_by: str


class TeamOut(BaseModel):
    employees: list[TeamMemberOut]
    #: The tiers the organization has granted models for, cheapest first.
    tiers: list[TierOut]
    #: The owner's cap on the headcount for this organization.
    max_employees: int
    count: int


class HireRequest(BaseModel):
    """A hire asks for a role and a tier. The model is chosen by the platform."""

    role: RoleKey
    tier: str = Field(min_length=1, max_length=20)
    name: str = Field(default="", max_length=120)


# ---------------------------------------------------------------------------
# Mapping
# ---------------------------------------------------------------------------

#: Neutral tier labels, cheapest first: the order the team view lists them in.
_TIER_LABELS: dict[str, str] = {"fast": "سریع", "balanced": "متعادل", "strong": "قوی"}


def _task_out(
    task: Task, state: TaskPipelineState | None, names: dict[str, str]
) -> EngineTaskOut:
    return EngineTaskOut(
        id=task.id,
        title=task.title,
        status=str(task.status),
        stage=str(state.stage) if state else None,
        pipeline=str(state.pipeline) if state else None,
        kind=str(state.kind) if state else None,
        cycle=state.cycle if state else 0,
        max_cycles=state.max_cycles if state else 0,
        thread_id=task.thread_id,
        assigned_employee_id=task.assigned_employee_id,
        assignee_name=names.get(task.assigned_employee_id or ""),
        created_at=task.created_at,
        updated_at=task.updated_at,
    )


def _turn_out(turn: EngineTurn, names: dict[str, str]) -> EngineTurnOut:
    return EngineTurnOut(
        id=turn.id,
        run_id=turn.run_id,
        role=str(turn.role),
        mode=str(turn.mode),
        stage=str(turn.stage),
        employee_name=names.get(turn.employee_id or ""),
        ok=turn.ok,
        error_kind=turn.error_kind,
        tool_calls=turn.tool_calls,
        reply_excerpt=turn.reply_excerpt[:_EXCERPT_CHARS],
        started_at=turn.started_at,
        finished_at=turn.finished_at,
    )


async def _employee_names(
    session: AsyncSession, organization_id: str, employee_ids: list[str | None]
) -> dict[str, str]:
    """Names of the employees a list refers to, read inside the organization only.

    One listing covers the team in the common case; an id outside that listing is
    looked up on its own, still scoped to the organization.
    """
    wanted = {eid for eid in employee_ids if eid}
    repo = EmployeeRepository(session)
    employees: list[Employee] = await repo.list_for_org(organization_id, limit=500)
    names = {e.id: e.name for e in employees if e.id in wanted}
    for employee_id in wanted:
        if employee_id not in names:
            employee = await repo.get_for_org(organization_id, employee_id)
            if employee is not None:
                names[employee_id] = employee.name
    return names


async def _task_titles(
    session: AsyncSession, organization_id: str, task_ids: list[str]
) -> dict[str, str]:
    """Titles of the tasks a list refers to, read inside the organization only."""
    wanted = set(task_ids)
    repo = TaskRepository(session)
    tasks: list[Task] = await repo.list_for_org(organization_id, limit=_BOARD_SCAN_LIMIT)
    titles = {t.id: t.title for t in tasks if t.id in wanted}
    for task_id in wanted:
        if task_id not in titles:
            task = await repo.get_for_org(organization_id, task_id)
            if task is not None:
                titles[task_id] = task.title
    return titles


async def _granted_tiers(
    session: AsyncSession, organization_id: str
) -> tuple[list[GrantedModel], dict[str, str]]:
    """The organization's granted models, and each one's neutral tier by model id."""
    granted = await build_model_grants(session).granted_models(organization_id=organization_id)
    return granted, {model.ai_model_id: model.tier for model in granted}


def _tier_list(granted: list[GrantedModel]) -> list[TierOut]:
    present = {model.tier for model in granted}
    return [
        TierOut(tier=tier, label=label) for tier, label in _TIER_LABELS.items() if tier in present
    ]


async def _member_out(
    session: AsyncSession,
    organization_id: str,
    employee: Employee,
    tiers: dict[str, str],
) -> TeamMemberOut:
    profile = await EmployeeProfileRepository(session).get(employee.id)
    tier = "unknown"
    busy_task_id: str | None = None
    busy_task_title: str | None = None
    hired_by = ""
    if profile is not None:
        tier = tiers.get(profile.ai_model_id or "", "unknown")
        busy_task_id = profile.busy_task_id
        hired_by = profile.hired_by
        if busy_task_id:
            busy = await TaskRepository(session).get_for_org(organization_id, busy_task_id)
            busy_task_title = busy.title if busy is not None else None
    return TeamMemberOut(
        id=employee.id,
        name=employee.name,
        role_key=str(employee.role_key),
        status=str(employee.status),
        tier=tier,
        busy_task_id=busy_task_id,
        busy_task_title=busy_task_title,
        hired_by=hired_by,
    )


# ---------------------------------------------------------------------------
# Routes
# ---------------------------------------------------------------------------


@router.get("/status", response_model=EngineStatusOut)
async def engine_status(
    context: OrgContext = Depends(get_org_context),
    session: AsyncSession = Depends(get_session),
) -> EngineStatusOut:
    """The credit balance, the owner's policy, the team size and the stage counts.

    ``credit_ok`` is false for a missing wallet as well as for a zero balance,
    because prepaid work needs credit to run (decision D-B).
    """
    organization_id = context.org_id
    wallet = await CreditWalletRepository(session).get_for_org(organization_id)
    policy = await build_model_grants(session).get_org_policy(organization_id)
    return EngineStatusOut(
        balance=wallet.balance if wallet is not None else 0.0,
        currency=wallet.currency if wallet is not None else "USD",
        credit_ok=wallet is not None and wallet.balance > 0,
        automation_enabled=policy.automation_enabled,
        max_employees=policy.max_employees,
        team_size=await EmployeeRepository(session).count_for_org(organization_id),
        stage_counts=await PipelineStateRepository(session).stage_counts(organization_id),
    )


@router.get("/tasks", response_model=list[EngineTaskOut])
async def list_engine_tasks(
    stage: Stage | None = Query(default=None),
    limit: int = Query(default=100, ge=1, le=500),
    context: OrgContext = Depends(get_org_context),
    session: AsyncSession = Depends(get_session),
) -> list[EngineTaskOut]:
    """The organization's tasks, newest first, with their stage when the engine owns them.

    With ``stage`` only the tasks at that stage are returned, found among the newest
    500 tasks of the organization.
    """
    organization_id = context.org_id
    states_repo = PipelineStateRepository(session)
    states = {
        s.task_id: s
        for s in await states_repo.list_for_org(organization_id, limit=_STATE_SCAN_LIMIT)
    }
    rows: list[tuple[Task, TaskPipelineState | None]] = []
    for task in await TaskRepository(session).list_for_org(
        organization_id, limit=_BOARD_SCAN_LIMIT
    ):
        state = states.get(task.id)
        if state is None:
            # A legacy task, or an engine task older than the newest states listed above.
            state = await states_repo.get(task.id)
        if stage is not None and (state is None or state.stage is not stage):
            continue
        rows.append((task, state))
        if len(rows) >= limit:
            break
    names = await _employee_names(
        session, organization_id, [task.assigned_employee_id for task, _ in rows]
    )
    return [_task_out(task, state, names) for task, state in rows]


@router.get("/tasks/{task_id}", response_model=EngineTaskDetailOut)
async def get_engine_task(
    task_id: str,
    context: OrgContext = Depends(get_org_context),
    session: AsyncSession = Depends(get_session),
) -> EngineTaskDetailOut:
    """One task with its stages, turns, gate verdicts, runs, diff and merged sha."""
    organization_id = context.org_id
    task = await TaskRepository(session).get_for_org(organization_id, task_id)
    if task is None:
        raise NotFoundError("کار یافت نشد.")
    state = await PipelineStateRepository(session).get(task.id)
    turns = await EngineTurnRepository(session).list_for_task(task.id, limit=_DETAIL_LIMIT)
    gates = await GateResultRepository(session).list_for_task(task.id, limit=_DETAIL_LIMIT)
    runs: list[Run] = await RunRepository(session).list_for_task(task.id)
    names = await _employee_names(
        session,
        organization_id,
        [task.assigned_employee_id, *(turn.employee_id for turn in turns)],
    )
    base = _task_out(task, state, names)
    return EngineTaskDetailOut(
        **base.model_dump(),
        description=task.description,
        stages=[str(stage) for stage in pipeline_stages(state.pipeline)] if state else [],
        waiver_reason=state.waiver_reason if state else "",
        branch=state.branch if state else "",
        depends_on=list(state.depends_on) if state else [],
        merged_sha=state.merged_sha if state else "",
        diff=state.last_diff if state else "",
        turns=[_turn_out(turn, names) for turn in turns],
        gates=[
            EngineGateOut(
                id=gate.id,
                gate=str(gate.gate),
                ok=gate.ok,
                detail=gate.detail,
                created_at=gate.created_at,
            )
            for gate in gates
        ],
        runs=[
            EngineRunOut(
                id=run.id,
                status=str(run.status),
                started_at=run.started_at,
                finished_at=run.finished_at,
            )
            for run in runs
        ],
    )


@router.get("/turns", response_model=list[EngineOrgTurnOut])
async def list_recent_turns(
    limit: int = Query(default=50, ge=1, le=200),
    context: OrgContext = Depends(get_org_context),
    session: AsyncSession = Depends(get_session),
) -> list[EngineOrgTurnOut]:
    """The organization's newest turns across all its tasks, with each task's title."""
    organization_id = context.org_id
    turns = await EngineTurnRepository(session).list_for_org(organization_id, limit=limit)
    names = await _employee_names(session, organization_id, [turn.employee_id for turn in turns])
    titles = await _task_titles(session, organization_id, [turn.task_id for turn in turns])
    return [
        EngineOrgTurnOut(
            **_turn_out(turn, names).model_dump(),
            task_id=turn.task_id,
            task_title=titles.get(turn.task_id, ""),
        )
        for turn in turns
    ]


# ---------------------------------------------------------------------------
# Actions on one task
# ---------------------------------------------------------------------------


async def _task_view(session: AsyncSession, organization_id: str, task_id: str) -> EngineTaskOut:
    task = await TaskRepository(session).get_for_org(organization_id, task_id)
    if task is None:
        raise NotFoundError("کار یافت نشد.")
    state = await PipelineStateRepository(session).get(task.id)
    names = await _employee_names(session, organization_id, [task.assigned_employee_id])
    return _task_out(task, state, names)


async def _control(
    *,
    action: Literal["pause", "resume", "cancel", "answer"],
    task_id: str,
    request: Request,
    context: OrgContext,
    session: AsyncSession,
    answer: str = "",
) -> EngineTaskOut:
    context.require(CAN_RUN_TASKS)
    # Imported inside the handler: the container wires the whole engine, and the read
    # routes above do not need it.
    from app.application.container import build_scheduler

    await build_scheduler(session).user_action(
        organization_id=context.org_id,
        task_id=task_id,
        action=action,
        user_id=context.user.id,
        answer=answer,
    )
    ip, ua = client_meta(request)
    # Only the action is written: an answer's text stays out of the audit trail.
    await record_audit(
        session,
        action=AuditAction.TASK_CONTROL,
        organization_id=context.org_id,
        actor_user_id=context.user.id,
        target_type="task",
        target_id=task_id,
        metadata={"action": action},
        ip=ip,
        user_agent=ua,
    )
    return await _task_view(session, context.org_id, task_id)


@router.post("/tasks/{task_id}/pause", response_model=EngineTaskOut)
async def pause_task(
    task_id: str,
    request: Request,
    context: OrgContext = Depends(get_org_context),
    session: AsyncSession = Depends(get_session),
) -> EngineTaskOut:
    """Stop a task before its next turn. Its stage is kept, so resume picks it up."""
    return await _control(
        action="pause", task_id=task_id, request=request, context=context, session=session
    )


@router.post("/tasks/{task_id}/resume", response_model=EngineTaskOut)
async def resume_task(
    task_id: str,
    request: Request,
    context: OrgContext = Depends(get_org_context),
    session: AsyncSession = Depends(get_session),
) -> EngineTaskOut:
    """Send a paused task back to triage. Refused while the organization has no credit."""
    return await _control(
        action="resume", task_id=task_id, request=request, context=context, session=session
    )


@router.post("/tasks/{task_id}/cancel", response_model=EngineTaskOut)
async def cancel_task(
    task_id: str,
    request: Request,
    context: OrgContext = Depends(get_org_context),
    session: AsyncSession = Depends(get_session),
) -> EngineTaskOut:
    """End a task. Its worktree is released and tasks waiting on it are parked."""
    return await _control(
        action="cancel", task_id=task_id, request=request, context=context, session=session
    )


@router.post("/tasks/{task_id}/answer", response_model=EngineTaskOut)
async def answer_task(
    task_id: str,
    body: EngineAnswerRequest,
    request: Request,
    context: OrgContext = Depends(get_org_context),
    session: AsyncSession = Depends(get_session),
) -> EngineTaskOut:
    """Answer the question a task waits on. Only a task at a decision accepts one."""
    return await _control(
        action="answer",
        task_id=task_id,
        request=request,
        context=context,
        session=session,
        answer=body.content,
    )


# ---------------------------------------------------------------------------
# Team: the view, hire by role and tier, fire
# ---------------------------------------------------------------------------


@router.get("/team", response_model=TeamOut)
async def get_team(
    context: OrgContext = Depends(get_org_context),
    session: AsyncSession = Depends(get_session),
) -> TeamOut:
    """The team, without the archived members, with each member's neutral tier."""
    organization_id = context.org_id
    granted, tiers = await _granted_tiers(session, organization_id)
    policy = await build_model_grants(session).get_org_policy(organization_id)
    employees = [
        employee
        for employee in await EmployeeRepository(session).list_for_org(organization_id, limit=500)
        if employee.status is not EmployeeStatus.ARCHIVED
    ]
    members = [await _member_out(session, organization_id, e, tiers) for e in employees]
    return TeamOut(
        employees=members,
        tiers=_tier_list(granted),
        max_employees=policy.max_employees,
        count=len(members),
    )


@router.post("/team/hire", response_model=TeamMemberOut, status_code=201)
async def hire_member(
    body: HireRequest,
    request: Request,
    context: OrgContext = Depends(get_org_context),
    session: AsyncSession = Depends(get_session),
) -> TeamMemberOut:
    """Hire one employee for a role on a granted model of the requested tier."""
    context.require(CAN_MANAGE_EMPLOYEES)
    # Imported inside the handler, as for the task actions above.
    from app.application.container import build_scheduler

    employee = await build_scheduler(session).manager.hire(
        organization_id=context.org_id,
        role=body.role,
        tier=body.tier,
        name=body.name,
        project_id=None,
    )
    ip, ua = client_meta(request)
    await record_audit(
        session,
        action=AuditAction.EMPLOYEE_CREATE,
        organization_id=context.org_id,
        actor_user_id=context.user.id,
        target_type="employee",
        target_id=employee.id,
        metadata={"role_key": str(body.role), "tier": body.tier},
        ip=ip,
        user_agent=ua,
    )
    _, tiers = await _granted_tiers(session, context.org_id)
    return await _member_out(session, context.org_id, employee, tiers)


@router.post("/team/{employee_id}/fire", response_model=TeamMemberOut)
async def fire_member(
    employee_id: str,
    request: Request,
    context: OrgContext = Depends(get_org_context),
    session: AsyncSession = Depends(get_session),
) -> TeamMemberOut:
    """Archive one employee. The last active manager cannot be fired."""
    context.require(CAN_MANAGE_EMPLOYEES)
    from app.application.container import build_scheduler

    employee = await build_scheduler(session).manager.fire(
        organization_id=context.org_id, employee_id=employee_id
    )
    ip, ua = client_meta(request)
    await record_audit(
        session,
        action=AuditAction.EMPLOYEE_UPDATE,
        organization_id=context.org_id,
        actor_user_id=context.user.id,
        target_type="employee",
        target_id=employee.id,
        metadata={"action": "fire"},
        ip=ip,
        user_agent=ua,
    )
    _, tiers = await _granted_tiers(session, context.org_id)
    return await _member_out(session, context.org_id, employee, tiers)