insights/products/wizard/backend/logic/sessions.py

"""
Business logic for Wizard sessions.
"""

import logging
from functools import partial

from django.db import transaction

from products.wizard.backend.facade.contracts import (
    UpsertWizardSessionInput,
    WizardSessionDTO,
    WizardSessionOwnershipError,
    WizardSessionUserDTO,
    WizardTaskDTO,
)
from products.wizard.backend.facade.enums import RunPhase, TaskStatus
from products.wizard.backend.logic.pubsub import publish_session_update
from products.wizard.backend.logic.utils import is_stale
from products.wizard.backend.metrics import report_session_upserted
from products.wizard.backend.models import WizardSession
from products.wizard.backend.tasks.tasks import sync_wizard_event_definitions

logger = logging.getLogger(__name__)


def upsert_session(params: UpsertWizardSessionInput) -> tuple[WizardSessionDTO, bool]:
    """Upsert a session row and return (dto, created).

    `defaults=` is applied on both insert AND update — each push replaces
    `tasks` / `run_phase` / `event_plan` / `error`, except that a completed
    push without an event plan preserves the plan from the running session.
    Concurrent POSTs for a brand-new session_id can race the unique constraint
    and surface as a 500; the CLI's normal HTTP retry handles that on the next
    attempt.
    """
    with transaction.atomic():
        previous_session = (
            WizardSession.objects.select_for_update()
            .filter(team_id=params.team_id, session_id=params.session_id)
            .first()
        )
        previous_run_phase = previous_session.run_phase if previous_session else None

        # A session belongs to the user who created it. A later push from a
        # different user would overwrite the run data while `created_by` stays
        # the original owner, so reject it. Legacy rows with a null
        # `created_by_id` predate attribution and stay updatable by anyone.
        if (
            previous_session is not None
            and previous_session.created_by_id is not None
            and previous_session.created_by_id != params.created_by_id
        ):
            raise WizardSessionOwnershipError("This wizard session belongs to a different user and can't be updated.")

        event_plan = params.event_plan
        if event_plan is None and params.run_phase == RunPhase.COMPLETED and previous_session:
            event_plan = previous_session.event_plan

        # Monotonic within a session: the doc arrives late in the run, so a push without it
        # (an ordering race between debounced snapshots, or an older CLI) must not wipe it.
        # A new run is a new session_id, so nothing ever needs to clear the field.
        handoff_text = params.handoff_text
        if not handoff_text and previous_session:
            handoff_text = previous_session.handoff_text

        defaults = {
            "workflow_id": params.workflow_id,
            "skill_id": params.skill_id,
            "started_at": params.started_at,
            "run_phase": params.run_phase.value,
            "tasks": [
                {
                    "id": task.id,
                    "title": task.title,
                    "status": task.status.value,
                }
                for task in params.tasks
            ],
            "event_plan": event_plan,
            "error": params.error,
            "pending_input": params.pending_input,
            "handoff_text": handoff_text,
        }
        # created_by only in create_defaults so a later push for the same run can't reattribute it.
        instance, created = WizardSession.objects.update_or_create(
            team_id=params.team_id,
            session_id=params.session_id,
            defaults=defaults,
            create_defaults={**defaults, "created_by_id": params.created_by_id},
        )
        if previous_run_phase != RunPhase.COMPLETED.value and params.run_phase == RunPhase.COMPLETED:
            transaction.on_commit(
                partial(_enqueue_event_definition_sync, params.team_id, params.session_id),
                robust=True,
            )
        dto = _to_dto(instance)

    report_session_upserted(previous_run_phase, dto)
    publish_session_update(dto)
    return dto, created


def _enqueue_event_definition_sync(team_id: int, session_id: str) -> None:
    try:
        sync_wizard_event_definitions.delay(team_id, session_id)
    except Exception:
        logger.exception("Failed to enqueue event definition sync for a completed wizard session")


def get_session(team_id: int, session_id: str) -> WizardSessionDTO | None:
    instance = WizardSession.objects.select_related("created_by").filter(team_id=team_id, session_id=session_id).first()
    return _to_dto(instance) if instance else None


def get_latest_session(team_id: int, workflow_id: str, skill_id: str | None = None) -> WizardSessionDTO | None:
    qs = WizardSession.objects.select_related("created_by").filter(team_id=team_id, workflow_id=workflow_id)
    if skill_id:
        qs = qs.filter(skill_id=skill_id)
    # created_at breaks ties on equal (client-supplied, second-granularity) started_at
    instance = qs.order_by("-started_at", "-created_at").first()
    return _to_dto(instance) if instance else None


def list_sessions(
    team_id: int,
    workflow_id: str | None = None,
    skill_id: str | None = None,
    *,
    offset: int = 0,
    limit: int | None = None,
) -> list[WizardSessionDTO]:
    """List sessions for a team, ordered by `started_at` desc.

    `offset`/`limit` are applied at the SQL layer (LIMIT/OFFSET) so the read
    cost stays bounded regardless of how many sessions the team has. The view
    layer should always pass a `limit`.
    """
    qs = WizardSession.objects.select_related("created_by").filter(team_id=team_id)
    if workflow_id:
        qs = qs.filter(workflow_id=workflow_id)
    if skill_id:
        qs = qs.filter(skill_id=skill_id)
    # created_at breaks ties on equal (client-supplied, second-granularity) started_at
    qs = qs.order_by("-started_at", "-created_at")
    if limit is not None:
        qs = qs[offset : offset + limit]
    elif offset:
        qs = qs[offset:]
    return [_to_dto(instance) for instance in qs]


def _to_dto(instance: WizardSession) -> WizardSessionDTO:
    run_phase = RunPhase(instance.run_phase)
    created_by = instance.created_by
    return WizardSessionDTO(
        session_id=instance.session_id,
        team_id=instance.team_id,
        workflow_id=instance.workflow_id,
        skill_id=instance.skill_id,
        started_at=instance.started_at,
        run_phase=run_phase,
        is_stale=is_stale(run_phase, instance.updated_at),
        tasks=tuple(
            WizardTaskDTO(
                id=task["id"],
                title=task["title"],
                status=TaskStatus(task["status"]),
            )
            for task in (instance.tasks or [])
        ),
        event_plan=instance.event_plan,
        error=instance.error,
        pending_input=instance.pending_input,
        handoff_text=instance.handoff_text,
        created_by=(
            WizardSessionUserDTO(id=created_by.id, first_name=created_by.first_name, email=created_by.email)
            if created_by is not None
            else None
        ),
        created_at=instance.created_at,
        updated_at=instance.updated_at,
    )