Merge branch 'feature/phase-7-polish-hardening'

This commit is contained in:
2026-03-11 10:21:16 -05:00
16 changed files with 1550 additions and 12 deletions

147
README.md Normal file
View File

@@ -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

View File

@@ -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)

View File

@@ -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:

View File

@@ -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__":

View File

@@ -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":

View File

@@ -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()

148
app/services/session.py Normal file
View File

@@ -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

View File

@@ -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."""

View File

@@ -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

240
docs/tools.md Normal file
View File

@@ -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

View File

@@ -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

View File

@@ -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()

View File

@@ -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

125
tests/unit/test_retry.py Normal file
View File

@@ -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

122
tests/unit/test_session.py Normal file
View File

@@ -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()

View File

@@ -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