Skip to content

Commit 53af51f

Browse files
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]>
1 parent 3d49f83 commit 53af51f

File tree

25 files changed

+626
-83
lines changed

25 files changed

+626
-83
lines changed

pyproject.toml

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,8 @@ dependencies = [
2323
"codeowners>=0.6.0",
2424
"unidiff>=0.7.5",
2525
"datamodel-code-generator>=0.26.5",
26-
"mcp[cli]",
26+
"mcp[cli]==1.9.4",
27+
"fastmcp>=2.9.0",
2728
# Utility dependencies
2829
"colorlog>=6.9.0",
2930
"psutil>=5.8.0",

src/codegen/cli/api/client.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -5,8 +5,6 @@
55
from pydantic import BaseModel
66
from rich import print as rprint
77

8-
from codegen.cli.api.endpoints import IDENTIFY_ENDPOINT
9-
from codegen.cli.api.schemas import IdentifyResponse
108
from codegen.cli.env.global_env import global_env
119
from codegen.cli.errors import InvalidTokenError, ServerError
1210

@@ -16,11 +14,13 @@
1614

1715
class AuthContext(BaseModel):
1816
"""Authentication context model."""
17+
1918
status: str
2019

2120

2221
class Identity(BaseModel):
2322
"""User identity model."""
23+
2424
auth_context: AuthContext
2525

2626

src/codegen/cli/auth/decorators.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -39,7 +39,7 @@ def wrapper(*args, **kwargs):
3939

4040
# Remove the session parameter from the wrapper's signature so Typer doesn't see it
4141
sig = inspect.signature(f)
42-
new_params = [param for name, param in sig.parameters.items() if name != 'session']
42+
new_params = [param for name, param in sig.parameters.items() if name != "session"]
4343
new_sig = sig.replace(parameters=new_params)
4444
wrapper.__signature__ = new_sig # type: ignore[attr-defined]
4545

src/codegen/cli/auth/session.ipynb

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@
1111
"\n",
1212
"\n",
1313
"# Create a session with the current directory as repo_path\n",
14-
"session = CodegenSession(repo_path=Path('.'))\n",
14+
"session = CodegenSession(repo_path=Path(\".\"))\n",
1515
"print(f\"Session: {session}\")\n",
1616
"print(f\"Repo path: {session.repo_path}\")\n",
1717
"print(f\"Config: {session.config}\")\n",

src/codegen/cli/cli.py

Lines changed: 15 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,19 @@
11
import typer
22
from rich.traceback import install
33

4+
from codegen import __version__
5+
46
# Import config command (still a Typer app)
57
from codegen.cli.commands.config.main import config_command
6-
from codegen import __version__
8+
9+
# Import the actual command functions
10+
from codegen.cli.commands.init.main import init
11+
from codegen.cli.commands.login.main import login
12+
from codegen.cli.commands.logout.main import logout
13+
from codegen.cli.commands.mcp.main import mcp
14+
from codegen.cli.commands.profile.main import profile
15+
from codegen.cli.commands.style_debug.main import style_debug
16+
from codegen.cli.commands.update.main import update
717

818
install(show_locals=True)
919

@@ -14,25 +24,15 @@ def version_callback(value: bool):
1424
print(__version__)
1525
raise typer.Exit()
1626

17-
# Create the main Typer app
18-
main = typer.Typer(
19-
name="codegen",
20-
help="Codegen CLI - Transform your code with AI.",
21-
rich_markup_mode="rich"
22-
)
2327

24-
# Import the actual command functions
25-
from codegen.cli.commands.init.main import init
26-
from codegen.cli.commands.login.main import login
27-
from codegen.cli.commands.logout.main import logout
28-
from codegen.cli.commands.profile.main import profile
29-
from codegen.cli.commands.style_debug.main import style_debug
30-
from codegen.cli.commands.update.main import update
28+
# Create the main Typer app
29+
main = typer.Typer(name="codegen", help="Codegen CLI - Transform your code with AI.", rich_markup_mode="rich")
3130

3231
# Add individual commands to the main app
3332
main.command("init", help="Initialize or update the Codegen folder.")(init)
3433
main.command("login", help="Store authentication token.")(login)
3534
main.command("logout", help="Clear stored authentication token.")(logout)
35+
main.command("mcp", help="Start the Codegen MCP server.")(mcp)
3636
main.command("profile", help="Display information about the currently authenticated user.")(profile)
3737
main.command("style-debug", help="Debug command to visualize CLI styling (spinners, etc).")(style_debug)
3838
main.command("update", help="Update Codegen to the latest or specified version")(update)
@@ -42,9 +42,7 @@ def version_callback(value: bool):
4242

4343

4444
@main.callback()
45-
def main_callback(
46-
version: bool = typer.Option(False, "--version", callback=version_callback, is_eager=True, help="Show version and exit")
47-
):
45+
def main_callback(version: bool = typer.Option(False, "--version", callback=version_callback, is_eager=True, help="Show version and exit")):
4846
"""Codegen CLI - Transform your code with AI."""
4947
pass
5048

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

Lines changed: 1 addition & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -90,10 +90,7 @@ def get_config(key: str = typer.Argument(..., help="Configuration key to get")):
9090

9191

9292
@config_command.command(name="set")
93-
def set_config(
94-
key: str = typer.Argument(..., help="Configuration key to set"),
95-
value: str = typer.Argument(..., help="Configuration value to set")
96-
):
93+
def set_config(key: str = typer.Argument(..., help="Configuration key to set"), value: str = typer.Argument(..., help="Configuration value to set")):
9794
"""Set a configuration value and write to .env"""
9895
config = _get_user_config()
9996
if not config.has_key(key):

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

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,3 @@
1-
import sys
21
from pathlib import Path
32
from typing import Optional
43

@@ -9,18 +8,19 @@
98
from codegen.cli.rich.codeblocks import format_command
109
from codegen.shared.path import get_git_root_path
1110

11+
1212
def init(
1313
path: Optional[str] = typer.Option(None, help="Path within a git repository. Defaults to the current directory."),
1414
token: Optional[str] = typer.Option(None, help="Access token for the git repository. Required for full functionality."),
1515
language: Optional[str] = typer.Option(None, help="Override automatic language detection (python or typescript)"),
16-
fetch_docs: bool = typer.Option(False, "--fetch-docs", help="Fetch docs and examples (requires auth)")
16+
fetch_docs: bool = typer.Option(False, "--fetch-docs", help="Fetch docs and examples (requires auth)"),
1717
):
1818
"""Initialize or update the Codegen folder."""
1919
# Validate language option
2020
if language and language.lower() not in ["python", "typescript"]:
2121
rich.print(f"[bold red]Error:[/bold red] Invalid language '{language}'. Must be 'python' or 'typescript'.")
2222
raise typer.Exit(1)
23-
23+
2424
# Print a message if not in a git repo
2525
path_obj = Path.cwd() if path is None else Path(path)
2626
repo_path = get_git_root_path(path_obj)

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

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,12 @@
11
from typing import Optional
2-
import typer
2+
33
import rich
4+
import typer
45

56
from codegen.cli.auth.login import login_routine
67
from codegen.cli.auth.token_manager import get_current_token
78

9+
810
def login(token: Optional[str] = typer.Option(None, help="API token for authentication")):
911
"""Store authentication token."""
1012
# Check if already authenticated

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,8 @@
11
import rich
2-
import typer
32

43
from codegen.cli.auth.token_manager import TokenManager
54

5+
66
def logout():
77
"""Clear stored authentication token."""
88
token_manager = TokenManager()
Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,2 @@
1+
2+

0 commit comments

Comments
 (0)