browser-use--browser-use/browser_use/llm/tests/test_anthropic_cache.py

import logging
from typing import cast

from browser_use.agent.service import Agent
from browser_use.llm.anthropic.chat import ChatAnthropic
from browser_use.llm.anthropic.serializer import AnthropicMessageSerializer, NonSystemMessage
from browser_use.llm.messages import (
	AssistantMessage,
	BaseMessage,
	ContentPartImageParam,
	ContentPartTextParam,
	Function,
	ImageURL,
	SystemMessage,
	ToolCall,
	UserMessage,
)

logger = logging.getLogger(__name__)


class TestAnthropicCache:
	"""Comprehensive test for Anthropic cache serialization."""

	def test_cache_basic_functionality(self):
		"""Test basic cache functionality for all message types."""
		# Test cache with different message types
		messages: list[BaseMessage] = [
			SystemMessage(content='System message!', cache=True),
			UserMessage(content='User message!', cache=True),
			AssistantMessage(content='Assistant message!', cache=False),
		]

		anthropic_messages, system_message = AnthropicMessageSerializer.serialize_messages(messages)

		assert len(anthropic_messages) == 2
		assert isinstance(system_message, list)
		assert isinstance(anthropic_messages[0]['content'], list)
		assert isinstance(anthropic_messages[1]['content'], str)

		# Test cache with assistant message
		agent_messages: list[BaseMessage] = [
			SystemMessage(content='System message!'),
			UserMessage(content='User message!'),
			AssistantMessage(content='Assistant message!', cache=True),
		]

		anthropic_messages, system_message = AnthropicMessageSerializer.serialize_messages(agent_messages)

		assert isinstance(system_message, str)
		assert isinstance(anthropic_messages[0]['content'], str)
		assert isinstance(anthropic_messages[1]['content'], list)

	def test_cache_with_tool_calls(self):
		"""Test cache functionality with tool calls."""
		tool_call = ToolCall(id='test_id', function=Function(name='test_function', arguments='{"arg": "value"}'))

		# Assistant with tool calls and cache
		assistant_with_tools = AssistantMessage(content='Assistant with tools', tool_calls=[tool_call], cache=True)
		messages, _ = AnthropicMessageSerializer.serialize_messages([assistant_with_tools])

		assert len(messages) == 1
		assert isinstance(messages[0]['content'], list)
		# Should have both text and tool_use blocks
		assert len(messages[0]['content']) >= 2

	def test_cache_with_images(self):
		"""Test cache functionality with image content."""
		user_with_image = UserMessage(
			content=[
				ContentPartTextParam(text='Here is an image:', type='text'),
				ContentPartImageParam(image_url=ImageURL(url='https://example.com/image.jpg'), type='image_url'),
			],
			cache=True,
		)

		messages, _ = AnthropicMessageSerializer.serialize_messages([user_with_image])

		assert len(messages) == 1
		assert isinstance(messages[0]['content'], list)
		assert len(messages[0]['content']) == 2

	def test_cache_with_base64_images(self):
		"""Test cache functionality with base64 images."""
		base64_url = 'data:image/jpeg;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg=='

		user_with_base64 = UserMessage(
			content=[
				ContentPartTextParam(text='Base64 image:', type='text'),
				ContentPartImageParam(image_url=ImageURL(url=base64_url), type='image_url'),
			],
			cache=True,
		)

		messages, _ = AnthropicMessageSerializer.serialize_messages([user_with_base64])

		assert len(messages) == 1
		assert isinstance(messages[0]['content'], list)

	def test_cache_content_types(self):
		"""Test different content types with cache."""
		# String content with cache should become list
		user_string_cached = UserMessage(content='String message', cache=True)
		messages, _ = AnthropicMessageSerializer.serialize_messages([user_string_cached])
		assert isinstance(messages[0]['content'], list)

		# String content without cache should remain string
		user_string_no_cache = UserMessage(content='String message', cache=False)
		messages, _ = AnthropicMessageSerializer.serialize_messages([user_string_no_cache])
		assert isinstance(messages[0]['content'], str)

		# List content maintains list format regardless of cache
		user_list_cached = UserMessage(content=[ContentPartTextParam(text='List message', type='text')], cache=True)
		messages, _ = AnthropicMessageSerializer.serialize_messages([user_list_cached])
		assert isinstance(messages[0]['content'], list)

		user_list_no_cache = UserMessage(content=[ContentPartTextParam(text='List message', type='text')], cache=False)
		messages, _ = AnthropicMessageSerializer.serialize_messages([user_list_no_cache])
		assert isinstance(messages[0]['content'], list)

	def test_assistant_cache_empty_content(self):
		"""Test AssistantMessage with empty content and cache."""
		# With cache
		assistant_empty_cached = AssistantMessage(content=None, cache=True)
		messages, _ = AnthropicMessageSerializer.serialize_messages([assistant_empty_cached])

		assert len(messages) == 1
		assert isinstance(messages[0]['content'], list)

		# Without cache
		assistant_empty_no_cache = AssistantMessage(content=None, cache=False)
		messages, _ = AnthropicMessageSerializer.serialize_messages([assistant_empty_no_cache])

		assert len(messages) == 1
		assert isinstance(messages[0]['content'], str)

	def test_mixed_cache_scenarios(self):
		"""Test various combinations of cached and non-cached messages."""
		messages_list: list[BaseMessage] = [
			SystemMessage(content='System with cache', cache=True),
			UserMessage(content='User with cache', cache=True),
			AssistantMessage(content='Assistant without cache', cache=False),
			UserMessage(content='User without cache', cache=False),
			AssistantMessage(content='Assistant with cache', cache=True),
		]

		serialized_messages, system_message = AnthropicMessageSerializer.serialize_messages(messages_list)

		# Check system message is cached (becomes list)
		assert isinstance(system_message, list)

		# Check serialized messages
		assert len(serialized_messages) == 4

		# User with cache should be string (cache was cleaned to False by _clean_cache_messages)
		# Only the last message with cache=True remains cached
		assert isinstance(serialized_messages[0]['content'], str)

		# Assistant without cache should be string
		assert isinstance(serialized_messages[1]['content'], str)

		# User without cache should be string
		assert isinstance(serialized_messages[2]['content'], str)

		# Assistant with cache should be list (this is the last cached message)
		assert isinstance(serialized_messages[3]['content'], list)

	def test_system_message_cache_behavior(self):
		"""Test SystemMessage specific cache behavior."""
		# With cache
		system_cached = SystemMessage(content='System message with cache', cache=True)
		result = AnthropicMessageSerializer.serialize(system_cached)
		assert isinstance(result, SystemMessage)

		# Test serialization to string format
		serialized_content = AnthropicMessageSerializer._serialize_content_to_str(result.content, use_cache=True)
		assert isinstance(serialized_content, list)

		# Without cache
		system_no_cache = SystemMessage(content='System message without cache', cache=False)
		result = AnthropicMessageSerializer.serialize(system_no_cache)
		assert isinstance(result, SystemMessage)

		serialized_content = AnthropicMessageSerializer._serialize_content_to_str(result.content, use_cache=False)
		assert isinstance(serialized_content, str)

	def test_agent_messages_integration(self):
		"""Test integration with actual agent messages."""
		agent = Agent(task='Hello, world!', llm=ChatAnthropic(''))

		messages = agent.message_manager.get_messages()
		anthropic_messages, system_message = AnthropicMessageSerializer.serialize_messages(messages)

		# System message should be properly handled
		assert system_message is not None

	def test_cache_cleaning_last_message_only(self):
		"""Test that only the last cache=True message remains cached."""
		# Create multiple messages with cache=True
		messages_list: list[BaseMessage] = [
			UserMessage(content='First user message', cache=True),
			AssistantMessage(content='First assistant message', cache=True),
			UserMessage(content='Second user message', cache=True),
			AssistantMessage(content='Second assistant message', cache=False),
			UserMessage(content='Third user message', cache=True),  # This should be the only one cached
		]

		# Test the cleaning method directly (only accepts non-system messages)
		normal_messages = cast(list[NonSystemMessage], [msg for msg in messages_list if not isinstance(msg, SystemMessage)])
		cleaned_messages = AnthropicMessageSerializer._clean_cache_messages(normal_messages)

		# Verify only the last cache=True message remains cached
		assert not cleaned_messages[0].cache  # First user message should be uncached
		assert not cleaned_messages[1].cache  # First assistant message should be uncached
		assert not cleaned_messages[2].cache  # Second user message should be uncached
		assert not cleaned_messages[3].cache  # Second assistant message was already uncached
		assert cleaned_messages[4].cache  # Third user message should remain cached

		# Test through serialize_messages
		serialized_messages, system_message = AnthropicMessageSerializer.serialize_messages(messages_list)

		# Count how many messages have list content (indicating caching)
		cached_content_count = sum(1 for msg in serialized_messages if isinstance(msg['content'], list))

		# Only one message should have cached content
		assert cached_content_count == 1

		# The last message should be the cached one
		assert isinstance(serialized_messages[-1]['content'], list)

	def test_cache_cleaning_with_system_message(self):
		"""Test that system messages are not affected by cache cleaning logic."""
		messages_list: list[BaseMessage] = [
			SystemMessage(content='System message', cache=True),  # System messages are handled separately
			UserMessage(content='First user message', cache=True),
			AssistantMessage(content='Assistant message', cache=True),  # This should be the only normal message cached
		]

		# Test through serialize_messages to see the full integration
		serialized_messages, system_message = AnthropicMessageSerializer.serialize_messages(messages_list)

		# System message should be cached
		assert isinstance(system_message, list)

		# Only one normal message should have cached content (the last one)
		cached_content_count = sum(1 for msg in serialized_messages if isinstance(msg['content'], list))
		assert cached_content_count == 1

		# The last message should be the cached one
		assert isinstance(serialized_messages[-1]['content'], list)

	def test_cache_cleaning_no_cached_messages(self):
		"""Test that messages without cache=True are not affected."""
		normal_messages_list = [
			UserMessage(content='User message 1', cache=False),
			AssistantMessage(content='Assistant message 1', cache=False),
			UserMessage(content='User message 2', cache=False),
		]

		cleaned_messages = AnthropicMessageSerializer._clean_cache_messages(normal_messages_list)

		# All messages should remain uncached
		for msg in cleaned_messages:
			assert not msg.cache

	def test_max_4_cache_blocks(self):
		"""Test that the max number of cache blocks is 4."""
		agent = Agent(task='Hello, world!', llm=ChatAnthropic(''))
		messages = agent.message_manager.get_messages()
		anthropic_messages, system_message = AnthropicMessageSerializer.serialize_messages(messages)

		logger.info(anthropic_messages)
		logger.info(system_message)

	def test_cache_only_last_block_in_message(self):
		"""Test that only the LAST block in a message gets cache_control when cache=True."""
		# Test UserMessage with multiple text parts
		user_msg = UserMessage(
			content=[
				ContentPartTextParam(text='Part 1', type='text'),
				ContentPartTextParam(text='Part 2', type='text'),
				ContentPartTextParam(text='Part 3', type='text'),
			],
			cache=True,
		)
		serialized = AnthropicMessageSerializer.serialize(user_msg)
		assert isinstance(serialized['content'], list)
		content_blocks = serialized['content']

		# Count blocks with cache_control
		# Note: content_blocks are dicts at runtime despite type annotations
		cache_count = sum(1 for block in content_blocks if block.get('cache_control') is not None)  # type: ignore[attr-defined]
		assert cache_count == 1, f'Expected 1 cache_control block, got {cache_count}'

		# Verify it's the last block
		assert content_blocks[-1].get('cache_control') is not None  # type: ignore[attr-defined]
		assert content_blocks[0].get('cache_control') is None  # type: ignore[attr-defined]
		assert content_blocks[1].get('cache_control') is None  # type: ignore[attr-defined]

	def test_cache_only_last_tool_call(self):
		"""Test that only the LAST tool_use block gets cache_control."""
		tool_calls = [
			ToolCall(id='id1', function=Function(name='func1', arguments='{"arg": "1"}')),
			ToolCall(id='id2', function=Function(name='func2', arguments='{"arg": "2"}')),
			ToolCall(id='id3', function=Function(name='func3', arguments='{"arg": "3"}')),
		]

		assistant_msg = AssistantMessage(content=None, tool_calls=tool_calls, cache=True)
		serialized = AnthropicMessageSerializer.serialize(assistant_msg)
		assert isinstance(serialized['content'], list)
		content_blocks = serialized['content']

		# Count tool_use blocks with cache_control
		# Note: content_blocks are dicts at runtime despite type annotations
		cache_count = sum(1 for block in content_blocks if block.get('cache_control') is not None)  # type: ignore[attr-defined]
		assert cache_count == 1, f'Expected 1 cache_control block, got {cache_count}'

		# Verify it's the last tool_use block
		assert content_blocks[-1].get('cache_control') is not None  # type: ignore[attr-defined]
		assert content_blocks[0].get('cache_control') is None  # type: ignore[attr-defined]
		assert content_blocks[1].get('cache_control') is None  # type: ignore[attr-defined]

	def test_cache_assistant_with_content_and_tools(self):
		"""Test AssistantMessage with both content and tool calls - only last tool gets cache."""
		tool_call = ToolCall(id='test_id', function=Function(name='test_function', arguments='{"arg": "value"}'))

		assistant_msg = AssistantMessage(
			content=[
				ContentPartTextParam(text='Text part 1', type='text'),
				ContentPartTextParam(text='Text part 2', type='text'),
			],
			tool_calls=[tool_call],
			cache=True,
		)
		serialized = AnthropicMessageSerializer.serialize(assistant_msg)
		assert isinstance(serialized['content'], list)
		content_blocks = serialized['content']

		# Should have 2 text blocks + 1 tool_use block = 3 blocks total
		assert len(content_blocks) == 3

		# Only the last block (tool_use) should have cache_control
		# Note: content_blocks are dicts at runtime despite type annotations
		cache_count = sum(1 for block in content_blocks if block.get('cache_control') is not None)  # type: ignore[attr-defined]
		assert cache_count == 1, f'Expected 1 cache_control block, got {cache_count}'
		assert content_blocks[-1].get('cache_control') is not None  # type: ignore[attr-defined]  # Last tool_use block
		assert content_blocks[0].get('cache_control') is None  # type: ignore[attr-defined]  # First text block
		assert content_blocks[1].get('cache_control') is None  # type: ignore[attr-defined]  # Second text block


if __name__ == '__main__':
	test_instance = TestAnthropicCache()
	test_instance.test_cache_basic_functionality()
	test_instance.test_cache_with_tool_calls()
	test_instance.test_cache_with_images()
	test_instance.test_cache_with_base64_images()
	test_instance.test_cache_content_types()
	test_instance.test_assistant_cache_empty_content()
	test_instance.test_mixed_cache_scenarios()
	test_instance.test_system_message_cache_behavior()
	test_instance.test_agent_messages_integration()
	test_instance.test_cache_cleaning_last_message_only()
	test_instance.test_cache_cleaning_with_system_message()
	test_instance.test_cache_cleaning_no_cached_messages()
	test_instance.test_max_4_cache_blocks()
	test_instance.test_cache_only_last_block_in_message()
	test_instance.test_cache_only_last_tool_call()
	test_instance.test_cache_assistant_with_content_and_tools()
	print('All cache tests passed!')