code-puppy 0.0.194__tar.gz → 0.0.195__tar.gz

{code_puppy-0.0.194 → code_puppy-0.0.195}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: code-puppy
-Version: 0.0.194
+Version: 0.0.195
 Summary: Code generation agent
 Project-URL: repository, https://github.com/mpfaffenberger/code_puppy
 Project-URL: HomePage, https://github.com/mpfaffenberger/code_puppy

{code_puppy-0.0.194 → code_puppy-0.0.195}/code_puppy/command_line/command_handler.py

@@ -1,9 +1,12 @@
 import os
+from datetime import datetime
+from pathlib import Path
 
 from code_puppy.command_line.model_picker_completion import update_model_in_input
 from code_puppy.command_line.motd import print_motd
 from code_puppy.command_line.utils import make_directory_table
-from code_puppy.config import get_config_keys
+from code_puppy.config import CONTEXTS_DIR, get_config_keys
+from code_puppy.session_storage import list_sessions, load_session, save_session
 from code_puppy.tools.tools_content import tools_content
 
 
@@ -76,18 +79,6 @@ def get_commands_help():
         Text("/load_context", style="cyan")
         + Text(" <name>  Load message history from file")
     )
-    help_lines.append(
-        Text("", style="cyan")
-        + Text("Session Management:", style="bold yellow")
-    )
-    help_lines.append(
-        Text("auto_save_session", style="cyan")
-        + Text("    Auto-save session after each response (true/false)")
-    )
-    help_lines.append(
-        Text("max_saved_sessions", style="cyan")
-        + Text("    Maximum number of sessions to keep (default: 20, 0 = unlimited)")
-    )
     help_lines.append(
         Text("/set", style="cyan")
         + Text(
@@ -367,8 +358,13 @@ def handle_command(command: str):
             config_keys = get_config_keys()
             if "compaction_strategy" not in config_keys:
                 config_keys.append("compaction_strategy")
+            session_help = (
+                "\n[yellow]Session Management[/yellow]"
+                "\n  [cyan]auto_save_session[/cyan]    Auto-save chat after every response (true/false)"
+                "\n  [cyan]max_saved_sessions[/cyan]  Cap how many auto-saves to keep (0 = unlimited)"
+            )
             emit_warning(
-                f"Usage: /set KEY=VALUE or /set KEY VALUE\nConfig keys: {', '.join(config_keys)}\n[dim]Note: compaction_strategy can be 'summarization' or 'truncation'[/dim]"
+                f"Usage: /set KEY=VALUE or /set KEY VALUE\nConfig keys: {', '.join(config_keys)}\n[dim]Note: compaction_strategy can be 'summarization' or 'truncation'[/dim]{session_help}"
             )
             return True
         if key:
@@ -655,14 +651,7 @@ def handle_command(command: str):
         return pr_prompt
 
     if command.startswith("/dump_context"):
-        import json
-        import pickle
-        from datetime import datetime
-        from pathlib import Path
-
-        # estimate_tokens_for_message has been moved to BaseAgent class
         from code_puppy.agents.agent_manager import get_current_agent
-        from code_puppy.config import CONFIG_DIR
 
         tokens = command.split()
         if len(tokens) != 2:
@@ -677,49 +666,26 @@ def handle_command(command: str):
             emit_warning("No message history to dump!")
             return True
 
-        # Create contexts directory inside CONFIG_DIR if it doesn't exist
-        contexts_dir = Path(CONFIG_DIR) / "contexts"
-        contexts_dir.mkdir(parents=True, exist_ok=True)
-
         try:
-            # Save as pickle for exact preservation
-            pickle_file = contexts_dir / f"{session_name}.pkl"
-            with open(pickle_file, "wb") as f:
-                pickle.dump(history, f)
-
-            # Also save metadata as JSON for readability
-            meta_file = contexts_dir / f"{session_name}_meta.json"
-            current_agent = get_current_agent()
-            metadata = {
-                "session_name": session_name,
-                "timestamp": datetime.now().isoformat(),
-                "message_count": len(history),
-                "total_tokens": sum(
-                    current_agent.estimate_tokens_for_message(m) for m in history
-                ),
-                "file_path": str(pickle_file),
-            }
-
-            with open(meta_file, "w") as f:
-                json.dump(metadata, f, indent=2)
-
+            metadata = save_session(
+                history=history,
+                session_name=session_name,
+                base_dir=Path(CONTEXTS_DIR),
+                timestamp=datetime.now().isoformat(),
+                token_estimator=agent.estimate_tokens_for_message,
+            )
             emit_success(
-                f"✅ Context saved: {len(history)} messages ({metadata['total_tokens']} tokens)\n"
-                f"📁 Files: {pickle_file}, {meta_file}"
+                f"✅ Context saved: {metadata.message_count} messages ({metadata.total_tokens} tokens)\n"
+                f"📁 Files: {metadata.pickle_path}, {metadata.metadata_path}"
             )
             return True
 
-        except Exception as e:
-            emit_error(f"Failed to dump context: {e}")
+        except Exception as exc:
+            emit_error(f"Failed to dump context: {exc}")
             return True
 
     if command.startswith("/load_context"):
-        import pickle
-        from pathlib import Path
-
-        # estimate_tokens_for_message has been moved to BaseAgent class
         from code_puppy.agents.agent_manager import get_current_agent
-        from code_puppy.config import CONFIG_DIR
 
         tokens = command.split()
         if len(tokens) != 2:
@@ -727,38 +693,30 @@ def handle_command(command: str):
             return True
 
         session_name = tokens[1]
-        contexts_dir = Path(CONFIG_DIR) / "contexts"
-        pickle_file = contexts_dir / f"{session_name}.pkl"
+        contexts_dir = Path(CONTEXTS_DIR)
+        session_path = contexts_dir / f"{session_name}.pkl"
 
-        if not pickle_file.exists():
-            emit_error(f"Context file not found: {pickle_file}")
-            # List available contexts
-            available = list(contexts_dir.glob("*.pkl"))
+        try:
+            history = load_session(session_name, contexts_dir)
+        except FileNotFoundError:
+            emit_error(f"Context file not found: {session_path}")
+            available = list_sessions(contexts_dir)
             if available:
-                names = [f.stem for f in available]
-                emit_info(f"Available contexts: {', '.join(names)}")
+                emit_info(f"Available contexts: {', '.join(available)}")
             return True
-
-        try:
-            with open(pickle_file, "rb") as f:
-                history = pickle.load(f)
-
-            agent = get_current_agent()
-            agent.set_message_history(history)
-            current_agent = get_current_agent()
-            total_tokens = sum(
-                current_agent.estimate_tokens_for_message(m) for m in history
-            )
-
-            emit_success(
-                f"✅ Context loaded: {len(history)} messages ({total_tokens} tokens)\n"
-                f"📁 From: {pickle_file}"
-            )
+        except Exception as exc:
+            emit_error(f"Failed to load context: {exc}")
             return True
 
-        except Exception as e:
-            emit_error(f"Failed to load context: {e}")
-            return True
+        agent = get_current_agent()
+        agent.set_message_history(history)
+        total_tokens = sum(agent.estimate_tokens_for_message(m) for m in history)
+
+        emit_success(
+            f"✅ Context loaded: {len(history)} messages ({total_tokens} tokens)\n"
+            f"📁 From: {session_path}"
+        )
+        return True
 
     if command.startswith("/truncate"):
         from code_puppy.agents.agent_manager import get_current_agent

{code_puppy-0.0.194 → code_puppy-0.0.195}/code_puppy/config.py

@@ -1,8 +1,11 @@
 import configparser
+import datetime
 import json
 import os
 import pathlib
 
+from code_puppy.session_storage import cleanup_sessions, save_session
+
 CONFIG_DIR = os.path.join(os.path.expanduser("~"), ".code_puppy")
 CONFIG_FILE = os.path.join(CONFIG_DIR, "puppy.cfg")
 MCP_SERVERS_FILE = os.path.join(CONFIG_DIR, "mcp_servers.json")
@@ -10,6 +13,8 @@ COMMAND_HISTORY_FILE = os.path.join(CONFIG_DIR, "command_history.txt")
 MODELS_FILE = os.path.join(CONFIG_DIR, "models.json")
 EXTRA_MODELS_FILE = os.path.join(CONFIG_DIR, "extra_models.json")
 AGENTS_DIR = os.path.join(CONFIG_DIR, "agents")
+CONTEXTS_DIR = os.path.join(CONFIG_DIR, "contexts")
+AUTOSAVE_DIR = os.path.join(CONFIG_DIR, "autosaves")
 
 DEFAULT_SECTION = "puppy"
 REQUIRED_KEYS = ["puppy_name", "owner_name"]
@@ -698,115 +703,63 @@ def set_max_saved_sessions(max_sessions: int):
 
 
 def _cleanup_old_sessions():
-    """Remove oldest sessions if we exceed the max_saved_sessions limit."""
+    """Remove oldest auto-saved sessions if we exceed the max_saved_sessions limit."""
     max_sessions = get_max_saved_sessions()
-    if max_sessions <= 0:  # 0 means unlimited
+    if max_sessions <= 0:
         return
-        
-    from pathlib import Path
-    
-    contexts_dir = Path(CONFIG_DIR) / "contexts"
-    if not contexts_dir.exists():
+
+    autosave_dir = pathlib.Path(AUTOSAVE_DIR)
+    removed_sessions = cleanup_sessions(autosave_dir, max_sessions)
+    if not removed_sessions:
         return
-        
-    # Get all .pkl files (session files) and sort by modification time
-    session_files = []
-    for pkl_file in contexts_dir.glob("*.pkl"):
-        try:
-            session_files.append((pkl_file.stat().st_mtime, pkl_file))
-        except OSError:
-            continue
-            
-    # Sort by modification time (oldest first)
-    session_files.sort(key=lambda x: x[0])
-    
-    # If we have more than max_sessions, remove the oldest ones
-    if len(session_files) > max_sessions:
-        files_to_remove = session_files[:-max_sessions]  # All except the last max_sessions
-        
-        from rich.console import Console
-        console = Console()
-        
-        for _, old_file in files_to_remove:
-            try:
-                # Remove the .pkl file
-                old_file.unlink()
-                
-                # Also remove the corresponding _meta.json file if it exists
-                meta_file = contexts_dir / f"{old_file.stem}_meta.json"
-                if meta_file.exists():
-                    meta_file.unlink()
-                    
-                console.print(f"[dim]🗑️  Removed old session: {old_file.name}[/dim]")
-                
-            except OSError as e:
-                console.print(f"[dim]❌ Failed to remove {old_file.name}: {e}[/dim]")
+
+    from rich.console import Console
+
+    console = Console()
+    for session_name in removed_sessions:
+        console.print(f"[dim]🗑️  Removed old session: {session_name}.pkl[/dim]")
 
 
 def auto_save_session_if_enabled() -> bool:
-    """Automatically save the current session if auto_save_session is enabled.
-    
-    Returns:
-        True if session was saved, False otherwise
-    """
+    """Automatically save the current session if auto_save_session is enabled."""
     if not get_auto_save_session():
         return False
-        
+
     try:
-        import datetime
-        import json
-        import pickle
-        from pathlib import Path
+        import pathlib
+        from rich.console import Console
+
         from code_puppy.agents.agent_manager import get_current_agent
-        
-        # Get current agent and message history
+
+        console = Console()
+
         current_agent = get_current_agent()
         history = current_agent.get_message_history()
-        
         if not history:
-            return False  # No history to save
-            
-        # Create timestamp-based session name
-        timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
-        session_name = f"auto_session_{timestamp}"
-        
-        # Create contexts directory if it doesn't exist
-        contexts_dir = Path(CONFIG_DIR) / "contexts"
-        contexts_dir.mkdir(parents=True, exist_ok=True)
-        
-        # Save as pickle for exact preservation
-        pickle_file = contexts_dir / f"{session_name}.pkl"
-        with open(pickle_file, "wb") as f:
-            pickle.dump(history, f)
-            
-        # Also save metadata as JSON for readability
-        meta_file = contexts_dir / f"{session_name}_meta.json"
-        metadata = {
-            "session_name": session_name,
-            "timestamp": datetime.datetime.now().isoformat(),
-            "message_count": len(history),
-            "total_tokens": sum(
-                current_agent.estimate_tokens_for_message(m) for m in history
-            ),
-            "file_path": str(pickle_file),
-            "auto_saved": True,
-        }
-        
-        with open(meta_file, "w") as f:
-            json.dump(metadata, f, indent=2)
-            
-        from rich.console import Console
-        console = Console()
+            return False
+
+        now = datetime.datetime.now()
+        session_name = f"auto_session_{now.strftime('%Y%m%d_%H%M%S')}"
+        autosave_dir = pathlib.Path(AUTOSAVE_DIR)
+
+        metadata = save_session(
+            history=history,
+            session_name=session_name,
+            base_dir=autosave_dir,
+            timestamp=now.isoformat(),
+            token_estimator=current_agent.estimate_tokens_for_message,
+            auto_saved=True,
+        )
+
         console.print(
-            f"🐾 [dim]Auto-saved session: {len(history)} messages ({metadata['total_tokens']} tokens)[/dim]"
+            f"🐾 [dim]Auto-saved session: {metadata.message_count} messages ({metadata.total_tokens} tokens)[/dim]"
         )
-        
-        # Cleanup old sessions if limit is set
+
         _cleanup_old_sessions()
         return True
-        
-    except Exception as e:
+
+    except Exception as exc:  # pragma: no cover - defensive logging
         from rich.console import Console
-        console = Console()
-        console.print(f"[dim]❌ Failed to auto-save session: {str(e)}[/dim]")
+
+        Console().print(f"[dim]❌ Failed to auto-save session: {exc}[/dim]")
         return False

{code_puppy-0.0.194 → code_puppy-0.0.195}/code_puppy/main.py

@@ -1,10 +1,13 @@
 import argparse
 import asyncio
+import json
 import os
 import subprocess
 import sys
 import time
 import webbrowser
+from datetime import datetime
+from pathlib import Path
 
 from rich.console import Console, ConsoleOptions, RenderResult
 from rich.markdown import CodeBlock, Markdown
@@ -18,11 +21,13 @@ from code_puppy.command_line.prompt_toolkit_completion import (
     get_prompt_with_active_model,
 )
 from code_puppy.config import (
+    AUTOSAVE_DIR,
     COMMAND_HISTORY_FILE,
     ensure_config_exists,
     initialize_command_history_file,
     save_command_to_history,
 )
+from code_puppy.session_storage import list_sessions, load_session, restore_autosave_interactively
 from code_puppy.http_utils import find_available_port
 from code_puppy.tools.common import console
 
@@ -288,6 +293,8 @@ async def interactive_mode(message_renderer, initial_command: str = None) -> Non
     from code_puppy.messaging import emit_info
 
     emit_info("[bold cyan]Initializing agent...[/bold cyan]")
+
+
     # Initialize the runtime agent manager
     if initial_command:
         from code_puppy.agents import get_current_agent
@@ -367,6 +374,8 @@ async def interactive_mode(message_renderer, initial_command: str = None) -> Non
             emit_error(f"Error installing prompt_toolkit: {e}")
             emit_warning("Falling back to basic input without tab completion")
 
+    await restore_autosave_interactively(Path(AUTOSAVE_DIR))
+
     while True:
         from code_puppy.agents.agent_manager import get_current_agent
         from code_puppy.messaging import emit_info

code_puppy-0.0.195/code_puppy/session_storage.py

@@ -0,0 +1,241 @@
+"""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 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
+
+
+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"))
+
+
+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)
+    top_entries = entries[:5]
+
+    emit_system_message("[bold magenta]Autosave Sessions Available:[/bold magenta]")
+    for index, (name, timestamp, message_count) in enumerate(top_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"  [{index}] {name} ({message_display}, saved at {timestamp_display})"
+        )
+
+    if len(entries) > len(top_entries):
+        emit_system_message(
+            f"  [dim]...and {len(entries) - len(top_entries)} more autosaves[/dim]"
+        )
+
+    try:
+        selection = await get_input_with_combined_completion(
+            FormattedText([("class:prompt", "Load autosave (number, name, or Enter to skip): ")])
+        )
+    except (KeyboardInterrupt, EOFError):
+        emit_warning("Autosave selection cancelled")
+        return
+
+    selection = selection.strip()
+    if not selection:
+        return
+
+    chosen_name = None
+    if selection.isdigit():
+        idx = int(selection) - 1
+        if 0 <= idx < len(top_entries):
+            chosen_name = top_entries[idx][0]
+    else:
+        for name, _, _ in entries:
+            if name == selection:
+                chosen_name = name
+                break
+
+    if not chosen_name:
+        emit_warning("No autosave loaded (invalid selection)")
+        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)
+    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-0.0.194 → code_puppy-0.0.195}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "code-puppy"
-version = "0.0.194"
+version = "0.0.195"
 description = "Code generation agent"
 readme = "README.md"
 requires-python = ">=3.11"