deerflow-code/offline-backend-20260512/backend/packages/harness/deerflow/browser_context/tasks.py

"""Small file-backed task queue for browser page fetch requests."""

from __future__ import annotations

import asyncio
import json
import logging
import re
import threading
import time
from datetime import UTC, datetime, timedelta
from pathlib import Path
from typing import Any
from uuid import uuid4

from deerflow.config.runtime_paths import runtime_home

logger = logging.getLogger(__name__)

_LOCK = threading.RLock()
_TASK_ID_RE = re.compile(r"^[a-f0-9]{32}$")
_TERMINAL_STATUSES = {"done", "error", "expired"}


def tasks_dir() -> Path:
    """Return the browser-context task directory, creating it if needed."""

    path = runtime_home() / "browser-context" / "tasks"
    path.mkdir(parents=True, exist_ok=True)
    return path


def _now() -> datetime:
    return datetime.now(UTC)


def _iso(value: datetime | None = None) -> str:
    return (value or _now()).isoformat()


def _parse_iso(value: Any) -> datetime | None:
    if not isinstance(value, str) or not value:
        return None
    try:
        parsed = datetime.fromisoformat(value.replace("Z", "+00:00"))
    except ValueError:
        return None
    if parsed.tzinfo is None:
        return parsed.replace(tzinfo=UTC)
    return parsed.astimezone(UTC)


def _task_path(task_id: str) -> Path | None:
    if not _TASK_ID_RE.fullmatch(task_id or ""):
        return None
    return tasks_dir() / f"{task_id}.json"


def _read_task(path: Path) -> dict[str, Any] | None:
    try:
        with path.open("r", encoding="utf-8") as fh:
            data = json.load(fh)
    except FileNotFoundError:
        return None
    except Exception:
        logger.warning("Ignoring unreadable browser task file: %s", path, exc_info=True)
        return None
    return data if isinstance(data, dict) else None


def _write_task(path: Path, data: dict[str, Any]) -> None:
    tmp = path.with_name(f"{path.name}.{uuid4().hex}.tmp")
    with tmp.open("w", encoding="utf-8") as fh:
        json.dump(data, fh, ensure_ascii=False, separators=(",", ":"))
    tmp.replace(path)


def _public_task(data: dict[str, Any]) -> dict[str, Any]:
    return {
        "id": data.get("id"),
        "kind": data.get("kind") or "fetch",
        "url": data.get("url") or "",
        "instruction": data.get("instruction") or "",
        "session_id": data.get("session_id") or "",
        "action": data.get("action") or None,
        "created_at": data.get("created_at") or "",
        "updated_at": data.get("updated_at") or "",
        "max_content_chars": int(data.get("max_content_chars") or 200_000),
    }


def _expire_if_needed(path: Path, data: dict[str, Any], now: datetime) -> dict[str, Any]:
    if data.get("status") in _TERMINAL_STATUSES:
        return data
    expires_at = _parse_iso(data.get("expires_at"))
    if expires_at is not None and now >= expires_at:
        data = dict(data)
        data["status"] = "expired"
        data["error"] = "Browser task expired before the extension returned a result."
        data["updated_at"] = _iso(now)
        _write_task(path, data)
    return data


def prune_browser_tasks(max_age_seconds: int = 86_400) -> int:
    """Delete old terminal tasks and return the number removed."""

    cutoff = _now() - timedelta(seconds=max(60, int(max_age_seconds)))
    removed = 0
    with _LOCK:
        for path in tasks_dir().glob("*.json"):
            data = _read_task(path)
            if not data or data.get("status") not in _TERMINAL_STATUSES:
                continue
            updated_at = _parse_iso(data.get("updated_at")) or _parse_iso(data.get("created_at"))
            if updated_at is None or updated_at < cutoff:
                try:
                    path.unlink()
                    removed += 1
                except OSError:
                    logger.debug("Failed to delete old browser task: %s", path, exc_info=True)
    return removed


def create_browser_task(
    *,
    user_id: str,
    url: str,
    instruction: str = "",
    timeout_seconds: int = 60,
    max_content_chars: int = 200_000,
    metadata: dict[str, Any] | None = None,
    kind: str = "fetch",
    session_id: str = "",
    action: dict[str, Any] | None = None,
) -> dict[str, Any]:
    """Create a browser task for the current user's extension to pick up.

    ``kind`` is ``"fetch"`` for one-shot page captures (the legacy behavior) or
    ``"act"`` for an interactive observe/act step driven by ``action`` within a
    long-lived ``session_id`` (a tab kept open by the extension).
    """

    now = _now()
    task_id = uuid4().hex
    task = {
        "id": task_id,
        "user_id": str(user_id or "default"),
        "status": "pending",
        "kind": kind if kind in {"fetch", "act"} else "fetch",
        "url": url,
        "instruction": instruction or "",
        "session_id": str(session_id or ""),
        "action": action or None,
        "created_at": _iso(now),
        "updated_at": _iso(now),
        "expires_at": _iso(now + timedelta(seconds=max(5, int(timeout_seconds)))),
        "max_content_chars": max(1_000, int(max_content_chars)),
        "worker_id": None,
        "claimed_at": None,
        "result": None,
        "error": None,
        "metadata": metadata or {},
    }
    path = tasks_dir() / f"{task_id}.json"
    with _LOCK:
        _write_task(path, task)
    return task


def get_browser_task(*, task_id: str, user_id: str) -> dict[str, Any] | None:
    """Read a task if it belongs to the requested user."""

    path = _task_path(task_id)
    if path is None:
        return None
    with _LOCK:
        data = _read_task(path)
        if not data or str(data.get("user_id") or "") != str(user_id or "default"):
            return None
        return _expire_if_needed(path, data, _now())


def claim_next_browser_task(
    *,
    user_id: str,
    worker_id: str,
    lease_seconds: int = 120,
) -> dict[str, Any] | None:
    """Claim the oldest pending task for a user."""

    now = _now()
    user_id = str(user_id or "default")
    worker_id = (worker_id or "chrome-extension").strip()[:128] or "chrome-extension"
    lease_seconds = max(10, int(lease_seconds))

    with _LOCK:
        paths = sorted(tasks_dir().glob("*.json"), key=lambda p: p.stat().st_mtime)
        for path in paths:
            data = _read_task(path)
            if not data or str(data.get("user_id") or "") != user_id:
                continue

            data = _expire_if_needed(path, data, now)
            status = data.get("status")
            if status == "running":
                updated_at = _parse_iso(data.get("updated_at")) or now
                if (now - updated_at).total_seconds() <= lease_seconds:
                    continue
                status = "pending"
                data = dict(data)
                data["status"] = status

            if status != "pending":
                continue

            data = dict(data)
            data["status"] = "running"
            data["worker_id"] = worker_id
            data["claimed_at"] = _iso(now)
            data["updated_at"] = _iso(now)
            _write_task(path, data)
            return _public_task(data)

    return None


def complete_browser_task(
    *,
    task_id: str,
    user_id: str,
    title: str = "",
    url: str = "",
    content: str = "",
    content_type: str = "markdown",
    error: str = "",
    metadata: dict[str, Any] | None = None,
    observation: dict[str, Any] | None = None,
) -> dict[str, Any] | None:
    """Mark a browser task complete or failed.

    For ``act`` tasks the extension returns a structured ``observation`` (final
    URL/title, interactive element list, page text) instead of a captured
    markdown document; it is stored alongside the result verbatim.
    """

    path = _task_path(task_id)
    if path is None:
        return None

    now = _now()
    with _LOCK:
        data = _read_task(path)
        if not data or str(data.get("user_id") or "") != str(user_id or "default"):
            return None

        data = dict(data)
        if error:
            data["status"] = "error"
            data["error"] = str(error).strip()[:4000]
            data["result"] = None
        else:
            raw_content = str(content or "").strip()
            max_content_chars = max(1_000, int(data.get("max_content_chars") or 200_000))
            truncated = len(raw_content) > max_content_chars
            data["status"] = "done"
            data["error"] = None
            data["result"] = {
                "title": str(title or "").strip()[:500],
                "url": str(url or data.get("url") or "").strip()[:4096],
                "content": raw_content[:max_content_chars] if truncated else raw_content,
                "content_type": content_type if content_type in {"markdown", "text"} else "markdown",
                "truncated": truncated,
                "observation": observation if isinstance(observation, dict) else None,
                "captured_at": _iso(now),
                "metadata": metadata or {},
            }
        data["updated_at"] = _iso(now)
        _write_task(path, data)
        return data


async def wait_for_browser_task_result(
    *,
    task_id: str,
    user_id: str,
    timeout_seconds: int,
    poll_interval_seconds: float = 1.0,
) -> dict[str, Any] | None:
    """Wait for a browser task to reach a terminal state."""

    deadline = time.monotonic() + max(1, int(timeout_seconds))
    interval = max(0.25, float(poll_interval_seconds))
    while True:
        task = get_browser_task(task_id=task_id, user_id=user_id)
        if task is None:
            return None
        if task.get("status") in _TERMINAL_STATUSES:
            return task
        if time.monotonic() >= deadline:
            return task
        await asyncio.sleep(interval)