Spaces:
Running
Running
Update main.py
Browse files
main.py
CHANGED
@@ -54,38 +54,6 @@ background_task_handle: Optional[asyncio.Task] = None
|
|
54 |
# --- Create FastAPI app ---
|
55 |
app = FastAPI(title="Pokemon Battle Livestream")
|
56 |
|
57 |
-
# --- WebSocket connection manager ---
|
58 |
-
class ConnectionManager:
|
59 |
-
def __init__(self):
|
60 |
-
self.active_connections: Set[WebSocket] = set()
|
61 |
-
self.current_html: str = create_idle_html("Initializing...", "Setting up Pokemon Battle Stream")
|
62 |
-
|
63 |
-
async def connect(self, websocket: WebSocket):
|
64 |
-
await websocket.accept()
|
65 |
-
self.active_connections.add(websocket)
|
66 |
-
# Send current state to newly connected client
|
67 |
-
await websocket.send_text(self.current_html)
|
68 |
-
|
69 |
-
def disconnect(self, websocket: WebSocket):
|
70 |
-
self.active_connections.remove(websocket)
|
71 |
-
|
72 |
-
async def update_all(self, html: str):
|
73 |
-
"""Update the current HTML and broadcast to all clients"""
|
74 |
-
self.current_html = html
|
75 |
-
if self.active_connections:
|
76 |
-
# Only log if there are connections to update
|
77 |
-
print(f"Broadcasting update to {len(self.active_connections)} clients")
|
78 |
-
|
79 |
-
# Make a copy of the connections set to avoid modification during iteration
|
80 |
-
connections_copy = self.active_connections.copy()
|
81 |
-
for connection in connections_copy:
|
82 |
-
try:
|
83 |
-
await connection.send_text(html)
|
84 |
-
except Exception as e:
|
85 |
-
print(f"Error sending to client: {e}")
|
86 |
-
# Don't remove here - will be handled by disconnect route
|
87 |
-
|
88 |
-
manager = ConnectionManager()
|
89 |
|
90 |
# --- Helper Functions ---
|
91 |
def get_active_battle(agent: Player) -> Optional[Battle]:
|
@@ -515,6 +483,39 @@ def log_task_exception(task: asyncio.Task):
|
|
515 |
print(f"Exception in background task {task.get_name()}: {e}")
|
516 |
traceback.print_exc()
|
517 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
518 |
# --- API Routes ---
|
519 |
@app.get("/", response_class=HTMLResponse)
|
520 |
async def get_homepage():
|
|
|
54 |
# --- Create FastAPI app ---
|
55 |
app = FastAPI(title="Pokemon Battle Livestream")
|
56 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
57 |
|
58 |
# --- Helper Functions ---
|
59 |
def get_active_battle(agent: Player) -> Optional[Battle]:
|
|
|
483 |
print(f"Exception in background task {task.get_name()}: {e}")
|
484 |
traceback.print_exc()
|
485 |
|
486 |
+
|
487 |
+
# --- WebSocket connection manager ---
|
488 |
+
class ConnectionManager:
|
489 |
+
def __init__(self):
|
490 |
+
self.active_connections: Set[WebSocket] = set()
|
491 |
+
self.current_html: str = create_idle_html("Initializing...", "Setting up Pokemon Battle Stream")
|
492 |
+
|
493 |
+
async def connect(self, websocket: WebSocket):
|
494 |
+
await websocket.accept()
|
495 |
+
self.active_connections.add(websocket)
|
496 |
+
# Send current state to newly connected client
|
497 |
+
await websocket.send_text(self.current_html)
|
498 |
+
|
499 |
+
def disconnect(self, websocket: WebSocket):
|
500 |
+
self.active_connections.remove(websocket)
|
501 |
+
|
502 |
+
async def update_all(self, html: str):
|
503 |
+
"""Update the current HTML and broadcast to all clients"""
|
504 |
+
self.current_html = html
|
505 |
+
if self.active_connections:
|
506 |
+
# Only log if there are connections to update
|
507 |
+
print(f"Broadcasting update to {len(self.active_connections)} clients")
|
508 |
+
|
509 |
+
# Make a copy of the connections set to avoid modification during iteration
|
510 |
+
connections_copy = self.active_connections.copy()
|
511 |
+
for connection in connections_copy:
|
512 |
+
try:
|
513 |
+
await connection.send_text(html)
|
514 |
+
except Exception as e:
|
515 |
+
print(f"Error sending to client: {e}")
|
516 |
+
# Don't remove here - will be handled by disconnect route
|
517 |
+
|
518 |
+
manager = ConnectionManager()
|
519 |
# --- API Routes ---
|
520 |
@app.get("/", response_class=HTMLResponse)
|
521 |
async def get_homepage():
|