Jofthomas commited on
Commit
cb75854
·
verified ·
1 Parent(s): 9f3492a

Update main.py

Browse files
Files changed (1) hide show
  1. main.py +228 -206
main.py CHANGED
@@ -6,14 +6,18 @@ import random
6
  import time
7
  import traceback
8
  import logging
9
- from typing import List, Dict, Optional, Set
 
 
 
 
10
 
11
  from fastapi import FastAPI, WebSocket, WebSocketDisconnect
12
  from fastapi.responses import HTMLResponse
13
  from fastapi.staticfiles import StaticFiles
14
 
15
  # --- Imports for poke_env and agents ---
16
- from poke_env.player import Player
17
  from poke_env import AccountConfiguration, ServerConfiguration
18
  from poke_env.environment.battle import Battle
19
 
@@ -50,10 +54,26 @@ active_agent_instance: Optional[Player] = None
50
  active_agent_task: Optional[asyncio.Task] = None
51
  current_battle_instance: Optional[Battle] = None
52
  background_task_handle: Optional[asyncio.Task] = None
 
 
 
53
 
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."""
@@ -77,11 +97,7 @@ def get_active_battle(agent: Player) -> Optional[Battle]:
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"
@@ -93,7 +109,6 @@ def create_battle_iframe(battle_id: str) -> str:
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">
@@ -105,7 +120,6 @@ def create_idle_html(status_message: str, instruction: str) -> str:
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">
@@ -115,9 +129,139 @@ def create_error_html(error_msg: str) -> str:
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
 
@@ -125,7 +269,13 @@ async def update_display_html(new_html_fragment: str) -> None:
125
  # --- Agent Lifecycle Management ---
126
  async def select_and_activate_new_agent():
127
  """Selects a random available agent, instantiates it, and starts its listening task."""
128
- global active_agent_name, active_agent_instance, active_agent_task
 
 
 
 
 
 
129
 
130
  if not AVAILABLE_AGENT_NAMES:
131
  print("Lifecycle: No available agents with passwords set.")
@@ -139,31 +289,29 @@ async def select_and_activate_new_agent():
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)
 
147
  agent = AgentClass(
148
  account_configuration=account_config,
149
  server_configuration=custom_config,
150
  battle_format=DEFAULT_BATTLE_FORMAT,
151
  log_level=logging.INFO,
152
- max_concurrent_battles=1
 
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
@@ -174,7 +322,6 @@ async def select_and_activate_new_agent():
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
180
  active_agent_task = None
@@ -182,38 +329,24 @@ async def select_and_activate_new_agent():
182
 
183
  async def check_for_new_battle():
184
  """Checks if the active agent has started a battle with a valid tag."""
185
- # --- FIX: Declare intention to use/modify global variables ---
186
  global active_agent_instance, current_battle_instance, active_agent_name, active_agent_task
187
- # -------------------------------------------------------------
188
 
189
  if active_agent_instance:
190
  battle = get_active_battle(active_agent_instance)
191
- # Check if battle exists AND has a valid battle_tag
192
  if battle and battle.battle_tag:
193
- # This line MODIFIES the global variable
194
  current_battle_instance = battle
195
  print(f"Lifecycle: Agent '{active_agent_name}' started battle: {battle.battle_tag}")
196
-
197
- # Stop the agent from listening for more challenges once a battle starts
198
  if active_agent_task and not active_agent_task.done():
199
  print(f"Lifecycle: Cancelling accept_challenges task for {active_agent_name} as battle started.")
200
  active_agent_task.cancel()
201
- # Optional: Wait briefly for cancellation confirmation, but don't block excessively
202
- # try:
203
- # await asyncio.wait_for(active_agent_task, timeout=0.5)
204
- # except (asyncio.CancelledError, asyncio.TimeoutError):
205
- # pass # Expected outcomes
206
- # else:
207
- # print(f"DEBUG: get_active_battle returned None or battle without tag.")
208
 
209
  async def deactivate_current_agent(reason: str = "cycle"):
210
  """Cleans up the currently active agent and resets state."""
211
  global active_agent_name, active_agent_instance, active_agent_task, current_battle_instance
212
 
213
- agent_name_to_deactivate = active_agent_name # Store before clearing
214
  print(f"Lifecycle: Deactivating agent '{agent_name_to_deactivate}' (Reason: {reason})...")
215
 
216
- # Display appropriate intermediate message
217
  if reason == "battle_end":
218
  await update_display_html(create_idle_html("Battle Finished!", f"Agent <strong>{agent_name_to_deactivate}</strong> completed the match."))
219
  elif reason == "cycle":
@@ -223,35 +356,22 @@ async def deactivate_current_agent(reason: str = "cycle"):
223
  else: # Generic reason or error
224
  await update_display_html(create_idle_html(f"Resetting Agent ({reason})", f"Cleaning up <strong>{agent_name_to_deactivate}</strong>..."))
225
 
226
- # Give users a moment to see the intermediate message
227
- await asyncio.sleep(3) # Adjust duration as needed
228
-
229
- # Show the "preparing next agent" message before lengthy cleanup
230
  await update_display_html(create_idle_html("Preparing Next Agent...", "Please wait..."))
231
 
232
-
233
  agent = active_agent_instance
234
  task = active_agent_task
235
 
236
- # Store a local copy of the battle instance before clearing it
237
- # last_battle_instance = current_battle_instance # Not strictly needed now
238
-
239
- # --- Crucial: Clear global state variables FIRST ---
240
- # This prevents race conditions where the lifecycle loop might try to
241
- # access the agent while it's being deactivated.
242
  active_agent_name = None
243
  active_agent_instance = None
244
  active_agent_task = None
245
  current_battle_instance = None
246
  print(f"Lifecycle: Global state cleared for '{agent_name_to_deactivate}'.")
247
 
248
- # --- Now perform cleanup actions ---
249
- # Cancel the accept_challenges task if it's still running (it might already be done/cancelled)
250
  if task and not task.done():
251
  print(f"Lifecycle: Ensuring task cancellation for {agent_name_to_deactivate} ({task.get_name()})...")
252
  task.cancel()
253
  try:
254
- # Wait briefly for the task to acknowledge cancellation
255
  await asyncio.wait_for(task, timeout=2.0)
256
  print(f"Lifecycle: Task cancellation confirmed for {agent_name_to_deactivate}.")
257
  except asyncio.CancelledError:
@@ -259,39 +379,32 @@ async def deactivate_current_agent(reason: str = "cycle"):
259
  except asyncio.TimeoutError:
260
  print(f"Lifecycle: Task did not confirm cancellation within timeout for {agent_name_to_deactivate}.")
261
  except Exception as e:
262
- # Catch other potential errors during task cleanup
263
  print(f"Lifecycle: Error during task cancellation wait for {agent_name_to_deactivate}: {e}")
264
 
265
- # Disconnect the player (ensure agent object exists)
266
  if agent:
267
  print(f"Lifecycle: Disconnecting player {agent.username}...")
268
  try:
269
- # Check websocket state before attempting disconnection
270
  if hasattr(agent, '_websocket') and agent._websocket and agent._websocket.open:
271
  await agent.disconnect()
272
  print(f"Lifecycle: Player {agent.username} disconnected successfully.")
273
  else:
274
  print(f"Lifecycle: Player {agent.username} already disconnected or websocket not available.")
275
  except Exception as e:
276
- # Log errors during disconnection but don't halt the process
277
  print(f"ERROR during agent disconnect ({agent.username}): {e}")
278
- traceback.print_exc() # Log full traceback for debugging
279
 
280
- # Add a brief delay AFTER deactivation before the loop potentially selects a new agent
281
- await asyncio.sleep(2) # Reduced from 3, adjust as needed
282
  print(f"Lifecycle: Agent '{agent_name_to_deactivate}' deactivation complete.")
283
 
284
  async def manage_agent_lifecycle():
285
  """Runs the main loop selecting, running, and cleaning up agents sequentially."""
286
- # --- FIX: Declare intention to use global variables ---
287
  global active_agent_name, active_agent_instance, active_agent_task, current_battle_instance
288
- # ------------------------------------------------------
289
 
290
  print("Background lifecycle manager started.")
291
- REFRESH_INTERVAL_SECONDS = 3 # How often to check state when idle/in battle
292
- LOOP_COOLDOWN_SECONDS = 1 # Small delay at end of loop if no other waits occurred
293
- ERROR_RETRY_DELAY_SECONDS = 10 # Longer delay after errors
294
- POST_BATTLE_DELAY_SECONDS = 5 # Delay after a battle finishes before selecting next agent
295
 
296
  loop_counter = 0
297
 
@@ -301,10 +414,6 @@ async def manage_agent_lifecycle():
301
  print(f"\n--- Lifecycle Check #{loop_counter} [{time.strftime('%H:%M:%S')}] ---")
302
 
303
  try:
304
- # ==================================
305
- # State 1: No agent active
306
- # ==================================
307
- # Now Python knows active_agent_instance refers to the global one
308
  if active_agent_instance is None:
309
  print(f"[{loop_counter}] State 1: No active agent. Selecting...")
310
  activated = await select_and_activate_new_agent()
@@ -312,82 +421,52 @@ async def manage_agent_lifecycle():
312
  print(f"[{loop_counter}] State 1: Activation failed. Waiting {ERROR_RETRY_DELAY_SECONDS}s before retry.")
313
  await asyncio.sleep(ERROR_RETRY_DELAY_SECONDS)
314
  else:
315
- # Now Python knows active_agent_name refers to the global one set by select_and_activate_new_agent
316
  print(f"[{loop_counter}] State 1: Agent '{active_agent_name}' activated successfully.")
317
- # No sleep here, proceed to next check immediately if needed
318
 
319
- # ==================================
320
- # State 2: Agent is active
321
- # ==================================
322
- else:
323
- # Now Python knows active_agent_name refers to the global one
324
- agent_name = active_agent_name # Cache for logging
325
  print(f"[{loop_counter}] State 2: Agent '{agent_name}' is active.")
326
 
327
- # --- Sub-state: Check for new battle if none is tracked ---
328
- # Now Python knows current_battle_instance refers to the global one
329
  if current_battle_instance is None:
330
  print(f"[{loop_counter}] State 2a: Checking for new battle for '{agent_name}'...")
331
- await check_for_new_battle() # This updates global current_battle_instance if found
332
 
333
- # Now Python knows current_battle_instance refers to the global one
334
  if current_battle_instance:
335
  battle_tag = current_battle_instance.battle_tag
336
  print(f"[{loop_counter}] State 2a: *** NEW BATTLE DETECTED: {battle_tag} for '{agent_name}' ***")
337
-
338
- # Check for non-public/suffixed format (heuristic: more than 3 parts, 3rd part is number)
339
  parts = battle_tag.split('-')
340
  is_suffixed_format = len(parts) > 3 and parts[2].isdigit()
341
 
342
  if is_suffixed_format:
343
- # Forfeit immediately if it looks like a private/suffixed battle ID
344
  print(f"[{loop_counter}] Detected potentially non-public battle format ({battle_tag}). Forfeiting.")
345
- # Don't update display yet, do it before deactivation
346
  try:
347
- # Now Python knows active_agent_instance refers to the global one
348
- if active_agent_instance: # Ensure agent still exists
349
  await active_agent_instance.forfeit(battle_tag)
350
- # await active_agent_instance.send_message("/forfeit", battle_tag) # Alternative
351
  print(f"[{loop_counter}] Sent forfeit command for {battle_tag}.")
352
- await asyncio.sleep(1.5) # Give forfeit time to register
353
  except Exception as forfeit_err:
354
  print(f"[{loop_counter}] ERROR sending forfeit for {battle_tag}: {forfeit_err}")
355
- # Deactivate agent after forfeit attempt
356
  await deactivate_current_agent(reason="forfeited_private_battle")
357
- continue # Skip rest of the loop for this iteration
358
-
359
  else:
360
- # Public battle format - display the iframe
361
  print(f"[{loop_counter}] Public battle format detected. Displaying battle {battle_tag}.")
362
  await update_display_html(create_battle_iframe(battle_tag))
363
- # Now fall through to monitor this battle in the next section
364
-
365
  else:
366
- # No new battle found, agent remains idle
367
  print(f"[{loop_counter}] State 2a: No new battle found. Agent '{agent_name}' remains idle, waiting for challenge.")
368
- # Periodically refresh idle screen to ensure consistency
369
  idle_html = create_idle_html(f"Agent Ready: <strong>{agent_name}</strong>",
370
  f"Please challenge <strong>{agent_name}</strong> to a <strong>{DEFAULT_BATTLE_FORMAT}</strong> battle.")
371
  await update_display_html(idle_html)
372
- await asyncio.sleep(REFRESH_INTERVAL_SECONDS) # Wait before next check if idle
373
-
374
 
375
- # --- Sub-state: Monitor ongoing battle ---
376
- # Now Python knows current_battle_instance refers to the global one
377
- if current_battle_instance is not None:
378
  battle_tag = current_battle_instance.battle_tag
379
  print(f"[{loop_counter}] State 2b: Monitoring battle {battle_tag} for '{agent_name}'")
380
 
381
- # Ensure agent instance still exists before accessing its battles
382
- # Now Python knows active_agent_instance refers to the global one
383
  if not active_agent_instance:
384
  print(f"[{loop_counter}] WARNING: Agent instance for '{agent_name}' disappeared while monitoring battle {battle_tag}! Deactivating.")
385
  await deactivate_current_agent(reason="agent_disappeared_mid_battle")
386
  continue
387
 
388
- # Get potentially updated battle object directly from agent's state
389
- # Use .get() for safety
390
- # Now Python knows active_agent_instance refers to the global one
391
  battle_obj = active_agent_instance._battles.get(battle_tag)
392
 
393
  if battle_obj and battle_obj.finished:
@@ -395,138 +474,109 @@ async def manage_agent_lifecycle():
395
  await deactivate_current_agent(reason="battle_end")
396
  print(f"[{loop_counter}] Waiting {POST_BATTLE_DELAY_SECONDS}s post-battle before selecting next agent.")
397
  await asyncio.sleep(POST_BATTLE_DELAY_SECONDS)
398
- continue # Start next loop iteration to select new agent
399
-
400
  elif not battle_obj:
401
- # This can happen briefly during transitions or if battle ends unexpectedly
402
  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.")
403
  await deactivate_current_agent(reason="battle_object_missing")
404
  continue
405
-
406
  else:
407
- # Battle is ongoing, battle object exists, iframe should be displayed
408
  print(f"[{loop_counter}] Battle {battle_tag} ongoing for '{agent_name}'.")
409
- # Optionally: Could re-send iframe HTML periodically if needed, but usually not necessary
410
- # await update_display_html(create_battle_iframe(battle_tag))
411
- await asyncio.sleep(REFRESH_INTERVAL_SECONDS) # Wait before next check
412
 
413
- # --- Global Exception Handling for the main loop ---
414
  except asyncio.CancelledError:
415
  print("Lifecycle manager task cancelled.")
416
- raise # Re-raise to ensure proper shutdown
417
  except Exception as e:
418
  print(f"!!! ERROR in main lifecycle loop #{loop_counter}: {e} !!!")
419
  traceback.print_exc()
420
- # Now Python knows active_agent_name refers to the global one
421
- current_agent_name = active_agent_name # Cache name before deactivation attempts
422
- # Now Python knows active_agent_instance refers to the global one
423
  if active_agent_instance:
424
- print(f"Attempting to deactivate agent '{current_agent_name}' due to loop error...")
425
  try:
426
  await deactivate_current_agent(reason="main_loop_error")
427
  except Exception as deactivation_err:
428
  print(f"Error during error-handling deactivation: {deactivation_err}")
429
- # Ensure state is cleared even if deactivation fails partially
430
  active_agent_name = None
431
  active_agent_instance = None
432
  active_agent_task = None
433
  current_battle_instance = None
434
  else:
435
- # Error happened potentially before agent activation or after clean deactivation
436
  print("No active agent instance during loop error.")
437
- # Show a generic error on the frontend
438
  await update_display_html(create_error_html(f"A server error occurred in the lifecycle manager. Please wait. ({e})"))
439
-
440
- # Wait longer after a major error before trying again
441
  print(f"Waiting {ERROR_RETRY_DELAY_SECONDS}s after loop error.")
442
  await asyncio.sleep(ERROR_RETRY_DELAY_SECONDS)
443
- continue # Go to next loop iteration after error handling
444
 
445
- # --- Delay at end of loop if no other significant waits happened ---
446
  elapsed_time = time.monotonic() - loop_start_time
447
  if elapsed_time < LOOP_COOLDOWN_SECONDS:
448
  await asyncio.sleep(LOOP_COOLDOWN_SECONDS - elapsed_time)
449
 
 
450
  def log_task_exception(task: asyncio.Task):
451
  """Callback to log exceptions from background tasks (like accept_challenges)."""
452
  try:
453
  if task.cancelled():
454
- # Don't log cancellation as an error, it's often expected
455
  print(f"Task '{task.get_name()}' was cancelled.")
456
  return
457
- # Accessing result will raise exception if task failed
458
  task.result()
459
  print(f"Task '{task.get_name()}' completed successfully.")
460
  except asyncio.CancelledError:
461
  print(f"Task '{task.get_name()}' confirmed cancelled (exception caught).")
462
- pass # Expected
463
  except Exception as e:
464
- # Log actual errors
465
  print(f"!!! Exception in background task '{task.get_name()}': {e} !!!")
466
  traceback.print_exc()
467
- # Optionally: Trigger some recovery or notification here if needed
468
-
469
 
470
  # --- WebSocket connection manager ---
471
  class ConnectionManager:
472
  def __init__(self):
473
  self.active_connections: Set[WebSocket] = set()
474
- # Initialize with the idle HTML fragment
475
  self.current_html_fragment: str = create_idle_html("Initializing...", "Setting up Pokémon Battle Stream")
476
 
477
  async def connect(self, websocket: WebSocket):
478
  await websocket.accept()
479
  self.active_connections.add(websocket)
480
  print(f"Client connected. Sending current state. Total clients: {len(self.active_connections)}")
481
- # Send current state (HTML fragment) to newly connected client
482
  try:
483
  await websocket.send_text(self.current_html_fragment)
484
  except Exception as e:
485
  print(f"Error sending initial state to new client: {e}")
486
- # Consider removing the connection if initial send fails
487
  await self.disconnect(websocket)
488
 
489
-
490
  async def disconnect(self, websocket: WebSocket):
491
- # Use discard() to safely remove even if not present
492
  self.active_connections.discard(websocket)
493
  print(f"Client disconnected. Total clients: {len(self.active_connections)}")
494
 
495
  async def update_all(self, html_fragment: str):
496
  """Update the current HTML fragment and broadcast to all clients."""
497
  if self.current_html_fragment == html_fragment:
498
- # print("Skipping broadcast, HTML fragment unchanged.")
499
- return # Avoid unnecessary updates if content is identical
500
 
501
  self.current_html_fragment = html_fragment
502
  if not self.active_connections:
503
- # print("No active connections to broadcast update to.")
504
  return
505
 
506
  print(f"Broadcasting update to {len(self.active_connections)} clients...")
507
-
508
- # Create a list of tasks to send updates concurrently
509
- # Make a copy of the set for safe iteration during potential disconnects
510
  send_tasks = [
511
  connection.send_text(html_fragment)
512
- for connection in list(self.active_connections) # Iterate over a copy
513
  ]
514
-
515
- # Use asyncio.gather to send to all clients, collecting results/exceptions
516
  results = await asyncio.gather(*send_tasks, return_exceptions=True)
517
-
518
- # Handle potential errors during broadcast (e.g., client disconnected abruptly)
519
- # Iterate over connections again, checking results
520
  connections_to_remove = set()
 
 
521
  for i, result in enumerate(results):
522
- connection = list(self.active_connections)[i] # Assumes order is maintained
523
- if isinstance(result, Exception):
524
- print(f"Error sending update to client: {result}. Marking for removal.")
525
- connections_to_remove.add(connection)
 
 
526
 
527
- # Disconnect clients that failed
528
  for connection in connections_to_remove:
529
- await self.disconnect(connection)
 
 
530
 
531
 
532
  manager = ConnectionManager()
@@ -535,8 +585,7 @@ manager = ConnectionManager()
535
  @app.get("/", response_class=HTMLResponse)
536
  async def get_homepage():
537
  """Serves the main HTML page with WebSocket connection and improved styling."""
538
- # NOTE: Ensure the static path '/static/pokemon_huggingface.png' is correct
539
- # and the image exists in a 'static' folder next to your main.py
540
  return """
541
  <!DOCTYPE html>
542
  <html lang="en">
@@ -671,7 +720,6 @@ async def get_homepage():
671
  let ws = null; // WebSocket instance
672
 
673
  function connectWebSocket() {
674
- // Use wss:// for https:// and ws:// for http://
675
  const wsProtocol = location.protocol === 'https:' ? 'wss' : 'ws';
676
  const wsUrl = `${wsProtocol}://${location.host}/ws`;
677
  ws = new WebSocket(wsUrl);
@@ -680,38 +728,28 @@ async def get_homepage():
680
 
681
  ws.onopen = (event) => {
682
  console.log('WebSocket connection established.');
683
- // Optional: Clear any 'connecting...' message if you have one
684
- // streamContainer.innerHTML = ''; // Clear container only if needed
685
  };
686
 
687
  ws.onmessage = (event) => {
688
- // console.log('Received update from server:', event.data);
689
- // Directly set the innerHTML with the fragment received from the server
690
  streamContainer.innerHTML = event.data;
691
  };
692
 
693
  ws.onclose = (event) => {
694
  console.log(`WebSocket connection closed. Code: ${event.code}, Reason: ${event.reason}. Attempting to reconnect in 5 seconds...`);
695
- ws = null; // Clear the instance
696
- // Clear the display or show a 'disconnected' message
697
  streamContainer.innerHTML = createReconnectMessage();
698
- setTimeout(connectWebSocket, 5000); // Retry connection after 5 seconds
699
  };
700
 
701
  ws.onerror = (event) => {
702
  console.error('WebSocket error:', event);
703
- // The onclose event will usually fire after an error,
704
- // so reconnection logic is handled there.
705
- // You might want to display an error message here briefly.
706
  streamContainer.innerHTML = createErrorMessage("WebSocket connection error. Attempting to reconnect...");
707
- // Optionally force close to trigger reconnect logic if onclose doesn't fire
708
  if (ws && ws.readyState !== WebSocket.CLOSED) {
709
  ws.close();
710
  }
711
  };
712
  }
713
 
714
- // Helper function to generate reconnecting message HTML (matches error style)
715
  function createReconnectMessage() {
716
  return `
717
  <div class="content-container error-container" style="background: #333;">
@@ -722,7 +760,6 @@ async def get_homepage():
722
  </div>`;
723
  }
724
 
725
- // Helper function to generate error message HTML
726
  function createErrorMessage(message) {
727
  return `
728
  <div class="content-container error-container">
@@ -733,10 +770,7 @@ async def get_homepage():
733
  </div>`;
734
  }
735
 
736
-
737
- // Initial connection attempt when the page loads
738
  connectWebSocket();
739
-
740
  </script>
741
  </body>
742
  </html>
@@ -747,32 +781,31 @@ async def websocket_endpoint(websocket: WebSocket):
747
  await manager.connect(websocket)
748
  try:
749
  while True:
750
- # Keep connection alive. Client doesn't send messages in this setup.
751
- # FastAPI's WebSocket implementation handles ping/pong internally usually.
752
- # If needed, you could implement explicit keepalive here.
753
  data = await websocket.receive_text()
754
- # We don't expect messages from the client in this design,
755
- # but log if received for debugging.
756
  print(f"Received unexpected message from client: {data}")
757
- # Or simply keep listening:
758
- # await asyncio.sleep(60) # Example keepalive interval if needed
759
  except WebSocketDisconnect as e:
760
  print(f"WebSocket disconnected: Code {e.code}, Reason: {getattr(e, 'reason', 'N/A')}")
761
- await manager.disconnect(websocket) # Use await here
762
  except Exception as e:
763
- # Catch other potential errors on the connection
764
  print(f"WebSocket error: {e}")
765
  traceback.print_exc()
766
- await manager.disconnect(websocket) # Ensure disconnect on error
767
-
768
-
 
 
 
 
 
 
 
 
 
769
  @app.on_event("startup")
770
  async def startup_event():
771
  """Start background tasks when the application starts."""
772
  global background_task_handle
773
 
774
- # Mount static files directory (make sure 'static' folder exists)
775
- # Place your 'pokemon_huggingface.png' inside this 'static' folder
776
  static_dir = "static"
777
  if not os.path.exists(static_dir):
778
  os.makedirs(static_dir)
@@ -783,9 +816,7 @@ async def startup_event():
783
  print(f"Mounted static directory '{static_dir}' at '/static'")
784
 
785
  print("🚀 Starting background tasks")
786
- # Start the main lifecycle manager task
787
  background_task_handle = asyncio.create_task(manage_agent_lifecycle(), name="LifecycleManager")
788
- # Add the exception logging callback
789
  background_task_handle.add_done_callback(log_task_exception)
790
  print("✅ Background tasks started")
791
 
@@ -796,7 +827,6 @@ async def shutdown_event():
796
 
797
  print("\n🔌 Shutting down application. Cleaning up...")
798
 
799
- # 1. Cancel the main lifecycle manager task
800
  if background_task_handle and not background_task_handle.done():
801
  print("Cancelling background task...")
802
  background_task_handle.cancel()
@@ -810,14 +840,11 @@ async def shutdown_event():
810
  except Exception as e:
811
  print(f"Error during background task cancellation: {e}")
812
 
813
- # 2. Deactivate and disconnect any currently active agent
814
- # Use a copy of the instance in case it gets cleared elsewhere during shutdown.
815
  agent_to_disconnect = active_agent_instance
816
  if agent_to_disconnect:
817
  agent_name = agent_to_disconnect.username if hasattr(agent_to_disconnect, 'username') else 'Unknown Agent'
818
  print(f"Disconnecting active agent '{agent_name}'...")
819
  try:
820
- # Check websocket status before disconnecting
821
  if hasattr(agent_to_disconnect, '_websocket') and agent_to_disconnect._websocket and agent_to_disconnect._websocket.open:
822
  await agent_to_disconnect.disconnect()
823
  print(f"Agent '{agent_name}' disconnected.")
@@ -826,32 +853,28 @@ async def shutdown_event():
826
  except Exception as e:
827
  print(f"Error during agent disconnect on shutdown for '{agent_name}': {e}")
828
 
829
- # 3. Close all active WebSocket connections cleanly
830
  print(f"Closing {len(manager.active_connections)} client WebSocket connections...")
831
- # Create tasks to close all connections concurrently
832
  close_tasks = [
833
- conn.close(code=1000, reason="Server shutting down") # 1000 = Normal Closure
834
- for conn in list(manager.active_connections) # Iterate over a copy
835
  ]
836
  if close_tasks:
837
- await asyncio.gather(*close_tasks, return_exceptions=True) # Log potential errors during close
838
 
839
  print("✅ Cleanup complete. Application shutdown.")
840
 
841
 
842
- # For direct script execution
843
  if __name__ == "__main__":
844
  import uvicorn
845
 
846
- # Configure logging
847
  logging.basicConfig(
848
  level=logging.INFO,
849
  format='%(asctime)s - %(name)s:%(lineno)d - %(levelname)s - %(message)s',
850
  datefmt='%Y-%m-%d %H:%M:%S'
851
  )
852
- # Reduce noise from poke_env's default INFO logging if desired
853
  logging.getLogger('poke_env').setLevel(logging.WARNING)
854
- logging.getLogger('websockets.client').setLevel(logging.INFO) # Show websocket connection attempts
855
 
856
  print("Starting Pokemon Battle Livestream Server...")
857
  print("="*60)
@@ -870,14 +893,13 @@ if __name__ == "__main__":
870
  print(f" - {name}")
871
  print("="*60)
872
  print(f"Server will run on http://0.0.0.0:7860")
 
873
  print("="*60)
874
 
875
- # Run with uvicorn
876
  uvicorn.run(
877
- "main:app", # Point to the FastAPI app instance
878
  host="0.0.0.0",
879
  port=7860,
880
- reload=False, # Disable reload for production/stable testing
881
- log_level="info" # Uvicorn's log level
882
- )
883
-
 
6
  import time
7
  import traceback
8
  import logging
9
+ # --- Additions for last_action route ---
10
+ import datetime
11
+ import html
12
+ from typing import List, Dict, Optional, Set, Callable # Added Callable
13
+ # ---------------------------------------
14
 
15
  from fastapi import FastAPI, WebSocket, WebSocketDisconnect
16
  from fastapi.responses import HTMLResponse
17
  from fastapi.staticfiles import StaticFiles
18
 
19
  # --- Imports for poke_env and agents ---
20
+ from poke_env.player import Player, ActionType, ForfeitAction, MoveOrder, SwitchOrder, DefaultOrder # Import action types
21
  from poke_env import AccountConfiguration, ServerConfiguration
22
  from poke_env.environment.battle import Battle
23
 
 
54
  active_agent_task: Optional[asyncio.Task] = None
55
  current_battle_instance: Optional[Battle] = None
56
  background_task_handle: Optional[asyncio.Task] = None
57
+ # --- NEW: Global variable for last action ---
58
+ last_llm_action: Optional[Dict] = None
59
+ # --------------------------------------------
60
 
61
  # --- Create FastAPI app ---
62
  app = FastAPI(title="Pokemon Battle Livestream")
63
 
64
+
65
+ # --- NEW: Callback function for agents ---
66
+ def update_last_action_callback(action_info: Dict):
67
+ """Callback for agents to report their chosen action."""
68
+ global last_llm_action
69
+ # Add a timestamp for context
70
+ action_info["timestamp_utc"] = datetime.datetime.utcnow().strftime('%Y-%m-%d %H:%M:%S UTC')
71
+ last_llm_action = action_info
72
+ # Optional: Log the update for debugging
73
+ print(f"ACTION_LOG: Agent '{action_info.get('agent', 'Unknown')}' chose action: {action_info.get('action_str', 'N/A')} (Turn: {action_info.get('turn', '?')})")
74
+ # ---------------------------------------
75
+
76
+
77
  # --- Helper Functions ---
78
  def get_active_battle(agent: Player) -> Optional[Battle]:
79
  """Returns the first non-finished battle for an agent."""
 
97
  def create_battle_iframe(battle_id: str) -> str:
98
  """Creates JUST the HTML for the battle iframe tag."""
99
  print("Creating iframe content for battle ID: ", battle_id)
 
 
100
  battle_url = f"https://jofthomas.com/play.pokemonshowdown.com/testclient.html#{battle_id}" # Using your custom URL
 
 
101
  return f"""
102
  <iframe
103
  id="battle-iframe"
 
109
 
110
  def create_idle_html(status_message: str, instruction: str) -> str:
111
  """Creates a visually appealing idle screen HTML fragment."""
 
112
  return f"""
113
  <div class="content-container idle-container">
114
  <div class="message-box">
 
120
 
121
  def create_error_html(error_msg: str) -> str:
122
  """Creates HTML fragment to display an error message."""
 
123
  return f"""
124
  <div class="content-container error-container">
125
  <div class="message-box">
 
129
  </div>
130
  """
131
 
132
+ # --- NEW: Helper function to create HTML for the last action page ---
133
+ def create_last_action_html(action_data: Optional[Dict]) -> str:
134
+ """Formats the last action data into an HTML page."""
135
+ if not action_data:
136
+ content = """
137
+ <div class="message-box">
138
+ <p class="status">No Action Recorded Yet</p>
139
+ <p class="instruction">Waiting for the agent to make its first move...</p>
140
+ </div>
141
+ """
142
+ else:
143
+ # Escape HTML characters in potentially user-generated content
144
+ raw_output_escaped = html.escape(action_data.get("raw_llm_output", "N/A"))
145
+ agent_name_escaped = html.escape(action_data.get('agent', 'Unknown'))
146
+ action_type_escaped = html.escape(action_data.get('action_type', 'N/A'))
147
+ action_str_escaped = html.escape(action_data.get('action_str', 'N/A'))
148
+
149
+ content = f"""
150
+ <div class="action-details">
151
+ <h2>Last Action by Agent: <span class="agent-name">{agent_name_escaped}</span></h2>
152
+ <p><strong>Timestamp:</strong> {action_data.get('timestamp_utc', 'N/A')}</p>
153
+ <p><strong>Battle Turn:</strong> {action_data.get('turn', '?')}</p>
154
+ <p><strong>Action Type:</strong> {action_type_escaped}</p>
155
+ <p><strong>Action Chosen:</strong> <span class="action-chosen">{action_str_escaped}</span></p>
156
+ <div class="raw-output">
157
+ <h3>Raw LLM Output (if available):</h3>
158
+ <pre>{raw_output_escaped}</pre>
159
+ </div>
160
+ </div>
161
+ """
162
+
163
+ # Basic HTML structure with auto-refresh and styling
164
+ return f"""
165
+ <!DOCTYPE html>
166
+ <html lang="en">
167
+ <head>
168
+ <meta charset="UTF-8">
169
+ <meta name="viewport" content="width=device-width, initial-scale=1.0">
170
+ <meta http-equiv="refresh" content="5">
171
+ <title>Last Agent Action</title>
172
+ <link rel="preconnect" href="https://fonts.googleapis.com">
173
+ <link rel="preconnect" href="https://fonts.gstatic.com" crossorigin>
174
+ <link href="https://fonts.googleapis.com/css2?family=Poppins:wght@400;700&family=Press+Start+2P&display=swap" rel="stylesheet">
175
+ <style>
176
+ body {{
177
+ font-family: 'Poppins', sans-serif;
178
+ background-color: #2c2f33;
179
+ color: #ffffff;
180
+ margin: 0;
181
+ padding: 20px;
182
+ line-height: 1.6;
183
+ }}
184
+ .container {{
185
+ max-width: 900px;
186
+ margin: 20px auto;
187
+ background-color: #3e4147;
188
+ padding: 25px;
189
+ border-radius: 10px;
190
+ box-shadow: 0 4px 15px rgba(0, 0, 0, 0.3);
191
+ }}
192
+ h1, h2 {{
193
+ font-family: 'Press Start 2P', cursive;
194
+ color: #ffcb05; /* Pokemon Yellow */
195
+ text-shadow: 2px 2px 0px #3b4cca; /* Pokemon Blue shadow */
196
+ margin-bottom: 20px;
197
+ text-align: center;
198
+ }}
199
+ .action-details p {{
200
+ margin: 10px 0;
201
+ font-size: 1.1em;
202
+ }}
203
+ .action-details strong {{
204
+ color: #ffcb05; /* Pokemon Yellow */
205
+ }}
206
+ .agent-name {{
207
+ color: #ff7f0f; /* Orange */
208
+ font-weight: bold;
209
+ }}
210
+ .action-chosen {{
211
+ font-weight: bold;
212
+ color: #76d7c4; /* Teal */
213
+ font-size: 1.2em;
214
+ }}
215
+ .raw-output {{
216
+ margin-top: 25px;
217
+ border-top: 1px solid #555;
218
+ padding-top: 15px;
219
+ }}
220
+ .raw-output h3 {{
221
+ margin-bottom: 10px;
222
+ color: #f0f0f0;
223
+ }}
224
+ pre {{
225
+ background-color: #23272a;
226
+ color: #dcdcdc;
227
+ padding: 15px;
228
+ border-radius: 5px;
229
+ white-space: pre-wrap; /* Wrap long lines */
230
+ word-wrap: break-word; /* Break words if necessary */
231
+ font-family: monospace;
232
+ font-size: 0.95em;
233
+ max-height: 300px; /* Limit height */
234
+ overflow-y: auto; /* Add scrollbar if needed */
235
+ }}
236
+ /* Styling for the 'No Action' message box */
237
+ .message-box {{
238
+ text-align: center;
239
+ padding: 30px;
240
+ }}
241
+ .message-box .status {{
242
+ font-family: 'Press Start 2P', cursive;
243
+ font-size: 1.8em;
244
+ color: #ff7f0f; /* Orange */
245
+ margin-bottom: 15px;
246
+ }}
247
+ .message-box .instruction {{
248
+ font-size: 1.1em;
249
+ color: #cccccc;
250
+ }}
251
+ </style>
252
+ </head>
253
+ <body>
254
+ <div class="container">
255
+ <h1>Agent Action Log</h1>
256
+ {content}
257
+ </div>
258
+ </body>
259
+ </html>
260
+ """
261
+ # --------------------------------------------------------------------
262
+
263
  async def update_display_html(new_html_fragment: str) -> None:
264
  """Updates the current display HTML fragment and broadcasts to all clients."""
 
265
  await manager.update_all(new_html_fragment)
266
  print("HTML Display FRAGMENT UPDATED and broadcasted.")
267
 
 
269
  # --- Agent Lifecycle Management ---
270
  async def select_and_activate_new_agent():
271
  """Selects a random available agent, instantiates it, and starts its listening task."""
272
+ # --- MODIFIED: Make sure globals are declared ---
273
+ global active_agent_name, active_agent_instance, active_agent_task, last_llm_action
274
+ # -----------------------------------------------
275
+
276
+ # --- MODIFIED: Reset last action when selecting a new agent ---
277
+ last_llm_action = None
278
+ # -----------------------------------------------------------
279
 
280
  if not AVAILABLE_AGENT_NAMES:
281
  print("Lifecycle: No available agents with passwords set.")
 
289
  agent_password = os.environ.get(password_env_var)
290
 
291
  print(f"Lifecycle: Activating agent '{selected_name}'...")
 
292
  await update_display_html(create_idle_html("Selecting Next Agent...", f"Preparing <strong>{selected_name}</strong>..."))
293
 
294
  try:
295
  account_config = AccountConfiguration(selected_name, agent_password)
296
+ # --- MODIFIED: Pass the action callback to the agent constructor ---
297
  agent = AgentClass(
298
  account_configuration=account_config,
299
  server_configuration=custom_config,
300
  battle_format=DEFAULT_BATTLE_FORMAT,
301
  log_level=logging.INFO,
302
+ max_concurrent_battles=1,
303
+ action_callback=update_last_action_callback # Pass the callback function
304
  )
305
+ # -------------------------------------------------------------------
306
 
 
 
307
  task = asyncio.create_task(agent.accept_challenges(None, 1), name=f"AcceptChallenge_{selected_name}")
308
+ task.add_done_callback(log_task_exception)
309
 
 
310
  active_agent_name = selected_name
311
  active_agent_instance = agent
312
  active_agent_task = task
313
 
314
  print(f"Lifecycle: Agent '{selected_name}' is active and listening for 1 challenge.")
 
315
  await update_display_html(create_idle_html(f"Agent Ready: <strong>{selected_name}</strong>",
316
  f"Please challenge <strong>{selected_name}</strong> to a <strong>{DEFAULT_BATTLE_FORMAT}</strong> battle."))
317
  return True
 
322
  traceback.print_exc()
323
  await update_display_html(create_error_html(f"Error activating {selected_name}. Please wait or check logs."))
324
 
 
325
  active_agent_name = None
326
  active_agent_instance = None
327
  active_agent_task = None
 
329
 
330
  async def check_for_new_battle():
331
  """Checks if the active agent has started a battle with a valid tag."""
 
332
  global active_agent_instance, current_battle_instance, active_agent_name, active_agent_task
 
333
 
334
  if active_agent_instance:
335
  battle = get_active_battle(active_agent_instance)
 
336
  if battle and battle.battle_tag:
 
337
  current_battle_instance = battle
338
  print(f"Lifecycle: Agent '{active_agent_name}' started battle: {battle.battle_tag}")
 
 
339
  if active_agent_task and not active_agent_task.done():
340
  print(f"Lifecycle: Cancelling accept_challenges task for {active_agent_name} as battle started.")
341
  active_agent_task.cancel()
 
 
 
 
 
 
 
342
 
343
  async def deactivate_current_agent(reason: str = "cycle"):
344
  """Cleans up the currently active agent and resets state."""
345
  global active_agent_name, active_agent_instance, active_agent_task, current_battle_instance
346
 
347
+ agent_name_to_deactivate = active_agent_name
348
  print(f"Lifecycle: Deactivating agent '{agent_name_to_deactivate}' (Reason: {reason})...")
349
 
 
350
  if reason == "battle_end":
351
  await update_display_html(create_idle_html("Battle Finished!", f"Agent <strong>{agent_name_to_deactivate}</strong> completed the match."))
352
  elif reason == "cycle":
 
356
  else: # Generic reason or error
357
  await update_display_html(create_idle_html(f"Resetting Agent ({reason})", f"Cleaning up <strong>{agent_name_to_deactivate}</strong>..."))
358
 
359
+ await asyncio.sleep(3)
 
 
 
360
  await update_display_html(create_idle_html("Preparing Next Agent...", "Please wait..."))
361
 
 
362
  agent = active_agent_instance
363
  task = active_agent_task
364
 
 
 
 
 
 
 
365
  active_agent_name = None
366
  active_agent_instance = None
367
  active_agent_task = None
368
  current_battle_instance = None
369
  print(f"Lifecycle: Global state cleared for '{agent_name_to_deactivate}'.")
370
 
 
 
371
  if task and not task.done():
372
  print(f"Lifecycle: Ensuring task cancellation for {agent_name_to_deactivate} ({task.get_name()})...")
373
  task.cancel()
374
  try:
 
375
  await asyncio.wait_for(task, timeout=2.0)
376
  print(f"Lifecycle: Task cancellation confirmed for {agent_name_to_deactivate}.")
377
  except asyncio.CancelledError:
 
379
  except asyncio.TimeoutError:
380
  print(f"Lifecycle: Task did not confirm cancellation within timeout for {agent_name_to_deactivate}.")
381
  except Exception as e:
 
382
  print(f"Lifecycle: Error during task cancellation wait for {agent_name_to_deactivate}: {e}")
383
 
 
384
  if agent:
385
  print(f"Lifecycle: Disconnecting player {agent.username}...")
386
  try:
 
387
  if hasattr(agent, '_websocket') and agent._websocket and agent._websocket.open:
388
  await agent.disconnect()
389
  print(f"Lifecycle: Player {agent.username} disconnected successfully.")
390
  else:
391
  print(f"Lifecycle: Player {agent.username} already disconnected or websocket not available.")
392
  except Exception as e:
 
393
  print(f"ERROR during agent disconnect ({agent.username}): {e}")
394
+ traceback.print_exc()
395
 
396
+ await asyncio.sleep(2)
 
397
  print(f"Lifecycle: Agent '{agent_name_to_deactivate}' deactivation complete.")
398
 
399
  async def manage_agent_lifecycle():
400
  """Runs the main loop selecting, running, and cleaning up agents sequentially."""
 
401
  global active_agent_name, active_agent_instance, active_agent_task, current_battle_instance
 
402
 
403
  print("Background lifecycle manager started.")
404
+ REFRESH_INTERVAL_SECONDS = 3
405
+ LOOP_COOLDOWN_SECONDS = 1
406
+ ERROR_RETRY_DELAY_SECONDS = 10
407
+ POST_BATTLE_DELAY_SECONDS = 5
408
 
409
  loop_counter = 0
410
 
 
414
  print(f"\n--- Lifecycle Check #{loop_counter} [{time.strftime('%H:%M:%S')}] ---")
415
 
416
  try:
 
 
 
 
417
  if active_agent_instance is None:
418
  print(f"[{loop_counter}] State 1: No active agent. Selecting...")
419
  activated = await select_and_activate_new_agent()
 
421
  print(f"[{loop_counter}] State 1: Activation failed. Waiting {ERROR_RETRY_DELAY_SECONDS}s before retry.")
422
  await asyncio.sleep(ERROR_RETRY_DELAY_SECONDS)
423
  else:
 
424
  print(f"[{loop_counter}] State 1: Agent '{active_agent_name}' activated successfully.")
 
425
 
426
+ else: # Agent is active
427
+ agent_name = active_agent_name
 
 
 
 
428
  print(f"[{loop_counter}] State 2: Agent '{agent_name}' is active.")
429
 
 
 
430
  if current_battle_instance is None:
431
  print(f"[{loop_counter}] State 2a: Checking for new battle for '{agent_name}'...")
432
+ await check_for_new_battle()
433
 
 
434
  if current_battle_instance:
435
  battle_tag = current_battle_instance.battle_tag
436
  print(f"[{loop_counter}] State 2a: *** NEW BATTLE DETECTED: {battle_tag} for '{agent_name}' ***")
 
 
437
  parts = battle_tag.split('-')
438
  is_suffixed_format = len(parts) > 3 and parts[2].isdigit()
439
 
440
  if is_suffixed_format:
 
441
  print(f"[{loop_counter}] Detected potentially non-public battle format ({battle_tag}). Forfeiting.")
 
442
  try:
443
+ if active_agent_instance:
 
444
  await active_agent_instance.forfeit(battle_tag)
 
445
  print(f"[{loop_counter}] Sent forfeit command for {battle_tag}.")
446
+ await asyncio.sleep(1.5)
447
  except Exception as forfeit_err:
448
  print(f"[{loop_counter}] ERROR sending forfeit for {battle_tag}: {forfeit_err}")
 
449
  await deactivate_current_agent(reason="forfeited_private_battle")
450
+ continue
 
451
  else:
 
452
  print(f"[{loop_counter}] Public battle format detected. Displaying battle {battle_tag}.")
453
  await update_display_html(create_battle_iframe(battle_tag))
 
 
454
  else:
 
455
  print(f"[{loop_counter}] State 2a: No new battle found. Agent '{agent_name}' remains idle, waiting for challenge.")
 
456
  idle_html = create_idle_html(f"Agent Ready: <strong>{agent_name}</strong>",
457
  f"Please challenge <strong>{agent_name}</strong> to a <strong>{DEFAULT_BATTLE_FORMAT}</strong> battle.")
458
  await update_display_html(idle_html)
459
+ await asyncio.sleep(REFRESH_INTERVAL_SECONDS)
 
460
 
461
+ if current_battle_instance is not None: # Check again in case a battle just started
 
 
462
  battle_tag = current_battle_instance.battle_tag
463
  print(f"[{loop_counter}] State 2b: Monitoring battle {battle_tag} for '{agent_name}'")
464
 
 
 
465
  if not active_agent_instance:
466
  print(f"[{loop_counter}] WARNING: Agent instance for '{agent_name}' disappeared while monitoring battle {battle_tag}! Deactivating.")
467
  await deactivate_current_agent(reason="agent_disappeared_mid_battle")
468
  continue
469
 
 
 
 
470
  battle_obj = active_agent_instance._battles.get(battle_tag)
471
 
472
  if battle_obj and battle_obj.finished:
 
474
  await deactivate_current_agent(reason="battle_end")
475
  print(f"[{loop_counter}] Waiting {POST_BATTLE_DELAY_SECONDS}s post-battle before selecting next agent.")
476
  await asyncio.sleep(POST_BATTLE_DELAY_SECONDS)
477
+ continue
 
478
  elif not battle_obj:
 
479
  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.")
480
  await deactivate_current_agent(reason="battle_object_missing")
481
  continue
 
482
  else:
 
483
  print(f"[{loop_counter}] Battle {battle_tag} ongoing for '{agent_name}'.")
484
+ await asyncio.sleep(REFRESH_INTERVAL_SECONDS)
 
 
485
 
 
486
  except asyncio.CancelledError:
487
  print("Lifecycle manager task cancelled.")
488
+ raise
489
  except Exception as e:
490
  print(f"!!! ERROR in main lifecycle loop #{loop_counter}: {e} !!!")
491
  traceback.print_exc()
492
+ current_agent_name_err = active_agent_name # Use different var name to avoid conflict
 
 
493
  if active_agent_instance:
494
+ print(f"Attempting to deactivate agent '{current_agent_name_err}' due to loop error...")
495
  try:
496
  await deactivate_current_agent(reason="main_loop_error")
497
  except Exception as deactivation_err:
498
  print(f"Error during error-handling deactivation: {deactivation_err}")
 
499
  active_agent_name = None
500
  active_agent_instance = None
501
  active_agent_task = None
502
  current_battle_instance = None
503
  else:
 
504
  print("No active agent instance during loop error.")
 
505
  await update_display_html(create_error_html(f"A server error occurred in the lifecycle manager. Please wait. ({e})"))
 
 
506
  print(f"Waiting {ERROR_RETRY_DELAY_SECONDS}s after loop error.")
507
  await asyncio.sleep(ERROR_RETRY_DELAY_SECONDS)
508
+ continue
509
 
 
510
  elapsed_time = time.monotonic() - loop_start_time
511
  if elapsed_time < LOOP_COOLDOWN_SECONDS:
512
  await asyncio.sleep(LOOP_COOLDOWN_SECONDS - elapsed_time)
513
 
514
+
515
  def log_task_exception(task: asyncio.Task):
516
  """Callback to log exceptions from background tasks (like accept_challenges)."""
517
  try:
518
  if task.cancelled():
 
519
  print(f"Task '{task.get_name()}' was cancelled.")
520
  return
 
521
  task.result()
522
  print(f"Task '{task.get_name()}' completed successfully.")
523
  except asyncio.CancelledError:
524
  print(f"Task '{task.get_name()}' confirmed cancelled (exception caught).")
525
+ pass
526
  except Exception as e:
 
527
  print(f"!!! Exception in background task '{task.get_name()}': {e} !!!")
528
  traceback.print_exc()
 
 
529
 
530
  # --- WebSocket connection manager ---
531
  class ConnectionManager:
532
  def __init__(self):
533
  self.active_connections: Set[WebSocket] = set()
 
534
  self.current_html_fragment: str = create_idle_html("Initializing...", "Setting up Pokémon Battle Stream")
535
 
536
  async def connect(self, websocket: WebSocket):
537
  await websocket.accept()
538
  self.active_connections.add(websocket)
539
  print(f"Client connected. Sending current state. Total clients: {len(self.active_connections)}")
 
540
  try:
541
  await websocket.send_text(self.current_html_fragment)
542
  except Exception as e:
543
  print(f"Error sending initial state to new client: {e}")
 
544
  await self.disconnect(websocket)
545
 
 
546
  async def disconnect(self, websocket: WebSocket):
 
547
  self.active_connections.discard(websocket)
548
  print(f"Client disconnected. Total clients: {len(self.active_connections)}")
549
 
550
  async def update_all(self, html_fragment: str):
551
  """Update the current HTML fragment and broadcast to all clients."""
552
  if self.current_html_fragment == html_fragment:
553
+ return
 
554
 
555
  self.current_html_fragment = html_fragment
556
  if not self.active_connections:
 
557
  return
558
 
559
  print(f"Broadcasting update to {len(self.active_connections)} clients...")
 
 
 
560
  send_tasks = [
561
  connection.send_text(html_fragment)
562
+ for connection in list(self.active_connections)
563
  ]
 
 
564
  results = await asyncio.gather(*send_tasks, return_exceptions=True)
 
 
 
565
  connections_to_remove = set()
566
+ # Need to iterate carefully if connections can change during gather
567
+ conn_list = list(self.active_connections)
568
  for i, result in enumerate(results):
569
+ # Ensure index is valid if connections changed mid-gather
570
+ if i < len(conn_list):
571
+ connection = conn_list[i]
572
+ if isinstance(result, Exception):
573
+ print(f"Error sending update to client: {result}. Marking for removal.")
574
+ connections_to_remove.add(connection)
575
 
 
576
  for connection in connections_to_remove:
577
+ # Check if connection still exists before disconnecting
578
+ if connection in self.active_connections:
579
+ await self.disconnect(connection)
580
 
581
 
582
  manager = ConnectionManager()
 
585
  @app.get("/", response_class=HTMLResponse)
586
  async def get_homepage():
587
  """Serves the main HTML page with WebSocket connection and improved styling."""
588
+ # ... (HTML remains the same as before)
 
589
  return """
590
  <!DOCTYPE html>
591
  <html lang="en">
 
720
  let ws = null; // WebSocket instance
721
 
722
  function connectWebSocket() {
 
723
  const wsProtocol = location.protocol === 'https:' ? 'wss' : 'ws';
724
  const wsUrl = `${wsProtocol}://${location.host}/ws`;
725
  ws = new WebSocket(wsUrl);
 
728
 
729
  ws.onopen = (event) => {
730
  console.log('WebSocket connection established.');
 
 
731
  };
732
 
733
  ws.onmessage = (event) => {
 
 
734
  streamContainer.innerHTML = event.data;
735
  };
736
 
737
  ws.onclose = (event) => {
738
  console.log(`WebSocket connection closed. Code: ${event.code}, Reason: ${event.reason}. Attempting to reconnect in 5 seconds...`);
739
+ ws = null;
 
740
  streamContainer.innerHTML = createReconnectMessage();
741
+ setTimeout(connectWebSocket, 5000);
742
  };
743
 
744
  ws.onerror = (event) => {
745
  console.error('WebSocket error:', event);
 
 
 
746
  streamContainer.innerHTML = createErrorMessage("WebSocket connection error. Attempting to reconnect...");
 
747
  if (ws && ws.readyState !== WebSocket.CLOSED) {
748
  ws.close();
749
  }
750
  };
751
  }
752
 
 
753
  function createReconnectMessage() {
754
  return `
755
  <div class="content-container error-container" style="background: #333;">
 
760
  </div>`;
761
  }
762
 
 
763
  function createErrorMessage(message) {
764
  return `
765
  <div class="content-container error-container">
 
770
  </div>`;
771
  }
772
 
 
 
773
  connectWebSocket();
 
774
  </script>
775
  </body>
776
  </html>
 
781
  await manager.connect(websocket)
782
  try:
783
  while True:
 
 
 
784
  data = await websocket.receive_text()
 
 
785
  print(f"Received unexpected message from client: {data}")
 
 
786
  except WebSocketDisconnect as e:
787
  print(f"WebSocket disconnected: Code {e.code}, Reason: {getattr(e, 'reason', 'N/A')}")
788
+ await manager.disconnect(websocket)
789
  except Exception as e:
 
790
  print(f"WebSocket error: {e}")
791
  traceback.print_exc()
792
+ await manager.disconnect(websocket)
793
+
794
+ # --- NEW: Route to display the last action ---
795
+ @app.get("/last_action", response_class=HTMLResponse)
796
+ async def get_last_action():
797
+ """Serves an HTML page displaying the last recorded agent action."""
798
+ global last_llm_action
799
+ # Return the formatted HTML page using the current state
800
+ return create_last_action_html(last_llm_action)
801
+ # --------------------------------------------
802
+
803
+ # --- Lifecyle Events ---
804
  @app.on_event("startup")
805
  async def startup_event():
806
  """Start background tasks when the application starts."""
807
  global background_task_handle
808
 
 
 
809
  static_dir = "static"
810
  if not os.path.exists(static_dir):
811
  os.makedirs(static_dir)
 
816
  print(f"Mounted static directory '{static_dir}' at '/static'")
817
 
818
  print("🚀 Starting background tasks")
 
819
  background_task_handle = asyncio.create_task(manage_agent_lifecycle(), name="LifecycleManager")
 
820
  background_task_handle.add_done_callback(log_task_exception)
821
  print("✅ Background tasks started")
822
 
 
827
 
828
  print("\n🔌 Shutting down application. Cleaning up...")
829
 
 
830
  if background_task_handle and not background_task_handle.done():
831
  print("Cancelling background task...")
832
  background_task_handle.cancel()
 
840
  except Exception as e:
841
  print(f"Error during background task cancellation: {e}")
842
 
 
 
843
  agent_to_disconnect = active_agent_instance
844
  if agent_to_disconnect:
845
  agent_name = agent_to_disconnect.username if hasattr(agent_to_disconnect, 'username') else 'Unknown Agent'
846
  print(f"Disconnecting active agent '{agent_name}'...")
847
  try:
 
848
  if hasattr(agent_to_disconnect, '_websocket') and agent_to_disconnect._websocket and agent_to_disconnect._websocket.open:
849
  await agent_to_disconnect.disconnect()
850
  print(f"Agent '{agent_name}' disconnected.")
 
853
  except Exception as e:
854
  print(f"Error during agent disconnect on shutdown for '{agent_name}': {e}")
855
 
 
856
  print(f"Closing {len(manager.active_connections)} client WebSocket connections...")
 
857
  close_tasks = [
858
+ conn.close(code=1000, reason="Server shutting down")
859
+ for conn in list(manager.active_connections)
860
  ]
861
  if close_tasks:
862
+ await asyncio.gather(*close_tasks, return_exceptions=True)
863
 
864
  print("✅ Cleanup complete. Application shutdown.")
865
 
866
 
867
+ # --- Main execution ---
868
  if __name__ == "__main__":
869
  import uvicorn
870
 
 
871
  logging.basicConfig(
872
  level=logging.INFO,
873
  format='%(asctime)s - %(name)s:%(lineno)d - %(levelname)s - %(message)s',
874
  datefmt='%Y-%m-%d %H:%M:%S'
875
  )
 
876
  logging.getLogger('poke_env').setLevel(logging.WARNING)
877
+ logging.getLogger('websockets.client').setLevel(logging.INFO)
878
 
879
  print("Starting Pokemon Battle Livestream Server...")
880
  print("="*60)
 
893
  print(f" - {name}")
894
  print("="*60)
895
  print(f"Server will run on http://0.0.0.0:7860")
896
+ print("Access the action log at http://0.0.0.0:7860/last_action") # Added info
897
  print("="*60)
898
 
 
899
  uvicorn.run(
900
+ "main:app",
901
  host="0.0.0.0",
902
  port=7860,
903
+ reload=False,
904
+ log_level="info"
905
+ )