"""Tests for AsMsgHandler.context_check method.""" import asyncio from agentscope.message import Msg from test_utils import get_token_counter from reme.core.utils import get_logger from reme.memory.file_based.utils import AsMsgHandler logger = get_logger() # ANSI color codes class Colors: """ANSI color codes for terminal output.""" GREEN = "\033[92m" RED = "\033[91m" YELLOW = "\033[93m" BLUE = "\033[94m" CYAN = "\033[96m" BOLD = "\033[1m" RESET = "\033[0m" def print_pass(test_name: str): """Print test passed message.""" print(f"{Colors.GREEN}{Colors.BOLD}✓ {test_name} PASSED{Colors.RESET}") def print_fail(test_name: str, error: str): """Print test failed message.""" print(f"{Colors.RED}{Colors.BOLD}✗ {test_name} FAILED: {error}{Colors.RESET}") def print_error(test_name: str, error: str): """Print test error message.""" print(f"{Colors.YELLOW}{Colors.BOLD}⚠ {test_name} ERROR: {error}{Colors.RESET}") def print_test_header(test_name: str): """Print test header.""" print(f"\n{Colors.CYAN}{'=' * 60}{Colors.RESET}") print(f"{Colors.BLUE}{Colors.BOLD}Running: {test_name}{Colors.RESET}") print(f"{Colors.CYAN}{'=' * 60}{Colors.RESET}") def create_handler() -> AsMsgHandler: """Create an AsMsgHandler instance for testing.""" return AsMsgHandler(token_counter=get_token_counter()) def verify_context_check_invariants( handler: AsMsgHandler, messages: list[Msg], to_compact: list[Msg], to_keep: list[Msg], memory_compact_threshold: int, memory_compact_reserve: int, test_name: str, ): """Verify that context_check results satisfy all invariants. This function checks: 1. Threshold requirement: If total tokens <= threshold, no compaction should occur 2. Reserve requirement: Kept messages' total tokens should not exceed reserve 3. Order requirement: Both to_compact and to_keep should preserve original order Args: handler: The AsMsgHandler instance messages: Original messages list to_compact: Messages to compact returned by context_check to_keep: Messages to keep returned by context_check memory_compact_threshold: The threshold parameter used memory_compact_reserve: The reserve parameter used test_name: Name of the test for error reporting Raises: AssertionError: If any invariant is violated """ # Calculate total tokens of original messages total_tokens = sum(asyncio.run(handler.stat_message(m)).total_tokens for m in messages) # 1. Threshold requirement check if total_tokens <= memory_compact_threshold: assert len(to_compact) == 0, ( f"[{test_name}] Threshold violation: total_tokens ({total_tokens}) <= " f"threshold ({memory_compact_threshold}), but to_compact is not empty " f"(has {len(to_compact)} messages)" ) assert to_keep == messages, ( f"[{test_name}] Threshold violation: total_tokens ({total_tokens}) <= " f"threshold ({memory_compact_threshold}), but to_keep differs from original messages" ) # 2. Reserve requirement check kept_tokens = sum(asyncio.run(handler.stat_message(m)).total_tokens for m in to_keep) assert kept_tokens <= memory_compact_reserve or len(to_keep) == 0, ( f"[{test_name}] Reserve violation: kept_tokens ({kept_tokens}) > " f"reserve ({memory_compact_reserve})" ) # 3. Order requirement check - both lists should preserve original order # Create a mapping of message id to original index msg_to_idx = {id(m): i for i, m in enumerate(messages)} # Check to_compact order compact_indices = [msg_to_idx.get(id(m), -1) for m in to_compact] for i in range(len(compact_indices) - 1): assert compact_indices[i] < compact_indices[i + 1], ( f"[{test_name}] Order violation in to_compact: message at original index " f"{compact_indices[i]} appears before message at index {compact_indices[i + 1]}" ) # Check to_keep order keep_indices = [msg_to_idx.get(id(m), -1) for m in to_keep] for i in range(len(keep_indices) - 1): assert keep_indices[i] < keep_indices[i + 1], ( f"[{test_name}] Order violation in to_keep: message at original index " f"{keep_indices[i]} appears before message at index {keep_indices[i + 1]}" ) # 4. Additional check: to_compact indices should all be less than to_keep indices # (compact messages come from the beginning, keep messages come from the end) if to_compact and to_keep: max_compact_idx = max(compact_indices) if compact_indices else -1 min_keep_idx = min(keep_indices) if keep_indices else len(messages) assert max_compact_idx < min_keep_idx, ( f"[{test_name}] Partition violation: max compact index ({max_compact_idx}) >= " f"min keep index ({min_keep_idx}). Compact and keep should be a clean partition." ) # 5. Check that all messages are accounted for (no duplicates, no missing) assert len(to_compact) + len(to_keep) == len(messages), ( f"[{test_name}] Count mismatch: to_compact ({len(to_compact)}) + " f"to_keep ({len(to_keep)}) != original ({len(messages)})" ) all_returned = set(id(m) for m in to_compact) | set(id(m) for m in to_keep) all_original = set(id(m) for m in messages) assert all_returned == all_original, f"[{test_name}] Message set mismatch: returned messages differ from original" def create_user_msg(content: str) -> Msg: """Create a user message.""" return Msg(name="user", role="user", content=content) def create_assistant_msg(content: str) -> Msg: """Create an assistant message.""" return Msg(name="assistant", role="assistant", content=content) def create_tool_use_msg(tool_id: str, tool_name: str, tool_input: dict) -> Msg: """Create a message with tool_use content block.""" return Msg( name="assistant", role="assistant", content=[ { "type": "tool_use", "id": tool_id, "name": tool_name, "input": tool_input, }, ], ) def create_tool_result_msg(tool_id: str, tool_name: str, output: str) -> Msg: """Create a message with tool_result content block.""" return Msg( name="tool", role="user", content=[ { "type": "tool_result", "id": tool_id, "name": tool_name, "output": output, }, ], ) def create_mixed_tool_msg( tool_use_id: str, tool_use_name: str, tool_use_input: dict, tool_result_id: str, tool_result_name: str, tool_result_output: str, ) -> Msg: """Create a message with both tool_use and tool_result blocks.""" return Msg( name="assistant", role="assistant", content=[ { "type": "tool_use", "id": tool_use_id, "name": tool_use_name, "input": tool_use_input, }, { "type": "tool_result", "id": tool_result_id, "name": tool_result_name, "output": tool_result_output, }, ], ) # ============================================================================= # Normal Cases # ============================================================================= def test_empty_messages(): """Test context_check with empty messages list.""" handler = create_handler() messages = [] threshold, reserve = 1000, 500 to_compact, to_keep, _ = asyncio.run( handler.context_check( messages=messages, memory_compact_threshold=threshold, memory_compact_reserve=reserve, ), ) assert not to_compact, f"Expected empty compact list, got: {to_compact}" assert to_keep == [], f"Expected empty keep list, got: {to_keep}" verify_context_check_invariants(handler, messages, to_compact, to_keep, threshold, reserve, "test_empty_messages") print_pass("test_empty_messages") def test_below_threshold_returns_all(): """Test that messages below threshold are all kept.""" handler = create_handler() messages = [ create_user_msg("Hello"), create_assistant_msg("Hi there!"), create_user_msg("How are you?"), ] threshold, reserve = 10000, 5000 to_compact, to_keep, _ = asyncio.run( handler.context_check( messages=messages, memory_compact_threshold=threshold, # Very high threshold memory_compact_reserve=reserve, ), ) assert not to_compact, f"Expected empty compact list, got: {len(to_compact)}" assert len(to_keep) == 3, f"Expected 3 messages to keep, got: {len(to_keep)}" assert to_keep == messages, "Messages to keep should be the original messages" verify_context_check_invariants( handler, messages, to_compact, to_keep, threshold, reserve, "test_below_threshold_returns_all", ) print_pass("test_below_threshold_returns_all") def test_above_threshold_triggers_compaction(): """Test that messages above threshold are split correctly.""" handler = create_handler() # Create messages that will exceed threshold messages = [ create_user_msg("First message " * 100), create_assistant_msg("Second message " * 100), create_user_msg("Third message " * 100), create_assistant_msg("Fourth message " * 100), ] threshold, reserve = 100, 200 to_compact, to_keep, _ = asyncio.run( handler.context_check( messages=messages, memory_compact_threshold=threshold, # Low threshold to trigger compaction memory_compact_reserve=reserve, ), ) # Should have some messages compacted and some kept assert len(to_compact) + len(to_keep) == len(messages), "Total messages should match" assert len(to_compact) > 0, "Expected some messages to be compacted" verify_context_check_invariants( handler, messages, to_compact, to_keep, threshold, reserve, "test_above_threshold_triggers_compaction", ) print_pass("test_above_threshold_triggers_compaction") def test_message_order_preserved(): """Test that message order is preserved in both lists.""" handler = create_handler() messages = [ create_user_msg("First " * 50), create_assistant_msg("Second " * 50), create_user_msg("Third " * 50), create_assistant_msg("Fourth " * 50), create_user_msg("Fifth " * 10), ] threshold, reserve = 100, 150 to_compact, to_keep, _ = asyncio.run( handler.context_check( messages=messages, memory_compact_threshold=threshold, # Low threshold memory_compact_reserve=reserve, ), ) # Check order preservation - compact messages should appear first in original all_messages = to_compact + to_keep for i, msg in enumerate(all_messages): assert msg in messages, f"Message {i} not found in original messages" verify_context_check_invariants( handler, messages, to_compact, to_keep, threshold, reserve, "test_message_order_preserved", ) print_pass("test_message_order_preserved") # ============================================================================= # Edge Cases - Threshold and Reserve Boundaries # ============================================================================= def test_single_message_below_threshold(): """Test single message below threshold.""" handler = create_handler() messages = [create_user_msg("Short message")] threshold, reserve = 1000, 500 to_compact, to_keep, _ = asyncio.run( handler.context_check( messages=messages, memory_compact_threshold=threshold, memory_compact_reserve=reserve, ), ) assert not to_compact, "Should not compact single message below threshold" assert len(to_keep) == 1, "Should keep the single message" verify_context_check_invariants( handler, messages, to_compact, to_keep, threshold, reserve, "test_single_message_below_threshold", ) print_pass("test_single_message_below_threshold") def test_single_message_above_threshold(): """Test single message that exceeds threshold - nothing can be kept in reserve.""" handler = create_handler() long_content = "Very long message " * 1000 messages = [create_user_msg(long_content)] threshold, reserve = 10, 5 to_compact, to_keep, _ = asyncio.run( handler.context_check( messages=messages, memory_compact_threshold=threshold, # Very low threshold memory_compact_reserve=reserve, # Even lower reserve ), ) # Message exceeds both threshold and reserve, so it's compacted assert len(to_compact) == 1, "Single large message should be compacted" assert len(to_keep) == 0, "Nothing can fit in reserve" verify_context_check_invariants( handler, messages, to_compact, to_keep, threshold, reserve, "test_single_message_above_threshold", ) print_pass("test_single_message_above_threshold") def test_reserve_zero(): """Test with reserve=0, no messages can be kept.""" handler = create_handler() messages = [ create_user_msg("Hello"), create_assistant_msg("Hi there!"), ] threshold, reserve = 1, 0 to_compact, to_keep, _ = asyncio.run( handler.context_check( messages=messages, memory_compact_threshold=threshold, # Trigger compaction memory_compact_reserve=reserve, # Zero reserve ), ) # All messages should be compacted since reserve is 0 assert len(to_compact) == 2, f"All messages should be compacted, got {len(to_compact)}" assert len(to_keep) == 0, f"No messages should be kept, got {len(to_keep)}" verify_context_check_invariants(handler, messages, to_compact, to_keep, threshold, reserve, "test_reserve_zero") print_pass("test_reserve_zero") def test_threshold_zero(): """Test with threshold=0, always triggers compaction.""" handler = create_handler() messages = [create_user_msg("A")] # Minimal message threshold, reserve = 0, 1000 to_compact, to_keep, _ = asyncio.run( handler.context_check( messages=messages, memory_compact_threshold=threshold, # Zero threshold - always triggers memory_compact_reserve=reserve, ), ) # Even minimal message triggers compaction with threshold=0 # But reserve is high so it should be kept assert len(to_compact) == 0 or len(to_keep) == 1, "Message should fit in reserve" verify_context_check_invariants(handler, messages, to_compact, to_keep, threshold, reserve, "test_threshold_zero") print_pass("test_threshold_zero") def test_exact_threshold_boundary(): """Test messages exactly at threshold boundary.""" handler = create_handler() messages = [create_user_msg("Test message")] # Get exact token count stat = asyncio.run(handler.stat_message(messages[0])) exact_tokens = stat.total_tokens threshold, reserve = exact_tokens, exact_tokens # Test at exact boundary to_compact, to_keep, _ = asyncio.run( handler.context_check( messages=messages, memory_compact_threshold=threshold, # Exactly at boundary memory_compact_reserve=reserve, ), ) # At exact boundary (<=), should not trigger compaction assert not to_compact, "Should not compact at exact boundary" assert len(to_keep) == 1, "Should keep message at exact boundary" verify_context_check_invariants( handler, messages, to_compact, to_keep, threshold, reserve, "test_exact_threshold_boundary", ) print_pass("test_exact_threshold_boundary") def test_reserve_larger_than_threshold(): """Test when reserve is larger than threshold (unusual but valid config).""" handler = create_handler() messages = [ create_user_msg("Message one " * 20), create_assistant_msg("Message two " * 20), ] threshold, reserve = 50, 10000 to_compact, to_keep, _ = asyncio.run( handler.context_check( messages=messages, memory_compact_threshold=threshold, # Low threshold memory_compact_reserve=reserve, # High reserve ), ) # Compaction triggered but reserve can hold everything # Total messages should be preserved assert len(to_compact) + len(to_keep) == 2 verify_context_check_invariants( handler, messages, to_compact, to_keep, threshold, reserve, "test_reserve_larger_than_threshold", ) print_pass("test_reserve_larger_than_threshold") # ============================================================================= # Edge Cases - Tool Use/Result Pairing # ============================================================================= def test_tool_use_result_paired(): """Test that tool_use and tool_result pairs are kept together.""" handler = create_handler() messages = [ create_user_msg("Please run the tool " * 50), create_tool_use_msg("call_001", "test_tool", {"arg": "value"}), create_tool_result_msg("call_001", "test_tool", "Tool output"), create_assistant_msg("The tool returned results"), ] threshold, reserve = 50, 1000 to_compact, to_keep, _ = asyncio.run( handler.context_check( messages=messages, memory_compact_threshold=threshold, # Trigger compaction memory_compact_reserve=reserve, # Enough for tool pair ), ) # If tool_result is kept, tool_use should also be kept tool_result_in_keep = any(any(b.get("type") == "tool_result" for b in m.get_content_blocks()) for m in to_keep) tool_use_in_keep = any(any(b.get("type") == "tool_use" for b in m.get_content_blocks()) for m in to_keep) if tool_result_in_keep: assert tool_use_in_keep, "tool_use should be kept when tool_result is kept" verify_context_check_invariants( handler, messages, to_compact, to_keep, threshold, reserve, "test_tool_use_result_paired", ) print_pass("test_tool_use_result_paired") def test_tool_use_without_result(): """Test tool_use message without corresponding tool_result.""" handler = create_handler() messages = [ create_user_msg("Run the tool"), create_tool_use_msg("call_orphan", "orphan_tool", {"arg": "value"}), create_assistant_msg("Something happened"), ] threshold, reserve = 10, 1000 to_compact, to_keep, _ = asyncio.run( handler.context_check( messages=messages, memory_compact_threshold=threshold, memory_compact_reserve=reserve, ), ) # Should not crash, just process normally assert len(to_compact) + len(to_keep) == 3 verify_context_check_invariants( handler, messages, to_compact, to_keep, threshold, reserve, "test_tool_use_without_result", ) print_pass("test_tool_use_without_result") def test_tool_result_without_use(): """Test tool_result message without corresponding tool_use.""" handler = create_handler() messages = [ create_user_msg("Here's a result"), create_tool_result_msg("call_orphan", "orphan_tool", "Some output"), create_assistant_msg("Got it"), ] threshold, reserve = 10, 1000 to_compact, to_keep, _ = asyncio.run( handler.context_check( messages=messages, memory_compact_threshold=threshold, memory_compact_reserve=reserve, ), ) # Should not crash even with orphan tool_result assert len(to_compact) + len(to_keep) == 3 verify_context_check_invariants( handler, messages, to_compact, to_keep, threshold, reserve, "test_tool_result_without_use", ) print_pass("test_tool_result_without_use") def test_multiple_tool_pairs(): """Test multiple tool_use/tool_result pairs.""" handler = create_handler() messages = [ create_user_msg("Start task " * 50), create_tool_use_msg("call_001", "tool_a", {"a": 1}), create_tool_result_msg("call_001", "tool_a", "Result A"), create_tool_use_msg("call_002", "tool_b", {"b": 2}), create_tool_result_msg("call_002", "tool_b", "Result B"), create_tool_use_msg("call_003", "tool_c", {"c": 3}), create_tool_result_msg("call_003", "tool_c", "Result C"), create_assistant_msg("All done"), ] threshold, reserve = 50, 500 to_compact, to_keep, _ = asyncio.run( handler.context_check( messages=messages, memory_compact_threshold=threshold, memory_compact_reserve=reserve, ), ) # Verify tool pairs integrity - for each kept tool_result, its tool_use should be kept for msg in to_keep: for block in msg.get_content_blocks("tool_result"): tool_id = block.get("id", "") if tool_id: # Find corresponding tool_use tool_use_found = False for keep_msg in to_keep: for use_block in keep_msg.get_content_blocks("tool_use"): if use_block.get("id") == tool_id: tool_use_found = True break assert tool_use_found, f"tool_use for {tool_id} should be kept with tool_result" verify_context_check_invariants( handler, messages, to_compact, to_keep, threshold, reserve, "test_multiple_tool_pairs", ) print_pass("test_multiple_tool_pairs") def test_tool_dependency_causes_extra_inclusion(): """Test that tool_use is included even if it exceeds simple reserve calculation.""" handler = create_handler() # Create a scenario where: # - First message (tool_use) is large # - Later message (tool_result) references it # - Reserve alone wouldn't fit tool_use, but dependency requires it large_tool_input = {"data": "x" * 200} messages = [ create_user_msg("Start " * 100), # Large message create_tool_use_msg("call_dep", "dep_tool", large_tool_input), # Medium create_user_msg("Middle " * 100), # Large message create_tool_result_msg("call_dep", "dep_tool", "Result"), # Small create_assistant_msg("End"), # Small ] threshold, reserve = 100, 500 to_compact, to_keep, _ = asyncio.run( handler.context_check( messages=messages, memory_compact_threshold=threshold, # Trigger compaction memory_compact_reserve=reserve, # Medium reserve ), ) # Check pair integrity result_kept = any( any(b.get("id") == "call_dep" and b.get("type") == "tool_result" for b in m.get_content_blocks()) for m in to_keep ) use_kept = any( any(b.get("id") == "call_dep" and b.get("type") == "tool_use" for b in m.get_content_blocks()) for m in to_keep ) if result_kept: assert use_kept, "Dependent tool_use should be included with tool_result" verify_context_check_invariants( handler, messages, to_compact, to_keep, threshold, reserve, "test_tool_dependency_causes_extra_inclusion", ) print_pass("test_tool_dependency_causes_extra_inclusion") def test_tool_dependency_exceeds_reserve(): """Test when tool_result + its tool_use dependency would exceed reserve.""" handler = create_handler() # tool_use is very large, making the pair not fit in reserve very_large_input = {"data": "x" * 2000} messages = [ create_user_msg("First"), create_tool_use_msg("call_big", "big_tool", very_large_input), # Very large create_tool_result_msg("call_big", "big_tool", "Small result"), create_assistant_msg("Last message"), ] threshold, reserve = 10, 100 to_compact, to_keep, _ = asyncio.run( handler.context_check( messages=messages, memory_compact_threshold=threshold, # Trigger compaction memory_compact_reserve=reserve, # Small reserve - can't fit the pair ), ) # The tool pair is too large, so it should be excluded or partially handled # Either both are compacted (pair excluded) or neither is kept result_kept = any( any(b.get("id") == "call_big" and b.get("type") == "tool_result" for b in m.get_content_blocks()) for m in to_keep ) if result_kept: # If result is kept, use must also be kept (pair integrity) use_kept = any( any(b.get("id") == "call_big" and b.get("type") == "tool_use" for b in m.get_content_blocks()) for m in to_keep ) assert use_kept, "Pair integrity violated" verify_context_check_invariants( handler, messages, to_compact, to_keep, threshold, reserve, "test_tool_dependency_exceeds_reserve", ) print_pass("test_tool_dependency_exceeds_reserve") def test_interleaved_tool_pairs(): """Test interleaved tool_use/tool_result (not strictly sequential).""" handler = create_handler() messages = [ create_user_msg("Multi-tool task " * 30), create_tool_use_msg("call_a", "tool_a", {"a": 1}), create_tool_use_msg("call_b", "tool_b", {"b": 2}), # Two uses before results create_tool_result_msg("call_a", "tool_a", "Result A"), create_tool_result_msg("call_b", "tool_b", "Result B"), create_assistant_msg("Both done"), ] threshold, reserve = 50, 1000 to_compact, to_keep, _ = asyncio.run( handler.context_check( messages=messages, memory_compact_threshold=threshold, memory_compact_reserve=reserve, ), ) # Verify pair integrity for interleaved pairs for msg in to_keep: for block in msg.get_content_blocks("tool_result"): tool_id = block.get("id", "") if tool_id: use_found = any( any(ub.get("id") == tool_id and ub.get("type") == "tool_use" for ub in km.get_content_blocks()) for km in to_keep ) assert use_found, f"Interleaved tool_use {tool_id} should be kept" verify_context_check_invariants( handler, messages, to_compact, to_keep, threshold, reserve, "test_interleaved_tool_pairs", ) print_pass("test_interleaved_tool_pairs") # ============================================================================= # Edge Cases - Message Content Variations # ============================================================================= def test_message_with_empty_content(): """Test message with empty string content.""" handler = create_handler() messages = [ create_user_msg(""), # Empty content create_assistant_msg("Response"), ] threshold, reserve = 1000, 500 to_compact, to_keep, _ = asyncio.run( handler.context_check( messages=messages, memory_compact_threshold=threshold, memory_compact_reserve=reserve, ), ) assert len(to_compact) + len(to_keep) == 2 verify_context_check_invariants( handler, messages, to_compact, to_keep, threshold, reserve, "test_message_with_empty_content", ) print_pass("test_message_with_empty_content") def test_message_with_whitespace_only(): """Test message with whitespace-only content.""" handler = create_handler() messages = [ create_user_msg(" \n\t "), # Whitespace only create_assistant_msg("Response"), ] threshold, reserve = 1000, 500 to_compact, to_keep, _ = asyncio.run( handler.context_check( messages=messages, memory_compact_threshold=threshold, memory_compact_reserve=reserve, ), ) assert len(to_compact) + len(to_keep) == 2 verify_context_check_invariants( handler, messages, to_compact, to_keep, threshold, reserve, "test_message_with_whitespace_only", ) print_pass("test_message_with_whitespace_only") def test_very_long_single_message(): """Test very long single message that exceeds any reasonable reserve.""" handler = create_handler() huge_content = "x" * 100000 # Very long messages = [create_user_msg(huge_content)] threshold, reserve = 100, 1000 to_compact, to_keep, _ = asyncio.run( handler.context_check( messages=messages, memory_compact_threshold=threshold, memory_compact_reserve=reserve, ), ) # Single huge message - either kept alone or compacted assert len(to_compact) + len(to_keep) == 1 verify_context_check_invariants( handler, messages, to_compact, to_keep, threshold, reserve, "test_very_long_single_message", ) print_pass("test_very_long_single_message") def test_many_small_messages(): """Test many small messages.""" handler = create_handler() messages = [create_user_msg(f"Msg {i}") for i in range(100)] threshold, reserve = 100, 200 to_compact, to_keep, _ = asyncio.run( handler.context_check( messages=messages, memory_compact_threshold=threshold, # Low threshold memory_compact_reserve=reserve, ), ) # Should compact older messages and keep recent ones assert len(to_compact) + len(to_keep) == 100 assert len(to_keep) > 0, "Should keep some messages" verify_context_check_invariants( handler, messages, to_compact, to_keep, threshold, reserve, "test_many_small_messages", ) print_pass("test_many_small_messages") def test_unicode_content(): """Test messages with unicode characters.""" handler = create_handler() messages = [ create_user_msg("你好世界！🎉 Emoji and 中文"), create_assistant_msg("مرحبا العالم 🌍 Arabic and more"), create_user_msg("日本語テスト 🇯🇵"), ] threshold, reserve = 1000, 500 to_compact, to_keep, _ = asyncio.run( handler.context_check( messages=messages, memory_compact_threshold=threshold, memory_compact_reserve=reserve, ), ) assert len(to_compact) + len(to_keep) == 3 verify_context_check_invariants(handler, messages, to_compact, to_keep, threshold, reserve, "test_unicode_content") print_pass("test_unicode_content") def test_special_characters_content(): """Test messages with special characters.""" handler = create_handler() messages = [ create_user_msg("Special chars: <>&\"'`~!@#$%^&*()[]{}|\\"), create_assistant_msg("More: \n\r\t\0 nulls and newlines"), ] threshold, reserve = 1000, 500 to_compact, to_keep, _ = asyncio.run( handler.context_check( messages=messages, memory_compact_threshold=threshold, memory_compact_reserve=reserve, ), ) assert len(to_compact) + len(to_keep) == 2 verify_context_check_invariants( handler, messages, to_compact, to_keep, threshold, reserve, "test_special_characters_content", ) print_pass("test_special_characters_content") # ============================================================================= # Edge Cases - Boundary Conditions # ============================================================================= def test_all_messages_fit_exactly_in_reserve(): """Test when all messages fit exactly in reserve after threshold exceeded.""" handler = create_handler() messages = [ create_user_msg("Message 1"), create_assistant_msg("Message 2"), ] # Calculate total tokens total = sum(asyncio.run(handler.stat_message(m)).total_tokens for m in messages) threshold, reserve = total - 1, total to_compact, to_keep, _ = asyncio.run( handler.context_check( messages=messages, memory_compact_threshold=threshold, # Just below total to trigger memory_compact_reserve=reserve, # Exactly fits all ), ) # All should be kept since reserve can hold everything assert len(to_keep) == 2, f"All messages should fit in reserve, got {len(to_keep)}" verify_context_check_invariants( handler, messages, to_compact, to_keep, threshold, reserve, "test_all_messages_fit_exactly_in_reserve", ) print_pass("test_all_messages_fit_exactly_in_reserve") def test_first_message_only_compacted(): """Test when only the first message is compacted.""" handler = create_handler() messages = [ create_user_msg("Large first message " * 100), # Large create_assistant_msg("Small"), # Small create_user_msg("Tiny"), # Tiny ] # Calculate tokens to set appropriate reserve small_msg_tokens = asyncio.run(handler.stat_message(messages[1])).total_tokens tiny_msg_tokens = asyncio.run(handler.stat_message(messages[2])).total_tokens threshold, reserve = 50, small_msg_tokens + tiny_msg_tokens + 10 to_compact, to_keep, _ = asyncio.run( handler.context_check( messages=messages, memory_compact_threshold=threshold, # Low to trigger memory_compact_reserve=reserve, # Fits last 2 ), ) assert len(to_compact) >= 1, "At least first message should be compacted" verify_context_check_invariants( handler, messages, to_compact, to_keep, threshold, reserve, "test_first_message_only_compacted", ) print_pass("test_first_message_only_compacted") def test_last_message_only_kept(): """Test when only the last message can be kept.""" handler = create_handler() messages = [ create_user_msg("Large " * 200), create_assistant_msg("Large " * 200), create_user_msg("Tiny"), # Only this fits ] tiny_tokens = asyncio.run(handler.stat_message(messages[2])).total_tokens threshold, reserve = 10, tiny_tokens + 5 to_compact, to_keep, _ = asyncio.run( handler.context_check( messages=messages, memory_compact_threshold=threshold, memory_compact_reserve=reserve, # Only fits last message ), ) if len(to_keep) == 1: # Last message should be the one kept assert to_keep[0] == messages[2], "Only last message should be kept" verify_context_check_invariants( handler, messages, to_compact, to_keep, threshold, reserve, "test_last_message_only_kept", ) print_pass("test_last_message_only_kept") def test_all_messages_compacted(): """Test when all messages need to be compacted (nothing fits in reserve).""" handler = create_handler() messages = [ create_user_msg("Large message " * 100), create_assistant_msg("Large message " * 100), ] threshold, reserve = 10, 1 to_compact, to_keep, _ = asyncio.run( handler.context_check( messages=messages, memory_compact_threshold=threshold, # Trigger compaction memory_compact_reserve=reserve, # Too small for anything ), ) assert len(to_compact) == 2, "All messages should be compacted" assert len(to_keep) == 0, "No messages should be kept" verify_context_check_invariants( handler, messages, to_compact, to_keep, threshold, reserve, "test_all_messages_compacted", ) print_pass("test_all_messages_compacted") # ============================================================================= # Edge Cases - Message Roles # ============================================================================= def test_system_message(): """Test handling of system role messages.""" handler = create_handler() system_msg = Msg(name="system", role="system", content="You are a helpful assistant.") messages = [ system_msg, create_user_msg("Hello"), create_assistant_msg("Hi there!"), ] threshold, reserve = 1000, 500 to_compact, to_keep, _ = asyncio.run( handler.context_check( messages=messages, memory_compact_threshold=threshold, memory_compact_reserve=reserve, ), ) assert len(to_compact) + len(to_keep) == 3 verify_context_check_invariants(handler, messages, to_compact, to_keep, threshold, reserve, "test_system_message") print_pass("test_system_message") def test_mixed_roles(): """Test messages with various roles (user, assistant, system).""" handler = create_handler() # agentscope.message.Msg only supports: user, assistant, system messages = [ Msg(name="system", role="system", content="System prompt"), Msg(name="user", role="user", content="User message"), Msg(name="assistant", role="assistant", content="Assistant response"), Msg(name="tool", role="user", content="Tool output as user role"), Msg(name="helper", role="assistant", content="Another assistant message"), ] threshold, reserve = 1000, 500 to_compact, to_keep, _ = asyncio.run( handler.context_check( messages=messages, memory_compact_threshold=threshold, memory_compact_reserve=reserve, ), ) assert len(to_compact) + len(to_keep) == 5 verify_context_check_invariants(handler, messages, to_compact, to_keep, threshold, reserve, "test_mixed_roles") print_pass("test_mixed_roles") # ============================================================================= # Edge Cases - Tool Block Variations # ============================================================================= def test_tool_use_with_empty_id(): """Test tool_use block with empty id.""" handler = create_handler() messages = [ create_user_msg("Run tool"), create_tool_use_msg("", "test_tool", {"arg": "value"}), # Empty ID create_assistant_msg("Done"), ] threshold, reserve = 10, 1000 to_compact, to_keep, _ = asyncio.run( handler.context_check( messages=messages, memory_compact_threshold=threshold, memory_compact_reserve=reserve, ), ) # Should handle gracefully assert len(to_compact) + len(to_keep) == 3 verify_context_check_invariants( handler, messages, to_compact, to_keep, threshold, reserve, "test_tool_use_with_empty_id", ) print_pass("test_tool_use_with_empty_id") def test_tool_result_with_empty_id(): """Test tool_result block with empty id.""" handler = create_handler() messages = [ create_user_msg("Got result"), create_tool_result_msg("", "test_tool", "Output"), # Empty ID create_assistant_msg("Noted"), ] threshold, reserve = 10, 1000 to_compact, to_keep, _ = asyncio.run( handler.context_check( messages=messages, memory_compact_threshold=threshold, memory_compact_reserve=reserve, ), ) # Should handle gracefully assert len(to_compact) + len(to_keep) == 3 verify_context_check_invariants( handler, messages, to_compact, to_keep, threshold, reserve, "test_tool_result_with_empty_id", ) print_pass("test_tool_result_with_empty_id") def test_duplicate_tool_ids(): """Test messages with duplicate tool IDs (unusual but possible).""" handler = create_handler() messages = [ create_tool_use_msg("call_dup", "tool_a", {"a": 1}), create_tool_result_msg("call_dup", "tool_a", "Result A"), create_tool_use_msg("call_dup", "tool_b", {"b": 2}), # Same ID, different tool create_tool_result_msg("call_dup", "tool_b", "Result B"), ] threshold, reserve = 10, 1000 to_compact, to_keep, _ = asyncio.run( handler.context_check( messages=messages, memory_compact_threshold=threshold, memory_compact_reserve=reserve, ), ) # Should not crash with duplicate IDs assert len(to_compact) + len(to_keep) == 4 verify_context_check_invariants( handler, messages, to_compact, to_keep, threshold, reserve, "test_duplicate_tool_ids", ) print_pass("test_duplicate_tool_ids") def test_message_with_multiple_tool_blocks(): """Test single message containing multiple tool blocks.""" handler = create_handler() msg_with_multiple_tools = Msg( name="assistant", role="assistant", content=[ {"type": "tool_use", "id": "call_1", "name": "tool1", "input": {}}, {"type": "tool_use", "id": "call_2", "name": "tool2", "input": {}}, {"type": "tool_use", "id": "call_3", "name": "tool3", "input": {}}, ], ) messages = [ create_user_msg("Do multiple things"), msg_with_multiple_tools, create_tool_result_msg("call_1", "tool1", "Result 1"), create_tool_result_msg("call_2", "tool2", "Result 2"), create_tool_result_msg("call_3", "tool3", "Result 3"), ] threshold, reserve = 10, 2000 to_compact, to_keep, _ = asyncio.run( handler.context_check( messages=messages, memory_compact_threshold=threshold, memory_compact_reserve=reserve, ), ) assert len(to_compact) + len(to_keep) == 5 verify_context_check_invariants( handler, messages, to_compact, to_keep, threshold, reserve, "test_message_with_multiple_tool_blocks", ) print_pass("test_message_with_multiple_tool_blocks") # ============================================================================= # Run All Tests # ============================================================================= def run_all_tests(): """Run all tests.""" tests = [ # Normal cases test_empty_messages, test_below_threshold_returns_all, test_above_threshold_triggers_compaction, test_message_order_preserved, # Edge cases - boundaries test_single_message_below_threshold, test_single_message_above_threshold, test_reserve_zero, test_threshold_zero, test_exact_threshold_boundary, test_reserve_larger_than_threshold, # Edge cases - tool pairing test_tool_use_result_paired, test_tool_use_without_result, test_tool_result_without_use, test_multiple_tool_pairs, test_tool_dependency_causes_extra_inclusion, test_tool_dependency_exceeds_reserve, test_interleaved_tool_pairs, # Edge cases - content variations test_message_with_empty_content, test_message_with_whitespace_only, test_very_long_single_message, test_many_small_messages, test_unicode_content, test_special_characters_content, # Edge cases - boundaries test_all_messages_fit_exactly_in_reserve, test_first_message_only_compacted, test_last_message_only_kept, test_all_messages_compacted, # Edge cases - roles test_system_message, test_mixed_roles, # Edge cases - tool blocks test_tool_use_with_empty_id, test_tool_result_with_empty_id, test_duplicate_tool_ids, test_message_with_multiple_tool_blocks, ] passed = 0 failed = 0 for test in tests: try: print_test_header(test.__name__) test() passed += 1 except AssertionError as e: print_fail(test.__name__, str(e)) failed += 1 except Exception as e: print_error(test.__name__, str(e)) failed += 1 # Print summary print(f"\n{Colors.CYAN}{'=' * 60}{Colors.RESET}") print(f"{Colors.BOLD}Test Results Summary{Colors.RESET}") print(f"{Colors.CYAN}{'=' * 60}{Colors.RESET}") print(f"{Colors.GREEN}{Colors.BOLD}✓ Passed: {passed}{Colors.RESET}") if failed > 0: print(f"{Colors.RED}{Colors.BOLD}✗ Failed: {failed}{Colors.RESET}") else: print(f"{Colors.GREEN}✗ Failed: {failed}{Colors.RESET}") print(f"{Colors.CYAN}{'=' * 60}{Colors.RESET}") if failed == 0: print(f"\n{Colors.GREEN}{Colors.BOLD}🎉 All tests passed!{Colors.RESET}") else: print(f"\n{Colors.RED}{Colors.BOLD}💥 Some tests failed!{Colors.RESET}") if __name__ == "__main__": run_all_tests()