feat: mcp server

Author: rushilpatel0

src/codegen/cli/api/client.py

@@ -15,6 +15,7 @@
     PR_LOOKUP_ENDPOINT,
     RUN_ENDPOINT,
     RUN_ON_PR_ENDPOINT,
+    IMPROVE_ENDPOINT,
 )
 from codegen.cli.api.schemas import (
     AskExpertInput,
@@ -36,12 +37,15 @@
     RunCodemodOutput,
     RunOnPRInput,
     RunOnPRResponse,
+    ImproveCodemodInput,
+    ImproveCodemodResponse,
 )
 from codegen.cli.auth.session import CodegenSession
 from codegen.cli.codemod.convert import convert_to_ui
 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 @@ class RestAPI:
 
     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]:
@@ -248,3 +252,13 @@ def lookup_pr(self, repo_full_name: str, github_pr_number: int) -> PRSchema:
             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,22 @@ 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,37 @@
+# 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
+```