my-pal-mcp-server/tests/test_conversation_history_bug.py

"""
Test suite for conversation history bug fix

This test verifies that the critical bug where conversation history
(including file context) was not included when using continuation_id
has been properly fixed.

The bug was that tools with continuation_id would not see previous
conversation turns, causing issues like Gemini not seeing files that
Claude had shared in earlier turns.
"""

import json
from unittest.mock import Mock, patch

import pytest
from pydantic import Field

from tools.base import BaseTool, ToolRequest
from utils.conversation_memory import ConversationTurn, ThreadContext


class FileContextRequest(ToolRequest):
    """Test request with file support"""

    prompt: str = Field(..., description="Test prompt")
    files: list[str] = Field(default_factory=list, description="Optional files")


class FileContextTool(BaseTool):
    """Test tool for file context verification"""

    def get_name(self) -> str:
        return "test_file_context"

    def get_description(self) -> str:
        return "Test tool for file context"

    def get_input_schema(self) -> dict:
        return {
            "type": "object",
            "properties": {
                "prompt": {"type": "string"},
                "files": {"type": "array", "items": {"type": "string"}},
                "continuation_id": {"type": "string", "required": False},
            },
        }

    def get_system_prompt(self) -> str:
        return "Test system prompt for file context"

    def get_request_model(self):
        return FileContextRequest

    async def prepare_prompt(self, request) -> str:
        # Simple prompt preparation that would normally read files
        # For this test, we're focusing on whether conversation history is included
        files_context = ""
        if request.files:
            files_context = f"\nFiles in current request: {', '.join(request.files)}"

        return f"System: {self.get_system_prompt()}\nUser: {request.prompt}{files_context}"


class TestConversationHistoryBugFix:
    """Test that conversation history is properly included with continuation_id"""

    def setup_method(self):
        self.tool = FileContextTool()

    @patch("tools.base.get_thread")
    @patch("tools.base.add_turn")
    async def test_conversation_history_included_with_continuation_id(self, mock_add_turn, mock_get_thread):
        """Test that conversation history (including file context) is included when using continuation_id"""

        # Create a thread context with previous turns including files
        thread_context = ThreadContext(
            thread_id="test-history-id",
            created_at="2023-01-01T00:00:00Z",
            last_updated_at="2023-01-01T00:02:00Z",
            tool_name="analyze",  # Started with analyze tool
            turns=[
                ConversationTurn(
                    role="assistant",
                    content="I've analyzed the authentication module and found several security issues.",
                    timestamp="2023-01-01T00:01:00Z",
                    tool_name="analyze",
                    files=["/src/auth.py", "/src/security.py"],  # Files from analyze tool
                ),
                ConversationTurn(
                    role="assistant",
                    content="The code review shows these files have critical vulnerabilities.",
                    timestamp="2023-01-01T00:02:00Z",
                    tool_name="codereview",
                    files=["/src/auth.py", "/tests/test_auth.py"],  # Files from codereview tool
                ),
            ],
            initial_context={"question": "Analyze authentication security"},
        )

        # Mock get_thread to return our test context
        mock_get_thread.return_value = thread_context
        # Mock add_turn to return success
        mock_add_turn.return_value = True

        # Mock the model to capture what prompt it receives
        captured_prompt = None

        with patch.object(self.tool, "create_model") as mock_create_model:
            mock_model = Mock()
            mock_response = Mock()
            mock_response.candidates = [
                Mock(
                    content=Mock(parts=[Mock(text="Response with conversation context")]),
                    finish_reason="STOP",
                )
            ]

            def capture_prompt(prompt):
                nonlocal captured_prompt
                captured_prompt = prompt
                return mock_response

            mock_model.generate_content.side_effect = capture_prompt
            mock_create_model.return_value = mock_model

            # Execute tool with continuation_id
            arguments = {
                "prompt": "What should we fix first?",
                "continuation_id": "test-history-id",
                "files": ["/src/utils.py"],  # New file for this turn
            }
            response = await self.tool.execute(arguments)

            # Verify response succeeded
            response_data = json.loads(response[0].text)
            assert response_data["status"] == "success"

            # Verify get_thread was called for history reconstruction
            mock_get_thread.assert_called_with("test-history-id")

            # Verify the prompt includes conversation history
            assert captured_prompt is not None

            # Check that conversation history is included
            assert "=== CONVERSATION HISTORY ===" in captured_prompt
            assert "Turn 1 (Gemini using analyze)" in captured_prompt
            assert "Turn 2 (Gemini using codereview)" in captured_prompt

            # Check that file context from previous turns is included
            assert "📁 Files referenced: /src/auth.py, /src/security.py" in captured_prompt
            assert "📁 Files referenced: /src/auth.py, /tests/test_auth.py" in captured_prompt

            # Check that previous turn content is included
            assert "I've analyzed the authentication module and found several security issues." in captured_prompt
            assert "The code review shows these files have critical vulnerabilities." in captured_prompt

            # Check that continuation instruction is included
            assert "Continue this conversation by building on the previous context." in captured_prompt

            # Check that current request is still included
            assert "What should we fix first?" in captured_prompt
            assert "Files in current request: /src/utils.py" in captured_prompt

    @patch("tools.base.get_thread")
    async def test_no_history_when_thread_not_found(self, mock_get_thread):
        """Test graceful handling when thread is not found"""

        # Mock get_thread to return None (thread not found)
        mock_get_thread.return_value = None

        captured_prompt = None

        with patch.object(self.tool, "create_model") as mock_create_model:
            mock_model = Mock()
            mock_response = Mock()
            mock_response.candidates = [
                Mock(
                    content=Mock(parts=[Mock(text="Response without history")]),
                    finish_reason="STOP",
                )
            ]

            def capture_prompt(prompt):
                nonlocal captured_prompt
                captured_prompt = prompt
                return mock_response

            mock_model.generate_content.side_effect = capture_prompt
            mock_create_model.return_value = mock_model

            # Execute tool with continuation_id for non-existent thread
            arguments = {"prompt": "Test without history", "continuation_id": "non-existent-thread-id"}
            response = await self.tool.execute(arguments)

            # Should still succeed but without history
            response_data = json.loads(response[0].text)
            assert response_data["status"] == "success"

            # Verify get_thread was called for non-existent thread
            mock_get_thread.assert_called_with("non-existent-thread-id")

            # Verify the prompt does NOT include conversation history
            assert captured_prompt is not None
            assert "=== CONVERSATION HISTORY ===" not in captured_prompt
            assert "Test without history" in captured_prompt

    async def test_no_history_for_new_conversations(self):
        """Test that new conversations (no continuation_id) don't get history"""

        captured_prompt = None

        with patch.object(self.tool, "create_model") as mock_create_model:
            mock_model = Mock()
            mock_response = Mock()
            mock_response.candidates = [
                Mock(
                    content=Mock(parts=[Mock(text="New conversation response")]),
                    finish_reason="STOP",
                )
            ]

            def capture_prompt(prompt):
                nonlocal captured_prompt
                captured_prompt = prompt
                return mock_response

            mock_model.generate_content.side_effect = capture_prompt
            mock_create_model.return_value = mock_model

            # Execute tool without continuation_id (new conversation)
            arguments = {"prompt": "Start new conversation", "files": ["/src/new_file.py"]}
            response = await self.tool.execute(arguments)

            # Should succeed (may offer continuation for new conversations)
            response_data = json.loads(response[0].text)
            assert response_data["status"] in ["success", "continuation_available"]

            # Verify the prompt does NOT include conversation history
            assert captured_prompt is not None
            assert "=== CONVERSATION HISTORY ===" not in captured_prompt
            assert "Start new conversation" in captured_prompt
            assert "Files in current request: /src/new_file.py" in captured_prompt

            # Should include follow-up instructions for new conversation
            # (This is the existing behavior for new conversations)
            assert "If you'd like to ask a follow-up question" in captured_prompt


if __name__ == "__main__":
    pytest.main([__file__])