Gemini MCP Server

""" Test model metadata preservation during conversation continuation. This test verifies that when using continuation_id without specifying a model, the system correctly retrieves and uses the model from the previous conversation turn instead of defaulting to DEFAULT_MODEL or the custom provider's default. Bug: https://github.com/BeehiveInnovations/zen-mcp-server/issues/111 """ from unittest.mock import MagicMock, patch import pytest from server import reconstruct_thread_context from utils.conversation_memory import add_turn, create_thread, get_thread from utils.model_context import ModelContext class TestModelMetadataContinuation: """Test model metadata preservation during conversation continuation.""" @pytest.mark.asyncio async def test_model_preserved_from_previous_turn(self): """Test that model is correctly retrieved from previous conversation turn.""" # Create a thread with a turn that has a specific model thread_id = create_thread("chat", {"prompt": "test"}) # Add an assistant turn with a specific model success = add_turn( thread_id, "assistant", "Here's my response", model_name="deepseek-r1-8b", model_provider="custom" ) assert success # Test continuation without model should use previous turn's model arguments = {"continuation_id": thread_id} # No model specified # Mock dependencies to avoid side effects with patch("utils.model_context.ModelContext.calculate_token_allocation") as mock_calc: mock_calc.return_value = MagicMock( total_tokens=200000, content_tokens=160000, response_tokens=40000, file_tokens=64000, history_tokens=64000, ) with patch("utils.conversation_memory.build_conversation_history") as mock_build: mock_build.return_value = ("=== CONVERSATION HISTORY ===\n", 1000) # Call the actual function enhanced_args = await reconstruct_thread_context(arguments) # Verify model was retrieved from thread assert enhanced_args.get("model") == "deepseek-r1-8b" # Verify ModelContext would use the correct model model_context = ModelContext.from_arguments(enhanced_args) assert model_context.model_name == "deepseek-r1-8b" @pytest.mark.asyncio async def test_reconstruct_thread_context_preserves_model(self): """Test that reconstruct_thread_context preserves model from previous turn.""" # Create thread with assistant turn thread_id = create_thread("chat", {"prompt": "initial"}) add_turn(thread_id, "assistant", "Initial response", model_name="o3-mini", model_provider="openai") # Test reconstruction without specifying model arguments = {"continuation_id": thread_id, "prompt": "follow-up question"} # Mock the model context to avoid initialization issues in tests with patch("utils.model_context.ModelContext.calculate_token_allocation") as mock_calc: mock_calc.return_value = MagicMock( total_tokens=200000, content_tokens=160000, response_tokens=40000, file_tokens=64000, history_tokens=64000, ) with patch("utils.conversation_memory.build_conversation_history") as mock_build: mock_build.return_value = ("=== CONVERSATION HISTORY ===\n", 1000) enhanced_args = await reconstruct_thread_context(arguments) # Verify model was retrieved from thread assert enhanced_args.get("model") == "o3-mini" @pytest.mark.asyncio async def test_multiple_turns_uses_last_assistant_model(self): """Test that with multiple turns, the last assistant turn's model is used.""" thread_id = create_thread("chat", {"prompt": "analyze this"}) # Add multiple turns with different models add_turn(thread_id, "assistant", "First response", model_name="gemini-2.5-flash", model_provider="google") add_turn(thread_id, "user", "Another question") add_turn(thread_id, "assistant", "Second response", model_name="o3", model_provider="openai") add_turn(thread_id, "user", "Final question") arguments = {"continuation_id": thread_id} # Mock dependencies with patch("utils.model_context.ModelContext.calculate_token_allocation") as mock_calc: mock_calc.return_value = MagicMock( total_tokens=200000, content_tokens=160000, response_tokens=40000, file_tokens=64000, history_tokens=64000, ) with patch("utils.conversation_memory.build_conversation_history") as mock_build: mock_build.return_value = ("=== CONVERSATION HISTORY ===\n", 1000) # Call the actual function enhanced_args = await reconstruct_thread_context(arguments) # Should use the most recent assistant model assert enhanced_args.get("model") == "o3" @pytest.mark.asyncio async def test_no_previous_assistant_turn_defaults(self): """Test behavior when there's no previous assistant turn.""" # Save and set DEFAULT_MODEL for test import importlib import os original_default = os.environ.get("DEFAULT_MODEL", "") os.environ["DEFAULT_MODEL"] = "auto" import config import utils.model_context importlib.reload(config) importlib.reload(utils.model_context) try: thread_id = create_thread("chat", {"prompt": "test"}) # Only add user turns add_turn(thread_id, "user", "First question") add_turn(thread_id, "user", "Second question") arguments = {"continuation_id": thread_id} # Mock dependencies with patch("utils.model_context.ModelContext.calculate_token_allocation") as mock_calc: mock_calc.return_value = MagicMock( total_tokens=200000, content_tokens=160000, response_tokens=40000, file_tokens=64000, history_tokens=64000, ) with patch("utils.conversation_memory.build_conversation_history") as mock_build: mock_build.return_value = ("=== CONVERSATION HISTORY ===\n", 1000) # Call the actual function enhanced_args = await reconstruct_thread_context(arguments) # Should not have set a model assert enhanced_args.get("model") is None # ModelContext should use DEFAULT_MODEL model_context = ModelContext.from_arguments(enhanced_args) from config import DEFAULT_MODEL assert model_context.model_name == DEFAULT_MODEL finally: # Restore original value if original_default: os.environ["DEFAULT_MODEL"] = original_default else: os.environ.pop("DEFAULT_MODEL", None) importlib.reload(config) importlib.reload(utils.model_context) @pytest.mark.asyncio async def test_explicit_model_overrides_previous_turn(self): """Test that explicitly specifying a model overrides the previous turn's model.""" thread_id = create_thread("chat", {"prompt": "test"}) add_turn(thread_id, "assistant", "Response", model_name="gemini-2.5-flash", model_provider="google") arguments = {"continuation_id": thread_id, "model": "o3"} # Explicitly specified # Mock dependencies with patch("utils.model_context.ModelContext.calculate_token_allocation") as mock_calc: mock_calc.return_value = MagicMock( total_tokens=200000, content_tokens=160000, response_tokens=40000, file_tokens=64000, history_tokens=64000, ) with patch("utils.conversation_memory.build_conversation_history") as mock_build: mock_build.return_value = ("=== CONVERSATION HISTORY ===\n", 1000) # Call the actual function enhanced_args = await reconstruct_thread_context(arguments) # Should keep the explicit model assert enhanced_args.get("model") == "o3" @pytest.mark.asyncio async def test_thread_chain_model_preservation(self): """Test model preservation across thread chains (parent-child relationships).""" # Create parent thread parent_id = create_thread("chat", {"prompt": "analyze"}) add_turn(parent_id, "assistant", "Analysis", model_name="gemini-2.5-pro", model_provider="google") # Create child thread using a simple tool instead of workflow tool child_id = create_thread("chat", {"prompt": "review"}, parent_thread_id=parent_id) # Child thread should be able to access parent's model through chain traversal # NOTE: Current implementation only checks current thread (not parent threads) context = get_thread(child_id) assert context.parent_thread_id == parent_id arguments = {"continuation_id": child_id} # Mock dependencies with patch("utils.model_context.ModelContext.calculate_token_allocation") as mock_calc: mock_calc.return_value = MagicMock( total_tokens=200000, content_tokens=160000, response_tokens=40000, file_tokens=64000, history_tokens=64000, ) with patch("utils.conversation_memory.build_conversation_history") as mock_build: mock_build.return_value = ("=== CONVERSATION HISTORY ===\n", 1000) # Call the actual function enhanced_args = await reconstruct_thread_context(arguments) # No turns in child thread yet, so model should not be set assert enhanced_args.get("model") is None