Minimal MCP Server for tool calling codeflash with --file --function

codeflash/cli_cmds/cli.py

@@ -5,7 +5,7 @@
 
 from codeflash.cli_cmds import logging_config
 from codeflash.cli_cmds.cli_common import apologize_and_exit
-from codeflash.cli_cmds.cmd_init import init_codeflash, install_github_actions
+from codeflash.cli_cmds.cmd_init import init_codeflash, install_github_actions, launch_mcp
 from codeflash.cli_cmds.console import logger
 from codeflash.code_utils import env_utils
 from codeflash.code_utils.code_utils import exit_with_message
@@ -24,6 +24,9 @@ def parse_args() -> Namespace:
     init_actions_parser = subparsers.add_parser("init-actions", help="Initialize GitHub Actions workflow")
     init_actions_parser.set_defaults(func=install_github_actions)
 
+    mcp_parser = subparsers.add_parser("mcp", help="Launches the codeflash mcp server")
+    mcp_parser.set_defaults(func=launch_mcp)
+
     trace_optimize = subparsers.add_parser("optimize", help="Trace and optimize a Python project.")
 
     from codeflash.tracer import main as tracer_main

codeflash/cli_cmds/cmd_init.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import ast
+import json
 import os
 import re
 import subprocess
@@ -68,6 +69,12 @@ class DependencyManager(Enum):
     UNKNOWN = auto()
 
 
+def launch_mcp() -> None:
+    from myserver import mcp
+
+    mcp.run(transport="stdio")
+
+
 def init_codeflash() -> None:
     try:
         welcome_panel = Panel(
@@ -130,6 +137,9 @@ def init_codeflash() -> None:
         )
         console.print(completion_panel)
 
+        # Ask about adding MCP server to Claude configuration
+        prompt_claude_mcp_setup()
+
         ph("cli-installation-successful", {"did_add_new_key": did_add_new_key})
         sys.exit(0)
     except KeyboardInterrupt:
@@ -1229,3 +1239,93 @@ def ask_for_telemetry() -> bool:
         default=True,
         show_default=True,
     )
+
+
+def prompt_claude_mcp_setup() -> None:
+    """Prompt user to add Codeflash MCP server to their Claude configuration."""
+    from rich.prompt import Confirm
+
+    mcp_panel = Panel(
+        Text(
+            "🤖 Claude Code Integration\n\n"
+            "Would you like to add the Codeflash MCP server to your Claude Code configuration?\n"
+            "This will allow Claude Code to use Codeflash's optimization tools directly.",
+            style="bright_blue",
+        ),
+        title="🚀 Claude Code MCP Setup",
+        border_style="bright_blue",
+    )
+    console.print(mcp_panel)
+    console.print()
+
+    setup_mcp = Confirm.ask("Add Codeflash MCP server to Claude Code configuration?", default=True, show_default=True)
+
+    if setup_mcp:
+        try:
+            add_mcp_server_to_claude_config()
+        except Exception as e:
+            logger.error(f"Failed to add MCP server to Claude configuration: {e}")
+            console.print(
+                Panel(
+                    Text(
+                        "❌ Failed to add MCP server to Claude configuration.\n\n"
+                        "You can manually add it later by updating your Claude Code settings.",
+                        style="red",
+                    ),
+                    title="⚠️ Setup Failed",
+                    border_style="red",
+                )
+            )
+    else:
+        skip_panel = Panel(
+            Text("⏩️ Skipping Claude Code MCP setup.", style="yellow"), title="⏩️ Skipped", border_style="yellow"
+        )
+        console.print(skip_panel)
+
+
+def add_mcp_server_to_claude_config() -> None:
+    """Add the Codeflash MCP server to Claude Code configuration."""
+    claude_config_dir = Path.cwd()
+    config_file = claude_config_dir / ".mcp.json"
+
+    # Create MCP server configuration
+    # TODO we assume uv exists,
+    codeflash_server_entry = {"codeflash": {"type": "stdio", "command": "codeflash", "args": ["mcp"], "env": {}}}
+
+    # Read existing config or create new one
+    if config_file.exists():
+        try:
+            with config_file.open("r", encoding="utf8") as f:
+                updated_config = json.load(f)
+            if "mcpServers" not in updated_config:
+                updated_config["mcpServers"] = {}
+            updated_config["mcpServers"].update(codeflash_server_entry)
+        except (json.JSONDecodeError, OSError) as e:
+            logger.warning(f"Could not read existing MCP config: {e}")
+            updated_config = {"mcpServers": codeflash_server_entry}
+    else:
+        updated_config = {"mcpServers": codeflash_server_entry}
+
+    # Write the updated configuration
+    try:
+        with config_file.open("w", encoding="utf8") as f:
+            json.dump(updated_config, f, indent=2)
+
+        success_panel = Panel(
+            Text(
+                f"✅ Successfully added Codeflash MCP server to Claude Code configuration!\n\n"
+                f"Configuration saved to: {config_file}\n\n"
+                f"You can now use Codeflash optimization tools directly in Claude Code.",
+                style="green",
+                justify="left",
+            ),
+            title="🎉 MCP Setup Complete!",
+            border_style="bright_green",
+        )
+        console.print(success_panel)
+
+    except OSError as e:
+        error_str = f"Failed to write Claude Code MCP configuration: {e}"
+        raise RuntimeError(error_str) from e
+
+    console.print()

myclient.py

@@ -0,0 +1,18 @@
+import anthropic
+from rich import print as rprint
+
+url = "https://0de03d07f8b9.ngrok-free.app"
+client = anthropic.Anthropic()
+response = client.beta.messages.create(
+    model="claude-sonnet-4-20250514",
+    max_tokens=1000,
+    messages=[
+        {
+            "role": "user",
+            "content": 'Optimize my codebase, the file is "/Users/codeflash/Downloads/codeflash-dev/codeflash/code_to_optimize/bubble_sort.py" and the function is "sorter"',
+        }
+    ],
+    mcp_servers=[{"type": "url", "url": f"{url}/mcp/", "name": "codeflash"}],
+    extra_headers={"anthropic-beta": "mcp-client-2025-04-04"},
+)
+rprint(response.content)

myserver.py

@@ -0,0 +1,44 @@
+from contextlib import asynccontextmanager
+from pathlib import Path
+from typing import Any, AsyncGenerator
+
+from fastmcp import FastMCP
+
+from tests.scripts.end_to_end_test_utilities import TestConfig, run_codeflash_command
+
+
+# Define lifespan context manager
+@asynccontextmanager
+async def lifespan(mcp: FastMCP) -> AsyncGenerator[None, Any]:
+    print("Starting up...")
+    print(mcp.name)
+    # Do startup work here (connect to DB, initialize cache, etc.)
+    yield
+    # Cleanup work after shutdown
+    print("Shutting down...")
+
+
+mcp = FastMCP(
+    name="codeflash",
+    instructions="""
+        This server provides code optimization tools.
+        Call optimize_code(file, function) to optimize your code.
+    """,
+    lifespan=lifespan,
+)
+
+
+@mcp.tool
+def optimize_code(file: str, function: str) -> str:
+    # TODO ask for pr or no pr if successful
+    config = TestConfig(file_path=Path(f"{file}"), function_name=f"{function}", test_framework="pytest")
+    cwd = Path(file).resolve().parent
+    status = run_codeflash_command(cwd, config, expected_improvement_pct=5)
+    if status:
+        return "Optimization Successful, file has been edited"
+    return "Codeflash run did not meet expected requirements for testing, reverting file changes."
+
+
+if __name__ == "__main__":
+    mcp.run(transport="stdio")
+    # Optimize my codebase, the file is "/Users/codeflash/Downloads/codeflash-dev/codeflash/code_to_optimize/bubble_sort.py" and the function is "sorter"

pyproject.toml

@@ -3,7 +3,7 @@ name = "codeflash"
 dynamic = ["version"]
 description = "Client for codeflash.ai - automatic code performance optimization, powered by AI"
 authors = [{ name = "CodeFlash Inc.", email = "[email protected]" }]
-requires-python = ">=3.9"
+requires-python = ">=3.10"
 readme = "README.md"
 license = {text = "BSL-1.1"}
 keywords = [
@@ -43,6 +43,7 @@ dependencies = [
     "platformdirs>=4.3.7",
     "pygls>=1.3.1",
     "codeflash-benchmark",
+    "fastmcp"
 ]
 
 [project.urls]