Spaces:
Sleeping
Sleeping
File size: 3,404 Bytes
0f336cf | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 | from datetime import datetime, timezone
from typing import Any, Dict, List, Optional
import uuid
from pydantic import BaseModel, Field
class BrowserInstance(BaseModel):
id: str = Field(default_factory=lambda: f"browser-{uuid.uuid4().hex[:8]}")
status: str = "WARM"
active_mission: Optional[str] = None
tabs_count: int = 1
tabs: List[Dict[str, Any]] = Field(default_factory=lambda: [{"tab_id": "tab-1", "url": "about:blank", "title": "Blank"}])
cookies: List[Dict[str, Any]] = Field(default_factory=list)
history: List[str] = Field(default_factory=list)
snapshots: List[Dict[str, Any]] = Field(default_factory=list)
last_screenshot_base64: Optional[str] = None
downloads: List[Dict[str, Any]] = Field(default_factory=list)
health_score: float = 100.0
memory_mb: float = 128.5
created_at: str = Field(default_factory=lambda: datetime.now(timezone.utc).isoformat())
class BrowserPoolManager:
"""Manages active browser instances, pool scaling, health monitoring, and session recycling."""
def __init__(self, max_pool_size: int = 5):
self.max_pool_size = max_pool_size
self.pool: Dict[str, BrowserInstance] = {}
self._initialize_pool()
def _initialize_pool(self):
for _ in range(min(2, self.max_pool_size)):
inst = BrowserInstance()
self.pool[inst.id] = inst
def get_status(self) -> Dict[str, Any]:
return {
"total_browsers": len(self.pool),
"max_pool_size": self.max_pool_size,
"active_browsers": len([b for b in self.pool.values() if b.active_mission]),
"browsers": list(self.pool.values()),
}
def acquire_instance(self, mission_id: str) -> BrowserInstance:
for inst in self.pool.values():
if not inst.active_mission:
inst.active_mission = mission_id
inst.status = "BUSY"
return inst
if len(self.pool) < self.max_pool_size:
new_inst = BrowserInstance(status="BUSY", active_mission=mission_id)
self.pool[new_inst.id] = new_inst
return new_inst
first_inst = list(self.pool.values())[0]
return first_inst
def release_instance(self, browser_id: str):
if browser_id in self.pool:
self.pool[browser_id].active_mission = None
self.pool[browser_id].status = "WARM"
def update_screenshot(self, browser_id: str, screenshot_base64: str, url: str):
if browser_id in self.pool:
inst = self.pool[browser_id]
inst.last_screenshot_base64 = screenshot_base64
if url not in inst.history:
inst.history.append(url)
def get_mission_screenshot(self, mission_id: str) -> Optional[str]:
for inst in self.pool.values():
if inst.active_mission == mission_id and inst.last_screenshot_base64:
return inst.last_screenshot_base64
return None
def get_all_thumbnails(self) -> List[Dict[str, Any]]:
return [
{
"browser_id": b.id,
"mission_id": b.active_mission,
"status": b.status,
"url": b.tabs[0]["url"] if b.tabs else "about:blank",
"has_screenshot": b.last_screenshot_base64 is not None,
}
for b in self.pool.values()
]
|