dify/api/tests/unit_tests/services/test_variable_truncator.py

import functools
import json
import uuid
from typing import Any
from uuid import uuid4

import pytest

from graphon.file import File, FileTransferMethod, FileType
from graphon.nodes.variable_assigner.common.helpers import UpdatedVariable
from graphon.variables.segments import (
    ArrayFileSegment,
    ArrayNumberSegment,
    ArraySegment,
    FileSegment,
    FloatSegment,
    IntegerSegment,
    NoneSegment,
    ObjectSegment,
    StringSegment,
)
from graphon.variables.types import SegmentType
from services import variable_truncator as truncator_module
from services.variable_truncator import (
    DummyVariableTruncator,
    MaxDepthExceededError,
    TruncationResult,
    UnknownTypeError,
    VariableTruncator,
)
from tests.helpers.unit.config_override import apply_config_overrides

"""
Comprehensive unit tests for VariableTruncator class based on current implementation.

This test suite covers all functionality of the current VariableTruncator including:
- JSON size calculation for different data types
- String, array, and object truncation logic
- Segment-based truncation interface
- Helper methods for budget-based truncation
- Edge cases and error handling
"""


_compact_json_dumps = functools.partial(json.dumps, separators=(",", ":"))


class TestVariableTruncatorCases:
    @staticmethod
    @pytest.fixture
    def file() -> File:
        return File(
            file_id=str(uuid4()),  # Generate new UUID for File.id
            tenant_id=str(uuid.uuid4()),
            file_type=FileType.DOCUMENT,
            transfer_method=FileTransferMethod.LOCAL_FILE,
            related_id=str(uuid.uuid4()),
            filename="test_file.txt",
            extension=".txt",
            mime_type="text/plain",
            size=1024,
            storage_key="initial_key",
        )

    class TestCalculateJsonSize:
        """Test calculate_json_size method with different data types."""

        @pytest.fixture
        def truncator(self) -> VariableTruncator:
            return VariableTruncator()

        def test_string_size_calculation(self) -> None:
            """Test JSON size calculation for strings."""
            # Simple ASCII string
            assert VariableTruncator.calculate_json_size("hello") == 7  # "hello" + 2 quotes

            # Empty string
            assert VariableTruncator.calculate_json_size("") == 2  # Just quotes

            # Unicode string
            assert VariableTruncator.calculate_json_size("你好") == 4

        def test_number_size_calculation(self, truncator: VariableTruncator) -> None:
            """Test JSON size calculation for numbers."""
            assert truncator.calculate_json_size(123) == 3
            assert truncator.calculate_json_size(12.34) == 5
            assert truncator.calculate_json_size(-456) == 4
            assert truncator.calculate_json_size(0) == 1

        def test_boolean_size_calculation(self, truncator: VariableTruncator) -> None:
            """Test JSON size calculation for booleans."""
            assert truncator.calculate_json_size(True) == 4  # "true"
            assert truncator.calculate_json_size(False) == 5  # "false"

        def test_null_size_calculation(self, truncator: VariableTruncator) -> None:
            """Test JSON size calculation for None/null."""
            assert truncator.calculate_json_size(None) == 4  # "null"

        def test_array_size_calculation(self, truncator: VariableTruncator) -> None:
            """Test JSON size calculation for arrays."""
            # Empty array
            assert truncator.calculate_json_size([]) == 2  # "[]"

            # Simple array
            simple_array = [1, 2, 3]
            # [1,2,3] = 1 + 1 + 1 + 1 + 1 + 2 = 7 (numbers + commas + brackets)
            assert truncator.calculate_json_size(simple_array) == 7

            # Array with strings
            string_array = ["a", "b"]
            # ["a","b"] = 3 + 3 + 1 + 2 = 9 (quoted strings + comma + brackets)
            assert truncator.calculate_json_size(string_array) == 9

        def test_object_size_calculation(self, truncator: VariableTruncator) -> None:
            """Test JSON size calculation for objects."""
            # Empty object
            assert truncator.calculate_json_size({}) == 2  # "{}"

            # Simple object
            simple_obj = {"a": 1}
            # {"a":1} = 3 + 1 + 1 + 2 = 7 (key + colon + value + brackets)
            assert truncator.calculate_json_size(simple_obj) == 7

            # Multiple keys
            multi_obj = {"a": 1, "b": 2}
            # {"a":1,"b":2} = 3 + 1 + 1 + 1 + 3 + 1 + 1 + 2 = 13
            assert truncator.calculate_json_size(multi_obj) == 13

        def test_nested_structure_size_calculation(self, truncator: VariableTruncator) -> None:
            """Test JSON size calculation for nested structures."""
            nested = {"items": [1, 2, {"nested": "value"}]}
            size = truncator.calculate_json_size(nested)
            assert size > 0  # Should calculate without error

            # Verify it matches actual JSON length roughly

            actual_json = _compact_json_dumps(nested)
            # Should be close but not exact due to UTF-8 encoding considerations
            assert abs(size - len(actual_json.encode())) <= 5

        def test_calculate_json_size_max_depth_exceeded(self, truncator: VariableTruncator) -> None:
            """Test that calculate_json_size handles deep nesting gracefully."""
            # Create deeply nested structure
            nested: dict[str, object] = {"level": 0}
            current: dict[str, object] = nested
            for i in range(105):  # Create deep nesting
                nxt: dict[str, object] = {"level": i + 1}
                current["next"] = nxt
                current = nxt

            # Should either raise an error or handle gracefully
            with pytest.raises(MaxDepthExceededError):
                truncator.calculate_json_size(nested)

        def test_calculate_json_size_unknown_type(self, truncator: VariableTruncator) -> None:
            """Test that calculate_json_size raises error for unknown types."""

            class CustomType:
                pass

            with pytest.raises(UnknownTypeError):
                truncator.calculate_json_size(CustomType())

    class TestStringTruncation:
        LENGTH_LIMIT = 10
        """Test string truncation functionality."""

        @pytest.fixture
        def small_truncator(self) -> VariableTruncator:
            return VariableTruncator(string_length_limit=10)

        def test_short_string_no_truncation(self, small_truncator: VariableTruncator) -> None:
            """Test that short strings are not truncated."""
            short_str = "hello"
            result = small_truncator._truncate_string(short_str, self.LENGTH_LIMIT)
            assert result.value == short_str
            assert result.truncated is False
            assert result.value_size == VariableTruncator.calculate_json_size(short_str)

        def test_long_string_truncation(self, small_truncator: VariableTruncator) -> None:
            """Test that long strings are truncated with ellipsis."""
            long_str = "this is a very long string that exceeds the limit"
            result = small_truncator._truncate_string(long_str, self.LENGTH_LIMIT)

            assert result.truncated is True
            assert result.value == long_str[:5] + "..."
            assert result.value_size == 10  # 10 chars + "..."

        def test_exact_limit_string(self, small_truncator: VariableTruncator) -> None:
            """Test string exactly at limit."""
            exact_str = "1234567890"  # Exactly 10 chars
            result = small_truncator._truncate_string(exact_str, self.LENGTH_LIMIT)
            assert result.value == "12345..."
            assert result.truncated is True
            assert result.value_size == 10

    class TestArrayTruncation:
        """Test array truncation functionality."""

        @pytest.fixture
        def small_truncator(self) -> VariableTruncator:
            return VariableTruncator(array_element_limit=3, max_size_bytes=100)

        def test_small_array_no_truncation(self, small_truncator: VariableTruncator) -> None:
            """Test that small arrays are not truncated."""
            small_array: list[object] = [1, 2]
            result = small_truncator._truncate_array(small_array, 1000)
            assert result.value == small_array
            assert result.truncated is False

        def test_array_element_limit_truncation(self, small_truncator: VariableTruncator) -> None:
            """Test that arrays over element limit are truncated."""
            large_array: list[object] = [1, 2, 3, 4, 5, 6]  # Exceeds limit of 3
            result = small_truncator._truncate_array(large_array, 1000)

            assert result.truncated is True
            assert result.value == [1, 2, 3]

        def test_array_size_budget_truncation(self, small_truncator: VariableTruncator) -> None:
            """Test array truncation due to size budget constraints."""
            # Create array with strings that will exceed size budget
            large_strings: list[object] = ["very long string " * 5, "another long string " * 5]
            result = small_truncator._truncate_array(large_strings, 50)

            assert result.truncated is True
            # Should have truncated the strings within the array
            for item in result.value:
                assert isinstance(item, str)
            assert VariableTruncator.calculate_json_size(result.value) <= 50

        def test_array_with_nested_objects(self, small_truncator: VariableTruncator) -> None:
            """Test array truncation with nested objects."""
            nested_array: list[object] = [
                {"name": "item1", "data": "some data"},
                {"name": "item2", "data": "more data"},
                {"name": "item3", "data": "even more data"},
            ]
            result = small_truncator._truncate_array(nested_array, 30)

            assert isinstance(result.value, list)
            assert len(result.value) <= 3
            for item in result.value:
                assert isinstance(item, dict)

    class TestObjectTruncation:
        """Test object truncation functionality."""

        @pytest.fixture
        def small_truncator(self) -> VariableTruncator:
            return VariableTruncator(max_size_bytes=100)

        def test_small_object_no_truncation(self, small_truncator: VariableTruncator) -> None:
            """Test that small objects are not truncated."""
            small_obj = {"a": 1, "b": 2}
            result = small_truncator._truncate_object(small_obj, 1000)
            assert result.value == small_obj
            assert result.truncated is False

        def test_empty_object_no_truncation(self, small_truncator: VariableTruncator) -> None:
            """Test that empty objects are not truncated."""
            empty_obj = {}
            result = small_truncator._truncate_object(empty_obj, 100)
            assert result.value == empty_obj
            assert result.truncated is False

        def test_object_value_truncation(self, small_truncator: VariableTruncator) -> None:
            """Test object truncation where values are truncated to fit budget."""
            obj_with_long_values = {
                "key1": "very long string " * 10,
                "key2": "another long string " * 10,
                "key3": "third long string " * 10,
            }
            result = small_truncator._truncate_object(obj_with_long_values, 80)

            assert result.truncated is True
            assert isinstance(result.value, dict)

            assert set(result.value.keys()).issubset(obj_with_long_values.keys())

            # Values should be truncated if they exist
            for key, value in result.value.items():
                assert isinstance(value, str)
                original_value = obj_with_long_values[key]
                # Value should be same or smaller
                assert len(value) <= len(original_value)

        def test_object_key_dropping(self, small_truncator: VariableTruncator) -> None:
            """Test object truncation where keys are dropped due to size constraints."""
            large_obj = {f"key{i:02d}": f"value{i}" for i in range(20)}
            result = small_truncator._truncate_object(large_obj, 50)

            assert result.truncated is True
            assert len(result.value) < len(large_obj)

            # Should maintain sorted key order
            result_keys = list(result.value.keys())
            assert result_keys == sorted(result_keys)

        def test_object_with_nested_structures(self, small_truncator: VariableTruncator) -> None:
            """Test object truncation with nested arrays and objects."""
            nested_obj = {
                "simple": "value",
                "array": [1, 2, 3, 4, 5],
                "nested": {"inner": "data", "more": ["a", "b", "c"]},
            }
            result = small_truncator._truncate_object(nested_obj, 60)

            assert isinstance(result.value, dict)

    class TestSegmentBasedTruncation:
        """Test the main truncate method that works with Segments."""

        @pytest.fixture
        def truncator(self) -> VariableTruncator:
            return VariableTruncator()

        @pytest.fixture
        def small_truncator(self) -> VariableTruncator:
            return VariableTruncator(string_length_limit=20, array_element_limit=3, max_size_bytes=200)

        def test_integer_segment_no_truncation(self, truncator: VariableTruncator) -> None:
            """Test that integer segments are never truncated."""
            segment = IntegerSegment(value=12345)
            result = truncator.truncate(segment)

            assert isinstance(result, TruncationResult)
            assert result.truncated is False
            assert result.result == segment

        def test_boolean_as_integer_segment(self, truncator: VariableTruncator) -> None:
            """Test boolean values in IntegerSegment are converted to int."""
            segment = IntegerSegment(value=True)
            result = truncator.truncate(segment)

            assert isinstance(result, TruncationResult)
            assert result.truncated is False
            assert isinstance(result.result, IntegerSegment)
            assert result.result.value == 1  # True converted to 1

        def test_float_segment_no_truncation(self, truncator: VariableTruncator) -> None:
            """Test that float segments are never truncated."""
            segment = FloatSegment(value=123.456)
            result = truncator.truncate(segment)

            assert isinstance(result, TruncationResult)
            assert result.truncated is False
            assert result.result == segment

        def test_none_segment_no_truncation(self, truncator: VariableTruncator) -> None:
            """Test that None segments are never truncated."""
            segment = NoneSegment()
            result = truncator.truncate(segment)

            assert isinstance(result, TruncationResult)
            assert result.truncated is False
            assert result.result == segment

        def test_file_segment_no_truncation(self, truncator: VariableTruncator, file: File) -> None:
            """Test that file segments are never truncated."""
            file_segment = FileSegment(value=file)
            result = truncator.truncate(file_segment)
            assert result.result == file_segment
            assert result.truncated is False

        def test_array_file_segment_no_truncation(self, truncator: VariableTruncator, file: File) -> None:
            """Test that array file segments are never truncated."""

            array_file_segment = ArrayFileSegment(value=[file] * 20)
            result = truncator.truncate(array_file_segment)
            assert result.result == array_file_segment
            assert result.truncated is False

        def test_string_segment_small_no_truncation(self, truncator: VariableTruncator) -> None:
            """Test small string segments are not truncated."""
            segment = StringSegment(value="hello world")
            result = truncator.truncate(segment)

            assert isinstance(result, TruncationResult)
            assert result.truncated is False
            assert result.result == segment

        def test_string_segment_large_truncation(self, small_truncator: VariableTruncator) -> None:
            """Test large string segments are truncated."""
            long_text = "this is a very long string that will definitely exceed the limit"
            segment = StringSegment(value=long_text)
            result = small_truncator.truncate(segment)

            assert isinstance(result, TruncationResult)
            assert result.truncated is True
            assert isinstance(result.result, StringSegment)
            assert len(result.result.value) < len(long_text)
            assert result.result.value.endswith("...")

        def test_array_segment_small_no_truncation(self, truncator: VariableTruncator) -> None:
            """Test small array segments are not truncated."""
            from factories.variable_factory import build_segment

            segment = build_segment([1, 2, 3])
            result = truncator.truncate(segment)

            assert isinstance(result, TruncationResult)
            assert result.truncated is False
            assert result.result == segment

        def test_array_segment_large_truncation(self, small_truncator: VariableTruncator) -> None:
            """Test large array segments are truncated."""
            from factories.variable_factory import build_segment

            large_array = list(range(10))  # Exceeds element limit of 3
            segment = build_segment(large_array)
            result = small_truncator.truncate(segment)

            assert isinstance(result, TruncationResult)
            assert result.truncated is True
            assert isinstance(result.result, ArraySegment)
            assert len(result.result.value) <= 3

        def test_object_segment_small_no_truncation(self, truncator: VariableTruncator) -> None:
            """Test small object segments are not truncated."""
            segment = ObjectSegment(value={"key": "value"})
            result = truncator.truncate(segment)

            assert isinstance(result, TruncationResult)
            assert result.truncated is False
            assert result.result == segment

        def test_object_segment_large_truncation(self, small_truncator: VariableTruncator) -> None:
            """Test large object segments are truncated."""
            large_obj = {f"key{i}": f"very long value {i}" * 5 for i in range(5)}
            segment = ObjectSegment(value=large_obj)
            result = small_truncator.truncate(segment)

            assert isinstance(result, TruncationResult)
            assert result.truncated is True
            assert isinstance(result.result, ObjectSegment)
            # Object should be smaller or equal than original
            original_size = small_truncator.calculate_json_size(large_obj)
            result_size = small_truncator.calculate_json_size(result.result.value)
            assert result_size <= original_size

        def test_final_size_fallback_to_json_string(self, small_truncator: VariableTruncator) -> None:
            """Test final fallback when truncated result still exceeds size limit."""
            # Create data that will still be large after initial truncation
            large_nested_data = {"data": ["very long string " * 5] * 5, "more": {"nested": "content " * 20}}
            segment = ObjectSegment(value=large_nested_data)

            # Use very small limit to force JSON string fallback
            tiny_truncator = VariableTruncator(max_size_bytes=50)
            result = tiny_truncator.truncate(segment)

            assert isinstance(result, TruncationResult)
            assert result.truncated is True
            assert isinstance(result.result, StringSegment)
            # Should be JSON string with possible truncation
            assert len(result.result.value) <= 53  # 50 + "..." = 53

        def test_final_size_fallback_string_truncation(self, small_truncator: VariableTruncator) -> None:
            """Test final fallback for string that still exceeds limit."""
            # Create very long string that exceeds string length limit
            very_long_string = "x" * 6000  # Exceeds default string_length_limit of 5000
            segment = StringSegment(value=very_long_string)

            # Use small limit to test string fallback path
            tiny_truncator = VariableTruncator(string_length_limit=100, max_size_bytes=50)
            result = tiny_truncator.truncate(segment)

            assert isinstance(result, TruncationResult)
            assert result.truncated is True
            assert isinstance(result.result, StringSegment)
            # Should be truncated due to string limit or final size limit
            assert len(result.result.value) <= 1000  # Much smaller than original

    class TestEdgeCases:
        """Test edge cases and error conditions."""

        def test_empty_inputs(self) -> None:
            """Test truncator with empty inputs."""
            truncator = VariableTruncator()

            # Empty string
            result = truncator.truncate(StringSegment(value=""))
            assert not result.truncated
            assert result.result.value == ""

            # Empty array
            from factories.variable_factory import build_segment

            result = truncator.truncate(build_segment([]))
            assert not result.truncated
            assert result.result.value == []

            # Empty object
            result = truncator.truncate(ObjectSegment(value={}))
            assert not result.truncated
            assert result.result.value == {}

        def test_zero_and_negative_limits(self) -> None:
            """Test truncator behavior with zero or very small limits."""
            # Zero string limit
            with pytest.raises(ValueError, match="string_length_limit should be greater than 3\\."):
                VariableTruncator(string_length_limit=3)

            with pytest.raises(ValueError, match="array_element_limit should be greater than 0\\."):
                VariableTruncator(array_element_limit=0)

            with pytest.raises(ValueError, match="max_size_bytes should be greater than 0\\."):
                VariableTruncator(max_size_bytes=0)

        def test_unicode_and_special_characters(self) -> None:
            """Test truncator with unicode and special characters."""
            truncator = VariableTruncator(string_length_limit=10)

            # Unicode characters
            unicode_text = "你好世界你好世界你好世界"  # Multi-byte UTF-8 characters
            result = truncator.truncate(StringSegment(value=unicode_text))
            assert result.truncated is True

            # Special JSON characters
            special_chars = '{"key": "value with \\"quotes\\" and \\n newlines"}'
            result = truncator.truncate(StringSegment(value=special_chars))
            assert isinstance(result.result, StringSegment)

    class TestTruncateJsonPrimitives:
        """Test _truncate_json_primitives method with different data types."""

        @pytest.fixture
        def truncator(self) -> VariableTruncator:
            return VariableTruncator()

        def test_truncate_json_primitives_file_type(self, truncator: VariableTruncator, file: File) -> None:
            """Test that File objects are handled correctly in _truncate_json_primitives."""
            # Test File object is returned as-is without truncation
            result = truncator._truncate_json_primitives(file, 1000)

            assert result.value == file
            assert result.truncated is False
            # Size should be calculated correctly
            expected_size = VariableTruncator.calculate_json_size(file)
            assert result.value_size == expected_size

        def test_truncate_json_primitives_file_type_small_budget(
            self, truncator: VariableTruncator, file: File
        ) -> None:
            """Test that File objects are returned as-is even with small budget."""
            # Even with a small size budget, File objects should not be truncated
            result = truncator._truncate_json_primitives(file, 10)

            assert result.value == file
            assert result.truncated is False

        def test_truncate_json_primitives_file_type_in_array(self, truncator: VariableTruncator, file: File) -> None:
            """Test File objects in arrays are handled correctly."""
            array_with_files: list[object] = [file, file]
            result = truncator._truncate_json_primitives(array_with_files, 1000)

            assert isinstance(result.value, list)
            assert len(result.value) == 2
            assert result.value[0] == file
            assert result.value[1] == file
            assert result.truncated is False

        def test_truncate_json_primitives_file_type_in_object(self, truncator: VariableTruncator, file: File) -> None:
            """Test File objects in objects are handled correctly."""
            result = truncator._truncate_json_primitives({"file1": file, "file2": file}, 1000)

            assert isinstance(result.value, dict)
            assert len(result.value) == 2
            assert result.value["file1"] == file
            assert result.value["file2"] == file
            assert result.truncated is False

    class TestIntegrationScenarios:
        """Test realistic integration scenarios."""

        def test_workflow_output_scenario(self) -> None:
            """Test truncation of typical workflow output data."""
            truncator = VariableTruncator()

            workflow_data = {
                "result": "success",
                "data": {
                    "users": [
                        {"id": 1, "name": "Alice", "email": "alice@example.com"},
                        {"id": 2, "name": "Bob", "email": "bob@example.com"},
                    ]
                    * 3,  # Multiply to make it larger
                    "metadata": {
                        "count": 6,
                        "processing_time": "1.23s",
                        "details": "x" * 200,  # Long string but not too long
                    },
                },
            }

            segment = ObjectSegment(value=workflow_data)
            result = truncator.truncate(segment)

            assert isinstance(result, TruncationResult)
            assert isinstance(result.result, (ObjectSegment, StringSegment))
            # Should handle complex nested structure appropriately

        def test_large_text_processing_scenario(self) -> None:
            """Test truncation of large text data."""
            truncator = VariableTruncator(string_length_limit=100)

            large_text = "This is a very long text document. " * 20  # Make it larger than limit

            segment = StringSegment(value=large_text)
            result = truncator.truncate(segment)

            assert isinstance(result, TruncationResult)
            assert result.truncated is True
            assert isinstance(result.result, StringSegment)
            assert len(result.result.value) <= 103  # 100 + "..."
            assert result.result.value.endswith("...")

        def test_mixed_data_types_scenario(self) -> None:
            """Test truncation with mixed data types in complex structure."""
            truncator = VariableTruncator(string_length_limit=30, array_element_limit=3, max_size_bytes=300)

            mixed_data = {
                "strings": ["short", "medium length", "very long string " * 3],
                "numbers": [1, 2.5, 999999],
                "booleans": [True, False, True],
                "nested": {
                    "more_strings": ["nested string " * 2],
                    "more_numbers": list(range(5)),
                    "deep": {"level": 3, "content": "deep content " * 3},
                },
                "nulls": [None, None],
            }

            segment = ObjectSegment(value=mixed_data)
            result = truncator.truncate(segment)

            assert isinstance(result, TruncationResult)
            assert result.truncated is True
            assert isinstance(result.result, ObjectSegment)
            # Verify the result is smaller or equal than original
            original_size = truncator.calculate_json_size(mixed_data)
            result_size = truncator.calculate_json_size(result.result.value)
            assert result_size <= original_size

        def test_file_and_array_file_variable_mapping(self, file: File) -> None:
            truncator = VariableTruncator(string_length_limit=30, array_element_limit=3, max_size_bytes=300)

            mapping = {"array_file": [file]}
            truncated_mapping, truncated = truncator.truncate_variable_mapping(mapping)
            assert truncated is False
            assert truncated_mapping == mapping

    @staticmethod
    def test_dummy_variable_truncator_methods() -> None:
        """Test DummyVariableTruncator methods work correctly."""
        truncator = DummyVariableTruncator()

        # Test truncate_variable_mapping
        test_data: dict[str, Any] = {
            "key1": "value1",
            "key2": ["item1", "item2"],
            "large_array": list(range(2000)),
        }
        result, is_truncated = truncator.truncate_variable_mapping(test_data)

        assert result == test_data
        assert not is_truncated

        # Test truncate method
        segment = StringSegment(value="test string")
        result = truncator.truncate(segment)
        assert isinstance(result, TruncationResult)
        assert result.result == segment
        assert result.truncated is False

        segment = ArrayNumberSegment(value=list(range(2000)))
        result = truncator.truncate(segment)
        assert isinstance(result, TruncationResult)
        assert result.result == segment
        assert result.truncated is False

    # ---------------------------------------------------------------------------
    # Regression tests for langgenius/dify#39218.
    #
    # Before the fix, ``_truncate_array`` had a "Dirty fix" branch that
    # unconditionally appended every ``File`` element to ``truncated_value``
    # *before* the count cap and the byte-budget check, and *before*
    # ``used_size`` was ever incremented. That made ``list[File]`` arrays:
    #   1. uncapped by ``array_element_limit``,
    #   2. uncounted against ``max_size_bytes``, and
    #   3. always reported ``truncated=False``.
    # The fix routes ``File`` through ``_truncate_json_primitives``'s dedicated
    # ``File`` branch, which returns the file as-is with its real serialized
    # size, while preserving the count cap and the byte budget.
    # ---------------------------------------------------------------------------

    class TestFileArrayTruncationRegression39218:
        """``list[File]`` must respect ``array_element_limit`` and the byte budget."""

        @pytest.fixture
        def truncator(self) -> VariableTruncator:
            return VariableTruncator(
                array_element_limit=3,
                max_size_bytes=1000,
                string_length_limit=50,
            )

        @staticmethod
        def _make_file(name: str = "f") -> File:
            return File(
                id=name,
                type=FileType.DOCUMENT,
                transfer_method=FileTransferMethod.REMOTE_URL,
                remote_url=f"https://example.com/{name}.txt",
                filename=f"{name}.txt",
                extension=".txt",
                mime_type="text/plain",
                size=1024,
            )

        def test_file_array_respects_element_count_cap(self, truncator: VariableTruncator) -> None:
            # Use a target_size larger than ``count * file_size`` so the byte
            # budget never binds — only the count cap should fire.
            # Each File serializes to ~237 bytes; 3 files = ~713 bytes.
            files: list[object] = [self._make_file(f"f{i}") for i in range(500)]

            result = truncator._truncate_array(files, target_size=10_000_000)

            # Before the fix, all 500 File entries survived (``len(value)==500``,
            # ``truncated==False``). After the fix, the array is capped at
            # ``array_element_limit=3`` and ``truncated`` flips to True.
            assert len(result.value) == 3
            assert result.truncated is True

        def test_file_array_reports_real_used_size(self, truncator: VariableTruncator) -> None:
            # Large budget so the count cap fires before the byte budget does.
            files: list[object] = [self._make_file(f"f{i}") for i in range(500)]

            result = truncator._truncate_array(files, target_size=10_000_000)

            # Before the fix, ``used_size`` for a File array was the empty-array
            # baseline of 2 bytes (``[]``), regardless of how many File entries
            # actually returned. After the fix, ``used_size`` reflects the real
            # serialized size of the returned ``File`` payload.
            assert result.value_size > 100
            assert result.truncated is True

        def test_file_array_respects_byte_budget(self, truncator: VariableTruncator) -> None:
            # Use a small ``target_size`` so the byte budget is the binding
            # constraint. Each File serializes to ~237 bytes, so even one File
            # blows the 200-byte budget.
            files: list[object] = [self._make_file(f"f{i}") for i in range(50)]

            result = truncator._truncate_array(files, target_size=200)

            # Before the fix, all 50 File entries survived and ``used_size``
            # reported ``2`` (the empty-array baseline). After the fix, the
            # loop sees the File payload: ``value_size`` reflects the real
            # serialized size, and the loop stops after the first File because
            # adding the next one would exceed ``target_size``.
            assert len(result.value) == 1
            assert result.value_size > 100  # the File's real serialized size
            assert result.value_size <= 250  # in the ballpark of the budget

        def test_mixed_array_counts_files_toward_cap(self, truncator: VariableTruncator) -> None:
            mixed: list[object] = [
                self._make_file("f0"),
                "a",
                self._make_file("f1"),
                "b",
                self._make_file("f2"),
                "c",
                self._make_file("f3"),
                "d",
            ]

            result = truncator._truncate_array(mixed, target_size=10_000_000)

            # 8 items, cap of 3 → exactly 3 items. Files and primitives are
            # counted together toward the cap.
            assert len(result.value) == 3
            assert result.truncated is True

        def test_single_file_in_array_is_preserved(self, truncator: VariableTruncator) -> None:
            result = truncator._truncate_array([self._make_file("only")], target_size=10_000_000)

            # The File itself is not truncated — the dedicated ``File`` branch
            # in ``_truncate_json_primitives`` returns the file untouched. Only
            # the array-shape accounting changes.
            assert len(result.value) == 1
            assert isinstance(result.value[0], File)
            assert result.value[0].id == "only"
            assert result.truncated is False


class TestVariableTruncatorAdditionalCases:
    class TestVariableTruncatorAdditionalBehavior:
        def test_default_should_use_dify_config_limits(self, monkeypatch: pytest.MonkeyPatch) -> None:
            apply_config_overrides(
                monkeypatch,
                WORKFLOW_VARIABLE_TRUNCATION_MAX_SIZE=111,
                WORKFLOW_VARIABLE_TRUNCATION_ARRAY_LENGTH=7,
                WORKFLOW_VARIABLE_TRUNCATION_STRING_LENGTH=33,
            )

            truncator = VariableTruncator.default()

            assert truncator._max_size_bytes == 111
            assert truncator._array_element_limit == 7
            assert truncator._string_length_limit == 33

        def test_truncate_variable_mapping_should_mark_over_budget_keys_with_ellipsis(self) -> None:
            truncator = VariableTruncator(max_size_bytes=5)
            mapping = {"very_long_key": "value"}

            result, truncated = truncator.truncate_variable_mapping(mapping)

            assert result == {"very_long_key": "..."}
            assert truncated is True

        def test_truncate_variable_mapping_should_handle_segment_values(self) -> None:
            truncator = VariableTruncator(max_size_bytes=100)
            mapping = {"seg": StringSegment(value="hello")}

            result, truncated = truncator.truncate_variable_mapping(mapping)

            assert isinstance(result["seg"], StringSegment)
            assert result["seg"].value == "hello"
            assert truncated is False

        @pytest.mark.parametrize(
            ("value", "expected"),
            [
                (None, False),
                (True, False),
                (1, False),
                (1.5, False),
                ("x", True),
                ({"k": "v"}, True),
            ],
        )
        def test_json_value_needs_truncation_should_match_expected_rules(
            self,
            value: object,
            expected: bool,
        ) -> None:
            result = VariableTruncator._json_value_needs_truncation(value)
            assert result is expected

        def test_truncate_should_use_string_fallback_when_truncated_value_size_exceeds_limit(
            self,
            monkeypatch: pytest.MonkeyPatch,
        ) -> None:
            truncator = VariableTruncator(max_size_bytes=10)
            forced_result = truncator_module._PartResult(
                value=StringSegment(value="this is too long"),
                value_size=100,
                truncated=True,
            )
            monkeypatch.setattr(truncator, "_truncate_segment", lambda *_args, **_kwargs: forced_result)

            result = truncator.truncate(StringSegment(value="input"))

            assert result.truncated is True
            assert isinstance(result.result, StringSegment)
            assert not result.result.value.startswith('"')

        def test_truncate_segment_should_raise_assertion_for_unexpected_truncatable_segment(
            self,
            monkeypatch: pytest.MonkeyPatch,
        ) -> None:
            truncator = VariableTruncator()
            monkeypatch.setattr(VariableTruncator, "_segment_need_truncation", lambda _segment: True)

            with pytest.raises(AssertionError):
                truncator._truncate_segment(IntegerSegment(value=1), 10)

        def test_calculate_json_size_should_unwrap_segment_values(self) -> None:
            segment = StringSegment(value="abc")

            size = VariableTruncator.calculate_json_size(segment)

            assert size == VariableTruncator.calculate_json_size("abc")

        def test_calculate_json_size_should_handle_updated_variable_instances(self) -> None:
            updated = UpdatedVariable(name="n", selector=["node", "var"], value_type=SegmentType.STRING, new_value="v")

            size = VariableTruncator.calculate_json_size(updated)

            assert size > 0

        def test_maybe_qa_structure_should_validate_shape(self) -> None:
            assert VariableTruncator._maybe_qa_structure({"qa_chunks": []}) is True
            assert VariableTruncator._maybe_qa_structure({"qa_chunks": "not-list"}) is False
            assert VariableTruncator._maybe_qa_structure({}) is False

        def test_maybe_parent_child_structure_should_validate_shape(self) -> None:
            assert (
                VariableTruncator._maybe_parent_child_structure({"parent_mode": "full", "parent_child_chunks": []})
                is True
            )
            assert (
                VariableTruncator._maybe_parent_child_structure({"parent_mode": 1, "parent_child_chunks": []}) is False
            )
            assert (
                VariableTruncator._maybe_parent_child_structure({"parent_mode": "full", "parent_child_chunks": "bad"})
                is False
            )

        def test_truncate_object_should_truncate_segment_values_inside_object(self) -> None:
            truncator = VariableTruncator(string_length_limit=8, max_size_bytes=30)
            mapping = {"s": StringSegment(value="long-content")}

            result = truncator._truncate_object(mapping, 20)

            assert result.truncated is True
            assert isinstance(result.value["s"], StringSegment)

        def test_truncate_json_primitives_should_handle_updated_variable_input(self) -> None:
            truncator = VariableTruncator(max_size_bytes=100)
            updated = UpdatedVariable(name="n", selector=["node", "var"], value_type=SegmentType.STRING, new_value="v")

            result = truncator._truncate_json_primitives(updated, 100)

            assert isinstance(result.value, dict)

        def test_truncate_json_primitives_should_raise_assertion_for_unsupported_value_type(self) -> None:
            truncator = VariableTruncator()

            with pytest.raises(AssertionError):
                truncator._truncate_json_primitives(object(), 100)  # type: ignore[arg-type]

        def test_truncate_should_apply_json_string_fallback_for_large_non_string_segment(
            self,
            monkeypatch: pytest.MonkeyPatch,
        ) -> None:
            truncator = VariableTruncator(max_size_bytes=10)
            forced_segment = ObjectSegment(value={"k": "v"})
            forced_result = truncator_module._PartResult(value=forced_segment, value_size=100, truncated=True)
            monkeypatch.setattr(truncator, "_truncate_segment", lambda *_args, **_kwargs: forced_result)

            result = truncator.truncate(ObjectSegment(value={"a": "b"}))

            assert result.truncated is True
            assert isinstance(result.result, StringSegment)