code-muse 0.0.1__py3-none-any.whl

code_muse/plugins/policy_engine/policy_evaluator.py

@@ -0,0 +1,75 @@
+import logging
+from fnmatch import fnmatch
+
+from code_muse.plugins.policy_engine.policy_toml_schema import Decision, ToolRule
+
+logger = logging.getLogger(__name__)
+
+
+def _tool_name_matches(tool_name: str, rule_tool_name: str) -> bool:
+    if rule_tool_name == "*":
+        return True
+    # Support fnmatch wildcards (e.g., "agent_run_*")
+    return fnmatch(tool_name, rule_tool_name)
+
+
+def evaluate_policy(
+    tool_name: str,
+    command: str | None,
+    rules: list[ToolRule],
+) -> tuple[Decision, ToolRule | None]:
+    matched: list[ToolRule] = []
+    for rule in rules:
+        if not _tool_name_matches(tool_name, rule.tool_name):
+            continue
+        if rule.command_prefix is not None:
+            if command is None:
+                # command_prefix specified but no command provided — skip
+                continue
+            if not command.startswith(rule.command_prefix):
+                continue
+        matched.append(rule)
+
+    if not matched:
+        logger.debug(
+            "Policy: no rules matched for tool=%s command=%s — default ALLOW",
+            tool_name,
+            command,
+        )
+        return (Decision.ALLOW, None)
+
+    # Sort by priority descending
+    matched.sort(key=lambda r: r.priority, reverse=True)
+    highest_priority = matched[0].priority
+    top_rules = [r for r in matched if r.priority == highest_priority]
+
+    if len(top_rules) > 1:
+        decisions = {r.decision for r in top_rules}
+        if len(decisions) > 1:
+            logger.warning(
+                "Policy conflict: %d rules at priority %d with different decisions "
+                "for tool=%s command=%s. Using first registered: %s",
+                len(top_rules),
+                highest_priority,
+                tool_name,
+                command,
+                top_rules[0],
+            )
+
+    chosen = top_rules[0]
+    logger.debug(
+        "Policy: matched rule for tool=%s command=%s → %s (priority=%d, source=%s)",
+        tool_name,
+        command,
+        chosen.decision.value,
+        chosen.priority,
+        chosen.source,
+    )
+    return (chosen.decision, chosen)
+
+
+def evaluate_tool_policy(
+    tool_name: str,
+    rules: list[ToolRule],
+) -> tuple[Decision, ToolRule | None]:
+    return evaluate_policy(tool_name, None, rules)

code_muse/plugins/policy_engine/policy_file_discovery.py

@@ -0,0 +1,90 @@
+import logging
+import os
+from pathlib import Path
+
+from code_muse.plugins.policy_engine.policy_toml_schema import (
+    ToolRule,
+    parse_policy_toml,
+)
+
+logger = logging.getLogger(__name__)
+
+# Simple cache: list of rules + dict of file mtimes
+_rule_cache: list[ToolRule] | None = None
+_file_mtimes: dict[str, float] = {}
+
+
+def _get_user_policies_dir() -> Path:
+    return Path.home() / ".muse" / "policies"
+
+
+def _get_project_policies_dir() -> Path:
+    return Path(os.getcwd()) / ".muse" / "policies"
+
+
+def discover_policy_files() -> list[Path]:
+    files: list[Path] = []
+    for directory in (_get_user_policies_dir(), _get_project_policies_dir()):
+        if not directory.exists():
+            continue
+        try:
+            toml_files = sorted(directory.glob("*.toml"))
+        except OSError as exc:
+            logger.warning("Cannot scan policy directory %s: %s", directory, exc)
+            continue
+        for f in toml_files:
+            if not f.is_file():
+                continue
+            try:
+                # Check readability
+                _ = f.stat()
+            except OSError as exc:
+                logger.warning("Unreadable policy file %s: %s", f, exc)
+                continue
+            files.append(f)
+    return files
+
+
+def _files_changed(files: list[Path]) -> bool:
+    global _file_mtimes
+    current: dict[str, float] = {}
+    for f in files:
+        try:
+            current[str(f)] = f.stat().st_mtime
+        except OSError:
+            return True
+    return current != _file_mtimes
+
+
+def load_all_policies(force_reload: bool = False) -> list[ToolRule]:
+    global _rule_cache, _file_mtimes
+
+    files = discover_policy_files()
+
+    if not force_reload and _rule_cache is not None and not _files_changed(files):
+        return _rule_cache
+
+    all_rules: list[ToolRule] = []
+    new_mtimes: dict[str, float] = {}
+
+    for f in files:
+        try:
+            rules = parse_policy_toml(f)
+            all_rules.extend(rules)
+            new_mtimes[str(f)] = f.stat().st_mtime
+        except ValueError as exc:
+            logger.warning("Skipping invalid policy file %s: %s", f, exc)
+        except OSError as exc:
+            logger.warning("Cannot read policy file %s: %s", f, exc)
+
+    _rule_cache = all_rules
+    _file_mtimes = new_mtimes
+    logger.info("Loaded %d policy rules from %d file(s)", len(all_rules), len(files))
+    return all_rules
+
+
+def clear_policy_cache() -> None:
+    global _rule_cache, _file_mtimes
+    _rule_cache = None
+    _file_mtimes = {}
+    logger.info("Policy cache cleared")

code_muse/plugins/policy_engine/policy_toml_schema.py

@@ -0,0 +1,115 @@
+import logging
+from dataclasses import dataclass, field
+from enum import Enum
+from pathlib import Path
+from typing import Any
+
+logger = logging.getLogger(__name__)
+
+
+class Decision(Enum):
+    ALLOW = "allow"
+    DENY = "deny"
+    ASK_USER = "ask_user"
+
+
+@dataclass
+class ToolRule:
+    tool_name: str
+    decision: Decision
+    command_prefix: str | None = None
+    priority: int = 0
+    description: str = ""
+    source: str = field(default="", repr=False)
+
+
+def _warn_unknown_fields(rule_table: dict[str, Any], path: str | Path) -> None:
+    known = {"toolName", "commandPrefix", "decision", "priority", "description"}
+    unknown = set(rule_table.keys()) - known
+    if unknown:
+        logger.warning("Unknown fields in rule from %s: %s", path, ", ".join(unknown))
+
+
+def parse_policy_toml(path: str | Path) -> list[ToolRule]:
+    path = Path(path)
+    import tomllib
+
+    raw = path.read_text(encoding="utf-8")
+    data = tomllib.loads(raw)
+
+    schema_version = data.get("schema_version")
+    if schema_version is not None and str(schema_version) != "1":
+        logger.warning(
+            "Policy file %s has schema_version %s; expected 1. "
+            "Forward compatibility assumed.",
+            path,
+            schema_version,
+        )
+
+    rules: list[ToolRule] = []
+    rule_tables = data.get("rule", [])
+    if isinstance(rule_tables, dict):
+        # Single [[rule]] can be parsed as dict by some TOML libs
+        rule_tables = [rule_tables]
+
+    for idx, rule_table in enumerate(rule_tables):
+        if not isinstance(rule_table, dict):
+            raise ValueError(
+                f"Invalid rule entry at index {idx} in {path}: expected table, got {type(rule_table).__name__}"
+            )
+
+        _warn_unknown_fields(rule_table, path)
+
+        tool_name = rule_table.get("toolName")
+        if tool_name is None or not isinstance(tool_name, str) or not tool_name.strip():
+            raise ValueError(f"Rule {idx + 1} in {path}: missing or invalid 'toolName'")
+
+        decision_str = rule_table.get("decision")
+        if decision_str is None or not isinstance(decision_str, str):
+            raise ValueError(f"Rule {idx + 1} in {path}: missing or invalid 'decision'")
+        try:
+            decision = Decision(decision_str)
+        except ValueError as exc:
+            raise ValueError(
+                f"Rule {idx + 1} in {path}: invalid decision '{decision_str}'. "
+                f"Must be one of: {', '.join(d.value for d in Decision)}"
+            ) from exc
+
+        command_prefix = rule_table.get("commandPrefix")
+        if command_prefix is not None and not isinstance(command_prefix, str):
+            raise ValueError(
+                f"Rule {idx + 1} in {path}: 'commandPrefix' must be a string"
+            )
+
+        priority = rule_table.get("priority", 0)
+        if not isinstance(priority, int):
+            raise ValueError(f"Rule {idx + 1} in {path}: 'priority' must be an integer")
+
+        description = rule_table.get("description", "")
+        if description is not None and not isinstance(description, str):
+            raise ValueError(
+                f"Rule {idx + 1} in {path}: 'description' must be a string"
+            )
+
+        rules.append(
+            ToolRule(
+                tool_name=tool_name,
+                decision=decision,
+                command_prefix=command_prefix if command_prefix else None,
+                priority=priority,
+                description=description,
+                source=str(path),
+            )
+        )
+
+    return rules
+
+
+def validate_rules(rules: list[ToolRule]) -> None:
+    for rule in rules:
+        if not rule.tool_name or not rule.tool_name.strip():
+            raise ValueError(f"Invalid rule: tool_name is empty (from {rule.source})")
+        if rule.decision == Decision.DENY and not rule.tool_name.strip():
+            raise ValueError(
+                f"Invalid rule: deny with empty tool_name (from {rule.source})"
+            )

code_muse/plugins/policy_engine/register_callbacks.py

@@ -0,0 +1,112 @@
+import logging
+from typing import Any
+
+from code_muse.callbacks import register_callback
+from code_muse.messaging import emit_info
+from code_muse.plugins.policy_engine.approval_flow_integration import (
+    integrate_policy_check,
+)
+from code_muse.plugins.policy_engine.policy_file_discovery import (
+    clear_policy_cache,
+    load_all_policies,
+)
+
+logger = logging.getLogger(__name__)
+
+# Load policies at import time so they're available immediately.
+# This is safe because load_all_policies caches results and logs warnings
+# for invalid files rather than raising.
+_INITIALIZED_RULES = load_all_policies()
+logger.info(
+    "Policy engine initialized with %d rule(s)",
+    len(_INITIALIZED_RULES),
+)
+
+
+async def _on_run_shell_command(
+    context: Any,
+    command: str,
+    cwd: str | None = None,
+    timeout: int = 60,
+) -> dict[str, Any | None]:
+    """Policy check for shell commands.
+
+    Returns:
+        - {"auto_approve": True} if policy says ALLOW
+        - {"blocked": True, "error_message": "..."} if policy says DENY
+        - None for ASK_USER or no match (normal confirmation flow)
+    """
+    return integrate_policy_check("agent_run_shell_command", command)
+
+
+async def _on_pre_tool_call(
+    tool_name: str,
+    tool_args: dict,
+    context: Any = None,
+) -> dict[str, Any | None]:
+    """Policy check for all tool calls.
+
+    Returns:
+        - {"blocked": True, "error_message": "..."} if policy says DENY
+        - None otherwise (ALLOW, ASK_USER, or no match)
+    """
+    result = integrate_policy_check(tool_name, None)
+    if result and result.get("blocked"):
+        return result
+    return None
+
+
+def _on_custom_command(command: str, name: str) -> Any:
+    """Handle /policies and /policies reload commands."""
+    if name == "policies":
+        parts = command.split(maxsplit=1)
+        subcommand = parts[1].strip().lower() if len(parts) > 1 else ""
+
+        if subcommand == "reload":
+            clear_policy_cache()
+            rules = load_all_policies(force_reload=True)
+            emit_info(f"[Run] Policy rules reloaded: {len(rules)} active rule(s)")
+            return True
+
+        # Default /policies — list loaded rules
+        rules = load_all_policies()
+        if not rules:
+            emit_info("[Warn] No policy rules loaded.")
+            return True
+
+        lines: list[str] = ["[Run] Active rules:"]
+        seen_sources: set[str] = set()
+        for rule in rules:
+            prefix = f"  [{rule.priority}] {rule.tool_name}"
+            if rule.command_prefix:
+                prefix += f" (prefix='{rule.command_prefix}')"
+            prefix += f" → {rule.decision.value}"
+            if rule.description:
+                prefix += f" — {rule.description}"
+            lines.append(prefix)
+            if rule.source:
+                seen_sources.add(rule.source)
+
+        if seen_sources:
+            lines.append("")
+            lines.append("Sources:")
+            for src in sorted(seen_sources):
+                lines.append(f"  • {src}")
+
+        emit_info("\n".join(lines))
+        return True
+
+    return None
+
+
+def _on_custom_command_help() -> list[tuple[str, str]]:
+    return [
+        ("policies", "List loaded policy rules"),
+        ("policies reload", "Reload policy rules from disk"),
+    ]
+
+
+register_callback("run_shell_command", _on_run_shell_command, priority=50)
+register_callback("pre_tool_call", _on_pre_tool_call, priority=50)
+register_callback("custom_command", _on_custom_command)
+register_callback("custom_command_help", _on_custom_command_help)

code_muse/plugins/pop_command/__init__.py

@@ -0,0 +1 @@
+"""/pop command plugin."""

code_muse/plugins/pop_command/register_callbacks.py

@@ -0,0 +1,189 @@
+"""Plugin that adds /pop for trimming recent conversation history."""
+
+from typing import Any
+
+from code_muse.callbacks import register_callback
+
+
+def emit_error(message: Any) -> None:
+    from code_muse.messaging import emit_error as _emit_error
+
+    _emit_error(message)
+
+
+def emit_info(message: Any) -> None:
+    from code_muse.messaging import emit_info as _emit_info
+
+    _emit_info(message)
+
+
+def emit_success(message: Any) -> None:
+    from code_muse.messaging import emit_success as _emit_success
+
+    _emit_success(message)
+
+
+def emit_warning(message: Any) -> None:
+    from code_muse.messaging import emit_warning as _emit_warning
+
+    _emit_warning(message)
+
+
+def _custom_help() -> list[tuple[str, str]]:
+    return [
+        (
+            "pop",
+            "Delete the N most-recent messages and prune broken tool-call fragments",
+        )
+    ]
+
+
+def _has_only_tool_returns(message: Any) -> bool:
+    """Return True when a request message contains only tool-return parts."""
+    try:
+        from pydantic_ai.messages import ModelRequest, ToolReturnPart
+
+        if not isinstance(message, ModelRequest):
+            return False
+
+        parts = getattr(message, "parts", []) or []
+        return bool(parts) and all(isinstance(part, ToolReturnPart) for part in parts)
+    except Exception:
+        return False
+
+
+def _has_unresolved_tool_calls(message: Any) -> bool:
+    """Return True when a response message still contains unresolved tool calls."""
+    try:
+        from pydantic_ai.messages import ModelResponse, ToolCallPart
+
+        if not isinstance(message, ModelResponse):
+            return False
+
+        parts = getattr(message, "parts", []) or []
+        return any(isinstance(part, ToolCallPart) for part in parts)
+    except Exception:
+        return False
+
+
+def _prune_dangling_tool_fragments(history: list[Any]) -> tuple[list[Any], int]:
+    """Strip incomplete tool-call sequences from the tail of history."""
+    pruned = 0
+
+    while history:
+        tail = history[-1]
+        if _has_only_tool_returns(tail):
+            history.pop()
+            pruned += 1
+            continue
+
+        if _has_unresolved_tool_calls(tail):
+            history.pop()
+            pruned += 1
+            continue
+
+        break
+
+    return history, pruned
+
+
+def _parse_pop_count(command: str) -> int | None:
+    tokens = command.split()
+    if len(tokens) < 2:
+        return 1
+
+    try:
+        count = int(tokens[1])
+    except ValueError:
+        emit_error(f"/pop: '{tokens[1]}' is not a valid integer – usage: /pop [N]")
+        return None
+
+    if count < 1:
+        emit_error("/pop: N must be a positive integer")
+        return None
+
+    return count
+
+
+def _handle_pop_command(command: str) -> bool:
+    from code_muse.agents.agent_manager import get_current_agent
+
+    count = _parse_pop_count(command)
+    if count is None:
+        return True
+
+    try:
+        agent = get_current_agent()
+    except Exception as exc:
+        emit_error(f"/pop: could not get current agent – {exc}")
+        return True
+
+    history: list[Any] = list(agent.get_message_history())
+    if not history:
+        emit_warning("/pop: conversation history is empty – nothing to remove")
+        return True
+
+    poppable = len(history) - 1
+    if poppable <= 0:
+        emit_warning("/pop: only the system prompt is in history – nothing to remove")
+        return True
+
+    if count > poppable:
+        emit_warning(
+            f"/pop: requested {count} but only {poppable} message(s) can be removed "
+            f"(the system prompt is always preserved). Removing {poppable}."
+        )
+        count = poppable
+
+    before_count = len(history)
+    history = history[: before_count - count]
+    history, extra_pruned = _prune_dangling_tool_fragments(history)
+    after_count = len(history)
+    total_removed = before_count - after_count
+
+    try:
+        agent.set_message_history(history)
+    except Exception as exc:
+        emit_error(f"/pop: failed to update message history – {exc}")
+        return True
+
+    summary_parts = [f":scissors: Popped {count} message(s)"]
+    if extra_pruned:
+        summary_parts.append(
+            f"and pruned {extra_pruned} extra incomplete tool-call fragment(s)"
+        )
+
+    remaining = max(after_count - 1, 0)
+    emit_success(
+        " ".join(summary_parts)
+        + ".\n"
+        + f":scroll: History: {before_count - 1} → {remaining} message(s) "
+        + f"(excluding system prompt, removed {total_removed} total)"
+    )
+
+    if after_count <= 1:
+        emit_info(":bulb: History is now empty (system prompt only). Starting fresh!")
+
+    return True
+
+
+def _handle_custom_command(command: str, name: str) -> bool | None:
+    if name != "pop":
+        return None
+
+    return _handle_pop_command(command)
+
+
+register_callback("custom_command_help", _custom_help)
+register_callback("custom_command", _handle_custom_command)
+
+
+__all__ = [
+    "_custom_help",
+    "_handle_custom_command",
+    "_handle_pop_command",
+    "_parse_pop_count",
+    "_prune_dangling_tool_fragments",
+    "_has_only_tool_returns",
+    "_has_unresolved_tool_calls",
+]

code_muse/plugins/prompt_newline/__init__.py

@@ -0,0 +1,13 @@
+"""Plugin: drop user input onto a fresh line below the prompt chrome.
+
+When enabled, transforms
+
+    [Run] prompt [agent] [model] (~/very/long/cwd) >>> typed text
+
+into
+
+    [Run] prompt [agent] [model] (~/very/long/cwd) >>>
+    typed text
+
+Toggle at runtime with ``/prompt_newline [on|off]``. Persisted in muse.cfg.
+"""

code_muse/plugins/prompt_newline/config.py

@@ -0,0 +1,19 @@
+"""Plugin-local config for the prompt_newline plugin."""
+
+from code_muse.config import get_value, set_config_value
+
+_CONFIG_KEY = "prompt_newline"
+_TRUTHY = ("true", "1", "yes", "on")
+
+
+def is_enabled() -> bool:
+    """Return True if the prompt_newline hack is enabled. Default: False."""
+    cfg_val = get_value(_CONFIG_KEY)
+    if cfg_val is None:
+        return False
+    return str(cfg_val).strip().lower() in _TRUTHY
+
+
+def set_enabled(enabled: bool) -> None:
+    """Persist the on/off switch to muse.cfg."""
+    set_config_value(_CONFIG_KEY, "true" if enabled else "false")