Add Phase 7: polish and hardening — retry, truncation, sessions, shutdown

- Config extensions: retry backoff, truncation threshold, session persistence
- LLM retry with exponential backoff + jitter on transient errors (5xx, connection)
- Conversation truncation: drops oldest messages preserving first user + recent N
- Session persistence: auto-save/restore with atomic writes, cleanup of old files
- Graceful shutdown: SIGTERM handler, cancel() on AgentLoop, save-on-exit
- Partial message recovery on mid-stream interruption
- New slash commands: /save, /session
- 18 new tests (5 retry, 5 truncation, 4 session, 4 integration workflows)
- README.md and docs/tools.md documentation

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

app/agent/context.py

@@ -77,3 +77,101 @@ class SessionContext:
     def start_time(self) -> datetime:
         """Session start timestamp (UTC)."""
         return self._start_time
+
+    def truncate_history(self, system_token_estimate: int = 0) -> int:
+        """Drop oldest messages to bring token usage under budget.
+
+        Preserves the first user message and the most recent N messages
+        (configured by ``truncation_keep_recent``). Cleans up orphaned tool
+        messages after truncation.
+
+        Args:
+            system_token_estimate: Estimated tokens used by the system prompt.
+
+        Returns:
+            Number of messages dropped.
+        """
+        budget = self._token_counter.budget
+        threshold = self._config.agent.truncation_threshold
+        keep_recent = self._config.agent.truncation_keep_recent
+
+        estimated = self._token_counter.estimate_messages_tokens(self._history) + system_token_estimate
+        if estimated < threshold * budget:
+            return 0
+
+        target = int(budget * 0.75)  # headroom
+        if len(self._history) <= keep_recent + 1:
+            return 0
+
+        # Split: first user message | droppable middle | recent tail
+        first_msg = self._history[0] if self._history and self._history[0].role == "user" else None
+        start_idx = 1 if first_msg else 0
+        tail_start = max(start_idx, len(self._history) - keep_recent)
+
+        dropped = 0
+        drop_indices: set[int] = set()
+
+        for i in range(start_idx, tail_start):
+            drop_indices.add(i)
+            dropped += 1
+            # Recalculate with remaining messages
+            remaining = [m for j, m in enumerate(self._history) if j not in drop_indices]
+            est = self._token_counter.estimate_messages_tokens(remaining) + system_token_estimate
+            if est < target:
+                break
+
+        if dropped == 0:
+            return 0
+
+        self._history = [m for j, m in enumerate(self._history) if j not in drop_indices]
+
+        # Clean up orphaned tool messages
+        self._cleanup_orphaned_tool_messages()
+
+        return dropped
+
+    def _cleanup_orphaned_tool_messages(self) -> None:
+        """Remove tool messages whose tool_call_id doesn't match any assistant tool_call."""
+        # Collect all tool_call IDs from assistant messages
+        valid_tc_ids: set[str] = set()
+        for msg in self._history:
+            if msg.role == "assistant" and msg.tool_calls:
+                for tc in msg.tool_calls:
+                    valid_tc_ids.add(tc.id)
+
+        # Remove tool messages referencing missing tool calls
+        self._history = [
+            msg for msg in self._history
+            if msg.role != "tool" or (msg.tool_call_id and msg.tool_call_id in valid_tc_ids)
+        ]
+
+    def to_serializable(self) -> dict:
+        """Export messages and token state for session persistence.
+
+        Returns:
+            Dict with messages and token usage data.
+        """
+        return {
+            "messages": [m.model_dump(exclude_none=True) for m in self._history],
+            "token_usage": self._token_counter.cumulative_usage.model_dump(),
+        }
+
+    def restore_from(self, data: dict) -> None:
+        """Clear and replay from serialized data.
+
+        Args:
+            data: Dict with messages and optional token_usage as produced by to_serializable().
+        """
+        self._history.clear()
+        self._message_count = 0
+
+        for msg_data in data.get("messages", []):
+            msg = Message(**msg_data)
+            self._history.append(msg)
+            self._message_count += 1
+
+        token_data = data.get("token_usage")
+        if token_data:
+            from app.utils.token_counter import TokenUsage
+            usage = TokenUsage(**token_data)
+            self._token_counter.count_usage(usage)

app/agent/loop.py

@@ -7,7 +7,7 @@ from app.agent.context import SessionContext
 from app.models.config import AppConfig
 from app.models.message import Message
 from app.models.tool_call import ToolCall, ToolResult, ToolResultStatus
-from app.services.llm import LLMClient, LLMConnectionError, LLMError
+from app.services.llm import LLMClient, LLMConnectionError, LLMError, LLMStreamError
 from app.services.permissions import PermissionsService
 from app.services.streaming import StreamHandler
 from app.tools.registry import ToolRegistry
@@ -51,6 +51,11 @@ class AgentLoop:
         self._permissions = permissions
         self._tools_schema = registry.get_openai_tools_schema()
         self._system_prompt = self._build_system_prompt()
+        self._cancelled = False
+
+    def cancel(self) -> None:
+        """Request cancellation of the current agent turn."""
+        self._cancelled = True
 
     def _build_system_prompt(self) -> str:
         """Build the system prompt including tool schemas and agent instructions."""
@@ -81,15 +86,25 @@ class AgentLoop:
             user_input: The user's message text.
         """
         self._ctx.add_message("user", user_input)
+        self._cancelled = False
 
         max_iter = self._config.agent.max_iterations
         reasoning_only_streak = 0
         for iteration in range(1, max_iter + 1):
-            # Check token budget
-            if self._ctx.token_counter.is_over_budget():
-                print_warning("Token budget exceeded. Stopping agent loop.")
+            if self._cancelled:
+                print_warning("Agent loop cancelled.")
                 break
 
+            # Check token budget — try truncation before giving up
+            if self._ctx.token_counter.is_over_budget():
+                system_tokens = self._ctx.token_counter.estimate_tokens(self._system_prompt)
+                dropped = self._ctx.truncate_history(system_tokens)
+                if dropped > 0:
+                    print_warning(f"Token budget pressure: dropped {dropped} oldest messages.")
+                else:
+                    print_warning("Token budget exceeded, cannot truncate further. Stopping.")
+                    break
+
             if iteration > 1:
                 print_iteration_header(iteration, max_iter)
 
@@ -169,18 +184,25 @@ class AgentLoop:
     async def _llm_step(self) -> Message | None:
         """Stream one LLM response and return the accumulated Message.
 
+        Uses retry-enabled streaming. On mid-stream errors, attempts to recover
+        partial content if available.
+
         Returns:
             The assistant Message, or None if an error occurred.
         """
         messages = self._get_messages_with_system_prompt()
         try:
-            chunk_iter = self._client.stream_chat(messages, tools=self._tools_schema)
+            chunk_iter = self._client.stream_chat_with_retry(messages, tools=self._tools_schema)
             return await self._handler.process_stream(chunk_iter)
         except KeyboardInterrupt:
             print_warning("Response interrupted.")
             self._handler.reset()
             return None
-        except LLMConnectionError as e:
+        except (LLMConnectionError, LLMStreamError) as e:
+            partial = self._handler.get_partial_message()
+            if partial is not None:
+                print_warning(f"Stream interrupted ({e}), returning partial response.")
+                return partial
             print_error(f"Connection error: {e}")
             return None
         except LLMError as e:

app/main.py

@@ -2,6 +2,7 @@
 
 import argparse
 import asyncio
+import signal
 import sys
 from pathlib import Path
 
@@ -12,6 +13,7 @@ from app.agent.loop import AgentLoop
 from app.models.config import AppConfig, load_config
 from app.services.llm import LLMClient, LLMConnectionError, LLMError
 from app.services.permissions import PermissionsService
+from app.services.session import SessionManager
 from app.services.streaming import StreamHandler
 from app.tools.registry import create_default_registry
 from app.utils.display import (
@@ -63,17 +65,62 @@ async def _preflight(config: AppConfig) -> None:
         await client.preflight_check()
 
 
+def _offer_session_resume(
+    session_mgr: SessionManager, ctx: SessionContext
+) -> bool:
+    """Check for a saved session and offer to resume it.
+
+    Args:
+        session_mgr: Session manager instance.
+        ctx: Session context to restore into.
+
+    Returns:
+        True if a session was restored.
+    """
+    saved = session_mgr.load_latest()
+    if saved is None:
+        return False
+
+    msg_count = len(saved.messages)
+    print_info(f"Found previous session ({msg_count} messages, model: {saved.model})")
+
+    try:
+        answer = console.input("[bold cyan]Resume previous session? [y/N] [/bold cyan]").strip().lower()
+    except (KeyboardInterrupt, EOFError):
+        return False
+
+    if answer in ("y", "yes"):
+        session_mgr.restore(saved, ctx)
+        print_success(f"Session restored ({msg_count} messages).")
+        return True
+
+    return False
+
+
+def _save_session_quiet(session_mgr: SessionManager, ctx: SessionContext) -> None:
+    """Save session without raising on errors."""
+    try:
+        if ctx.message_count > 0:
+            session_mgr.save(ctx)
+    except OSError as e:
+        print_warning(f"Could not save session: {e}")
+
+
 async def _run_repl(
     ctx: SessionContext,
     config: AppConfig,
+    session_mgr: SessionManager,
     logger: structlog.stdlib.BoundLogger,
+    shutdown_event: asyncio.Event,
 ) -> None:
     """Run the interactive REPL loop with streaming LLM responses.
 
     Args:
         ctx: Session context for conversation state.
         config: Application configuration.
+        session_mgr: Session manager for auto-save.
         logger: Structured logger instance.
+        shutdown_event: Event signalling graceful shutdown.
     """
     registry = create_default_registry(config.agent.workspace_root, config)
     permissions = PermissionsService(config.permissions)
@@ -82,11 +129,12 @@ async def _run_repl(
         handler = StreamHandler(config.display)
         agent = AgentLoop(config, ctx, client, handler, registry, permissions)
 
-        while True:
+        while not shutdown_event.is_set():
             try:
                 user_input = console.input("[bold cyan]> [/bold cyan]")
             except (KeyboardInterrupt, EOFError):
-                console.print("\n[dim]Goodbye![/dim]")
+                _save_session_quiet(session_mgr, ctx)
+                console.print("\n[dim]Session saved. Goodbye![/dim]")
                 break
 
             user_input = user_input.strip()
@@ -97,13 +145,24 @@ async def _run_repl(
             if user_input.startswith("/"):
                 command = user_input.lower()
                 if command == "/quit":
-                    console.print("[dim]Goodbye![/dim]")
+                    _save_session_quiet(session_mgr, ctx)
+                    console.print("[dim]Session saved. Goodbye![/dim]")
                     break
                 elif command == "/history":
                     print_history(ctx.get_history())
                 elif command == "/clear":
                     ctx.clear_history()
                     print_success("Conversation history cleared.")
+                elif command == "/save":
+                    try:
+                        path = session_mgr.save(ctx)
+                        print_success(f"Session saved to {path}")
+                    except OSError as e:
+                        print_error(f"Failed to save session: {e}")
+                elif command == "/session":
+                    print_info(f"Messages: {ctx.message_count}")
+                    print_info(f"Tokens: ~{ctx.estimated_tokens:,} / {ctx.token_counter.budget:,}")
+                    print_info(f"Started: {ctx.start_time.isoformat()}")
                 else:
                     print_warning(f"Unknown command: {user_input}")
                 continue
@@ -112,6 +171,15 @@ async def _run_repl(
             await agent.run_turn(user_input)
             logger.debug("turn_complete", message_count=ctx.message_count)
 
+            # Auto-save after each turn
+            if config.session.auto_save:
+                _save_session_quiet(session_mgr, ctx)
+
+        # Shutdown triggered by signal
+        if shutdown_event.is_set():
+            _save_session_quiet(session_mgr, ctx)
+            console.print("\n[dim]Session saved. Shutting down gracefully.[/dim]")
+
 
 def main() -> None:
     """Main entrypoint: load config, setup logging, launch interactive REPL."""
@@ -154,12 +222,37 @@ def main() -> None:
 
     print_success("Ollama connected, model ready.")
 
-    # Create session and start REPL
+    # Create session and session manager
     ctx = SessionContext(config)
+    session_mgr = SessionManager(config.session, config.agent.workspace_root, config.llm.model)
+
+    # Clean up old session files
+    cleaned = session_mgr.cleanup_old()
+    if cleaned > 0:
+        logger.info("old_sessions_cleaned", count=cleaned)
+
+    # Offer to resume previous session
+    if config.session.offer_resume:
+        _offer_session_resume(session_mgr, ctx)
+
     logger.info("startup_complete")
 
-    print_info("Commands: /quit, /history, /clear")
-    asyncio.run(_run_repl(ctx, config, logger))
+    # Setup shutdown event and SIGTERM handler
+    shutdown_event = asyncio.Event()
+
+    original_sigterm = signal.getsignal(signal.SIGTERM)
+
+    def _sigterm_handler(signum: int, frame: object) -> None:
+        shutdown_event.set()
+
+    signal.signal(signal.SIGTERM, _sigterm_handler)
+
+    print_info("Commands: /quit, /history, /clear, /save, /session")
+
+    try:
+        asyncio.run(_run_repl(ctx, config, session_mgr, logger, shutdown_event))
+    finally:
+        signal.signal(signal.SIGTERM, original_sigterm)
 
 
 if __name__ == "__main__":

app/models/config.py

@@ -16,6 +16,9 @@ class LLMConfig(BaseModel):
     temperature: float = Field(default=0.1, description="Sampling temperature")
     max_tokens: int = Field(default=4096, description="Maximum tokens in LLM response")
     timeout: int = Field(default=120, description="Request timeout in seconds")
+    max_retries: int = Field(default=3, description="Max retry attempts on transient errors")
+    retry_backoff_base: float = Field(default=1.0, description="Base seconds for exponential backoff")
+    retry_backoff_max: float = Field(default=30.0, description="Maximum backoff seconds")
 
 
 class AgentConfig(BaseModel):
@@ -28,6 +31,12 @@ class AgentConfig(BaseModel):
     workspace_root: Path = Field(
         default=Path("."), description="Root directory for file operations"
     )
+    truncation_keep_recent: int = Field(
+        default=10, description="Number of recent messages to preserve during truncation"
+    )
+    truncation_threshold: float = Field(
+        default=0.85, description="Token budget fraction that triggers truncation"
+    )
 
 
 class PermissionsConfig(BaseModel):
@@ -60,6 +69,19 @@ class ToolsConfig(BaseModel):
     filesystem: FilesystemToolConfig = Field(default_factory=FilesystemToolConfig)
 
 
+class SessionConfig(BaseModel):
+    """Session persistence configuration."""
+
+    session_dir: Path = Field(
+        default=Path(".sneakycode/sessions"), description="Directory for session files"
+    )
+    auto_save: bool = Field(default=True, description="Auto-save session after each turn")
+    max_session_age_hours: int = Field(
+        default=72, description="Max age in hours before session files are cleaned up"
+    )
+    offer_resume: bool = Field(default=True, description="Offer to resume previous sessions on startup")
+
+
 class DisplayConfig(BaseModel):
     """Terminal display preferences."""
 
@@ -76,6 +98,7 @@ class AppConfig(BaseModel):
     permissions: PermissionsConfig = Field(default_factory=PermissionsConfig)
     tools: ToolsConfig = Field(default_factory=ToolsConfig)
     display: DisplayConfig = Field(default_factory=DisplayConfig)
+    session: SessionConfig = Field(default_factory=SessionConfig)
 
     @model_validator(mode="after")
     def resolve_workspace_root(self) -> "AppConfig":

app/services/llm.py

@@ -1,6 +1,8 @@
 """LLM client wrapper for Ollama / OpenAI-compatible endpoints."""
 
+import asyncio
 import json
+import random
 from collections.abc import AsyncIterator
 from typing import Any, Self
 
@@ -162,6 +164,61 @@ class LLMClient:
         except httpx.HTTPError as e:
             raise LLMError(f"HTTP error communicating with LLM: {e}") from e
 
+    async def stream_chat_with_retry(
+        self,
+        messages: list[Message],
+        tools: list[dict[str, Any]] | None = None,
+    ) -> AsyncIterator[dict]:
+        """Stream chat with automatic retry on transient errors.
+
+        Retries on LLMConnectionError and LLMResponseError with status >= 500.
+        Does NOT retry on 4xx errors (client-side, not transient).
+        Uses exponential backoff with jitter.
+
+        Args:
+            messages: Conversation history to send to the model.
+            tools: Optional OpenAI function-calling tool schemas.
+
+        Yields:
+            Parsed JSON dicts from each SSE data line.
+
+        Raises:
+            LLMConnectionError: After exhausting retries on connection failures.
+            LLMResponseError: After exhausting retries on server errors, or immediately on 4xx.
+        """
+        max_retries = self._config.max_retries
+        last_exception: LLMError | None = None
+
+        for attempt in range(max_retries + 1):
+            try:
+                async for chunk in self.stream_chat(messages, tools=tools):
+                    yield chunk
+                return
+            except LLMConnectionError as e:
+                last_exception = e
+            except LLMResponseError as e:
+                if e.status_code is not None and e.status_code < 500:
+                    raise
+                last_exception = e
+            except LLMStreamError as e:
+                last_exception = e
+
+            if attempt < max_retries:
+                backoff = min(
+                    self._config.retry_backoff_base * (2 ** attempt) + random.uniform(0, 1),
+                    self._config.retry_backoff_max,
+                )
+                logger.warning(
+                    "llm_retry",
+                    attempt=attempt + 1,
+                    max_retries=max_retries,
+                    backoff_seconds=round(backoff, 2),
+                    error=str(last_exception),
+                )
+                await asyncio.sleep(backoff)
+
+        raise last_exception  # type: ignore[misc]
+
     async def close(self) -> None:
         """Close the underlying HTTP client."""
         await self._client.aclose()

app/services/session.py

@@ -0,0 +1,148 @@
+"""Session persistence — auto-save and restore conversation state."""
+
+import hashlib
+import json
+from datetime import UTC, datetime
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+from pydantic import BaseModel, Field
+
+from app.models.config import SessionConfig
+from app.utils.logging import get_logger
+
+if TYPE_CHECKING:
+    from app.agent.context import SessionContext
+
+logger = get_logger(__name__)
+
+
+class SessionData(BaseModel):
+    """Serialized session state for persistence."""
+
+    version: int = Field(default=1, description="Schema version for forward compatibility")
+    session_id: str = Field(description="Unique session identifier")
+    created_at: str = Field(description="ISO timestamp of session creation")
+    updated_at: str = Field(description="ISO timestamp of last update")
+    model: str = Field(description="LLM model name used in session")
+    workspace_root: str = Field(description="Workspace root path")
+    messages: list[dict] = Field(default_factory=list, description="Serialized messages")
+    token_usage: dict = Field(default_factory=dict, description="Cumulative token usage")
+
+
+class SessionManager:
+    """Manages session file I/O: save, load, restore, and cleanup.
+
+    Session files are keyed by a hash of the workspace root path so that
+    each project directory has its own session history.
+    """
+
+    def __init__(self, config: SessionConfig, workspace_root: Path, model: str) -> None:
+        """Initialize session manager.
+
+        Args:
+            config: Session configuration.
+            workspace_root: Absolute path to workspace root.
+            model: LLM model name for session metadata.
+        """
+        self._config = config
+        self._workspace_root = workspace_root
+        self._model = model
+        self._workspace_hash = hashlib.sha256(str(workspace_root).encode()).hexdigest()[:12]
+        self._session_dir = workspace_root / config.session_dir
+        self._session_id = f"{self._workspace_hash}_{datetime.now(UTC).strftime('%Y%m%d_%H%M%S')}"
+
+    def save(self, ctx: "SessionContext") -> Path:
+        """Save session state to a JSON file via atomic write.
+
+        Args:
+            ctx: Session context to persist.
+
+        Returns:
+            Path to the saved session file.
+        """
+        self._session_dir.mkdir(parents=True, exist_ok=True)
+
+        serialized = ctx.to_serializable()
+        data = SessionData(
+            session_id=self._session_id,
+            created_at=ctx.start_time.isoformat(),
+            updated_at=datetime.now(UTC).isoformat(),
+            model=self._model,
+            workspace_root=str(self._workspace_root),
+            messages=serialized["messages"],
+            token_usage=serialized["token_usage"],
+        )
+
+        file_path = self._session_dir / f"{self._session_id}.json"
+        tmp_path = file_path.with_suffix(".tmp")
+
+        tmp_path.write_text(data.model_dump_json(indent=2), encoding="utf-8")
+        tmp_path.rename(file_path)
+
+        logger.debug("session_saved", path=str(file_path))
+        return file_path
+
+    def load_latest(self) -> SessionData | None:
+        """Find and load the newest session file for this workspace.
+
+        Returns:
+            SessionData if a valid session is found, None otherwise.
+        """
+        if not self._session_dir.exists():
+            return None
+
+        session_files = sorted(
+            self._session_dir.glob(f"{self._workspace_hash}_*.json"),
+            key=lambda p: p.stat().st_mtime,
+            reverse=True,
+        )
+
+        for path in session_files:
+            try:
+                raw = json.loads(path.read_text(encoding="utf-8"))
+                return SessionData(**raw)
+            except (json.JSONDecodeError, ValueError, OSError) as e:
+                logger.warning("session_load_error", path=str(path), error=str(e))
+                continue
+
+        return None
+
+    def restore(self, data: SessionData, ctx: "SessionContext") -> None:
+        """Replay session data into a SessionContext.
+
+        Args:
+            data: Saved session data to restore.
+            ctx: Session context to populate.
+        """
+        ctx.restore_from({
+            "messages": data.messages,
+            "token_usage": data.token_usage,
+        })
+        # Preserve the original session ID for continuity
+        self._session_id = data.session_id
+        logger.info("session_restored", session_id=data.session_id, messages=len(data.messages))
+
+    def cleanup_old(self) -> int:
+        """Delete session files older than max_session_age_hours.
+
+        Returns:
+            Number of files deleted.
+        """
+        if not self._session_dir.exists():
+            return 0
+
+        cutoff = datetime.now(UTC).timestamp() - (self._config.max_session_age_hours * 3600)
+        deleted = 0
+
+        for path in self._session_dir.glob("*.json"):
+            try:
+                if path.stat().st_mtime < cutoff:
+                    path.unlink()
+                    deleted += 1
+            except OSError:
+                continue
+
+        if deleted > 0:
+            logger.info("sessions_cleaned", deleted=deleted)
+        return deleted

app/services/streaming.py

@@ -142,6 +142,24 @@ class StreamHandler:
             )
         return result
 
+    def get_partial_message(self) -> Message | None:
+        """Return whatever content/tool_calls have been accumulated so far.
+
+        Useful for mid-stream interruption recovery — returns None if nothing
+        has been accumulated yet.
+
+        Returns:
+            Partial assistant Message, or None if no content accumulated.
+        """
+        tool_calls = self._build_tool_calls() or None
+        if not self._accumulated_content and not tool_calls:
+            return None
+        return Message(
+            role="assistant",
+            content=self._accumulated_content or None,
+            tool_calls=tool_calls,
+        )
+
     @property
     def usage(self) -> TokenUsage | None:
         """Token usage reported by the API, if available."""