prism-code-review 2.0.0__py3-none-any.whl

prism/ai_test.py

@@ -0,0 +1,65 @@
+"""
+AI Test Agent - Uses Gemini for test coverage review.
+
+Catches: missing test coverage, weak test cases, untested code paths,
+and test quality issues.
+
+Rule IDs follow SonarQube convention: TST-XXX
+"""
+
+import logging
+from typing import List
+from prism.agents import BaseAgent, Finding, register_agent
+from prism.ai_common import review_with_agent
+
+logger = logging.getLogger(__name__)
+
+SYSTEM_PROMPT = """
+You are a software testing expert with 10+ years of experience in test automation.
+You review code for test coverage gaps and test quality issues.
+
+Focus on:
+1. Untested code paths (functions without corresponding tests)
+2. Weak test cases (tests that don't actually verify behavior)
+3. Missing edge cases (null inputs, empty inputs, boundary values)
+4. Missing error handling tests (what happens when things go wrong?)
+5. Test isolation issues (tests that depend on shared state)
+6. Slow tests (tests that perform I/O, network calls, or database access)
+7. Missing assertions (tests that don't verify expected outcomes)
+8. Untested public APIs (exported functions with no test coverage)
+
+Rules:
+1. Return ONLY valid JSON. No markdown, no extra text.
+2. Be specific. Each finding must include: severity, category, message, file, line, suggestion, suggested_code.
+3. For suggested_code, provide the actual test code that fixes the issue.
+4. Focus on actual coverage gaps, not test style preferences.
+5. If you find no issues, return empty findings array.
+
+Schema:
+{
+  "findings": [
+    {
+      "severity": "critical" | "high" | "medium" | "low",
+      "category": "string",
+      "message": "string",
+      "file": "string",
+      "line": number,
+      "suggestion": "string",
+      "suggested_code": "string or null"
+    }
+  ],
+  "summary": "2-3 sentence overall test coverage assessment"
+}
+"""
+
+
+@register_agent
+class AITestAgent(BaseAgent):
+    name = "ai_test"
+    description = "AI-powered test coverage reviewer using Gemini"
+
+    def _check(self, file_path: str, content: str, diff: str,
+               skeleton: str = "", line_count: int = None,
+               fast: bool = False) -> List[Finding]:
+        return review_with_agent(SYSTEM_PROMPT, file_path, content,
+                                 self.name, skeleton, line_count, diff, fast)

prism/autofix.py

@@ -0,0 +1,360 @@
+"""
+Auto-fix engine - Generates corrected code for auto-fixable findings.
+
+Based on GitHub Copilot's "suggested changes" and CodeRabbit's "One-Click Fix":
+- AI generates corrected code snippets
+- Developer reviews before applying
+- Changes are attributed to the developer, not committed automatically
+"""
+
+import re
+from typing import List, Optional, Tuple
+from prism.agents import Finding
+from prism.langs import lang_for
+from prism.logging_config import logger
+
+# A declaration is a name preceded by modifiers, a type, or a qualifier
+# (`const token`, `String apiKey`, `self.password`). Taking the name from the
+# start of the line got the boundary wrong: `String apiKey` produced
+# `StringapiKey` and `\tapiKey :=` produced the single character `y`.
+_SECRET_ASSIGN = re.compile(
+    r'^(?P<head>\s*(?:[\w$]+(?:[\s*&.]*|::))*?)'
+    r'(?P<name>\$?[A-Za-z_]\w*)'
+    r'(?P<op>\s*(?::=|(?<![=!<>:])=(?!=)|:(?!=))\s*)'
+    r'(?P<q>["\'])(?P<lit>.*?)(?P=q)(?P<tail>.*)$')
+
+
+def _is_valid_python(source: str) -> bool:
+    """True when source parses as a complete Python module."""
+    try:
+        compile(source, "<prism-autofix>", "exec")
+    except (SyntaxError, ValueError):
+        return False
+    return True
+
+
+def generate_fix(finding: Finding, original_content: str) -> Optional[str]:
+    """
+    Generate a corrected version of the line(s) for a finding.
+
+    Returns the fixed code string, or None if no fix is possible.
+    """
+    lines = original_content.split("\n")
+    line_idx = (finding.line or 1) - 1
+
+    if line_idx < 0 or line_idx >= len(lines):
+        return None
+
+    original_line = lines[line_idx]
+    fix_type = _classify_fix(finding)
+    spec = _spec_for(finding)
+
+    if fix_type == "remove_line":
+        return _remove_line(lines, line_idx)
+
+    if fix_type == "replace_secret":
+        return _replace_secret(original_line, finding, spec)
+
+    if fix_type == "remove_eval":
+        return _remove_eval(original_line, finding, spec)
+
+    if fix_type == "remove_exec":
+        return _remove_exec(original_line, finding, spec)
+
+    if fix_type == "add_logging":
+        return _replace_print_with_logging(original_line, finding, spec)
+
+    if fix_type == "remove_debugger":
+        return _remove_debugger(original_line, finding, spec)
+
+    if fix_type == "remove_console_log":
+        return _comment_out(original_line, finding, spec)
+
+    if fix_type == "add_test_docstring":
+        return _add_test_docstring(lines, line_idx, finding)
+
+    return None
+
+
+def _spec_for(finding):
+    """The language of the file a finding came from, or None if it is not source."""
+    return lang_for(getattr(finding, "file", "") or "")
+
+
+def _comment_prefix(spec):
+    """How a language starts a comment, for fixes that must disable a line."""
+    if spec is None or not spec.line_comments:
+        return "#"
+    return spec.line_comments[0]
+
+
+def _is_python(spec):
+    return spec is None or spec.name == "python"
+
+
+def _classify_fix(finding: Finding) -> str:
+    """Classify what kind of fix is needed.
+
+    Rule IDs are authoritative when present. The message fallbacks use word
+    boundaries so "evaluation" no longer looks like eval() and "keyword" no
+    longer looks like an API key.
+    """
+    rule_id = finding.rule_id or ""
+    msg = finding.message.lower()
+
+    if rule_id == "SECU-201":
+        return "remove_eval"
+    if rule_id == "SECU-202":
+        return "remove_exec"
+    if re.match(r"^SECU-1\d\d$", rule_id):
+        return "replace_secret"
+    if rule_id == "POL-203":
+        return "remove_debugger"
+    if rule_id == "POL-202":
+        return "add_logging"
+    # QUAL-110 is "missing docstring": a real docstring needs the function's
+    # purpose, which cannot be invented from the name. `add_test_docstring`
+    # only ever applied to TST-102, and routing QUAL-110 here produced a
+    # suggestion that just echoed the def line back.
+    if rule_id == "TST-102":
+        return "add_test_docstring"
+
+    if re.search(r"\beval\s*\(", msg):
+        return "remove_eval"
+    if re.search(r"\bexec\s*\(", msg):
+        return "remove_exec"
+    if re.search(r"\bdebugger\b", msg):
+        return "remove_debugger"
+    if re.search(r"console\.log", msg):
+        return "remove_console_log"
+    if re.search(r"\bprint\s*\(", msg):
+        return "add_logging"
+    if re.search(r"\b(api[_-]?secret|secret|password|passwd|pwd|token)\b", msg):
+        return "replace_secret"
+    # "docstring" alone is not enough: only test docstrings can be generated.
+    if re.search(r"\bdocstring\b", msg) and "test" in msg:
+        return "add_test_docstring"
+
+    return "unknown"
+
+
+def _remove_line(lines: list, idx: int) -> str:
+    """Remove a line entirely."""
+    return ""
+
+
+def _replace_secret(line: str, finding: Finding, spec=None) -> str:
+    """Replace a hardcoded secret with a read from the environment.
+
+    The declaration is left as it is and only the literal is swapped, so
+    ``apiKey := "..."`` keeps its ``:=`` and its trailing ``;`` while the value
+    becomes the language's own lookup (``os.Getenv("APIKEY")`` in Go,
+    ``process.env.APIKEY`` in JS). Emitting ``os.getenv`` into a Go file would
+    not compile, which is worse than offering no fix at all.
+    """
+    env_read = spec.env_read if spec is not None else 'os.getenv("{name}")'
+    match = _SECRET_ASSIGN.match(line)
+    if not match:
+        return line
+    # PHP declares `$key`; the sigil is not part of the variable's name.
+    name = re.sub(r"[^A-Za-z0-9]", "_", match.group("name").lstrip("$")).upper()
+    return (match.group("head") + match.group("name") + match.group("op")
+            + env_read.format(name=name) + match.group("tail"))
+
+
+def _remove_dangerous_call(line: str, name: str, spec=None) -> str:
+    """Replace a dangerous call with a syntactically valid, inert statement.
+
+    Commenting out only the call left fragments such as
+    ``return # eval() removed`` followed by ``# x)``, which is confusing and
+    silently changes control flow. Emit a complete statement instead.
+
+    ``pass`` and ``return None`` are Python, so elsewhere the line is disabled
+    with the language's own comment marker: it cannot be wrong about syntax
+    because nothing new is executed.
+    """
+    indent = re.match(r"^(\s*)", line).group(1)
+    stripped = line.strip()
+    if not _is_python(spec):
+        return "{}{} {}() removed - security risk: {}".format(
+            indent, _comment_prefix(spec), name, stripped)
+    if re.match(r"^return\s+" + name + r"\s*\(", stripped):
+        return "{}return None  # {}() removed - security risk".format(indent, name)
+    match = re.match(r"^([A-Za-z_]\w*)\s*=\s*" + name + r"\s*\(", stripped)
+    if match:
+        return "{}{} = None  # {}() removed - security risk".format(
+            indent, match.group(1), name)
+    if re.match(r"^" + name + r"\s*\(", stripped):
+        return "{}pass  # {}() removed - security risk".format(indent, name)
+    return "{}# {}() removed - security risk: {}".format(indent, name, stripped)
+
+
+def _remove_eval(line: str, finding: Finding, spec=None) -> str:
+    """Replace eval() with an inert statement."""
+    return _remove_dangerous_call(line, "eval", spec)
+
+
+def _remove_exec(line: str, finding: Finding, spec=None) -> str:
+    """Replace exec() with an inert statement."""
+    return _remove_dangerous_call(line, "exec", spec)
+
+
+def _replace_print_with_logging(line: str, finding: Finding, spec=None) -> str:
+    """Replace print() with logging.info() (import is ensured later).
+
+    ``logging`` is Python's module; in another language the line is disabled and
+    the developer is pointed at the language's own logger.
+    """
+    match = re.match(r'^(\s*)print\((.*)\)', line)
+    if match and _is_python(spec):
+        indent, content = match.groups()
+        return '{}logging.info({})'.format(indent, content)
+    if match:
+        indent, content = match.groups()
+        return "{}{} print() removed - use a logging framework: {}".format(
+            indent, _comment_prefix(spec), content)
+    return line
+
+
+def _remove_debugger(line: str, finding: Finding, spec=None) -> str:
+    """Comment out a debugger statement.
+
+    Only matches `debugger;` when it is the statement itself, not when the
+    text appears inside a string or regex literal (e.g. a rule definition
+    containing the word "debugger"), which would corrupt the line.
+    """
+    # A quoted literal before the keyword means it is data, not code.
+    stripped = line.lstrip()
+    if stripped.startswith(('"', "'", "r'", 'r"', "b'", 'b"', "f'", 'f"')):
+        return line
+    prefix = _comment_prefix(spec)
+    return re.sub(r'^(\s*)debugger;', r'\g<1>' + prefix + ' debugger removed', line)
+
+
+def _comment_out(line: str, finding=None, spec=None) -> str:
+    """Comment out a line, using the language's own comment marker."""
+    if spec is None:
+        spec = _spec_for(finding)
+    return re.sub(r'^', _comment_prefix(spec) + " ", line, count=1)
+
+
+def _add_test_docstring(lines: list, idx: int, finding: Finding) -> str:
+    """Add a docstring to a test function."""
+    line = lines[idx]
+    # Match def test_name(...): optionally followed by content on same line
+    match = re.match(r'^(\s*def\s+test_\w+\(.*\):)', line)
+    if match:
+        rest = line[match.end():].strip()
+        indent = "    "
+        docstring = '\n{}"""Add a description of what this test validates."""'.format(indent)
+        if rest:
+            return match.group(1) + " " + rest + docstring
+        return match.group(1) + docstring
+    return line
+
+
+def apply_fixes(fixes: List[Tuple[Finding, str]], content: str) -> str:
+    """
+    Apply multiple fixes to content.
+
+    Args:
+        fixes: List of (finding, fixed_line) tuples
+        content: Original file content
+
+    Returns:
+        Modified content
+    """
+    lines = content.split("\n")
+    for finding, fixed_line in fixes:
+        line_idx = (finding.line or 1) - 1
+        if 0 <= line_idx < len(lines):
+            lines[line_idx] = fixed_line
+    # `import os` is Python's way of getting the lookup; adding it to a Go or
+    # JS file would be a syntax error in a language it does not belong to.
+    spec = _spec_for(fixes[0][0]) if fixes else None
+    if _is_python(spec):
+        fixed = _ensure_imports("\n".join(lines))
+    elif spec is not None and spec.name == "go":
+        # `os.Getenv` needs `import "os"`. A single import statement is legal
+        # anywhere after the package clause, so this cannot break the file.
+        fixed = _ensure_go_imports("\n".join(lines))
+    else:
+        fixed = "\n".join(lines)
+
+    # The heuristics are line-local, so a fix can still produce unparsable
+    # code (e.g. inside a multi-line expression). Never hand back a broken
+    # file: fall back to the original and let the developer fix it manually.
+    if _is_python(spec) and _is_valid_python(content) \
+            and not _is_valid_python(fixed):
+        logger.warning(
+            "Auto-fix rejected: generated code for %s would not compile",
+            ", ".join(str((f.line or 0)) for f, _ in fixes),
+        )
+        return content
+    return fixed
+
+
+def _ensure_go_imports(content: str) -> str:
+    """Add ``import "os"`` when a generated fix calls ``os.Getenv``.
+
+    Go has no implicit imports, so a fix that calls ``os.Getenv`` in a file
+    that does not already import ``os`` would not build. Any ``import`` after
+    the package clause is legal, including a second one, so this is a plain
+    insertion rather than an edit of an existing import block.
+    """
+    if not re.search(r"\bos\.Getenv\b", content):
+        return content
+    # Either form: `import "os"`, or the spec on its own line inside `import (...)`.
+    # Importing the same package twice is a compile error in Go, so this guard
+    # has to see through the parenthesised form.
+    spec_line = r'^\s*(?:\w+\s+)?"os"\s*$'
+    if re.search(r'^\s*import\s+(?:\w+\s+)?"os"\s*$', content, re.MULTILINE) or \
+            re.search(spec_line, content, re.MULTILINE):
+        return content
+    lines = content.split("\n")
+    for i, line in enumerate(lines):
+        if re.match(r"^\s*package\s+\w+", line):
+            lines.insert(i + 1, "")
+            lines.insert(i + 2, 'import "os"')
+            return "\n".join(lines)
+    return content
+
+
+def _ensure_imports(content: str) -> str:
+    """Add imports required by generated fixes (e.g. os for os.getenv)."""
+    needs = []
+    if re.search(r"\bos\.(?:getenv|environ)\b", content) and \
+            not re.search(r"^\s*import\s+os\b", content, re.MULTILINE):
+        needs.append("import os")
+    if re.search(r"\blogging\.\w+", content) and \
+            not re.search(r"^\s*import\s+logging\b", content, re.MULTILINE):
+        needs.append("import logging")
+    if not needs:
+        return content
+
+    lines = content.split("\n")
+    insert_at = 0
+    if lines and lines[0].startswith("#!"):
+        insert_at = 1
+    if insert_at < len(lines):
+        stripped = lines[insert_at].lstrip()
+        if stripped.startswith(('"""', "'''")):
+            quote = stripped[:3]
+            if stripped.count(quote) >= 2 and len(stripped) > 3:
+                insert_at += 1
+            else:
+                for j in range(insert_at + 1, len(lines)):
+                    if quote in lines[j]:
+                        insert_at = j + 1
+                        break
+    while insert_at < len(lines) and (
+        lines[insert_at].strip() == "" or lines[insert_at].lstrip().startswith("#")
+    ):
+        insert_at += 1
+    lines[insert_at:insert_at] = needs
+    return "\n".join(lines)
+
+
+def get_auto_fixable(findings: List[Finding]) -> List[Finding]:
+    """Filter findings that can be auto-fixed."""
+    return [f for f in findings if _classify_fix(f) != "unknown"]

prism/cli.py

@@ -0,0 +1,248 @@
+"""PRISM CLI entry point.
+
+This module is a thin facade: it builds the argparse parser and dispatches to
+the command handlers. The implementation is split across:
+  - cli_config.py      : config, git helpers, deviation policy
+  - cli_review.py      : review orchestration, report formatting, JSON output
+  - cli_commands.py    : check / check-staged / check-diff, self-tests
+  - commands_common.py : helpers shared by the command handlers
+  - commands_reports.py: reports, cache, config, stats, dashboard
+  - commands_setup.py  : key, init, hook, version, help, agents
+
+Importing this module runs the @register_agent decorators for the
+deterministic regex agents so they are available alongside the AI agents.
+"""
+
+import argparse
+import sys
+from pathlib import Path
+
+# Add parent directory to path so prism imports work when run as executable
+sys.path.insert(0, str(Path(__file__).parent.parent))
+
+# Load .env file if it exists (for GEMINI_API_KEY and other settings).
+# prism/__init__.py already does this at package import; the call is kept
+# explicit (and idempotent) so this module works even if imported oddly.
+from prism import load_env as _load_env
+_load_env()
+
+# Importing these modules runs their @register_agent decorators so the
+# deterministic regex agents are available alongside the AI agents.
+from prism import agents_policy, agents_quality, agents_security, agents_test
+
+_registered_agents = (agents_policy, agents_quality, agents_security,
+                      agents_test)
+
+from prism.cli_commands import (
+    _run_self_tests, _run_benchmark, _cmd_check, _cmd_check_staged,
+    _cmd_check_diff, _cmd_scan,
+)
+from prism.ai_providers import PROVIDERS as _PROVIDERS
+from prism.commands_reports import (
+    _cmd_reports, _cmd_report, _cmd_delete_report, _cmd_clear_reports,
+    _cmd_dashboard, _cmd_stats, _cmd_config, _cmd_policy, _cmd_feedback,
+    _cmd_suppressed, _cmd_cache,
+)
+from prism.commands_setup import (
+    _cmd_key, _cmd_init, _cmd_install_hook, _cmd_uninstall_hook, _cmd_version,
+    _cmd_help, _cmd_agents, _cmd_ignore, _cmd_mode, _cmd_doctor,
+)
+
+
+def build_parser():
+    """Build the argparse parser with all subcommands."""
+    parser = argparse.ArgumentParser(
+        prog="prism", description="PRISM - Pre-Commit Code Review"
+    )
+    parser.add_argument(
+        "--effort", choices=["lite", "balanced", "max"],
+        default="balanced", help="AI review effort level",
+    )
+    parser.add_argument(
+        "--version", action="store_true",
+        help="Show the PRISM version and exit (same as `prism version`)",
+    )
+    sub = parser.add_subparsers(dest="command")
+
+    check = sub.add_parser("check", help="Review a file")
+    check.add_argument("file", help="File to review")
+    check.add_argument("--fix", action="store_true", help="Apply auto-fixable changes")
+    check.add_argument("--yes", action="store_true",
+                       help="Apply --fix changes without the confirmation prompt")
+    check.add_argument("--json", action="store_true", help="Output results as JSON")
+    check.add_argument("--prompt", action="store_true",
+                       help="Ask developer about medium/low findings")
+    check.add_argument("--effort", choices=["lite", "balanced", "max"],
+                       default="balanced", help="AI review effort level")
+    check.add_argument("--fast", action="store_true",
+                       help="Use two-pass summarization for large files")
+    check.add_argument("--no-dashboard", action="store_true",
+                       help="Do not open the HTML dashboard after review")
+
+    check_staged = sub.add_parser("check-staged", help="Review staged git changes")
+    check_staged.add_argument("--json", action="store_true", help="Output results as JSON")
+    check_staged.add_argument("--prompt", action="store_true",
+                              help="Ask developer about medium/low findings")
+    check_staged.add_argument("--effort", choices=["lite", "balanced", "max"],
+                              default="balanced", help="AI review effort level")
+    check_staged.add_argument("--fast", action="store_true",
+                              help="Use two-pass summarization for large files")
+    check_staged.add_argument("--dashboard", action="store_true",
+                              help="Open visual dashboard after review")
+    check_staged.add_argument("--test", action="store_true",
+                              help="Run pytest on changed Python files before review")
+    check_staged.add_argument("--no-dashboard", action="store_true",
+                              help="Do not open the HTML dashboard after review")
+
+    check_diff = sub.add_parser(
+        "check-diff", aliases=["diff"],
+        help="Review the diff against a base branch (PR review)")
+    check_diff.add_argument("--base", default="origin/main",
+                            help="Base ref to diff against (default: origin/main)")
+    check_diff.add_argument("--json", action="store_true", help="Output results as JSON")
+    check_diff.add_argument("--effort", choices=["lite", "balanced", "max"],
+                            default="balanced", help="AI review effort level")
+    check_diff.add_argument("--fast", action="store_true",
+                            help="Use two-pass summarization for large files")
+    check_diff.add_argument("--no-dashboard", action="store_true",
+                            help="Do not open the HTML dashboard after review")
+
+    sub.add_parser("init", help="Create prism.json config")
+    sub.add_parser("install-hook", help="Install pre-commit hook")
+    sub.add_parser("policy", help="Show deviation policy")
+    sub.add_parser("feedback", help="Show feedback loop stats")
+    sub.add_parser("suppressed", help="Show historically suppressed rules")
+    sub.add_parser("reports", help="Show review history")
+    sub.add_parser("report", help="Show a specific report by ID").add_argument("id")
+    sub.add_parser("dashboard", help="Open visual report dashboard in browser")
+    sub.add_parser("doctor", help="Check environment and provider setup")
+
+    scan_cmd = sub.add_parser(
+        "scan",
+        help="Review every file under a path (whole-project audit)")
+    scan_cmd.add_argument("path", nargs="?", default=".",
+                          help="Root directory to scan (default: .)")
+    scan_cmd.add_argument("--json", action="store_true",
+                          help="Output results as JSON")
+    scan_cmd.add_argument("--effort", choices=["lite", "balanced", "max"],
+                          default="balanced", help="AI review effort level")
+    scan_cmd.add_argument("--fast", action="store_true",
+                          help="Use two-pass summarization for large files")
+    scan_cmd.add_argument("--no-dashboard", action="store_true",
+                          help="Do not open the HTML dashboard after review")
+
+    delete_cmd = sub.add_parser("delete-report", help="Delete a specific report")
+    delete_cmd.add_argument("id", help="Report ID to delete")
+    clear_cmd = sub.add_parser("clear-reports", help="Delete ALL reports")
+    clear_cmd.add_argument("--force", action="store_true", help="Skip confirmation")
+
+    agents_cmd = sub.add_parser("agents", help="List available agents")
+    agents_cmd.add_argument("--enable", nargs="+", metavar="NAME",
+                            help="Enable one or more agents")
+    agents_cmd.add_argument("--disable", nargs="+", metavar="NAME",
+                            help="Disable one or more agents")
+
+    sub.add_parser("help", help="Show all commands and their uses")
+    sub.add_parser("version", help="Show PRISM version")
+    sub.add_parser("uninstall-hook", help="Remove pre-commit hook")
+    sub.add_parser("stats", help="Show review statistics")
+
+    ignore_cmd = sub.add_parser(
+        "ignore",
+        help="Manage files skipped by pre-commit review (.prismignore)")
+    ignore_cmd.add_argument("action",
+                            choices=["list", "add", "remove", "reset"],
+                            nargs="?", default="list",
+                            help="What to do (default: list)")
+    ignore_cmd.add_argument("patterns", nargs="*",
+                            help="Pattern(s) to add or remove")
+
+    mode_cmd = sub.add_parser(
+        "mode",
+        help="Show or switch the agent mode (all / regex / none)")
+    mode_cmd.add_argument("mode", nargs="?",
+                          choices=["all", "regex", "none"],
+                          help="Mode to switch to (omit to show the current one)")
+
+    config_cmd = sub.add_parser("config", help="Show or manage configuration")
+    config_cmd.add_argument("--set", nargs=2, metavar=("KEY", "VALUE"),
+                            help="Set a config value")
+    config_cmd.add_argument("--get", metavar="KEY", help="Get a config value")
+
+    test_cmd = sub.add_parser("test", help="Run self-tests")
+    test_cmd.add_argument("--quick", action="store_true",
+                          help="Run only quick tests (skip API calls)")
+
+    benchmark_cmd = sub.add_parser("benchmark", help="Benchmark agent performance")
+    benchmark_cmd.add_argument("--file", help="File to benchmark against")
+    benchmark_cmd.add_argument("--iterations", type=int, default=3,
+                               help="Number of iterations")
+
+    cache_cmd = sub.add_parser("cache", help="Manage the codebase cache")
+    cache_cmd.add_argument("action", choices=["stats", "clear", "status"],
+                           help="Cache action to perform")
+
+    key_cmd = sub.add_parser("key", help="Manage API keys")
+    key_cmd.add_argument("action", choices=["set", "delete", "status"],
+                         help="Key action to perform")
+    key_cmd.add_argument("--provider", choices=list(_PROVIDERS),
+                         default=None,
+                         help="AI provider (default: gemini for set)")
+    key_cmd.add_argument("--key", help="API key value (required for set)")
+    key_cmd.add_argument("--global", dest="global_", action="store_true",
+                         help="Write the key to the install folder's .env "
+                              "(tool-wide, visible from every project) "
+                              "instead of the current project")
+
+    return parser
+
+
+_COMMANDS = {
+    "agents": _cmd_agents,
+    "help": _cmd_help,
+    "version": _cmd_version,
+    "uninstall-hook": _cmd_uninstall_hook,
+    "stats": _cmd_stats,
+    "config": _cmd_config,
+    "init": _cmd_init,
+    "install-hook": _cmd_install_hook,
+    "policy": _cmd_policy,
+    "feedback": _cmd_feedback,
+    "suppressed": _cmd_suppressed,
+    "reports": _cmd_reports,
+    "dashboard": _cmd_dashboard,
+    "report": _cmd_report,
+    "delete-report": _cmd_delete_report,
+    "clear-reports": _cmd_clear_reports,
+    "cache": _cmd_cache,
+    "check": _cmd_check,
+    "check-staged": _cmd_check_staged,
+    "check-diff": _cmd_check_diff,
+    "diff": _cmd_check_diff,  # documented alias for check-diff
+    "scan": _cmd_scan,
+    "doctor": _cmd_doctor,
+    "test": _run_self_tests,
+    "benchmark": _run_benchmark,
+    "key": _cmd_key,
+    "ignore": _cmd_ignore,
+    "mode": _cmd_mode,
+}
+
+
+def main():
+    parser = build_parser()
+    args = parser.parse_args()
+
+    if getattr(args, "version", False):
+        _cmd_version(args)
+        return
+
+    handler = _COMMANDS.get(args.command)
+    if handler is None:
+        parser.print_help()
+        return
+    handler(args)
+
+
+if __name__ == "__main__":
+    main()