From f36afd7f1beac73477cff7fb6fa5f90a31be365d Mon Sep 17 00:00:00 2001
From: Wei Ouyang <oeway007@gmail.com>
Date: Sat, 1 Feb 2025 00:38:10 +0800
Subject: [PATCH] Update interactive server cli

---
 hypha/__main__.py    | 10 +++++++---
 hypha/interactive.py | 30 +++++++++++++++++++++++++++++-
 hypha/server.py      | 14 +++-----------
 3 files changed, 39 insertions(+), 15 deletions(-)

diff --git a/hypha/__main__.py b/hypha/__main__.py
index eecdfc17..a63a780a 100644
--- a/hypha/__main__.py
+++ b/hypha/__main__.py
@@ -16,7 +16,7 @@ def run_interactive_cli():
     opt = arg_parser.parse_args(sys.argv[1:])
 
     # Force interactive server mode
-    if not opt.interactive_server:
+    if not opt.interactive:
         opt.interactive = True
 
     app = create_application(opt)
@@ -25,16 +25,20 @@ def run_interactive_cli():
     asyncio.run(start_interactive_shell(app, opt))
 
 
-if __name__ == "__main__":
+def main():
     # Create the app instance
     arg_parser = get_argparser()
     opt = arg_parser.parse_args()
     app = create_application(opt)
 
-    if opt.interactive or opt.interactive_server:
+    if opt.interactive:
         asyncio.run(start_interactive_shell(app, opt))
     else:
         uvicorn.run(app, host=opt.host, port=int(opt.port))
+
+
+if __name__ == "__main__":
+    main()
 else:
     # Create the app instance when imported by uvicorn
     arg_parser = get_argparser(add_help=False)
diff --git a/hypha/interactive.py b/hypha/interactive.py
index 80dca5bd..1906d20b 100644
--- a/hypha/interactive.py
+++ b/hypha/interactive.py
@@ -40,6 +40,29 @@ def configure_ptpython(repl: PythonRepl) -> None:
     print(welcome_message)
 
 
+def is_port_in_use(host: str, port: int) -> bool:
+    """Check if a port is in use on the specified interface by attempting to bind to it.
+
+    Args:
+        host: Interface/host to check (e.g. 'localhost', '0.0.0.0')
+        port: Port number to check
+
+    Returns:
+        bool: True if port is in use, False otherwise
+    """
+    import socket
+
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+        # Allow reuse of address to avoid "Address already in use" after testing
+        s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+        try:
+            s.bind((host, port))
+            # Close and unbind the socket by returning False
+            return False
+        except socket.error:
+            return True
+
+
 async def start_interactive_shell(app: FastAPI, args: Any) -> None:
     """Start an interactive shell with the hypha store and server app."""
 
@@ -52,6 +75,11 @@ async def start_interactive_shell(app: FastAPI, args: Any) -> None:
     print("Initializing interactive shell...\n")
 
     async def start_server():
+        if is_port_in_use(args.host, int(args.port)):
+            print(
+                f"\nFailed to start server on port {args.port}. Port is already in use."
+            )
+            sys.exit(1)
         nonlocal server, server_task
         config = uvicorn.Config(app, host=args.host, port=int(args.port))
         server = uvicorn.Server(config)
@@ -60,7 +88,7 @@ async def start_server():
         await store.get_event_bus().wait_for_local("startup")
         print(f"\nServer started at http://{args.host}:{args.port}")
 
-    if args.interactive_server:
+    if args.enable_server:
         # Start the server in the background
         await start_server()
 
diff --git a/hypha/server.py b/hypha/server.py
index 468f2f60..4cccbacc 100644
--- a/hypha/server.py
+++ b/hypha/server.py
@@ -506,15 +506,10 @@ def get_argparser(add_help=True):
         action="store_true",
         help="start an interactive shell with the hypha store",
     )
-    return parser
-
-
-def add_interactive_arguments(parser):
-    """Add interactive-specific arguments to the parser."""
     parser.add_argument(
-        "--interactive-server",
+        "--enable-server",
         action="store_true",
-        help="start an interactive server with the hypha store",
+        help="enable server in interactive mode",
     )
     return parser
 
@@ -524,9 +519,6 @@ def add_interactive_arguments(parser):
 
     arg_parser = get_argparser()
 
-    # Only add interactive server argument when running as main script
-    add_interactive_arguments(arg_parser)
-
     opt = arg_parser.parse_args()
 
     # Apply database migrations
@@ -541,7 +533,7 @@ def add_interactive_arguments(parser):
         command.upgrade(alembic_cfg, "head")
 
     app = create_application(opt)
-    if opt.interactive or opt.interactive_server:
+    if opt.interactive:
         from hypha.interactive import start_interactive_shell
 
         asyncio.run(start_interactive_shell(app, opt))