| 1 | import initialize |
| 2 | from helpers import dotenv, extension, runtime |
| 3 | from helpers.api import csrf_protect, requires_auth |
| 4 | from helpers.print_style import PrintStyle |
| 5 | from helpers.server_startup import run_uvicorn_with_retries |
| 6 | from helpers.ui_server import UiServerRuntime, configure_process_environment |
| 7 | |
| 8 | |
| 9 | def run(): |
| 10 | configure_process_environment() |
| 11 | PrintStyle().print("Initializing Python framework...") |
| 12 | PrintStyle().print("Checking for data migration...") |
| 13 | run_migration_checks() |
| 14 | |
| 15 | PrintStyle().print("Preparing web server runtime...") |
| 16 | server_runtime, host, port = prepare_web_runtime() |
| 17 | |
| 18 | PrintStyle().print("Initializing Agent Zero components...") |
| 19 | init_a0() |
| 20 | |
| 21 | PrintStyle().print("Starting UI/API server...") |
| 22 | start_web_server(server_runtime, host, port) |
| 23 | |
| 24 | |
| 25 | def run_migration_checks() -> None: |
| 26 | initialize.initialize_migration() |
| 27 | |
| 28 | |
| 29 | def prepare_web_runtime() -> tuple[UiServerRuntime, str, int]: |
| 30 | host = ( |
| 31 | runtime.get_arg("host") or dotenv.get_dotenv_value("WEB_UI_HOST") or "localhost" |
| 32 | ) |
| 33 | port = runtime.get_web_ui_port() |
| 34 | server_runtime = UiServerRuntime.create() |
| 35 | server_runtime.register_http_routes() |
| 36 | server_runtime.register_transport_handlers() |
| 37 | |
| 38 | return server_runtime, host, port |
| 39 | |
| 40 | |
| 41 | def start_web_server(server_runtime: UiServerRuntime, host: str, port: int) -> None: |
| 42 | run_uvicorn_with_retries( |
| 43 | host=host, |
| 44 | port=port, |
| 45 | build_asgi_app=server_runtime.build_asgi_app, |
| 46 | flush_callback=create_flush_callback(), |
| 47 | access_log=server_runtime.access_log_enabled(), |
| 48 | ws="wsproto", |
| 49 | ) |
| 50 | |
| 51 | |
| 52 | def create_flush_callback(): |
| 53 | def flush_and_shutdown_callback() -> None: |
| 54 | """ |
| 55 | TODO(dev): add cleanup + flush-to-disk logic here. |
| 56 | """ |
| 57 | return |
| 58 | |
| 59 | flush_ran = False |
| 60 | |
| 61 | def _run_flush(reason: str) -> None: |
| 62 | nonlocal flush_ran |
| 63 | if flush_ran: |
| 64 | return |
| 65 | flush_ran = True |
| 66 | try: |
| 67 | flush_and_shutdown_callback() |
| 68 | except Exception as e: |
| 69 | PrintStyle.warning(f"Shutdown flush failed ({reason}): {e}") |
| 70 | |
| 71 | return _run_flush |
| 72 | |
| 73 | |
| 74 | @extension.extensible |
| 75 | def init_a0(): |
| 76 | init_chats = initialize.initialize_chats() |
| 77 | init_chats.result_sync() |
| 78 | |
| 79 | initialize.initialize_mcp() |
| 80 | initialize.initialize_job_loop() |
| 81 | initialize.initialize_preload() |
| 82 | |
| 83 | |
| 84 | if __name__ == "__main__": |
| 85 | runtime.initialize() |
| 86 | dotenv.load_dotenv() |
| 87 | run() |