Jofthomas commited on
Commit
4cbc162
Β·
verified Β·
1 Parent(s): 42e69ad

Update main.py

Browse files
Files changed (1) hide show
  1. main.py +563 -359
main.py CHANGED
@@ -54,181 +54,73 @@ background_task_handle: Optional[asyncio.Task] = None
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]:
60
  """Returns the first non-finished battle for an agent."""
61
  if agent and agent._battles:
62
  active_battles = [b for b in agent._battles.values() if not b.finished]
63
  if active_battles:
64
- if active_battles[0].battle_tag:
65
- return active_battles[0]
 
 
 
 
 
 
 
66
  else:
67
- print(f"WARN: Found active battle for {agent.username} but it has no battle_tag yet.")
68
  return None
69
  return None
70
 
71
  def create_battle_iframe(battle_id: str) -> str:
72
- """Creates the HTML for the battle iframe."""
73
- print("CURRENT BATTLE ID: ", battle_id)
74
- battle_url = f"https://jofthomas.com/play.pokemonshowdown.com/testclient.html#{battle_id}"
75
-
 
 
 
76
  return f"""
77
- <html>
78
- <head>
79
- <title>Pokemon Battle Livestream</title>
80
- <style>
81
- body, html {{
82
- margin: 0;
83
- padding: 0;
84
- height: 100%;
85
- overflow: hidden;
86
- }}
87
- iframe {{
88
- position: fixed;
89
- top: 0;
90
- left: 0;
91
- bottom: 0;
92
- right: 0;
93
- width: 100%;
94
- height: 100%;
95
- border: none;
96
- margin: 0;
97
- padding: 0;
98
- overflow: hidden;
99
- z-index: 999999;
100
- }}
101
- </style>
102
- </head>
103
- <body>
104
- <iframe
105
- src="{battle_url}"
106
- allowfullscreen
107
- ></iframe>
108
- </body>
109
- </html>
110
  """
111
 
112
  def create_idle_html(status_message: str, instruction: str) -> str:
113
- """Creates a visually appealing idle screen HTML."""
 
114
  return f"""
115
- <html>
116
- <head>
117
- <title>Pokemon Battle Livestream</title>
118
- <style>
119
- body, html {{
120
- margin: 0;
121
- padding: 0;
122
- height: 100%;
123
- overflow: hidden;
124
- font-family: Arial, sans-serif;
125
- }}
126
- .container {{
127
- position: fixed;
128
- top: 0;
129
- left: 0;
130
- width: 100%;
131
- height: 100%;
132
- display: flex;
133
- flex-direction: column;
134
- justify-content: center;
135
- align-items: center;
136
- background-image: url('/static/pokemon_huggingface.png');
137
- background-size: cover;
138
- background-position: center;
139
- z-index: 999999;
140
- }}
141
- .message-box {{
142
- background-color: rgba(0, 0, 0, 0.65);
143
- padding: 30px;
144
- border-radius: 15px;
145
- max-width: 80%;
146
- text-align: center;
147
- }}
148
- .status {{
149
- font-size: 2.5em;
150
- margin-bottom: 20px;
151
- text-shadow: 2px 2px 4px rgba(0, 0, 0, 0.8);
152
- color: white;
153
- }}
154
- .instruction {{
155
- font-size: 1.5em;
156
- text-shadow: 1px 1px 3px rgba(0, 0, 0, 0.8);
157
- color: white;
158
- }}
159
- </style>
160
- </head>
161
- <body>
162
- <div class="container">
163
- <div class="message-box">
164
- <p class="status">{status_message}</p>
165
- <p class="instruction">{instruction}</p>
166
- </div>
167
  </div>
168
- </body>
169
- </html>
170
  """
171
 
172
  def create_error_html(error_msg: str) -> str:
173
- """Creates HTML to display an error message."""
 
174
  return f"""
175
- <html>
176
- <head>
177
- <title>Pokemon Battle Livestream - Error</title>
178
- <style>
179
- body, html {{
180
- margin: 0;
181
- padding: 0;
182
- height: 100%;
183
- overflow: hidden;
184
- font-family: Arial, sans-serif;
185
- }}
186
- .container {{
187
- position: fixed;
188
- top: 0;
189
- left: 0;
190
- width: 100%;
191
- height: 100%;
192
- display: flex;
193
- flex-direction: column;
194
- justify-content: center;
195
- align-items: center;
196
- background-color: #330000;
197
- z-index: 999999;
198
- }}
199
- .message-box {{
200
- background-color: rgba(200, 0, 0, 0.7);
201
- padding: 30px;
202
- border-radius: 15px;
203
- max-width: 80%;
204
- text-align: center;
205
- }}
206
- .title {{
207
- font-size: 2em;
208
- margin-bottom: 20px;
209
- color: white;
210
- }}
211
- .message {{
212
- font-size: 1.2em;
213
- color: #ffdddd;
214
- }}
215
- </style>
216
- </head>
217
- <body>
218
- <div class="container">
219
- <div class="message-box">
220
- <p class="title">An Error Occurred</p>
221
- <p class="message">{error_msg}</p>
222
- </div>
223
  </div>
224
- </body>
225
- </html>
226
  """
227
 
228
- async def update_display_html(new_html: str) -> None:
229
- """Updates the current display HTML and broadcasts to all clients."""
230
- await manager.update_all(new_html)
231
- print("HTML Display UPDATED and broadcasted.")
 
 
232
 
233
  # --- Agent Lifecycle Management ---
234
  async def select_and_activate_new_agent():
@@ -237,7 +129,7 @@ async def select_and_activate_new_agent():
237
 
238
  if not AVAILABLE_AGENT_NAMES:
239
  print("Lifecycle: No available agents with passwords set.")
240
- await update_display_html(create_error_html("No agents available (check password env vars)."))
241
  return False
242
 
243
  selected_name = random.choice(AVAILABLE_AGENT_NAMES)
@@ -247,7 +139,8 @@ async def select_and_activate_new_agent():
247
  agent_password = os.environ.get(password_env_var)
248
 
249
  print(f"Lifecycle: Activating agent '{selected_name}'...")
250
- await update_display_html(create_idle_html("Selecting Next Agent...", f"Preparing {selected_name}..."))
 
251
 
252
  try:
253
  account_config = AccountConfiguration(selected_name, agent_password)
@@ -260,16 +153,18 @@ async def select_and_activate_new_agent():
260
  )
261
 
262
  # Start the task to accept exactly one battle challenge
263
- task = asyncio.create_task(agent.accept_challenges(None, 1), name=f"accept_challenge_{selected_name}")
264
- task.add_done_callback(log_task_exception)
 
265
 
266
  # Update global state
267
  active_agent_name = selected_name
268
  active_agent_instance = agent
269
  active_agent_task = task
270
-
271
  print(f"Lifecycle: Agent '{selected_name}' is active and listening for 1 challenge.")
272
- await update_display_html(create_idle_html(f"Agent <strong>{selected_name}</strong> is ready!",
 
273
  f"Please challenge <strong>{selected_name}</strong> to a <strong>{DEFAULT_BATTLE_FORMAT}</strong> battle."))
274
  return True
275
 
@@ -277,7 +172,8 @@ async def select_and_activate_new_agent():
277
  error_msg = f"Failed to activate agent '{selected_name}': {e}"
278
  print(error_msg)
279
  traceback.print_exc()
280
-
 
281
  # Clear state if activation failed
282
  active_agent_name = None
283
  active_agent_instance = None
@@ -285,74 +181,103 @@ async def select_and_activate_new_agent():
285
  return False
286
 
287
  async def check_for_new_battle():
288
- """Checks if the active agent has started a battle."""
289
  global active_agent_instance, current_battle_instance
290
-
291
  if active_agent_instance:
292
  battle = get_active_battle(active_agent_instance)
 
293
  if battle and battle.battle_tag:
294
  print(f"Lifecycle: Agent '{active_agent_name}' started battle: {battle.battle_tag}")
295
  current_battle_instance = battle
296
-
297
- # Prevent accepting more challenges
298
  if active_agent_task and not active_agent_task.done():
299
  print(f"Lifecycle: Cancelling accept_challenges task for {active_agent_name} as battle started.")
300
  active_agent_task.cancel()
 
 
 
 
 
 
 
 
301
 
302
  async def deactivate_current_agent(reason: str = "cycle"):
303
  """Cleans up the currently active agent and resets state."""
304
  global active_agent_name, active_agent_instance, active_agent_task, current_battle_instance
305
 
306
- print(f"Lifecycle: Deactivating agent '{active_agent_name}' (Reason: {reason})...")
307
-
308
- # Different messages based on deactivation reason
 
309
  if reason == "battle_end":
310
- await update_display_html(create_idle_html("Battle Finished", "Preparing to select next agent..."))
311
  elif reason == "cycle":
312
- await update_display_html(create_idle_html("Cycling Agents", "Preparing for next selection..."))
313
- else:
314
- await update_display_html(create_idle_html(f"Agent Reset ({reason})", "Preparing for next selection..."))
 
 
 
 
 
 
 
 
 
315
 
316
  agent = active_agent_instance
317
  task = active_agent_task
318
-
319
  # Store a local copy of the battle instance before clearing it
320
- last_battle_instance = current_battle_instance
321
 
322
- # Clear state first to prevent race conditions
 
 
323
  active_agent_name = None
324
  active_agent_instance = None
325
  active_agent_task = None
326
  current_battle_instance = None
 
327
 
328
- # Cancel the accept_challenges task if it's still running
 
329
  if task and not task.done():
330
- print(f"Lifecycle: Cancelling task for {agent.username if agent else 'unknown agent'}...")
331
  task.cancel()
332
  try:
 
333
  await asyncio.wait_for(task, timeout=2.0)
 
334
  except asyncio.CancelledError:
335
- print(f"Lifecycle: Task cancellation confirmed.")
336
  except asyncio.TimeoutError:
337
- print(f"Lifecycle: Task did not confirm cancellation within timeout.")
338
  except Exception as e:
339
- print(f"Lifecycle: Error during task cancellation wait: {e}")
 
340
 
341
- # Disconnect the player
342
  if agent:
343
  print(f"Lifecycle: Disconnecting player {agent.username}...")
344
  try:
 
345
  if hasattr(agent, '_websocket') and agent._websocket and agent._websocket.open:
346
- await agent.disconnect()
347
- print(f"Lifecycle: Player {agent.username} disconnected.")
348
  else:
349
- print(f"Lifecycle: Player {agent.username} already disconnected or websocket not ready.")
350
  except Exception as e:
 
351
  print(f"ERROR during agent disconnect ({agent.username}): {e}")
 
352
 
353
- # Add a brief delay after deactivation to ensure clean state
354
- await asyncio.sleep(3)
355
- print(f"Lifecycle: Agent deactivation complete.")
356
 
357
  # --- Main Background Task ---
358
  async def manage_agent_lifecycle():
@@ -360,14 +285,19 @@ async def manage_agent_lifecycle():
360
  global active_agent_instance, active_agent_task, current_battle_instance
361
 
362
  print("Background lifecycle manager started.")
363
- REFRESH_INTERVAL_SECONDS = 3
 
 
 
 
364
  loop_counter = 0
365
 
366
  while True:
367
  loop_counter += 1
368
- try:
369
- print(f"\n--- Lifecycle Check #{loop_counter} [{time.strftime('%H:%M:%S')}] ---")
370
 
 
371
  # ==================================
372
  # State 1: No agent active
373
  # ==================================
@@ -375,214 +305,425 @@ async def manage_agent_lifecycle():
375
  print(f"[{loop_counter}] State 1: No active agent. Selecting...")
376
  activated = await select_and_activate_new_agent()
377
  if not activated:
378
- print(f"[{loop_counter}] State 1: Activation failed. Waiting before retry.")
379
- await asyncio.sleep(10)
380
- continue
381
  else:
382
- print(f"[{loop_counter}] State 1: Agent '{active_agent_name}' activated.")
 
383
 
384
  # ==================================
385
  # State 2: Agent is active
386
  # ==================================
387
  else:
388
- print(f"[{loop_counter}] State 2: Agent '{active_agent_name}' active.")
389
-
390
- # --- Sub-state: Monitoring an ongoing battle ---
391
- if current_battle_instance is not None:
392
- battle_tag = current_battle_instance.battle_tag
393
- print(f"[{loop_counter}] State 2: Monitoring battle {battle_tag}")
394
-
395
- # Get the potentially updated battle object from the agent's internal list
396
- battle_obj = active_agent_instance._battles.get(battle_tag)
397
-
398
- if battle_obj and battle_obj.finished:
399
- # Battle is finished, deactivate agent
400
- print(f"[{loop_counter}] Battle {battle_tag} is FINISHED. Deactivating agent.")
401
- await deactivate_current_agent(reason="battle_end")
402
- await asyncio.sleep(5) # Brief cooldown between battles
403
- continue
404
 
405
- elif not battle_obj:
406
- # Battle object missing unexpectedly, deactivate agent
407
- print(f"[{loop_counter}] WARNING: Battle object for {battle_tag} MISSING from agent's list! Deactivating.")
408
- await deactivate_current_agent(reason="battle_object_missing")
409
- continue
410
- else:
411
- # Battle is ongoing and battle object exists
412
- pass
413
-
414
- # --- Sub-state: Checking for a new battle / Agent is idle ---
415
- elif current_battle_instance is None:
416
- print(f"[{loop_counter}] State 2: Checking for new battle...")
417
- await check_for_new_battle()
418
 
419
- # --- If a NEW battle was just detected by check_for_new_battle() ---
420
  if current_battle_instance:
421
  battle_tag = current_battle_instance.battle_tag
422
- print(f"[{loop_counter}] State 2: *** NEW BATTLE DETECTED: {battle_tag} ***")
423
 
424
- # Check if the battle ID has the non-public/suffixed format
425
  parts = battle_tag.split('-')
426
  is_suffixed_format = len(parts) > 3 and parts[2].isdigit()
427
 
428
  if is_suffixed_format:
429
- # Non-public format detected, forfeit immediately
430
- print(f"[{loop_counter}] Detected non-public battle format ({battle_tag}). Forfeiting.")
431
- await update_display_html(create_idle_html(f"Forfeiting non-public battle...", f"Agent {active_agent_name} detected private room."))
432
  try:
433
- # Send forfeit command TO THE SPECIFIC BATTLE ROOM
434
- await active_agent_instance.send_message("/forfeit", battle_tag)
435
- print(f"[{loop_counter}] Sent /forfeit command for {battle_tag}.")
436
- await asyncio.sleep(1) # Give a moment for forfeit to process server-side
 
437
  except Exception as forfeit_err:
438
- print(f"[{loop_counter}] ERROR sending forfeit command for {battle_tag}: {forfeit_err}")
439
- # Deactivate the agent as the battle is effectively over
440
  await deactivate_current_agent(reason="forfeited_private_battle")
441
- continue
442
 
443
  else:
444
- # Public battle format detected, display the battle iframe
445
- print(f"[{loop_counter}] Public battle format detected. Displaying battle.")
446
  await update_display_html(create_battle_iframe(battle_tag))
 
447
 
448
- # --- If no new battle was found, agent remains idle ---
449
  else:
450
- print(f"[{loop_counter}] State 2: No new battle found. Agent '{active_agent_name}' remains idle.")
451
- expected_idle_html = create_idle_html(f"Agent <strong>{active_agent_name}</strong> is ready!",
452
- f"Please challenge <strong>{active_agent_name}</strong> to a <strong>{DEFAULT_BATTLE_FORMAT}</strong> battle.")
453
- await update_display_html(expected_idle_html)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
454
 
455
  # --- Global Exception Handling for the main loop ---
 
 
 
456
  except Exception as e:
457
- print(f"ERROR in main lifecycle loop #{loop_counter}: {e}")
458
  traceback.print_exc()
 
459
  if active_agent_instance:
460
- # If an error happened while an agent was active, try to deactivate it
461
- await deactivate_current_agent(reason="main_loop_error")
 
 
 
 
 
 
 
 
462
  else:
463
- # If error happened before agent activation, show generic error
464
- await update_display_html(create_error_html(f"Error in lifecycle manager: {e}"))
 
 
 
465
  # Wait longer after a major error before trying again
466
- await asyncio.sleep(10)
 
 
 
 
 
 
 
467
 
468
- # --- Sleep at the end of each loop iteration (if no 'continue' was hit) ---
469
- print(f"[{loop_counter}] Sleeping {REFRESH_INTERVAL_SECONDS}s before next check.")
470
- await asyncio.sleep(REFRESH_INTERVAL_SECONDS)
471
 
472
  def log_task_exception(task: asyncio.Task):
473
- """Callback to log exceptions from background tasks."""
474
  try:
475
  if task.cancelled():
476
- print(f"Task {task.get_name()} was cancelled.")
 
477
  return
478
- task.result() # Raises exception if task failed
479
- print(f"Task {task.get_name()} finished cleanly.")
 
480
  except asyncio.CancelledError:
481
- pass
 
482
  except Exception as e:
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():
522
- """Serves the main HTML page with WebSocket connection."""
 
 
523
  return """
524
  <!DOCTYPE html>
525
- <html>
526
  <head>
 
 
527
  <title>Pokemon Battle Livestream</title>
 
 
 
528
  <style>
529
- body, html {
 
 
 
 
 
530
  margin: 0;
531
  padding: 0;
532
  height: 100%;
533
- overflow: hidden;
 
 
 
 
534
  }
 
 
535
  #stream-container {
536
- position: fixed;
537
  top: 0;
538
  left: 0;
539
- bottom: 0;
540
- right: 0;
541
  width: 100%;
542
  height: 100%;
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
543
  }
 
544
  </style>
545
  </head>
546
  <body>
547
- <div id="stream-container"></div>
548
-
 
549
  <script>
550
  const streamContainer = document.getElementById('stream-container');
551
-
552
- // Create WebSocket connection
553
- const ws = new WebSocket(`${location.protocol === 'https:' ? 'wss' : 'ws'}://${location.host}/ws`);
554
-
555
- // Connection opened
556
- ws.addEventListener('open', (event) => {
557
- console.log('Connected to WebSocket server');
558
- });
559
-
560
- // Listen for messages
561
- ws.addEventListener('message', (event) => {
562
- console.log('Received update from server');
563
- // Set inner HTML directly - complete document replacement
564
- const parser = new DOMParser();
565
- const doc = parser.parseFromString(event.data, 'text/html');
566
- const bodyContent = doc.body.innerHTML;
567
-
568
- // Update just the body content to avoid page refresh flicker
569
- if (bodyContent) {
570
- streamContainer.innerHTML = bodyContent;
571
- }
572
- });
573
-
574
- // Handle connection close/errors
575
- ws.addEventListener('close', (event) => {
576
- console.log('WebSocket connection closed');
577
- setTimeout(() => {
578
- console.log('Attempting to reconnect...');
579
- location.reload();
580
- }, 3000);
581
- });
582
-
583
- ws.addEventListener('error', (event) => {
584
- console.error('WebSocket error:', event);
585
- });
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
586
  </script>
587
  </body>
588
  </html>
@@ -593,25 +734,45 @@ async def websocket_endpoint(websocket: WebSocket):
593
  await manager.connect(websocket)
594
  try:
595
  while True:
596
- # Just keep the connection open
597
- # The updates are pushed from the lifecycle manager
598
- await websocket.receive_text()
599
- except WebSocketDisconnect:
600
- manager.disconnect(websocket)
 
 
 
 
 
 
 
601
  except Exception as e:
 
602
  print(f"WebSocket error: {e}")
603
- manager.disconnect(websocket)
 
 
604
 
605
  @app.on_event("startup")
606
  async def startup_event():
607
  """Start background tasks when the application starts."""
608
  global background_task_handle
609
-
610
- # Mount static files directory for images
611
- app.mount("/static", StaticFiles(directory="static"), name="static")
612
-
 
 
 
 
 
 
 
 
613
  print("πŸš€ Starting background tasks")
614
- background_task_handle = asyncio.create_task(manage_agent_lifecycle(), name="lifecycle_manager")
 
 
615
  background_task_handle.add_done_callback(log_task_exception)
616
  print("βœ… Background tasks started")
617
 
@@ -619,48 +780,91 @@ async def startup_event():
619
  async def shutdown_event():
620
  """Clean up tasks when shutting down."""
621
  global background_task_handle, active_agent_instance
622
-
623
- print("\nShutting down app. Cleaning up...")
624
-
625
- # Cancel background task
626
  if background_task_handle and not background_task_handle.done():
 
627
  background_task_handle.cancel()
628
  try:
629
  await asyncio.wait_for(background_task_handle, timeout=5.0)
630
- except (asyncio.CancelledError, asyncio.TimeoutError):
631
- pass
632
-
633
- # Disconnect agent if active
634
- if active_agent_instance:
 
 
 
 
 
 
 
 
 
635
  try:
636
- if hasattr(active_agent_instance, '_websocket') and active_agent_instance._websocket and active_agent_instance._websocket.open:
637
- await active_agent_instance.disconnect()
 
 
 
 
638
  except Exception as e:
639
- print(f"Error during agent disconnect on shutdown: {e}")
640
-
641
- print("Cleanup complete. Shutting down.")
 
 
 
 
 
 
 
 
 
 
 
642
 
643
  # For direct script execution
644
  if __name__ == "__main__":
645
  import uvicorn
646
-
647
- logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
 
 
 
 
 
 
648
  logging.getLogger('poke_env').setLevel(logging.WARNING)
649
-
 
650
  print("Starting Pokemon Battle Livestream Server...")
651
- print("="*50)
652
-
653
  if not AVAILABLE_AGENT_NAMES:
654
- print("FATAL: No agents found with configured passwords. Please set environment variables:")
 
 
655
  for name, cfg in AGENT_CONFIGS.items():
656
- print(f"- {cfg.get('password_env_var', 'N/A')} (for agent: {name})")
657
- print("="*50)
658
  exit("Exiting due to missing agent passwords.")
659
  else:
660
- print("Found available agents:")
661
  for name in AVAILABLE_AGENT_NAMES:
662
- print(f"- {name}")
663
- print("="*50)
664
-
 
 
665
  # Run with uvicorn
666
- uvicorn.run(app, host="0.0.0.0", port=7860)
 
 
 
 
 
 
 
 
54
  # --- Create FastAPI app ---
55
  app = FastAPI(title="Pokemon Battle Livestream")
56
 
 
57
  # --- Helper Functions ---
58
  def get_active_battle(agent: Player) -> Optional[Battle]:
59
  """Returns the first non-finished battle for an agent."""
60
  if agent and agent._battles:
61
  active_battles = [b for b in agent._battles.values() if not b.finished]
62
  if active_battles:
63
+ # Ensure the battle object has a battle_tag before returning
64
+ if hasattr(active_battles[0], 'battle_tag') and active_battles[0].battle_tag:
65
+ # Check if the battle_tag has the expected format (starts with 'battle-')
66
+ if active_battles[0].battle_tag.startswith("battle-"):
67
+ return active_battles[0]
68
+ else:
69
+ # This handles cases where the battle object might exist but tag isn't ready
70
+ # print(f"DEBUG: Found active battle for {agent.username} but tag '{active_battles[0].battle_tag}' not ready.")
71
+ return None
72
  else:
73
+ # print(f"DEBUG: Found active battle for {agent.username} but it has no battle_tag attribute yet.")
74
  return None
75
  return None
76
 
77
  def create_battle_iframe(battle_id: str) -> str:
78
+ """Creates JUST the HTML for the battle iframe tag."""
79
+ print("Creating iframe content for battle ID: ", battle_id)
80
+ # Use the official client URL unless you specifically need the test client
81
+ # battle_url = f"https://play.pokemonshowdown.com/{battle_id}"
82
+ battle_url = f"https://jofthomas.com/play.pokemonshowdown.com/testclient.html#{battle_id}" # Using your custom URL
83
+
84
+ # Return ONLY the iframe tag with a class for styling
85
  return f"""
86
+ <iframe
87
+ id="battle-iframe"
88
+ class="battle-iframe"
89
+ src="{battle_url}"
90
+ allowfullscreen
91
+ ></iframe>
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
92
  """
93
 
94
  def create_idle_html(status_message: str, instruction: str) -> str:
95
+ """Creates a visually appealing idle screen HTML fragment."""
96
+ # Returns ONLY the content div, not the full HTML page
97
  return f"""
98
+ <div class="content-container idle-container">
99
+ <div class="message-box">
100
+ <p class="status">{status_message}</p>
101
+ <p class="instruction">{instruction}</p>
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
102
  </div>
103
+ </div>
 
104
  """
105
 
106
  def create_error_html(error_msg: str) -> str:
107
+ """Creates HTML fragment to display an error message."""
108
+ # Returns ONLY the content div, not the full HTML page
109
  return f"""
110
+ <div class="content-container error-container">
111
+ <div class="message-box">
112
+ <p class="status">🚨 Error 🚨</p>
113
+ <p class="instruction">{error_msg}</p>
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
114
  </div>
115
+ </div>
 
116
  """
117
 
118
+ async def update_display_html(new_html_fragment: str) -> None:
119
+ """Updates the current display HTML fragment and broadcasts to all clients."""
120
+ # Pass the fragment directly
121
+ await manager.update_all(new_html_fragment)
122
+ print("HTML Display FRAGMENT UPDATED and broadcasted.")
123
+
124
 
125
  # --- Agent Lifecycle Management ---
126
  async def select_and_activate_new_agent():
 
129
 
130
  if not AVAILABLE_AGENT_NAMES:
131
  print("Lifecycle: No available agents with passwords set.")
132
+ await update_display_html(create_error_html("No agents available. Check server logs/environment variables."))
133
  return False
134
 
135
  selected_name = random.choice(AVAILABLE_AGENT_NAMES)
 
139
  agent_password = os.environ.get(password_env_var)
140
 
141
  print(f"Lifecycle: Activating agent '{selected_name}'...")
142
+ # Use HTML tags for slight emphasis if desired
143
+ await update_display_html(create_idle_html("Selecting Next Agent...", f"Preparing <strong>{selected_name}</strong>..."))
144
 
145
  try:
146
  account_config = AccountConfiguration(selected_name, agent_password)
 
153
  )
154
 
155
  # Start the task to accept exactly one battle challenge
156
+ # Setting name for easier debugging
157
+ task = asyncio.create_task(agent.accept_challenges(None, 1), name=f"AcceptChallenge_{selected_name}")
158
+ task.add_done_callback(log_task_exception) # Add callback for errors
159
 
160
  # Update global state
161
  active_agent_name = selected_name
162
  active_agent_instance = agent
163
  active_agent_task = task
164
+
165
  print(f"Lifecycle: Agent '{selected_name}' is active and listening for 1 challenge.")
166
+ # Use HTML tags for slight emphasis
167
+ await update_display_html(create_idle_html(f"Agent Ready: <strong>{selected_name}</strong>",
168
  f"Please challenge <strong>{selected_name}</strong> to a <strong>{DEFAULT_BATTLE_FORMAT}</strong> battle."))
169
  return True
170
 
 
172
  error_msg = f"Failed to activate agent '{selected_name}': {e}"
173
  print(error_msg)
174
  traceback.print_exc()
175
+ await update_display_html(create_error_html(f"Error activating {selected_name}. Please wait or check logs."))
176
+
177
  # Clear state if activation failed
178
  active_agent_name = None
179
  active_agent_instance = None
 
181
  return False
182
 
183
  async def check_for_new_battle():
184
+ """Checks if the active agent has started a battle with a valid tag."""
185
  global active_agent_instance, current_battle_instance
186
+
187
  if active_agent_instance:
188
  battle = get_active_battle(active_agent_instance)
189
+ # Check if battle exists AND has a valid battle_tag
190
  if battle and battle.battle_tag:
191
  print(f"Lifecycle: Agent '{active_agent_name}' started battle: {battle.battle_tag}")
192
  current_battle_instance = battle
193
+
194
+ # Stop the agent from listening for more challenges once a battle starts
195
  if active_agent_task and not active_agent_task.done():
196
  print(f"Lifecycle: Cancelling accept_challenges task for {active_agent_name} as battle started.")
197
  active_agent_task.cancel()
198
+ # Optional: Wait briefly for cancellation confirmation, but don't block excessively
199
+ # try:
200
+ # await asyncio.wait_for(active_agent_task, timeout=0.5)
201
+ # except (asyncio.CancelledError, asyncio.TimeoutError):
202
+ # pass # Expected outcomes
203
+ # else:
204
+ # print(f"DEBUG: get_active_battle returned None or battle without tag.")
205
+
206
 
207
  async def deactivate_current_agent(reason: str = "cycle"):
208
  """Cleans up the currently active agent and resets state."""
209
  global active_agent_name, active_agent_instance, active_agent_task, current_battle_instance
210
 
211
+ agent_name_to_deactivate = active_agent_name # Store before clearing
212
+ print(f"Lifecycle: Deactivating agent '{agent_name_to_deactivate}' (Reason: {reason})...")
213
+
214
+ # Display appropriate intermediate message
215
  if reason == "battle_end":
216
+ await update_display_html(create_idle_html("Battle Finished!", f"Agent <strong>{agent_name_to_deactivate}</strong> completed the match."))
217
  elif reason == "cycle":
218
+ await update_display_html(create_idle_html("Cycling Agents", f"Switching from <strong>{agent_name_to_deactivate}</strong>..."))
219
+ elif reason == "forfeited_private_battle":
220
+ await update_display_html(create_idle_html("Switching Agent", f"Agent <strong>{agent_name_to_deactivate}</strong> forfeited a private battle."))
221
+ else: # Generic reason or error
222
+ await update_display_html(create_idle_html(f"Resetting Agent ({reason})", f"Cleaning up <strong>{agent_name_to_deactivate}</strong>..."))
223
+
224
+ # Give users a moment to see the intermediate message
225
+ await asyncio.sleep(3) # Adjust duration as needed
226
+
227
+ # Show the "preparing next agent" message before lengthy cleanup
228
+ await update_display_html(create_idle_html("Preparing Next Agent...", "Please wait..."))
229
+
230
 
231
  agent = active_agent_instance
232
  task = active_agent_task
233
+
234
  # Store a local copy of the battle instance before clearing it
235
+ # last_battle_instance = current_battle_instance # Not strictly needed now
236
 
237
+ # --- Crucial: Clear global state variables FIRST ---
238
+ # This prevents race conditions where the lifecycle loop might try to
239
+ # access the agent while it's being deactivated.
240
  active_agent_name = None
241
  active_agent_instance = None
242
  active_agent_task = None
243
  current_battle_instance = None
244
+ print(f"Lifecycle: Global state cleared for '{agent_name_to_deactivate}'.")
245
 
246
+ # --- Now perform cleanup actions ---
247
+ # Cancel the accept_challenges task if it's still running (it might already be done/cancelled)
248
  if task and not task.done():
249
+ print(f"Lifecycle: Ensuring task cancellation for {agent_name_to_deactivate} ({task.get_name()})...")
250
  task.cancel()
251
  try:
252
+ # Wait briefly for the task to acknowledge cancellation
253
  await asyncio.wait_for(task, timeout=2.0)
254
+ print(f"Lifecycle: Task cancellation confirmed for {agent_name_to_deactivate}.")
255
  except asyncio.CancelledError:
256
+ print(f"Lifecycle: Task cancellation confirmation (CancelledError) for {agent_name_to_deactivate}.")
257
  except asyncio.TimeoutError:
258
+ print(f"Lifecycle: Task did not confirm cancellation within timeout for {agent_name_to_deactivate}.")
259
  except Exception as e:
260
+ # Catch other potential errors during task cleanup
261
+ print(f"Lifecycle: Error during task cancellation wait for {agent_name_to_deactivate}: {e}")
262
 
263
+ # Disconnect the player (ensure agent object exists)
264
  if agent:
265
  print(f"Lifecycle: Disconnecting player {agent.username}...")
266
  try:
267
+ # Check websocket state before attempting disconnection
268
  if hasattr(agent, '_websocket') and agent._websocket and agent._websocket.open:
269
+ await agent.disconnect()
270
+ print(f"Lifecycle: Player {agent.username} disconnected successfully.")
271
  else:
272
+ print(f"Lifecycle: Player {agent.username} already disconnected or websocket not available.")
273
  except Exception as e:
274
+ # Log errors during disconnection but don't halt the process
275
  print(f"ERROR during agent disconnect ({agent.username}): {e}")
276
+ traceback.print_exc() # Log full traceback for debugging
277
 
278
+ # Add a brief delay AFTER deactivation before the loop potentially selects a new agent
279
+ await asyncio.sleep(2) # Reduced from 3, adjust as needed
280
+ print(f"Lifecycle: Agent '{agent_name_to_deactivate}' deactivation complete.")
281
 
282
  # --- Main Background Task ---
283
  async def manage_agent_lifecycle():
 
285
  global active_agent_instance, active_agent_task, current_battle_instance
286
 
287
  print("Background lifecycle manager started.")
288
+ REFRESH_INTERVAL_SECONDS = 3 # How often to check state when idle/in battle
289
+ LOOP_COOLDOWN_SECONDS = 1 # Small delay at end of loop if no other waits occurred
290
+ ERROR_RETRY_DELAY_SECONDS = 10 # Longer delay after errors
291
+ POST_BATTLE_DELAY_SECONDS = 5 # Delay after a battle finishes before selecting next agent
292
+
293
  loop_counter = 0
294
 
295
  while True:
296
  loop_counter += 1
297
+ loop_start_time = time.monotonic()
298
+ print(f"\n--- Lifecycle Check #{loop_counter} [{time.strftime('%H:%M:%S')}] ---")
299
 
300
+ try:
301
  # ==================================
302
  # State 1: No agent active
303
  # ==================================
 
305
  print(f"[{loop_counter}] State 1: No active agent. Selecting...")
306
  activated = await select_and_activate_new_agent()
307
  if not activated:
308
+ print(f"[{loop_counter}] State 1: Activation failed. Waiting {ERROR_RETRY_DELAY_SECONDS}s before retry.")
309
+ await asyncio.sleep(ERROR_RETRY_DELAY_SECONDS)
 
310
  else:
311
+ print(f"[{loop_counter}] State 1: Agent '{active_agent_name}' activated successfully.")
312
+ # No sleep here, proceed to next check immediately if needed
313
 
314
  # ==================================
315
  # State 2: Agent is active
316
  # ==================================
317
  else:
318
+ agent_name = active_agent_name # Cache for logging
319
+ print(f"[{loop_counter}] State 2: Agent '{agent_name}' is active.")
 
 
 
 
 
 
 
 
 
 
 
 
 
 
320
 
321
+ # --- Sub-state: Check for new battle if none is tracked ---
322
+ if current_battle_instance is None:
323
+ print(f"[{loop_counter}] State 2a: Checking for new battle for '{agent_name}'...")
324
+ await check_for_new_battle() # This updates current_battle_instance if found
 
 
 
 
 
 
 
 
 
325
 
 
326
  if current_battle_instance:
327
  battle_tag = current_battle_instance.battle_tag
328
+ print(f"[{loop_counter}] State 2a: *** NEW BATTLE DETECTED: {battle_tag} for '{agent_name}' ***")
329
 
330
+ # Check for non-public/suffixed format (heuristic: more than 3 parts, 3rd part is number)
331
  parts = battle_tag.split('-')
332
  is_suffixed_format = len(parts) > 3 and parts[2].isdigit()
333
 
334
  if is_suffixed_format:
335
+ # Forfeit immediately if it looks like a private/suffixed battle ID
336
+ print(f"[{loop_counter}] Detected potentially non-public battle format ({battle_tag}). Forfeiting.")
337
+ # Don't update display yet, do it before deactivation
338
  try:
339
+ if active_agent_instance: # Ensure agent still exists
340
+ await active_agent_instance.forfeit(battle_tag)
341
+ # await active_agent_instance.send_message("/forfeit", battle_tag) # Alternative
342
+ print(f"[{loop_counter}] Sent forfeit command for {battle_tag}.")
343
+ await asyncio.sleep(1.5) # Give forfeit time to register
344
  except Exception as forfeit_err:
345
+ print(f"[{loop_counter}] ERROR sending forfeit for {battle_tag}: {forfeit_err}")
346
+ # Deactivate agent after forfeit attempt
347
  await deactivate_current_agent(reason="forfeited_private_battle")
348
+ continue # Skip rest of the loop for this iteration
349
 
350
  else:
351
+ # Public battle format - display the iframe
352
+ print(f"[{loop_counter}] Public battle format detected. Displaying battle {battle_tag}.")
353
  await update_display_html(create_battle_iframe(battle_tag))
354
+ # Now fall through to monitor this battle in the next section
355
 
 
356
  else:
357
+ # No new battle found, agent remains idle
358
+ print(f"[{loop_counter}] State 2a: No new battle found. Agent '{agent_name}' remains idle, waiting for challenge.")
359
+ # Periodically refresh idle screen to ensure consistency
360
+ idle_html = create_idle_html(f"Agent Ready: <strong>{agent_name}</strong>",
361
+ f"Please challenge <strong>{agent_name}</strong> to a <strong>{DEFAULT_BATTLE_FORMAT}</strong> battle.")
362
+ await update_display_html(idle_html)
363
+ await asyncio.sleep(REFRESH_INTERVAL_SECONDS) # Wait before next check if idle
364
+
365
+
366
+ # --- Sub-state: Monitor ongoing battle ---
367
+ if current_battle_instance is not None:
368
+ battle_tag = current_battle_instance.battle_tag
369
+ print(f"[{loop_counter}] State 2b: Monitoring battle {battle_tag} for '{agent_name}'")
370
+
371
+ # Ensure agent instance still exists before accessing its battles
372
+ if not active_agent_instance:
373
+ print(f"[{loop_counter}] WARNING: Agent instance for '{agent_name}' disappeared while monitoring battle {battle_tag}! Deactivating.")
374
+ await deactivate_current_agent(reason="agent_disappeared_mid_battle")
375
+ continue
376
+
377
+ # Get potentially updated battle object directly from agent's state
378
+ # Use .get() for safety
379
+ battle_obj = active_agent_instance._battles.get(battle_tag)
380
+
381
+ if battle_obj and battle_obj.finished:
382
+ print(f"[{loop_counter}] Battle {battle_tag} is FINISHED. Deactivating agent '{agent_name}'.")
383
+ await deactivate_current_agent(reason="battle_end")
384
+ print(f"[{loop_counter}] Waiting {POST_BATTLE_DELAY_SECONDS}s post-battle before selecting next agent.")
385
+ await asyncio.sleep(POST_BATTLE_DELAY_SECONDS)
386
+ continue # Start next loop iteration to select new agent
387
+
388
+ elif not battle_obj:
389
+ # This can happen briefly during transitions or if battle ends unexpectedly
390
+ print(f"[{loop_counter}] WARNING: Battle object for {battle_tag} not found in agent's list for '{agent_name}'. Battle might have ended abruptly. Deactivating.")
391
+ await deactivate_current_agent(reason="battle_object_missing")
392
+ continue
393
+
394
+ else:
395
+ # Battle is ongoing, battle object exists, iframe should be displayed
396
+ print(f"[{loop_counter}] Battle {battle_tag} ongoing for '{agent_name}'.")
397
+ # Optionally: Could re-send iframe HTML periodically if needed, but usually not necessary
398
+ # await update_display_html(create_battle_iframe(battle_tag))
399
+ await asyncio.sleep(REFRESH_INTERVAL_SECONDS) # Wait before next check
400
 
401
  # --- Global Exception Handling for the main loop ---
402
+ except asyncio.CancelledError:
403
+ print("Lifecycle manager task cancelled.")
404
+ raise # Re-raise to ensure proper shutdown
405
  except Exception as e:
406
+ print(f"!!! ERROR in main lifecycle loop #{loop_counter}: {e} !!!")
407
  traceback.print_exc()
408
+ current_agent_name = active_agent_name # Cache name before deactivation attempts
409
  if active_agent_instance:
410
+ print(f"Attempting to deactivate agent '{current_agent_name}' due to loop error...")
411
+ try:
412
+ await deactivate_current_agent(reason="main_loop_error")
413
+ except Exception as deactivation_err:
414
+ print(f"Error during error-handling deactivation: {deactivation_err}")
415
+ # Ensure state is cleared even if deactivation fails partially
416
+ active_agent_name = None
417
+ active_agent_instance = None
418
+ active_agent_task = None
419
+ current_battle_instance = None
420
  else:
421
+ # Error happened potentially before agent activation or after clean deactivation
422
+ print("No active agent instance during loop error.")
423
+ # Show a generic error on the frontend
424
+ await update_display_html(create_error_html(f"A server error occurred in the lifecycle manager. Please wait. ({e})"))
425
+
426
  # Wait longer after a major error before trying again
427
+ print(f"Waiting {ERROR_RETRY_DELAY_SECONDS}s after loop error.")
428
+ await asyncio.sleep(ERROR_RETRY_DELAY_SECONDS)
429
+ continue # Go to next loop iteration after error handling
430
+
431
+ # --- Delay at end of loop if no other significant waits happened ---
432
+ elapsed_time = time.monotonic() - loop_start_time
433
+ if elapsed_time < LOOP_COOLDOWN_SECONDS:
434
+ await asyncio.sleep(LOOP_COOLDOWN_SECONDS - elapsed_time)
435
 
 
 
 
436
 
437
  def log_task_exception(task: asyncio.Task):
438
+ """Callback to log exceptions from background tasks (like accept_challenges)."""
439
  try:
440
  if task.cancelled():
441
+ # Don't log cancellation as an error, it's often expected
442
+ print(f"Task '{task.get_name()}' was cancelled.")
443
  return
444
+ # Accessing result will raise exception if task failed
445
+ task.result()
446
+ print(f"Task '{task.get_name()}' completed successfully.")
447
  except asyncio.CancelledError:
448
+ print(f"Task '{task.get_name()}' confirmed cancelled (exception caught).")
449
+ pass # Expected
450
  except Exception as e:
451
+ # Log actual errors
452
+ print(f"!!! Exception in background task '{task.get_name()}': {e} !!!")
453
  traceback.print_exc()
454
+ # Optionally: Trigger some recovery or notification here if needed
455
 
456
 
457
  # --- WebSocket connection manager ---
458
  class ConnectionManager:
459
  def __init__(self):
460
  self.active_connections: Set[WebSocket] = set()
461
+ # Initialize with the idle HTML fragment
462
+ self.current_html_fragment: str = create_idle_html("Initializing...", "Setting up PokΓ©mon Battle Stream")
463
 
464
  async def connect(self, websocket: WebSocket):
465
  await websocket.accept()
466
  self.active_connections.add(websocket)
467
+ print(f"Client connected. Sending current state. Total clients: {len(self.active_connections)}")
468
+ # Send current state (HTML fragment) to newly connected client
469
+ try:
470
+ await websocket.send_text(self.current_html_fragment)
471
+ except Exception as e:
472
+ print(f"Error sending initial state to new client: {e}")
473
+ # Consider removing the connection if initial send fails
474
+ await self.disconnect(websocket)
475
+
476
+
477
+ async def disconnect(self, websocket: WebSocket):
478
+ # Use discard() to safely remove even if not present
479
+ self.active_connections.discard(websocket)
480
+ print(f"Client disconnected. Total clients: {len(self.active_connections)}")
481
+
482
+ async def update_all(self, html_fragment: str):
483
+ """Update the current HTML fragment and broadcast to all clients."""
484
+ if self.current_html_fragment == html_fragment:
485
+ # print("Skipping broadcast, HTML fragment unchanged.")
486
+ return # Avoid unnecessary updates if content is identical
487
+
488
+ self.current_html_fragment = html_fragment
489
+ if not self.active_connections:
490
+ # print("No active connections to broadcast update to.")
491
+ return
492
+
493
+ print(f"Broadcasting update to {len(self.active_connections)} clients...")
494
+
495
+ # Create a list of tasks to send updates concurrently
496
+ # Make a copy of the set for safe iteration during potential disconnects
497
+ send_tasks = [
498
+ connection.send_text(html_fragment)
499
+ for connection in list(self.active_connections) # Iterate over a copy
500
+ ]
501
+
502
+ # Use asyncio.gather to send to all clients, collecting results/exceptions
503
+ results = await asyncio.gather(*send_tasks, return_exceptions=True)
504
+
505
+ # Handle potential errors during broadcast (e.g., client disconnected abruptly)
506
+ # Iterate over connections again, checking results
507
+ connections_to_remove = set()
508
+ for i, result in enumerate(results):
509
+ connection = list(self.active_connections)[i] # Assumes order is maintained
510
+ if isinstance(result, Exception):
511
+ print(f"Error sending update to client: {result}. Marking for removal.")
512
+ connections_to_remove.add(connection)
513
+
514
+ # Disconnect clients that failed
515
+ for connection in connections_to_remove:
516
+ await self.disconnect(connection)
517
+
518
 
519
  manager = ConnectionManager()
520
+
521
  # --- API Routes ---
522
  @app.get("/", response_class=HTMLResponse)
523
  async def get_homepage():
524
+ """Serves the main HTML page with WebSocket connection and improved styling."""
525
+ # NOTE: Ensure the static path '/static/pokemon_huggingface.png' is correct
526
+ # and the image exists in a 'static' folder next to your main.py
527
  return """
528
  <!DOCTYPE html>
529
+ <html lang="en">
530
  <head>
531
+ <meta charset="UTF-8">
532
+ <meta name="viewport" content="width=device-width, initial-scale=1.0">
533
  <title>Pokemon Battle Livestream</title>
534
+ <link rel="preconnect" href="https://fonts.googleapis.com">
535
+ <link rel="preconnect" href="https://fonts.gstatic.com" crossorigin>
536
+ <link href="https://fonts.googleapis.com/css2?family=Poppins:wght@400;700&family=Press+Start+2P&display=swap" rel="stylesheet">
537
  <style>
538
+ /* Basic Reset */
539
+ * {
540
+ box-sizing: border-box;
541
+ }
542
+
543
+ html, body {
544
  margin: 0;
545
  padding: 0;
546
  height: 100%;
547
+ width: 100%;
548
+ overflow: hidden; /* Prevent scrollbars on body */
549
+ font-family: 'Poppins', sans-serif; /* Default font */
550
+ color: #ffffff; /* Default text color */
551
+ background-color: #1a1a1a; /* Dark background */
552
  }
553
+
554
+ /* Container for dynamic content */
555
  #stream-container {
556
+ position: fixed; /* Use fixed to ensure it covers viewport */
557
  top: 0;
558
  left: 0;
 
 
559
  width: 100%;
560
  height: 100%;
561
+ display: flex; /* Use flexbox for centering content */
562
+ justify-content: center;
563
+ align-items: center;
564
+ }
565
+
566
+ /* Iframe Styling */
567
+ .battle-iframe {
568
+ width: 100%;
569
+ height: 100%;
570
+ border: none; /* Remove default border */
571
+ display: block; /* Prevents potential extra space below iframe */
572
+ }
573
+
574
+ /* Base Content Container Styling (used by idle/error) */
575
+ .content-container {
576
+ width: 100%;
577
+ height: 100%;
578
+ display: flex;
579
+ flex-direction: column;
580
+ justify-content: center;
581
+ align-items: center;
582
+ padding: 20px;
583
+ text-align: center;
584
+ }
585
+
586
+ /* Idle Screen Specific Styling */
587
+ .idle-container {
588
+ /* Ensure the background covers the entire container */
589
+ background-image: url('/static/pokemon_huggingface.png');
590
+ background-size: cover;
591
+ background-position: center;
592
+ background-repeat: no-repeat;
593
+ }
594
+
595
+ /* Error Screen Specific Styling */
596
+ .error-container {
597
+ background: linear-gradient(135deg, #4d0000, #1a0000); /* Dark red gradient */
598
+ }
599
+
600
+ /* Message Box Styling (shared by idle/error) */
601
+ .message-box {
602
+ background-color: rgba(0, 0, 0, 0.75); /* Darker, more opaque */
603
+ padding: 40px 50px; /* More padding */
604
+ border-radius: 20px; /* More rounded */
605
+ max-width: 70%; /* Max width */
606
+ box-shadow: 0 8px 25px rgba(0, 0, 0, 0.5); /* Softer shadow */
607
+ border: 1px solid rgba(255, 255, 255, 0.1); /* Subtle border */
608
+ }
609
+
610
+ .status {
611
+ font-family: 'Press Start 2P', cursive; /* Pixel font for status */
612
+ font-size: clamp(1.5em, 4vw, 2.5em); /* Responsive font size */
613
+ margin-bottom: 25px;
614
+ color: #ffcb05; /* Pokemon Yellow */
615
+ text-shadow: 3px 3px 0px #3b4cca; /* Pokemon Blue shadow */
616
+ /* Subtle pulse animation for idle status */
617
+ animation: pulse 2s infinite ease-in-out;
618
+ }
619
+
620
+ .instruction {
621
+ font-size: clamp(1em, 2.5vw, 1.4em); /* Responsive font size */
622
+ color: #f0f0f0; /* Light grey for readability */
623
+ line-height: 1.6;
624
+ text-shadow: 1px 1px 3px rgba(0, 0, 0, 0.7);
625
+ }
626
+ .instruction strong {
627
+ color: #ff7f0f; /* A contrasting color like orange */
628
+ font-weight: 700; /* Ensure Poppins bold is used */
629
+ }
630
+
631
+
632
+ /* Error Screen Specific Text Styling */
633
+ .error-container .status {
634
+ color: #ff4d4d; /* Bright Red for error status */
635
+ text-shadow: 2px 2px 0px #800000; /* Darker red shadow */
636
+ animation: none; /* No pulse on error */
637
+ }
638
+ .error-container .instruction {
639
+ color: #ffdddd; /* Lighter red for error details */
640
+ }
641
+
642
+
643
+ /* Pulse Animation */
644
+ @keyframes pulse {
645
+ 0% { transform: scale(1); }
646
+ 50% { transform: scale(1.03); }
647
+ 100% { transform: scale(1); }
648
  }
649
+
650
  </style>
651
  </head>
652
  <body>
653
+ <div id="stream-container">
654
+ </div>
655
+
656
  <script>
657
  const streamContainer = document.getElementById('stream-container');
658
+ let ws = null; // WebSocket instance
659
+
660
+ function connectWebSocket() {
661
+ // Use wss:// for https:// and ws:// for http://
662
+ const wsProtocol = location.protocol === 'https:' ? 'wss' : 'ws';
663
+ const wsUrl = `${wsProtocol}://${location.host}/ws`;
664
+ ws = new WebSocket(wsUrl);
665
+
666
+ console.log('Attempting to connect to WebSocket server...');
667
+
668
+ ws.onopen = (event) => {
669
+ console.log('WebSocket connection established.');
670
+ // Optional: Clear any 'connecting...' message if you have one
671
+ // streamContainer.innerHTML = ''; // Clear container only if needed
672
+ };
673
+
674
+ ws.onmessage = (event) => {
675
+ // console.log('Received update from server:', event.data);
676
+ // Directly set the innerHTML with the fragment received from the server
677
+ streamContainer.innerHTML = event.data;
678
+ };
679
+
680
+ ws.onclose = (event) => {
681
+ console.log(`WebSocket connection closed. Code: ${event.code}, Reason: ${event.reason}. Attempting to reconnect in 5 seconds...`);
682
+ ws = null; // Clear the instance
683
+ // Clear the display or show a 'disconnected' message
684
+ streamContainer.innerHTML = createReconnectMessage();
685
+ setTimeout(connectWebSocket, 5000); // Retry connection after 5 seconds
686
+ };
687
+
688
+ ws.onerror = (event) => {
689
+ console.error('WebSocket error:', event);
690
+ // The onclose event will usually fire after an error,
691
+ // so reconnection logic is handled there.
692
+ // You might want to display an error message here briefly.
693
+ streamContainer.innerHTML = createErrorMessage("WebSocket connection error. Attempting to reconnect...");
694
+ // Optionally force close to trigger reconnect logic if onclose doesn't fire
695
+ if (ws && ws.readyState !== WebSocket.CLOSED) {
696
+ ws.close();
697
+ }
698
+ };
699
+ }
700
+
701
+ // Helper function to generate reconnecting message HTML (matches error style)
702
+ function createReconnectMessage() {
703
+ return `
704
+ <div class="content-container error-container" style="background: #333;">
705
+ <div class="message-box" style="background-color: rgba(0,0,0,0.6);">
706
+ <p class="status" style="color: #ffcb05; text-shadow: none; animation: none;">πŸ”Œ Disconnected πŸ”Œ</p>
707
+ <p class="instruction" style="color: #eee;">Connection lost. Attempting to reconnect automatically...</p>
708
+ </div>
709
+ </div>`;
710
+ }
711
+
712
+ // Helper function to generate error message HTML
713
+ function createErrorMessage(message) {
714
+ return `
715
+ <div class="content-container error-container">
716
+ <div class="message-box">
717
+ <p class="status">🚨 Error 🚨</p>
718
+ <p class="instruction">${message}</p>
719
+ </div>
720
+ </div>`;
721
+ }
722
+
723
+
724
+ // Initial connection attempt when the page loads
725
+ connectWebSocket();
726
+
727
  </script>
728
  </body>
729
  </html>
 
734
  await manager.connect(websocket)
735
  try:
736
  while True:
737
+ # Keep connection alive. Client doesn't send messages in this setup.
738
+ # FastAPI's WebSocket implementation handles ping/pong internally usually.
739
+ # If needed, you could implement explicit keepalive here.
740
+ data = await websocket.receive_text()
741
+ # We don't expect messages from the client in this design,
742
+ # but log if received for debugging.
743
+ print(f"Received unexpected message from client: {data}")
744
+ # Or simply keep listening:
745
+ # await asyncio.sleep(60) # Example keepalive interval if needed
746
+ except WebSocketDisconnect as e:
747
+ print(f"WebSocket disconnected: Code {e.code}, Reason: {getattr(e, 'reason', 'N/A')}")
748
+ await manager.disconnect(websocket) # Use await here
749
  except Exception as e:
750
+ # Catch other potential errors on the connection
751
  print(f"WebSocket error: {e}")
752
+ traceback.print_exc()
753
+ await manager.disconnect(websocket) # Ensure disconnect on error
754
+
755
 
756
  @app.on_event("startup")
757
  async def startup_event():
758
  """Start background tasks when the application starts."""
759
  global background_task_handle
760
+
761
+ # Mount static files directory (make sure 'static' folder exists)
762
+ # Place your 'pokemon_huggingface.png' inside this 'static' folder
763
+ static_dir = "static"
764
+ if not os.path.exists(static_dir):
765
+ os.makedirs(static_dir)
766
+ print(f"Created static directory at: {os.path.abspath(static_dir)}")
767
+ print("!!! Please add 'pokemon_huggingface.png' to this directory! !!!")
768
+
769
+ app.mount("/static", StaticFiles(directory=static_dir), name="static")
770
+ print(f"Mounted static directory '{static_dir}' at '/static'")
771
+
772
  print("πŸš€ Starting background tasks")
773
+ # Start the main lifecycle manager task
774
+ background_task_handle = asyncio.create_task(manage_agent_lifecycle(), name="LifecycleManager")
775
+ # Add the exception logging callback
776
  background_task_handle.add_done_callback(log_task_exception)
777
  print("βœ… Background tasks started")
778
 
 
780
  async def shutdown_event():
781
  """Clean up tasks when shutting down."""
782
  global background_task_handle, active_agent_instance
783
+
784
+ print("\nπŸ”Œ Shutting down application. Cleaning up...")
785
+
786
+ # 1. Cancel the main lifecycle manager task
787
  if background_task_handle and not background_task_handle.done():
788
+ print("Cancelling background task...")
789
  background_task_handle.cancel()
790
  try:
791
  await asyncio.wait_for(background_task_handle, timeout=5.0)
792
+ print("Background task cancelled successfully.")
793
+ except asyncio.CancelledError:
794
+ print("Background task cancellation confirmed (CancelledError).")
795
+ except asyncio.TimeoutError:
796
+ print("Background task did not finish cancelling within timeout.")
797
+ except Exception as e:
798
+ print(f"Error during background task cancellation: {e}")
799
+
800
+ # 2. Deactivate and disconnect any currently active agent
801
+ # Use a copy of the instance in case it gets cleared elsewhere during shutdown.
802
+ agent_to_disconnect = active_agent_instance
803
+ if agent_to_disconnect:
804
+ agent_name = agent_to_disconnect.username if hasattr(agent_to_disconnect, 'username') else 'Unknown Agent'
805
+ print(f"Disconnecting active agent '{agent_name}'...")
806
  try:
807
+ # Check websocket status before disconnecting
808
+ if hasattr(agent_to_disconnect, '_websocket') and agent_to_disconnect._websocket and agent_to_disconnect._websocket.open:
809
+ await agent_to_disconnect.disconnect()
810
+ print(f"Agent '{agent_name}' disconnected.")
811
+ else:
812
+ print(f"Agent '{agent_name}' already disconnected or websocket not available.")
813
  except Exception as e:
814
+ print(f"Error during agent disconnect on shutdown for '{agent_name}': {e}")
815
+
816
+ # 3. Close all active WebSocket connections cleanly
817
+ print(f"Closing {len(manager.active_connections)} client WebSocket connections...")
818
+ # Create tasks to close all connections concurrently
819
+ close_tasks = [
820
+ conn.close(code=1000, reason="Server shutting down") # 1000 = Normal Closure
821
+ for conn in list(manager.active_connections) # Iterate over a copy
822
+ ]
823
+ if close_tasks:
824
+ await asyncio.gather(*close_tasks, return_exceptions=True) # Log potential errors during close
825
+
826
+ print("βœ… Cleanup complete. Application shutdown.")
827
+
828
 
829
  # For direct script execution
830
  if __name__ == "__main__":
831
  import uvicorn
832
+
833
+ # Configure logging
834
+ logging.basicConfig(
835
+ level=logging.INFO,
836
+ format='%(asctime)s - %(name)s:%(lineno)d - %(levelname)s - %(message)s',
837
+ datefmt='%Y-%m-%d %H:%M:%S'
838
+ )
839
+ # Reduce noise from poke_env's default INFO logging if desired
840
  logging.getLogger('poke_env').setLevel(logging.WARNING)
841
+ logging.getLogger('websockets.client').setLevel(logging.INFO) # Show websocket connection attempts
842
+
843
  print("Starting Pokemon Battle Livestream Server...")
844
+ print("="*60)
845
+
846
  if not AVAILABLE_AGENT_NAMES:
847
+ print("β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆ FATAL ERROR β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆ")
848
+ print(" No agents found with configured passwords!")
849
+ print(" Please set the required environment variables:")
850
  for name, cfg in AGENT_CONFIGS.items():
851
+ print(f" - {cfg.get('password_env_var', 'N/A')} (for agent: {name})")
852
+ print("="*60)
853
  exit("Exiting due to missing agent passwords.")
854
  else:
855
+ print("✨ Available Agents Found:")
856
  for name in AVAILABLE_AGENT_NAMES:
857
+ print(f" - {name}")
858
+ print("="*60)
859
+ print(f"Server will run on http://0.0.0.0:7860")
860
+ print("="*60)
861
+
862
  # Run with uvicorn
863
+ uvicorn.run(
864
+ "main:app", # Point to the FastAPI app instance
865
+ host="0.0.0.0",
866
+ port=7860,
867
+ reload=False, # Disable reload for production/stable testing
868
+ log_level="info" # Uvicorn's log level
869
+ )
870
+