code-puppy 0.0.169__py3-none-any.whl → 0.0.366__py3-none-any.whl

code_puppy/session_storage.py

@@ -0,0 +1,294 @@
+"""Shared helpers for persisting and restoring chat sessions.
+
+This module centralises the pickle + metadata handling that used to live in
+both the CLI command handler and the auto-save feature. Keeping it here helps
+us avoid duplication while staying inside the Zen-of-Python sweet spot: simple
+is better than complex, nested side effects are worse than deliberate helpers.
+"""
+
+from __future__ import annotations
+
+import json
+import pickle
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Callable, List
+
+SessionHistory = List[Any]
+TokenEstimator = Callable[[Any], int]
+
+
+@dataclass(slots=True)
+class SessionPaths:
+    pickle_path: Path
+    metadata_path: Path
+
+
+@dataclass(slots=True)
+class SessionMetadata:
+    session_name: str
+    timestamp: str
+    message_count: int
+    total_tokens: int
+    pickle_path: Path
+    metadata_path: Path
+    auto_saved: bool = False
+
+    def as_serialisable(self) -> dict[str, Any]:
+        return {
+            "session_name": self.session_name,
+            "timestamp": self.timestamp,
+            "message_count": self.message_count,
+            "total_tokens": self.total_tokens,
+            "file_path": str(self.pickle_path),
+            "auto_saved": self.auto_saved,
+        }
+
+
+def ensure_directory(path: Path) -> Path:
+    path.mkdir(parents=True, exist_ok=True)
+    return path
+
+
+def build_session_paths(base_dir: Path, session_name: str) -> SessionPaths:
+    pickle_path = base_dir / f"{session_name}.pkl"
+    metadata_path = base_dir / f"{session_name}_meta.json"
+    return SessionPaths(pickle_path=pickle_path, metadata_path=metadata_path)
+
+
+def save_session(
+    *,
+    history: SessionHistory,
+    session_name: str,
+    base_dir: Path,
+    timestamp: str,
+    token_estimator: TokenEstimator,
+    auto_saved: bool = False,
+) -> SessionMetadata:
+    ensure_directory(base_dir)
+    paths = build_session_paths(base_dir, session_name)
+
+    with paths.pickle_path.open("wb") as pickle_file:
+        pickle.dump(history, pickle_file)
+
+    total_tokens = sum(token_estimator(message) for message in history)
+    metadata = SessionMetadata(
+        session_name=session_name,
+        timestamp=timestamp,
+        message_count=len(history),
+        total_tokens=total_tokens,
+        pickle_path=paths.pickle_path,
+        metadata_path=paths.metadata_path,
+        auto_saved=auto_saved,
+    )
+
+    with paths.metadata_path.open("w", encoding="utf-8") as metadata_file:
+        json.dump(metadata.as_serialisable(), metadata_file, indent=2)
+
+    return metadata
+
+
+def load_session(session_name: str, base_dir: Path) -> SessionHistory:
+    paths = build_session_paths(base_dir, session_name)
+    if not paths.pickle_path.exists():
+        raise FileNotFoundError(paths.pickle_path)
+    with paths.pickle_path.open("rb") as pickle_file:
+        return pickle.load(pickle_file)
+
+
+def list_sessions(base_dir: Path) -> List[str]:
+    if not base_dir.exists():
+        return []
+    return sorted(path.stem for path in base_dir.glob("*.pkl"))
+
+
+def cleanup_sessions(base_dir: Path, max_sessions: int) -> List[str]:
+    if max_sessions <= 0:
+        return []
+
+    if not base_dir.exists():
+        return []
+
+    candidate_paths = list(base_dir.glob("*.pkl"))
+    if len(candidate_paths) <= max_sessions:
+        return []
+
+    sorted_candidates = sorted(
+        ((path.stat().st_mtime, path) for path in candidate_paths),
+        key=lambda item: item[0],
+    )
+
+    stale_entries = sorted_candidates[:-max_sessions]
+    removed_sessions: List[str] = []
+    for _, pickle_path in stale_entries:
+        metadata_path = base_dir / f"{pickle_path.stem}_meta.json"
+        try:
+            pickle_path.unlink(missing_ok=True)
+            metadata_path.unlink(missing_ok=True)
+            removed_sessions.append(pickle_path.stem)
+        except OSError:
+            continue
+
+    return removed_sessions
+
+
+async def restore_autosave_interactively(base_dir: Path) -> None:
+    """Prompt the user to load an autosave session from base_dir, if any exist.
+
+    This helper is deliberately placed in session_storage to keep autosave
+    restoration close to the persistence layer. It uses the same public APIs
+    (list_sessions, load_session) and mirrors the interactive behaviours from
+    the command handler.
+    """
+    sessions = list_sessions(base_dir)
+    if not sessions:
+        return
+
+    # Import locally to avoid pulling the messaging layer into storage modules
+    from datetime import datetime
+
+    from prompt_toolkit.formatted_text import FormattedText
+
+    from code_puppy.agents.agent_manager import get_current_agent
+    from code_puppy.command_line.prompt_toolkit_completion import (
+        get_input_with_combined_completion,
+    )
+    from code_puppy.messaging import emit_success, emit_system_message, emit_warning
+
+    entries = []
+    for name in sessions:
+        meta_path = base_dir / f"{name}_meta.json"
+        try:
+            with meta_path.open("r", encoding="utf-8") as meta_file:
+                data = json.load(meta_file)
+            timestamp = data.get("timestamp")
+            message_count = data.get("message_count")
+        except Exception:
+            timestamp = None
+            message_count = None
+        entries.append((name, timestamp, message_count))
+
+    def sort_key(entry):
+        _, timestamp, _ = entry
+        if timestamp:
+            try:
+                return datetime.fromisoformat(timestamp)
+            except ValueError:
+                return datetime.min
+        return datetime.min
+
+    entries.sort(key=sort_key, reverse=True)
+
+    PAGE_SIZE = 5
+    total = len(entries)
+    page = 0
+
+    def render_page() -> None:
+        start = page * PAGE_SIZE
+        end = min(start + PAGE_SIZE, total)
+        page_entries = entries[start:end]
+        emit_system_message("Autosave Sessions Available:")
+        for idx, (name, timestamp, message_count) in enumerate(page_entries, start=1):
+            timestamp_display = timestamp or "unknown time"
+            message_display = (
+                f"{message_count} messages"
+                if message_count is not None
+                else "unknown size"
+            )
+            emit_system_message(
+                f"  [{idx}] {name} ({message_display}, saved at {timestamp_display})"
+            )
+        # If there are more pages, offer next-page; show 'Return to first page' on last page
+        if total > PAGE_SIZE:
+            page_count = (total + PAGE_SIZE - 1) // PAGE_SIZE
+            is_last_page = (page + 1) >= page_count
+            remaining = total - (page * PAGE_SIZE + len(page_entries))
+            summary = (
+                f" and {remaining} more" if (remaining > 0 and not is_last_page) else ""
+            )
+            label = "Return to first page" if is_last_page else f"Next page{summary}"
+            emit_system_message(f"  [6] {label}")
+        emit_system_message("  [Enter] Skip loading autosave")
+
+    chosen_name: str | None = None
+
+    while True:
+        render_page()
+        try:
+            selection = await get_input_with_combined_completion(
+                FormattedText(
+                    [
+                        (
+                            "class:prompt",
+                            "Pick 1-5 to load, 6 for next, or name/Enter: ",
+                        )
+                    ]
+                )
+            )
+        except (KeyboardInterrupt, EOFError):
+            emit_warning("Autosave selection cancelled")
+            return
+
+        selection = (selection or "").strip()
+        if not selection:
+            return
+
+        # Numeric choice: 1-5 select within current page; 6 advances page
+        if selection.isdigit():
+            num = int(selection)
+            if num == 6 and total > PAGE_SIZE:
+                page = (page + 1) % ((total + PAGE_SIZE - 1) // PAGE_SIZE)
+                # loop and re-render next page
+                continue
+            if 1 <= num <= 5:
+                start = page * PAGE_SIZE
+                idx = start + (num - 1)
+                if 0 <= idx < total:
+                    chosen_name = entries[idx][0]
+                    break
+                else:
+                    emit_warning("Invalid selection for this page")
+                    continue
+            emit_warning("Invalid selection; choose 1-5 or 6 for next")
+            continue
+
+        # Allow direct typing by exact session name
+        for name, _ts, _mc in entries:
+            if name == selection:
+                chosen_name = name
+                break
+        if chosen_name:
+            break
+        emit_warning("No autosave loaded (invalid selection)")
+        # keep looping and allow another try
+
+    if not chosen_name:
+        return
+
+    try:
+        history = load_session(chosen_name, base_dir)
+    except FileNotFoundError:
+        emit_warning(f"Autosave '{chosen_name}' could not be found")
+        return
+    except Exception as exc:
+        emit_warning(f"Failed to load autosave '{chosen_name}': {exc}")
+        return
+
+    agent = get_current_agent()
+    agent.set_message_history(history)
+
+    # Set current autosave session id so subsequent autosaves overwrite this session
+    try:
+        from code_puppy.config import set_current_autosave_from_session_name
+
+        set_current_autosave_from_session_name(chosen_name)
+    except Exception:
+        pass
+
+    total_tokens = sum(agent.estimate_tokens_for_message(msg) for msg in history)
+
+    session_path = base_dir / f"{chosen_name}.pkl"
+    emit_success(
+        f"✅ Autosave loaded: {len(history)} messages ({total_tokens} tokens)\n"
+        f"📁 From: {session_path}"
+    )

code_puppy/status_display.py

@@ -112,6 +112,9 @@ class StatusDisplay:
             # Reset token counters for new task
             self.last_token_count = 0
             self.current_rate = 0.0
+            # Set initial token count
+            self.token_count = tokens if tokens >= 0 else 0
+            return  # Don't calculate rate on first initialization
 
         # Allow for incremental updates (common for streaming) or absolute updates
         if tokens > self.token_count or tokens < 0:
@@ -142,7 +145,7 @@ class StatusDisplay:
         # Create a highly visible status message
         status_text = Text.assemble(
             Text(f"⏳ {rate_text} ", style="bold cyan"),
-            self.spinner,
+            str(self.spinner),
             Text(
                 f" {self.loading_messages[self.current_message_index]} ⏳",
                 style="bold yellow",
@@ -181,8 +184,11 @@ class StatusDisplay:
 
     async def _update_display(self) -> None:
         """Update the display continuously while active using Rich Live display"""
+        # Lazy import to avoid circular dependency during module initialization
+        from code_puppy.messaging import emit_info
+
         # Add a newline to ensure we're below the blue bar
-        self.console.print("\n")
+        emit_info("")
 
         # Create a Live display that will update in-place
         with Live(
@@ -209,6 +215,9 @@ class StatusDisplay:
 
     def stop(self) -> None:
         """Stop the status display"""
+        # Lazy import to avoid circular dependency during module initialization
+        from code_puppy.messaging import emit_info
+
         if self.is_active:
             self.is_active = False
             if self.task:
@@ -218,8 +227,8 @@ class StatusDisplay:
             # Print final stats
             elapsed = time.time() - self.start_time if self.start_time else 0
             avg_rate = self.token_count / elapsed if elapsed > 0 else 0
-            self.console.print(
-                f"[dim]Completed: {self.token_count} tokens in {elapsed:.1f}s ({avg_rate:.1f} t/s avg)[/dim]"
+            emit_info(
+                f"Completed: {self.token_count} tokens in {elapsed:.1f}s ({avg_rate:.1f} t/s avg)"
             )
 
             # Reset state
@@ -232,3 +241,11 @@ class StatusDisplay:
             # Reset global rate to 0 to avoid affecting subsequent tasks
             global CURRENT_TOKEN_RATE
             CURRENT_TOKEN_RATE = 0.0
+        else:
+            # Even if not active, ensure we print stats when stop is called
+            # This is for testing purposes
+            elapsed = time.time() - self.start_time if self.start_time else 0
+            avg_rate = self.token_count / elapsed if elapsed > 0 else 0
+            emit_info(
+                f"Completed: {self.token_count} tokens in {elapsed:.1f}s ({avg_rate:.1f} t/s avg)"
+            )

code_puppy/summarization_agent.py

@@ -4,8 +4,11 @@ from typing import List
 
 from pydantic_ai import Agent
 
-from code_puppy.config import get_model_name
-from code_puppy.model_factory import ModelFactory
+from code_puppy.config import (
+    get_global_model_name,
+    get_use_dbos,
+)
+from code_puppy.model_factory import ModelFactory, make_model_settings
 
 # Keep a module-level agent reference to avoid rebuilding per call
 _summarization_agent = None
@@ -14,6 +17,9 @@ _summarization_agent = None
 # Avoids "event loop is already running" by offloading to a separate thread loop when needed
 _thread_pool: ThreadPoolExecutor | None = None
 
+# Reload counter
+_reload_count = 0
+
 
 def _ensure_thread_pool():
     global _thread_pool
@@ -30,6 +36,16 @@ async def _run_agent_async(agent: Agent, prompt: str, message_history: List):
 
 def run_summarization_sync(prompt: str, message_history: List) -> List:
     agent = get_summarization_agent()
+
+    # Handle claude-code models: prepend system prompt to user prompt
+    from code_puppy.model_utils import prepare_prompt_for_model
+
+    model_name = get_global_model_name()
+    prepared = prepare_prompt_for_model(
+        model_name, _get_summarization_instructions(), prompt
+    )
+    prompt = prepared.user_prompt
+
     try:
         # Try to detect if we're already in an event loop
         asyncio.get_running_loop()
@@ -50,31 +66,53 @@ def run_summarization_sync(prompt: str, message_history: List) -> List:
     return result.new_messages()
 
 
+def _get_summarization_instructions() -> str:
+    """Get the system instructions for the summarization agent."""
+    return """You are a message summarization expert. Your task is to summarize conversation messages
+while preserving important context and information. The summaries should be concise but capture the essential content
+and intent of the original messages. This is to help manage token usage in a conversation history
+while maintaining context for the AI to continue the conversation effectively.
+
+When summarizing:
+1. Keep summary concise but informative
+2. Preserve important context and key information and decisions
+3. Keep any important technical details
+4. Don't summarize the system message
+5. Make sure all tool calls and responses are summarized, as they are vital
+6. Focus on token usage efficiency and system message preservation"""
+
+
 def reload_summarization_agent():
     """Create a specialized agent for summarizing messages when context limit is reached."""
+    from code_puppy.model_utils import prepare_prompt_for_model
+
     models_config = ModelFactory.load_config()
-    model_name = get_model_name()
+    model_name = get_global_model_name()
     model = ModelFactory.get_model(model_name, models_config)
 
-    # Specialized instructions for summarization
-    instructions = """You are a message summarization expert. Your task is to summarize conversation messages
-while preserving important context and information. The summaries should be concise but capture the essential
-content and intent of the original messages. This is to help manage token usage in a conversation history
-while maintaining context for the AI to continue the conversation effectively.
+    # Handle claude-code models: swap instructions (prompt prepending happens in run_summarization_sync)
+    instructions = _get_summarization_instructions()
+    prepared = prepare_prompt_for_model(
+        model_name, instructions, "", prepend_system_to_user=False
+    )
+    instructions = prepared.instructions
 
-When summarizing:
-1. Keep summary brief but informative
-2. Preserve key information and decisions
-3. Keep any important technical details
-4. Don't summarize the system message
-5. Make sure all tool calls and responses are summarized, as they are vital"""
+    model_settings = make_model_settings(model_name)
 
     agent = Agent(
         model=model,
         instructions=instructions,
         output_type=str,
         retries=1,  # Fewer retries for summarization
+        model_settings=model_settings,
     )
+    if get_use_dbos():
+        from pydantic_ai.durable_exec.dbos import DBOSAgent
+
+        global _reload_count
+        _reload_count += 1
+        dbos_agent = DBOSAgent(agent, name=f"summarization-agent-{_reload_count}")
+        return dbos_agent
     return agent