diff --git a/.mcp.json b/.mcp.json new file mode 100644 index 00000000..56c5f62b --- /dev/null +++ b/.mcp.json @@ -0,0 +1,20 @@ +{ + "mcpServers": { + "codeflash": { + "type": "stdio", + "command": "codeflash", + "args": [ + "mcp" + ], + "env": {} + }, + "playwright": { + "type": "stdio", + "command": "npx", + "args": [ + "@playwright/mcp@latest" + ], + "env": {} + } + } +} \ No newline at end of file diff --git a/codeflash/cli_cmds/cli.py b/codeflash/cli_cmds/cli.py index 3f8f8767..02abdefa 100644 --- a/codeflash/cli_cmds/cli.py +++ b/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 diff --git a/codeflash/cli_cmds/cmd_init.py b/codeflash/cli_cmds/cmd_init.py index 61a949d0..1a8996a4 100644 --- a/codeflash/cli_cmds/cmd_init.py +++ b/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() diff --git a/myclient.py b/myclient.py new file mode 100644 index 00000000..784c9af5 --- /dev/null +++ b/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) diff --git a/myserver.py b/myserver.py new file mode 100644 index 00000000..50bda397 --- /dev/null +++ b/myserver.py @@ -0,0 +1,105 @@ +import sys +from contextlib import asynccontextmanager +from pathlib import Path +from typing import Any, AsyncGenerator, Optional + +from fastmcp import FastMCP +from lsprotocol.types import TextDocumentItem + +from codeflash.lsp.server import CodeflashLanguageServer +from codeflash.lsp.beta import perform_function_optimization, FunctionOptimizationParams, \ + initialize_function_optimization, validate_project, discover_function_tests +from tests.scripts.end_to_end_test_utilities import TestConfig, run_codeflash_command +from lsprotocol import types + +# dummy method for getting pyproject.toml path +def _find_pyproject_toml(workspace_path: str) -> Optional[Path]: + workspace_path_obj = Path(workspace_path) + for file_path in workspace_path_obj.rglob("pyproject.toml"): + return file_path.resolve() + return None + + +# 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.) + # server = CodeflashLanguageServer(name = "codeflash", version = "0.0.1") + # config_file = Path("/Users/codeflash/Downloads/codeflash-dev/codeflash/pyproject.toml") + # file = "/Users/codeflash/Downloads/codeflash-dev/codeflash/code_to_optimize/bubble_sort.py" + # function = "sorter" + # params = FunctionOptimizationParams(functionName=function, textDocument=types.TextDocumentIdentifier(Path(file).as_uri())) + # server.prepare_optimizer_arguments(config_file) + # initialize_function_optimization(server, params) + # perform_function_optimization(server, params) + # #optimize_code(file, function) + + #################### initialize the server ############################# + server = CodeflashLanguageServer("codeflash-language-server", "v1.0") + # suppose the pyproject.toml is in the current directory + server.prepare_optimizer_arguments(_find_pyproject_toml(".")) + result = validate_project(server, None) + if result["status"] == "error": + # handle if the project is not valid, it can be because pyproject.toml is not valid or the repository is in bare state or the repository has no commits, which will stop the worktree from working + print(result["message"]) + sys.exit(1) + + #################### start the optimization for file, function ############################# + file_path = "/Users/codeflash/Downloads/codeflash-dev/codeflash/code_to_optimize/bubble_sort.py" + function_name = "sorter" + + # This is not necessary, just for testing + server.args.module_root = Path("/Users/codeflash/Downloads/codeflash-dev/codeflash/code_to_optimize") + result = initialize_function_optimization(server, FunctionOptimizationParams( + functionName=function_name, + textDocument=TextDocumentItem( + uri=file_path, + language_id="python", + version=1, + text="" + ) + ) + ) + if result["status"] == "error": + # handle if the function is not optimizable + print(result["message"]) + sys.exit(1) + + discover_function_tests(server, FunctionOptimizationParams(functionName=function_name, textDocument=None)) + final_result = perform_function_optimization(server, FunctionOptimizationParams(functionName=function_name, + textDocument=None)) + if final_result["status"] == "success": + print(final_result) + yield + # Cleanup work after shutdown + print("Shutting down...") + server.cleanup_the_optimizer() + server.shutdown() + +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" + diff --git a/pyproject.toml b/pyproject.toml index efcacfcf..449fab54 100644 --- a/pyproject.toml +++ b/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 = "contact@codeflash.ai" }] -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]