"""Tests for MCP (Model Context Protocol) integration. These tests verify that Headroom correctly compresses MCP tool outputs while preserving 100% of critical data (errors, anomalies). """ import json import random from datetime import datetime, timedelta import pytest from headroom.integrations.mcp import ( HeadroomMCPClientWrapper, HeadroomMCPCompressor, MCPCompressionResult, MCPToolProfile, compress_tool_result, compress_tool_result_with_metrics, ) from headroom.providers import OpenAIProvider from headroom.transforms.smart_crusher import strip_ccr_sentinels # ============================================================================ # Test Fixtures # ============================================================================ @pytest.fixture def mcp_compressor(): """Create MCP compressor with default settings.""" return HeadroomMCPCompressor() @pytest.fixture def tokenizer(): """Create token counter.""" provider = OpenAIProvider() return provider.get_token_counter("gpt-4o") def generate_slack_messages(count: int, error_rate: float = 0.1) -> str: """Generate realistic Slack search results.""" channels = ["#engineering", "#incidents", "#support", "#general", "#alerts"] users = ["alice", "bob", "charlie", "diana", "eve"] messages = [] for i in range(count): is_error = random.random() < error_rate if is_error: text = random.choice( [ "ERROR: Database connection failed at 2:30am", "CRITICAL: API latency spike detected", "Exception: NullPointerException in AuthService", "FAILED: Build pipeline broke on main branch", "BUG: Users can't login - investigating now", ] ) else: text = random.choice( [ "Reviewed the PR, looks good to merge", "Updated the docs with new API endpoints", "Meeting notes from standup attached", "Can someone review my changes?", "Deployed v2.3.1 to staging", ] ) messages.append( { "id": f"msg_{i}", "channel": random.choice(channels), "user": random.choice(users), "text": text, "timestamp": (datetime.now() - timedelta(hours=i)).isoformat(), "reactions": random.randint(0, 10), } ) return json.dumps({"messages": messages, "total": count}) def generate_database_results(count: int, null_rate: float = 0.1) -> str: """Generate realistic database query results.""" rows = [] for i in range(count): has_null = random.random() < null_rate has_error = random.random() < 0.05 # 5% error rate row = { "id": i + 1, "user_id": f"user_{random.randint(1000, 9999)}", "email": f"user{i}@example.com", "status": "ERROR: validation failed" if has_error else random.choice(["active", "inactive", "pending"]), "created_at": (datetime.now() - timedelta(days=random.randint(1, 365))).isoformat(), "balance": None if has_null else round(random.uniform(0, 10000), 2), } rows.append(row) return json.dumps({"rows": rows, "count": count}) def generate_log_entries(count: int, error_rate: float = 0.15) -> str: """Generate realistic log entries.""" services = ["api-gateway", "auth-service", "payment-service", "user-service"] entries = [] for i in range(count): if random.random() < error_rate: level = random.choice(["ERROR", "FATAL"]) message = random.choice( [ "Connection timeout to database", "Failed to process payment: insufficient funds", "Authentication failed for user", "Memory limit exceeded", "Unhandled exception in request handler", ] ) else: level = random.choice(["DEBUG", "INFO", "INFO", "INFO", "WARN"]) message = random.choice( [ "Request processed successfully", "Cache hit for user data", "Starting health check", "Connection pool recycled", "Metrics exported", ] ) entries.append( { "timestamp": (datetime.now() - timedelta(minutes=i)).isoformat(), "level": level, "service": random.choice(services), "message": message, "trace_id": f"trace_{random.randint(100000, 999999)}", } ) return json.dumps({"entries": entries}) def generate_github_issues(count: int, bug_rate: float = 0.2) -> str: """Generate realistic GitHub issues.""" labels_pool = ["enhancement", "documentation", "question", "good first issue"] bug_labels = ["bug", "critical", "urgent", "blocker"] issues = [] for i in range(count): is_bug = random.random() < bug_rate labels = ( random.sample(bug_labels, k=random.randint(1, 2)) if is_bug else random.sample(labels_pool, k=random.randint(0, 2)) ) issues.append( { "number": i + 1, "title": f"{'BUG: ' if is_bug else ''}{random.choice(['Fix login flow', 'Update API docs', 'Add dark mode', 'Improve performance'])}", "state": random.choice(["open", "closed"]), "labels": labels, "author": f"user{random.randint(1, 100)}", "created_at": (datetime.now() - timedelta(days=random.randint(1, 30))).isoformat(), "comments": random.randint(0, 20), } ) return json.dumps({"issues": issues, "total_count": count}) # ============================================================================ # Test Classes # ============================================================================ class TestMCPToolProfiles: """Test tool profile matching.""" def test_slack_profile_match(self, mcp_compressor): """Slack tools should match slack profile.""" profile = mcp_compressor.get_profile("mcp__slack__search") assert "slack" in profile.tool_name_pattern assert profile.max_items == 25 def test_database_profile_match(self, mcp_compressor): """Database tools should match database profile.""" profile = mcp_compressor.get_profile("mcp__database__query") assert "database" in profile.tool_name_pattern or "sql" in profile.tool_name_pattern assert profile.max_items == 30 def test_github_profile_match(self, mcp_compressor): """GitHub tools should match github profile.""" profile = mcp_compressor.get_profile("mcp__github__list_issues") assert "github" in profile.tool_name_pattern def test_log_profile_match(self, mcp_compressor): """Log tools should match log profile with higher max_items.""" profile = mcp_compressor.get_profile("search_logs") assert "log" in profile.tool_name_pattern assert profile.max_items == 40 # Logs get more items def test_fallback_profile(self, mcp_compressor): """Unknown tools should get fallback profile.""" profile = mcp_compressor.get_profile("some_random_tool") assert profile.max_items == 20 # Default class TestMCPCompressionBasics: """Test basic compression functionality.""" def test_compress_returns_result(self, mcp_compressor): """Compression should return MCPCompressionResult.""" content = generate_slack_messages(100) result = mcp_compressor.compress( content=content, tool_name="slack_search", user_query="find errors", ) assert isinstance(result, MCPCompressionResult) def test_compress_reduces_tokens(self, mcp_compressor): """Compression should reduce token count.""" content = generate_slack_messages(200) result = mcp_compressor.compress( content=content, tool_name="slack_search", user_query="find errors", ) assert result.compressed_tokens < result.original_tokens assert result.tokens_saved > 0 def test_compress_tracks_metrics(self, mcp_compressor): """Compression should track items before/after.""" content = generate_slack_messages(100) result = mcp_compressor.compress( content=content, tool_name="slack_search", ) assert result.items_before == 100 assert result.items_after is not None assert result.items_after < result.items_before def test_small_content_not_compressed(self, mcp_compressor): """Small content should not be compressed.""" content = generate_slack_messages(5) # Very small result = mcp_compressor.compress( content=content, tool_name="slack_search", ) assert result.was_compressed is False assert result.compressed_content == content class TestMCPErrorPreservation: """Test error preservation - SmartCrusher detects errors via field patterns.""" def test_all_log_errors_preserved(self, mcp_compressor): """100% of ERROR/FATAL log entries must be preserved. SmartCrusher detects errors via the 'level' field containing ERROR/FATAL. This is the strongest error preservation guarantee. """ random.seed(43) content = generate_log_entries(200, error_rate=0.25) data = json.loads(content) original_errors = [e for e in data["entries"] if e["level"] in ["ERROR", "FATAL"]] result = mcp_compressor.compress( content=content, tool_name="search_logs", user_query="find errors", ) compressed_data = json.loads(result.compressed_content) compressed_errors = [ e for e in strip_ccr_sentinels(compressed_data["entries"]) if e["level"] in ["ERROR", "FATAL"] ] # CRITICAL: 100% of errors must be preserved assert len(compressed_errors) >= len(original_errors), ( f"Lost errors: {len(original_errors)} -> {len(compressed_errors)}" ) def test_slack_significant_compression_with_content(self, mcp_compressor): """Slack messages should compress while preserving error keywords in text.""" random.seed(42) content = generate_slack_messages(200, error_rate=0.2) result = mcp_compressor.compress( content=content, tool_name="slack_search", user_query="find errors", ) # Should achieve significant compression assert result.compression_ratio > 0.5 compressed_data = json.loads(result.compressed_content) # Should preserve some messages with error keywords (SmartCrusher detects these) error_msgs = [ m for m in strip_ccr_sentinels(compressed_data["messages"]) if any(kw in m["text"].lower() for kw in ["error", "failed", "exception"]) ] assert len(error_msgs) > 0, "Should preserve some error messages" def test_database_error_status_preserved(self, mcp_compressor): """Database rows with ERROR status should be preserved.""" random.seed(44) content = generate_database_results(150, null_rate=0.15) data = json.loads(content) original_errors = [r for r in data["rows"] if "error" in str(r["status"]).lower()] result = mcp_compressor.compress( content=content, tool_name="database_query", user_query="find errors", ) compressed_data = json.loads(result.compressed_content) compressed_errors = [ r for r in strip_ccr_sentinels(compressed_data["rows"]) if "error" in str(r["status"]).lower() ] # Should preserve most error rows assert len(compressed_errors) >= len(original_errors) * 0.8, ( f"Lost too many errors: {len(original_errors)} -> {len(compressed_errors)}" ) def test_github_bugs_partial_preservation(self, mcp_compressor): """GitHub bug issues should have partial preservation.""" random.seed(45) content = generate_github_issues(100, bug_rate=0.3) result = mcp_compressor.compress( content=content, tool_name="github_issues", user_query="find bugs", ) compressed_data = json.loads(result.compressed_content) # Should preserve at least some bugs compressed_bugs = [ i for i in strip_ccr_sentinels(compressed_data["issues"]) if any(label in ["bug", "critical", "urgent", "blocker"] for label in i["labels"]) ] # At least 5 bugs should be preserved assert len(compressed_bugs) >= 5, "Should preserve at least 5 bug issues" class TestMCPStandaloneFunction: """Test the standalone compress_tool_result function.""" def test_standalone_returns_string(self): """Standalone function should return compressed string.""" content = generate_slack_messages(100) result = compress_tool_result( content=content, tool_name="slack_search", tool_args={"query": "errors"}, user_query="find errors in slack", ) assert isinstance(result, str) # Should be valid JSON json.loads(result) def test_standalone_with_metrics(self): """Standalone function with metrics should return MCPCompressionResult.""" content = generate_log_entries(150) result = compress_tool_result_with_metrics( content=content, tool_name="search_logs", tool_args={"service": "api"}, user_query="find errors", ) assert isinstance(result, MCPCompressionResult) assert result.tool_name == "search_logs" class TestMCPClientWrapper: """Test the async client wrapper.""" @pytest.fixture def mock_mcp_client(self): """Create a mock MCP client.""" class MockMCPClient: async def call_tool(self, name: str, arguments: dict | None = None) -> str: if "slack" in name: return generate_slack_messages(100) elif "log" in name: return generate_log_entries(150) else: return generate_database_results(80) return MockMCPClient() @pytest.mark.asyncio async def test_wrapper_compresses_automatically(self, mock_mcp_client): """Wrapper should automatically compress tool results.""" wrapper = HeadroomMCPClientWrapper(mock_mcp_client) result = await wrapper.call_tool("slack_search", {"query": "test"}) # Result should be valid JSON data = json.loads(result) # Should be compressed (fewer items) assert len(data["messages"]) < 100 @pytest.mark.asyncio async def test_wrapper_tracks_metrics(self, mock_mcp_client): """Wrapper should track compression metrics.""" wrapper = HeadroomMCPClientWrapper(mock_mcp_client) await wrapper.call_tool("slack_search", {"query": "test"}) await wrapper.call_tool("search_logs", {"service": "api"}) metrics = wrapper.get_metrics() assert len(metrics) == 2 assert metrics[0].tool_name == "slack_search" assert metrics[1].tool_name == "search_logs" @pytest.mark.asyncio async def test_wrapper_total_tokens_saved(self, mock_mcp_client): """Wrapper should track total tokens saved.""" wrapper = HeadroomMCPClientWrapper(mock_mcp_client) await wrapper.call_tool("slack_search", {"query": "test"}) await wrapper.call_tool("search_logs", {"service": "api"}) total_saved = wrapper.get_total_tokens_saved() assert total_saved > 0 class TestMCPCompressionRatio: """Test compression efficiency.""" def test_significant_compression_slack(self, mcp_compressor): """Slack results should compress well (>50%).""" content = generate_slack_messages(200) result = mcp_compressor.compress( content=content, tool_name="slack_search", ) assert result.compression_ratio > 0.5, ( f"Compression ratio too low: {result.compression_ratio:.2%}" ) def test_significant_compression_logs(self, mcp_compressor): """Log entries should compress well (>50%).""" content = generate_log_entries(200) result = mcp_compressor.compress( content=content, tool_name="search_logs", ) assert result.compression_ratio > 0.5, ( f"Compression ratio too low: {result.compression_ratio:.2%}" ) def test_compression_efficiency_increases_with_size(self, mcp_compressor): """Larger outputs should compress more efficiently.""" small = generate_slack_messages(50) large = generate_slack_messages(500) small_result = mcp_compressor.compress(small, "slack_search") large_result = mcp_compressor.compress(large, "slack_search") # Large should have higher compression ratio assert large_result.compression_ratio >= small_result.compression_ratio class TestMCPSchemaPreservation: """Test that JSON schema is preserved.""" def test_schema_preserved_slack(self, mcp_compressor): """Slack message schema should be preserved.""" content = generate_slack_messages(100) result = mcp_compressor.compress(content, "slack_search") compressed_data = json.loads(result.compressed_content) assert "messages" in compressed_data assert len(compressed_data["messages"]) > 0 # Check first message has all fields msg = compressed_data["messages"][0] assert "id" in msg assert "channel" in msg assert "user" in msg assert "text" in msg def test_schema_preserved_logs(self, mcp_compressor): """Log entry schema should be preserved.""" content = generate_log_entries(100) result = mcp_compressor.compress(content, "search_logs") compressed_data = json.loads(result.compressed_content) assert "entries" in compressed_data assert len(compressed_data["entries"]) > 0 entry = compressed_data["entries"][0] assert "timestamp" in entry assert "level" in entry assert "service" in entry assert "message" in entry class TestMCPEdgeCases: """Test edge cases.""" def test_empty_array(self, mcp_compressor): """Empty array should pass through unchanged.""" content = json.dumps({"messages": []}) result = mcp_compressor.compress(content, "slack_search") assert result.was_compressed is False def test_single_item(self, mcp_compressor): """Single item should pass through unchanged.""" content = json.dumps({"messages": [{"id": 1, "text": "test"}]}) result = mcp_compressor.compress(content, "slack_search") compressed = json.loads(result.compressed_content) assert len(compressed["messages"]) == 1 def test_non_json_passthrough(self, mcp_compressor): """Non-JSON content should pass through unchanged.""" content = "This is plain text, not JSON" result = mcp_compressor.compress(content, "some_tool") assert result.compressed_content == content assert result.was_compressed is False def test_malformed_json_passthrough(self, mcp_compressor): """Malformed JSON should pass through unchanged.""" content = '{"messages": [broken json' result = mcp_compressor.compress(content, "slack_search") assert result.compressed_content == content assert result.was_compressed is False class TestMCPContextUsage: """Test context extraction for relevance.""" def test_context_from_user_query(self, mcp_compressor): """User query should be used for context.""" content = generate_slack_messages(100) result = mcp_compressor.compress( content=content, tool_name="slack_search", user_query="find authentication errors", ) assert "authentication errors" in result.context_used def test_context_from_tool_args(self, mcp_compressor): """Tool args should be included in context.""" content = generate_slack_messages(100) result = mcp_compressor.compress( content=content, tool_name="slack_search", tool_args={"channel": "#incidents", "query": "outage"}, ) assert "incidents" in result.context_used or "outage" in result.context_used def test_combined_context(self, mcp_compressor): """Both user query and tool args should be combined.""" content = generate_slack_messages(100) result = mcp_compressor.compress( content=content, tool_name="slack_search", tool_args={"channel": "#alerts"}, user_query="find database errors", ) assert "database errors" in result.context_used class TestMCPCustomProfiles: """Test custom tool profiles.""" def test_custom_profile(self): """Custom profiles should override defaults.""" custom_profiles = [ MCPToolProfile( tool_name_pattern=r".*custom.*", max_items=10, min_tokens_to_compress=100, ), ] compressor = HeadroomMCPCompressor(profiles=custom_profiles) profile = compressor.get_profile("custom_tool") assert profile.max_items == 10 def test_profile_disabled(self): """Disabled profiles should not compress.""" custom_profiles = [ MCPToolProfile( tool_name_pattern=r".*", enabled=False, ), ] compressor = HeadroomMCPCompressor(profiles=custom_profiles) content = generate_slack_messages(200) result = compressor.compress(content, "any_tool") assert result.was_compressed is False