# -*- coding: utf-8 -*- """ Worker startup entrypoint. Usage: python -m worker.main [--server ws://IP:8000/ws] [--worker-id pc-a] [--worker-name name] """ from __future__ import annotations import argparse import asyncio import logging import os os.environ.setdefault("DJANGO_SETTINGS_MODULE", "server.settings") import django # noqa: E402 django.setup() # noqa: E402 from tunnel.client import TunnelClient from worker import config from worker.recruit_filter_sync import bootstrap_recruit_filter_snapshot from worker.tasks.registry import register_all_handlers from worker.ws_client import WorkerWSClient logging.basicConfig( level=logging.INFO, format="%(asctime)s %(name)-28s %(levelname)-5s %(message)s", datefmt="%Y-%m-%d %H:%M:%S", ) logger = logging.getLogger("worker.main") def parse_args(): parser = argparse.ArgumentParser(description="Browser Control Worker Agent") parser.add_argument("--worker", action="store_true", help=argparse.SUPPRESS) parser.add_argument("--server", default=config.SERVER_WS_URL, help=f"WebSocket server URL (default: {config.SERVER_WS_URL})") parser.add_argument("--worker-id", default=config.WORKER_ID, help=f"Worker ID (default: {config.WORKER_ID})") parser.add_argument("--worker-name", default=config.WORKER_NAME, help=f"Worker name (default: {config.WORKER_NAME})") parser.add_argument("--bit-api", default=config.BIT_API_BASE, help=f"BitBrowser local API URL (default: {config.BIT_API_BASE})") parser.add_argument("--no-tunnel", action="store_true", help="Disable tunnel client") return parser.parse_args() def _local_port_from_bit_api(bit_api_base: str) -> int: try: from urllib.parse import urlparse p = urlparse(bit_api_base if "://" in bit_api_base else "http://" + bit_api_base) return p.port or 54345 except Exception: return config.TUNNEL_LOCAL_PORT def _extract_host_from_ws_url(ws_url: str) -> str: try: from urllib.parse import urlparse p = urlparse(ws_url) return p.hostname or "127.0.0.1" except Exception: return "127.0.0.1" async def run(args): register_all_handlers() logger.info("Task handlers registered") # Startup bootstrap: fetch site filters (main1 logic) and persist for frontend dropdown. try: bootstrap_result = bootstrap_recruit_filter_snapshot( worker_id=args.worker_id, bit_api_base=args.bit_api, logger=logger, ) if bootstrap_result: logger.info( "Recruit filters synced on startup: account=%s groups=%s options=%s", bootstrap_result.get("account_name", ""), bootstrap_result.get("groups", 0), bootstrap_result.get("flat_options", 0), ) except Exception as e: logger.warning("Recruit filter startup sync failed: %s", e) client = WorkerWSClient( server_url=args.server, worker_id=args.worker_id, worker_name=args.worker_name, bit_api_base=args.bit_api, ) tunnel_enabled = config.TUNNEL_ENABLED and not args.no_tunnel tunnel_client = None if tunnel_enabled: tunnel_host = _extract_host_from_ws_url(args.server) tunnel_client = TunnelClient( server_host=tunnel_host, control_port=config.TUNNEL_CONTROL_PORT, stream_port=config.TUNNEL_STREAM_PORT, worker_id=args.worker_id, local_port=_local_port_from_bit_api(args.bit_api), ) logger.info( "Tunnel enabled: expose local %s -> %s (worker_id=%s)", _local_port_from_bit_api(args.bit_api), tunnel_host, args.worker_id, ) logger.info( "Worker startup: id=%s, name=%s, server=%s", args.worker_id, args.worker_name, args.server, ) async def run_worker(): await client.run() async def run_tunnel(): if tunnel_client: await tunnel_client.run() worker_task = asyncio.create_task(run_worker()) tunnel_task = asyncio.create_task(run_tunnel()) if tunnel_client else None try: if tunnel_task is not None: await asyncio.gather(worker_task, tunnel_task) else: await worker_task except KeyboardInterrupt: logger.info("Keyboard interrupt received, stopping...") worker_task.cancel() if tunnel_task is not None: tunnel_task.cancel() try: await worker_task except asyncio.CancelledError: pass if tunnel_task is not None: try: await tunnel_task except asyncio.CancelledError: pass await tunnel_client.stop() await client.stop() def main(): args = parse_args() try: asyncio.run(run(args)) except KeyboardInterrupt: logger.info("Worker exited") if __name__ == "__main__": main()