runtime_agent.py

"""
AgentRuntime-backed agent with optional vision executor fallback.

This module intentionally keeps the control plane verification-first:
- Actions may be proposed by either a structured executor (DOM snapshot prompt)
  or a vision executor (screenshot prompt).
- Verification is always executed via AgentRuntime predicates.
"""

from __future__ import annotations

import base64
import inspect
import re
from collections.abc import Callable
from dataclasses import dataclass, field
from typing import Any, Literal

from .agent_runtime import AgentRuntime
from .backends import actions as backend_actions
from .llm_interaction_handler import LLMInteractionHandler
from .llm_provider import LLMProvider
from .models import BBox, Snapshot, StepHookContext
from .verification import Predicate


@dataclass(frozen=True)
class StepVerification:
    predicate: Predicate
    label: str
    required: bool = True
    eventually: bool = True
    timeout_s: float = 10.0
    poll_s: float = 0.25
    max_snapshot_attempts: int = 3
    min_confidence: float | None = None


@dataclass(frozen=True)
class RuntimeStep:
    goal: str
    intent: str | None = None
    verifications: list[StepVerification] = field(default_factory=list)

    # Snapshot quality policy (handled at agent layer; SDK core unchanged).
    snapshot_limit_base: int = 60
    snapshot_limit_step: int = 40
    snapshot_limit_max: int = 220
    max_snapshot_attempts: int = 3
    min_confidence: float | None = None
    min_actionables: int | None = None

    # Vision executor fallback (bounded).
    vision_executor_enabled: bool = True
    max_vision_executor_attempts: int = 1


@dataclass(frozen=True)
class ActOnceResult:
    action: str
    snap: Snapshot
    used_vision: bool


@dataclass(frozen=True)
class PreActionAuthorityDecision:
    allowed: bool
    reason: str | None = None


class RuntimeAgent:
    """
    A thin orchestration layer over AgentRuntime:
    - snapshot (with limit ramp)
    - propose action (structured executor; optionally vision executor fallback)
    - execute action (backend-agnostic primitives)
    - verify (AgentRuntime predicates)
    """

    def __init__(
        self,
        *,
        runtime: AgentRuntime,
        executor: LLMProvider,
        vision_executor: LLMProvider | None = None,
        vision_verifier: LLMProvider | None = None,
        short_circuit_canvas: bool = True,
        pre_action_authorizer: Callable[[Any], Any] | None = None,
        authority_principal_id: str | None = None,
        authority_tenant_id: str | None = None,
        authority_session_id: str | None = None,
        authority_fail_closed: bool = True,
    ) -> None:
        self.runtime = runtime
        self.executor = executor
        self.vision_executor = vision_executor
        self.vision_verifier = vision_verifier
        self.short_circuit_canvas = short_circuit_canvas
        self.pre_action_authorizer = pre_action_authorizer
        self.authority_principal_id = authority_principal_id
        self.authority_tenant_id = authority_tenant_id
        self.authority_session_id = authority_session_id
        self.authority_fail_closed = authority_fail_closed

        self._structured_llm = LLMInteractionHandler(executor)

    async def run_step(
        self,
        *,
        task_goal: str,
        step: RuntimeStep,
        on_step_start: Callable[[StepHookContext], Any] | None = None,
        on_step_end: Callable[[StepHookContext], Any] | None = None,
    ) -> bool:
        step_id = self.runtime.begin_step(step.goal)
        await self._run_hook(
            on_step_start,
            StepHookContext(
                step_id=step_id,
                step_index=self.runtime.step_index,
                goal=step.goal,
                url=getattr(self.runtime.last_snapshot, "url", None),
            ),
        )
        emitted = False
        ok = False
        error_msg: str | None = None
        outcome: str | None = None
        try:
            snap = await self._snapshot_with_ramp(step=step)

            if await self._should_short_circuit_to_vision(step=step, snap=snap):
                ok = await self._vision_executor_attempt(task_goal=task_goal, step=step, snap=snap)
                outcome = "ok" if ok else "verification_failed"
                return ok

            # 1) Structured executor attempt.
            action = self._propose_structured_action(task_goal=task_goal, step=step, snap=snap)
            await self._execute_action(action=action, snap=snap, step_goal=step.goal)
            ok = await self._apply_verifications(step=step)
            if ok:
                outcome = "ok"
                return True

            # 2) Optional vision executor fallback (bounded).
            if step.vision_executor_enabled and step.max_vision_executor_attempts > 0:
                ok = await self._vision_executor_attempt(task_goal=task_goal, step=step, snap=snap)
                outcome = "ok" if ok else "verification_failed"
                return ok

            outcome = "verification_failed"
            return False
        except Exception as exc:
            error_msg = str(exc)
            outcome = "exception"
            try:
                await self.runtime.emit_step_end(
                    success=False,
                    error=str(exc),
                    outcome="exception",
                    verify_passed=False,
                )
                emitted = True
            except Exception:
                pass
            raise
        finally:
            if not emitted:
                try:
                    await self.runtime.emit_step_end(
                        success=ok,
                        outcome=("ok" if ok else "verification_failed"),
                        verify_passed=ok,
                    )
                except Exception:
                    pass
            await self._run_hook(
                on_step_end,
                StepHookContext(
                    step_id=step_id,
                    step_index=self.runtime.step_index,
                    goal=step.goal,
                    url=getattr(self.runtime.last_snapshot, "url", None),
                    success=ok,
                    outcome=outcome,
                    error=error_msg,
                ),
            )

    async def act_once(
        self,
        *,
        task_goal: str,
        step: RuntimeStep,
        allow_vision_fallback: bool = True,
        history_summary: str = "",
        compact_prompt_builder: Callable[
            [str, str, str, Snapshot, str], tuple[str, str]
        ]
        | None = None,
        dom_context_postprocessor: Callable[[str], str] | None = None,
    ) -> str:
        """
        Execute exactly one action for a step without owning step lifecycle.

        This helper is designed for orchestration layers (e.g. WebBench) that already
        call `runtime.begin_step(...)` / `runtime.emit_step_end(...)` and want to
        reuse RuntimeAgent's snapshot-first action proposal + execution logic without:
        - double-counting step budgets
        - emitting duplicate step_start/step_end events

        Returns:
            Action string (e.g. "CLICK(123)", "TYPE(5, \"foo\")", "PRESS(\"Enter\")", "FINISH()")
        """
        res = await self.act_once_result(
            task_goal=task_goal,
            step=step,
            allow_vision_fallback=allow_vision_fallback,
            history_summary=history_summary,
            compact_prompt_builder=compact_prompt_builder,
            dom_context_postprocessor=dom_context_postprocessor,
        )
        return res.action

    async def act_once_with_snapshot(
        self,
        *,
        task_goal: str,
        step: RuntimeStep,
        allow_vision_fallback: bool = True,
        history_summary: str = "",
        compact_prompt_builder: Callable[
            [str, str, str, Snapshot, str], tuple[str, str]
        ]
        | None = None,
        dom_context_postprocessor: Callable[[str], str] | None = None,
    ) -> tuple[str, Snapshot]:
        """
        Like `act_once`, but also returns the pre-action snapshot used for proposal.
        """
        res = await self.act_once_result(
            task_goal=task_goal,
            step=step,
            allow_vision_fallback=allow_vision_fallback,
            history_summary=history_summary,
            compact_prompt_builder=compact_prompt_builder,
            dom_context_postprocessor=dom_context_postprocessor,
        )
        return res.action, res.snap

    async def act_once_result(
        self,
        *,
        task_goal: str,
        step: RuntimeStep,
        allow_vision_fallback: bool = True,
        history_summary: str = "",
        compact_prompt_builder: Callable[
            [str, str, str, Snapshot, str], tuple[str, str]
        ]
        | None = None,
        dom_context_postprocessor: Callable[[str], str] | None = None,
    ) -> ActOnceResult:
        """
        Like `act_once`, but returns action + proposal snapshot + whether vision was used.
        """
        snap = await self._snapshot_with_ramp(step=step)

        # Optional short-circuit to vision (bounded by caller).
        if allow_vision_fallback and await self._should_short_circuit_to_vision(step=step, snap=snap):
            if self.vision_executor and self.vision_executor.supports_vision():
                url = await self._get_url_for_prompt()
                image_b64 = await self._screenshot_base64_png()
                system_prompt, user_prompt = self._vision_executor_prompts(
                    task_goal=task_goal,
                    step=step,
                    url=url,
                    snap=snap,
                )
                resp = self.vision_executor.generate_with_image(
                    system_prompt,
                    user_prompt,
                    image_b64,
                    temperature=0.0,
                )
                action = self._extract_action_from_text(resp.content)
                await self._execute_action(action=action, snap=snap, step_goal=step.goal)
                return ActOnceResult(action=action, snap=snap, used_vision=True)

        # Structured snapshot-first proposal.
        dom_context = self._structured_llm.build_context(snap, step.goal)
        if dom_context_postprocessor is not None:
            dom_context = dom_context_postprocessor(dom_context)

        if compact_prompt_builder is not None:
            system_prompt, user_prompt = compact_prompt_builder(
                task_goal, step.goal, dom_context, snap, history_summary or ""
            )
            resp = self.executor.generate(system_prompt, user_prompt, temperature=0.0)
            action = self._structured_llm.extract_action(resp.content)
        else:
            combined_goal = task_goal
            if history_summary:
                combined_goal = f"{task_goal}\n\nRECENT STEPS:\n{history_summary}"
            combined_goal = f"{combined_goal}\n\nSTEP: {step.goal}"
            resp = self._structured_llm.query_llm(dom_context, combined_goal)
            action = self._structured_llm.extract_action(resp.content)

        await self._execute_action(action=action, snap=snap, step_goal=step.goal)
        return ActOnceResult(action=action, snap=snap, used_vision=False)

    async def _run_hook(
        self,
        hook: Callable[[StepHookContext], Any] | None,
        ctx: StepHookContext,
    ) -> None:
        if hook is None:
            return
        result = hook(ctx)
        if inspect.isawaitable(result):
            await result

    async def _snapshot_with_ramp(self, *, step: RuntimeStep) -> Snapshot:
        limit = step.snapshot_limit_base
        last: Snapshot | None = None

        for _attempt in range(max(1, step.max_snapshot_attempts)):
            last = await self.runtime.snapshot(limit=limit, goal=step.goal)
            if last is None:
                limit = min(step.snapshot_limit_max, limit + step.snapshot_limit_step)
                continue

            if step.min_confidence is not None:
                conf = getattr(getattr(last, "diagnostics", None), "confidence", None)
                if isinstance(conf, (int, float)) and conf < step.min_confidence:
                    limit = min(step.snapshot_limit_max, limit + step.snapshot_limit_step)
                    continue

            if step.min_actionables is not None:
                if self._count_actionables(last) < step.min_actionables:
                    limit = min(step.snapshot_limit_max, limit + step.snapshot_limit_step)
                    continue

            return last

        # If we didn't return early, use last snapshot (may be low quality).
        if last is None:
            raise RuntimeError("snapshot() returned None repeatedly")
        return last

    def _propose_structured_action(
        self, *, task_goal: str, step: RuntimeStep, snap: Snapshot
    ) -> str:
        dom_context = self._structured_llm.build_context(snap, step.goal)
        combined_goal = f"{task_goal}\n\nSTEP: {step.goal}"
        resp = self._structured_llm.query_llm(dom_context, combined_goal)
        return self._structured_llm.extract_action(resp.content)

    async def _vision_executor_attempt(
        self,
        *,
        task_goal: str,
        step: RuntimeStep,
        snap: Snapshot | None,
    ) -> bool:
        if not self.vision_executor or not self.vision_executor.supports_vision():
            return False

        url = await self._get_url_for_prompt()
        image_b64 = await self._screenshot_base64_png()
        system_prompt, user_prompt = self._vision_executor_prompts(
            task_goal=task_goal,
            step=step,
            url=url,
            snap=snap,
        )

        resp = self.vision_executor.generate_with_image(
            system_prompt,
            user_prompt,
            image_b64,
            temperature=0.0,
        )

        action = self._extract_action_from_text(resp.content)
        await self._execute_action(action=action, snap=snap, step_goal=step.goal)
        # Important: vision executor fallback is a *retry* of the same step.
        # Clear prior step assertions so required_assertions_passed reflects the final attempt.
        self.runtime.flush_assertions()
        return await self._apply_verifications(step=step)

    async def _apply_verifications(self, *, step: RuntimeStep) -> bool:
        if not step.verifications:
            # No explicit verifications provided: treat as pass.
            return True

        all_ok = True
        for v in step.verifications:
            if v.eventually:
                ok = await self.runtime.check(
                    v.predicate, label=v.label, required=v.required
                ).eventually(
                    timeout_s=v.timeout_s,
                    poll_s=v.poll_s,
                    max_snapshot_attempts=v.max_snapshot_attempts,
                    min_confidence=v.min_confidence,
                    vision_provider=self.vision_verifier,
                )
            else:
                ok = self.runtime.assert_(v.predicate, label=v.label, required=v.required)
            all_ok = all_ok and ok

        # Respect required verifications semantics.
        return self.runtime.required_assertions_passed() and all_ok

    async def _execute_action(self, *, action: str, snap: Snapshot | None, step_goal: str | None) -> None:
        url = None
        try:
            url = await self.runtime.get_url()
        except Exception:
            url = getattr(snap, "url", None)

        # Coordinate-backed execution (by snapshot id or explicit coordinates).
        kind, payload = self._parse_action(action)

        if kind == "finish":
            await self.runtime.record_action(action, url=url)
            return

        await self._authorize_pre_action_or_raise(
            action=action,
            kind=kind,
            url=url,
            step_goal=step_goal,
        )
        await self.runtime.record_action(action, url=url)

        if kind == "press":
            await self._press_key_best_effort(payload["key"])
            await self._stabilize_best_effort()
            return

        if kind == "click_xy":
            await backend_actions.click(self.runtime.backend, (payload["x"], payload["y"]))
            await self._stabilize_best_effort()
            return

        if kind == "click_rect":
            bbox = BBox(x=payload["x"], y=payload["y"], width=payload["w"], height=payload["h"])
            await backend_actions.click(self.runtime.backend, bbox)
            await self._stabilize_best_effort()
            return

        if snap is None:
            raise RuntimeError("Cannot execute CLICK(id)/TYPE(id, ...) without a snapshot")

        if kind == "click":
            el = self._find_element(snap, payload["id"])
            if el is None:
                raise RuntimeError(f"Element id {payload['id']} not found in snapshot")
            await backend_actions.click(self.runtime.backend, el.bbox)
            await self._stabilize_best_effort()
            return

        if kind == "type":
            el = self._find_element(snap, payload["id"])
            if el is None:
                raise RuntimeError(f"Element id {payload['id']} not found in snapshot")
            await backend_actions.type_text(self.runtime.backend, payload["text"], target=el.bbox)
            await self._stabilize_best_effort()
            return

        raise ValueError(f"Unknown action kind: {kind}")

    async def _authorize_pre_action_or_raise(
        self,
        *,
        action: str,
        kind: str,
        url: str | None,
        step_goal: str | None,
    ) -> None:
        if self.pre_action_authorizer is None:
            return
        principal_id = self.authority_principal_id or "agent:sdk-python"
        action_name = self._authority_action_name(kind)
        resource = url or "about:blank"
        intent = step_goal or action

        try:
            request = self.runtime.build_authority_action_request(
                principal_id=principal_id,
                action=action_name,
                resource=resource,
                intent=intent,
                tenant_id=self.authority_tenant_id,
                session_id=self.authority_session_id,
            )
            decision_raw = self.pre_action_authorizer(request)
            if inspect.isawaitable(decision_raw):
                decision_raw = await decision_raw
            decision = self._normalize_authority_decision(decision_raw)
            if decision.allowed:
                return
            raise RuntimeError(
                f"pre_action_authority_denied: {decision.reason or 'denied_by_authority'}"
            )
        except Exception:
            if self.authority_fail_closed:
                raise
            return

    def _normalize_authority_decision(self, value: Any) -> PreActionAuthorityDecision:
        if isinstance(value, PreActionAuthorityDecision):
            return value
        allowed_attr = getattr(value, "allowed", None)
        if isinstance(allowed_attr, bool):
            reason_attr = getattr(value, "reason", None)
            reason = str(reason_attr) if isinstance(reason_attr, str) and reason_attr else None
            return PreActionAuthorityDecision(allowed=allowed_attr, reason=reason)
        if isinstance(value, bool):
            return PreActionAuthorityDecision(allowed=value)
        raise RuntimeError("invalid_pre_action_authority_decision")

    def _authority_action_name(self, kind: str) -> str:
        if kind in {"click", "click_xy", "click_rect"}:
            return "browser.click"
        if kind == "type":
            return "browser.type"
        if kind == "press":
            return "browser.press"
        return "browser.unknown"

    async def _stabilize_best_effort(self) -> None:
        try:
            await self.runtime.backend.wait_ready_state(state="interactive", timeout_ms=15000)
        except Exception:
            return

    async def _press_key_best_effort(self, key: str) -> None:
        # BrowserBackend does not expose a dedicated keypress primitive; do best-effort JS events.
        key_esc = key.replace("\\", "\\\\").replace("'", "\\'")
        await self.runtime.backend.eval(
            f"""
            (() => {{
              const el = document.activeElement || document.body;
              const down = new KeyboardEvent('keydown', {{key: '{key_esc}', bubbles: true}});
              const up = new KeyboardEvent('keyup', {{key: '{key_esc}', bubbles: true}});
              el.dispatchEvent(down);
              el.dispatchEvent(up);
              return true;
            }})()
            """
        )

    async def _screenshot_base64_png(self) -> str:
        png = await self.runtime.backend.screenshot_png()
        return base64.b64encode(png).decode("utf-8")

    async def _get_url_for_prompt(self) -> str | None:
        try:
            return await self.runtime.get_url()
        except Exception:
            return getattr(self.runtime.last_snapshot, "url", None)

    async def _should_short_circuit_to_vision(
        self, *, step: RuntimeStep, snap: Snapshot | None
    ) -> bool:
        if not (
            step.vision_executor_enabled
            and self.vision_executor
            and self.vision_executor.supports_vision()
        ):
            return False

        if snap is None:
            return True

        if (
            step.min_actionables is not None
            and self._count_actionables(snap) < step.min_actionables
        ):
            if self.short_circuit_canvas:
                try:
                    n_canvas = await self.runtime.backend.eval(
                        "document.querySelectorAll('canvas').length"
                    )
                    if isinstance(n_canvas, (int, float)) and n_canvas > 0:
                        return True
                except Exception:
                    pass

        return False

    def _vision_executor_prompts(
        self,
        *,
        task_goal: str,
        step: RuntimeStep,
        url: str | None,
        snap: Snapshot | None,
    ) -> tuple[str, str]:
        # Include URL as text: screenshots generally don't include browser chrome reliably.
        verify_targets = self._verification_targets_human(step.verifications)

        snapshot_summary = ""
        if snap is not None:
            snapshot_summary = (
                f"\n\nStructured snapshot summary:\n"
                f"- url: {getattr(snap, 'url', None)}\n"
                f"- elements: {len(getattr(snap, 'elements', []) or [])}\n"
            )

        system_prompt = f"""You are a vision-capable web automation executor.

TASK GOAL:
{task_goal}

STEP GOAL:
{step.goal}

CURRENT URL (text):
{url or "(unknown)"}

VERIFICATION TARGETS (text):
{verify_targets or "(none provided)"}
{snapshot_summary}

RESPONSE FORMAT:
Return ONLY ONE of:
- CLICK(id)
- TYPE(id, "text")
- CLICK_XY(x, y)
- CLICK_RECT(x, y, w, h)
- PRESS("key")
- FINISH()

No explanations, no markdown.
"""

        user_prompt = "From the screenshot, return the single best next action:"
        return system_prompt, user_prompt

    def _verification_targets_human(self, verifications: list[StepVerification]) -> str:
        if not verifications:
            return ""
        lines: list[str] = []
        for v in verifications:
            req = "required" if v.required else "optional"
            lines.append(f"- {v.label} ({req})")
        return "\n".join(lines)

    def _count_actionables(self, snap: Snapshot) -> int:
        n = 0
        for el in snap.elements or []:
            cues = getattr(el, "visual_cues", None)
            clickable = bool(getattr(cues, "is_clickable", False))
            if clickable:
                n += 1
        return n

    def _find_element(self, snap: Snapshot, element_id: int) -> Any | None:
        for el in snap.elements or []:
            if getattr(el, "id", None) == element_id:
                return el
        return None

    def _parse_action(
        self,
        action: str,
    ) -> tuple[
        Literal["click", "type", "press", "finish", "click_xy", "click_rect"], dict[str, Any]
    ]:
        action = action.strip()

        if re.match(r"FINISH\s*\(\s*\)\s*$", action, re.IGNORECASE):
            return "finish", {}

        if m := re.match(
            r"CLICK_XY\s*\(\s*(-?\d+(?:\.\d+)?)\s*,\s*(-?\d+(?:\.\d+)?)\s*\)\s*$",
            action,
            re.IGNORECASE,
        ):
            return "click_xy", {"x": float(m.group(1)), "y": float(m.group(2))}

        if m := re.match(
            r"CLICK_RECT\s*\(\s*(-?\d+(?:\.\d+)?)\s*,\s*(-?\d+(?:\.\d+)?)\s*,\s*(-?\d+(?:\.\d+)?)\s*,\s*(-?\d+(?:\.\d+)?)\s*\)\s*$",
            action,
            re.IGNORECASE,
        ):
            return "click_rect", {
                "x": float(m.group(1)),
                "y": float(m.group(2)),
                "w": float(m.group(3)),
                "h": float(m.group(4)),
            }

        if m := re.match(r"CLICK\s*\(\s*(\d+)\s*\)\s*$", action, re.IGNORECASE):
            return "click", {"id": int(m.group(1))}

        if m := re.match(
            r'TYPE\s*\(\s*(\d+)\s*,\s*["\']([^"\']*)["\']\s*\)\s*$',
            action,
            re.IGNORECASE,
        ):
            return "type", {"id": int(m.group(1)), "text": m.group(2)}

        if m := re.match(r'PRESS\s*\(\s*["\']([^"\']+)["\']\s*\)\s*$', action, re.IGNORECASE):
            return "press", {"key": m.group(1)}

        raise ValueError(f"Unknown action format: {action}")

    def _extract_action_from_text(self, text: str) -> str:
        # Keep consistent with LLMInteractionHandler.extract_action, but without DOM context dependency.
        text = re.sub(r"```[\w]*\n?", "", text).strip()
        pat = r'(CLICK_XY\s*\(\s*-?\d+(?:\.\d+)?\s*,\s*-?\d+(?:\.\d+)?\s*\)|CLICK_RECT\s*\(\s*-?\d+(?:\.\d+)?\s*,\s*-?\d+(?:\.\d+)?\s*,\s*-?\d+(?:\.\d+)?\s*,\s*-?\d+(?:\.\d+)?\s*\)|CLICK\s*\(\s*\d+\s*\)|TYPE\s*\(\s*\d+\s*,\s*["\'].*?["\']\s*\)|PRESS\s*\(\s*["\'].*?["\']\s*\)|FINISH\s*\(\s*\))'
        m = re.search(pat, text, re.IGNORECASE)
        return m.group(1) if m else text