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>
2026-03-11 10:20:16 -05:00
parent 82846d6236
commit 76ba490aa2
16 changed files with 1550 additions and 12 deletions
--- a/README.md
+++ b/README.md
@@ -0,0 +1,147 @@
+# SneakyCode
+
+A privacy-first, locally-running Python coding agent that uses a local LLM (via Ollama) to perform autonomous coding tasks inside a project directory.
+
+SneakyCode accepts natural language tasks and executes them using a defined toolset for filesystem operations, shell execution, code search, and file manipulation. It runs a ReAct-style tool-call loop: send conversation history to the LLM, receive tool calls, execute them with permission checks, and feed results back until the task is complete.
+
+## Prerequisites
+
+- Python 3.11+
+- [Ollama](https://ollama.ai/) running locally with a model that supports function calling (e.g., `qwen3.5`, `llama3.1`, `mistral-nemo`)
+- [uv](https://docs.astral.sh/uv/) (recommended) or pip
+
+## Installation
+
+```bash
+# Clone the repository
+git clone <repo-url>
+cd SneakyCode
+
+# Install dependencies
+uv sync --dev
+
+# Or with pip
+pip install -e ".[dev]"
+```
+
+## Configuration
+
+Edit `config/config.yaml` to configure the agent. Key settings:
+
+```yaml
+llm:
+  model: "qwen3.5:latest"        # Ollama model name
+  endpoint: "http://localhost:11434"  # Ollama endpoint
+  max_retries: 3                   # Retry attempts on transient errors
+  retry_backoff_base: 1.0          # Exponential backoff base (seconds)
+
+agent:
+  max_iterations: 25               # Max tool-call iterations per turn
+  max_conversation_tokens: 32000   # Token budget for conversation
+  workspace_root: "."              # Project directory for file operations
+  truncation_keep_recent: 10       # Messages preserved during truncation
+  truncation_threshold: 0.85       # Budget fraction that triggers truncation
+
+session:
+  auto_save: true                  # Save session after each turn
+  max_session_age_hours: 72        # Auto-cleanup old sessions
+  offer_resume: true               # Offer to resume on startup
+
+permissions:
+  auto_approve: [read_file, list_dir, grep_files, find_files, finish]
+  prompt_user: [write_file, delete_file, run_command, str_replace, patch_apply, make_dir]
+  deny: []
+```
+
+Environment variable `SNEAKYCODE_CONFIG` can override the config file path.
+
+## Usage
+
+```bash
+# Start the interactive REPL
+sneakycode
+
+# Or run directly
+python -m app.main
+
+# With options
+sneakycode --config path/to/config.yaml --verbose --log-file sneakycode.log
+```
+
+### REPL Commands
+
+| Command    | Description                          |
+|------------|--------------------------------------|
+| `/quit`    | Save session and exit                |
+| `/history` | Show conversation history            |
+| `/clear`   | Clear conversation history           |
+| `/save`    | Manually save session                |
+| `/session` | Show session info (messages, tokens) |
+
+### Session Persistence
+
+Sessions are automatically saved after each agent turn and on exit. On startup, SneakyCode offers to resume the most recent session for the current workspace.
+
+Session files are stored in `.sneakycode/sessions/` within the workspace root.
+
+## Available Tools
+
+SneakyCode provides 11 tools across 5 categories. See [docs/tools.md](docs/tools.md) for the full reference.
+
+| Category   | Tools                                           | Permission    |
+|------------|-------------------------------------------------|---------------|
+| Read       | `read_file`, `list_dir`                         | Auto-approved |
+| Search     | `grep_files`, `find_files`                      | Auto-approved |
+| Write      | `write_file`, `make_dir`, `delete_file`         | User confirm  |
+| Edit       | `str_replace`, `patch_apply`                    | User confirm  |
+| Shell      | `run_command`                                   | User confirm  |
+| Control    | `finish`                                        | Auto-approved |
+
+## Development
+
+```bash
+# Run tests
+.venv/bin/python -m pytest tests/ -v
+
+# Run with coverage
+.venv/bin/python -m pytest tests/ --cov=app
+
+# Lint
+.venv/bin/ruff check app/ tests/
+
+# Format
+.venv/bin/ruff format app/ tests/
+```
+
+### Project Structure
+
+```
+app/
+├── agent/           # Agent loop and session context
+├── models/          # Pydantic config and message schemas
+├── services/        # LLM client, streaming, permissions, session persistence
+├── tools/           # Tool implementations (one file per group)
+└── utils/           # Logging, display, file helpers, token counter
+config/
+└── config.yaml      # Application configuration
+tests/
+├── unit/            # Unit tests for individual components
+└── integration/     # End-to-end workflow tests with mocked LLM
+```
+
+## Architecture
+
+SneakyCode follows a **ReAct-style** agent pattern:
+
+1. User provides a task in natural language
+2. Agent sends conversation history + tool schemas to the LLM
+3. LLM responds with either text (task complete) or tool calls
+4. Agent executes tool calls with permission checks
+5. Results are fed back to the LLM for the next iteration
+6. Loop continues until the LLM produces a plain-text response or calls `finish`
+
+The LLM client is abstracted behind an OpenAI-compatible interface, so any endpoint implementing the `/v1/chat/completions` SSE streaming protocol works as a backend.
+
+## License
+
+MIT
--- a/app/agent/context.py
+++ b/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)
--- a/app/agent/loop.py
+++ b/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:
--- a/app/main.py
+++ b/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__":
--- a/app/models/config.py
+++ b/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":
--- a/app/services/llm.py
+++ b/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()
--- a/app/services/session.py
+++ b/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
--- a/app/services/streaming.py
+++ b/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."""
--- a/config/config.yaml
+++ b/config/config.yaml
@@ -7,11 +7,16 @@ llm:
  temperature: 0.1
  max_tokens: 4096
  timeout: 120
+  max_retries: 3
+  retry_backoff_base: 1.0
+  retry_backoff_max: 30.0

 agent:
  max_iterations: 25
  max_conversation_tokens: 32000
  workspace_root: "."
+  truncation_keep_recent: 10
+  truncation_threshold: 0.85

 permissions:
  auto_approve:
@@ -56,6 +61,12 @@ tools:
    max_file_size_bytes: 1048576  # 1 MB
    binary_detection: true

+session:
+  session_dir: ".sneakycode/sessions"
+  auto_save: true
+  max_session_age_hours: 72
+  offer_resume: true
+
 display:
  show_tool_calls: true
  show_token_usage: true
--- a/docs/tools.md
+++ b/docs/tools.md
@@ -0,0 +1,240 @@
+# Tool Reference
+
+SneakyCode provides 11 agent-callable tools organized into 5 categories. All file path arguments must be **relative to the workspace root**.
+
+## Permission Tiers
+
+| Tier          | Behavior                              | Tools                                                          |
+|---------------|---------------------------------------|----------------------------------------------------------------|
+| Auto-approved | Executed without user confirmation    | `read_file`, `list_dir`, `grep_files`, `find_files`, `finish` |
+| User confirm  | Prompts user before execution         | `write_file`, `make_dir`, `delete_file`, `str_replace`, `patch_apply`, `run_command` |
+| Denied        | Blocked entirely (configurable)       | Any tool added to `permissions.deny` in config                 |
+
+---
+
+## Read Tools
+
+### read_file
+
+Read the full contents of a text file.
+
+| Parameter   | Type | Required | Description                                    |
+|-------------|------|----------|------------------------------------------------|
+| `file_path` | str  | Yes      | Path to the file to read (relative to workspace) |
+
+**Permission:** Auto-approved
+
+**Example:**
+```json
+{"file_path": "app/main.py"}
+```
+
+**Notes:** Binary files are detected and rejected. Files exceeding `max_file_size_bytes` (default 1 MB) are rejected.
+
+---
+
+### list_dir
+
+List the contents of a directory. Directories are suffixed with `/`. Results are sorted with directories first, then files.
+
+| Parameter        | Type | Required | Default | Description                          |
+|------------------|------|----------|---------|--------------------------------------|
+| `directory_path` | str  | No       | `"."`   | Path to directory (relative)         |
+| `recursive`      | bool | No       | `false` | If true, list entries recursively    |
+
+**Permission:** Auto-approved
+
+**Example:**
+```json
+{"directory_path": "app/tools", "recursive": true}
+```
+
+---
+
+## Search Tools
+
+### grep_files
+
+Search for a regex pattern in file contents. Returns matching lines with file paths and line numbers.
+
+| Parameter      | Type       | Required | Default | Description                              |
+|----------------|------------|----------|---------|------------------------------------------|
+| `pattern`      | str        | Yes      |         | Regular expression pattern to search for |
+| `path`         | str        | No       | `"."`   | Directory or file to search in           |
+| `file_pattern` | str\|null  | No       | `null`  | Glob pattern to filter files (e.g. `*.py`) |
+
+**Permission:** Auto-approved
+
+**Example:**
+```json
+{"pattern": "def main", "path": "app/", "file_pattern": "*.py"}
+```
+
+---
+
+### find_files
+
+Search for files matching a name pattern. Returns relative file paths.
+
+| Parameter | Type | Required | Default | Description                                       |
+|-----------|------|----------|---------|---------------------------------------------------|
+| `pattern` | str  | Yes      |         | File name pattern (e.g. `*.py`, `config.yaml`)    |
+| `path`    | str  | No       | `"."`   | Directory to search in                             |
+
+**Permission:** Auto-approved
+
+**Example:**
+```json
+{"pattern": "*.yaml", "path": "config/"}
+```
+
+---
+
+## Write Tools
+
+### write_file
+
+Write text content to a file. Creates parent directories if needed. Overwrites existing file content.
+
+| Parameter   | Type | Required | Description                     |
+|-------------|------|----------|---------------------------------|
+| `file_path` | str  | Yes      | Path to the file to write       |
+| `content`   | str  | Yes      | Content to write to the file    |
+
+**Permission:** User confirmation required
+
+**Example:**
+```json
+{"file_path": "app/utils/helpers.py", "content": "def greet():\n    return 'hello'\n"}
+```
+
+---
+
+### make_dir
+
+Create a directory and any necessary parent directories.
+
+| Parameter        | Type | Required | Description                      |
+|------------------|------|----------|----------------------------------|
+| `directory_path` | str  | Yes      | Path to the directory to create  |
+
+**Permission:** User confirmation required
+
+**Example:**
+```json
+{"directory_path": "app/services/new_module"}
+```
+
+---
+
+### delete_file
+
+Delete a single file. Does not delete directories.
+
+| Parameter   | Type | Required | Description                     |
+|-------------|------|----------|---------------------------------|
+| `file_path` | str  | Yes      | Path to the file to delete      |
+
+**Permission:** User confirmation required
+
+**Example:**
+```json
+{"file_path": "app/utils/deprecated.py"}
+```
+
+---
+
+## Edit Tools
+
+### str_replace
+
+Replace exactly one occurrence of `old_str` with `new_str` in a file. Fails if `old_str` is not found or appears more than once.
+
+| Parameter   | Type | Required | Description                                        |
+|-------------|------|----------|----------------------------------------------------|
+| `file_path` | str  | Yes      | Path to the file to edit                           |
+| `old_str`   | str  | Yes      | The exact string to find and replace (must be unique) |
+| `new_str`   | str  | Yes      | The replacement string                             |
+
+**Permission:** User confirmation required
+
+**Example:**
+```json
+{
+  "file_path": "app/main.py",
+  "old_str": "def old_function():",
+  "new_str": "def new_function():"
+}
+```
+
+---
+
+### patch_apply
+
+Apply a unified diff (patch) to a file. The patch must be in standard unified diff format.
+
+| Parameter   | Type | Required | Description                                |
+|-------------|------|----------|--------------------------------------------|
+| `file_path` | str  | Yes      | Path to the file to patch                  |
+| `patch`     | str  | Yes      | Unified diff format patch to apply         |
+
+**Permission:** User confirmation required
+
+**Example:**
+```json
+{
+  "file_path": "app/main.py",
+  "patch": "--- a/app/main.py\n+++ b/app/main.py\n@@ -1,3 +1,3 @@\n-old line\n+new line\n"
+}
+```
+
+---
+
+## Shell Tools
+
+### run_command
+
+Run a shell command in the workspace directory. Only allowed commands may be executed; dangerous commands are blocked.
+
+| Parameter | Type     | Required | Default | Description                       |
+|-----------|----------|----------|---------|-----------------------------------|
+| `command` | str      | Yes      |         | Shell command to execute          |
+| `timeout` | int\|null | No      | `30`    | Timeout in seconds                |
+
+**Permission:** User confirmation required. Subject to `tools.shell.allowed_commands` and `tools.shell.denied_commands` in config.
+
+**Example:**
+```json
+{"command": "git status", "timeout": 10}
+```
+
+**Notes:** Output is truncated to `max_output_bytes` (default 64 KB). The command's first word is checked against allow/deny lists.
+
+---
+
+## Control Tools
+
+### finish
+
+Signal that the task is complete. Terminates the agent loop.
+
+| Parameter | Type | Required | Default            | Description                  |
+|-----------|------|----------|--------------------|------------------------------|
+| `message` | str  | No       | `"Task complete."` | Final message to the user    |
+
+**Permission:** Auto-approved
+
+**Example:**
+```json
+{"message": "Created the new module with tests."}
+```
+
+---
+
+## Security Notes
+
+- All file paths are resolved against `workspace_root` with path traversal protection
+- Binary file detection prevents reading/writing binary files
+- File size limits prevent reading/writing excessively large files
+- Shell commands are validated against configurable allow/deny lists
+- Tool call arguments from the LLM are validated against JSON schema before execution
--- a/tests/integration/conftest.py
+++ b/tests/integration/conftest.py
@@ -0,0 +1,150 @@
+"""Shared fixtures for integration tests."""
+
+import json
+from collections.abc import AsyncIterator
+from pathlib import Path
+from typing import Any
+from unittest.mock import AsyncMock
+
+import pytest
+
+from app.models.config import AgentConfig, AppConfig, DisplayConfig, LLMConfig, PermissionsConfig, SessionConfig
+from app.models.message import Message
+from app.services.llm import LLMClient
+
+
+@pytest.fixture
+def tmp_workspace(tmp_path: Path) -> Path:
+    """Create a temporary workspace directory with a sample file."""
+    ws = tmp_path / "workspace"
+    ws.mkdir()
+    (ws / "hello.txt").write_text("Hello, world!")
+    return ws
+
+
+@pytest.fixture
+def test_config(tmp_workspace: Path) -> AppConfig:
+    """AppConfig suitable for integration tests."""
+    return AppConfig(
+        llm=LLMConfig(
+            model="test-model",
+            endpoint="http://localhost:11434",
+            max_retries=2,
+            retry_backoff_base=0.01,
+            retry_backoff_max=0.02,
+        ),
+        agent=AgentConfig(
+            max_iterations=10,
+            max_conversation_tokens=32000,
+            workspace_root=tmp_workspace,
+            truncation_keep_recent=4,
+            truncation_threshold=0.85,
+        ),
+        permissions=PermissionsConfig(
+            auto_approve=["read_file", "list_dir", "grep_files", "find_files", "finish"],
+        ),
+        display=DisplayConfig(
+            show_tool_calls=False,
+            show_token_usage=False,
+            stream_output=False,
+        ),
+        session=SessionConfig(
+            session_dir=tmp_workspace / ".sneakycode" / "sessions",
+        ),
+    )
+
+
+def make_text_chunks(content: str) -> list[dict[str, Any]]:
+    """Create SSE chunk dicts for a plain text response."""
+    chunks = []
+    for char in content:
+        chunks.append({
+            "choices": [{"delta": {"content": char}, "index": 0}]
+        })
+    chunks.append({
+        "choices": [{"delta": {}, "finish_reason": "stop", "index": 0}],
+        "usage": {"prompt_tokens": 10, "completion_tokens": len(content), "total_tokens": 10 + len(content)},
+    })
+    return chunks
+
+
+def make_tool_call_chunks(name: str, args: dict[str, Any], tc_id: str = "call_001") -> list[dict[str, Any]]:
+    """Create SSE chunk dicts for a tool call response."""
+    args_str = json.dumps(args)
+    chunks = [
+        {
+            "choices": [{
+                "delta": {
+                    "tool_calls": [{
+                        "index": 0,
+                        "id": tc_id,
+                        "function": {"name": name, "arguments": ""},
+                    }]
+                },
+                "index": 0,
+            }]
+        },
+        {
+            "choices": [{
+                "delta": {
+                    "tool_calls": [{
+                        "index": 0,
+                        "function": {"arguments": args_str},
+                    }]
+                },
+                "index": 0,
+            }]
+        },
+        {
+            "choices": [{"delta": {}, "finish_reason": "tool_calls", "index": 0}],
+            "usage": {"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30},
+        },
+    ]
+    return chunks
+
+
+class MockLLMClient:
+    """LLM client that returns scripted SSE chunk sequences."""
+
+    def __init__(self, responses: list[list[dict[str, Any]]]) -> None:
+        self._responses = list(responses)
+        self._call_count = 0
+
+    async def stream_chat(
+        self,
+        messages: list[Message],
+        tools: list[dict[str, Any]] | None = None,
+    ) -> AsyncIterator[dict]:
+        if self._call_count >= len(self._responses):
+            raise RuntimeError("MockLLMClient ran out of scripted responses")
+        chunks = self._responses[self._call_count]
+        self._call_count += 1
+        for chunk in chunks:
+            yield chunk
+
+    async def stream_chat_with_retry(
+        self,
+        messages: list[Message],
+        tools: list[dict[str, Any]] | None = None,
+    ) -> AsyncIterator[dict]:
+        async for chunk in self.stream_chat(messages, tools=tools):
+            yield chunk
+
+    @property
+    def call_count(self) -> int:
+        return self._call_count
+
+    async def close(self) -> None:
+        pass
+
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, *exc):
+        pass
+
+
+@pytest.fixture
+def mock_llm_client():
+    """Factory fixture for creating MockLLMClient instances."""
+    return MockLLMClient
--- a/tests/integration/test_agent_workflows.py
+++ b/tests/integration/test_agent_workflows.py
@@ -0,0 +1,155 @@
+"""Integration tests for end-to-end agent workflows with mocked LLM."""
+
+from pathlib import Path
+from unittest.mock import AsyncMock, patch
+
+import pytest
+
+from app.agent.context import SessionContext
+from app.agent.loop import AgentLoop
+from app.models.config import AgentConfig, AppConfig, LLMConfig
+from app.services.llm import LLMConnectionError
+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 .conftest import MockLLMClient, make_text_chunks, make_tool_call_chunks
+
+
+@pytest.fixture
+def agent_factory(test_config: AppConfig, tmp_workspace: Path):
+    """Factory that creates an AgentLoop wired to a MockLLMClient."""
+
+    def create(responses):
+        ctx = SessionContext(test_config)
+        mock_client = MockLLMClient(responses)
+        handler = StreamHandler(test_config.display)
+        registry = create_default_registry(test_config.agent.workspace_root, test_config)
+        permissions = PermissionsService(test_config.permissions)
+        agent = AgentLoop(test_config, ctx, mock_client, handler, registry, permissions)
+        return agent, ctx, mock_client
+
+    return create
+
+
+class TestAgentWorkflows:
+    @pytest.mark.asyncio
+    async def test_multi_turn_read_workflow(self, agent_factory, tmp_workspace: Path) -> None:
+        """Agent reads a file then responds with text — full 2-turn workflow."""
+        responses = [
+            make_tool_call_chunks("read_file", {"file_path": "hello.txt"}),
+            make_text_chunks("The file contains: Hello, world!"),
+        ]
+        agent, ctx, mock_client = agent_factory(responses)
+
+        await agent.run_turn("What's in hello.txt?")
+
+        assert mock_client.call_count == 2
+        history = ctx.get_history()
+        # user, assistant (tool_call), tool (result), assistant (text)
+        assert len(history) == 4
+        assert history[0].role == "user"
+        assert history[1].role == "assistant"
+        assert history[1].tool_calls is not None
+        assert history[2].role == "tool"
+        assert "Hello, world!" in (history[2].content or "")
+        assert history[3].role == "assistant"
+        assert "Hello, world!" in (history[3].content or "")
+
+    @pytest.mark.asyncio
+    async def test_token_budget_truncation(self, test_config: AppConfig, tmp_workspace: Path) -> None:
+        """When token budget is exceeded, truncation drops messages instead of stopping."""
+        # Use a tiny budget to trigger truncation
+        test_config.agent.max_conversation_tokens = 100
+        test_config.agent.truncation_keep_recent = 2
+        test_config.agent.truncation_threshold = 0.5
+
+        ctx = SessionContext(test_config)
+
+        # Fill history with enough to exceed budget
+        for i in range(10):
+            ctx.add_message("user", f"Message {i} " * 20)
+            ctx.add_message("assistant", f"Response {i} " * 20)
+
+        # Force token counter over budget
+        from app.utils.token_counter import TokenUsage
+        ctx.token_counter.count_usage(TokenUsage(total_tokens=100))
+
+        original_count = len(ctx.get_history())
+        dropped = ctx.truncate_history()
+
+        assert dropped > 0
+        assert len(ctx.get_history()) < original_count
+        # Recent messages should still be present
+        assert len(ctx.get_history()) >= 2
+
+    @pytest.mark.asyncio
+    async def test_session_save_and_restore(self, test_config: AppConfig, tmp_workspace: Path) -> None:
+        """Session can be saved after a turn and restored into a fresh context."""
+        responses = [make_text_chunks("Hello from the agent!")]
+        ctx = SessionContext(test_config)
+        mock_client = MockLLMClient(responses)
+        handler = StreamHandler(test_config.display)
+        registry = create_default_registry(test_config.agent.workspace_root, test_config)
+        permissions = PermissionsService(test_config.permissions)
+        agent = AgentLoop(test_config, ctx, mock_client, handler, registry, permissions)
+
+        await agent.run_turn("Hi")
+
+        # Save session
+        session_mgr = SessionManager(
+            test_config.session, test_config.agent.workspace_root, test_config.llm.model
+        )
+        path = session_mgr.save(ctx)
+        assert path.exists()
+
+        # Restore into fresh context
+        fresh_ctx = SessionContext(test_config)
+        loaded = session_mgr.load_latest()
+        assert loaded is not None
+        session_mgr.restore(loaded, fresh_ctx)
+
+        assert fresh_ctx.message_count == ctx.message_count
+        original_history = ctx.get_history()
+        restored_history = fresh_ctx.get_history()
+        for orig, restored in zip(original_history, restored_history):
+            assert orig.role == restored.role
+            assert orig.content == restored.content
+
+    @pytest.mark.asyncio
+    async def test_retry_on_transient_error(self, test_config: AppConfig, tmp_workspace: Path) -> None:
+        """Agent recovers from transient LLM errors via retry."""
+        from app.services.llm import LLMClient
+
+        ctx = SessionContext(test_config)
+
+        # Create a real client but mock stream_chat to fail then succeed
+        client = LLMClient(test_config.llm)
+        call_count = 0
+
+        async def flaky_stream(*args, **kwargs):
+            nonlocal call_count
+            call_count += 1
+            if call_count == 1:
+                raise LLMConnectionError("Temporary failure")
+            for chunk in make_text_chunks("Recovered!"):
+                yield chunk
+
+        client.stream_chat = flaky_stream  # type: ignore[assignment]
+
+        handler = StreamHandler(test_config.display)
+        registry = create_default_registry(test_config.agent.workspace_root, test_config)
+        permissions = PermissionsService(test_config.permissions)
+        agent = AgentLoop(test_config, ctx, client, handler, registry, permissions)
+
+        with patch("app.services.llm.asyncio.sleep", new_callable=AsyncMock):
+            await agent.run_turn("Test retry")
+
+        history = ctx.get_history()
+        # Should have succeeded: user + assistant
+        assert len(history) == 2
+        assert history[1].content == "Recovered!"
+        assert call_count == 2
+
+        await client.close()
--- a/tests/unit/test_agent_loop.py
+++ b/tests/unit/test_agent_loop.py
@@ -62,6 +62,7 @@ def handler() -> MagicMock:
    mock.usage = None
    mock.had_reasoning_only = False
    mock.reset = MagicMock()
+    mock.get_partial_message = MagicMock(return_value=None)
    return mock


--- a/tests/unit/test_retry.py
+++ b/tests/unit/test_retry.py
@@ -0,0 +1,125 @@
+"""Unit tests for LLM retry with exponential backoff."""
+
+from unittest.mock import AsyncMock, patch
+
+import pytest
+
+from app.models.config import LLMConfig
+from app.models.message import Message
+from app.services.llm import LLMClient, LLMConnectionError, LLMResponseError
+
+
+@pytest.fixture
+def llm_config() -> LLMConfig:
+    return LLMConfig(
+        model="test-model",
+        endpoint="http://localhost:11434",
+        max_retries=3,
+        retry_backoff_base=0.01,
+        retry_backoff_max=0.05,
+    )
+
+
+@pytest.fixture
+def client(llm_config: LLMConfig) -> LLMClient:
+    return LLMClient(llm_config)
+
+
+@pytest.fixture
+def messages() -> list[Message]:
+    return [Message(role="user", content="Hello")]
+
+
+class TestRetry:
+    @pytest.mark.asyncio
+    async def test_succeeds_without_retry(self, client: LLMClient, messages: list[Message]) -> None:
+        """Successful stream doesn't retry."""
+        call_count = 0
+
+        async def fake_stream(*args, **kwargs):
+            nonlocal call_count
+            call_count += 1
+            yield {"choices": [{"delta": {"content": "Hi"}}]}
+
+        client.stream_chat = fake_stream  # type: ignore[assignment]
+
+        collected = []
+        async for chunk in client.stream_chat_with_retry(messages):
+            collected.append(chunk)
+
+        assert len(collected) == 1
+        assert call_count == 1
+
+    @pytest.mark.asyncio
+    async def test_retries_on_connection_error(self, client: LLMClient, messages: list[Message]) -> None:
+        """Retries on LLMConnectionError, then succeeds."""
+        call_count = 0
+
+        async def flaky_stream(*args, **kwargs):
+            nonlocal call_count
+            call_count += 1
+            if call_count < 3:
+                raise LLMConnectionError("Connection refused")
+            yield {"choices": [{"delta": {"content": "OK"}}]}
+
+        client.stream_chat = flaky_stream  # type: ignore[assignment]
+
+        with patch("app.services.llm.asyncio.sleep", new_callable=AsyncMock):
+            collected = []
+            async for chunk in client.stream_chat_with_retry(messages):
+                collected.append(chunk)
+
+        assert len(collected) == 1
+        assert call_count == 3
+
+    @pytest.mark.asyncio
+    async def test_retries_on_5xx(self, client: LLMClient, messages: list[Message]) -> None:
+        """Retries on 5xx LLMResponseError."""
+        call_count = 0
+
+        async def server_error_stream(*args, **kwargs):
+            nonlocal call_count
+            call_count += 1
+            if call_count < 2:
+                raise LLMResponseError("Internal Server Error", status_code=500)
+            yield {"choices": [{"delta": {"content": "OK"}}]}
+
+        client.stream_chat = server_error_stream  # type: ignore[assignment]
+
+        with patch("app.services.llm.asyncio.sleep", new_callable=AsyncMock):
+            collected = []
+            async for chunk in client.stream_chat_with_retry(messages):
+                collected.append(chunk)
+
+        assert len(collected) == 1
+        assert call_count == 2
+
+    @pytest.mark.asyncio
+    async def test_no_retry_on_4xx(self, client: LLMClient, messages: list[Message]) -> None:
+        """Does NOT retry on 4xx errors — raises immediately."""
+        async def bad_request_stream(*args, **kwargs):
+            raise LLMResponseError("Bad Request", status_code=400)
+            yield  # pragma: no cover — make this an async generator
+
+        client.stream_chat = bad_request_stream  # type: ignore[assignment]
+
+        with pytest.raises(LLMResponseError, match="Bad Request"):
+            async for _ in client.stream_chat_with_retry(messages):
+                pass  # pragma: no cover
+
+    @pytest.mark.asyncio
+    async def test_respects_max_retries(self, client: LLMClient, messages: list[Message]) -> None:
+        """After exhausting retries, re-raises the last exception."""
+        async def always_fail(*args, **kwargs):
+            raise LLMConnectionError("Down forever")
+            yield  # pragma: no cover
+
+        client.stream_chat = always_fail  # type: ignore[assignment]
+
+        with patch("app.services.llm.asyncio.sleep", new_callable=AsyncMock) as mock_sleep:
+            with pytest.raises(LLMConnectionError, match="Down forever"):
+                async for _ in client.stream_chat_with_retry(messages):
+                    pass  # pragma: no cover
+
+        # Should have slept max_retries times (3 retries after initial attempt)
+        assert mock_sleep.call_count == 3
--- a/tests/unit/test_session.py
+++ b/tests/unit/test_session.py
@@ -0,0 +1,122 @@
+"""Unit tests for session persistence."""
+
+import json
+import time
+from pathlib import Path
+
+import pytest
+
+from app.agent.context import SessionContext
+from app.models.config import AgentConfig, AppConfig, LLMConfig, SessionConfig
+from app.services.session import SessionManager
+
+
+@pytest.fixture
+def tmp_workspace(tmp_path: Path) -> Path:
+    return tmp_path / "workspace"
+
+
+@pytest.fixture
+def session_config(tmp_workspace: Path) -> SessionConfig:
+    return SessionConfig(
+        session_dir=tmp_workspace / ".sneakycode" / "sessions",
+        auto_save=True,
+        max_session_age_hours=72,
+        offer_resume=True,
+    )
+
+
+@pytest.fixture
+def config(tmp_workspace: Path) -> AppConfig:
+    return AppConfig(
+        llm=LLMConfig(model="test-model", endpoint="http://localhost:11434"),
+        agent=AgentConfig(workspace_root=tmp_workspace),
+    )
+
+
+@pytest.fixture
+def ctx(config: AppConfig) -> SessionContext:
+    return SessionContext(config)
+
+
+@pytest.fixture
+def session_mgr(session_config: SessionConfig, tmp_workspace: Path) -> SessionManager:
+    tmp_workspace.mkdir(parents=True, exist_ok=True)
+    return SessionManager(session_config, tmp_workspace, "test-model")
+
+
+class TestSessionPersistence:
+    def test_save_creates_file(self, session_mgr: SessionManager, ctx: SessionContext, session_config: SessionConfig) -> None:
+        """Saving a session creates a JSON file in the session directory."""
+        ctx.add_message("user", "Hello")
+        ctx.add_message("assistant", "Hi there!")
+
+        path = session_mgr.save(ctx)
+
+        assert path.exists()
+        assert path.suffix == ".json"
+
+        data = json.loads(path.read_text())
+        assert data["model"] == "test-model"
+        assert len(data["messages"]) == 2
+
+    def test_load_latest_returns_newest(self, session_mgr: SessionManager, ctx: SessionContext, session_config: SessionConfig) -> None:
+        """load_latest returns the most recently modified session."""
+        ctx.add_message("user", "First session")
+        session_mgr.save(ctx)
+
+        # Create a second session manager (simulates a new startup)
+        mgr2 = SessionManager(session_config, session_mgr._workspace_root, "test-model")
+        ctx.add_message("assistant", "Second response")
+        path2 = mgr2.save(ctx)
+
+        loaded = mgr2.load_latest()
+        assert loaded is not None
+        assert loaded.session_id == mgr2._session_id
+        assert len(loaded.messages) == 2
+
+    def test_restore_populates_context(self, session_mgr: SessionManager, ctx: SessionContext, config: AppConfig) -> None:
+        """Restoring a session populates the context with saved messages."""
+        ctx.add_message("user", "Hello")
+        ctx.add_message("assistant", "World")
+        session_mgr.save(ctx)
+
+        # Load and restore into a fresh context
+        fresh_ctx = SessionContext(config)
+        loaded = session_mgr.load_latest()
+        assert loaded is not None
+
+        session_mgr.restore(loaded, fresh_ctx)
+
+        history = fresh_ctx.get_history()
+        assert len(history) == 2
+        assert history[0].role == "user"
+        assert history[0].content == "Hello"
+        assert history[1].role == "assistant"
+        assert history[1].content == "World"
+
+    def test_cleanup_removes_old_files(self, session_config: SessionConfig, tmp_workspace: Path) -> None:
+        """cleanup_old deletes files older than max_session_age_hours."""
+        tmp_workspace.mkdir(parents=True, exist_ok=True)
+
+        # Create session config with very short max age
+        short_config = SessionConfig(
+            session_dir=session_config.session_dir,
+            max_session_age_hours=0,  # 0 hours = everything is old
+        )
+        mgr = SessionManager(short_config, tmp_workspace, "test-model")
+
+        # Create a session file manually with old timestamp
+        session_dir = tmp_workspace / short_config.session_dir
+        session_dir.mkdir(parents=True, exist_ok=True)
+        old_file = session_dir / "old_session.json"
+        old_file.write_text('{"version": 1}')
+
+        # Set mtime to the past
+        import os
+        old_time = time.time() - 3600  # 1 hour ago
+        os.utime(old_file, (old_time, old_time))
+
+        deleted = mgr.cleanup_old()
+        assert deleted == 1
+        assert not old_file.exists()
--- a/tests/unit/test_truncation.py
+++ b/tests/unit/test_truncation.py
@@ -0,0 +1,128 @@
+"""Unit tests for conversation truncation logic."""
+
+from pathlib import Path
+
+import pytest
+
+from app.agent.context import SessionContext
+from app.models.config import AgentConfig, AppConfig, LLMConfig
+
+
+@pytest.fixture
+def config() -> AppConfig:
+    return AppConfig(
+        llm=LLMConfig(model="test-model", endpoint="http://localhost:11434"),
+        agent=AgentConfig(
+            max_conversation_tokens=200,
+            truncation_keep_recent=3,
+            truncation_threshold=0.85,
+            workspace_root=Path("/tmp/test"),
+        ),
+    )
+
+
+@pytest.fixture
+def ctx(config: AppConfig) -> SessionContext:
+    return SessionContext(config)
+
+
+class TestTruncation:
+    def test_no_truncation_under_threshold(self, ctx: SessionContext) -> None:
+        """No messages dropped when under threshold."""
+        ctx.add_message("user", "Hello")
+        ctx.add_message("assistant", "Hi there!")
+
+        dropped = ctx.truncate_history()
+        assert dropped == 0
+        assert ctx.message_count == 2
+
+    def test_drops_oldest_messages(self, ctx: SessionContext) -> None:
+        """Drops middle messages when over budget."""
+        # Fill with enough content to exceed the small 200-token budget
+        ctx.add_message("user", "First message " * 20)
+        for i in range(8):
+            ctx.add_message("assistant", f"Response {i} " * 15)
+            ctx.add_message("user", f"Follow-up {i} " * 15)
+
+        # Force the token counter to report over budget
+        from app.utils.token_counter import TokenUsage
+        ctx.token_counter.count_usage(TokenUsage(total_tokens=200))
+
+        original_count = len(ctx.get_history())
+        dropped = ctx.truncate_history()
+
+        assert dropped > 0
+        assert len(ctx.get_history()) < original_count
+
+    def test_preserves_recent_messages(self, ctx: SessionContext) -> None:
+        """The most recent N messages are always preserved."""
+        ctx.add_message("user", "First message " * 20)
+        for i in range(10):
+            ctx.add_message("assistant", f"Response {i} " * 10)
+            ctx.add_message("user", f"Follow-up {i} " * 10)
+
+        from app.utils.token_counter import TokenUsage
+        ctx.token_counter.count_usage(TokenUsage(total_tokens=200))
+
+        history_before = ctx.get_history()
+        recent_before = history_before[-3:]  # keep_recent=3
+
+        ctx.truncate_history()
+
+        history_after = ctx.get_history()
+        recent_after = history_after[-3:]
+
+        # Recent messages should be preserved
+        for before, after in zip(recent_before, recent_after):
+            assert before.content == after.content
+
+    def test_preserves_first_user_message(self, ctx: SessionContext) -> None:
+        """First user message is always kept."""
+        first_content = "This is the very first user message"
+        ctx.add_message("user", first_content)
+        for i in range(10):
+            ctx.add_message("assistant", f"Response {i} " * 10)
+            ctx.add_message("user", f"Follow-up {i} " * 10)
+
+        from app.utils.token_counter import TokenUsage
+        ctx.token_counter.count_usage(TokenUsage(total_tokens=200))
+
+        ctx.truncate_history()
+
+        history = ctx.get_history()
+        assert history[0].role == "user"
+        assert history[0].content == first_content
+
+    def test_orphaned_tool_messages_cleaned(self, ctx: SessionContext) -> None:
+        """Tool messages without matching tool_call are cleaned up."""
+        from app.models.tool_call import ToolCall, ToolCallFunction
+
+        ctx.add_message("user", "Do something " * 20)
+        # Assistant with tool call
+        ctx.add_message(
+            "assistant",
+            None,
+            tool_calls=[ToolCall(id="tc_1", type="function", function=ToolCallFunction(name="read_file", arguments='{"path": "x"}'))],
+        )
+        # Tool result for tc_1
+        ctx.add_message("tool", "file contents " * 20, tool_call_id="tc_1", name="read_file")
+        # More padding to push over budget
+        for i in range(8):
+            ctx.add_message("assistant", f"Analysis {i} " * 15)
+            ctx.add_message("user", f"Next {i} " * 15)
+
+        from app.utils.token_counter import TokenUsage
+        ctx.token_counter.count_usage(TokenUsage(total_tokens=200))
+
+        ctx.truncate_history()
+
+        history = ctx.get_history()
+        # If the assistant message with tc_1 was dropped, the orphaned tool message should also be gone
+        has_tc1_assistant = any(
+            m.role == "assistant" and m.tool_calls and any(tc.id == "tc_1" for tc in m.tool_calls)
+            for m in history
+        )
+        has_tc1_tool = any(m.role == "tool" and m.tool_call_id == "tc_1" for m in history)
+
+        # Either both exist or neither exists
+        assert has_tc1_assistant == has_tc1_tool