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>
This commit is contained in:
150
tests/integration/conftest.py
Normal file
150
tests/integration/conftest.py
Normal 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
|
||||
155
tests/integration/test_agent_workflows.py
Normal file
155
tests/integration/test_agent_workflows.py
Normal 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()
|
||||
@@ -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
125
tests/unit/test_retry.py
Normal 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
122
tests/unit/test_session.py
Normal 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()
|
||||
128
tests/unit/test_truncation.py
Normal file
128
tests/unit/test_truncation.py
Normal 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
|
||||
Reference in New Issue
Block a user