llmcode-cli 1.0.0__py3-none-any.whl

llm_code/runtime/checkpoint.py

@@ -0,0 +1,70 @@
+"""Git-based checkpoint manager for undoable tool operations."""
+
+from __future__ import annotations
+
+import subprocess
+from dataclasses import dataclass
+from datetime import datetime, timezone
+from pathlib import Path
+
+
+@dataclass(frozen=True)
+class Checkpoint:
+    id: str               # incrementing "cp-001", "cp-002", …
+    timestamp: str        # ISO format
+    tool_name: str
+    tool_args_summary: str  # short display string (first 80 chars of str(tool_args))
+    git_sha: str
+
+
+class CheckpointManager:
+    def __init__(self, cwd: Path) -> None:
+        self._cwd = cwd
+        self._stack: list[Checkpoint] = []
+        self._counter = 0
+
+    def create(self, tool_name: str, tool_args: dict) -> Checkpoint:
+        """Commit the current working-tree state and push a Checkpoint onto the stack."""
+        subprocess.run(["git", "add", "-A"], cwd=self._cwd, capture_output=True)
+        subprocess.run(
+            ["git", "commit", "--allow-empty", "-m", f"llm-code checkpoint: before {tool_name}"],
+            cwd=self._cwd,
+            capture_output=True,
+        )
+
+        sha_result = subprocess.run(
+            ["git", "rev-parse", "HEAD"],
+            capture_output=True,
+            text=True,
+            cwd=self._cwd,
+        )
+        git_sha = sha_result.stdout.strip()
+
+        self._counter += 1
+        cp = Checkpoint(
+            id=f"cp-{self._counter:03d}",
+            timestamp=datetime.now(tz=timezone.utc).isoformat(),
+            tool_name=tool_name,
+            tool_args_summary=str(tool_args)[:80],
+            git_sha=git_sha,
+        )
+        self._stack.append(cp)
+        return cp
+
+    def undo(self) -> Checkpoint | None:
+        """Pop the last checkpoint and hard-reset the repo to that SHA."""
+        if not self._stack:
+            return None
+        cp = self._stack.pop()
+        subprocess.run(
+            ["git", "reset", "--hard", cp.git_sha],
+            cwd=self._cwd,
+            capture_output=True,
+        )
+        return cp
+
+    def list_checkpoints(self) -> list[Checkpoint]:
+        return list(self._stack)
+
+    def can_undo(self) -> bool:
+        return len(self._stack) > 0

llm_code/runtime/checkpoint_recovery.py

@@ -0,0 +1,142 @@
+"""Session checkpoint recovery: save/load full session state for crash recovery."""
+from __future__ import annotations
+
+import asyncio
+import json
+import logging
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from llm_code.runtime.session import Session
+
+logger = logging.getLogger(__name__)
+
+_CHECKPOINTS_DIR_NAME = "checkpoints"
+
+
+class CheckpointRecovery:
+    """Persist and restore full session state for crash recovery.
+
+    Checkpoints are stored as JSON files under
+    ``~/.llm-code/checkpoints/<session_id>.json`` (or a custom *checkpoints_dir*).
+    """
+
+    def __init__(self, checkpoints_dir: Path) -> None:
+        self._dir = checkpoints_dir
+        self._dir.mkdir(parents=True, exist_ok=True)
+        self._auto_save_task: asyncio.Task | None = None
+
+    # ------------------------------------------------------------------
+    # Core persistence
+    # ------------------------------------------------------------------
+
+    def save_checkpoint(self, session: "Session") -> Path:
+        """Serialize *session* to disk and return the checkpoint file path."""
+
+        data = session.to_dict()
+        data["checkpoint_saved_at"] = datetime.now(timezone.utc).isoformat()
+
+        path = self._dir / f"{session.id}.json"
+        path.write_text(json.dumps(data, indent=2), encoding="utf-8")
+        logger.debug("Checkpoint saved: %s", path)
+        return path
+
+    def load_checkpoint(self, session_id: str) -> "Session | None":
+        """Deserialize a checkpoint by *session_id*, or return None."""
+        from llm_code.runtime.session import Session  # local import to avoid cycles
+
+        path = self._dir / f"{session_id}.json"
+        if not path.exists():
+            return None
+        try:
+            data = json.loads(path.read_text(encoding="utf-8"))
+            # checkpoint_saved_at is extra metadata; Session.from_dict ignores unknown keys
+            # but we strip it to keep from_dict clean
+            data.pop("checkpoint_saved_at", None)
+            return Session.from_dict(data)
+        except (json.JSONDecodeError, KeyError, TypeError) as exc:
+            logger.warning("Failed to load checkpoint %s: %s", session_id, exc)
+            return None
+
+    def list_checkpoints(self) -> list[dict]:
+        """Return checkpoint descriptors sorted by modification time (newest first).
+
+        Each dict has: ``session_id``, ``saved_at``, ``message_count``,
+        ``project_path``, ``updated_at``.
+        """
+        results: list[dict] = []
+        for path in sorted(
+            self._dir.glob("*.json"),
+            key=lambda p: p.stat().st_mtime,
+            reverse=True,
+        ):
+            try:
+                data = json.loads(path.read_text(encoding="utf-8"))
+                results.append({
+                    "session_id": data.get("id", path.stem),
+                    "saved_at": data.get("checkpoint_saved_at", ""),
+                    "message_count": len(data.get("messages", [])),
+                    "project_path": data.get("project_path", ""),
+                    "updated_at": data.get("updated_at", ""),
+                })
+            except (json.JSONDecodeError, OSError):
+                continue
+        return results
+
+    def delete_checkpoint(self, session_id: str) -> bool:
+        """Delete a checkpoint file; returns True if it existed."""
+        path = self._dir / f"{session_id}.json"
+        if path.exists():
+            path.unlink()
+            return True
+        return False
+
+    # ------------------------------------------------------------------
+    # Auto-save background task
+    # ------------------------------------------------------------------
+
+    def start_auto_save(self, get_session_fn, interval: int = 60) -> None:
+        """Start a background asyncio task that saves a checkpoint every *interval* seconds.
+
+        *get_session_fn* is a zero-argument callable that returns the current
+        :class:`~llm_code.runtime.session.Session` (or None to skip).
+        """
+        if self._auto_save_task is not None and not self._auto_save_task.done():
+            return  # already running
+
+        async def _loop():
+            while True:
+                await asyncio.sleep(interval)
+                try:
+                    session = get_session_fn()
+                    if session is not None:
+                        self.save_checkpoint(session)
+                except Exception as exc:
+                    logger.debug("Auto-save checkpoint error: %s", exc)
+
+        try:
+            loop = asyncio.get_event_loop()
+        except RuntimeError:
+            return
+
+        self._auto_save_task = loop.create_task(_loop())
+        logger.debug("Checkpoint auto-save started (interval=%ds)", interval)
+
+    def stop_auto_save(self) -> None:
+        """Cancel the auto-save background task if running."""
+        if self._auto_save_task is not None and not self._auto_save_task.done():
+            self._auto_save_task.cancel()
+            self._auto_save_task = None
+
+    # ------------------------------------------------------------------
+    # Startup detection
+    # ------------------------------------------------------------------
+
+    def detect_last_checkpoint(self) -> "Session | None":
+        """Return the most recently modified checkpoint session, or None."""
+        entries = self.list_checkpoints()
+        if not entries:
+            return None
+        return self.load_checkpoint(entries[0]["session_id"])

llm_code/runtime/compaction.py

@@ -0,0 +1,35 @@
+"""Context compaction: trim old session messages when the context grows too large."""
+from __future__ import annotations
+
+import dataclasses
+
+from llm_code.api.types import Message, TextBlock
+from llm_code.runtime.session import Session
+
+
+def needs_compaction(session: Session, threshold: int = 80000) -> bool:
+    """Return True when the session's estimated token count exceeds *threshold*."""
+    return session.estimated_tokens() > threshold
+
+
+def compact_session(
+    session: Session,
+    keep_recent: int = 4,
+    summary: str = "",
+) -> Session:
+    """Return a compacted session keeping only the most recent *keep_recent* messages.
+
+    If the session has <= keep_recent messages, the original session is returned
+    unchanged.  Otherwise a single summary message is prepended to the last
+    *keep_recent* messages.
+    """
+    if len(session.messages) <= keep_recent:
+        return session
+
+    summary_msg = Message(
+        role="user",
+        content=(TextBlock(text=f"[Previous conversation summary]\n{summary}"),),
+    )
+    recent = session.messages[-keep_recent:]
+    new_messages = (summary_msg,) + recent
+    return dataclasses.replace(session, messages=new_messages)

llm_code/runtime/compressor.py

@@ -0,0 +1,415 @@
+"""ContextCompressor: 5-level progressive context compression."""
+from __future__ import annotations
+
+import dataclasses
+import logging
+from typing import TYPE_CHECKING
+
+from llm_code.api.types import Message, MessageRequest, TextBlock, ToolResultBlock, ToolUseBlock
+from llm_code.runtime.session import Session
+
+if TYPE_CHECKING:
+    from llm_code.api.provider import LLMProvider
+
+_log = logging.getLogger(__name__)
+
+_SUMMARIZE_SYSTEM_PROMPT = """\
+You are a context compression agent. Given conversation messages from a coding \
+session, produce a concise summary preserving:
+
+1. What files were read, created, or modified (exact paths)
+2. Key decisions made and their rationale
+3. Current state of the task (what's done, what's pending)
+4. Any errors encountered and how they were resolved
+
+Be factual. Use bullet points. Do not include code blocks unless critical.
+"""
+
+
+class ContextCompressor:
+    """Progressively compress a Session context through 5 escalating levels.
+
+    Level 1 — snip_compact: Truncate oversized ToolResultBlock content.
+    Level 2 — micro_compact: Remove stale read_file results (keep only latest per path).
+    Level 3 — context_collapse: Replace old tool_call+result pairs with one-line summaries.
+    Level 4 — auto_compact: Discard all old messages, keep a summary + recent tail.
+    Level 5 — llm_summarize: (async only) Replace Level 4 placeholder with LLM-generated summary.
+
+    Cache-aware: tracks which message indices have been sent to the API (cached).
+    Compression levels prefer removing non-cached messages first to preserve
+    API-side prompt cache hits.
+    """
+
+    def __init__(
+        self,
+        max_result_chars: int = 2000,
+        provider: "LLMProvider | None" = None,
+        summarize_model: str = "",
+        max_summary_tokens: int = 1000,
+    ) -> None:
+        self._max_result_chars = max_result_chars
+        self._cached_indices: set[int] = set()
+        self._provider = provider
+        self._summarize_model = summarize_model
+        self._max_summary_tokens = max_summary_tokens
+
+    # ------------------------------------------------------------------
+    # Cache tracking
+    # ------------------------------------------------------------------
+
+    def mark_as_cached(self, message_indices: set[int]) -> None:
+        """Mark which message indices have been sent to the API (cache hits)."""
+        self._cached_indices.update(message_indices)
+
+    def _is_cached(self, index: int) -> bool:
+        """Return True if the message at *index* has been sent to the API."""
+        return index in self._cached_indices
+
+    # ------------------------------------------------------------------
+    # Public entry point
+    # ------------------------------------------------------------------
+
+    def compress(self, session: Session, max_tokens: int) -> Session:
+        """Compress *session* until estimated_tokens() <= max_tokens.
+
+        Applies levels in order, stopping as soon as the budget is met.
+        If all 4 levels still cannot reach the budget, the Level-4 result
+        is returned (best-effort).
+
+        Resets cached indices after compression since message indices change.
+        """
+        if session.estimated_tokens() <= max_tokens:
+            return session
+        # Reset stale cache indices — message positions change after compression
+        self._cached_indices.clear()
+
+        session = self._snip_compact(session)
+        if session.estimated_tokens() <= max_tokens:
+            return session
+
+        session = self._micro_compact(session)
+        if session.estimated_tokens() <= max_tokens:
+            return session
+
+        session = self._context_collapse(session, keep_recent=6)
+        if session.estimated_tokens() <= max_tokens:
+            return session
+
+        session = self._auto_compact(session, keep_recent=4)
+        return session
+
+    async def compress_async(self, session: Session, max_tokens: int) -> Session:
+        """Async compress with optional Level 5 LLM summarization."""
+        result = self.compress(session, max_tokens)
+        if self._provider is not None:
+            result = await self._llm_summarize(result)
+        return result
+
+    # ------------------------------------------------------------------
+    # Level 5 (async only)
+    # ------------------------------------------------------------------
+
+    async def _llm_summarize(self, session: Session) -> Session:
+        """Replace Level 4 placeholder with LLM-generated summary."""
+        placeholder_idx = None
+        for i, msg in enumerate(session.messages):
+            for block in msg.content:
+                if isinstance(block, TextBlock) and "[Previous conversation summary]" in block.text:
+                    placeholder_idx = i
+                    break
+            if placeholder_idx is not None:
+                break
+
+        if placeholder_idx is None:
+            return session
+
+        # Build context from remaining messages
+        context_parts: list[str] = []
+        for i, msg in enumerate(session.messages):
+            if i == placeholder_idx:
+                continue
+            for block in msg.content:
+                if isinstance(block, TextBlock):
+                    context_parts.append(f"[{msg.role}] {block.text[:500]}")
+                elif isinstance(block, ToolUseBlock):
+                    context_parts.append(f"[tool_call] {block.name}({str(block.input)[:200]})")
+                elif isinstance(block, ToolResultBlock):
+                    context_parts.append(f"[tool_result] {block.content[:200]}")
+
+        if not context_parts:
+            return session
+
+        try:
+            request = MessageRequest(
+                model=self._summarize_model,
+                system=_SUMMARIZE_SYSTEM_PROMPT,
+                messages=(
+                    Message(
+                        role="user",
+                        content=(TextBlock(text="Summarize this conversation:\n\n" + "\n".join(context_parts)),),
+                    ),
+                ),
+                max_tokens=self._max_summary_tokens,
+            )
+            response = await self._provider.complete(request)
+            summary_text = response.content if isinstance(response.content, str) else str(response.content)
+        except Exception:
+            _log.warning("Level 5 LLM summarization failed, keeping placeholder", exc_info=True)
+            return session
+
+        summary_msg = Message(
+            role="user",
+            content=(TextBlock(text=f"[Conversation summary]\n{summary_text}"),),
+        )
+        messages = list(session.messages)
+        messages[placeholder_idx] = summary_msg
+        return dataclasses.replace(session, messages=tuple(messages))
+
+    # ------------------------------------------------------------------
+    # Level 1
+    # ------------------------------------------------------------------
+
+    def _snip_compact(self, session: Session) -> Session:
+        """Truncate each ToolResultBlock's content to *max_result_chars*.
+
+        Cache-aware: truncate non-cached messages first.  If no non-cached
+        messages are over-budget, fall through to truncating cached ones too.
+        """
+        new_messages: list[Message] = []
+        changed = False
+
+        # First pass: truncate only non-cached oversized results
+        for idx, msg in enumerate(session.messages):
+            new_blocks: list = []
+            msg_changed = False
+            for block in msg.content:
+                if (
+                    isinstance(block, ToolResultBlock)
+                    and len(block.content) > self._max_result_chars
+                    and not self._is_cached(idx)
+                ):
+                    truncated = block.content[: self._max_result_chars]
+                    new_blocks.append(dataclasses.replace(block, content=truncated))
+                    msg_changed = True
+                else:
+                    new_blocks.append(block)
+            if msg_changed:
+                new_messages.append(dataclasses.replace(msg, content=tuple(new_blocks)))
+                changed = True
+            else:
+                new_messages.append(msg)
+
+        interim = dataclasses.replace(session, messages=tuple(new_messages)) if changed else session
+
+        # Second pass: also truncate cached oversized results (fallback)
+        final_messages: list[Message] = []
+        second_changed = False
+        for idx, msg in enumerate(interim.messages):
+            new_blocks = []
+            msg_changed = False
+            for block in msg.content:
+                if isinstance(block, ToolResultBlock) and len(block.content) > self._max_result_chars:
+                    truncated = block.content[: self._max_result_chars]
+                    new_blocks.append(dataclasses.replace(block, content=truncated))
+                    msg_changed = True
+                else:
+                    new_blocks.append(block)
+            if msg_changed:
+                final_messages.append(dataclasses.replace(msg, content=tuple(new_blocks)))
+                second_changed = True
+            else:
+                final_messages.append(msg)
+
+        if not changed and not second_changed:
+            return session
+        if second_changed:
+            return dataclasses.replace(session, messages=tuple(final_messages))
+        return interim
+
+    # ------------------------------------------------------------------
+    # Level 2
+    # ------------------------------------------------------------------
+
+    def _micro_compact(self, session: Session) -> Session:
+        """For the same file read multiple times, keep only the latest read_file result.
+
+        Strategy: build a mapping from tool_use_id → file path for all read_file
+        ToolUseBlocks.  Then, for each file path, collect the tool_use_ids in order
+        and mark all but the last one for removal.  Finally rebuild messages, dropping
+        ToolResultBlocks whose tool_use_id is marked.
+
+        Cache-aware: prefer removing non-cached stale reads first.  If no
+        non-cached duplicates exist, fall back to removing cached ones.
+        """
+        # Pass 1: map tool_use_id → (path, message_index) for read_file calls
+        id_to_path: dict[str, str] = {}
+        id_to_msg_index: dict[str, int] = {}
+        for msg_idx, msg in enumerate(session.messages):
+            for block in msg.content:
+                if isinstance(block, ToolUseBlock) and block.name == "read_file":
+                    path = block.input.get("path", "")
+                    if path:
+                        id_to_path[block.id] = path
+                        id_to_msg_index[block.id] = msg_idx
+
+        # For each path, keep only the last tool_use_id
+        path_to_ids: dict[str, list[str]] = {}
+        for tid, path in id_to_path.items():
+            path_to_ids.setdefault(path, []).append(tid)
+
+        stale_ids: set[str] = set()
+        for path, ids in path_to_ids.items():
+            if len(ids) > 1:
+                candidate_stale = ids[:-1]  # all but the last
+                # Prefer removing non-cached first; only include cached if necessary
+                non_cached_stale = [t for t in candidate_stale if not self._is_cached(id_to_msg_index.get(t, -1))]
+                if non_cached_stale:
+                    stale_ids.update(non_cached_stale)
+                else:
+                    # Fallback: remove cached stale reads when no non-cached option exists
+                    stale_ids.update(candidate_stale)
+
+        if not stale_ids:
+            return session
+
+        # Pass 2: rebuild messages, dropping stale ToolResultBlocks (and their paired ToolUseBlocks)
+        new_messages: list[Message] = []
+        for msg in session.messages:
+            new_blocks = []
+            for block in msg.content:
+                if isinstance(block, ToolResultBlock) and block.tool_use_id in stale_ids:
+                    continue  # drop stale result
+                if isinstance(block, ToolUseBlock) and block.id in stale_ids:
+                    continue  # drop stale use block too
+                new_blocks.append(block)
+            if new_blocks:
+                new_messages.append(dataclasses.replace(msg, content=tuple(new_blocks)))
+            # If a message becomes empty (all blocks dropped), skip it entirely
+
+        return dataclasses.replace(session, messages=tuple(new_messages))
+
+    # ------------------------------------------------------------------
+    # Level 3
+    # ------------------------------------------------------------------
+
+    def _context_collapse(self, session: Session, keep_recent: int = 6) -> Session:
+        """Replace old tool_call+result pairs with one-line summary text.
+
+        Messages in the *keep_recent* tail are kept intact.  Earlier messages
+        are converted: ToolUseBlock/ToolResultBlock → summary TextBlock.
+
+        Cache-aware: collapse non-cached messages first.  Cached messages in
+        the old section are passed through as-is; only when there are no
+        non-cached messages to collapse do we fall back to collapsing cached ones.
+        """
+        if len(session.messages) <= keep_recent:
+            return session
+
+        old_messages = session.messages[:-keep_recent]
+        recent_messages = session.messages[-keep_recent:]
+
+        # Separate old messages into non-cached (collapse) and cached (preserve when possible)
+        non_cached_old: list[tuple[int, Message]] = []
+        cached_old: list[tuple[int, Message]] = []
+        for rel_idx, msg in enumerate(old_messages):
+            abs_idx = rel_idx  # old_messages starts at index 0
+            if self._is_cached(abs_idx):
+                cached_old.append((abs_idx, msg))
+            else:
+                non_cached_old.append((abs_idx, msg))
+
+        # Collapse non-cached old messages into summary lines
+        summary_lines: list[str] = []
+        for _idx, msg in non_cached_old:
+            for block in msg.content:
+                if isinstance(block, ToolUseBlock):
+                    args_summary = ", ".join(
+                        f"{k}={v!r}" for k, v in list(block.input.items())[:3]
+                    )
+                    summary_lines.append(f"Used {block.name}({args_summary})")
+                elif isinstance(block, ToolResultBlock):
+                    pass
+                elif isinstance(block, TextBlock) and block.text.strip():
+                    excerpt = block.text[:80].replace("\n", " ")
+                    summary_lines.append(f"[msg] {excerpt}")
+
+        # If non-cached messages produced summary lines, keep cached old messages intact
+        if summary_lines or cached_old:
+            # Build the new old section: cached messages preserved + summary of non-cached
+            preserved_cached = tuple(msg for _idx, msg in cached_old)
+            if summary_lines:
+                summary_text = "\n".join(summary_lines)
+                summary_msg = Message(
+                    role="user",
+                    content=(TextBlock(text=f"[Context summary]\n{summary_text}"),),
+                )
+                new_old_section = preserved_cached + (summary_msg,)
+            else:
+                new_old_section = preserved_cached
+
+            if not new_old_section:
+                return dataclasses.replace(session, messages=recent_messages)
+            return dataclasses.replace(
+                session,
+                messages=new_old_section + recent_messages,
+            )
+
+        # Fallback: collapse all old messages (including cached) — no non-cached existed
+        all_summary_lines: list[str] = []
+        for msg in old_messages:
+            for block in msg.content:
+                if isinstance(block, ToolUseBlock):
+                    args_summary = ", ".join(
+                        f"{k}={v!r}" for k, v in list(block.input.items())[:3]
+                    )
+                    all_summary_lines.append(f"Used {block.name}({args_summary})")
+                elif isinstance(block, ToolResultBlock):
+                    pass
+                elif isinstance(block, TextBlock) and block.text.strip():
+                    excerpt = block.text[:80].replace("\n", " ")
+                    all_summary_lines.append(f"[msg] {excerpt}")
+
+        if not all_summary_lines:
+            return dataclasses.replace(session, messages=recent_messages)
+
+        summary_text = "\n".join(all_summary_lines)
+        summary_msg = Message(
+            role="user",
+            content=(TextBlock(text=f"[Context summary]\n{summary_text}"),),
+        )
+        return dataclasses.replace(
+            session,
+            messages=(summary_msg,) + recent_messages,
+        )
+
+    # ------------------------------------------------------------------
+    # Level 4
+    # ------------------------------------------------------------------
+
+    def _auto_compact(self, session: Session, keep_recent: int = 4) -> Session:
+        """Replace all old messages with a single summary placeholder + keep tail.
+
+        This mirrors the logic in :func:`llm_code.runtime.compaction.compact_session`.
+
+        Cache-aware: cached messages from the old section are preserved before
+        the summary placeholder so they remain available for API cache hits.
+        """
+        if len(session.messages) <= keep_recent:
+            return session
+
+        old_messages = session.messages[:-keep_recent]
+        recent = session.messages[-keep_recent:]
+
+        # Preserve cached messages from the old section
+        preserved_cached = tuple(
+            msg for idx, msg in enumerate(old_messages) if self._is_cached(idx)
+        )
+
+        summary_msg = Message(
+            role="user",
+            content=(TextBlock(text="[Previous conversation summary]\n"),),
+        )
+        return dataclasses.replace(
+            session,
+            messages=preserved_cached + (summary_msg,) + recent,
+        )