zk-data-agent/src/bundled_skills.py

"""Bundled skill definitions — prompt-type skills invocable via the Skill tool.

Mirrors the npm ``src/skills/bundled/`` module.

Bundled skills differ from slash commands:
- They generate AI prompts sent to the model (prompt-type).
- They carry ``when_to_use`` guidance for model auto-invocation.
- They can restrict ``allowed_tools`` during execution.
- They appear in system-reminder skill listings for model discovery.
"""

from __future__ import annotations

import subprocess
from dataclasses import dataclass, field
from pathlib import Path
from typing import TYPE_CHECKING, Callable

if TYPE_CHECKING:
    from .agent_runtime import LocalCodingAgent


@dataclass(frozen=True)
class BundledSkill:
    """A bundled skill definition."""

    name: str
    description: str
    when_to_use: str = ''
    aliases: tuple[str, ...] = ()
    allowed_tools: tuple[str, ...] = ()
    user_invocable: bool = True
    source: str = 'python'
    path: str | None = None
    get_prompt: Callable[['LocalCodingAgent', str], str] = lambda _a, _args: ''


# ---------------------------------------------------------------------------
# Skill prompt generators
# ---------------------------------------------------------------------------

def _simplify_prompt(agent: 'LocalCodingAgent', args: str) -> str:
    """Generate the simplify review prompt."""
    cwd = str(agent.runtime_config.cwd)
    diff = ''
    try:
        proc = subprocess.run(
            ['git', 'diff', 'HEAD'],
            cwd=cwd, capture_output=True, text=True, timeout=30,
        )
        diff = proc.stdout.strip()
    except Exception:
        pass

    if not diff:
        try:
            proc = subprocess.run(
                ['git', 'diff'],
                cwd=cwd, capture_output=True, text=True, timeout=30,
            )
            diff = proc.stdout.strip()
        except Exception:
            pass

    diff_section = f'```diff\n{diff}\n```' if diff else 'No changes detected — run `git diff` to verify.'

    return f"""Review the changed code for reuse, quality, and efficiency, then fix any issues found.

## Changed Code

{diff_section}

## Review Checklist

### Code Reuse
- Duplicated functions or logic that should be consolidated
- Redundant utilities or helpers
- Patterns that should use existing abstractions

### Code Quality
- Redundant state or unnecessary variables
- Parameter sprawl (too many function arguments)
- Copy-paste code that should be refactored
- Leaky abstractions or wrong abstraction level
- Stringly-typed code that should use enums/types
- Overly nested conditionals

### Efficiency
- N+1 query or call patterns
- Missed concurrency opportunities (parallel I/O)
- Hot-path bloat (unnecessary work in tight loops)
- Memory leaks or unbounded growth

## Instructions

1. Read the full diff above
2. For each category, identify concrete issues with file paths and line numbers
3. Fix each issue directly — do not just report them
4. After fixing, verify the changes compile/run correctly

{f'Additional context: {args}' if args.strip() else ''}"""


def _debug_prompt(agent: 'LocalCodingAgent', args: str) -> str:
    """Generate the debug prompt."""
    import os

    lines = ['## Debug Session Info', '']
    lines.append(f'Working directory: {agent.runtime_config.cwd}')
    lines.append(f'Model: {agent.model_config.model}')
    lines.append(f'Base URL: {agent.model_config.base_url}')

    session = agent.last_session
    if session:
        lines.append(f'Messages in session: {len(session.messages)}')
        tool_count = sum(1 for m in session.messages if m.role == 'tool')
        lines.append(f'Tool results: {tool_count}')

    usage = agent.cumulative_usage
    lines.append(f'Total tokens used: {usage.total_tokens:,}')
    lines.append(f'Cost: ${agent.cumulative_cost_usd:.4f}')

    # Check for debug log
    debug_log = os.environ.get('CLAUDE_CODE_DEBUG_LOG', '')
    if debug_log:
        lines.append(f'Debug log: {debug_log}')
    else:
        lines.append('Debug logging: not enabled (set CLAUDE_CODE_DEBUG_LOG to enable)')

    return '\n'.join(lines)


def _directory_skill_prompt(body: str) -> Callable[['LocalCodingAgent', str], str]:
    def _prompt(_agent: 'LocalCodingAgent', args: str) -> str:
        if args.strip():
            return f'{body.rstrip()}\n\n## Invocation Arguments\n\n{args.strip()}'
        return body

    return _prompt


def _parse_front_matter(text: str) -> tuple[dict[str, str], str]:
    if not text.startswith('---\n'):
        return {}, text
    end_marker = text.find('\n---\n', 4)
    if end_marker == -1:
        return {}, text
    raw_meta = text[4:end_marker]
    body = text[end_marker + len('\n---\n') :]
    metadata: dict[str, str] = {}
    for line in raw_meta.splitlines():
        if not line.strip() or line.lstrip().startswith('#') or ':' not in line:
            continue
        key, value = line.split(':', 1)
        metadata[key.strip()] = value.strip()
    return metadata, body


def _split_csv(value: str | None) -> tuple[str, ...]:
    if not value:
        return ()
    return tuple(item.strip() for item in value.split(',') if item.strip())


def _parse_bool(value: str | None, *, default: bool = True) -> bool:
    if value is None or not value.strip():
        return default
    return value.strip().lower() not in {'0', 'false', 'no', 'off'}


def _load_directory_skill(path: Path, *, source: str = 'directory') -> BundledSkill | None:
    try:
        text = path.read_text(encoding='utf-8')
    except OSError:
        return None
    metadata, body = _parse_front_matter(text)
    name = metadata.get('name') or path.parent.name
    description = metadata.get('description', '').strip()
    if not name.strip() or not description:
        return None
    return BundledSkill(
        name=name.strip(),
        description=description,
        when_to_use=metadata.get('when_to_use', '').strip(),
        aliases=_split_csv(metadata.get('aliases')),
        allowed_tools=_split_csv(metadata.get('allowed_tools')),
        user_invocable=_parse_bool(metadata.get('user_invocable'), default=True),
        source=source,
        path=str(path),
        get_prompt=_directory_skill_prompt(body.strip()),
    )


def load_directory_skills(
    root: Path | None = None,
    *,
    source: str = 'directory',
) -> tuple[BundledSkill, ...]:
    skill_root = root or (Path(__file__).resolve().parent / 'skills' / 'bundled')
    if not skill_root.exists() or not skill_root.is_dir():
        return ()
    skills: list[BundledSkill] = []
    for path in sorted(skill_root.glob('*/SKILL.md')):
        skill = _load_directory_skill(path, source=source)
        if skill is not None:
            skills.append(skill)
    return tuple(skills)


def load_project_skills(cwd: Path | str | None) -> tuple[BundledSkill, ...]:
    if cwd is None:
        return ()
    return load_directory_skills(Path(cwd).resolve() / 'skills', source='project')


# ---------------------------------------------------------------------------
# Skill registry
# ---------------------------------------------------------------------------

PYTHON_BUNDLED_SKILLS: tuple[BundledSkill, ...] = (
    BundledSkill(
        name='simplify',
        description='Review changed code for reuse, quality, and efficiency, then fix issues.',
        when_to_use='When the user asks to review, simplify, or clean up recent code changes.',
        allowed_tools=('read_file', 'edit_file', 'write_file', 'bash', 'grep_search', 'glob_search'),
        get_prompt=_simplify_prompt,
    ),
    BundledSkill(
        name='debug',
        description='Debug the current session — show diagnostics, token usage, and config.',
        when_to_use='When the user asks to debug the agent session or check diagnostics.',
        user_invocable=True,
        get_prompt=_debug_prompt,
    ),
)


def get_bundled_skills(cwd: Path | str | None = None) -> tuple[BundledSkill, ...]:
    """Return registered skills, with project skills taking precedence."""
    skills = [
        *load_project_skills(cwd),
        *load_directory_skills(),
        *PYTHON_BUNDLED_SKILLS,
    ]
    seen: set[str] = set()
    deduped: list[BundledSkill] = []
    for skill in skills:
        lowered = skill.name.lower()
        if lowered in seen:
            continue
        seen.add(lowered)
        deduped.append(skill)
    return tuple(deduped)


def find_bundled_skill(name: str, cwd: Path | str | None = None) -> BundledSkill | None:
    """Look up a bundled skill by name or alias."""
    lowered = name.lower()
    for skill in get_bundled_skills(cwd):
        if lowered == skill.name or lowered in skill.aliases:
            return skill
    return None


def format_skills_for_system_prompt(
    max_chars: int = 8000,
    cwd: Path | str | None = None,
) -> str:
    """Format bundled skills for inclusion in system-reminder messages.

    The model discovers available skills through this listing.
    """
    lines = ['Available skills (invoke via Skill tool):']
    char_count = len(lines[0])

    for skill in get_bundled_skills(cwd):
        if not skill.user_invocable:
            continue
        entry = f'- {skill.name}: {skill.description}'
        if skill.when_to_use:
            entry += f' When to use: {skill.when_to_use}'
        if len(entry) > 250:
            entry = entry[:247] + '...'
        if char_count + len(entry) + 1 > max_chars:
            break
        lines.append(entry)
        char_count += len(entry) + 1

    return '\n'.join(lines)