Update server/app.py

server/app.py

@@ -1049,20 +1049,22 @@ async def get_schemas() -> SchemaResponse:
         state=AetherTaskFlowState.model_json_schema(),
     )
 
-
-
-def main(host: str = "0.0.0.0", port: int = 7860) -> None:
+def _run_server(host: str, port: int, difficulty: str) -> None:
+    if difficulty != _DIFFICULTY:
+        os.environ["AETHER_DIFFICULTY"] = difficulty
     import uvicorn
     uvicorn.run(app, host=host, port=port)
 
 
-if __name__ == "__main__":
+def main() -> None:
     import argparse
     parser = argparse.ArgumentParser(description="AETHER-TaskFlow server")
     parser.add_argument("--host", default="0.0.0.0")
     parser.add_argument("--port", type=int, default=7860)
     parser.add_argument("--difficulty", choices=["easy", "medium", "hard"], default=_DIFFICULTY)
     args = parser.parse_args()
-    if args.difficulty != _DIFFICULTY:
-        os.environ["AETHER_DIFFICULTY"] = args.difficulty
-    main(host=args.host, port=args.port)
+    _run_server(host=args.host, port=args.port, difficulty=args.difficulty)
+
+
+if __name__ == "__main__":
+    main()