EvoScientist-Multi/EvoScientist/middleware/dynamic_review.py

"""Run-scoped automatic review for Ai4Sci Web executions."""

from __future__ import annotations

import os
from collections.abc import Mapping
from typing import Annotated, Any, NotRequired

import httpx

from EvoScientist.internal_service import internal_service_headers
from langchain.agents.middleware import HumanInTheLoopMiddleware, hook_config
from langchain.agents.middleware.types import AgentState, OmitFromSchema
from langgraph.config import get_config


class DynamicReviewState(AgentState[Any]):
    _verified_review_mode: NotRequired[
        Annotated[dict[str, Any], OmitFromSchema(input=True, output=True)]
    ]


class AutoReviewVerificationError(RuntimeError):
    """Raised when an automatic review request cannot be verified."""


def _review_context() -> tuple[str, Mapping[str, Any] | None]:
    try:
        config = get_config()
    except RuntimeError:
        return "", None
    configurable = config.get("configurable") if isinstance(config, Mapping) else None
    if not isinstance(configurable, Mapping):
        return "", None
    run_id = str(configurable.get("ai4sci_run_id") or "")
    review = configurable.get("ai4sci_review_mode")
    return run_id, review if isinstance(review, Mapping) else None


def _manual_state(run_id: str, review: Mapping[str, Any] | None) -> dict[str, Any]:
    revision = review.get("review_mode_revision") if review is not None else 0
    return {
        "protocol": "verified-review-mode-state-v1",
        "execution_run_id": run_id,
        "mode": "manual",
        "revision": revision if isinstance(revision, int) and revision >= 0 else 0,
    }


def _request_payload(
    run_id: str, review: Mapping[str, Any]
) -> tuple[str, dict[str, str]]:
    configured_url = str(review.get("gateway_url") or "").rstrip("/")
    gateway_url = os.environ.get("GATEWAY_INTERNAL_URL", "").strip().rstrip("/")
    gateway_url = gateway_url or configured_url
    payload = {
        "run_id": run_id,
        "envelope_digest": str(review.get("envelope_digest") or ""),
        "envelope_signature": str(review.get("envelope_signature") or ""),
    }
    if not gateway_url or not all(payload.values()):
        raise AutoReviewVerificationError("AUTO_REVIEW_CONTEXT_INVALID")
    return gateway_url, payload


def _service_headers() -> dict[str, str]:
    return internal_service_headers()


def _validated_auto_state(
    run_id: str,
    review: Mapping[str, Any],
    resolved: Mapping[str, Any],
) -> dict[str, Any]:
    requested_revision = review.get("review_mode_revision")
    if (
        not run_id
        or review.get("requested_mode") != "auto"
        or not isinstance(requested_revision, int)
        or requested_revision < 0
        or resolved.get("protocol") != "resolved-review-mode-v1"
        or str(resolved.get("run_id") or "") != run_id
        or str(resolved.get("envelope_digest") or "")
        != str(review.get("envelope_digest") or "")
        or resolved.get("mode") != "auto"
        or resolved.get("revision") != requested_revision
    ):
        raise AutoReviewVerificationError("AUTO_REVIEW_RESPONSE_INVALID")
    return {
        "protocol": "verified-review-mode-state-v1",
        "execution_run_id": run_id,
        "mode": "auto",
        "revision": requested_revision,
    }


def _resolve_sync(run_id: str, review: Mapping[str, Any]) -> dict[str, Any]:
    gateway_url, payload = _request_payload(run_id, review)
    try:
        response = httpx.post(
            f"{gateway_url}/api/internal/recoverable-runs/review-mode/resolve",
            json=payload,
            headers=_service_headers(),
            timeout=httpx.Timeout(10.0, connect=3.0),
        )
        response.raise_for_status()
        resolved = response.json()
    except (httpx.HTTPError, ValueError, TypeError) as exc:
        raise AutoReviewVerificationError("AUTO_REVIEW_VERIFICATION_FAILED") from exc
    if not isinstance(resolved, Mapping):
        raise AutoReviewVerificationError("AUTO_REVIEW_RESPONSE_INVALID")
    return _validated_auto_state(run_id, review, resolved)


async def _resolve_async(run_id: str, review: Mapping[str, Any]) -> dict[str, Any]:
    gateway_url, payload = _request_payload(run_id, review)
    try:
        async with httpx.AsyncClient(
            timeout=httpx.Timeout(10.0, connect=3.0)
        ) as client:
            response = await client.post(
                f"{gateway_url}/api/internal/recoverable-runs/review-mode/resolve",
                json=payload,
                headers=_service_headers(),
            )
        response.raise_for_status()
        resolved = response.json()
    except (httpx.HTTPError, ValueError, TypeError) as exc:
        raise AutoReviewVerificationError("AUTO_REVIEW_VERIFICATION_FAILED") from exc
    if not isinstance(resolved, Mapping):
        raise AutoReviewVerificationError("AUTO_REVIEW_RESPONSE_INVALID")
    return _validated_auto_state(run_id, review, resolved)


class DynamicReviewMiddleware(HumanInTheLoopMiddleware):
    """Use HITL for manual Runs and bypass it only for verified automatic Runs."""

    state_schema = DynamicReviewState

    @staticmethod
    def _abort_requested() -> bool:
        """本轮是否被要求"终止"（网关放弃某个待审批时注入的标记）。

        方案 A：中断仍由库正常消费（该工具不执行），但随后不再发起模型调用，
        直接把本轮跳到 end —— 即用户口径里的"不批准 → 停止这条对话"。
        """

        try:
            config = get_config()
        except RuntimeError:
            return False
        configurable = config.get("configurable") if isinstance(config, Mapping) else None
        if not isinstance(configurable, Mapping):
            return False
        return bool(configurable.get("ai4sci_abort_turn"))

    @classmethod
    def _finish_turn(cls, update: dict[str, Any] | None) -> dict[str, Any] | None:
        """终止本轮：保留库的决议结果（中断被消费掉），再结束本轮。"""

        if not cls._abort_requested():
            return update
        return {**(update or {}), "jump_to": "end"}

    def before_agent(self, state: DynamicReviewState, runtime: Any) -> dict[str, Any]:
        del state, runtime
        run_id, review = _review_context()
        if review is None or review.get("requested_mode") != "auto":
            return {"_verified_review_mode": _manual_state(run_id, review)}
        return {"_verified_review_mode": _resolve_sync(run_id, review)}

    async def abefore_agent(
        self, state: DynamicReviewState, runtime: Any
    ) -> dict[str, Any]:
        del state, runtime
        run_id, review = _review_context()
        if review is None or review.get("requested_mode") != "auto":
            return {"_verified_review_mode": _manual_state(run_id, review)}
        return {"_verified_review_mode": await _resolve_async(run_id, review)}

    @hook_config(can_jump_to=["end"])
    def after_model(
        self, state: DynamicReviewState, runtime: Any
    ) -> dict[str, Any] | None:
        verified = state.get("_verified_review_mode")
        if not isinstance(verified, Mapping):
            raise AutoReviewVerificationError("REVIEW_MODE_STATE_MISSING")
        mode = verified.get("mode")
        current_run_id, review = _review_context()
        if mode == "auto":
            if not current_run_id:
                raise AutoReviewVerificationError("REVIEW_MODE_RUN_MISMATCH")
            if verified.get("execution_run_id") == current_run_id:
                return None
            # A LangGraph resume continues at the interrupted node and does not
            # re-run before_agent, so execution_run_id still names the parent
            # Run that established the verified auto-mode state. Re-verify auto
            # approval against the review context the gateway injected for the
            # resume child.
            if review is None:
                # Legacy resume without injected context: trust the parent's
                # already-verified auto approval for the remainder of the turn.
                return None
            if review.get("requested_mode") != "auto":
                # The gateway now requires manual approval for this turn.
                return self._finish_turn(super().after_model(state, runtime))
            try:
                _resolve_sync(current_run_id, review)
            except AutoReviewVerificationError:
                # Auto approval could not be re-verified; fall back to review.
                return self._finish_turn(super().after_model(state, runtime))
            return None
        if mode == "manual":
            # A LangGraph resume continues at this interrupted node and does not
            # re-run before_agent. Reusing a manual state is restrictive and is
            # required for the existing resume child Run to complete.
            # EXCEPTION: the gateway snapshots the thread's LIVE review mode
            # into each resume child's envelope. When the user flipped the
            # thread (or the current interrupt) to auto AFTER the parent run
            # started, the injected ai4sci_review_mode context is "auto" —
            # verify it and bypass HITL instead of interrupting again.
            if review is not None and review.get("requested_mode") == "auto":
                try:
                    _resolve_sync(current_run_id, review)
                except AutoReviewVerificationError:
                    return self._finish_turn(super().after_model(state, runtime))
                return None
            return self._finish_turn(super().after_model(state, runtime))
        raise AutoReviewVerificationError("REVIEW_MODE_STATE_INVALID")

    @hook_config(can_jump_to=["end"])
    async def aafter_model(
        self, state: DynamicReviewState, runtime: Any
    ) -> dict[str, Any] | None:
        verified = state.get("_verified_review_mode")
        if not isinstance(verified, Mapping):
            raise AutoReviewVerificationError("REVIEW_MODE_STATE_MISSING")
        mode = verified.get("mode")
        current_run_id, review = _review_context()
        if mode == "auto":
            if not current_run_id:
                raise AutoReviewVerificationError("REVIEW_MODE_RUN_MISMATCH")
            if verified.get("execution_run_id") == current_run_id:
                return None
            if review is None:
                return None
            if review.get("requested_mode") != "auto":
                return self._finish_turn(super().after_model(state, runtime))
            try:
                await _resolve_async(current_run_id, review)
            except AutoReviewVerificationError:
                return self._finish_turn(super().after_model(state, runtime))
            return None
        if mode == "manual":
            # Mirror the sync path: an injected auto context on a resume child
            # (thread flipped to auto after the parent started) bypasses HITL
            # after successful re-verification.
            if review is not None and review.get("requested_mode") == "auto":
                try:
                    await _resolve_async(current_run_id, review)
                except AutoReviewVerificationError:
                    return self._finish_turn(super().after_model(state, runtime))
                return None
            return self._finish_turn(super().after_model(state, runtime))
        raise AutoReviewVerificationError("REVIEW_MODE_STATE_INVALID")