traceforge-toolkit 0.1.0__py3-none-any.whl

traceforge/classify/shell.py

@@ -0,0 +1,503 @@
+"""Shell command classification via tree-sitter AST analysis."""
+
+from __future__ import annotations
+
+import os
+import re
+from collections import defaultdict
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+import tree_sitter as ts
+import tree_sitter_bash as tsbash
+
+from traceforge.classify.core import (
+    Classification,
+    Effect,
+    PhaseSegment,
+    Structure,
+    aggregate_effect,
+)
+from traceforge.classify.coding import (
+    CodingAction,
+    CodingMechanism,
+    ShellStructure,
+)
+from traceforge.classify.rules import (
+    SHELL_DELIVERY,
+    SHELL_SETUP,
+    SHELL_VERIFICATION,
+    ShellActivity,
+    classify_binary,
+    effect_for_binary,
+    match_rule,
+)
+from traceforge.classify.workflow import Phase
+
+if TYPE_CHECKING:
+    from traceforge.classify.config import ClassificationEngine
+
+_BASH_LANGUAGE = ts.Language(tsbash.language())
+_parser = ts.Parser(_BASH_LANGUAGE)
+_Q_COMMANDS = ts.Query(_BASH_LANGUAGE, "(command) @cmd")
+
+_IS_ENV_VAR = re.compile(r"^[a-zA-Z_][a-zA-Z0-9_]*=")
+_STRIP_BINARY_EXT = re.compile(r"\.(exe|cmd|bat|ps1|sh)$", re.IGNORECASE)
+
+
+# ── AST helpers ──
+
+
+def _inside_command_substitution(node: ts.Node) -> bool:
+    p = node.parent
+    while p:
+        if p.type == "command_substitution":
+            return True
+        p = p.parent
+    return False
+
+
+def _words_from_command_node(node: ts.Node) -> list[str]:
+    """Extract word tokens from a command AST node, including quoted strings."""
+    _EXTRACTABLE = frozenset({"word", "raw_string", "string", "simple_expansion", "concatenation"})
+    words: list[str] = []
+    for child in node.children:
+        if child.type == "command_name":
+            for sub in child.children:
+                if sub.type in _EXTRACTABLE and sub.text:
+                    # Strip surrounding quotes from string literals
+                    text = sub.text.decode("utf-8")
+                    if len(text) >= 2 and text[0] in ("'", '"') and text[-1] == text[0]:
+                        text = text[1:-1]
+                    words.append(text)
+        elif child.type in _EXTRACTABLE and child.text:
+            text = child.text.decode("utf-8")
+            if len(text) >= 2 and text[0] in ("'", '"') and text[-1] == text[0]:
+                text = text[1:-1]
+            words.append(text)
+    return words
+
+
+# Wrapper flags that consume exactly one argument (e.g., sudo -u <user>)
+_WRAPPER_FLAGS_WITH_ARG: dict[str, frozenset[str]] = {
+    "sudo": frozenset(
+        {"-u", "--user", "-g", "--group", "-C", "-D", "-R", "--role", "-t", "--type"}
+    ),
+    "env": frozenset({"-u", "--unset", "-S", "--split-string"}),
+    "timeout": frozenset({"--signal", "-s", "-k", "--kill-after"}),
+    "nice": frozenset({"-n", "--adjustment"}),
+    "ionice": frozenset({"-c", "--class", "-n", "--classdata", "-p", "--pid"}),
+    "chroot": frozenset({"--userspec", "--groups"}),
+    "runuser": frozenset({"-u", "--user", "-g", "--group", "-G", "--supp-group"}),
+}
+
+
+def _looks_like_command(token: str) -> bool:
+    return bool(token) and not token[0].isdigit() and "/" not in token
+
+
+def _unwrap_binary(
+    words: list[str],
+    engine: ClassificationEngine,
+) -> tuple[str, str | None, list[str], frozenset[str]]:
+    """Extract binary, subcmd, flags, and wrapper capabilities from a word list.
+
+    Returns (binary, subcmd, flags, wrapper_caps) where wrapper_caps contains
+    capabilities derived from stripped wrappers (e.g., elevated_privilege from sudo).
+    """
+    transparent = engine.transparent_wrappers
+    idx = 0
+    wrapper_caps: set[str] = set()
+
+    while idx < len(words) and _IS_ENV_VAR.match(words[idx]):
+        idx += 1
+
+    limit = 5
+    while limit > 0 and idx < len(words):
+        binary = _STRIP_BINARY_EXT.sub("", os.path.basename(words[idx]).lower())
+
+        if binary not in transparent:
+            break
+
+        # Track capabilities from wrappers
+        if binary == "sudo":
+            wrapper_caps.add("elevated_privilege")
+
+        known_arg_flags = _WRAPPER_FLAGS_WITH_ARG.get(binary, frozenset())
+        idx += 1
+        # Skip env-style VAR=val assignments
+        while idx < len(words) and "=" in words[idx] and not words[idx].startswith("-"):
+            idx += 1
+        # Skip wrapper flags
+        while idx < len(words) and words[idx].startswith("-"):
+            flag = words[idx]
+            idx += 1
+            # Flags with = already have their value inline (e.g., --timeout=5)
+            if "=" in flag:
+                continue
+            # Known flags that take an argument: always consume next token
+            if flag in known_arg_flags:
+                if idx < len(words) and not words[idx].startswith("-"):
+                    idx += 1
+                continue
+            # Short flags that take an argument: consume next token if it
+            # doesn't look like another flag or a command name
+            if (
+                idx < len(words)
+                and not words[idx].startswith("-")
+                and "=" not in words[idx]
+                and not _looks_like_command(words[idx])
+            ):
+                idx += 1
+        # For 'timeout', consume the duration positional before the command
+        if binary == "timeout" and idx < len(words) and not words[idx].startswith("-"):
+            # Duration arg: skip if it looks like a number/duration (e.g., 5, 5s, 10m)
+            token = words[idx]
+            if token and (token[0].isdigit() or token in ("infinity",)):
+                idx += 1
+        limit -= 1
+
+    if idx >= len(words):
+        return "", None, [], frozenset(wrapper_caps)
+
+    binary = _STRIP_BINARY_EXT.sub("", os.path.basename(words[idx]).lower())
+
+    remaining = words[idx + 1 :]
+    subcmd = remaining[0] if remaining and not remaining[0].startswith("-") else None
+    flags = [w for w in remaining if w.startswith("-")]
+    return binary, subcmd, flags, frozenset(wrapper_caps)
+
+
+# ── Public API ──
+
+
+def _extract_commands_from_ast(command: str) -> list[str]:
+    """Extract individual command texts from a shell string via tree-sitter query."""
+    if not command or not command.strip():
+        return []
+
+    tree = _parser.parse(command.encode("utf-8"))
+    cursor = ts.QueryCursor(_Q_COMMANDS)
+    matches = cursor.matches(tree.root_node)
+
+    commands: list[str] = []
+    for _pat, captures in matches:
+        for node in captures.get("cmd", []):
+            if _inside_command_substitution(node):
+                continue
+            text = node.text.decode("utf-8").strip() if node.text else ""
+            if text:
+                commands.append(text)
+
+    if not commands:
+        commands.append(command.strip())
+
+    return commands
+
+
+@dataclass(frozen=True)
+class _CommandClassification:
+    """Per-command classification result (internal helper)."""
+
+    binary: str
+    activity: ShellActivity
+    action: str
+    scope: str
+    role: str
+    phase: str
+    capabilities: frozenset[str]
+    effect: str | None
+
+
+def classify_single_command(
+    binary: str,
+    subcmd: str | None,
+    flags: list[str],
+    words: list[str] | None = None,
+    *,
+    engine: ClassificationEngine,
+) -> _CommandClassification:
+    """Classify a single binary invocation into its dimensions.
+
+    Shared across bash/powershell/cmd classifiers so all produce consistent
+    Classification objects with phase_map.
+    """
+    activity = classify_binary(binary, subcmd, flags, words, engine=engine)
+    rule = match_rule(binary, subcmd, flags, engine=engine)
+
+    bi = engine.binary_info
+    act_to_action = engine.activity_to_action
+    act_to_scope = engine.activity_to_scope
+    act_to_phase = engine.activity_to_phase
+    git_subcmd_action = engine.git_subcmd_actions
+    verif_role_action = engine.verification_role_actions
+
+    # Role (rule-based, fallback to binary info)
+    cmd_role = ""
+    if rule and rule.role:
+        cmd_role = rule.role
+    else:
+        info = bi.get(binary)
+        if info:
+            cmd_role = info.role
+
+    # Action: rule override > per-subcommand precision > activity default
+    cmd_action = ""
+    if rule and rule.action:
+        cmd_action = rule.action
+    elif activity == SHELL_DELIVERY and binary == "git" and subcmd:
+        cmd_action = git_subcmd_action.get(subcmd, CodingAction.COMMIT)
+    elif activity == SHELL_VERIFICATION and cmd_role:
+        cmd_action = verif_role_action.get(cmd_role, CodingAction.TEST)
+    else:
+        cmd_action = act_to_action.get(activity, "")
+
+    # Scope: rule override > activity default
+    cmd_scope = ""
+    if rule and rule.scope:
+        cmd_scope = rule.scope
+    else:
+        cmd_scope = act_to_scope.get(activity, "")
+
+    # Phase: rule override > activity default
+    cmd_phase = ""
+    if rule and rule.phase:
+        cmd_phase = rule.phase
+    else:
+        cmd_phase = act_to_phase.get(activity, Phase.IMPLEMENTATION)
+
+    # Capabilities
+    caps: set[str] = set()
+    info = bi.get(binary)
+    if info and info.network:
+        caps.add("network_outbound")
+    if binary == "git" and subcmd in ("push", "pull", "fetch", "clone"):
+        caps.add("network_outbound")
+    if activity == SHELL_SETUP:
+        caps.add("filesystem_write")
+        caps.add("network_outbound")
+    if binary == "sudo":
+        caps.add("elevated_privilege")
+
+    # Effect: rule override > effect_overrides > binary_info
+    if rule and rule.effect:
+        effect = rule.effect
+    else:
+        effect = effect_for_binary(binary, subcmd, flags, engine=engine, all_words=words)
+
+    return _CommandClassification(
+        binary=binary,
+        activity=activity,
+        action=cmd_action,
+        scope=cmd_scope,
+        role=cmd_role,
+        phase=cmd_phase,
+        capabilities=frozenset(caps),
+        effect=effect,
+    )
+
+
+def build_classification_from_commands(
+    command_results: list[_CommandClassification],
+    structure: frozenset[str] = frozenset(),
+    shell_dialect: str | None = None,
+) -> Classification:
+    """Build a Classification from a list of per-command results.
+
+    Shared across bash/powershell/cmd classifiers.
+    """
+    all_roles: set[str] = set()
+    all_actions: set[str] = set()
+    all_scopes: set[str] = set()
+    all_capabilities: set[str] = {"subprocess"}
+    all_effects: list[str] = []
+    all_binaries: list[str] = []
+    phase_actions: dict[str, set[str]] = defaultdict(set)
+    phase_scopes: dict[str, set[str]] = defaultdict(set)
+    phase_roles: dict[str, set[str]] = defaultdict(set)
+
+    for cmd in command_results:
+        all_binaries.append(cmd.binary)
+        if cmd.role:
+            all_roles.add(cmd.role)
+        if cmd.action:
+            all_actions.add(cmd.action)
+        if cmd.scope:
+            all_scopes.add(cmd.scope)
+        all_capabilities.update(cmd.capabilities)
+        if cmd.effect:
+            all_effects.append(cmd.effect)
+
+        # Group by phase
+        if cmd.action:
+            phase_actions[cmd.phase].add(cmd.action)
+        if cmd.scope:
+            phase_scopes[cmd.phase].add(cmd.scope)
+        if cmd.role:
+            phase_roles[cmd.phase].add(cmd.role)
+
+    # Aggregate effect
+    agg_effect = aggregate_effect(*all_effects) if all_effects else None
+
+    # Filesystem capabilities from aggregate effect
+    if agg_effect in (Effect.MUTATING, Effect.DESTRUCTIVE):
+        all_capabilities.add("filesystem_write")
+    else:
+        all_capabilities.add("filesystem_read")
+
+    # Build phase_map
+    all_phases = set(phase_actions.keys()) | set(phase_scopes.keys()) | set(phase_roles.keys())
+    phase_map = tuple(
+        PhaseSegment(
+            phase=phase,
+            actions=frozenset(phase_actions.get(phase, set())),
+            scopes=frozenset(phase_scopes.get(phase, set())),
+            roles=frozenset(phase_roles.get(phase, set())),
+        )
+        for phase in sorted(all_phases)
+    )
+
+    return Classification(
+        mechanism=CodingMechanism.PROCESS_SHELL,
+        effect=agg_effect,
+        scope=frozenset(all_scopes),
+        role=frozenset(all_roles),
+        action=frozenset(all_actions),
+        capability=frozenset(all_capabilities),
+        structure=structure,
+        shell_dialect=shell_dialect,
+        binaries=tuple(dict.fromkeys(all_binaries)),
+        phase_map=phase_map,
+    )
+
+
+def _detect_structure(tree: ts.Tree) -> frozenset[str]:
+    """Detect structural properties from the AST."""
+    structures: set[str] = set()
+    root = tree.root_node
+
+    for child in root.children:
+        if child.type == "list":
+            structures.add(Structure.SEQUENTIAL)
+        elif child.type == "pipeline":
+            if child.child_count > 1:
+                structures.add(ShellStructure.PIPED)
+
+    _check_tree_structure(root, structures)
+    return frozenset(structures)
+
+
+_OUTPUT_REDIRECT_OPS = frozenset({">", ">>", ">&", "&>", "&>>"})
+
+
+def _has_output_redirect(node: ts.Node) -> bool:
+    """Check if a redirected_statement contains an output redirect (not input-only)."""
+    for child in node.children:
+        if child.type == "file_redirect":
+            for op in child.children:
+                if op.type in _OUTPUT_REDIRECT_OPS:
+                    return True
+        elif child.type == "heredoc_redirect":
+            # Heredocs are input, not output
+            continue
+    return False
+
+
+def _check_tree_structure(node: ts.Node, structures: set[str]) -> None:
+    """Recursively check for structural patterns."""
+    if node.type == "redirected_statement":
+        if _has_output_redirect(node):
+            structures.add(ShellStructure.REDIRECTED)
+    elif node.type == "pipeline" and node.child_count > 1:
+        structures.add(ShellStructure.PIPED)
+    elif node.type in ("if_statement", "case_statement"):
+        structures.add(Structure.CONDITIONAL)
+
+    for child in node.children:
+        _check_tree_structure(child, structures)
+
+
+def classify_shell(
+    command: str,
+    *,
+    engine: ClassificationEngine,
+) -> Classification:
+    """Classify a bash shell command into a Classification object.
+
+    For compound commands (e.g. `pytest && git push`), each command is classified
+    independently. Actions, scopes, and roles are grouped by derived phase in
+    `phase_map`, and also unioned into the top-level aggregate sets.
+    """
+    if not command or not command.strip():
+        return Classification(
+            mechanism=CodingMechanism.PROCESS_SHELL,
+            effect=None,
+        )
+
+    tree = _parser.parse(command.encode("utf-8"))
+    cursor = ts.QueryCursor(_Q_COMMANDS)
+    matches = cursor.matches(tree.root_node)
+
+    if not matches:
+        return Classification(
+            mechanism=CodingMechanism.PROCESS_SHELL,
+            effect=None,
+            capability=frozenset({"subprocess"}),
+        )
+
+    command_results: list[_CommandClassification] = []
+    all_wrapper_caps: set[str] = set()
+
+    for _pattern_idx, captures in matches:
+        for node in captures.get("cmd", []):
+            if _inside_command_substitution(node):
+                continue
+
+            words = _words_from_command_node(node)
+            if not words:
+                continue
+
+            binary, subcmd, flags, wrapper_caps = _unwrap_binary(words, engine=engine)
+            all_wrapper_caps.update(wrapper_caps)
+            if not binary:
+                continue
+
+            cmd_cls = classify_single_command(binary, subcmd, flags, words, engine=engine)
+            command_results.append(cmd_cls)
+
+    if not command_results:
+        return Classification(
+            mechanism=CodingMechanism.PROCESS_SHELL,
+            effect=None,
+            capability=frozenset({"subprocess"} | all_wrapper_caps),
+        )
+
+    structure = _detect_structure(tree)
+    cls = build_classification_from_commands(command_results, structure, "bash")
+    # Merge wrapper capabilities (e.g., elevated_privilege from sudo)
+    extra_caps: set[str] = set(all_wrapper_caps)
+    extra_effect: str | None = None
+    # Redirections (>, >>, 2>) indicate filesystem writes
+    if ShellStructure.REDIRECTED in structure:
+        extra_caps.add("filesystem_write")
+        extra_effect = "mutating"
+    if extra_caps or extra_effect:
+        from traceforge.classify.core import aggregate_effect as _agg
+
+        merged_effect = cls.effect
+        if extra_effect:
+            effects = [e for e in (cls.effect, extra_effect) if e]
+            merged_effect = _agg(*effects) if effects else cls.effect
+        cls = Classification(
+            mechanism=cls.mechanism,
+            effect=merged_effect,
+            scope=cls.scope,
+            role=cls.role,
+            action=cls.action,
+            capability=cls.capability | frozenset(extra_caps),
+            structure=cls.structure,
+            shell_dialect=cls.shell_dialect,
+            binaries=cls.binaries,
+            phase_map=cls.phase_map,
+        )
+    return cls

traceforge/classify/tools.py

@@ -0,0 +1,91 @@
+"""Tool name normalization and classification."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from traceforge.classify.core import Classification, Mechanism
+from traceforge.classify.phases import with_phase_map as _with_phase_map
+
+if TYPE_CHECKING:
+    from traceforge.classify.config import ClassificationEngine
+
+
+def normalize_tool_name(
+    raw_name: str,
+    engine: ClassificationEngine,
+) -> str:
+    """Normalize a raw tool name to its canonical form."""
+    if not raw_name:
+        return raw_name
+
+    name = raw_name.strip()
+
+    if name.startswith("mcp__"):
+        parts = name.split("__", 2)
+        if len(parts) == 3:
+            name = parts[2]
+    elif "." in name:
+        dot_idx = name.index(".")
+        prefix = name[:dot_idx]
+        if prefix.replace("_", "").isalpha() and prefix.islower():
+            name = name[dot_idx + 1 :]
+
+    lowered = name.lower().replace("-", "_")
+    canonical_map = engine.canonical_tools
+    return canonical_map.get(lowered, lowered)
+
+
+def classify_tool(
+    tool_name: str,
+    custom_classifications: dict[str, Classification] | None = None,
+    *,
+    engine: ClassificationEngine,
+) -> Classification:
+    """Classify a tool name into a full Classification object.
+
+    All returned Classifications carry phase_map — same system as shell commands.
+
+    Classification priority:
+    1. Custom user-provided classifications
+    2. MCP server profiles (checked on raw name BEFORE canonical normalization,
+       so that MCP-specific suffixes like 'search' don't collide with first-party
+       canonical aliases like 'grep')
+    3. Built-in canonical tool classifications
+    4. MCP verb inference for MCP-formatted names with unknown namespaces
+    5. UNKNOWN mechanism fallback
+    """
+    from traceforge.classify.mcp import _infer_from_verb, classify_mcp_tool
+
+    if not tool_name:
+        fallback = Classification(mechanism=Mechanism.UNKNOWN, effect=None)
+        return _with_phase_map(fallback)
+
+    canonical = normalize_tool_name(tool_name, engine=engine)
+    if custom_classifications:
+        lower = canonical.lower()
+        for key, cls in custom_classifications.items():
+            if key.lower() == lower or normalize_tool_name(key, engine=engine) == canonical:
+                if not cls.phase_map:
+                    return _with_phase_map(cls)
+                return cls
+
+    mcp_result = classify_mcp_tool(tool_name, engine=engine)
+    if mcp_result is not None:
+        return mcp_result
+
+    tool_cls_map = engine.tool_classifications
+    result = tool_cls_map.get(canonical)
+    if result is not None:
+        return result
+
+    verb_effect, verb_action = _infer_from_verb(canonical, engine=engine)
+    if verb_effect is not None or verb_action is not None:
+        cls = Classification(
+            mechanism=Mechanism.UNKNOWN,
+            effect=verb_effect,
+            action=frozenset({verb_action}) if verb_action else frozenset(),
+        )
+        return _with_phase_map(cls)
+
+    return _with_phase_map(Classification(mechanism=Mechanism.UNKNOWN, effect=None))

traceforge/classify/workflow.py

@@ -0,0 +1,32 @@
+"""Session workflow dimensions — phase and visibility.
+
+These are derived/presentation concerns, not intrinsic semantic properties
+of an action. They live separately from the classification taxonomy.
+"""
+
+from __future__ import annotations
+
+from enum import StrEnum
+
+
+class Phase(StrEnum):
+    """Workflow stage inferred from classification dimensions of recent events."""
+
+    PLANNING = "planning"
+    IMPLEMENTATION = "implementation"
+    VERIFICATION = "verification"
+    EXPLORATION = "exploration"
+    REVIEW = "review"
+
+
+class Visibility(StrEnum):
+    """Display audience for an event in a trace view.
+
+    SYSTEM = bookkeeping, not shown to end users.
+    VISIBLE = shown by default.
+    COLLAPSED = shown but minimized (e.g. repeated similar events).
+    """
+
+    VISIBLE = "visible"
+    SYSTEM = "system"
+    COLLAPSED = "collapsed"

traceforge/cli/__init__.py

@@ -0,0 +1,42 @@
+"""Traceforge CLI — Click application entry point."""
+
+from __future__ import annotations
+
+import sys
+
+import click
+
+from traceforge.cli.detect import detect
+from traceforge.cli.config_cmd import config
+from traceforge.cli.watch import watch
+from traceforge.cli.score import score
+from traceforge.cli.replay import replay
+from traceforge.cli.status import status
+from traceforge.cli.gate_cmd import gate
+from traceforge.cli.init_cmd import init
+from traceforge.cli.download_cmd import download_model
+
+
+@click.group()
+@click.version_option(package_name="traceforge-toolkit")
+def main() -> None:
+    """Traceforge — governance pipeline for AI coding agents."""
+    # Force UTF-8 on stdout/stderr so the CLI's Unicode glyphs (box rules,
+    # ✓/✗) never raise UnicodeEncodeError on a non-UTF-8 console such as a
+    # Windows cp1252 stdout. Guarded because streams may lack reconfigure
+    # (e.g. when replaced by an in-memory buffer under test).
+    for stream in (sys.stdout, sys.stderr):
+        reconfigure = getattr(stream, "reconfigure", None)
+        if reconfigure is not None:
+            reconfigure(encoding="utf-8")
+
+
+main.add_command(watch)
+main.add_command(score)
+main.add_command(detect)
+main.add_command(replay)
+main.add_command(config)
+main.add_command(status)
+main.add_command(gate)
+main.add_command(init)
+main.add_command(download_model)