code-muse 0.0.1__py3-none-any.whl

code_muse/plugins/checkpointing/conversation_snapshots.py

@@ -0,0 +1,117 @@
+"""Conversation snapshot serialization for checkpointing."""
+
+import json
+import logging
+from datetime import UTC, datetime
+from pathlib import Path
+from typing import Any
+
+from code_muse.agents import get_current_agent
+
+logger = logging.getLogger(__name__)
+
+
+def create_snapshot(
+    agent: Any, tool_name: str, tool_call_id: str | None = None
+) -> Path | None:
+    """Serialize the current conversation state to a JSON snapshot."""
+    try:
+        if agent is None:
+            agent = get_current_agent()
+
+        messages = list(agent.get_message_history())
+        timestamp = datetime.now(UTC).isoformat()
+
+        snapshot: dict[str, Any] = {
+            "turn_id": f"{tool_name}_{timestamp}",
+            "timestamp": timestamp,
+            "tool_name": tool_name,
+            "tool_call_id": tool_call_id or "",
+            "messages": _serialize_messages(messages),
+            "agent_state": {
+                "name": getattr(agent, "name", "unknown"),
+                "model": getattr(agent, "model", "unknown"),
+            },
+        }
+
+        # Store in shadow git repo area
+        project_root = Path.cwd()
+        project_hash = _hash_project_root(str(project_root))
+        repo_path = Path.home() / ".muse" / "history" / project_hash
+        snapshot_dir = repo_path / "snapshots"
+        snapshot_dir.mkdir(parents=True, exist_ok=True)
+
+        snapshot_path = snapshot_dir / f"snapshot_{timestamp.replace(':', '_')}.json"
+        snapshot_path.write_text(json.dumps(snapshot, indent=2), encoding="utf-8")
+
+        logger.info(f"Conversation snapshot saved to {snapshot_path}")
+        return snapshot_path
+    except Exception as exc:
+        logger.error(f"Failed to create snapshot: {exc}")
+        return None
+
+
+def load_snapshot(path: Path) -> dict[str, Any] | None:
+    """Load and validate a snapshot JSON file."""
+    try:
+        data = json.loads(path.read_text(encoding="utf-8"))
+        required_keys = {"turn_id", "timestamp", "tool_name", "messages", "agent_state"}
+        if not required_keys.issubset(data.keys()):
+            logger.warning(f"Snapshot at {path} is missing required keys")
+            return None
+        return data
+    except Exception as exc:
+        logger.error(f"Failed to load snapshot {path}: {exc}")
+        return None
+
+
+def list_snapshots(repo_path: Path) -> list[dict[str, Any]]:
+    """Scan snapshot directory and return metadata sorted newest-first."""
+    snapshot_dir = repo_path / "snapshots"
+    if not snapshot_dir.exists():
+        return []
+
+    results: list[dict[str, Any]] = []
+    for path in snapshot_dir.glob("snapshot_*.json"):
+        data = load_snapshot(path)
+        if data is None:
+            continue
+        results.append(
+            {
+                "timestamp": data.get("timestamp", ""),
+                "tool_name": data.get("tool_name", ""),
+                "tool_call_id": data.get("tool_call_id", ""),
+                "path": str(path),
+            }
+        )
+
+    results.sort(key=lambda x: x["timestamp"], reverse=True)
+    return results
+
+
+def _hash_project_root(project_root: str) -> str:
+    import hashlib
+
+    return hashlib.sha256(project_root.encode()).hexdigest()
+
+
+def _serialize_messages(messages: list[Any]) -> list[dict[str, Any]]:
+    """Best-effort serialization of message history to JSON-safe dicts."""
+    serialized: list[dict[str, Any]] = []
+    for msg in messages:
+        try:
+            if hasattr(msg, "model_dump"):
+                serialized.append(msg.model_dump())
+            elif hasattr(msg, "dict"):
+                serialized.append(msg.dict())
+            else:
+                serialized.append(
+                    {
+                        "type": type(msg).__name__,
+                        "repr": repr(msg),
+                    }
+                )
+        except Exception as exc:
+            logger.warning(f"Could not serialize message: {exc}")
+            serialized.append({"type": "unserializable", "error": str(exc)})
+    return serialized

code_muse/plugins/checkpointing/register_callbacks.py

@@ -0,0 +1,51 @@
+"""Register checkpointing callbacks and commands."""
+
+import logging
+
+from code_muse.callbacks import register_callback
+from code_muse.plugins.checkpointing.checkpoint_hook import (
+    on_pre_tool_call_checkpoint,
+)
+from code_muse.plugins.checkpointing.restore_command import (
+    _custom_help,
+    _handle_custom_command,
+)
+from code_muse.plugins.checkpointing.rewind_shortcut import RewindKeyListener
+
+logger = logging.getLogger(__name__)
+
+_rewind_listener: RewindKeyListener | None = None
+
+
+def _start_rewind_listener() -> None:
+    global _rewind_listener
+    if _rewind_listener is not None:
+        return
+
+    def _on_double_esc() -> None:
+        from code_muse.plugins.checkpointing.restore_command import (
+            _handle_restore_command,
+        )
+
+        _handle_restore_command("/restore")
+
+    _rewind_listener = RewindKeyListener(_on_double_esc)
+    _rewind_listener.start()
+
+
+def _stop_rewind_listener() -> None:
+    global _rewind_listener
+    if _rewind_listener is not None:
+        _rewind_listener.stop()
+        _rewind_listener = None
+
+
+def _on_shutdown() -> None:
+    _stop_rewind_listener()
+
+
+# Register callbacks
+register_callback("pre_tool_call", on_pre_tool_call_checkpoint)
+register_callback("custom_command_help", _custom_help)
+register_callback("custom_command", _handle_custom_command)
+register_callback("shutdown", _on_shutdown)

code_muse/plugins/checkpointing/restore_command.py

@@ -0,0 +1,263 @@
+"""/restore slash command for checkpointing rewind."""
+
+import logging
+import subprocess
+from pathlib import Path
+from typing import Any
+
+from code_muse.tools.common import get_user_approval_async
+
+logger = logging.getLogger(__name__)
+
+
+def _custom_help() -> list[tuple[str, str]]:
+    return [
+        (
+            "restore",
+            "List checkpoints or revert to a previous checkpoint (files and/or conversation)",
+        )
+    ]
+
+
+def _handle_restore_command(command: str) -> bool:
+    from code_muse.messaging import emit_error, emit_info
+    from code_muse.plugins.checkpointing.conversation_snapshots import (
+        list_snapshots,
+        load_snapshot,
+    )
+
+    tokens = command.split()
+    if len(tokens) < 2:
+        # List all checkpoints
+        project_root = Path.cwd()
+        project_hash = _hash_project_root(str(project_root))
+        repo_path = Path.home() / ".muse" / "history" / project_hash
+        snapshots = list_snapshots(repo_path)
+        if not snapshots:
+            emit_info("No checkpoints available yet.")
+            return True
+
+        lines = [":rewind: Available checkpoints:"]
+        for idx, snap in enumerate(snapshots, start=1):
+            ts = snap.get("timestamp", "")
+            tool = snap.get("tool_name", "")
+            lines.append(f"  {idx}. [{ts}] {tool}")
+
+        emit_info("\n".join(lines))
+        emit_info("Usage: /restore <index> [full|files|conversation]")
+        return True
+
+    try:
+        index = int(tokens[1])
+    except ValueError:
+        emit_error("/restore: index must be an integer")
+        return True
+
+    project_root = Path.cwd()
+    project_hash = _hash_project_root(str(project_root))
+    repo_path = Path.home() / ".muse" / "history" / project_hash
+    snapshots = list_snapshots(repo_path)
+    if not snapshots:
+        emit_info("No checkpoints available yet.")
+        return True
+
+    if index < 1 or index > len(snapshots):
+        emit_error(f"/restore: index {index} out of range (1–{len(snapshots)})")
+        return True
+
+    selected = snapshots[index - 1]
+    snapshot = load_snapshot(Path(selected["path"]))
+    if snapshot is None:
+        emit_error("/restore: could not load selected snapshot")
+        return True
+
+    scope = "preview"
+    if len(tokens) >= 3:
+        scope = tokens[2].lower()
+        if scope not in ("full", "files", "conversation", "preview"):
+            emit_error(
+                "/restore: scope must be one of full, files, conversation, preview"
+            )
+            return True
+
+    if scope == "preview":
+        _preview_checkpoint(selected, snapshot, project_root)
+        return True
+
+    # Confirmation dialog before destructive restore
+    _run_restore(scope, selected, snapshot, project_root)
+    return True
+
+
+def _preview_checkpoint(
+    selected: dict[str, Any], snapshot: dict[str, Any], project_root: Path
+) -> None:
+    from code_muse.messaging import emit_info
+
+    commit_hash = _get_commit_hash_for_snapshot(selected, project_root)
+    lines = [":mag: Checkpoint preview:"]
+    lines.append(f"  Timestamp: {selected.get('timestamp', '')}")
+    lines.append(f"  Tool:      {selected.get('tool_name', '')}")
+    lines.append(f"  Messages:  {len(snapshot.get('messages', []))} at checkpoint")
+
+    if commit_hash:
+        try:
+            diff_result = subprocess.run(
+                [
+                    "git",
+                    "-C",
+                    str(project_root),
+                    "diff",
+                    f"{commit_hash}..HEAD",
+                    "--stat",
+                ],
+                capture_output=True,
+                text=True,
+            )
+            if diff_result.returncode == 0 and diff_result.stdout.strip():
+                lines.append("  Diff since checkpoint:")
+                for line in diff_result.stdout.strip().splitlines():
+                    lines.append(f"    {line}")
+            else:
+                lines.append("  No file changes since checkpoint.")
+        except Exception as exc:
+            lines.append(f"  Could not compute diff: {exc}")
+    else:
+        lines.append("  Commit hash not available.")
+
+    lines.append("")
+    lines.append("Usage: /restore <index> [full|files|conversation]")
+    emit_info("\n".join(lines))
+
+
+def _run_restore(
+    scope: str,
+    selected: dict[str, Any],
+    snapshot: dict[str, Any],
+    project_root: Path,
+) -> None:
+    import asyncio
+
+    from code_muse.messaging import emit_error, emit_info, emit_success
+
+    async def _do_restore() -> bool:
+        # TODO: PEP 734 async bridge — _get_commit_hash_for_snapshot uses sync subprocess
+        commit_hash = await asyncio.to_thread(
+            _get_commit_hash_for_snapshot, selected, project_root
+        )
+        restore_parts: list[str] = []
+        if scope in ("full", "files"):
+            restore_parts.append("files")
+        if scope in ("full", "conversation"):
+            restore_parts.append("conversation")
+
+        preview_text = f"Revert {', '.join(restore_parts)} to checkpoint [{selected.get('timestamp', '')}]?"
+
+        confirmed, feedback = await get_user_approval_async(
+            title="Restore Checkpoint",
+            content=preview_text,
+            preview=None,
+            border_style="dim white",
+            agent_name="Muse",
+        )
+        if not confirmed:
+            emit_info("Restore cancelled.")
+            return False
+
+        if scope in ("full", "files") and commit_hash:
+            try:
+                checkout_proc = await asyncio.create_subprocess_exec(
+                    "git",
+                    "-C",
+                    str(project_root),
+                    "checkout",
+                    commit_hash,
+                    "--",
+                    ".",
+                    stdout=asyncio.subprocess.PIPE,
+                    stderr=asyncio.subprocess.PIPE,
+                )
+                checkout_stdout, checkout_stderr = await checkout_proc.communicate()
+                if checkout_proc.returncode != 0:
+                    emit_error(f"git checkout failed: {checkout_stderr.decode()}")
+                    return False
+                emit_success(":white_check_mark: Files restored.")
+            except Exception as exc:
+                emit_error(f"File restore failed: {exc}")
+                return False
+
+        if scope in ("full", "conversation"):
+            try:
+                from code_muse.agents import get_current_agent
+
+                agent = get_current_agent()
+                messages = snapshot.get("messages", [])
+                if messages:
+                    agent.set_message_history(messages)
+                    emit_success(":white_check_mark: Conversation restored.")
+                else:
+                    emit_info("No conversation state in snapshot.")
+            except Exception as exc:
+                emit_error(f"Conversation restore failed: {exc}")
+                return False
+
+        return True
+
+    try:
+        asyncio.get_running_loop()
+        asyncio.create_task(_do_restore())
+    except RuntimeError:
+        asyncio.run(_do_restore())
+
+
+def _get_commit_hash_for_snapshot(
+    selected: dict[str, Any], project_root: Path
+) -> str | None:
+    """Find the commit hash that matches this snapshot's timestamp/tool."""
+    try:
+        timestamp = selected.get("timestamp", "")
+        tool_name = selected.get("tool_name", "")
+        log_result = subprocess.run(
+            [
+                "git",
+                "-C",
+                str(project_root),
+                "log",
+                "--oneline",
+                "--format=%H %s",
+            ],
+            capture_output=True,
+            text=True,
+        )
+        if log_result.returncode != 0:
+            return None
+        for line in log_result.stdout.strip().splitlines():
+            parts = line.split(None, 1)
+            if len(parts) < 2:
+                continue
+            commit_hash, message = parts
+            if f"checkpoint: {tool_name}" in message and timestamp in message:
+                return commit_hash
+        # Fallback: return HEAD if no exact match
+        head_result = subprocess.run(
+            ["git", "-C", str(project_root), "rev-parse", "HEAD"],
+            capture_output=True,
+            text=True,
+        )
+        if head_result.returncode == 0:
+            return head_result.stdout.strip()
+    except Exception as exc:
+        logger.warning(f"Could not resolve commit hash: {exc}")
+    return None
+
+
+def _hash_project_root(project_root: str) -> str:
+    import hashlib
+
+    return hashlib.sha256(project_root.encode()).hexdigest()
+
+
+def _handle_custom_command(command: str, name: str) -> bool | None:
+    if name != "restore":
+        return None
+    return _handle_restore_command(command)

code_muse/plugins/checkpointing/rewind_shortcut.py

@@ -0,0 +1,88 @@
+"""Esc×2 rewind keyboard shortcut."""
+
+import logging
+import sys
+import threading
+import time
+from collections.abc import Callable
+
+logger = logging.getLogger(__name__)
+
+
+class DoublePressDetector:
+    """Detects double-press of a key within a time window."""
+
+    def __init__(self, window_ms: int = 500) -> None:
+        self._window_ms = window_ms
+        self._last_press: float = 0.0
+
+    def press(self) -> bool:
+        now = time.monotonic()
+        delta_ms = (now - self._last_press) * 1000
+        self._last_press = now
+        return 0 < delta_ms < self._window_ms
+
+
+class RewindKeyListener:
+    """Daemon thread that listens for raw ESC keycodes and triggers rewind on double-press."""
+
+    def __init__(self, on_double_esc: Callable[[], None]) -> None:
+        self._detector = DoublePressDetector()
+        self._on_double_esc = on_double_esc
+        self._stop = threading.Event()
+        self._thread = threading.Thread(target=self._listen, daemon=True)
+
+    def start(self) -> None:
+        self._thread.start()
+        logger.info("RewindKeyListener started")
+
+    def stop(self) -> None:
+        self._stop.set()
+        logger.info("RewindKeyListener stopped")
+
+    def _listen(self) -> None:
+        if sys.platform == "win32":
+            self._listen_windows()
+        else:
+            self._listen_posix()
+
+    def _listen_posix(self) -> None:
+        try:
+            import select
+            import termios
+            import tty
+
+            fd = sys.stdin.fileno()
+            old_settings = termios.tcgetattr(fd)
+            tty.setcbreak(fd)
+            try:
+                while not self._stop.is_set():
+                    ready, _, _ = select.select([sys.stdin], [], [], 0.1)
+                    if ready:
+                        char = sys.stdin.read(1)
+                        if char == "\x1b":
+                            if self._detector.press():
+                                try:
+                                    self._on_double_esc()
+                                except Exception as exc:
+                                    logger.error(f"Double-esc handler failed: {exc}")
+            finally:
+                termios.tcsetattr(fd, termios.TCSADRAIN, old_settings)
+        except Exception as exc:
+            logger.debug(f"POSIX rewind listener could not start: {exc}")
+
+    def _listen_windows(self) -> None:
+        try:
+            import msvcrt
+
+            while not self._stop.is_set():
+                if msvcrt.kbhit():
+                    char = msvcrt.getch()
+                    if char == b"\x1b" and self._detector.press():
+                        try:
+                            self._on_double_esc()
+                        except Exception as exc:
+                            logger.error(f"Double-esc handler failed: {exc}")
+                time.sleep(0.05)
+        except Exception as exc:
+            logger.debug(f"Windows rewind listener could not start: {exc}")

code_muse/plugins/checkpointing/shadow_git.py

@@ -0,0 +1,90 @@
+"""Shadow git repository for checkpointing file snapshots."""
+
+import hashlib
+import logging
+import subprocess
+from datetime import UTC, datetime
+from pathlib import Path
+
+logger = logging.getLogger(__name__)
+
+
+class ShadowGit:
+    """Manages a bare shadow git repository for project checkpointing."""
+
+    def __init__(self, project_root: str) -> None:
+        self.project_root = Path(project_root).resolve()
+        project_hash = hashlib.sha256(str(self.project_root).encode()).hexdigest()
+        self.repo_path = Path.home() / ".muse" / "history" / project_hash
+        self._init_repo()
+
+    def _init_repo(self) -> None:
+        if not self.repo_path.exists():
+            try:
+                self.repo_path.mkdir(parents=True, exist_ok=True)
+                subprocess.run(
+                    ["git", "init", "--bare", str(self.repo_path)],
+                    capture_output=True,
+                    text=True,
+                    check=True,
+                )
+                logger.info(f"Initialized shadow git repo at {self.repo_path}")
+            except Exception as exc:
+                logger.warning(f"Failed to init shadow git repo: {exc}")
+
+    def create_checkpoint(
+        self, tool_name: str, affected_files: list[str] | None = None
+    ) -> str | None:
+        """Create a git checkpoint and return the commit hash."""
+        timestamp = datetime.now(UTC).isoformat()
+        message = f"checkpoint: {tool_name} {timestamp}"
+        if affected_files:
+            files_str = ", ".join(affected_files)
+            message += f" ({files_str})"
+
+        try:
+            # Stage everything in the project root
+            add_result = subprocess.run(
+                ["git", "-C", str(self.project_root), "add", "-A"],
+                capture_output=True,
+                text=True,
+            )
+            if add_result.returncode != 0:
+                logger.warning(f"git add failed: {add_result.stderr}")
+                return None
+
+            # Commit
+            commit_result = subprocess.run(
+                [
+                    "git",
+                    "-C",
+                    str(self.project_root),
+                    "commit",
+                    "--allow-empty",
+                    "-m",
+                    message,
+                ],
+                capture_output=True,
+                text=True,
+            )
+            if commit_result.returncode != 0:
+                # No changes to commit is acceptable
+                if "nothing to commit" in commit_result.stdout.lower():
+                    pass
+                else:
+                    logger.warning(f"git commit failed: {commit_result.stderr}")
+                    return None
+
+            # Get commit hash
+            hash_result = subprocess.run(
+                ["git", "-C", str(self.project_root), "rev-parse", "HEAD"],
+                capture_output=True,
+                text=True,
+                check=True,
+            )
+            commit_hash = hash_result.stdout.strip()
+            logger.info(f"Created checkpoint {commit_hash} for {tool_name}")
+            return commit_hash
+        except Exception as exc:
+            logger.warning(f"Checkpoint creation failed: {exc}")
+            return None

code_muse/plugins/claude_code_hooks/__init__.py

@@ -0,0 +1 @@
+"""Claude Code hooks plugin."""