333 lines
11 KiB
Python
333 lines
11 KiB
Python
"""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))
|
|
queue_dir = tasks_dir()
|
|
|
|
with _LOCK:
|
|
paths = sorted(queue_dir.glob("*.json"), key=lambda p: p.stat().st_mtime)
|
|
matching_user = 0
|
|
matching_status_counts: dict[str, int] = {}
|
|
for path in paths:
|
|
data = _read_task(path)
|
|
if not data or str(data.get("user_id") or "") != user_id:
|
|
continue
|
|
matching_user += 1
|
|
|
|
data = _expire_if_needed(path, data, now)
|
|
status = data.get("status")
|
|
status_key = str(status or "unknown")
|
|
matching_status_counts[status_key] = matching_status_counts.get(status_key, 0) + 1
|
|
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)
|
|
logger.info(
|
|
"Browser task claim scan user_id=%s worker_id=%s tasks_dir=%s total_files=%d matching_user=%d "
|
|
"matching_statuses=%s claimed_task_id=%s claimed_kind=%s",
|
|
user_id,
|
|
worker_id,
|
|
queue_dir,
|
|
len(paths),
|
|
matching_user,
|
|
matching_status_counts,
|
|
data.get("id"),
|
|
data.get("kind") or "fetch",
|
|
)
|
|
return _public_task(data)
|
|
|
|
logger.info(
|
|
"Browser task claim scan user_id=%s worker_id=%s tasks_dir=%s total_files=%d matching_user=%d "
|
|
"matching_statuses=%s claimed_task_id=%s claimed_kind=%s",
|
|
user_id,
|
|
worker_id,
|
|
queue_dir,
|
|
len(paths),
|
|
matching_user,
|
|
matching_status_counts,
|
|
None,
|
|
None,
|
|
)
|
|
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)
|