Artemis-Eapil/artemis/llm/final_review.py

"""One bounded, observable budget for a read-only final acceptance attempt."""

import asyncio
from collections.abc import Awaitable, Callable
from contextvars import ContextVar
import logging
import sys
import time
from uuid import uuid4

from artemis.data_engine.context_vars import CURRENT_TRACE_ID
from artemis.llm.reliability import LLMCallError, LLMExhaustedError

logger = logging.getLogger(__name__)

# Allow evidence selection, a verdict, and (when needed) structured formatting,
# plus one shared retry allowance. Do not multiply this by the checker's loop
# limit: every request, tool read, backoff and breaker wait spends the same pool.
FINAL_REVIEW_REQUEST_SLOTS = 3
FINAL_REVIEW_TOOL_ALLOWANCE_SECONDS = 60.0
MAX_FINAL_REVIEW_SECONDS = 3600.0

ACTIVE_FINAL_REVIEW: ContextVar["FinalReviewBudget | None"] = ContextVar(
    "final_review_budget", default=None
)


class FinalReviewBudgetExceeded(TimeoutError):
    def __init__(self, seconds: float):
        super().__init__(f"最终验收总预算已用尽（{seconds:g} 秒），尚未取得完整结论")


class FinalReviewBudget:
    def __init__(
        self,
        engine,
        *,
        attempt_id: str,
        request_timeout: float,
        max_retries: int,
        retry_delay: float,
    ):
        self.engine = engine
        self.trace_id = uuid4()
        self.parent_id = CURRENT_TRACE_ID.get()
        self.total_seconds = min(
            MAX_FINAL_REVIEW_SECONDS,
            request_timeout * (FINAL_REVIEW_REQUEST_SLOTS + max_retries)
            + retry_delay * max_retries
            + FINAL_REVIEW_TOOL_ALLOWANCE_SECONDS,
        )
        self.clock = time.monotonic()
        self.timer: asyncio.Timeout | None = None
        self.data = {
            "review_id": str(self.trace_id),
            "attempt_id": attempt_id,
            "status": "running",
            "stage": "preparing_evidence",
            "iteration": 0,
            "request_count": 0,
            "retry_count": 0,
            "sequence": 0,
            "total_seconds": self.total_seconds,
            "request_timeout_seconds": request_timeout,
            "max_retries": max_retries,
            "retry_allowance_seconds": retry_delay * max_retries,
        }

    @property
    def expired(self) -> bool:
        """Only our timer, never a user/guidance cancellation, exhausts the budget."""
        return self.timer is not None and self.timer.expired()

    def snapshot(self) -> dict:
        elapsed = self.data.get("elapsed_seconds", time.monotonic() - self.clock)
        if self.data["status"] == "running":
            elapsed = time.monotonic() - self.clock
        return {
            **self.data,
            "elapsed_seconds": round(elapsed, 3),
            "remaining_seconds": round(max(0.0, self.total_seconds - elapsed), 3),
        }

    def publish(self) -> None:
        self.data["sequence"] += 1
        self.data["updated_at"] = time.time()
        if not self.engine or not getattr(self.engine, "current_session_id", None):
            return
        try:
            self.engine.record_trace(
                type="final_review_diagnostic",
                name="final_review",
                trace_id=self.trace_id,
                step_id=getattr(self.engine, "current_step_id", None),
                parent_trace_id=self.parent_id,
                payload=self.snapshot(),
                status="running"
                if self.data["status"] == "running"
                else "success"
                if self.data["status"] == "completed"
                else "failed",
                duration=self.snapshot()["elapsed_seconds"],
            )
        except (OSError, RuntimeError, ValueError, TypeError, AttributeError):
            logger.warning("Could not persist final acceptance diagnostics", exc_info=True)

    def stage(self, stage: str, *, iteration: int | None = None) -> None:
        self.data["stage"] = stage
        if iteration is not None:
            self.data["iteration"] = iteration
        self.publish()

    def request_started(self, request_id: str, attempt: int) -> None:
        self.data["request_count" if attempt == 1 else "retry_count"] += 1
        self.data.update(
            request_id=request_id, request_attempt=attempt, request_phase="waiting_response"
        )
        self.publish()

    def request_finished(self, phase: str, category: str | None = None) -> None:
        self.data["request_phase"] = phase
        if category:
            self.data["last_failure_category"] = category
        self.publish()

    def finish(self, status: str, *, error: BaseException | None = None) -> None:
        self.data.update(
            status=status,
            ended_at=time.time(),
            elapsed_seconds=time.monotonic() - self.clock,
        )
        if error is not None:
            self.data["error_type"] = type(error).__name__
            if isinstance(error, LLMCallError):
                self.data["last_failure_category"] = error.failure.category.value
        self.publish()

    async def run[T](self, operation: Callable[[], Awaitable[T]]) -> T:
        self.clock = time.monotonic()
        self.data["started_at"] = time.time()
        self.data["deadline_at"] = self.data["started_at"] + self.total_seconds
        token = ACTIVE_FINAL_REVIEW.set(self)
        self.timer = asyncio.timeout(self.total_seconds)
        self.publish()
        try:
            async with self.timer:
                result = await operation()
        except (KeyboardInterrupt, asyncio.CancelledError) as exc:
            self.finish("cancelled", error=exc)
            raise
        except TimeoutError as exc:
            if self.expired:
                error = FinalReviewBudgetExceeded(self.total_seconds)
                self.finish("budget_exhausted", error=error)
                raise error from exc
            raise
        except LLMExhaustedError as exc:
            self.finish("request_retries_exhausted", error=exc)
            raise
        else:
            self.finish("completed")
            return result
        finally:
            if self.data["status"] == "running":
                self.finish("failed", error=sys.exception())
            ACTIVE_FINAL_REVIEW.reset(token)


def final_review_stage(stage: str, *, iteration: int | None = None) -> None:
    if review := ACTIVE_FINAL_REVIEW.get():
        review.stage(stage, iteration=iteration)