ai-cli-toolkit 0.2.0__py3-none-any.whl

ai_cli/completion_gen.py

@@ -0,0 +1,680 @@
+"""Generate concise shell completions for ai-cli and wrapped tools.
+
+Design goals:
+- Keep completion scripts human-readable and fast (no huge autogenerated blobs)
+- Distinguish wrapped aliases (~/.ai-cli/bin/<tool>) from native binaries
+- Offer wrapper-specific flags only when command is wrapped
+- Capture baseline completion scripts from tools that expose built-in generators
+"""
+
+from __future__ import annotations
+
+import argparse
+import re
+import subprocess
+import sys
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Optional
+
+from ai_cli.config import ensure_config, get_tool_config
+from ai_cli.tools import load_registry
+
+WRAPPER_FLAGS = [
+    "--ai-cli-system-instructions-file",
+    "--ai-cli-system-instructions-text",
+    "--ai-cli-canary-rule",
+    "--ai-cli-passthrough",
+    "--ai-cli-debug-requests",
+    "--ai-cli-developer-instructions-mode",
+    "--ai-cli-no-startup-context",
+]
+
+
+@dataclass
+class ToolCompletionData:
+    name: str
+    native_flags: list[str]
+    native_subcommands: list[str]
+    file_flags: list[str]
+    dir_flags: list[str]
+
+
+def _run(cmd: list[str], timeout: int = 5) -> tuple[int, str]:
+    try:
+        proc = subprocess.run(
+            cmd, check=False, capture_output=True, text=True,
+            timeout=timeout, stdin=subprocess.DEVNULL,
+        )
+    except subprocess.TimeoutExpired:
+        return 1, ""
+    except OSError:
+        return 1, ""
+    output = (proc.stdout or "") + (proc.stderr or "")
+    return proc.returncode, output
+
+
+def _looks_like_script(text: str, shell: str) -> bool:
+    if not text.strip():
+        return False
+    if shell == "bash":
+        return "complete " in text or "compgen" in text
+    return "compdef" in text or "_arguments" in text
+
+
+def _baseline_completion_command(tool: str, binary: str, shell: str) -> Optional[list[str]]:
+    if tool == "codex":
+        return [binary, "completion", shell]
+    if tool in {"copilot", "gemini", "claude"}:
+        return [binary, "completion", shell]
+    return None
+
+
+def _extract_flags(help_text: str) -> list[tuple[Optional[str], str, bool]]:
+    flags: list[tuple[Optional[str], str, bool]] = []
+    seen: set[str] = set()
+    regex = re.compile(r"^\s*(?:(-[A-Za-z0-9]),\s*)?(--[A-Za-z0-9][A-Za-z0-9-]*)(.*)$")
+
+    for raw in help_text.splitlines():
+        line = raw.rstrip("\n")
+        m = regex.match(line)
+        if not m:
+            continue
+        short = m.group(1)
+        long_opt = m.group(2)
+        tail = m.group(3) or ""
+        if long_opt in seen:
+            continue
+        seen.add(long_opt)
+        expects_value = any(tok in tail for tok in ("<", ":", "="))
+        flags.append((short, long_opt, expects_value))
+
+    return flags
+
+
+def _extract_commands(help_text: str) -> list[str]:
+    commands: list[str] = []
+    seen: set[str] = set()
+    in_block = False
+    cmd_line = re.compile(r"^  ([A-Za-z][A-Za-z0-9-]*)\s{2,}.*$")
+
+    for raw in help_text.splitlines():
+        line = raw.rstrip("\n")
+        low = line.strip().lower()
+
+        if not in_block and (
+            low.startswith("commands:") or low.startswith("available commands:")
+        ):
+            in_block = True
+            continue
+
+        if not in_block:
+            continue
+
+        if not line.strip():
+            continue
+
+        if low.startswith(("options:", "flags:", "arguments:", "positionals:", "examples:")):
+            break
+
+        m = cmd_line.match(raw)
+        if not m:
+            continue
+
+        token = m.group(1).strip("[]")
+        if not token or token.startswith("-") or token in {"[query..]", "[command]"}:
+            continue
+        if token in seen:
+            continue
+        seen.add(token)
+        commands.append(token)
+
+    return commands
+
+
+def _classify_flags(flags: list[tuple[Optional[str], str, bool]]) -> tuple[list[str], list[str]]:
+    file_flags: list[str] = []
+    dir_flags: list[str] = []
+    for _, long_opt, expects in flags:
+        if not expects:
+            continue
+        low = long_opt.lower()
+        if any(k in low for k in ("dir", "directory", "worktree")):
+            dir_flags.append(long_opt)
+        elif any(k in low for k in ("file", "path", "config", "schema", "settings", "ca")):
+            file_flags.append(long_opt)
+    return sorted(set(file_flags)), sorted(set(dir_flags))
+
+
+_responsive_tools: set[str] = set()
+"""Tools whose --help responded within timeout (skip slow ones in baseline)."""
+
+
+def _tool_data() -> list[ToolCompletionData]:
+    registry = load_registry()
+    cfg = ensure_config()
+    result: list[ToolCompletionData] = []
+
+    for name, spec in registry.items():
+        tool_cfg = get_tool_config(cfg, name)
+        binary = spec.resolve_binary(tool_cfg.get("binary", ""))
+        code, help_text = _run([binary, "--help"])
+        if code != 0 and not help_text.strip():
+            help_text = ""
+        else:
+            _responsive_tools.add(name)
+
+        parsed_flags = _extract_flags(help_text)
+        commands = _extract_commands(help_text)
+        file_flags, dir_flags = _classify_flags(parsed_flags)
+
+        native_flags: list[str] = []
+        for short, long_opt, _ in parsed_flags:
+            if short:
+                native_flags.append(short)
+            native_flags.append(long_opt)
+
+        result.append(
+            ToolCompletionData(
+                name=name,
+                native_flags=sorted(set(native_flags)),
+                native_subcommands=commands,
+                file_flags=file_flags,
+                dir_flags=dir_flags,
+            )
+        )
+
+    return result
+
+
+def _capture_baselines(root: Path, tools: list[ToolCompletionData], shell: str) -> None:
+    registry = load_registry()
+    cfg = ensure_config()
+    dest = root / "completions" / "generated" / "baseline"
+    dest.mkdir(parents=True, exist_ok=True)
+
+    for tool in tools:
+        if tool.name not in _responsive_tools:
+            continue
+        spec = registry[tool.name]
+        tool_cfg = get_tool_config(cfg, tool.name)
+        binary = spec.resolve_binary(tool_cfg.get("binary", ""))
+        cmd = _baseline_completion_command(tool.name, binary, shell)
+        if not cmd:
+            continue
+        rc, out = _run(cmd)
+        if rc == 0 and _looks_like_script(out, shell):
+            ext = "bash" if shell == "bash" else "zsh"
+            (dest / f"{tool.name}.{ext}").write_text(out, encoding="utf-8")
+
+
+def _join(values: list[str]) -> str:
+    return " ".join(values)
+
+
+def _render_bash(tools: list[ToolCompletionData]) -> str:
+    case_blocks = []
+    for t in tools:
+        case_blocks.append(
+            f"""    {t.name})
+      native_flags=\"{_join(t.native_flags)}\"
+      native_subcommands=\"{_join(t.native_subcommands)}\"
+      file_flags=\"{_join(t.file_flags)}\"
+      dir_flags=\"{_join(t.dir_flags)}\"
+      ;;
+"""
+        )
+
+    ai_subcommands = "claude codex copilot gemini menu status system prompt-edit session traffic update completions help"
+    tool_names = "claude codex copilot gemini"
+    wrapper_flags = _join(WRAPPER_FLAGS)
+    traffic_callers = "claude copilot codex gemini"
+    traffic_sort_opts = "date address"
+
+    return f'''#!/usr/bin/env bash
+# Bash completion for ai-cli and wrapped/native tool aliases.
+# Generated by ai_cli.completion_gen
+
+_ai_cli_is_wrapped_cmd() {{
+  local cmd="$1"
+  local resolved
+  resolved="$(command -v "$cmd" 2>/dev/null || true)"
+  [[ "$resolved" == "$HOME/.ai-cli/bin/$cmd" ]]
+}}
+
+_ai_cli_fill_tool_data() {{
+  local tool="$1"
+  native_flags=""
+  native_subcommands=""
+  file_flags=""
+  dir_flags=""
+  case "$tool" in
+{''.join(case_blocks)}    *) ;;
+  esac
+}}
+
+_ai_cli_tool_completion_impl() {{
+  local tool="$1"
+  local cur prev words cword
+  _init_completion 2>/dev/null || {{
+    COMPREPLY=()
+    cur="${{COMP_WORDS[COMP_CWORD]}}"
+    prev="${{COMP_WORDS[COMP_CWORD-1]}}"
+    words=("${{COMP_WORDS[@]}}")
+    cword=$COMP_CWORD
+  }}
+
+  _ai_cli_fill_tool_data "$tool"
+
+  local flags="$native_flags"
+  if _ai_cli_is_wrapped_cmd "$tool"; then
+    flags="$flags {wrapper_flags}"
+  fi
+
+  case " $file_flags " in
+    *" $prev "*) COMPREPLY=($(compgen -f -- "$cur")); return ;;
+  esac
+  case " $dir_flags " in
+    *" $prev "*) COMPREPLY=($(compgen -d -- "$cur")); return ;;
+  esac
+
+  if [[ "$cur" == -* ]]; then
+    COMPREPLY=($(compgen -W "$flags" -- "$cur"))
+    return
+  fi
+
+  local i subcommand=""
+  for (( i=1; i < cword; i++ )); do
+    if [[ "${{words[i]}}" != -* ]]; then
+      subcommand="${{words[i]}}"
+      break
+    fi
+  done
+
+  if [[ -z "$subcommand" ]]; then
+    COMPREPLY=($(compgen -W "$native_subcommands" -- "$cur"))
+    COMPREPLY+=($(compgen -d -- "$cur"))
+  fi
+}}
+
+_ai_cli_completion() {{
+  local cur prev words cword
+  _init_completion 2>/dev/null || {{
+    COMPREPLY=()
+    cur="${{COMP_WORDS[COMP_CWORD]}}"
+    prev="${{COMP_WORDS[COMP_CWORD-1]}}"
+    words=("${{COMP_WORDS[@]}}")
+    cword=$COMP_CWORD
+  }}
+
+  # Top-level subcommand completions with descriptions
+  if [[ $cword -eq 1 ]]; then
+    COMPREPLY=($(compgen -W "{ai_subcommands}" -- "$cur"))
+    return
+  fi
+
+  local cmd="${{words[1]}}"
+    case "$cmd" in
+    claude|codex|copilot|gemini)
+      _ai_cli_tool_completion_impl "$cmd"
+      return ;;
+    system)
+      if [[ $cword -eq 2 ]]; then
+        COMPREPLY=($(compgen -W "prompt {tool_names}" -- "$cur"))
+      else
+        COMPREPLY=($(compgen -W "{tool_names}" -- "$cur"))
+      fi
+      return ;;
+    session)
+      case "$prev" in
+        --agent) COMPREPLY=($(compgen -W "all {tool_names}" -- "$cur")); return ;;
+        --tail) return ;;
+      esac
+      if [[ "$cur" == -* ]]; then
+        COMPREPLY=($(compgen -W "--agent --all --list --grep --tail --tools --raw" -- "$cur"))
+      else
+        COMPREPLY=($(compgen -f -- "$cur"))
+      fi
+      return ;;
+    traffic)
+      case "$prev" in
+        --caller|-c) COMPREPLY=($(compgen -W "{traffic_callers}" -- "$cur")); return ;;
+        --sort) COMPREPLY=($(compgen -W "{traffic_sort_opts}" -- "$cur")); return ;;
+        --db) COMPREPLY=($(compgen -f -- "$cur")); return ;;
+        --host|--search|-s|--limit|-n|--detail|-d) return ;;
+      esac
+      if [[ "$cur" == -* ]]; then
+        COMPREPLY=($(compgen -W "--caller --host --search --api --sort --limit --db --plain --detail" -- "$cur"))
+      fi
+      return ;;
+    update)
+      if [[ "$cur" == -* ]]; then
+        COMPREPLY=($(compgen -W "--all --dry-run --list --method --list-methods" -- "$cur"))
+      else
+        COMPREPLY=($(compgen -W "{tool_names}" -- "$cur"))
+      fi
+      return ;;
+    completions)
+      if [[ $cword -eq 2 ]]; then
+        COMPREPLY=($(compgen -W "generate" -- "$cur"))
+      elif [[ "$cur" == -* ]]; then
+        COMPREPLY=($(compgen -W "--shell" -- "$cur"))
+      fi
+      return ;;
+  esac
+}}
+
+_ai_cli_claude_completion() {{ _ai_cli_tool_completion_impl claude; }}
+_ai_cli_codex_completion() {{ _ai_cli_tool_completion_impl codex; }}
+_ai_cli_copilot_completion() {{ _ai_cli_tool_completion_impl copilot; }}
+_ai_cli_gemini_completion() {{ _ai_cli_tool_completion_impl gemini; }}
+
+complete -F _ai_cli_completion ai-cli
+complete -F _ai_cli_claude_completion claude
+complete -F _ai_cli_codex_completion codex
+complete -F _ai_cli_copilot_completion copilot
+complete -F _ai_cli_gemini_completion gemini
+'''
+
+
+def _render_zsh(tools: list[ToolCompletionData]) -> str:
+    case_blocks = []
+    for t in tools:
+        case_blocks.append(
+            f"""    {t.name})
+      native_flags=({_join(t.native_flags)})
+      native_subcommands=({_join(t.native_subcommands)})
+      file_flags=({_join(t.file_flags)})
+      dir_flags=({_join(t.dir_flags)})
+      ;;
+"""
+        )
+
+    ai_subcommands = "claude codex copilot gemini menu status system prompt-edit session traffic update completions help"
+    tool_names = "claude codex copilot gemini"
+    wrapper_flags = _join(WRAPPER_FLAGS)
+    traffic_callers = "claude copilot codex gemini"
+    traffic_sort_opts = "date address"
+
+    return f'''#compdef ai-cli claude codex copilot gemini
+# Zsh completion for ai-cli and wrapped/native tool aliases.
+# Generated by ai_cli.completion_gen
+
+autoload -U is-at-least
+
+_ai_cli_is_wrapped_cmd() {{
+  local cmd="$1"
+  local resolved
+  resolved="$(whence -p "$cmd" 2>/dev/null)"
+  [[ "$resolved" == "$HOME/.ai-cli/bin/$cmd" ]]
+}}
+
+_ai_cli_fill_tool_data() {{
+  local tool="$1"
+  native_flags=()
+  native_subcommands=()
+  file_flags=()
+  dir_flags=()
+  case "$tool" in
+{''.join(case_blocks)}    *) ;;
+  esac
+}}
+
+_ai_cli_tool_completion_impl() {{
+  local tool="$1"
+  _ai_cli_fill_tool_data "$tool"
+
+  local -a flags
+  flags=(${{native_flags[@]}})
+  if _ai_cli_is_wrapped_cmd "$tool"; then
+    flags+=({wrapper_flags})
+  fi
+
+  local prev="${{words[CURRENT-1]}}"
+  if (( ${{file_flags[(Ie)$prev]}} )); then
+    _files
+    return
+  fi
+  if (( ${{dir_flags[(Ie)$prev]}} )); then
+    _directories
+    return
+  fi
+
+  if [[ "${{words[CURRENT]}}" == -* ]]; then
+    _describe 'option' flags
+    return
+  fi
+
+  local i subcommand=""
+  for ((i=2; i<CURRENT; i++)); do
+    if [[ "${{words[i]}}" != -* ]]; then
+      subcommand="${{words[i]}}"
+      break
+    fi
+  done
+
+  if [[ -z "$subcommand" ]]; then
+    _describe 'command' native_subcommands
+    _directories
+  fi
+}}
+
+_ai_cli_main_completion() {{
+  local cur="${{words[CURRENT]}}"
+  local cmd="${{words[2]}}"
+
+  if (( CURRENT == 2 )); then
+    local -a sc
+    sc=(
+      "claude:Launch Claude Code agent"
+      "codex:Launch OpenAI Codex agent"
+      "copilot:Launch GitHub Copilot CLI"
+      "gemini:Launch Google Gemini CLI"
+      "menu:Interactive tool manager (TUI)"
+      "status:Show installed tools and versions"
+      "system:Edit or view system instructions"
+      "prompt-edit:Edit global or tool prompt file"
+      "session:Browse agent conversation history"
+      "traffic:Browse and search proxied API traffic"
+      "update:Install or update wrapped tools"
+      "completions:Generate shell completion scripts"
+      "help:Show usage information"
+    )
+    _describe 'command' sc
+    return
+  fi
+
+  case "$cmd" in
+    claude|codex|copilot|gemini)
+      _ai_cli_tool_completion_impl "$cmd"
+      return ;;
+    system)
+      if (( CURRENT == 3 )); then
+        local -a sa
+        sa=(
+          "prompt:Show captured system prompt for a model"
+          "claude:Edit Claude instructions"
+          "codex:Edit Codex instructions"
+          "copilot:Edit Copilot instructions"
+          "gemini:Edit Gemini instructions"
+        )
+        _describe 'subcommand' sa
+      else
+        local -a tn
+        tn=({tool_names})
+        _describe 'tool' tn
+      fi
+      return ;;
+    prompt-edit)
+      if (( CURRENT == 3 )); then
+        local -a sa
+        sa=(
+          "global:Edit global prompt file"
+          "tool:Edit tool-specific prompt file"
+        )
+        _describe 'scope' sa
+      elif (( CURRENT == 4 )); then
+        local -a tn
+        tn=({tool_names})
+        _describe 'tool' tn
+      fi
+      return ;;
+    session)
+      case "${{words[CURRENT-1]}}" in
+        --agent)
+          local -a a
+          a=(
+            "all:All tools"
+            "claude:Claude sessions"
+            "codex:Codex sessions"
+            "copilot:Copilot sessions"
+            "gemini:Gemini sessions"
+          )
+          _describe 'agent' a
+          return ;;
+        --tail)
+          _message 'number'
+          return ;;
+      esac
+      if [[ "$cur" == -* ]]; then
+        local -a sf
+        sf=(
+          "--agent:Filter by tool name"
+          "--all:Show all sessions"
+          "--list:List sessions"
+          "--grep:Search session content"
+          "--tail:Show last N entries"
+          "--tools:Show tool metadata"
+          "--raw:Raw output format"
+        )
+        _describe 'flag' sf
+      else
+        _files
+      fi
+      return ;;
+    traffic)
+      case "${{words[CURRENT-1]}}" in
+        --caller|-c)
+          local -a tc
+          tc=(
+            "claude:Show Claude traffic"
+            "copilot:Show Copilot traffic"
+            "codex:Show Codex traffic"
+            "gemini:Show Gemini traffic"
+          )
+          _describe 'caller' tc
+          return ;;
+        --sort)
+          local -a so
+          so=(
+            "date:Sort by timestamp (newest first)"
+            "address:Sort by host and path"
+          )
+          _describe 'sort order' so
+          return ;;
+        --db)
+          _files
+          return ;;
+        --host|--search|-s|--limit|-n|--detail|-d)
+          return ;;
+      esac
+      if [[ "$cur" == -* ]]; then
+        local -a tf
+        tf=(
+          "--caller:Filter by caller tool"
+          "--host:Filter by host substring"
+          "--search:Search request/response bodies"
+          "--api:Show only confirmed API calls"
+          "--sort:Sort order (date or address)"
+          "--limit:Maximum rows to show"
+          "--db:Path to traffic database"
+          "--plain:Plain text output (no curses)"
+          "--detail:Show detail for a specific row ID"
+        )
+        _describe 'flag' tf
+      fi
+      return ;;
+    update)
+      if [[ "$cur" == -* ]]; then
+        local -a uf
+        uf=(
+          "--all:Update all tools"
+          "--dry-run:Show what would be done"
+          "--list:List available tools"
+          "--method:Installation method"
+          "--list-methods:List install methods"
+        )
+        _describe 'flag' uf
+      else
+        local -a tn
+        tn=({tool_names})
+        _describe 'tool' tn
+      fi
+      return ;;
+    completions)
+      if (( CURRENT == 3 )); then
+        local -a acts
+        acts=("generate:Generate completion scripts")
+        _describe 'action' acts
+      elif [[ "$cur" == -* ]]; then
+        local -a cf
+        cf=("--shell:Shell type (bash, zsh, all)")
+        _describe 'flag' cf
+      fi
+      return ;;
+  esac
+}}
+
+_ai_cli_command_completion() {{ _ai_cli_main_completion "$@"; }}
+_ai_cli_claude_completion() {{ _ai_cli_tool_completion_impl claude; }}
+_ai_cli_codex_completion() {{ _ai_cli_tool_completion_impl codex; }}
+_ai_cli_copilot_completion() {{ _ai_cli_tool_completion_impl copilot; }}
+_ai_cli_gemini_completion() {{ _ai_cli_tool_completion_impl gemini; }}
+
+compdef _ai_cli_command_completion ai-cli
+compdef _ai_cli_claude_completion claude
+compdef _ai_cli_codex_completion codex
+compdef _ai_cli_copilot_completion copilot
+compdef _ai_cli_gemini_completion gemini
+'''
+
+
+def generate(shell: str = "all") -> int:
+    root = Path(__file__).resolve().parent.parent
+    comp_dir = root / "completions"
+    comp_dir.mkdir(parents=True, exist_ok=True)
+
+    tools = _tool_data()
+
+    if shell in ("all", "bash"):
+        _capture_baselines(root, tools, "bash")
+        (comp_dir / "ai-cli.bash").write_text(_render_bash(tools), encoding="utf-8")
+
+    if shell in ("all", "zsh"):
+        _capture_baselines(root, tools, "zsh")
+        (comp_dir / "_ai-cli").write_text(_render_zsh(tools), encoding="utf-8")
+
+    print("Generated completions:")
+    print(f"- {comp_dir / '_ai-cli'}")
+    print(f"- {comp_dir / 'ai-cli.bash'}")
+    print(f"- {comp_dir / 'generated' / 'baseline'}")
+    return 0
+
+
+def main(argv: Optional[list[str]] = None) -> int:
+    parser = argparse.ArgumentParser(prog="ai-cli completions")
+    sub = parser.add_subparsers(dest="action")
+    gen = sub.add_parser("generate")
+    gen.add_argument("--shell", choices=["bash", "zsh", "all"], default="all")
+    args = parser.parse_args(argv)
+
+    action = args.action or "generate"
+    if action == "generate":
+        return generate(shell=getattr(args, "shell", "all"))
+
+    parser.print_help()
+    return 1
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())