Jofthomas HF Staff commited on
Commit
f0266b8
·
verified ·
1 Parent(s): 60ec1bc

Update main.py

Browse files
Files changed (1) hide show
  1. main.py +33 -32
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():