"""Tests for rolling window transform.""" import pytest from headroom import OpenAIProvider, RollingWindowConfig, Tokenizer from headroom.parser import find_tool_units from headroom.transforms import RollingWindow # Create a shared provider for tests _provider = OpenAIProvider() def get_tokenizer(model: str = "gpt-4o") -> Tokenizer: """Get a tokenizer for tests using OpenAI provider.""" token_counter = _provider.get_token_counter(model) return Tokenizer(token_counter, model) # Fixtures for realistic message scenarios @pytest.fixture def messages_with_system(): """Messages with a system prompt.""" return [ { "role": "system", "content": "You are a helpful assistant. You help users with their tasks.", }, {"role": "user", "content": "Hello, can you help me?"}, {"role": "assistant", "content": "Of course! What do you need help with?"}, {"role": "user", "content": "I need to analyze some data."}, { "role": "assistant", "content": "I'd be happy to help analyze your data. What kind of data do you have?", }, ] @pytest.fixture def messages_with_tool_calls(): """Messages with tool calls and responses.""" return [ {"role": "system", "content": "You are a helpful assistant with access to tools."}, {"role": "user", "content": "Find user with ID 12345"}, { "role": "assistant", "content": None, "tool_calls": [ { "id": "call_abc123", "type": "function", "function": {"name": "get_user", "arguments": '{"user_id": "12345"}'}, } ], }, { "role": "tool", "tool_call_id": "call_abc123", "content": '{"id": "12345", "name": "Alice", "email": "alice@example.com", "status": "active"}', }, { "role": "assistant", "content": "I found the user. Alice (ID: 12345) is an active user with email alice@example.com.", }, {"role": "user", "content": "Can you also find user 67890?"}, { "role": "assistant", "content": None, "tool_calls": [ { "id": "call_def456", "type": "function", "function": {"name": "get_user", "arguments": '{"user_id": "67890"}'}, } ], }, { "role": "tool", "tool_call_id": "call_def456", "content": '{"id": "67890", "name": "Bob", "email": "bob@example.com", "status": "inactive"}', }, {"role": "assistant", "content": "Found Bob (ID: 67890). This user is currently inactive."}, {"role": "user", "content": "Thanks for the help!"}, {"role": "assistant", "content": "You're welcome! Let me know if you need anything else."}, ] @pytest.fixture def messages_multiple_tool_calls(): """Messages with multiple tool calls in a single assistant message.""" return [ {"role": "system", "content": "You are a helpful assistant."}, {"role": "user", "content": "Search for users Alice and Bob"}, { "role": "assistant", "content": None, "tool_calls": [ { "id": "call_multi_1", "type": "function", "function": {"name": "search_user", "arguments": '{"name": "Alice"}'}, }, { "id": "call_multi_2", "type": "function", "function": {"name": "search_user", "arguments": '{"name": "Bob"}'}, }, ], }, {"role": "tool", "tool_call_id": "call_multi_1", "content": '{"id": "1", "name": "Alice"}'}, {"role": "tool", "tool_call_id": "call_multi_2", "content": '{"id": "2", "name": "Bob"}'}, {"role": "assistant", "content": "I found both users."}, ] @pytest.fixture def long_conversation(): """A longer conversation to test window dropping.""" messages = [ {"role": "system", "content": "You are a helpful assistant. " * 50}, # ~250 tokens ] # Add 20 turns of conversation for i in range(20): messages.append( {"role": "user", "content": f"This is user message number {i}. " * 10} ) # ~50 tokens each messages.append( {"role": "assistant", "content": f"This is assistant response number {i}. " * 10} ) return messages class TestRollingWindowProtection: """Tests for protected message handling.""" def test_never_drops_system_prompt(self, messages_with_system): """System prompt should never be dropped even under tight budget.""" config = RollingWindowConfig( enabled=True, keep_system=True, keep_last_turns=0, # Don't protect any turns output_buffer_tokens=0, ) window = RollingWindow(config) tokenizer = get_tokenizer() # Apply with very small budget (smaller than system prompt itself isn't realistic, # but we test that system is retained even when budget is tight) result = window.apply( messages_with_system, tokenizer, model_limit=200, # Very small budget output_buffer=0, ) # System message should still be present system_messages = [m for m in result.messages if m.get("role") == "system"] assert len(system_messages) == 1 assert "You are a helpful assistant" in system_messages[0]["content"] def test_never_drops_last_n_turns(self, messages_with_system): """Last N turns should be protected from dropping.""" config = RollingWindowConfig( enabled=True, keep_system=True, keep_last_turns=2, # Protect last 2 turns output_buffer_tokens=0, ) window = RollingWindow(config) tokenizer = get_tokenizer() # Apply with tight budget result = window.apply( messages_with_system, tokenizer, model_limit=300, # Tight budget output_buffer=0, ) # Last 2 turns (4 messages) should be preserved # The last messages in the conversation are about analyzing data last_user = [m for m in result.messages if m.get("role") == "user"][-1] last_assistant = [m for m in result.messages if m.get("role") == "assistant"][-1] assert "analyze" in last_user["content"].lower() or "data" in last_user["content"].lower() assert "data" in last_assistant["content"].lower() def test_protects_tool_responses_for_protected_assistant(self, messages_with_tool_calls): """Tool responses for protected assistant messages should also be protected.""" config = RollingWindowConfig( enabled=True, keep_system=True, keep_last_turns=1, # Protect last turn output_buffer_tokens=0, ) window = RollingWindow(config) tokenizer = get_tokenizer() # Get original last assistant with tool calls if any # In our fixture, the last assistant doesn't have tool calls, but second-to-last does # Modify fixture for this test - use a smaller subset messages = [ {"role": "system", "content": "You are a helpful assistant."}, {"role": "user", "content": "Old message"}, {"role": "assistant", "content": "Old response"}, {"role": "user", "content": "Find user 999"}, { "role": "assistant", "content": None, "tool_calls": [ { "id": "call_protected", "type": "function", "function": {"name": "get_user", "arguments": '{"id": "999"}'}, } ], }, { "role": "tool", "tool_call_id": "call_protected", "content": '{"id": "999", "name": "Protected User"}', }, {"role": "user", "content": "Thanks!"}, {"role": "assistant", "content": "You're welcome!"}, ] window.apply( messages, tokenizer, model_limit=500, output_buffer=0, ) # Check that if the assistant with tool_calls is protected, its tool response is too protected_indices = window._get_protected_indices(messages) # Find the assistant message with tool_calls for i, msg in enumerate(messages): if msg.get("role") == "assistant" and msg.get("tool_calls"): if i in protected_indices: # The tool response should also be protected for tc in msg.get("tool_calls", []): tc_id = tc.get("id") for j, other_msg in enumerate(messages): if ( other_msg.get("role") == "tool" and other_msg.get("tool_call_id") == tc_id ): assert j in protected_indices, ( f"Tool response at {j} should be protected" ) class TestDropPriority: """Tests for drop priority ordering.""" def test_drops_oldest_tool_units_first(self, messages_with_tool_calls): """Tool units should be dropped before regular turns.""" config = RollingWindowConfig( enabled=True, keep_system=True, keep_last_turns=1, # Protect only the last turn output_buffer_tokens=0, ) window = RollingWindow(config) get_tokenizer() # Check drop candidates ordering protected = window._get_protected_indices(messages_with_tool_calls) tool_units = find_tool_units(messages_with_tool_calls) candidates = window._build_drop_candidates(messages_with_tool_calls, protected, tool_units) # First candidates should be tool units (priority 1) tool_candidates = [c for c in candidates if c["type"] == "tool_unit"] turn_candidates = [c for c in candidates if c["type"] in ("turn", "single")] if tool_candidates and turn_candidates: # Tool units should come before turns in the sorted list first_tool_idx = candidates.index(tool_candidates[0]) first_turn_idx = candidates.index(turn_candidates[0]) assert first_tool_idx < first_turn_idx, "Tool units should be dropped before turns" def test_drops_oldest_turns_second(self, long_conversation): """After tool units, oldest turns should be dropped.""" config = RollingWindowConfig( enabled=True, keep_system=True, keep_last_turns=2, output_buffer_tokens=0, ) window = RollingWindow(config) tokenizer = get_tokenizer() # Apply with budget that forces dropping some turns result = window.apply( long_conversation, tokenizer, model_limit=2000, # Will need to drop many turns output_buffer=0, ) # The newest messages (user message 19, assistant 19) should still be there # Check that at least some oldest messages were dropped remaining_content = " ".join(m.get("content", "") or "" for m in result.messages) # Last turn should be preserved assert "number 19" in remaining_content # Old turns should be dropped (message 0 or 1) # Due to marker insertion, check that we have fewer messages assert len(result.messages) < len(long_conversation) def test_tool_unit_atomic(self, messages_with_tool_calls): """Tool calls and their responses should be dropped together.""" config = RollingWindowConfig( enabled=True, keep_system=True, keep_last_turns=1, output_buffer_tokens=0, ) window = RollingWindow(config) tokenizer = get_tokenizer() # Find tool units before applying tool_units_before = find_tool_units(messages_with_tool_calls) assert len(tool_units_before) > 0, "Test requires messages with tool calls" # Apply transform with tight budget result = window.apply( messages_with_tool_calls, tokenizer, model_limit=500, output_buffer=0, ) # For each remaining assistant with tool_calls, verify its tool responses exist for msg in result.messages: if msg.get("role") == "assistant" and msg.get("tool_calls"): tool_call_ids = {tc.get("id") for tc in msg.get("tool_calls", [])} # Find matching tool responses tool_responses = [ m for m in result.messages if m.get("role") == "tool" and m.get("tool_call_id") in tool_call_ids ] # All tool calls should have their responses response_ids = {m.get("tool_call_id") for m in tool_responses} assert tool_call_ids == response_ids, "Tool calls must have matching responses" def test_never_orphans_tool_result(self, messages_multiple_tool_calls): """Tool results should never be left without their assistant message.""" config = RollingWindowConfig( enabled=True, keep_system=True, keep_last_turns=1, output_buffer_tokens=0, ) window = RollingWindow(config) tokenizer = get_tokenizer() result = window.apply( messages_multiple_tool_calls, tokenizer, model_limit=300, output_buffer=0, ) # Check that no tool message is orphaned for msg in result.messages: if msg.get("role") == "tool": tool_call_id = msg.get("tool_call_id") # Find the corresponding assistant message found_assistant = False for other_msg in result.messages: if other_msg.get("role") == "assistant" and other_msg.get("tool_calls"): for tc in other_msg.get("tool_calls", []): if tc.get("id") == tool_call_id: found_assistant = True break assert found_assistant, f"Tool response {tool_call_id} is orphaned" class TestTokenBudget: """Tests for token budget handling.""" def test_no_drops_under_budget(self, messages_with_system): """No messages should be dropped when under budget.""" config = RollingWindowConfig( enabled=True, keep_system=True, keep_last_turns=2, output_buffer_tokens=0, ) window = RollingWindow(config) tokenizer = get_tokenizer() # Apply with large budget result = window.apply( messages_with_system, tokenizer, model_limit=100000, # Plenty of room output_buffer=0, ) # No messages should be dropped assert len(result.messages) == len(messages_with_system) assert len(result.transforms_applied) == 0 assert result.tokens_before == result.tokens_after def test_drops_until_under_budget(self, long_conversation): """Should drop messages until under budget.""" config = RollingWindowConfig( enabled=True, keep_system=True, keep_last_turns=2, output_buffer_tokens=0, ) window = RollingWindow(config) tokenizer = get_tokenizer() model_limit = 1500 # Tight budget result = window.apply( long_conversation, tokenizer, model_limit=model_limit, output_buffer=0, ) # Should be under budget after transform assert result.tokens_after <= model_limit # Should have dropped something assert result.tokens_after < result.tokens_before def test_respects_output_buffer(self, long_conversation): """Output buffer should be subtracted from available budget.""" config = RollingWindowConfig( enabled=True, keep_system=True, keep_last_turns=2, output_buffer_tokens=4000, # Default buffer ) window = RollingWindow(config) tokenizer = get_tokenizer() model_limit = 5000 output_buffer = 2000 result = window.apply( long_conversation, tokenizer, model_limit=model_limit, output_buffer=output_buffer, ) # Should be under (model_limit - output_buffer) available = model_limit - output_buffer assert result.tokens_after <= available def test_model_limit_parameter(self, long_conversation): """Model limit parameter should control the budget.""" config = RollingWindowConfig( enabled=True, keep_system=True, keep_last_turns=2, output_buffer_tokens=0, ) window = RollingWindow(config) tokenizer = get_tokenizer() # Test with different model limits result_small = window.apply( long_conversation, tokenizer, model_limit=1000, output_buffer=0, ) result_large = window.apply( long_conversation, tokenizer, model_limit=3000, output_buffer=0, ) # Smaller limit should result in fewer tokens assert result_small.tokens_after <= result_large.tokens_after class TestMarkers: """Tests for dropped context markers.""" def test_inserts_dropped_context_marker(self, long_conversation): """A marker should be inserted when content is dropped.""" config = RollingWindowConfig( enabled=True, keep_system=True, keep_last_turns=2, output_buffer_tokens=0, ) window = RollingWindow(config) tokenizer = get_tokenizer() result = window.apply( long_conversation, tokenizer, model_limit=1500, output_buffer=0, ) # Should have a marker in the result assert len(result.markers_inserted) > 0 assert "dropped_context" in result.markers_inserted[0] # Marker should be in the messages marker_found = False for msg in result.messages: content = msg.get("content", "") if content and " 0 marker = result.markers_inserted[0] assert 'count="' in marker # Transforms applied should indicate count assert len(result.transforms_applied) > 0 assert "window_cap:" in result.transforms_applied[0] class TestBuildDropCandidates: """Tests for the _build_drop_candidates method.""" def test_tool_units_have_priority_1(self, messages_with_tool_calls): """Tool units should have priority 1.""" config = RollingWindowConfig(enabled=True) window = RollingWindow(config) protected = window._get_protected_indices(messages_with_tool_calls) tool_units = find_tool_units(messages_with_tool_calls) candidates = window._build_drop_candidates(messages_with_tool_calls, protected, tool_units) tool_candidates = [c for c in candidates if c["type"] == "tool_unit"] for tc in tool_candidates: assert tc["priority"] == 1 def test_turns_have_priority_2(self, messages_with_system): """Regular turns should have priority 2.""" config = RollingWindowConfig(enabled=True, keep_last_turns=0) window = RollingWindow(config) protected = window._get_protected_indices(messages_with_system) tool_units = find_tool_units(messages_with_system) # Empty for this fixture candidates = window._build_drop_candidates(messages_with_system, protected, tool_units) turn_candidates = [c for c in candidates if c["type"] in ("turn", "single")] for tc in turn_candidates: assert tc["priority"] == 2 def test_candidates_sorted_by_age(self, long_conversation): """Candidates should be sorted by priority then by position (oldest first).""" config = RollingWindowConfig(enabled=True, keep_last_turns=0) window = RollingWindow(config) protected = {0} # Only protect system tool_units = [] candidates = window._build_drop_candidates(long_conversation, protected, tool_units) # Check that candidates are sorted: first by priority, then by position for i in range(1, len(candidates)): prev = candidates[i - 1] curr = candidates[i] # Priority should be non-decreasing assert prev["priority"] <= curr["priority"] # Within same priority, position should be increasing if prev["priority"] == curr["priority"]: assert prev["position"] <= curr["position"] class TestEdgeCases: """Tests for edge cases.""" def test_empty_messages(self): """Empty message list should be handled gracefully.""" config = RollingWindowConfig(enabled=True) window = RollingWindow(config) tokenizer = get_tokenizer() result = window.apply( [], tokenizer, model_limit=1000, output_buffer=0, ) assert result.messages == [] # Tokenizer may return a small overhead even for empty messages assert result.tokens_before == result.tokens_after assert len(result.transforms_applied) == 0 def test_system_only(self): """Conversation with only system message should not drop anything.""" messages = [{"role": "system", "content": "You are a helpful assistant."}] config = RollingWindowConfig(enabled=True, keep_system=True) window = RollingWindow(config) tokenizer = get_tokenizer() result = window.apply( messages, tokenizer, model_limit=1000, output_buffer=0, ) assert len(result.messages) == 1 assert result.messages[0]["role"] == "system" def test_large_conversation(self, long_conversation): """Large conversations should be handled without errors.""" config = RollingWindowConfig( enabled=True, keep_system=True, keep_last_turns=3, output_buffer_tokens=0, ) window = RollingWindow(config) tokenizer = get_tokenizer() result = window.apply( long_conversation, tokenizer, model_limit=2000, output_buffer=0, ) # Should complete without errors assert result.tokens_after <= 2000 # System should be preserved system_msgs = [m for m in result.messages if m.get("role") == "system"] assert len(system_msgs) == 1 def test_all_protected(self): """When everything is protected, nothing should be dropped.""" messages = [ {"role": "system", "content": "You are helpful."}, {"role": "user", "content": "Hello"}, {"role": "assistant", "content": "Hi there!"}, ] config = RollingWindowConfig( enabled=True, keep_system=True, keep_last_turns=10, # More than we have output_buffer_tokens=0, ) window = RollingWindow(config) tokenizer = get_tokenizer() result = window.apply( messages, tokenizer, model_limit=50, # Very tight, but everything is protected output_buffer=0, ) # Nothing should be dropped since everything is protected # The result might still be over budget assert len(result.messages) == len(messages) assert len(result.transforms_applied) == 0 class TestShouldApply: """Tests for the should_apply method.""" def test_disabled_config_returns_false(self, messages_with_system): """should_apply returns False when disabled in config.""" config = RollingWindowConfig(enabled=False) window = RollingWindow(config) tokenizer = get_tokenizer() assert window.should_apply(messages_with_system, tokenizer, model_limit=100) is False def test_under_budget_returns_false(self, messages_with_system): """should_apply returns False when under budget.""" config = RollingWindowConfig(enabled=True) window = RollingWindow(config) tokenizer = get_tokenizer() # Large budget - should not need to apply assert window.should_apply(messages_with_system, tokenizer, model_limit=100000) is False def test_over_budget_returns_true(self, long_conversation): """should_apply returns True when over budget.""" config = RollingWindowConfig(enabled=True) window = RollingWindow(config) tokenizer = get_tokenizer() # Small budget - should need to apply assert window.should_apply(long_conversation, tokenizer, model_limit=500) is True class TestConvenienceFunction: """Tests for the apply_rolling_window convenience function.""" def test_convenience_function(self, long_conversation): """The convenience function should work correctly.""" from headroom.transforms.rolling_window import apply_rolling_window messages, transforms = apply_rolling_window( long_conversation, model_limit=1500, output_buffer=500, keep_last_turns=2, ) # Should have applied transform assert len(messages) < len(long_conversation) assert len(transforms) > 0 def test_convenience_function_with_config(self, long_conversation): """The convenience function should accept a config.""" from headroom.transforms.rolling_window import apply_rolling_window config = RollingWindowConfig( enabled=True, keep_system=True, keep_last_turns=3, ) messages, transforms = apply_rolling_window( long_conversation, model_limit=1500, output_buffer=500, keep_last_turns=3, config=config, ) assert len(messages) < len(long_conversation) class TestTransformResult: """Tests for TransformResult fields.""" def test_tokens_before_after(self, long_conversation): """tokens_before and tokens_after should be correct.""" config = RollingWindowConfig(enabled=True, keep_last_turns=2) window = RollingWindow(config) tokenizer = get_tokenizer() tokens_before = tokenizer.count_messages(long_conversation) result = window.apply( long_conversation, tokenizer, model_limit=1500, output_buffer=0, ) assert result.tokens_before == tokens_before assert result.tokens_after <= result.tokens_before assert result.tokens_after == tokenizer.count_messages(result.messages) def test_transforms_applied_field(self, long_conversation): """transforms_applied should contain window_cap info.""" config = RollingWindowConfig(enabled=True, keep_last_turns=2) window = RollingWindow(config) tokenizer = get_tokenizer() result = window.apply( long_conversation, tokenizer, model_limit=1500, output_buffer=0, ) assert len(result.transforms_applied) > 0 assert any("window_cap" in t for t in result.transforms_applied) def test_warnings_field(self, messages_with_system): """warnings field should be present (may be empty).""" config = RollingWindowConfig(enabled=True) window = RollingWindow(config) tokenizer = get_tokenizer() result = window.apply( messages_with_system, tokenizer, model_limit=100000, output_buffer=0, ) # warnings should be a list (possibly empty) assert isinstance(result.warnings, list)