wuyang6
416 lines
16 KiB
Python
416 lines
16 KiB
Python
"""Tests for src/compact.py – the conversation compaction service."""
|
||
|
||
from __future__ import annotations
|
||
|
||
import tempfile
|
||
import unittest
|
||
from dataclasses import replace
|
||
from pathlib import Path
|
||
from unittest.mock import MagicMock
|
||
|
||
from src.agent_runtime import LocalCodingAgent
|
||
from src.agent_session import AgentMessage, AgentSessionState
|
||
from src.agent_types import AgentRuntimeConfig, ModelConfig
|
||
from src.compact import (
|
||
AUTOCOMPACT_BUFFER_TOKENS,
|
||
ERROR_INCOMPLETE_RESPONSE,
|
||
ERROR_NOT_ENOUGH_MESSAGES,
|
||
CompactionResult,
|
||
compact_conversation,
|
||
format_compact_summary,
|
||
get_compact_prompt,
|
||
get_compact_user_summary_message,
|
||
)
|
||
|
||
|
||
class TestGetCompactPrompt(unittest.TestCase):
|
||
"""Tests for the compact prompt builder."""
|
||
|
||
def test_basic_prompt_contains_all_nine_sections(self) -> None:
|
||
prompt = get_compact_prompt()
|
||
for section in [
|
||
'1. Primary Request and Intent',
|
||
'2. Key Technical Concepts',
|
||
'3. Files and Code Sections',
|
||
'4. Errors and fixes',
|
||
'5. Problem Solving',
|
||
'6. All user messages',
|
||
'7. Pending Tasks',
|
||
'8. Current Work',
|
||
'9. Optional Next Step',
|
||
]:
|
||
self.assertIn(section, prompt, f'Missing section: {section}')
|
||
|
||
def test_no_tools_preamble_present(self) -> None:
|
||
prompt = get_compact_prompt()
|
||
self.assertIn('CRITICAL: Respond with TEXT ONLY', prompt)
|
||
self.assertIn('Do NOT call any tools', prompt)
|
||
|
||
def test_no_tools_trailer_present(self) -> None:
|
||
prompt = get_compact_prompt()
|
||
self.assertIn('REMINDER: Do NOT call any tools', prompt)
|
||
|
||
def test_analysis_and_summary_example_tags_present(self) -> None:
|
||
prompt = get_compact_prompt()
|
||
self.assertIn('<analysis>', prompt)
|
||
self.assertIn('</analysis>', prompt)
|
||
self.assertIn('<summary>', prompt)
|
||
self.assertIn('</summary>', prompt)
|
||
|
||
def test_custom_instructions_appended(self) -> None:
|
||
prompt = get_compact_prompt('Focus on database changes.')
|
||
self.assertIn('Additional Instructions:', prompt)
|
||
self.assertIn('Focus on database changes.', prompt)
|
||
|
||
def test_empty_custom_instructions_ignored(self) -> None:
|
||
prompt_no_custom = get_compact_prompt()
|
||
prompt_empty = get_compact_prompt('')
|
||
prompt_whitespace = get_compact_prompt(' ')
|
||
self.assertEqual(prompt_no_custom, prompt_empty)
|
||
self.assertEqual(prompt_no_custom, prompt_whitespace)
|
||
|
||
def test_none_custom_instructions_ignored(self) -> None:
|
||
prompt_none = get_compact_prompt(None)
|
||
prompt_no_arg = get_compact_prompt()
|
||
self.assertEqual(prompt_none, prompt_no_arg)
|
||
|
||
|
||
class TestFormatCompactSummary(unittest.TestCase):
|
||
"""Tests for the summary formatting / XML stripping."""
|
||
|
||
def test_strips_analysis_block(self) -> None:
|
||
raw = '<analysis>thinking here</analysis>\n\n<summary>result</summary>'
|
||
formatted = format_compact_summary(raw)
|
||
self.assertNotIn('<analysis>', formatted)
|
||
self.assertNotIn('thinking here', formatted)
|
||
self.assertIn('result', formatted)
|
||
|
||
def test_unwraps_summary_tags(self) -> None:
|
||
raw = '<summary>The main points.\n1. First</summary>'
|
||
formatted = format_compact_summary(raw)
|
||
self.assertNotIn('<summary>', formatted)
|
||
self.assertNotIn('</summary>', formatted)
|
||
self.assertIn('Summary:', formatted)
|
||
self.assertIn('The main points.', formatted)
|
||
|
||
def test_handles_no_xml_tags(self) -> None:
|
||
raw = 'Plain text summary without any tags.'
|
||
formatted = format_compact_summary(raw)
|
||
self.assertEqual(formatted, raw)
|
||
|
||
def test_collapses_excess_blank_lines(self) -> None:
|
||
raw = '<analysis>x</analysis>\n\n\n\n<summary>y</summary>'
|
||
formatted = format_compact_summary(raw)
|
||
self.assertNotIn('\n\n\n', formatted)
|
||
|
||
def test_multiline_analysis_stripped(self) -> None:
|
||
raw = (
|
||
'<analysis>\nLine 1\nLine 2\nLine 3\n</analysis>\n'
|
||
'<summary>Final summary</summary>'
|
||
)
|
||
formatted = format_compact_summary(raw)
|
||
self.assertNotIn('Line 1', formatted)
|
||
self.assertIn('Final summary', formatted)
|
||
|
||
def test_summary_content_keeps_regex_backslashes_literal(self) -> None:
|
||
raw = r'<summary>Regex examples: \s+ and \1 should stay literal.</summary>'
|
||
formatted = format_compact_summary(raw)
|
||
self.assertIn(r'\s+', formatted)
|
||
self.assertIn(r'\1', formatted)
|
||
|
||
|
||
class TestGetCompactUserSummaryMessage(unittest.TestCase):
|
||
"""Tests for the post-compact user message builder."""
|
||
|
||
def test_basic_message_structure(self) -> None:
|
||
msg = get_compact_user_summary_message('<summary>overview</summary>')
|
||
self.assertIn('continued from a previous conversation', msg)
|
||
self.assertIn('overview', msg)
|
||
|
||
def test_transcript_path_appended(self) -> None:
|
||
msg = get_compact_user_summary_message(
|
||
'<summary>ok</summary>',
|
||
transcript_path='/tmp/transcript.json',
|
||
)
|
||
self.assertIn('/tmp/transcript.json', msg)
|
||
|
||
def test_suppress_follow_up(self) -> None:
|
||
msg = get_compact_user_summary_message(
|
||
'<summary>ok</summary>',
|
||
suppress_follow_up=True,
|
||
)
|
||
self.assertIn('without asking the user any further questions', msg)
|
||
|
||
def test_no_suppress_follow_up_default(self) -> None:
|
||
msg = get_compact_user_summary_message('<summary>ok</summary>')
|
||
self.assertNotIn('without asking the user any further questions', msg)
|
||
|
||
|
||
class TestCompactConversation(unittest.TestCase):
|
||
"""Tests for the core compact_conversation() function."""
|
||
|
||
def _make_agent(self, tmp_dir: str) -> LocalCodingAgent:
|
||
"""Create a minimal agent with a session loaded."""
|
||
agent = LocalCodingAgent(
|
||
model_config=ModelConfig(model='test-model'),
|
||
runtime_config=AgentRuntimeConfig(cwd=Path(tmp_dir)),
|
||
)
|
||
return agent
|
||
|
||
def _set_session(
|
||
self, agent: LocalCodingAgent, messages: list[AgentMessage]
|
||
) -> None:
|
||
session = AgentSessionState(
|
||
system_prompt_parts=('You are a helpful assistant.',),
|
||
messages=messages,
|
||
)
|
||
agent.last_session = session
|
||
|
||
def _make_messages(self, count: int) -> list[AgentMessage]:
|
||
msgs: list[AgentMessage] = []
|
||
for i in range(count):
|
||
role = 'user' if i % 2 == 0 else 'assistant'
|
||
msgs.append(
|
||
AgentMessage(
|
||
role=role,
|
||
content=f'Message {i} content. ' * 10,
|
||
message_id=f'msg_{i}',
|
||
)
|
||
)
|
||
return msgs
|
||
|
||
def test_no_session_returns_error(self) -> None:
|
||
with tempfile.TemporaryDirectory() as tmp_dir:
|
||
agent = self._make_agent(tmp_dir)
|
||
agent.last_session = None
|
||
result = compact_conversation(agent)
|
||
self.assertIsNotNone(result.error)
|
||
self.assertIn('Not enough', result.error)
|
||
|
||
def test_empty_session_returns_error(self) -> None:
|
||
with tempfile.TemporaryDirectory() as tmp_dir:
|
||
agent = self._make_agent(tmp_dir)
|
||
self._set_session(agent, [])
|
||
result = compact_conversation(agent)
|
||
self.assertIsNotNone(result.error)
|
||
|
||
def test_too_few_messages_returns_error(self) -> None:
|
||
with tempfile.TemporaryDirectory() as tmp_dir:
|
||
agent = self._make_agent(tmp_dir)
|
||
# With only 2 messages and preserve_count=4, nothing to compact
|
||
self._set_session(agent, self._make_messages(2))
|
||
result = compact_conversation(agent)
|
||
self.assertIsNotNone(result.error)
|
||
|
||
def test_successful_compaction(self) -> None:
|
||
"""Simulate a successful model call and verify session is compacted."""
|
||
with tempfile.TemporaryDirectory() as tmp_dir:
|
||
agent = self._make_agent(tmp_dir)
|
||
msgs = self._make_messages(10)
|
||
self._set_session(agent, msgs)
|
||
|
||
# Mock the client's complete method
|
||
from src.openai_compat import AssistantTurn
|
||
from src.agent_types import UsageStats
|
||
|
||
mock_turn = AssistantTurn(
|
||
content=(
|
||
'<analysis>Thinking through the conversation...</analysis>\n'
|
||
'<summary>\n1. Primary Request and Intent:\n'
|
||
' User wanted to test compaction.\n'
|
||
'2. Key Technical Concepts:\n - Testing\n'
|
||
'3. Files and Code Sections:\n - test.py\n'
|
||
'4. Errors and fixes:\n - None\n'
|
||
'5. Problem Solving:\n Basic testing.\n'
|
||
'6. All user messages:\n - "test compaction"\n'
|
||
'7. Pending Tasks:\n - None\n'
|
||
'8. Current Work:\n Testing compaction.\n'
|
||
'9. Optional Next Step:\n Verify it works.\n'
|
||
'</summary>'
|
||
),
|
||
tool_calls=(),
|
||
finish_reason='stop',
|
||
raw_message={},
|
||
usage=UsageStats(),
|
||
)
|
||
agent.client = MagicMock()
|
||
agent.client.complete.return_value = mock_turn
|
||
|
||
result = compact_conversation(agent)
|
||
|
||
self.assertIsNone(result.error)
|
||
self.assertGreater(result.pre_compact_token_count, 0)
|
||
# Session should have fewer messages than original 10
|
||
self.assertLess(
|
||
len(agent.last_session.messages), 10,
|
||
'Session should have fewer messages after compaction',
|
||
)
|
||
# Should contain a compact_boundary message
|
||
boundary_msgs = [
|
||
m for m in agent.last_session.messages
|
||
if m.metadata.get('kind') == 'compact_boundary'
|
||
]
|
||
self.assertEqual(len(boundary_msgs), 1)
|
||
# Should contain a compact_summary message
|
||
summary_msgs = [
|
||
m for m in agent.last_session.messages
|
||
if m.metadata.get('kind') == 'compact_summary'
|
||
]
|
||
self.assertEqual(len(summary_msgs), 1)
|
||
# Summary should not contain <analysis> block
|
||
self.assertNotIn('<analysis>', result.summary_text)
|
||
# Summary should contain the actual summary content
|
||
self.assertIn('User wanted to test compaction', result.summary_text)
|
||
# UI display transcript must remain append-only and not be replaced by
|
||
# compact_summary / compact_boundary model-context messages.
|
||
display_contents = [m.content for m in agent.last_session.display_messages]
|
||
self.assertEqual(display_contents, [m.content for m in msgs])
|
||
|
||
def test_api_error_returns_compaction_error(self) -> None:
|
||
with tempfile.TemporaryDirectory() as tmp_dir:
|
||
agent = self._make_agent(tmp_dir)
|
||
self._set_session(agent, self._make_messages(10))
|
||
|
||
agent.client = MagicMock()
|
||
agent.client.complete.side_effect = RuntimeError('API down')
|
||
|
||
result = compact_conversation(agent)
|
||
self.assertIsNotNone(result.error)
|
||
self.assertIn('API down', result.error)
|
||
|
||
def test_empty_model_response_returns_error(self) -> None:
|
||
with tempfile.TemporaryDirectory() as tmp_dir:
|
||
agent = self._make_agent(tmp_dir)
|
||
self._set_session(agent, self._make_messages(10))
|
||
|
||
from src.openai_compat import AssistantTurn
|
||
from src.agent_types import UsageStats
|
||
|
||
agent.client = MagicMock()
|
||
agent.client.complete.return_value = AssistantTurn(
|
||
content='',
|
||
tool_calls=(),
|
||
finish_reason='stop',
|
||
raw_message={},
|
||
usage=UsageStats(),
|
||
)
|
||
|
||
result = compact_conversation(agent)
|
||
self.assertIsNotNone(result.error)
|
||
|
||
def test_preserves_tail_messages(self) -> None:
|
||
"""The most recent messages should survive compaction."""
|
||
with tempfile.TemporaryDirectory() as tmp_dir:
|
||
agent = self._make_agent(tmp_dir)
|
||
msgs = self._make_messages(12)
|
||
self._set_session(agent, msgs)
|
||
|
||
from src.openai_compat import AssistantTurn
|
||
from src.agent_types import UsageStats
|
||
|
||
agent.client = MagicMock()
|
||
agent.client.complete.return_value = AssistantTurn(
|
||
content='<summary>Summarised.</summary>',
|
||
tool_calls=(),
|
||
finish_reason='stop',
|
||
raw_message={},
|
||
usage=UsageStats(),
|
||
)
|
||
|
||
result = compact_conversation(agent)
|
||
|
||
self.assertIsNone(result.error)
|
||
# The last 4 messages (default preserve_count) should still be present
|
||
session_contents = [m.content for m in agent.last_session.messages]
|
||
for original_msg in msgs[-4:]:
|
||
self.assertIn(
|
||
original_msg.content,
|
||
session_contents,
|
||
f'Tail message "{original_msg.message_id}" should be preserved',
|
||
)
|
||
|
||
def test_custom_instructions_passed_to_prompt(self) -> None:
|
||
with tempfile.TemporaryDirectory() as tmp_dir:
|
||
agent = self._make_agent(tmp_dir)
|
||
self._set_session(agent, self._make_messages(10))
|
||
|
||
from src.openai_compat import AssistantTurn
|
||
from src.agent_types import UsageStats
|
||
|
||
agent.client = MagicMock()
|
||
agent.client.complete.return_value = AssistantTurn(
|
||
content='<summary>Custom summary.</summary>',
|
||
tool_calls=(),
|
||
finish_reason='stop',
|
||
raw_message={},
|
||
usage=UsageStats(),
|
||
)
|
||
|
||
compact_conversation(agent, custom_instructions='Focus on CSS.')
|
||
|
||
# Check that the API was called with custom instructions in the prompt
|
||
call_args = agent.client.complete.call_args
|
||
messages = call_args[0][0]
|
||
last_user_msg = messages[-1]['content']
|
||
self.assertIn('Focus on CSS.', last_user_msg)
|
||
self.assertIn('Additional Instructions:', last_user_msg)
|
||
|
||
|
||
class TestCompactSlashCommand(unittest.TestCase):
|
||
"""Test the /compact slash command handler end-to-end."""
|
||
|
||
def test_slash_compact_returns_success_message(self) -> None:
|
||
with tempfile.TemporaryDirectory() as tmp_dir:
|
||
agent = LocalCodingAgent(
|
||
model_config=ModelConfig(model='test-model'),
|
||
runtime_config=AgentRuntimeConfig(cwd=Path(tmp_dir)),
|
||
)
|
||
# Set up a session with enough messages
|
||
session = AgentSessionState(
|
||
system_prompt_parts=('You are a helper.',),
|
||
messages=[
|
||
AgentMessage(role='user', content='Hello', message_id=f'u{i}')
|
||
if i % 2 == 0
|
||
else AgentMessage(role='assistant', content='Hi', message_id=f'a{i}')
|
||
for i in range(10)
|
||
],
|
||
)
|
||
agent.last_session = session
|
||
|
||
from src.openai_compat import AssistantTurn
|
||
from src.agent_types import UsageStats
|
||
|
||
agent.client = MagicMock()
|
||
agent.client.complete.return_value = AssistantTurn(
|
||
content='<summary>Session summarised.</summary>',
|
||
tool_calls=(),
|
||
finish_reason='stop',
|
||
raw_message={},
|
||
usage=UsageStats(),
|
||
)
|
||
|
||
result = agent.run('/compact')
|
||
|
||
self.assertIn('Conversation compacted', result.final_output)
|
||
self.assertIn('Tokens before', result.final_output)
|
||
|
||
def test_slash_compact_no_session_returns_error(self) -> None:
|
||
with tempfile.TemporaryDirectory() as tmp_dir:
|
||
agent = LocalCodingAgent(
|
||
model_config=ModelConfig(model='test-model'),
|
||
runtime_config=AgentRuntimeConfig(cwd=Path(tmp_dir)),
|
||
)
|
||
result = agent.run('/compact')
|
||
self.assertIn('failed', result.final_output.lower())
|
||
|
||
|
||
class TestConstants(unittest.TestCase):
|
||
"""Verify key constants match the npm reference."""
|
||
|
||
def test_autocompact_buffer_tokens(self) -> None:
|
||
self.assertEqual(AUTOCOMPACT_BUFFER_TOKENS, 13_000)
|
||
|
||
|
||
if __name__ == '__main__':
|
||
unittest.main()
|