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 <[email protected]>

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/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

@@ -10,6 +10,7 @@
 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
@@ -31,6 +32,7 @@ def version_callback(value: bool):
 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)

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/update/main.py

@@ -43,7 +43,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."""