headroom/tests/test_strands_tokenizer.py

"""Tests for Strands SDK content block tokenization (#111).

Strands SDK sends content blocks without a "type" field:
  {"text": "..."} instead of {"type": "text", "text": "..."}
  {"toolUse": {...}} instead of {"type": "tool_use", ...}
  {"toolResult": {...}} instead of {"type": "tool_result", ...}

The tokenizer must count these correctly.
"""

from headroom.tokenizers import get_tokenizer


def _get_counter():
    return get_tokenizer("claude-sonnet-4-6")


class TestStrandsTextBlocks:
    """Strands text blocks: {"text": "..."} without "type" field."""

    def test_strands_text_matches_anthropic_text(self):
        """Strands {"text": ...} should count same as Anthropic {"type": "text", "text": ...}."""
        t = _get_counter()
        text = "Hello world this is a test message " * 50

        anthropic = [{"role": "user", "content": [{"type": "text", "text": text}]}]
        strands = [{"role": "user", "content": [{"text": text}]}]

        a = t.count_messages(anthropic)
        s = t.count_messages(strands)
        assert a == s, f"Anthropic={a}, Strands={s}"

    def test_strands_text_matches_plain_string(self):
        """Strands text block should count same as plain string content."""
        t = _get_counter()
        text = "Some question " * 1000

        plain = [{"role": "user", "content": text}]
        strands = [{"role": "user", "content": [{"text": text}]}]

        p = t.count_messages(plain)
        s = t.count_messages(strands)
        assert p == s, f"Plain={p}, Strands={s}"

    def test_strands_multiple_text_blocks(self):
        """Multiple Strands text blocks should all be counted."""
        t = _get_counter()

        msg = [
            {
                "role": "user",
                "content": [
                    {"text": "First block " * 100},
                    {"text": "Second block " * 100},
                ],
            }
        ]
        count = t.count_messages(msg)

        # Should be roughly 2x a single block
        single = [{"role": "user", "content": [{"text": "First block " * 100}]}]
        single_count = t.count_messages(single)

        assert count > single_count * 1.5, f"Multiple blocks={count}, single={single_count}"

    def test_strands_system_message(self):
        """System message with Strands text blocks."""
        t = _get_counter()
        text = "You are a helpful assistant. " * 200

        strands = [{"role": "system", "content": [{"text": text}]}]
        plain = [{"role": "system", "content": text}]

        s = t.count_messages(strands)
        p = t.count_messages(plain)
        assert s == p, f"Strands={s}, Plain={p}"


class TestStrandsToolBlocks:
    """Strands tool blocks: {"toolUse": {...}} and {"toolResult": {...}}."""

    def test_strands_tool_use_counted(self):
        """Strands toolUse block should be counted, not zero."""
        t = _get_counter()

        msg = [
            {
                "role": "assistant",
                "content": [
                    {
                        "toolUse": {
                            "toolUseId": "t1",
                            "name": "read_file",
                            "input": {"path": "/src/main.py"},
                        }
                    }
                ],
            }
        ]
        count = t.count_messages(msg)
        # Should include the tool name and input, not just message overhead
        assert count > 15, f"toolUse count too low: {count}"

    def test_strands_tool_result_counted(self):
        """Strands toolResult with nested text content should be counted."""
        t = _get_counter()
        big_content = "File contents here. " * 500

        msg = [
            {
                "role": "user",
                "content": [
                    {
                        "toolResult": {
                            "toolUseId": "t1",
                            "content": [{"text": big_content}],
                        }
                    }
                ],
            }
        ]
        count = t.count_messages(msg)

        # Should reflect the size of the content, not just overhead
        plain_count = t.count_messages([{"role": "user", "content": big_content}])
        assert count > plain_count * 0.5, (
            f"toolResult count={count} should be close to plain={plain_count}"
        )

    def test_strands_tool_result_string_content(self):
        """Strands toolResult with string content."""
        t = _get_counter()

        msg = [
            {
                "role": "user",
                "content": [
                    {
                        "toolResult": {
                            "toolUseId": "t1",
                            "content": "Simple string result " * 100,
                        }
                    }
                ],
            }
        ]
        count = t.count_messages(msg)
        assert count > 50, f"toolResult string count too low: {count}"


class TestMixedFormats:
    """Messages mixing Anthropic and Strands formats."""

    def test_mixed_conversation(self):
        """Full conversation with mixed Strands and Anthropic blocks."""
        t = _get_counter()
        messages = [
            # Strands system
            {"role": "system", "content": [{"text": "You are helpful. " * 50}]},
            # Strands user
            {"role": "user", "content": [{"text": "Fix the bug in auth.py"}]},
            # Strands assistant with toolUse
            {
                "role": "assistant",
                "content": [
                    {
                        "toolUse": {
                            "toolUseId": "t1",
                            "name": "read_file",
                            "input": {"path": "auth.py"},
                        }
                    }
                ],
            },
            # Strands tool result
            {
                "role": "user",
                "content": [
                    {
                        "toolResult": {
                            "toolUseId": "t1",
                            "content": [{"text": "def authenticate():\n    pass\n" * 100}],
                        }
                    }
                ],
            },
            # Anthropic-style text (for comparison)
            {"role": "assistant", "content": [{"type": "text", "text": "I found the issue."}]},
        ]
        count = t.count_messages(messages)
        # Should be substantial — the tool result alone is ~700 tokens
        assert count > 500, f"Mixed conversation count too low: {count}"