refac

src/mcpo/__init__.py

@@ -33,7 +33,10 @@ def main(
         Optional[str],
         typer.Option("--env-path", help="Path to environment variables file"),
     ] = None,
-    config: Annotated[
+    server_type: Annotated[
+        Optional[str], typer.Option("--type", help="Server type")
+    ] = None,
+    config_path: Annotated[
         Optional[str], typer.Option("--config", "-c", help="Config file path")
     ] = None,
     name: Annotated[
@@ -56,7 +59,7 @@ def main(
     ] = None,
 ):
     server_command = None
-    if not config:
+    if not config_path:
         # Find the position of "--"
         if "--" not in sys.argv:
             typer.echo("Usage: mcpo --host 0.0.0.0 --port 8000 -- your_mcp_command")
@@ -71,8 +74,8 @@ def main(
 
     from mcpo.main import run
 
-    if config:
-        print("Starting MCP OpenAPI Proxy with config file:", config)
+    if config_path:
+        print("Starting MCP OpenAPI Proxy with config file:", config_path)
     else:
         print(
             f"Starting MCP OpenAPI Proxy on {host}:{port} with command: {' '.join(server_command)}"
@@ -114,7 +117,8 @@ def main(
             port,
             api_key=api_key,
             cors_allow_origins=cors_allow_origins,
-            config=config,
+            server_type=server_type,
+            config_path=config_path,
             name=name,
             description=description,
             version=version,

src/mcpo/main.py

@@ -65,16 +65,18 @@ async def create_dynamic_endpoints(app: FastAPI, api_dependency=None):
 
 @asynccontextmanager
 async def lifespan(app: FastAPI):
+    server_type = getattr(app.state, "server_type", "stdio")
     command = getattr(app.state, "command", None)
     args = getattr(app.state, "args", [])
     env = getattr(app.state, "env", {})
 
-    mcptype = "sse" if command == "sse" else "stdio"
     args = args if isinstance(args, list) else [args]
-
     api_dependency = getattr(app.state, "api_dependency", None)
 
-    if (mcptype == "stdio" and not command) or (mcptype == "sse" and not args[0]):
+    if (server_type == "stdio" and not command) or (
+        server_type == "sse" and not args[0]
+    ):
+        # Main app lifespan (when config_path is provided)
         async with AsyncExitStack() as stack:
             for route in app.routes:
                 if isinstance(route, Mount) and isinstance(route.app, FastAPI):
@@ -82,9 +84,8 @@ async def lifespan(app: FastAPI):
                         route.app.router.lifespan_context(route.app),  # noqa
                     )
             yield
-
     else:
-        if mcptype == "stdio":
+        if server_type == "stdio":
             server_params = StdioServerParameters(
                 command=command,
                 args=args,
@@ -96,7 +97,7 @@ async def lifespan(app: FastAPI):
                     app.state.session = session
                     await create_dynamic_endpoints(app, api_dependency=api_dependency)
                     yield
-        if mcptype == "sse":
+        if server_type == "sse":
             async with sse_client(url=args[0]) as (reader, writer):
                 async with ClientSession(reader, writer) as session:
                     app.state.session = session
@@ -114,10 +115,14 @@ async def run(
     # Server API Key
     api_dependency = get_verify_api_key(api_key) if api_key else None
 
-    # MCP Config
-    config_path = kwargs.get("config")
+    # MCP Server
+    server_type = kwargs.get("server_type")  # "stdio" or "sse" or "http"
     server_command = kwargs.get("server_command")
 
+    # MCP Config
+    config_path = kwargs.get("config_path")
+
+    # mcpo server
     name = kwargs.get("name") or "MCP OpenAPI Proxy"
     description = (
         kwargs.get("description") or "Automatically generated API from MCP Tool Schemas"
@@ -145,12 +150,18 @@ async def run(
         allow_headers=["*"],
     )
 
-    if server_command:
+    if server_type == "sse":
+        main_app.state.server_type = "sse"
+        main_app.state.args = server_command[0]
+        main_app.state.api_dependency = api_dependency
+
+    elif server_command:
         main_app.state.command = server_command[0]
         main_app.state.args = server_command[1:]
         main_app.state.env = os.environ.copy()
 
         main_app.state.api_dependency = api_dependency
+
     elif config_path:
         with open(config_path, "r") as f:
             config_data = json.load(f)
@@ -183,7 +194,7 @@ async def run(
                 sub_app.state.env = {**os.environ, **server_cfg.get("env", {})}
             if server_cfg.get("url"):
                 # SSE
-                sub_app.state.command = "sse"
+                sub_app.state.server_type = "sse"
                 sub_app.state.args = server_cfg["url"]
 
             sub_app.state.api_dependency = api_dependency