Yellhorn MCP

"""Tests for token limit enforcement in context fetching and workplan processing."""

from pathlib import Path
from unittest.mock import AsyncMock, MagicMock, patch

import pytest

from yellhorn_mcp.formatters.context_fetcher import apply_token_limit, get_codebase_context
from yellhorn_mcp.utils.token_utils import TokenCounter


class TestTokenLimitEnforcement:
    """Test suite for token limit enforcement with safety margins."""

    @pytest.mark.asyncio
    async def test_get_codebase_context_applies_safety_margin(self):
        """Test that get_codebase_context applies 10% safety margin to token limit."""
        repo_path = Path("/test/repo")
        token_limit = 10000
        model = "gpt-4o"

        # Create mock file content that will exceed the limit
        large_content = "x" * 50000  # Large content to trigger truncation
        mock_files = ["file1.py", "file2.py", "file3.py"]
        mock_file_contents = {f: large_content for f in mock_files}

        with patch(
            "yellhorn_mcp.formatters.context_fetcher.get_codebase_snapshot"
        ) as mock_snapshot:
            mock_snapshot.return_value = (mock_files, mock_file_contents)

            with patch(
                "yellhorn_mcp.formatters.context_fetcher.format_codebase_for_prompt"
            ) as mock_format:
                # Return very large content that needs truncation
                mock_format.return_value = "x" * 100000

                # Capture log messages
                log_messages = []

                def log_func(msg):
                    log_messages.append(msg)

                # Call with token limit
                content, paths = await get_codebase_context(
                    repo_path, "full", log_function=log_func, token_limit=token_limit, model=model
                )

                # Verify context truncation happened
                assert any("Context exceeds token limit" in msg for msg in log_messages)

                # Verify content was truncated
                token_counter = TokenCounter()
                actual_tokens = token_counter.count_tokens(content, model)

                # Should be within the effective limit (9000 tokens)
                # We now reserve 50 tokens for truncation notice, so effective limit is 8950
                assert (
                    actual_tokens <= 9000
                ), f"Content has {actual_tokens} tokens, should be <= 9000"

    @pytest.mark.asyncio
    async def test_get_codebase_context_without_token_limit(self):
        """Test that get_codebase_context works normally without token limit."""
        repo_path = Path("/test/repo")
        model = "gpt-4o"

        mock_files = ["file1.py", "file2.py"]
        mock_file_contents = {"file1.py": "content1", "file2.py": "content2"}

        with patch(
            "yellhorn_mcp.formatters.context_fetcher.get_codebase_snapshot"
        ) as mock_snapshot:
            mock_snapshot.return_value = (mock_files, mock_file_contents)

            with patch(
                "yellhorn_mcp.formatters.context_fetcher.format_codebase_for_prompt"
            ) as mock_format:
                expected_content = "formatted content"
                mock_format.return_value = expected_content

                # Call without token limit
                content, paths = await get_codebase_context(
                    repo_path, "full", token_limit=None, model=model
                )

                # Should return full content without truncation
                assert content == expected_content
                assert paths == mock_files

    @pytest.mark.asyncio
    async def test_get_codebase_context_validates_model_param(self):
        """Test that get_codebase_context requires model when token_limit is set."""
        repo_path = Path("/test/repo")

        # Should raise ValueError when token_limit is set but model is not
        with pytest.raises(
            ValueError, match="Model name is required when token_limit is specified"
        ):
            await get_codebase_context(repo_path, "full", token_limit=10000, model=None)

    def test_apply_token_limit_file_based_truncation(self):
        """Test that apply_token_limit truncates by complete files when possible."""
        model = "gpt-4o"
        token_limit = 200  # Small limit to force truncation

        # Create content with multiple files
        file_paths = ["file1.py", "file2.py", "file3.py", "file4.py"]
        file_contents = {
            "file1.py": "small content",
            "file2.py": "x" * 2000,  # Large file
            "file3.py": "x" * 3000,  # Very large file
            "file4.py": "small content 2",
        }

        # Construct formatted content
        content = ""
        for path in file_paths:
            content += f"\n--- File: {path} ---\n{file_contents[path]}\n"

        log_messages = []

        def log_func(msg):
            log_messages.append(msg)

        # Apply token limit
        truncated, included_paths = apply_token_limit(
            content,
            token_limit,
            model,
            log_func,
            file_paths=file_paths,
            file_contents=file_contents,
        )

        # Verify truncation occurred
        assert "Content truncated due to token limit" in truncated
        assert len(included_paths) < len(file_paths)

        # Verify complete files were included (not partial)
        for path in included_paths:
            assert f"--- File: {path} ---" in truncated

        # Verify token count is within limit (with reserved buffer)
        token_counter = TokenCounter()
        actual_tokens = token_counter.count_tokens(truncated, model)
        assert actual_tokens <= token_limit

    def test_apply_token_limit_character_based_fallback(self):
        """Test that apply_token_limit falls back to character truncation without file info."""
        model = "gpt-4o"
        token_limit = 100  # Small limit to force truncation

        # Create content without file structure
        content = "x" * 10000  # Large content without file markers

        log_messages = []

        def log_func(msg):
            log_messages.append(msg)

        # Apply token limit without file information
        truncated, included_paths = apply_token_limit(
            content, token_limit, model, log_func, file_paths=None, file_contents=None
        )

        # Verify truncation occurred
        assert "Content truncated due to token limit" in truncated
        assert len(truncated) < len(content)
        assert included_paths == []  # No file paths when using character truncation

        # Verify token count is within limit
        token_counter = TokenCounter()
        actual_tokens = token_counter.count_tokens(truncated, model)
        # Should be within limit since we reserve 50 tokens for truncation notice
        assert actual_tokens <= token_limit

        # Verify truncation message was logged
        assert any("Context truncated from" in msg for msg in log_messages)

    def test_apply_token_limit_no_truncation_needed(self):
        """Test that apply_token_limit returns content unchanged when within limit."""
        model = "gpt-4o"
        token_limit = 10000

        # Create small content that fits within limit
        content = "This is a small content that fits within the token limit."
        file_paths = ["file1.py"]
        file_contents = {"file1.py": content}

        log_messages = []

        def log_func(msg):
            log_messages.append(msg)

        # Apply token limit
        result_content, result_paths = apply_token_limit(
            content,
            token_limit,
            model,
            log_func,
            file_paths=file_paths,
            file_contents=file_contents,
        )

        # Content should be unchanged
        assert result_content == content
        assert result_paths == file_paths

        # No truncation message should be logged
        assert not any("truncating" in msg.lower() for msg in log_messages)

    @pytest.mark.asyncio
    async def test_get_codebase_context_lsp_mode_with_token_limit(self):
        """Test token limit enforcement in LSP mode."""
        repo_path = Path("/test/repo")
        token_limit = 5000
        model = "gemini-2.0-flash-exp"

        mock_files = ["file1.py", "file2.py"]

        with patch(
            "yellhorn_mcp.formatters.context_fetcher.get_codebase_snapshot"
        ) as mock_snapshot:
            mock_snapshot.return_value = (mock_files, {})

            with patch("yellhorn_mcp.formatters.context_fetcher.get_lsp_snapshot") as mock_lsp:
                # Return large LSP content
                lsp_content = {f: "x" * 5000 for f in mock_files}
                mock_lsp.return_value = (mock_files, lsp_content)

                with patch(
                    "yellhorn_mcp.formatters.context_fetcher.format_codebase_for_prompt"
                ) as mock_format:
                    mock_format.return_value = "x" * 20000  # Large content needing truncation

                    log_messages = []

                    def log_func(msg):
                        log_messages.append(msg)

                    content, paths = await get_codebase_context(
                        repo_path,
                        "lsp",
                        log_function=log_func,
                        token_limit=token_limit,
                        model=model,
                    )

                    # Verify truncation if needed
                    if len(content) < len("x" * 20000):
                        assert any("Context exceeds token limit" in msg for msg in log_messages)

                    # Verify content was truncated to within limit
                    token_counter = TokenCounter()
                    actual_tokens = token_counter.count_tokens(content, model)
                    assert actual_tokens <= 5000  # Should be within original limit

    @pytest.mark.asyncio
    async def test_get_codebase_context_file_structure_mode_with_token_limit(self):
        """Test token limit enforcement in file_structure mode."""
        repo_path = Path("/test/repo")
        token_limit = 2000
        model = "gpt-4o-mini"

        # Create many files to exceed limit
        mock_files = [f"dir{i}/file{j}.py" for i in range(10) for j in range(10)]

        with patch(
            "yellhorn_mcp.formatters.context_fetcher.get_codebase_snapshot"
        ) as mock_snapshot:
            mock_snapshot.return_value = (mock_files, {})

            with patch(
                "yellhorn_mcp.formatters.context_fetcher.build_file_structure_context"
            ) as mock_build:
                # Return large tree structure
                mock_build.return_value = "\n".join([f"├── {f}" for f in mock_files]) * 10

                log_messages = []

                def log_func(msg):
                    log_messages.append(msg)

                content, paths = await get_codebase_context(
                    repo_path,
                    "file_structure",
                    log_function=log_func,
                    token_limit=token_limit,
                    model=model,
                )

                # Verify truncation if the content was too large
                if "Content truncated due to token limit" in content:
                    assert any("Context exceeds token limit" in msg for msg in log_messages)

                # Verify truncation occurred
                assert "Content truncated due to token limit" in content

    @pytest.mark.asyncio
    async def test_get_codebase_context_none_mode_returns_empty(self):
        """Test that 'none' mode returns empty content regardless of token limit."""
        repo_path = Path("/test/repo")

        content, paths = await get_codebase_context(
            repo_path, "none", token_limit=1000, model="gpt-4o"
        )

        assert content == ""
        assert paths == []


class TestWorkplanTokenLimits:
    """Test token limit handling in workplan processing."""

    @pytest.mark.asyncio
    async def test_workplan_respects_codebase_token_limit(self):
        """Test that workplan processor calculates appropriate token limit for codebase."""
        from yellhorn_mcp.processors.workplan_processor import process_workplan_async

        repo_path = Path("/test/repo")
        model = "gpt-4o"  # 128k limit

        # Create mocks
        mock_llm_manager = MagicMock()
        mock_llm_manager._is_openai_model.return_value = True
        mock_llm_manager.call_llm_with_usage = AsyncMock(
            return_value={
                "content": "Generated workplan",
                "usage_metadata": MagicMock(
                    prompt_tokens=1000, completion_tokens=500, total_tokens=1500
                ),
            }
        )

        with patch(
            "yellhorn_mcp.processors.workplan_processor.get_codebase_context"
        ) as mock_get_context:
            mock_get_context.return_value = ("codebase content", ["file1.py"])

            with patch(
                "yellhorn_mcp.processors.workplan_processor.update_issue_with_workplan"
            ) as mock_update:
                mock_update.return_value = None

                with patch(
                    "yellhorn_mcp.processors.workplan_processor.add_issue_comment"
                ) as mock_comment:
                    mock_comment.return_value = None

                    await process_workplan_async(
                        repo_path=repo_path,
                        llm_manager=mock_llm_manager,
                        model=model,
                        title="Test Workplan",
                        issue_number="123",
                        codebase_reasoning="full",
                        detailed_description="Test description",
                        debug=False,
                        disable_search_grounding=True,
                        github_command_func=AsyncMock(),
                        git_command_func=AsyncMock(),
                    )

                    # Verify get_codebase_context was called with appropriate token limit
                    mock_get_context.assert_called_once()
                    call_args = mock_get_context.call_args

                    # Should have token_limit set
                    assert call_args[1]["token_limit"] is not None

                    # Token limit should be 70% of (model_limit - 5500)
                    # For gpt-4o: (128000 - 5500) * 0.7 = 85750
                    expected_limit = int((128000 - 5500) * 0.7)
                    assert call_args[1]["token_limit"] == expected_limit
                    assert call_args[1]["model"] == model

    @pytest.mark.asyncio
    async def test_revision_respects_codebase_token_limit(self):
        """Test that revision processor calculates appropriate token limit for codebase."""
        from yellhorn_mcp.processors.workplan_processor import process_revision_async

        repo_path = Path("/test/repo")
        model = "gemini-2.0-flash-exp"  # 1M limit

        # Create mocks
        mock_llm_manager = MagicMock()
        mock_llm_manager._is_openai_model.return_value = False
        mock_llm_manager.call_llm_with_citations = AsyncMock(
            return_value={
                "content": "Revised workplan",
                "usage_metadata": MagicMock(
                    prompt_tokens=2000, completion_tokens=1000, total_tokens=3000
                ),
                "grounding_metadata": None,
            }
        )

        with patch(
            "yellhorn_mcp.processors.workplan_processor.get_codebase_context"
        ) as mock_get_context:
            mock_get_context.return_value = ("codebase content", ["file1.py"])

            with patch(
                "yellhorn_mcp.processors.workplan_processor.update_issue_with_workplan"
            ) as mock_update:
                mock_update.return_value = None

                with patch(
                    "yellhorn_mcp.processors.workplan_processor.add_issue_comment"
                ) as mock_comment:
                    mock_comment.return_value = None

                    await process_revision_async(
                        repo_path=repo_path,
                        llm_manager=mock_llm_manager,
                        model=model,
                        issue_number="123",
                        original_workplan="# Original\n## Summary\nOriginal workplan",
                        revision_instructions="Add more detail",
                        codebase_reasoning="lsp",
                        debug=False,
                        disable_search_grounding=True,
                        github_command_func=AsyncMock(),
                        git_command_func=AsyncMock(),
                    )

                    # Verify get_codebase_context was called with token limit
                    mock_get_context.assert_called_once()
                    call_args = mock_get_context.call_args

                    # Should have token_limit set
                    assert call_args[1]["token_limit"] is not None

                    # Token limit should be 70% of (model_limit - 5500)
                    # For gemini-2.0-flash-exp: (1048576 - 5500) * 0.7 = 730153
                    expected_limit = int((1048576 - 5500) * 0.7)
                    assert call_args[1]["token_limit"] == expected_limit
                    assert call_args[1]["model"] == model


class TestJudgementTokenLimits:
    """Test token limit handling in judgement processing."""

    @pytest.mark.asyncio
    async def test_judgement_handles_large_diffs(self):
        """Test that judgement processor handles large diffs appropriately."""
        from yellhorn_mcp.processors.judgement_processor import process_judgement_async

        repo_path = Path("/test/repo")
        model = "gpt-4o"

        # Create mocks
        mock_llm_manager = MagicMock()
        mock_llm_manager._is_openai_model.return_value = True
        mock_llm_manager.call_llm_with_usage = AsyncMock(
            return_value={
                "content": "Judgement content",
                "usage_metadata": MagicMock(
                    prompt_tokens=1000, completion_tokens=500, total_tokens=1500
                ),
            }
        )

        # Large diff content
        large_diff = "+" + "x" * 100000  # Simulate large diff

        with patch(
            "yellhorn_mcp.processors.judgement_processor.create_judgement_subissue"
        ) as mock_create:
            mock_create.return_value = "https://github.com/test/repo/issues/124"

            with patch(
                "yellhorn_mcp.processors.judgement_processor.add_issue_comment"
            ) as mock_comment:
                mock_comment.return_value = None

                # Should complete without token limit errors
                await process_judgement_async(
                    repo_path=repo_path,
                    llm_manager=mock_llm_manager,
                    model=model,
                    workplan_content="Original workplan",
                    diff_content=large_diff,
                    base_ref="main",
                    head_ref="feature",
                    base_commit_hash="abc123",
                    head_commit_hash="def456",
                    parent_workplan_issue_number="123",
                    debug=False,
                    codebase_reasoning="full",
                )

                # Verify LLM was called
                mock_llm_manager.call_llm_with_usage.assert_called_once()

                # Get the prompt that was passed
                call_args = mock_llm_manager.call_llm_with_usage.call_args
                prompt = call_args[1]["prompt"]

                # Prompt should contain the diff
                assert large_diff in prompt