snake-python/server/blueprints/battlesnake.py

from typing import cast
import asyncio, json, time, os

from quart import Blueprint, request, jsonify

from quart_common.web.decorators import require_user_agent
from quart_common.web.logger import await_log
from snakes import DEFAULT_SNAKE_CONFIG
from server.GameBoard import GameBoard
from server.Files import read_file
from server.services.move_safety import fallback_move, run_budgeted_move
from server.services.move_trace import MoveTrace

lazy from server.Server import Server

def create_battlesnake_blueprint(server:'Server') -> Blueprint:
  blueprint = Blueprint('battlesnake', __name__)

  async def _override_snake_config_with_environment_variables(config:dict[str, str]) -> dict[str, str]:
    print(config)
    config['version'] = server.snake_version

    for key in ('author', 'color', 'head', 'tail'):
      value = os.environ.get(f'SNAKE_{key.upper()}')
      if value is not None:
        config[key] = value

    version_override = os.environ.get('SNAKE_VERSION')
    if version_override is not None:
      config['version'] = version_override

    return config

  @blueprint.get('/')
  async def on_info():
    server.metrics_collector.record_http_request('info')

    snake_config = cast(dict[str, str]|None, await read_file(server.config_file, json.load))
    if not snake_config:
      snake_json = await _override_snake_config_with_environment_variables(DEFAULT_SNAKE_CONFIG)
    else:
      snake_json = await _override_snake_config_with_environment_variables(snake_config)

    await await_log(server.logger.info(f'INFO Snake: {snake_json}'))
    return snake_json

  @blueprint.post('/start')
  @require_user_agent(("BattlesnakeEngine", "BattlesnakeArena"), abort_code=404)
  async def on_start():
    server.metrics_collector.record_http_request('start')
    game_state = await request.get_json()
    game_board = await server.game_runtime.create_game_board(game_state, publish_metrics=False)
    active_count = len(server.game_runtime.running_games)
    server.game_lifecycle.submit(
      game_state['game']['id'],
      lambda: server.game_lifecycle.started(game_state, game_board, active_count),
    )
    return 'ok'

  @blueprint.post('/move')
  @require_user_agent(("BattlesnakeEngine", "BattlesnakeArena"), abort_code=404)
  async def on_move():
    move_started = time.perf_counter()
    server.metrics_collector.record_http_request('move')
    game_state = await request.get_json()
    trace = MoveTrace(game_state['game']['id'], game_state.get('turn'))
    trace.set(parse_ms=round((time.perf_counter() - move_started) * 1000, 3), timed_out=False)

    game_id = game_state['game']['id']
    timeout_ms = int(game_state.get('game', {}).get('timeout', 500))
    budget_sec = max(0.05, (timeout_ms - 50) / 1000.0)

    next_move = fallback_move(game_state)
    game_board = None

    try:
      async with asyncio.timeout(budget_sec):
        with trace.stage('runtime'):
          game_board = cast(GameBoard, await server.game_runtime.get_game_board(game_state))
        loop = asyncio.get_running_loop()
        submitted = time.perf_counter()
        next_move = await loop.run_in_executor(
          None, run_budgeted_move, game_board, move_started + budget_sec, next_move, trace, submitted,
        )
    except TimeoutError:
      trace.set(timed_out=True)

    # Never persist another turn's history when the worker has timed out.
    thinking = None if trace.snapshot()['timed_out'] else server.gameplay_tracking._extract_latest_snake_thinking(game_board)
    thinking = dict(thinking or {})
    trace.set(handler_ms=round((time.perf_counter() - move_started) * 1000, 3),
              move=next_move, budget_ms=round(budget_sec * 1000, 3))
    thinking['move_trace'] = trace.snapshot()
    server.metrics_collector.record_move_local(next_move, trace.snapshot()['handler_ms'])

    bookkeeping_submitted = time.perf_counter()

    async def finish_move():
      trace.set(bookkeeping_queue_ms=round((time.perf_counter() - bookkeeping_submitted) * 1000, 3))
      try:
        try:
          with trace.stage('metrics_publish'):
            await server.metrics_collector.publish_move_metrics()
        except Exception:
          trace.set(metrics_publish_failed=True)
          server.app.logger.exception('Move metrics publish failed; trace=%s', trace.id)
        with trace.stage('persistence'):
          thinking['move_trace'] = trace.snapshot()
          await server.gameplay_tracking.record_gameplay_turn_snapshot(game_state, next_move, thinking)
        with trace.stage('dashboard_publish'):
          await server.dashboard_query.push_dashboard_game_replay_update(game_id)
      finally:
        trace.log(server.app.logger)

    server.game_lifecycle.submit(game_id, finish_move)

    if server.debug:
      await await_log(server.logger.debug(f'TURN: {game_state['turn']:3}, MOVE: {next_move:5}'))

    with trace.stage('response_encode'):
      response = jsonify(move=next_move)
    trace.set(handler_ms=round((time.perf_counter() - move_started) * 1000, 3))
    response.headers['X-Move-Trace-ID'] = trace.id
    return response

  @blueprint.post('/end')
  @require_user_agent(("BattlesnakeEngine", "BattlesnakeArena"), abort_code=404)
  async def on_end():
    server.metrics_collector.record_http_request('end')
    game_state = await request.get_json()
    game_id = game_state['game']['id']
    game_board = server.game_runtime.running_games.get(game_id)
    if game_board is not None:
      game_board.read_game_data(game_state)
      game_board.end_game(game_state)
    await server.game_runtime.delete_game_board(game_state)
    server.game_lifecycle.submit(
      game_id, lambda: server.game_lifecycle.ended(game_state, game_board),
    )
    return 'ok'

  # @blueprint.get('/cleanup')
  # async def cleanup():
  #   results = server._cleanup_database()
  #   return jsonify(data=json.loads(results), status=200)

  return blueprint