deer-flow/backend/tests/test_write_file_tool_budget_annotation.py

"""Tests for write_file tool budget annotation and assembly isolation.

Verifies fixes for reviewer findings on PR #5569:
- [P1] Guarded max_tokens extraction avoiding AttributeError on ModelConfig without max_tokens
- [P2] Cloning write_file tool to preserve module-level singleton immutability across assemblies
  and prevent cross-model guidance leakage.
"""

from types import SimpleNamespace
from unittest.mock import MagicMock

import pytest
from pydantic import BaseModel

from deerflow.config.app_config import AppConfig, ModelConfig, SandboxConfig, ToolConfig
from deerflow.sandbox.tools import write_file_tool
from deerflow.tools.tools import (
    _clone_tool_with_description,
    _extract_max_tokens,
    get_available_tools,
)


def test_extract_max_tokens_various_inputs():
    """Verify _extract_max_tokens safely handles all expected and edge-case inputs."""
    # None and empty
    assert _extract_max_tokens(None) is None

    # ModelConfig with and without max_tokens
    mc_without = ModelConfig(name="test", model="m", use="u")
    assert _extract_max_tokens(mc_without) is None

    mc_with = ModelConfig(name="test", model="m", use="u", max_tokens=4096)
    assert _extract_max_tokens(mc_with) == 4096

    mc_zero = ModelConfig(name="test", model="m", use="u", max_tokens=0)
    assert _extract_max_tokens(mc_zero) is None

    mc_neg = ModelConfig(name="test", model="m", use="u", max_tokens=-100)
    assert _extract_max_tokens(mc_neg) is None

    # Dictionaries
    assert _extract_max_tokens({}) is None
    assert _extract_max_tokens({"max_tokens": 2048}) == 2048
    assert _extract_max_tokens({"max_tokens": "8192"}) == 8192
    assert _extract_max_tokens({"max_tokens": None}) is None
    assert _extract_max_tokens({"max_tokens": 0}) is None

    # SimpleNamespace
    assert _extract_max_tokens(SimpleNamespace()) is None
    assert _extract_max_tokens(SimpleNamespace(max_tokens=1024)) == 1024

    # Booleans (must NOT be treated as 1 or 0)
    assert _extract_max_tokens({"max_tokens": True}) is None
    assert _extract_max_tokens({"max_tokens": False}) is None
    assert _extract_max_tokens(SimpleNamespace(max_tokens=True)) is None

    # Floats
    assert _extract_max_tokens({"max_tokens": 4096.0}) == 4096

    # Unparseable strings
    assert _extract_max_tokens({"max_tokens": "unlimited"}) is None

    # MagicMock (in Python unittest.mock, int(MagicMock()) defaults to 1; must be rejected)
    mock_without = MagicMock(spec=[])
    assert _extract_max_tokens(mock_without) is None

    mock_with = MagicMock()
    mock_with.max_tokens = 8000
    assert _extract_max_tokens(mock_with) == 8000


def test_clone_tool_with_description_preserves_singleton():
    """Verify _clone_tool_with_description returns an isolated copy and keeps original unchanged."""
    original_desc = write_file_tool.description
    assert "CUSTOM_TEST_BUDGET" not in original_desc

    cloned = _clone_tool_with_description(write_file_tool, original_desc + "\n\nCUSTOM_TEST_BUDGET")

    assert "CUSTOM_TEST_BUDGET" in cloned.description
    assert write_file_tool.description == original_desc
    assert cloned.name == write_file_tool.name
    assert cloned.func is write_file_tool.func
    assert cloned.coroutine is write_file_tool.coroutine
    assert cloned.args_schema is write_file_tool.args_schema
    assert isinstance(cloned, BaseModel)


def _build_minimal_app_config(models: list[ModelConfig]) -> AppConfig:
    return AppConfig(
        models=models,
        sandbox=SandboxConfig(use="deerflow.sandbox.local:LocalSandboxProvider"),
        tools=[
            ToolConfig(name="write_file", group="file:write", use="deerflow.sandbox.tools:write_file_tool"),
        ],
    )


def test_get_available_tools_with_model_config_lacking_max_tokens():
    """Verify get_available_tools does not raise AttributeError when max_tokens is omitted."""
    model_without_max_tokens = ModelConfig(name="capped-model", model="m", use="u")
    config = _build_minimal_app_config([model_without_max_tokens])

    tools = get_available_tools(model_name="capped-model", app_config=config, include_mcp=False)
    write_tool = next((t for t in tools if t.name == "write_file"), None)
    assert write_tool is not None
    assert "PER-RESPONSE BUDGET:" not in write_tool.description


def test_write_file_singleton_remains_unmutated_across_assemblies():
    """Verify write_file_tool process singleton is never mutated during tool assembly."""
    baseline_desc = write_file_tool.description

    model_with_budget = ModelConfig(name="budget-model", model="m", use="u", max_tokens=4096)
    config = _build_minimal_app_config([model_with_budget])

    tools = get_available_tools(model_name="budget-model", app_config=config, include_mcp=False)
    assembled_write_file = next(t for t in tools if t.name == "write_file")

    assert "PER-RESPONSE BUDGET: your output limit is 4096 tokens" in assembled_write_file.description
    # The process-wide singleton must remain pristine
    assert write_file_tool.description == baseline_desc
    assert "PER-RESPONSE BUDGET:" not in write_file_tool.description


def test_repeated_assembly_cross_model_isolation():
    """Verify consecutive tool assemblies for different models do not leak guidance or duplicate notes."""
    baseline_desc = write_file_tool.description

    model_4k = ModelConfig(name="model-4k", model="m", use="u", max_tokens=4096)
    model_32k = ModelConfig(name="model-32k", model="m", use="u", max_tokens=32768)
    model_none = ModelConfig(name="model-none", model="m", use="u")
    config = _build_minimal_app_config([model_4k, model_32k, model_none])

    # Assembly 1: 4K model
    tools_4k = get_available_tools(model_name="model-4k", app_config=config, include_mcp=False)
    wf_4k = next(t for t in tools_4k if t.name == "write_file")
    assert "your output limit is 4096 tokens" in wf_4k.description
    assert "32768" not in wf_4k.description
    assert write_file_tool.description == baseline_desc

    # Assembly 2: 32K model (must not carry 4096 note)
    tools_32k = get_available_tools(model_name="model-32k", app_config=config, include_mcp=False)
    wf_32k = next(t for t in tools_32k if t.name == "write_file")
    assert "your output limit is 32768 tokens" in wf_32k.description
    assert "4096" not in wf_32k.description
    assert wf_32k.description.count("PER-RESPONSE BUDGET:") == 1
    assert write_file_tool.description == baseline_desc

    # Assembly 3: model without max_tokens (must have NO budget note at all)
    tools_none = get_available_tools(model_name="model-none", app_config=config, include_mcp=False)
    wf_none = next(t for t in tools_none if t.name == "write_file")
    assert "PER-RESPONSE BUDGET:" not in wf_none.description
    assert wf_none.description == baseline_desc
    assert write_file_tool.description == baseline_desc


@pytest.mark.parametrize(
    ("profile_overrides", "agent_settings", "thinking_enabled", "bootstrap", "expected"),
    [
        ({}, {"max_tokens": 1024}, False, False, 1024),
        ({"when_thinking_enabled": {"max_tokens": 1024}}, {}, True, False, 1024),
        ({"when_thinking_disabled": {"max_tokens": 1024}}, {}, False, False, 1024),
        ({"when_thinking_enabled": {"max_tokens": 1024}}, {"max_tokens": 2048}, True, False, 1024),
        ({"when_thinking_disabled": {"max_tokens": None}}, {}, False, False, None),
        ({"when_thinking_enabled": {"max_tokens": 1024}}, {}, True, True, 1024),
        ({"when_thinking_disabled": {"max_tokens": 1024}}, {}, False, True, 1024),
    ],
    ids=["custom-agent", "thinking-on", "thinking-off", "thinking-over-agent", "uncapped", "bootstrap-thinking-on", "bootstrap-thinking-off"],
)
def test_lead_write_file_budget_matches_constructed_model(monkeypatch, profile_overrides, agent_settings, thinking_enabled, bootstrap, expected):
    """Exercise real model and tool assembly, including override precedence."""
    from deerflow.agents.lead_agent import agent as lead_agent_module
    from deerflow.config.agents_config import AgentConfig
    from deerflow.config.extensions_config import ExtensionsConfig

    model = ModelConfig(
        name="budget-model",
        model="budget-model",
        use="langchain_openai:ChatOpenAI",
        api_key="test-key",
        max_tokens=32768,
        supports_thinking=True,
        **profile_overrides,
    )
    app_config = _build_minimal_app_config([model])
    agent_config = AgentConfig(name="researcher", model="budget-model", model_settings=agent_settings)
    monkeypatch.setattr(lead_agent_module, "load_agent_config", lambda *args, **kwargs: agent_config)
    monkeypatch.setattr(lead_agent_module, "_load_enabled_available_skills", lambda *args, **kwargs: [])
    monkeypatch.setattr(lead_agent_module, "build_middlewares", lambda *args, **kwargs: [])
    monkeypatch.setattr(lead_agent_module, "apply_prompt_template", lambda **kwargs: "system prompt")
    monkeypatch.setattr(lead_agent_module, "create_agent", lambda **kwargs: kwargs)
    monkeypatch.setattr(lead_agent_module, "build_tracing_callbacks", lambda: [])
    monkeypatch.setattr(ExtensionsConfig, "from_file", lambda *args, **kwargs: ExtensionsConfig())

    graph = lead_agent_module._make_lead_agent(
        {"context": {"agent_name": "researcher", "thinking_enabled": thinking_enabled, "is_bootstrap": bootstrap}},
        app_config=app_config,
    )

    assert graph["model"].max_tokens == expected
    write_tool = next(tool for tool in graph["tools"] if tool.name == "write_file")
    if expected is None:
        assert "PER-RESPONSE BUDGET:" not in write_tool.description
    else:
        assert f"your output limit is {expected} tokens" in write_tool.description
        assert "32768 tokens" not in write_tool.description
    assert "PER-RESPONSE BUDGET:" not in write_file_tool.description