Implement a debug mode with coredumpy and viztracer

.github/workflows/test_and_cov.yml

@@ -32,13 +32,13 @@ jobs:
       - name: install pdm and dependencies
         if: matrix.os != 'macos-14'
         run: |
-          pdm lock --group dev --group lsp --group mcp
+          pdm lock --group dev --group lsp --group mcp --group debug
           pdm install
 
       - name: install pdm and dependencies for legacy system
         if: matrix.os == 'macos-14'
         run: |
-          pdm lock --group dev --group lsp --group mcp --group legacy
+          pdm lock --group dev --group lsp --group mcp --group legacy --group debug
           pdm install
 
       - name: Set custom HF cache directory

Makefile

@@ -1,7 +1,7 @@
 .PHONY: multitest
 
 deps:
-	pdm lock --group dev --group lsp --group mcp; \
+	pdm lock --group dev --group lsp --group mcp --group debug; \
 	pdm install
 
 test:

pyproject.toml

@@ -44,6 +44,7 @@ omit = [
     "./tests/*",
     "src/vectorcode/_version.py",
     "src/vectorcode/__init__.py",
+    "src/vectorcode/debugging.py",
     "/tmp/*",
 ]
 include = ['src/vectorcode/**/*.py']
@@ -61,14 +62,12 @@ write_template = "__version__ = '{}' # pragma: no cover"
 dev = [
     "ipython>=8.31.0",
     "ruff>=0.9.1",
-    "viztracer>=1.0.0",
     "pre-commit>=4.0.1",
     "pytest>=8.3.4",
     "pdm-backend>=2.4.3",
     "coverage>=7.6.12",
     "pytest-asyncio>=0.25.3",
     "debugpy>=1.8.12",
-    "coredumpy>=0.4.1",
     "basedpyright>=1.29.2",
 ]
 
@@ -77,6 +76,7 @@ legacy = ["numpy<2.0.0", "torch==2.2.2", "transformers<=4.49.0"]
 intel = ['optimum[openvino]', 'openvino']
 lsp = ['pygls<2.0.0', 'lsprotocol']
 mcp = ['mcp<2.0.0', 'pydantic']
+debug = ["coredumpy>=0.4.1", "viztracer>=1.0.0"]
 
 [tool.basedpyright]
 typeCheckingMode = "standard"

src/vectorcode/cli_utils.py

@@ -64,6 +64,7 @@ class CliAction(Enum):
 
 @dataclass
 class Config:
+    debug: bool = False
     no_stderr: bool = False
     recursive: bool = False
     include_hidden: bool = False
@@ -189,6 +190,12 @@ async def merge_from(self, other: "Config") -> "Config":
 def get_cli_parser():
     __default_config = Config()
     shared_parser = argparse.ArgumentParser(add_help=False)
+    shared_parser.add_argument(
+        "--debug",
+        default=False,
+        action="store_true",
+        help="Enable debug mode (requires vectorcode[debug]).",
+    )
     chunking_parser = argparse.ArgumentParser(add_help=False)
     chunking_parser.add_argument(
         "--overlap",
@@ -397,6 +404,7 @@ async def parse_cli_args(args: Optional[Sequence[str]] = None):
         "action": CliAction(main_args.action),
         "project_root": main_args.project_root,
         "pipe": main_args.pipe,
+        "debug": main_args.debug,
     }
 
     match main_args.action:

src/vectorcode/debugging.py

@@ -0,0 +1,38 @@
+import logging
+import os
+from datetime import datetime
+
+# import atexit
+
+__LOG_DIR = os.path.expanduser("~/.local/share/vectorcode/logs/")
+
+logger = logging.getLogger(name=__name__)
+
+__tracer = None
+
+
+def finish():
+    if __tracer is not None:
+        __tracer.stop()
+        __tracer.save(
+            output_file=os.path.join(
+                __LOG_DIR,
+                f"viztracer-{datetime.now().strftime('%Y-%m-%d_%H-%M-%S')}.json",
+            )
+        )
+
+
+def enable():
+    global __tracer
+    try:
+        import coredumpy
+        # import viztracer
+    except ModuleNotFoundError:
+        logger.warning("Failed to import modules. Please install vectorcode[debug]")
+        return
+
+    coredumpy.patch_except(directory=__LOG_DIR)
+    # if __tracer is None:
+    #     __tracer = viztracer.VizTracer(log_async=True)
+    # __tracer.start()
+    # atexit.register(finish)

src/vectorcode/main.py

@@ -4,7 +4,7 @@
 import sys
 import traceback
 
-from vectorcode import __version__
+from vectorcode import __version__, debugging
 from vectorcode.cli_utils import (
     CliAction,
     config_logging,
@@ -20,6 +20,10 @@ async def async_main():
     cli_args = await parse_cli_args()
     if cli_args.no_stderr:
         sys.stderr = open(os.devnull, "w")
+
+    if cli_args.debug:
+        debugging.enable()
+
     logger.info("Collected CLI arguments: %s", cli_args)
 
     if cli_args.project_root is None: