spark_colony / browser /browser_pool.py
diwash-barla1's picture
refactor: decompose app into modular domain packages for v2.5
0f336cf
Raw
History Blame Contribute Delete
3.4 kB
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()
]