From 5439c2d966a9d2cdfee5c411d8325d4cc29da6e0 Mon Sep 17 00:00:00 2001 From: aseembits93 Date: Tue, 12 Aug 2025 16:59:27 -0700 Subject: [PATCH 01/17] works for hardcoded file and fn --- myclient.py | 11 +++++++++++ myserver.py | 29 +++++++++++++++++++++++++++++ 2 files changed, 40 insertions(+) create mode 100644 myclient.py create mode 100644 myserver.py diff --git a/myclient.py b/myclient.py new file mode 100644 index 000000000..fcc0f71fc --- /dev/null +++ b/myclient.py @@ -0,0 +1,11 @@ +import asyncio +from fastmcp import Client + +client = Client("myserver.py") + +async def call_tool(file: str, function: str)-> None: + async with client: + result = await client.call_tool("optimize_code", {"file": f"{file}", "function":f"{function}"}) + print(result) + +asyncio.run(call_tool("Ford","Mustang")) diff --git a/myserver.py b/myserver.py new file mode 100644 index 000000000..1f34d95a3 --- /dev/null +++ b/myserver.py @@ -0,0 +1,29 @@ +import os +import pathlib + +from fastmcp import FastMCP + +from tests.scripts.end_to_end_test_utilities import CoverageExpectation, TestConfig, run_codeflash_command + +mcp = FastMCP("My MCP Server") + +@mcp.tool +def optimize_code(file: str, function: str) -> bool: #todo add file and function name as arguments + config = TestConfig( + file_path=pathlib.Path("/Users/codeflash/Downloads/codeflash-dev/codeflash/code_to_optimize/bubble_sort.py"), + function_name="sorter", + test_framework="pytest", + min_improvement_x=1.0, + coverage_expectations=[ + CoverageExpectation( + function_name="sorter", expected_coverage=100.0, expected_lines=[2, 3, 4, 5, 6, 7, 8, 9, 10] + ) + ], + ) + cwd = (pathlib.Path(__file__).parent/ "code_to_optimize").resolve() + return run_codeflash_command( + cwd, config, 100, ['print("codeflash stdout: Sorting list")', 'print(f"result: {arr}")'] + ) + +if __name__ == "__main__": + mcp.run() \ No newline at end of file From 36c83b333ad836555e0c63125ae8805958a3e502 Mon Sep 17 00:00:00 2001 From: aseembits93 Date: Tue, 12 Aug 2025 17:39:57 -0700 Subject: [PATCH 02/17] works without hardcoded, need to remove cwd line --- myclient.py | 71 +++++++++++++++++++++++++++++++++++++++++++++++------ myserver.py | 19 ++++++++------ 2 files changed, 74 insertions(+), 16 deletions(-) diff --git a/myclient.py b/myclient.py index fcc0f71fc..58a17e691 100644 --- a/myclient.py +++ b/myclient.py @@ -1,11 +1,66 @@ -import asyncio -from fastmcp import Client +# import asyncio +# from fastmcp import Client +# +# client = Client("myserver.py") +# +# async def call_tool(file: str, function: str)-> None: +# async with client: +# result = await client.call_tool("optimize_code", {"file": f"{file}", "function":f"{function}"}) +# print(result) +# +# asyncio.run(call_tool("Ford","Mustang")) -client = Client("myserver.py") -async def call_tool(file: str, function: str)-> None: - async with client: - result = await client.call_tool("optimize_code", {"file": f"{file}", "function":f"{function}"}) - print(result) +import anthropic +from rich import print -asyncio.run(call_tool("Ford","Mustang")) +# Your server URL (replace with your actual URL) +url = 'https://0de03d07f8b9.ngrok-free.app' + +client = anthropic.Anthropic() + +#file and fn names are hard coded for now +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": "HelpfulAssistant", + } + ], + extra_headers={ + "anthropic-beta": "mcp-client-2025-04-04" + } +) + +print(response.content) + + +# import anthropic +# from rich import print +# +# # Your server URL (replace with your actual URL) +# 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": "Roll a few dice!"}], +# mcp_servers=[ +# { +# "type": "url", +# "url": f"{url}/mcp/", +# "name": "Dice Roller", +# } +# ], +# extra_headers={ +# "anthropic-beta": "mcp-client-2025-04-04" +# } +# ) +# +# print(response.content) \ No newline at end of file diff --git a/myserver.py b/myserver.py index 1f34d95a3..71a6b8a41 100644 --- a/myserver.py +++ b/myserver.py @@ -1,29 +1,32 @@ -import os import pathlib from fastmcp import FastMCP from tests.scripts.end_to_end_test_utilities import CoverageExpectation, TestConfig, run_codeflash_command -mcp = FastMCP("My MCP Server") +mcp = FastMCP(name="HelpfulAssistant", + instructions=""" + This server provides code optimization tools. + Call optimize_code(file, function) to optimize your code. + """,) @mcp.tool -def optimize_code(file: str, function: str) -> bool: #todo add file and function name as arguments +def optimize_code(file: str, function: str) -> bool: config = TestConfig( - file_path=pathlib.Path("/Users/codeflash/Downloads/codeflash-dev/codeflash/code_to_optimize/bubble_sort.py"), - function_name="sorter", + file_path=pathlib.Path(file), + function_name=function, test_framework="pytest", min_improvement_x=1.0, coverage_expectations=[ CoverageExpectation( - function_name="sorter", expected_coverage=100.0, expected_lines=[2, 3, 4, 5, 6, 7, 8, 9, 10] + function_name=function, expected_coverage=100.0, expected_lines=[2, 3, 4, 5, 6, 7, 8, 9, 10] ) ], ) - cwd = (pathlib.Path(__file__).parent/ "code_to_optimize").resolve() + cwd = (pathlib.Path(__file__).parent/ "code_to_optimize").resolve() # todo remove it return run_codeflash_command( cwd, config, 100, ['print("codeflash stdout: Sorting list")', 'print(f"result: {arr}")'] ) if __name__ == "__main__": - mcp.run() \ No newline at end of file + mcp.run(transport="http", port=8000) \ No newline at end of file From f65bfc893f8fd3c65eb565a1104887201169503e Mon Sep 17 00:00:00 2001 From: aseembits93 Date: Tue, 12 Aug 2025 17:42:52 -0700 Subject: [PATCH 03/17] cleaning up --- myclient.py | 50 ++------------------------------------------------ myserver.py | 9 +++------ 2 files changed, 5 insertions(+), 54 deletions(-) diff --git a/myclient.py b/myclient.py index 58a17e691..7076d31ef 100644 --- a/myclient.py +++ b/myclient.py @@ -1,25 +1,7 @@ -# import asyncio -# from fastmcp import Client -# -# client = Client("myserver.py") -# -# async def call_tool(file: str, function: str)-> None: -# async with client: -# result = await client.call_tool("optimize_code", {"file": f"{file}", "function":f"{function}"}) -# print(result) -# -# asyncio.run(call_tool("Ford","Mustang")) - - import anthropic from rich import print - -# Your server URL (replace with your actual URL) url = 'https://0de03d07f8b9.ngrok-free.app' - client = anthropic.Anthropic() - -#file and fn names are hard coded for now response = client.beta.messages.create( model="claude-sonnet-4-20250514", max_tokens=1000, @@ -28,39 +10,11 @@ { "type": "url", "url": f"{url}/mcp/", - "name": "HelpfulAssistant", + "name": "Code Optimization Assistant", } ], extra_headers={ "anthropic-beta": "mcp-client-2025-04-04" } ) - -print(response.content) - - -# import anthropic -# from rich import print -# -# # Your server URL (replace with your actual URL) -# 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": "Roll a few dice!"}], -# mcp_servers=[ -# { -# "type": "url", -# "url": f"{url}/mcp/", -# "name": "Dice Roller", -# } -# ], -# extra_headers={ -# "anthropic-beta": "mcp-client-2025-04-04" -# } -# ) -# -# print(response.content) \ No newline at end of file +print(response.content) \ No newline at end of file diff --git a/myserver.py b/myserver.py index 71a6b8a41..abcad4856 100644 --- a/myserver.py +++ b/myserver.py @@ -1,14 +1,11 @@ import pathlib - from fastmcp import FastMCP - from tests.scripts.end_to_end_test_utilities import CoverageExpectation, TestConfig, run_codeflash_command - -mcp = FastMCP(name="HelpfulAssistant", +mcp = FastMCP(name="Code Optimization Assistant", instructions=""" This server provides code optimization tools. Call optimize_code(file, function) to optimize your code. - """,) + """) @mcp.tool def optimize_code(file: str, function: str) -> bool: @@ -23,7 +20,7 @@ def optimize_code(file: str, function: str) -> bool: ) ], ) - cwd = (pathlib.Path(__file__).parent/ "code_to_optimize").resolve() # todo remove it + cwd = (pathlib.Path(__file__).parent/ "code_to_optimize").resolve() # TODO remove it return run_codeflash_command( cwd, config, 100, ['print("codeflash stdout: Sorting list")', 'print(f"result: {arr}")'] ) From f41af3d650f7c964d28c48dd314fbd1376009f92 Mon Sep 17 00:00:00 2001 From: aseembits93 Date: Tue, 12 Aug 2025 17:53:36 -0700 Subject: [PATCH 04/17] precommit mypy fix --- myclient.py | 20 +++++++++----------- myserver.py | 15 +++++++++++---- 2 files changed, 20 insertions(+), 15 deletions(-) diff --git a/myclient.py b/myclient.py index 7076d31ef..faf8517ed 100644 --- a/myclient.py +++ b/myclient.py @@ -1,20 +1,18 @@ import anthropic -from rich import print -url = 'https://0de03d07f8b9.ngrok-free.app' +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=[ + messages=[ { - "type": "url", - "url": f"{url}/mcp/", - "name": "Code Optimization Assistant", + "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"', } ], - extra_headers={ - "anthropic-beta": "mcp-client-2025-04-04" - } + mcp_servers=[{"type": "url", "url": f"{url}/mcp/", "name": "Code Optimization Assistant"}], + extra_headers={"anthropic-beta": "mcp-client-2025-04-04"}, ) -print(response.content) \ No newline at end of file +rprint(response.content) diff --git a/myserver.py b/myserver.py index abcad4856..c4545d6b6 100644 --- a/myserver.py +++ b/myserver.py @@ -1,11 +1,17 @@ import pathlib + from fastmcp import FastMCP + from tests.scripts.end_to_end_test_utilities import CoverageExpectation, TestConfig, run_codeflash_command -mcp = FastMCP(name="Code Optimization Assistant", + +mcp = FastMCP( + name="Code Optimization Assistant", instructions=""" This server provides code optimization tools. Call optimize_code(file, function) to optimize your code. - """) + """, +) + @mcp.tool def optimize_code(file: str, function: str) -> bool: @@ -20,10 +26,11 @@ def optimize_code(file: str, function: str) -> bool: ) ], ) - cwd = (pathlib.Path(__file__).parent/ "code_to_optimize").resolve() # TODO remove it + cwd = (pathlib.Path(__file__).parent / "code_to_optimize").resolve() # TODO remove it return run_codeflash_command( cwd, config, 100, ['print("codeflash stdout: Sorting list")', 'print(f"result: {arr}")'] ) + if __name__ == "__main__": - mcp.run(transport="http", port=8000) \ No newline at end of file + mcp.run(transport="http", port=8000) From 25e215720b9f8542a08b199b97437e90f72b861e Mon Sep 17 00:00:00 2001 From: aseembits93 Date: Thu, 14 Aug 2025 16:12:54 -0700 Subject: [PATCH 05/17] quick fix --- myserver.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/myserver.py b/myserver.py index c4545d6b6..f34bc1ffd 100644 --- a/myserver.py +++ b/myserver.py @@ -5,7 +5,7 @@ from tests.scripts.end_to_end_test_utilities import CoverageExpectation, TestConfig, run_codeflash_command mcp = FastMCP( - name="Code Optimization Assistant", + name="codeflash", instructions=""" This server provides code optimization tools. Call optimize_code(file, function) to optimize your code. From ffbd894af930abc94128842bd91aa0977a79e273 Mon Sep 17 00:00:00 2001 From: aseembits93 Date: Thu, 14 Aug 2025 17:14:28 -0700 Subject: [PATCH 06/17] typo --- myclient.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/myclient.py b/myclient.py index faf8517ed..784c9af5c 100644 --- a/myclient.py +++ b/myclient.py @@ -12,7 +12,7 @@ "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": "Code Optimization Assistant"}], + mcp_servers=[{"type": "url", "url": f"{url}/mcp/", "name": "codeflash"}], extra_headers={"anthropic-beta": "mcp-client-2025-04-04"}, ) rprint(response.content) From e8b9202ffe6188726693d78d377160d375e293c7 Mon Sep 17 00:00:00 2001 From: Aseem Saxena Date: Wed, 20 Aug 2025 12:01:21 -0700 Subject: [PATCH 07/17] stdio server --- codeflash/version.py | 2 +- myserver.py | 45 +++++++++++++++++++++++--------------------- pyproject.toml | 3 ++- 3 files changed, 27 insertions(+), 23 deletions(-) diff --git a/codeflash/version.py b/codeflash/version.py index 42ce5de3f..741acda63 100644 --- a/codeflash/version.py +++ b/codeflash/version.py @@ -1,2 +1,2 @@ # These version placeholders will be replaced by uv-dynamic-versioning during build. -__version__ = "0.16.3" +__version__ = "0.16.5" diff --git a/myserver.py b/myserver.py index f34bc1ffd..972ab8ac4 100644 --- a/myserver.py +++ b/myserver.py @@ -1,9 +1,5 @@ -import pathlib - from fastmcp import FastMCP -from tests.scripts.end_to_end_test_utilities import CoverageExpectation, TestConfig, run_codeflash_command - mcp = FastMCP( name="codeflash", instructions=""" @@ -14,23 +10,30 @@ @mcp.tool -def optimize_code(file: str, function: str) -> bool: - config = TestConfig( - file_path=pathlib.Path(file), - function_name=function, - test_framework="pytest", - min_improvement_x=1.0, - coverage_expectations=[ - CoverageExpectation( - function_name=function, expected_coverage=100.0, expected_lines=[2, 3, 4, 5, 6, 7, 8, 9, 10] - ) - ], - ) - cwd = (pathlib.Path(__file__).parent / "code_to_optimize").resolve() # TODO remove it - return run_codeflash_command( - cwd, config, 100, ['print("codeflash stdout: Sorting list")', 'print(f"result: {arr}")'] - ) +def optimize_code(file: str, function: str) -> dict[str, str]: + # config = TestConfig( + # file_path=pathlib.Path(file), + # function_name=function, + # test_framework="pytest", + # min_improvement_x=1.0, + # coverage_expectations=[ + # CoverageExpectation( + # function_name=function, expected_coverage=100.0, expected_lines=[2, 3, 4, 5, 6, 7, 8, 9, 10] + # ) + # ], + # ) + # cwd = pathlib.Path("/Users/aseemsaxena/Downloads/codeflash_dev/codeflash/code_to_optimize") # TODO remove it + print(file, function) + return { + "code": """def sorter(arr): + print("codeflash stdout: Sorting list") + arr.sort() + print(f"result: {arr}") + return arr +""", + "explanation": "A faster version of the code is using python's in-built timsort function", + } if __name__ == "__main__": - mcp.run(transport="http", port=8000) + mcp.run(transport="stdio") diff --git a/pyproject.toml b/pyproject.toml index 3333a3b68..cf52b2911 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] From 75c919b5e702d2ef18d9a10665b93739f48ab1f2 Mon Sep 17 00:00:00 2001 From: Aseem Saxena Date: Wed, 20 Aug 2025 12:19:01 -0700 Subject: [PATCH 08/17] vibecoded by claude code --- codeflash/cli_cmds/cmd_init.py | 114 +++++++++++++++++++++++++++++++++ 1 file changed, 114 insertions(+) diff --git a/codeflash/cli_cmds/cmd_init.py b/codeflash/cli_cmds/cmd_init.py index 57133fa16..5e15ddfd7 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 @@ -130,6 +131,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: @@ -1215,3 +1219,113 @@ 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.""" + # Try to find Claude Code config directory + home_dir = Path.home() + claude_config_dirs = [ + home_dir / ".config" / "claude-code", + home_dir / ".claude-code", + home_dir / "Library" / "Application Support" / "claude-code", # macOS + ] + + claude_config_dir = None + for config_dir in claude_config_dirs: + if config_dir.exists(): + claude_config_dir = config_dir + break + + # If no existing config directory found, create one + if not claude_config_dir: + claude_config_dir = home_dir / ".config" / "claude-code" + claude_config_dir.mkdir(parents=True, exist_ok=True) + console.print(f"📁 Created Claude Code config directory: {claude_config_dir}") + + config_file = claude_config_dir / "mcp_servers.json" + + # Get the absolute path to myserver.py + myserver_path = Path.cwd() / "myserver.py" + + # Create MCP server configuration + server_config = {"codeflash": {"command": "python", "args": [str(myserver_path.absolute())], "env": {}}} + + # Read existing config or create new one + if config_file.exists(): + try: + with config_file.open("r", encoding="utf8") as f: + existing_config = json.load(f) + existing_config.update(server_config) + updated_config = existing_config + except (json.JSONDecodeError, OSError) as e: + logger.warning(f"Could not read existing MCP config: {e}") + updated_config = server_config + else: + updated_config = server_config + + # 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() From 81b8dab718198498c65d0ce9e60376b74873b430 Mon Sep 17 00:00:00 2001 From: Aseem Saxena Date: Wed, 20 Aug 2025 12:53:58 -0700 Subject: [PATCH 09/17] vibecoded by gemini code assist --- codeflash/cli_cmds/cmd_init.py | 47 ++++++++++++---------------------- 1 file changed, 17 insertions(+), 30 deletions(-) diff --git a/codeflash/cli_cmds/cmd_init.py b/codeflash/cli_cmds/cmd_init.py index 5e15ddfd7..7703c2c61 100644 --- a/codeflash/cli_cmds/cmd_init.py +++ b/codeflash/cli_cmds/cmd_init.py @@ -1265,46 +1265,33 @@ def prompt_claude_mcp_setup() -> None: def add_mcp_server_to_claude_config() -> None: """Add the Codeflash MCP server to Claude Code configuration.""" - # Try to find Claude Code config directory - home_dir = Path.home() - claude_config_dirs = [ - home_dir / ".config" / "claude-code", - home_dir / ".claude-code", - home_dir / "Library" / "Application Support" / "claude-code", # macOS - ] - - claude_config_dir = None - for config_dir in claude_config_dirs: - if config_dir.exists(): - claude_config_dir = config_dir - break - - # If no existing config directory found, create one - if not claude_config_dir: - claude_config_dir = home_dir / ".config" / "claude-code" - claude_config_dir.mkdir(parents=True, exist_ok=True) - console.print(f"📁 Created Claude Code config directory: {claude_config_dir}") - - config_file = claude_config_dir / "mcp_servers.json" - - # Get the absolute path to myserver.py - myserver_path = Path.cwd() / "myserver.py" + claude_config_dir = Path.cwd() + config_file = claude_config_dir / ".mcp.json" # Create MCP server configuration - server_config = {"codeflash": {"command": "python", "args": [str(myserver_path.absolute())], "env": {}}} + # TODO we assume uv exists, + codeflash_server_entry = { + "codeflash": { + "type": "stdio", + "command": "uv", + "args": ["run", "--directory", str(claude_config_dir), "myserver.py"], + "env": {}, + } + } # Read existing config or create new one if config_file.exists(): try: with config_file.open("r", encoding="utf8") as f: - existing_config = json.load(f) - existing_config.update(server_config) - updated_config = existing_config + 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 = server_config + updated_config = {"mcpServers": codeflash_server_entry} else: - updated_config = server_config + updated_config = {"mcpServers": codeflash_server_entry} # Write the updated configuration try: From e3697ff0abf74f71760c2bdf19cef720193dbbee Mon Sep 17 00:00:00 2001 From: Aseem Saxena Date: Wed, 20 Aug 2025 21:49:37 -0700 Subject: [PATCH 10/17] codeflash itself can launch mcp --- codeflash/cli_cmds/cli.py | 5 ++++- codeflash/cli_cmds/cmd_init.py | 15 +++++++-------- myserver.py | 12 ++---------- 3 files changed, 13 insertions(+), 19 deletions(-) diff --git a/codeflash/cli_cmds/cli.py b/codeflash/cli_cmds/cli.py index 5aa331f36..76a5e9954 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 @@ -23,6 +23,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 7703c2c61..1d41e04dd 100644 --- a/codeflash/cli_cmds/cmd_init.py +++ b/codeflash/cli_cmds/cmd_init.py @@ -69,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( @@ -1270,14 +1276,7 @@ def add_mcp_server_to_claude_config() -> None: # Create MCP server configuration # TODO we assume uv exists, - codeflash_server_entry = { - "codeflash": { - "type": "stdio", - "command": "uv", - "args": ["run", "--directory", str(claude_config_dir), "myserver.py"], - "env": {}, - } - } + codeflash_server_entry = {"codeflash": {"type": "stdio", "command": "codeflash", "args": ["mcp"], "env": {}}} # Read existing config or create new one if config_file.exists(): diff --git a/myserver.py b/myserver.py index 972ab8ac4..dc08f01d1 100644 --- a/myserver.py +++ b/myserver.py @@ -10,7 +10,7 @@ @mcp.tool -def optimize_code(file: str, function: str) -> dict[str, str]: +def optimize_code(file: str, function: str) -> str: # config = TestConfig( # file_path=pathlib.Path(file), # function_name=function, @@ -24,15 +24,7 @@ def optimize_code(file: str, function: str) -> dict[str, str]: # ) # cwd = pathlib.Path("/Users/aseemsaxena/Downloads/codeflash_dev/codeflash/code_to_optimize") # TODO remove it print(file, function) - return { - "code": """def sorter(arr): - print("codeflash stdout: Sorting list") - arr.sort() - print(f"result: {arr}") - return arr -""", - "explanation": "A faster version of the code is using python's in-built timsort function", - } + return "the function is already optimal" if __name__ == "__main__": From a4b12b45865eac1f95d82daf589a512b93ce326f Mon Sep 17 00:00:00 2001 From: aseembits93 Date: Thu, 21 Aug 2025 13:35:07 -0700 Subject: [PATCH 11/17] works with absolute paths inside codeflash --- myserver.py | 30 ++++++++++++++++-------------- 1 file changed, 16 insertions(+), 14 deletions(-) diff --git a/myserver.py b/myserver.py index dc08f01d1..23040b474 100644 --- a/myserver.py +++ b/myserver.py @@ -1,4 +1,7 @@ from fastmcp import FastMCP +from pathlib import Path + +from tests.scripts.end_to_end_test_utilities import TestConfig, run_codeflash_command mcp = FastMCP( name="codeflash", @@ -11,21 +14,20 @@ @mcp.tool def optimize_code(file: str, function: str) -> str: - # config = TestConfig( - # file_path=pathlib.Path(file), - # function_name=function, - # test_framework="pytest", - # min_improvement_x=1.0, - # coverage_expectations=[ - # CoverageExpectation( - # function_name=function, expected_coverage=100.0, expected_lines=[2, 3, 4, 5, 6, 7, 8, 9, 10] - # ) - # ], - # ) - # cwd = pathlib.Path("/Users/aseemsaxena/Downloads/codeflash_dev/codeflash/code_to_optimize") # TODO remove it - print(file, function) - return "the function is already optimal" + # 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: + "Optimization Successful, file has been edited" + else: + 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" \ No newline at end of file From d23c73922bf066a27cc2ceb6ed008e3f57da485b Mon Sep 17 00:00:00 2001 From: aseembits93 Date: Thu, 21 Aug 2025 13:42:32 -0700 Subject: [PATCH 12/17] minor bugfix --- myserver.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/myserver.py b/myserver.py index 23040b474..fed2dc432 100644 --- a/myserver.py +++ b/myserver.py @@ -23,7 +23,7 @@ def optimize_code(file: str, function: str) -> str: cwd = Path(file).resolve().parent status = run_codeflash_command(cwd, config, expected_improvement_pct=5) if status: - "Optimization Successful, file has been edited" + return "Optimization Successful, file has been edited" else: return "Codeflash run did not meet expected requirements for testing, reverting file changes." From c51782c6c6bb1352b1c2d58e01f6bc2038f636b9 Mon Sep 17 00:00:00 2001 From: aseembits93 Date: Wed, 27 Aug 2025 16:31:35 -0700 Subject: [PATCH 13/17] lifespan for graceful entry and exit of mcp server --- myserver.py | 28 +++++++++++++++++++--------- 1 file changed, 19 insertions(+), 9 deletions(-) diff --git a/myserver.py b/myserver.py index fed2dc432..4dc771746 100644 --- a/myserver.py +++ b/myserver.py @@ -1,33 +1,43 @@ -from fastmcp import FastMCP +from contextlib import asynccontextmanager from pathlib import Path +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) -> None: + 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", - ) + 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" - else: - return "Codeflash run did not meet expected requirements for testing, reverting file changes." + 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" \ No newline at end of file + # Optimize my codebase, the file is "/Users/codeflash/Downloads/codeflash-dev/codeflash/code_to_optimize/bubble_sort.py" and the function is "sorter" From 96313f2cbfc5554cda039f72f9a64c7887a9303c Mon Sep 17 00:00:00 2001 From: aseembits93 Date: Wed, 27 Aug 2025 16:37:03 -0700 Subject: [PATCH 14/17] type annotation fix --- myserver.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/myserver.py b/myserver.py index 4dc771746..48920cc4a 100644 --- a/myserver.py +++ b/myserver.py @@ -1,5 +1,6 @@ from contextlib import asynccontextmanager from pathlib import Path +from typing import Any, AsyncGenerator from fastmcp import FastMCP @@ -8,7 +9,7 @@ # Define lifespan context manager @asynccontextmanager -async def lifespan(mcp: FastMCP) -> None: +async def lifespan(mcp: FastMCP) -> AsyncGenerator[None, Any]: print("Starting up...") print(mcp.name) # Do startup work here (connect to DB, initialize cache, etc.) From eebc4dcb034f6acc4d10f52c3c8efb23387382ed Mon Sep 17 00:00:00 2001 From: aseembits93 Date: Thu, 28 Aug 2025 13:21:34 -0700 Subject: [PATCH 15/17] interface with lsp --- myserver.py | 19 +++++++++++++++++-- 1 file changed, 17 insertions(+), 2 deletions(-) diff --git a/myserver.py b/myserver.py index 48920cc4a..7a476d8cf 100644 --- a/myserver.py +++ b/myserver.py @@ -4,7 +4,11 @@ from fastmcp import FastMCP +from codeflash.lsp.server import CodeflashLanguageServer +from codeflash.lsp.beta import perform_function_optimization, FunctionOptimizationParams, \ + initialize_function_optimization from tests.scripts.end_to_end_test_utilities import TestConfig, run_codeflash_command +from lsprotocol import types # Define lifespan context manager @@ -13,10 +17,20 @@ 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) yield # Cleanup work after shutdown print("Shutting down...") - + server.cleanup_the_optimizer() + server.shutdown() mcp = FastMCP( name="codeflash", @@ -28,7 +42,7 @@ async def lifespan(mcp: FastMCP) -> AsyncGenerator[None, Any]: ) -@mcp.tool +#@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") @@ -42,3 +56,4 @@ def optimize_code(file: str, function: str) -> str: 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" + From d510627a3161775b9099a47881958343c1913c31 Mon Sep 17 00:00:00 2001 From: aseembits93 Date: Thu, 28 Aug 2025 13:31:40 -0700 Subject: [PATCH 16/17] bugfix --- myserver.py | 70 ++++++++++++++++++++++++++++++++++++++++++++--------- 1 file changed, 58 insertions(+), 12 deletions(-) diff --git a/myserver.py b/myserver.py index 7a476d8cf..50bda397b 100644 --- a/myserver.py +++ b/myserver.py @@ -1,31 +1,77 @@ +import sys from contextlib import asynccontextmanager from pathlib import Path -from typing import Any, AsyncGenerator +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 + 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) + # # 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...") From 2a21f56d3dae8ce91416ca90e2efcb00d9ebcbde Mon Sep 17 00:00:00 2001 From: aseembits93 Date: Thu, 28 Aug 2025 13:32:51 -0700 Subject: [PATCH 17/17] add mcp json --- .mcp.json | 20 ++++++++++++++++++++ 1 file changed, 20 insertions(+) create mode 100644 .mcp.json diff --git a/.mcp.json b/.mcp.json new file mode 100644 index 000000000..56c5f62b1 --- /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