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