Implement MCP Server functionality for Codegen CLI (#1153)

Co-authored-by: codegen-sh[bot] <131295404+codegen-sh[bot]@users.noreply.github.com>
Co-authored-by: Jay Hack <jayhack@users.noreply.github.com>

Author: codegen-sh[bot]

pyproject.toml

@@ -23,7 +23,8 @@ dependencies = [
   "codeowners>=0.6.0",
   "unidiff>=0.7.5",
   "datamodel-code-generator>=0.26.5",
-  "mcp[cli]",
+  "mcp[cli]==1.9.4",
+  "fastmcp>=2.9.0",
   # Utility dependencies
   "colorlog>=6.9.0",
   "psutil>=5.8.0",

src/codegen/cli/api/client.py

@@ -5,8 +5,6 @@
 from pydantic import BaseModel
 from rich import print as rprint
 
-from codegen.cli.api.endpoints import IDENTIFY_ENDPOINT
-from codegen.cli.api.schemas import IdentifyResponse
 from codegen.cli.env.global_env import global_env
 from codegen.cli.errors import InvalidTokenError, ServerError
 
@@ -16,11 +14,13 @@
 
 class AuthContext(BaseModel):
     """Authentication context model."""
+
     status: str
 
 
 class Identity(BaseModel):
     """User identity model."""
+
     auth_context: AuthContext
 
 

src/codegen/cli/auth/decorators.py

@@ -39,7 +39,7 @@ def wrapper(*args, **kwargs):
 
     # Remove the session parameter from the wrapper's signature so Typer doesn't see it
     sig = inspect.signature(f)
-    new_params = [param for name, param in sig.parameters.items() if name != 'session']
+    new_params = [param for name, param in sig.parameters.items() if name != "session"]
     new_sig = sig.replace(parameters=new_params)
     wrapper.__signature__ = new_sig  # type: ignore[attr-defined]
 

src/codegen/cli/auth/session.ipynb

@@ -11,7 +11,7 @@
     "\n",
     "\n",
     "# Create a session with the current directory as repo_path\n",
-    "session = CodegenSession(repo_path=Path('.'))\n",
+    "session = CodegenSession(repo_path=Path(\".\"))\n",
     "print(f\"Session: {session}\")\n",
     "print(f\"Repo path: {session.repo_path}\")\n",
     "print(f\"Config: {session.config}\")\n",

src/codegen/cli/cli.py

@@ -1,9 +1,19 @@
 import typer
 from rich.traceback import install
 
+from codegen import __version__
+
 # Import config command (still a Typer app)
 from codegen.cli.commands.config.main import config_command
-from codegen import __version__
+
+# Import the actual command functions
+from codegen.cli.commands.init.main import init
+from codegen.cli.commands.login.main import login
+from codegen.cli.commands.logout.main import logout
+from codegen.cli.commands.mcp.main import mcp
+from codegen.cli.commands.profile.main import profile
+from codegen.cli.commands.style_debug.main import style_debug
+from codegen.cli.commands.update.main import update
 
 install(show_locals=True)
 
@@ -14,25 +24,15 @@ def version_callback(value: bool):
         print(__version__)
         raise typer.Exit()
 
-# Create the main Typer app
-main = typer.Typer(
-    name="codegen",
-    help="Codegen CLI - Transform your code with AI.",
-    rich_markup_mode="rich"
-)
 
-# Import the actual command functions
-from codegen.cli.commands.init.main import init
-from codegen.cli.commands.login.main import login
-from codegen.cli.commands.logout.main import logout
-from codegen.cli.commands.profile.main import profile
-from codegen.cli.commands.style_debug.main import style_debug
-from codegen.cli.commands.update.main import update
+# Create the main Typer app
+main = typer.Typer(name="codegen", help="Codegen CLI - Transform your code with AI.", rich_markup_mode="rich")
 
 # Add individual commands to the main app
 main.command("init", help="Initialize or update the Codegen folder.")(init)
 main.command("login", help="Store authentication token.")(login)
 main.command("logout", help="Clear stored authentication token.")(logout)
+main.command("mcp", help="Start the Codegen MCP server.")(mcp)
 main.command("profile", help="Display information about the currently authenticated user.")(profile)
 main.command("style-debug", help="Debug command to visualize CLI styling (spinners, etc).")(style_debug)
 main.command("update", help="Update Codegen to the latest or specified version")(update)
@@ -42,9 +42,7 @@ def version_callback(value: bool):
 
 
 @main.callback()
-def main_callback(
-    version: bool = typer.Option(False, "--version", callback=version_callback, is_eager=True, help="Show version and exit")
-):
+def main_callback(version: bool = typer.Option(False, "--version", callback=version_callback, is_eager=True, help="Show version and exit")):
     """Codegen CLI - Transform your code with AI."""
     pass
 

src/codegen/cli/commands/config/main.py

@@ -90,10 +90,7 @@ def get_config(key: str = typer.Argument(..., help="Configuration key to get")):
 
 
 @config_command.command(name="set")
-def set_config(
-    key: str = typer.Argument(..., help="Configuration key to set"),
-    value: str = typer.Argument(..., help="Configuration value to set")
-):
+def set_config(key: str = typer.Argument(..., help="Configuration key to set"), value: str = typer.Argument(..., help="Configuration value to set")):
     """Set a configuration value and write to .env"""
     config = _get_user_config()
     if not config.has_key(key):

src/codegen/cli/commands/init/main.py

@@ -1,4 +1,3 @@
-import sys
 from pathlib import Path
 from typing import Optional
 
@@ -9,18 +8,19 @@
 from codegen.cli.rich.codeblocks import format_command
 from codegen.shared.path import get_git_root_path
 
+
 def init(
     path: Optional[str] = typer.Option(None, help="Path within a git repository. Defaults to the current directory."),
     token: Optional[str] = typer.Option(None, help="Access token for the git repository. Required for full functionality."),
     language: Optional[str] = typer.Option(None, help="Override automatic language detection (python or typescript)"),
-    fetch_docs: bool = typer.Option(False, "--fetch-docs", help="Fetch docs and examples (requires auth)")
+    fetch_docs: bool = typer.Option(False, "--fetch-docs", help="Fetch docs and examples (requires auth)"),
 ):
     """Initialize or update the Codegen folder."""
     # Validate language option
     if language and language.lower() not in ["python", "typescript"]:
         rich.print(f"[bold red]Error:[/bold red] Invalid language '{language}'. Must be 'python' or 'typescript'.")
         raise typer.Exit(1)
-    
+
     # Print a message if not in a git repo
     path_obj = Path.cwd() if path is None else Path(path)
     repo_path = get_git_root_path(path_obj)

src/codegen/cli/commands/login/main.py

@@ -1,10 +1,12 @@
 from typing import Optional
-import typer
+
 import rich
+import typer
 
 from codegen.cli.auth.login import login_routine
 from codegen.cli.auth.token_manager import get_current_token
 
+
 def login(token: Optional[str] = typer.Option(None, help="API token for authentication")):
     """Store authentication token."""
     # Check if already authenticated

src/codegen/cli/commands/logout/main.py

@@ -1,8 +1,8 @@
 import rich
-import typer
 
 from codegen.cli.auth.token_manager import TokenManager
 
+
 def logout():
     """Clear stored authentication token."""
     token_manager = TokenManager()

src/codegen/cli/commands/mcp/__init__.py

@@ -0,0 +1,2 @@
+
+

src/codegen/cli/commands/mcp/main.py

@@ -0,0 +1,38 @@
+"""MCP server command for the Codegen CLI."""
+
+from typing import Optional
+
+import typer
+from rich.console import Console
+
+console = Console()
+
+
+def mcp(
+    host: str = typer.Option("localhost", help="Host to bind the MCP server to"),
+    port: Optional[int] = typer.Option(None, help="Port to bind the MCP server to (default: stdio transport)"),
+    transport: str = typer.Option("stdio", help="Transport protocol to use (stdio or http)"),
+):
+    """Start the Codegen MCP server."""
+    console.print("🚀 Starting Codegen MCP server...", style="bold green")
+
+    if transport == "stdio":
+        console.print("📡 Using stdio transport", style="dim")
+    else:
+        console.print(f"📡 Using HTTP transport on {host}:{port}", style="dim")
+
+    # Validate transport
+    if transport not in ["stdio", "http"]:
+        console.print(f"❌ Invalid transport: {transport}. Must be 'stdio' or 'http'", style="bold red")
+        raise typer.Exit(1)
+
+    # Import here to avoid circular imports and ensure dependencies are available
+    from codegen.cli.mcp.server import run_server
+
+    try:
+        run_server(transport=transport, host=host, port=port)
+    except KeyboardInterrupt:
+        console.print("\n👋 MCP server stopped", style="yellow")
+    except Exception as e:
+        console.print(f"❌ Error starting MCP server: {e}", style="bold red")
+        raise typer.Exit(1)

src/codegen/cli/commands/profile/main.py

@@ -1,5 +1,4 @@
 import rich
-import typer
 from rich import box
 from rich.panel import Panel
 
@@ -13,6 +12,7 @@ def requires_init(func):
     """Simple stub decorator that does nothing."""
     return func
 
+
 @requires_auth
 @requires_init
 def profile(session: CodegenSession):

src/codegen/cli/commands/style_debug/main.py

@@ -6,6 +6,7 @@
 
 from codegen.cli.rich.spinners import create_spinner
 
+
 def style_debug(text: str = typer.Option("Loading...", help="Text to show in the spinner")):
     """Debug command to visualize CLI styling (spinners, etc)."""
     try:

src/codegen/cli/commands/update/main.py

@@ -33,7 +33,7 @@ def install_package(package: str, *args: str) -> None:
 
 def update(
     list_: bool = typer.Option(False, "--list", "-l", help="List all supported versions of the codegen"),
-    version: Optional[str] = typer.Option(None, "--version", "-v", help="Update to a specific version of the codegen")
+    version: Optional[str] = typer.Option(None, "--version", "-v", help="Update to a specific version of the codegen"),
 ):
     """Update Codegen to the latest or specified version
 
@@ -44,7 +44,8 @@ def update(
         rich.print("[red]Error:[/red] Cannot specify both --list and --version")
         raise typer.Exit(1)
 
-    package_info = distribution(codegen.__package__)
+    package_name = codegen.__package__ or "codegen"
+    package_info = distribution(package_name)
     current_version = Version(package_info.version)
 
     if list_:

src/codegen/cli/mcp/__init__.py

@@ -0,0 +1 @@
+"""MCP (Model Context Protocol) server for Codegen."""