rename run_step_headless to run

2 files changed, 8 insertions, 9 deletions

diff --git a/continuedev/src/continuedev/__main__.py b/continuedev/src/continuedev/__main__.py
index efd62d64..52fb85b7 100644
--- a/continuedev/src/continuedev/__main__.py
+++ b/continuedev/src/continuedev/__main__.py
@@ -1,9 +1,8 @@
-import asyncio
 from typing import Optional
 
 import typer
 
-from .headless import start_headless_session
+from .headless import run
 from .server.main import run_server
 
 app = typer.Typer()
@@ -19,10 +18,7 @@ def main(
     headless: bool = typer.Option(False, help="Run in headless mode"),
 ):
     if headless:
-        loop = asyncio.get_event_loop()
-        loop.run_until_complete(start_headless_session(config=config))
-        tasks = asyncio.all_tasks(loop)
-        loop.run_until_complete(asyncio.gather(*tasks))
+        run(config)
     else:
         run_server(port=port, host=host)
 
diff --git a/continuedev/src/continuedev/headless/__init__.py b/continuedev/src/continuedev/headless/__init__.py
index 27722ee7..fb2c9768 100644
--- a/continuedev/src/continuedev/headless/__init__.py
+++ b/continuedev/src/continuedev/headless/__init__.py
@@ -22,9 +22,12 @@ async def start_headless_session(
     return await session_manager.new_session(ide, config=config)
 
 
-def run_step_headless(step: Step):
-    config = ContinueConfig()
-    config.steps_on_startup = [step]
+def run(step_or_config: Union[Step, ContinueConfig]):
+    if isinstance(step_or_config, ContinueConfig):
+        config = step_or_config
+    else:
+        config = ContinueConfig()
+        config.steps_on_startup = [step_or_config]
 
     loop = asyncio.get_event_loop()
     loop.run_until_complete(start_headless_session(config=config))