[INIT] install VSCode extension

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, install_vscode_extension
 from codeflash.cli_cmds.console import logger
 from codeflash.code_utils import env_utils
 from codeflash.code_utils.code_utils import exit_with_message
@@ -21,6 +21,9 @@ def parse_args() -> Namespace:
     init_parser = subparsers.add_parser("init", help="Initialize Codeflash for a Python project.")
     init_parser.set_defaults(func=init_codeflash)
 
+    install_vscode_parser = subparsers.add_parser("vscode-install", help="Install the Codeflash VSCode extension")
+    install_vscode_parser.set_defaults(func=install_vscode_extension)
+
     init_actions_parser = subparsers.add_parser("init-actions", help="Initialize GitHub Actions workflow")
     init_actions_parser.set_defaults(func=install_github_actions)
 

codeflash/cli_cmds/cmd_init.py

@@ -3,6 +3,7 @@
 import ast
 import os
 import re
+import shutil
 import subprocess
 import sys
 from enum import Enum, auto
@@ -23,7 +24,7 @@
 
 from codeflash.api.cfapi import is_github_app_installed_on_repo
 from codeflash.cli_cmds.cli_common import apologize_and_exit
-from codeflash.cli_cmds.console import console, logger
+from codeflash.cli_cmds.console import console, logger, progress_bar
 from codeflash.code_utils.compat import LF
 from codeflash.code_utils.config_parser import parse_config_file
 from codeflash.code_utils.env_utils import check_formatter_installed, get_codeflash_api_key
@@ -98,6 +99,8 @@ def init_codeflash() -> None:
 
         install_github_actions(override_formatter_check=True)
 
+        install_vscode_extension()
+
         module_string = ""
         if "setup_info" in locals():
             module_string = f" you selected ({setup_info.module_root})"
@@ -802,6 +805,41 @@ def install_github_actions(override_formatter_check: bool = False) -> None:  # n
         apologize_and_exit()
 
 
+def install_vscode_extension() -> None:
+    vscode_path = shutil.which("code")
+    if not vscode_path:
+        return
+
+    error = ""
+    with progress_bar("Installing Codeflash for VSCode…"):
+        try:
+            result = subprocess.run(
+                [vscode_path, "--install-extension", "codeflash.codeflash", "--force"],
+                check=True,
+                text=True,
+                timeout=60,
+                capture_output=True,
+            )
+        except subprocess.TimeoutExpired:
+            error = "Installation timed out."
+        except subprocess.CalledProcessError as e:
+            error = e.stderr or "Unknown error."
+
+    if error:
+        ph("vscode-extension-install-failed", {"error": error.strip()})
+        click.echo(
+            "Failed to install Codeflash for VSCode. Please try installing it manually from the Marketplace: https://marketplace.visualstudio.com/items?itemName=codeflash.codeflash"
+        )
+        click.echo(error.strip())
+    else:
+        output = (result.stdout or "").lower()
+        if "already installed" in output:
+            click.echo("✅ Codeflash for VSCode is already installed.")
+            return
+        ph("vscode-extension-installed")
+        click.echo("✅ Installed the latest version of Codeflash for VSCode.")
+
+
 def determine_dependency_manager(pyproject_data: dict[str, Any]) -> DependencyManager:  # noqa: PLR0911
     """Determine which dependency manager is being used based on pyproject.toml contents."""
     if (Path.cwd() / "poetry.lock").exists():