re-use same db connection and app lifespan

Author: mohammedahmed18

mcp_server/db.py

@@ -14,6 +14,7 @@
 
 _DEFAULT_DB_PATH = Path.home() / ".codeflash" / "mcp_results.db"
 
+db_conn = None
 
 def get_db_path() -> Path:
     path = Path(os.environ.get("CODEFLASH_MCP_DB_PATH", str(_DEFAULT_DB_PATH)))
@@ -22,14 +23,26 @@ def get_db_path() -> Path:
 
 
 def get_connection() -> sqlite3.Connection:
+    global db_conn
+    if db_conn is not None:
+        return db_conn
     db_path = get_db_path()
+    print(db_path)
     conn = sqlite3.connect(str(db_path))
     conn.execute("PRAGMA journal_mode=WAL")
     conn.execute("PRAGMA foreign_keys=ON")
     _create_tables(conn)
+    db_conn = conn
     return conn
 
 
+def close_conn() -> None:
+    global db_conn
+    if db_conn is not None:
+        db_conn.close()
+        db_conn = None
+
+
 def _create_tables(conn: sqlite3.Connection) -> None:
     conn.executescript("""
         CREATE TABLE IF NOT EXISTS runs (

mcp_server/server.py

@@ -3,10 +3,24 @@
 from typing import Any
 
 from fastmcp import FastMCP
+from fastmcp.server.mixins import lifespan
+
+from mcp_server.db import close_conn
+
+
+@lifespan
+async def app_lifespan(server) -> None:
+    print("codeflash-mcp is up")
+    yield
+    print("shutting dowm the server")
+    close_conn()
+
+
 
 mcp = FastMCP(
     "codeflash-mcp",
     instructions="Run behavioral tests, compare results, and benchmark performance for code optimization.",
+    lifespan=app_lifespan,
 )
 
 

mcp_server/test_mcp_workflow.py

@@ -21,7 +21,7 @@ def get_e2e_test_config_by_language(language: str) -> dict[str, str]:
             "module_path": str(CODE_DIR / "bubble_sort.py"),
             "function_name": "sorter",
             "root": str(CODE_DIR),
-            "optimized_code": "def sorter(arr):\n    return sorted(arr)",
+            "optimized_code": "def sorter(arr):\n    result = sorted(arr)\n    print('codeflash stdout: Sorting list')\n    print(f'result: {result}')\n    return result",
         }
     if language == "javascript":
         return {
@@ -203,6 +203,7 @@ def main() -> int:
     print("=" * 60)
     print("\n")
     langs = ["python", "javascript"]
+    langs = ["python"]
     for lang in langs:
         exit_code = doTest(get_e2e_test_config_by_language(lang))
         if exit_code != 0:

mcp_server/tools/behavioral.py

@@ -37,19 +37,16 @@ def run_behavioral_tests(
     )
 
     conn = get_connection()
-    try:
-        store_run(
-            conn=conn,
-            run_id=run_id,
-            run_type="behavioral",
-            project_root=project_root,
-            test_files=test_files,
-            test_results=test_results,
-            raw_stdout=run_result.stdout or "",
-            raw_stderr=run_result.stderr or "",
-        )
-    finally:
-        conn.close()
+    store_run(
+        conn=conn,
+        run_id=run_id,
+        run_type="behavioral",
+        project_root=project_root,
+        test_files=test_files,
+        test_results=test_results,
+        raw_stdout=run_result.stdout or "",
+        raw_stderr=run_result.stderr or "",
+    )
 
     invocation_results = []
     errors = []

mcp_server/tools/benchmarking.py

@@ -46,23 +46,20 @@ def run_benchmarking_tests(
     )
 
     conn = get_connection()
-    try:
-        store_run(
-            conn=conn,
-            run_id=run_id,
-            run_type="benchmarking",
-            project_root=project_root,
-            test_files=test_files,
-            test_results=test_results,
-            raw_stdout=run_result.stdout or "",
-            raw_stderr=run_result.stderr or "",
-        )
+    store_run(
+        conn=conn,
+        run_id=run_id,
+        run_type="benchmarking",
+        project_root=project_root,
+        test_files=test_files,
+        test_results=test_results,
+        raw_stdout=run_result.stdout or "",
+        raw_stderr=run_result.stderr or "",
+    )
 
-        speedup_info = None
-        if baseline_run_id:
-            speedup_info = _compute_speedup(conn, baseline_run_id, test_results)
-    finally:
-        conn.close()
+    speedup_info = None
+    if baseline_run_id:
+        speedup_info = _compute_speedup(conn, baseline_run_id, test_results)
 
     best_summed_runtime_ns = test_results.total_passed_runtime() if test_results else 0
     loops_executed = test_results.effective_loop_count() if test_results else 0

mcp_server/tools/compare.py

@@ -7,11 +7,8 @@
 
 def compare_test_results(original_run_id: str, candidate_run_id: str, pass_fail_only: bool = False) -> dict[str, Any]:
     conn = get_connection()
-    try:
-        original_results = load_test_results(conn, original_run_id)
-        candidate_results = load_test_results(conn, candidate_run_id)
-    finally:
-        conn.close()
+    original_results = load_test_results(conn, original_run_id)
+    candidate_results = load_test_results(conn, candidate_run_id)
 
     if not original_results:
         return {"error": f"No results found for run_id: {original_run_id}"}