beaver_project/app-instance/backend/beaver/skills/learning/pipeline.py

"""Manual skill learning pipeline orchestration."""

from __future__ import annotations

from typing import Any, Callable

from beaver.engine.providers import ProviderBundle
from beaver.memory.skills import SkillDraftEvalReport, SkillDraftSafetyReport, SkillLearningCandidate, SkillLearningStore
from beaver.skills.drafts import DraftService
from beaver.skills.learning.eval import SkillDraftEvaluator
from beaver.skills.learning.replay import ReplayRunner
from beaver.skills.learning.service import NoDraftChanges, SkillLearningService
from beaver.skills.learning.safety import SkillDraftSafetyChecker
from beaver.skills.publisher import SkillPublisher
from beaver.skills.reviews import ReviewService
from beaver.skills.specs import SkillDraft, SkillReviewRecord, SkillReviewState, SkillSpec, SkillVersion

_REJECTABLE_DRAFT_STATUSES = {
    SkillReviewState.DRAFT.value,
    SkillReviewState.IN_REVIEW.value,
    SkillReviewState.APPROVED.value,
}


class DraftSynthesisInProgress(RuntimeError):
    """Raised when another request already claimed the candidate for synthesis."""


class DraftHasNoChanges(RuntimeError):
    """Raised when synthesis produced no effective changes from the base skill."""


class SkillLearningPipelineService:
    """Coordinates candidate -> draft -> review -> publish lifecycle."""

    def __init__(
        self,
        *,
        learning_store: SkillLearningStore,
        learning_service: SkillLearningService,
        draft_service: DraftService,
        review_service: ReviewService,
        publisher: SkillPublisher,
        safety_checker: SkillDraftSafetyChecker | None = None,
        evaluator: SkillDraftEvaluator | None = None,
        publish_observer: Callable[[SkillDraft, SkillVersion | SkillSpec], None] | None = None,
    ) -> None:
        self.learning_store = learning_store
        self.learning_service = learning_service
        self.draft_service = draft_service
        self.review_service = review_service
        self.publisher = publisher
        self.safety_checker = safety_checker or SkillDraftSafetyChecker()
        self.evaluator = evaluator
        self.publish_observer = publish_observer

    def list_candidates(self, status: str | None = None) -> list[SkillLearningCandidate]:
        return self.learning_store.list_learning_candidates(status=status)

    def get_candidate(self, candidate_id: str) -> SkillLearningCandidate:
        for candidate in self.learning_store.list_learning_candidates():
            if candidate.candidate_id == candidate_id:
                return candidate
        raise ValueError(f"Unknown learning candidate: {candidate_id}")

    async def synthesize_draft(
        self,
        candidate_id: str,
        *,
        provider_bundle: ProviderBundle,
        force: bool = False,
    ) -> SkillDraft:
        if not force:
            existing = self._draft_for_candidate(candidate_id)
            if existing is not None:
                return existing
        claimed = self.learning_store.claim_learning_candidate_for_synthesis(candidate_id, force=force)
        if claimed is None:
            existing = self._draft_for_candidate(candidate_id)
            if existing is not None:
                return existing
            raise DraftSynthesisInProgress(f"Draft synthesis is already in progress for candidate: {candidate_id}")
        try:
            draft = await self.learning_service.synthesize_draft(candidate_id, provider_bundle)
        except NoDraftChanges as exc:
            self.mark_candidate_superseded(candidate_id, str(exc))
            raise DraftHasNoChanges(str(exc)) from exc
        self.mark_draft_synthesized(candidate_id, draft)
        return draft

    async def regenerate_draft(
        self,
        candidate_id: str,
        *,
        provider_bundle: ProviderBundle,
    ) -> SkillDraft:
        return await self.synthesize_draft(candidate_id, provider_bundle=provider_bundle, force=True)

    def mark_candidate_queued(self, candidate_id: str) -> SkillLearningCandidate:
        return self._require_updated(
            self.learning_store.transition_learning_candidate(
                candidate_id,
                "queued",
                event_type="candidate_queued",
                last_error=None,
            ),
            candidate_id,
        )

    def mark_candidate_synthesizing(self, candidate_id: str) -> SkillLearningCandidate:
        return self._require_updated(
            self.learning_store.transition_learning_candidate(
                candidate_id,
                "synthesizing",
                event_type="draft_synthesis_started",
                last_error=None,
            ),
            candidate_id,
        )

    def mark_draft_synthesized(self, candidate_id: str, draft: SkillDraft) -> SkillLearningCandidate:
        candidate = self.get_candidate(candidate_id)
        evidence = dict(candidate.evidence)
        evidence["draft_id"] = draft.draft_id
        evidence["draft_skill_name"] = draft.skill_name
        return self._require_updated(
            self.learning_store.transition_learning_candidate(
                candidate_id,
                "draft_ready",
                event_type="draft_synthesis_completed",
                evidence=evidence,
                draft_id=draft.draft_id,
                draft_skill_name=draft.skill_name,
                risk_level=candidate.risk_level,
                last_error=None,
                payload={"draft_id": draft.draft_id, "skill_name": draft.skill_name},
            ),
            candidate_id,
        )

    def mark_candidate_failed(
        self,
        candidate_id: str,
        error: str,
        *,
        retry_count: int,
        terminal: bool,
    ) -> SkillLearningCandidate:
        return self._require_updated(
            self.learning_store.transition_learning_candidate(
                candidate_id,
                "failed" if terminal else "open",
                event_type="failed",
                retry_count=retry_count,
                last_error=error,
                payload={"error": error, "terminal": terminal, "retry_count": retry_count},
            ),
            candidate_id,
        )

    def mark_candidate_superseded(self, candidate_id: str, reason: str) -> SkillLearningCandidate:
        return self._require_updated(
            self.learning_store.transition_learning_candidate(
                candidate_id,
                "superseded",
                event_type="superseded",
                last_error=reason,
                payload={"reason": reason},
            ),
            candidate_id,
        )

    def list_drafts(self, skill_name: str | None = None) -> list[SkillDraft]:
        return self.draft_service.list_drafts(skill_name)

    def get_draft(self, skill_name: str, draft_id: str) -> SkillDraft:
        draft = self.draft_service.get_draft(skill_name, draft_id)
        if draft is None:
            raise ValueError(f"Draft not found: {skill_name}/{draft_id}")
        return draft

    def _draft_for_candidate(self, candidate_id: str) -> SkillDraft | None:
        candidate = self.get_candidate(candidate_id)
        if not candidate.draft_skill_name or not candidate.draft_id:
            return None
        return self.draft_service.get_draft(candidate.draft_skill_name, candidate.draft_id)

    def submit_review(
        self,
        skill_name: str,
        draft_id: str,
        *,
        requested_by: str = "system",
        notes: str = "",
    ) -> SkillReviewRecord:
        draft = self.get_draft(skill_name, draft_id)
        if draft.status != SkillReviewState.DRAFT.value:
            raise ValueError("Draft must be in draft status before review submission")
        safety = self.get_safety_report(skill_name, draft_id)
        if safety is not None and (not safety.passed or safety.risk_level == "critical"):
            raise ValueError("Draft cannot enter review because safety check failed")
        review = self.review_service.submit_for_review(
            skill_name,
            draft_id,
            reviewer_request=notes,
            requested_by=requested_by,
        )
        self._mark_candidate_by_draft(
            skill_name,
            draft_id,
            "review_pending",
            "review_submitted",
            last_error=None,
        )
        return review

    def approve(
        self,
        skill_name: str,
        draft_id: str,
        *,
        reviewer: str = "system",
        notes: str = "",
    ) -> SkillReviewRecord:
        draft = self.get_draft(skill_name, draft_id)
        if draft.status != SkillReviewState.IN_REVIEW.value:
            raise ValueError("Draft must be in review before approval")
        safety = self.get_safety_report(skill_name, draft_id)
        if safety is not None and (not safety.passed or safety.risk_level == "critical"):
            raise ValueError("Draft cannot be approved because safety check failed")
        review = self.review_service.approve(skill_name, draft_id, reviewer=reviewer, notes=notes)
        self._mark_candidate_by_draft(skill_name, draft_id, "approved", "approved")
        return review

    def reject(
        self,
        skill_name: str,
        draft_id: str,
        *,
        reviewer: str = "system",
        notes: str = "",
    ) -> SkillReviewRecord:
        draft = self.get_draft(skill_name, draft_id)
        if draft.status not in _REJECTABLE_DRAFT_STATUSES:
            raise ValueError("Draft is not rejectable from its current status")
        review = self.review_service.reject(skill_name, draft_id, reviewer=reviewer, notes=notes)
        self._mark_candidate_by_draft(skill_name, draft_id, "rejected", "rejected")
        return review

    def publish(
        self,
        skill_name: str,
        draft_id: str,
        *,
        publisher: str = "system",
        notes: str = "",
        confirm_high_risk: bool = False,
    ) -> SkillVersion | SkillSpec:
        draft = self.get_draft(skill_name, draft_id)
        self._validate_publish_gates(draft, confirm_high_risk=confirm_high_risk)
        if draft.proposal_kind == "retire_skill":
            result = self.publisher.apply_retire_proposal(skill_name, draft_id, actor=publisher, notes=notes)
        else:
            result = self.publisher.publish(skill_name, draft_id, publisher=publisher, notes=notes)
        self._mark_candidate_by_draft(skill_name, draft_id, "published", "published")
        if self.publish_observer is not None:
            try:
                self.publish_observer(draft, result)
            except Exception as exc:  # noqa: BLE001 - observer is best effort after successful publish.
                candidate = self._candidate_by_draft(skill_name, draft_id)
                self.learning_store.append_audit_event(
                    candidate.candidate_id if candidate is not None else f"draft:{draft_id}",
                    "plugin_publish_ack_failed",
                    {"error": str(exc), "skill_name": skill_name, "draft_id": draft_id},
                )
        return result

    def rollback(
        self,
        skill_name: str,
        target_version: str,
        *,
        actor: str = "system",
        reason: str = "",
    ) -> SkillSpec:
        return self.publisher.rollback(skill_name, target_version, actor=actor, reason=reason or "manual rollback")

    def disable(
        self,
        skill_name: str,
        *,
        actor: str = "system",
        reason: str = "",
    ) -> SkillSpec:
        return self.publisher.disable(skill_name, actor=actor, reason=reason or "manual disable")

    def reviews_for_draft(self, skill_name: str, draft_id: str) -> list[SkillReviewRecord]:
        return self.review_service.store.list_reviews(skill_name, draft_id=draft_id)

    def check_safety(self, skill_name: str, draft_id: str) -> SkillDraftSafetyReport:
        draft = self.get_draft(skill_name, draft_id)
        report = self.safety_checker.check(draft)
        self.learning_store.write_safety_report(report)
        status = (
            "safety_failed"
            if not report.passed or report.risk_level == "critical"
            else self._candidate_status_for_draft(draft)
        )
        current = self._candidate_by_draft(skill_name, draft_id)
        if current is not None and current.status == "eval_failed" and status != "safety_failed":
            status = "eval_failed"
        self._mark_candidate_by_draft(
            skill_name,
            draft_id,
            status,
            "safety_checked",
            safety_report_id=report.report_id,
            risk_level=report.risk_level,
            last_error="; ".join(report.blocked_reasons) if status == "safety_failed" else None,
        )
        return report

    def get_safety_report(self, skill_name: str, draft_id: str) -> SkillDraftSafetyReport | None:
        return self.learning_store.get_safety_report(skill_name, draft_id)

    def get_eval_report(self, skill_name: str, draft_id: str) -> SkillDraftEvalReport | None:
        return self.learning_store.get_eval_report(skill_name, draft_id)

    async def evaluate_draft(
        self,
        candidate_id: str,
        skill_name: str,
        draft_id: str,
        *,
        provider_bundle: ProviderBundle | None,
        replay_runner: ReplayRunner | None = None,
        progress_callback: Callable[[dict[str, Any]], None] | None = None,
    ) -> SkillDraftEvalReport:
        draft = self.get_draft(skill_name, draft_id)
        candidate = self.get_candidate(candidate_id)
        evaluator = self.evaluator or SkillDraftEvaluator(
            self.learning_service.run_store,
            skill_store=self.draft_service.store,
        )
        report = await evaluator.evaluate(
            candidate=candidate,
            draft=draft,
            provider_bundle=provider_bundle,
            replay_runner=replay_runner,
            progress_callback=progress_callback,
        )
        self.learning_store.write_eval_report(report)
        if report.status == "skipped_provider_unavailable":
            status = self._candidate_status_for_draft(draft)
            error = "eval skipped: provider unavailable"
        elif report.passed:
            status = self._candidate_status_for_draft(draft)
            error = None
        else:
            status = "eval_failed"
            error = "eval failed"
        current = self._candidate_by_draft(skill_name, draft_id)
        if current is not None and current.status == "safety_failed" and status == "draft_ready":
            status = "safety_failed"
            error = current.last_error
        self.learning_store.transition_learning_candidate(
            candidate_id,
            status,
            event_type="eval_completed",
            eval_report_id=report.report_id,
            eval_progress={
                "phase": "completed",
                "completed_arms": len(report.cases) * 2 if report.mode == "replay" else 0,
                "total_arms": len(report.cases) * 2 if report.mode == "replay" else 0,
                "completed_cases": len(report.cases),
                "total_cases": len(report.cases),
            },
            last_error=error,
            payload=report.to_dict(),
        )
        return report

    def mark_eval_progress(self, candidate_id: str, progress: dict[str, Any]) -> SkillLearningCandidate:
        return self._require_updated(
            self.learning_store.update_learning_candidate(
                candidate_id,
                eval_progress=dict(progress),
            ),
            candidate_id,
        )

    def mark_eval_failed(self, candidate_id: str, error: str) -> SkillLearningCandidate:
        candidate = self.get_candidate(candidate_id)
        progress = dict(candidate.eval_progress)
        progress["phase"] = "failed"
        return self._require_updated(
            self.learning_store.transition_learning_candidate(
                candidate_id,
                "eval_failed",
                eval_progress=progress,
                event_type="eval_failed",
                last_error=error,
                payload={"error": error},
            ),
            candidate_id,
        )

    def _validate_publish_gates(self, draft: SkillDraft, *, confirm_high_risk: bool) -> None:
        reviews = self.reviews_for_draft(draft.skill_name, draft.draft_id)
        if not any(review.status in {SkillReviewState.IN_REVIEW.value, SkillReviewState.APPROVED.value} for review in reviews):
            raise ValueError("Draft must be submitted for review before publish")
        safety = self.get_safety_report(draft.skill_name, draft.draft_id)
        if safety is None:
            raise ValueError("Draft requires a passing safety report before publish")
        if not safety.passed:
            raise ValueError("Draft safety report did not pass")
        if safety.risk_level == "critical":
            raise ValueError("Critical risk drafts cannot be published")
        if safety.risk_level == "high" and not confirm_high_risk:
            raise ValueError("High risk draft publish requires confirm_high_risk=true")
        eval_report = self.get_eval_report(draft.skill_name, draft.draft_id)
        if eval_report is not None and eval_report.status != "skipped_provider_unavailable" and not eval_report.passed:
            raise ValueError("Draft eval report did not pass")
        if eval_report is not None and eval_report.mode == "replay":
            if eval_report.confidence == "low":
                raise ValueError("Draft replay eval has low confidence and requires revision before publish")
            if eval_report.blocked_coverage >= 1.0:
                raise ValueError("Draft replay eval blocked all important tool calls")
            preservation = eval_report.preservation_report or {}
            if preservation.get("passed") is False:
                raise ValueError("Draft preservation check did not pass")
            if draft.proposal_kind == "plugin_skill_update":
                if draft.provenance.get("merge_mode") == "three_way" and preservation.get("mode") != "plugin_three_way":
                    raise ValueError("Plugin update requires a three-way preservation report")
                if preservation.get("unresolved_conflicts"):
                    raise ValueError("Plugin update has unresolved merge conflicts")
                supporting_plan = draft.provenance.get("supporting_file_plan")
                if isinstance(supporting_plan, dict) and supporting_plan.get("conflicts"):
                    raise ValueError("Plugin update has unresolved supporting-file conflicts")

    def _mark_candidate_by_draft(
        self,
        skill_name: str,
        draft_id: str,
        status: str,
        event_type: str,
        **updates: object,
    ) -> SkillLearningCandidate | None:
        candidate = self._candidate_by_draft(skill_name, draft_id)
        if candidate is None:
            return None
        if candidate.status in {"safety_failed", "eval_failed"} and status in {"review_pending", "approved"}:
            return candidate
        return self.learning_store.transition_learning_candidate(
            candidate.candidate_id,
            status,
            event_type=event_type,
            **updates,
        )

    def _candidate_by_draft(self, skill_name: str, draft_id: str) -> SkillLearningCandidate | None:
        for candidate in self.learning_store.list_learning_candidates():
            if candidate.draft_skill_name == skill_name and candidate.draft_id == draft_id:
                return candidate
        return None

    @staticmethod
    def _candidate_status_for_draft(draft: SkillDraft) -> str:
        if draft.status == SkillReviewState.APPROVED.value:
            return "approved"
        if draft.status == SkillReviewState.IN_REVIEW.value:
            return "review_pending"
        return "draft_ready"

    @staticmethod
    def _require_updated(candidate: SkillLearningCandidate | None, candidate_id: str) -> SkillLearningCandidate:
        if candidate is None:
            raise ValueError(f"Unknown learning candidate: {candidate_id}")
        return candidate


def model_to_dict(value: Any) -> dict[str, Any]:
    if hasattr(value, "to_dict"):
        return value.to_dict()
    if isinstance(value, dict):
        return dict(value)
    raise TypeError(f"Cannot convert {type(value).__name__} to dict")