feat: [CG-10632] mcp server

docs/introduction/ide-usage.mdx

@@ -51,6 +51,46 @@ Codegen creates a custom Python environment in `.codegen/.venv`. Configure your
 </AccordionGroup>
 
 
+## MCP Server Setup
+This is an optional step but highly recommended if your IDE supports MCP support and you use AI Agents.
+The MCP server is a local server that allows your AI Agent to interact with the Codegen specific tools,
+it will allow an agent to:
+- ask an expert to create a codemod
+- improve a codemod
+- get setup instructions
+
+### Configuration
+#### Usage with Cline:
+Add this to your cline_mcp_settings.json:
+```
+{
+  "mcpServers": {
+    "codegen-cli": {
+        "command": "uv",
+        "args": [
+            "--directory",
+            "<path to codegen installation>/codegen-sdk/src/codegen/cli/mcp",
+            "run",
+            "server.py"
+        ]
+    }
+  }
+}
+```
+
+
+#### Usage with Cursor:
+Under the `Settings` > `Feature` > `MCP Servers` section, click "Add New MCP Server" and add the following:
+
+```
+Name: codegen-mcp
+Type: Command
+Command: uv --directory <path to codegen installation>/codegen-sdk/src/codegen/cli/mcp run server.py
+```
+
+
+
+
 ## Create a New Codemod
 
 Generate the boilerplate for a new code manipulation program using [codegen create](/cli/create):

src/codegen/cli/api/client.py

@@ -11,6 +11,7 @@
     DOCS_ENDPOINT,
     EXPERT_ENDPOINT,
     IDENTIFY_ENDPOINT,
+    IMPROVE_ENDPOINT,
     LOOKUP_ENDPOINT,
     PR_LOOKUP_ENDPOINT,
     RUN_ENDPOINT,
@@ -27,6 +28,8 @@
     DocsInput,
     DocsResponse,
     IdentifyResponse,
+    ImproveCodemodInput,
+    ImproveCodemodResponse,
     LookupInput,
     LookupOutput,
     PRLookupInput,
@@ -42,6 +45,7 @@
 from codegen.cli.env.global_env import global_env
 from codegen.cli.errors import InvalidTokenError, ServerError
 from codegen.cli.utils.codemods import Codemod
+from codegen.cli.utils.constants import ProgrammingLanguage
 from codegen.cli.utils.function_finder import DecoratedFunction
 
 InputT = TypeVar("InputT", bound=BaseModel)
@@ -55,7 +59,7 @@
 
     auth_token: str | None = None
 
-    def __init__(self, auth_token: str):
+    def __init__(self, auth_token: str | None = None):
         self.auth_token = auth_token
 
     def _get_headers(self) -> dict[str, str]:
@@ -144,7 +148,7 @@
 
         # Add template context if provided
         if template_context:
             base_input["template_context"] = template_context
 
         input_data = RunCodemodInput(input=RunCodemodInput.BaseRunCodemodInput(**base_input))
         return self._make_request(
@@ -157,7 +161,7 @@
     def get_docs(self) -> dict:
         """Search documentation."""
         session = CodegenSession()
         return self._make_request(
             "GET",
             DOCS_ENDPOINT,
             DocsInput(docs_input=DocsInput.BaseDocsInput(repo_full_name=session.repo_name)),
@@ -242,9 +246,18 @@
 
     def lookup_pr(self, repo_full_name: str, github_pr_number: int) -> PRSchema:
         """Look up a PR by repository and PR number."""
         return self._make_request(
             "GET",
             PR_LOOKUP_ENDPOINT,
             PRLookupInput(input=PRLookupInput.BasePRLookupInput(repo_full_name=repo_full_name, github_pr_number=github_pr_number)),
             PRLookupResponse,
         )
+
+    def improve_codemod(self, codemod: str, task: str, concerns: list[str], context: dict[str, str], language: ProgrammingLanguage) -> ImproveCodemodResponse:
+        """Improve a codemod."""
+        return self._make_request(
+            "GET",
+            IMPROVE_ENDPOINT,
+            ImproveCodemodInput(input=ImproveCodemodInput.BaseImproveCodemodInput(codemod=codemod, task=task, concerns=concerns, context=context, language=language)),
+            ImproveCodemodResponse,
+        )

src/codegen/cli/api/endpoints.py

@@ -10,3 +10,4 @@
 RUN_ON_PR_ENDPOINT = f"https://{MODAL_PREFIX}--cli-run-on-pull-request.modal.run"
 PR_LOOKUP_ENDPOINT = f"https://{MODAL_PREFIX}--cli-pr-lookup.modal.run"
 CODEGEN_SYSTEM_PROMPT_URL = "https://gist.githubusercontent.com/jayhack/15681a2ceaccd726f19e6fdb3a44738b/raw/17c08054e3931b3b7fdf424458269c9e607541e8/codegen-system-prompt.txt"
+IMPROVE_ENDPOINT = f"https://{MODAL_PREFIX}--cli-improve.modal.run"

src/codegen/cli/api/schemas.py

@@ -234,3 +234,24 @@ class RunOnPRResponse(BaseModel):
     codemod_id: int = Field(..., description="ID of the codemod")
     codemod_run_id: int = Field(..., description="ID of the codemod run")
     web_url: str = Field(..., description="URL to view the test results")
+
+
+###########################################################################
+# IMPROVE
+###########################################################################
+
+
+class ImproveCodemodInput(BaseModel):
+    class BaseImproveCodemodInput(BaseModel):
+        codemod: str = Field(..., description="Source code of the codemod to improve")
+        task: str = Field(..., description="Task to which the codemod should implement to solve")
+        concerns: list[str] = Field(..., description="A list of issues that were discovered with the current codemod that need to be considered in the next iteration")
+        context: dict[str, str] = Field(..., description="Additional context for the codemod this can be a list of files that are related, additional information about the task, etc.")
+        language: ProgrammingLanguage = Field(..., description="Language of the codemod")
+
+    input: BaseImproveCodemodInput = Field(..., description="Input data for improvement")
+
+
+class ImproveCodemodResponse(BaseModel):
+    success: bool = Field(..., description="Whether the improvement was successful")
+    codemod_source: str = Field(..., description="Source code of the improved codemod")

src/codegen/cli/mcp/README.md

@@ -0,0 +1,40 @@
+# Codegen MCP server
+
+A MCP server implementation that provides tools and resources for using and working with the Codegen CLI and SDK, enabling AI agents to iterate quickly on writing codemods with the codegen sdk.
+
+### Dependencies
+
+- [fastmcp](https://github.com/codegen-sh/fastmcp)
+
+## Usage
+
+Most AI Agents that support MCP will have some way to configure the server startup.
+
+### Cline
+
+Add this to your `cline_mcp_settings.json` file to get started:
+
+```
+{
+  "mcpServers": {
+    "codegen-cli": {
+        "command": "uv",
+        "args": [
+            "--directory",
+            "<path to codegen installation>/codegen-sdk/src/codegen/cli/mcp",
+            "run",
+            "server.py"
+        ]
+    }
+  }
+}
+```
+
+Cursor:
+Under the `Settings` > `Feature` > `MCP Servers` section, click "Add New MCP Server" and add the following:
+
+```
+Name: codegen-mcp
+Type: Command
+Command: uv --directory <path to codegen installation>/codegen-sdk/src/codegen/cli/mcp run server.py
+```