302 lines
12 KiB
Python
302 lines
12 KiB
Python
import asyncio
|
|
import sys
|
|
import time
|
|
from typing import Dict, List, Optional, Set, Any
|
|
from fastapi import WebSocket
|
|
|
|
from crawler.config import CrawlerConfig
|
|
from crawler.engine import Crawler
|
|
|
|
|
|
class CrawlSessionManager:
|
|
"""Manages active crawl sessions, WebSocket broadcasting, metrics, and error aggregation."""
|
|
|
|
def __init__(self):
|
|
self.active_connections: Set[WebSocket] = set()
|
|
self.crawler: Optional[Crawler] = None
|
|
self.task: Optional[asyncio.Task] = None
|
|
self.status: str = "idle" # "idle" | "running" | "stopping" | "finished" | "error"
|
|
self.status_message: str = "Ready to start"
|
|
self.start_time: Optional[float] = None
|
|
self.end_time: Optional[float] = None
|
|
self.current_iteration: int = 1
|
|
self.total_iterations: int = 1
|
|
self.iteration_mode: str = "all"
|
|
self.max_visits_per_iteration: int = 10
|
|
|
|
self.stats: Dict[str, int] = {
|
|
"internal": 0,
|
|
"external": 0,
|
|
"cookies_accepted": 0,
|
|
}
|
|
self.queue_size: int = 0
|
|
self.visited_count: int = 0
|
|
self.total_visited_all: int = 0
|
|
self.total_errors_all: int = 0
|
|
self.tor_identity: Optional[Dict[str, Any]] = None
|
|
|
|
self.errors: List[Dict[str, Any]] = []
|
|
self.logs: List[Dict[str, Any]] = []
|
|
self._log_counter: int = 0
|
|
self._error_counter: int = 0
|
|
|
|
async def connect_ws(self, websocket: WebSocket):
|
|
"""Registers a new WebSocket connection and transmits initial state snapshot."""
|
|
await websocket.accept()
|
|
self.active_connections.add(websocket)
|
|
# Send current state snapshot immediately upon connection
|
|
snapshot = self.get_snapshot()
|
|
await websocket.send_json({
|
|
"type": "init",
|
|
"data": snapshot,
|
|
})
|
|
|
|
def disconnect_ws(self, websocket: WebSocket):
|
|
"""Removes a disconnected WebSocket."""
|
|
self.active_connections.discard(websocket)
|
|
|
|
async def broadcast(self, message: Dict[str, Any]):
|
|
"""Dispatches an event payload to all connected WebSockets."""
|
|
disconnected = set()
|
|
for ws in self.active_connections:
|
|
try:
|
|
await ws.send_json(message)
|
|
except Exception:
|
|
disconnected.add(ws)
|
|
|
|
for ws in disconnected:
|
|
self.active_connections.discard(ws)
|
|
|
|
def get_snapshot(self) -> Dict[str, Any]:
|
|
"""Generates a complete snapshot of current crawler session."""
|
|
elapsed = 0.0
|
|
if self.start_time:
|
|
if self.end_time:
|
|
elapsed = self.end_time - self.start_time
|
|
else:
|
|
elapsed = time.time() - self.start_time
|
|
|
|
return {
|
|
"status": self.status,
|
|
"status_message": self.status_message,
|
|
"elapsed_seconds": round(elapsed, 1),
|
|
"current_iteration": self.current_iteration,
|
|
"total_iterations": self.total_iterations,
|
|
"iteration_mode": self.iteration_mode,
|
|
"max_visits_per_iteration": self.max_visits_per_iteration,
|
|
"stats": self.stats,
|
|
"queue_size": self.queue_size,
|
|
"visited_count": self.visited_count,
|
|
"total_visited_all": self.total_visited_all,
|
|
"total_errors_all": max(self.total_errors_all, len(self.errors)),
|
|
"tor_identity": self.tor_identity,
|
|
"error_count": len(self.errors),
|
|
"errors": self.errors[-100:], # last 100 errors for quick load
|
|
"logs": self.logs[-500:], # last 500 logs for quick load
|
|
}
|
|
|
|
async def handle_crawler_event(self, event: Dict[str, Any]):
|
|
"""Processes events dispatched by the core crawler engine."""
|
|
event_type = event.get("type")
|
|
|
|
if event_type == "log":
|
|
self._log_counter += 1
|
|
log_entry = {
|
|
"id": self._log_counter,
|
|
"message": event.get("message", ""),
|
|
"level": event.get("level", "info"),
|
|
"timestamp": event.get("timestamp", time.time()),
|
|
}
|
|
self.logs.append(log_entry)
|
|
if len(self.logs) > 2000:
|
|
self.logs = self.logs[-1500:]
|
|
await self.broadcast({"type": "log", "data": log_entry})
|
|
|
|
elif event_type == "stats":
|
|
self.stats = event.get("stats", self.stats)
|
|
self.queue_size = event.get("queue_size", self.queue_size)
|
|
self.visited_count = event.get("visited_count", self.visited_count)
|
|
if "total_visited_all" in event:
|
|
self.total_visited_all = event["total_visited_all"]
|
|
else:
|
|
self.total_visited_all = max(self.total_visited_all, self.visited_count)
|
|
|
|
if "total_errors_all" in event:
|
|
self.total_errors_all = event["total_errors_all"]
|
|
else:
|
|
self.total_errors_all = len(self.errors)
|
|
|
|
await self.broadcast({
|
|
"type": "stats",
|
|
"data": {
|
|
"stats": self.stats,
|
|
"queue_size": self.queue_size,
|
|
"visited_count": self.visited_count,
|
|
"error_count": event.get("error_count", 0),
|
|
"total_visited_all": self.total_visited_all,
|
|
"total_errors_all": self.total_errors_all,
|
|
},
|
|
})
|
|
|
|
elif event_type == "tor_identity":
|
|
self.tor_identity = event.get("data", {})
|
|
await self.broadcast({
|
|
"type": "tor_identity",
|
|
"data": self.tor_identity,
|
|
})
|
|
|
|
elif event_type == "error":
|
|
self._error_counter += 1
|
|
error_entry = {
|
|
"id": self._error_counter,
|
|
"url": event.get("url", ""),
|
|
"source_url": event.get("source_url") or "Root / Direct Input",
|
|
"error": event.get("error", "Unknown error"),
|
|
"timestamp": event.get("timestamp", time.time()),
|
|
}
|
|
self.errors.append(error_entry)
|
|
self.total_errors_all = len(self.errors)
|
|
await self.broadcast({
|
|
"type": "error",
|
|
"data": {
|
|
**error_entry,
|
|
"total_errors_all": self.total_errors_all,
|
|
},
|
|
})
|
|
|
|
elif event_type == "iteration":
|
|
self.current_iteration = event.get("iteration", 1)
|
|
self.total_iterations = event.get("total", 1)
|
|
if "iteration_mode" in event:
|
|
self.iteration_mode = event["iteration_mode"]
|
|
if "max_visits_per_iteration" in event:
|
|
self.max_visits_per_iteration = event["max_visits_per_iteration"]
|
|
if "total_visited_all" in event:
|
|
self.total_visited_all = event["total_visited_all"]
|
|
if "total_errors_all" in event:
|
|
self.total_errors_all = event["total_errors_all"]
|
|
await self.broadcast({
|
|
"type": "iteration",
|
|
"data": {
|
|
"current": self.current_iteration,
|
|
"total": self.total_iterations,
|
|
"iteration_mode": self.iteration_mode,
|
|
"max_visits_per_iteration": self.max_visits_per_iteration,
|
|
"total_visited_all": self.total_visited_all,
|
|
"total_errors_all": self.total_errors_all,
|
|
},
|
|
})
|
|
|
|
elif event_type == "status":
|
|
self.status = event.get("status", self.status)
|
|
self.status_message = event.get("message", "")
|
|
if "total_visited_all" in event:
|
|
self.total_visited_all = event["total_visited_all"]
|
|
if "total_errors_all" in event:
|
|
self.total_errors_all = event["total_errors_all"]
|
|
if self.status in ["finished", "stopped", "error"]:
|
|
self.end_time = time.time()
|
|
|
|
await self.broadcast({
|
|
"type": "status",
|
|
"data": {
|
|
"status": self.status,
|
|
"message": self.status_message,
|
|
"elapsed": round((self.end_time or time.time()) - (self.start_time or time.time()), 1),
|
|
"total_visited_all": self.total_visited_all,
|
|
"total_errors_all": self.total_errors_all,
|
|
},
|
|
})
|
|
|
|
async def start_crawl(self, config_data: Dict[str, Any]) -> Dict[str, Any]:
|
|
"""Configures and initiates a new crawl in a background task."""
|
|
if self.status == "running":
|
|
return {"success": False, "message": "Crawler is already running."}
|
|
|
|
# Safeguard: Verify Windows event loop supports subprocesses for Playwright
|
|
if sys.platform == "win32":
|
|
loop = asyncio.get_running_loop()
|
|
if not isinstance(loop, getattr(asyncio, "ProactorEventLoop", ())):
|
|
err_msg = (
|
|
"Current asyncio event loop does not support subprocesses on Windows "
|
|
"(ProactorEventLoop is required for Playwright). "
|
|
"Please launch Websurfer using 'python web.py' or ensure Uvicorn uses '--loop asyncio:ProactorEventLoop'."
|
|
)
|
|
return {"success": False, "message": err_msg}
|
|
|
|
# Clear session data
|
|
self.stats = {"internal": 0, "external": 0, "cookies_accepted": 0}
|
|
self.queue_size = 0
|
|
self.visited_count = 0
|
|
self.total_visited_all = 0
|
|
self.total_errors_all = 0
|
|
self.tor_identity = None
|
|
|
|
self.errors = []
|
|
self.logs = []
|
|
self._log_counter = 0
|
|
self._error_counter = 0
|
|
self.start_time = time.time()
|
|
self.end_time = None
|
|
self.status = "running"
|
|
self.status_message = "Initializing crawler..."
|
|
self.current_iteration = 1
|
|
self.total_iterations = int(config_data.get("num_iterations", 1))
|
|
|
|
# Build configuration object
|
|
config = CrawlerConfig.from_dict(config_data)
|
|
self.iteration_mode = config.iteration_mode
|
|
self.max_visits_per_iteration = config.max_visits_per_iteration
|
|
|
|
# Instantiate crawler with event callback
|
|
self.crawler = Crawler(config, event_callback=self.handle_crawler_event)
|
|
|
|
async def _run_task():
|
|
try:
|
|
await self.crawler.run()
|
|
except Exception as e:
|
|
await self.handle_crawler_event({
|
|
"type": "status",
|
|
"status": "error",
|
|
"message": f"Execution error: {e}",
|
|
})
|
|
finally:
|
|
if self.status == "running":
|
|
self.status = "finished"
|
|
self.end_time = time.time()
|
|
await self.broadcast({
|
|
"type": "status",
|
|
"data": {
|
|
"status": self.status,
|
|
"message": "Completed",
|
|
"elapsed": round(self.end_time - self.start_time, 1),
|
|
},
|
|
})
|
|
|
|
self.task = asyncio.create_task(_run_task())
|
|
await self.broadcast({"type": "init", "data": self.get_snapshot()})
|
|
return {"success": True, "message": "Crawl started"}
|
|
|
|
async def stop_crawl(self) -> Dict[str, Any]:
|
|
"""Stops the active crawler job."""
|
|
if self.status != "running" or not self.crawler:
|
|
return {"success": False, "message": "No active crawl to stop."}
|
|
|
|
self.status = "stopping"
|
|
self.status_message = "Cancelling workers..."
|
|
await self.broadcast({
|
|
"type": "status",
|
|
"data": {"status": "stopping", "message": "Cancelling workers..."},
|
|
})
|
|
|
|
self.crawler.stop()
|
|
return {"success": True, "message": "Stop signal sent."}
|
|
|
|
def clear_logs(self):
|
|
"""Clears in-memory logs."""
|
|
self.logs = []
|
|
self._log_counter = 0
|
|
|
|
|
|
# Global singleton manager instance
|
|
session_manager = CrawlSessionManager()
|