diff --git a/letta/schemas/message.py b/letta/schemas/message.py index 9027c0ea..0af6d2f0 100644 --- a/letta/schemas/message.py +++ b/letta/schemas/message.py @@ -2139,69 +2139,6 @@ class Message(BaseMessage): for idx in reversed(messages_to_filter): # reverse to avoid index shift messages.remove(idx) - # Filter orphaned approval_request messages that have tool_calls but no corresponding tool_return. - # This can happen due to race conditions in checkpoint, client disconnect, or other edge cases. - # We scan for ALL such messages (not just the last one) to handle sporadic state corruption. - # - # Performance: O(n) pre-build lookup + O(n) scan = O(n) total - # This scales well to large context windows (10K-100K+ messages). - - # Pre-build map of tool_call_id -> minimum position where a tool_return appears. - # We need position-awareness because tool_returns must come AFTER the tool_use (approval_request). - # This converts O(n × a) inner loop to O(1) lookups per approval_request. - tool_return_positions: dict = {} # tool_call_id -> min position with tool_return - for idx, msg in enumerate(messages): - if msg.role == MessageRole.tool: - # Check tool_returns array - if msg.tool_returns: - for tr in msg.tool_returns: - if tr.tool_call_id: - # Keep the minimum (earliest) position for each tool_call_id - if tr.tool_call_id not in tool_return_positions: - tool_return_positions[tr.tool_call_id] = idx - # Check single tool_call_id (legacy) - elif msg.tool_call_id: - if msg.tool_call_id not in tool_return_positions: - tool_return_positions[msg.tool_call_id] = idx - # Check approval messages with approvals array containing ToolReturn entries. - # NOTE: Only ToolReturn entries count (they have tool results), not ApprovalReturn - # (which only has approve/deny status without tool_result content). - elif msg.role == MessageRole.approval and msg.approvals: - for approval in msg.approvals: - # Only count ToolReturn entries, not ApprovalReturn: - # - ApprovalReturn has 'approve' attr (bool) but no tool result content - # - ToolReturn has 'func_response' attr with actual tool results - is_tool_return = hasattr(approval, "func_response") and not hasattr(approval, "approve") - if is_tool_return: - tcid = getattr(approval, "tool_call_id", None) - if tcid and tcid not in tool_return_positions: - tool_return_positions[tcid] = idx - - # Now scan for orphaned approval_requests using O(1) lookups - i = 0 - while i < len(messages): - msg = messages[i] - # Check if this is an approval_request (has tool_calls) - if msg.is_approval_request(): - tool_call_ids = {tc.id for tc in msg.tool_calls if tc.id} - # Check if any tool_calls have returns AFTER this position - has_tool_returns = any(tcid in tool_return_positions and tool_return_positions[tcid] > i for tcid in tool_call_ids) - - if not has_tool_returns: - # Remove orphaned approval_request - logger.warning(f"Filtering orphaned approval_request message {msg.id} with tool_calls but no tool_return") - messages.pop(i) - # Also remove preceding assistant message with tool_calls if present (they share tool_use blocks) - if i > 0 and messages[i - 1].role == MessageRole.assistant and messages[i - 1].tool_calls: - # Check if the assistant's tool_calls overlap with the orphaned approval's tool_calls - assistant_tc_ids = {tc.id for tc in messages[i - 1].tool_calls if tc.id} - if tool_call_ids & assistant_tc_ids: - logger.warning(f"Filtering related assistant message {messages[i - 1].id} with orphaned tool_calls") - messages.pop(i - 1) - i -= 1 - continue # Don't increment i, check new message at this index - i += 1 - # Filter last message if it is a lone approval request without a response - this only occurs for token counting if messages[-1].role == "approval" and messages[-1].tool_calls is not None and len(messages[-1].tool_calls) > 0: messages.remove(messages[-1]) diff --git a/tests/test_message_serialization.py b/tests/test_message_serialization.py index 2d6fc80b..c0124a25 100644 --- a/tests/test_message_serialization.py +++ b/tests/test_message_serialization.py @@ -1,10 +1,7 @@ -from openai.types.chat.chat_completion_message_tool_call import ChatCompletionMessageToolCall as OpenAIToolCall, Function as OpenAIFunction - from letta.llm_api.openai_client import fill_image_content_in_responses_input from letta.schemas.enums import MessageRole -from letta.schemas.letta_message import ApprovalReturn from letta.schemas.letta_message_content import Base64Image, ImageContent, TextContent -from letta.schemas.message import Message, ToolReturn +from letta.schemas.message import Message def _user_message_with_image_first(text: str) -> Message: @@ -33,288 +30,3 @@ def test_to_openai_responses_dicts_handles_image_only_content(): serialized = Message.to_openai_responses_dicts_from_list([message]) parts = serialized[0]["content"] assert parts[0]["type"] == "input_image" - - -def _create_tool_call(tool_call_id: str, name: str, arguments: str = "{}") -> OpenAIToolCall: - """Helper to create an OpenAI-style tool call.""" - return OpenAIToolCall( - id=tool_call_id, - type="function", - function=OpenAIFunction(name=name, arguments=arguments), - ) - - -def test_filter_orphaned_approval_request_at_end(): - """ - Test that a lone approval_request at the end of messages is filtered out. - This is the original behavior that should still work. - """ - system_msg = Message(role=MessageRole.system, content=[TextContent(text="You are a helpful assistant.")]) - user_msg = Message(role=MessageRole.user, content=[TextContent(text="Do something")]) - # Approval request with tool_calls but no response - approval_request = Message( - role=MessageRole.approval, - tool_calls=[_create_tool_call("tc_123", "some_tool")], - ) - - messages = [system_msg, user_msg, approval_request] - filtered = Message.filter_messages_for_llm_api(messages) - - # The orphaned approval_request should be removed - assert len(filtered) == 2 - assert filtered[0].role == MessageRole.system - assert filtered[1].role == MessageRole.user - - -def test_filter_orphaned_approval_request_in_middle(): - """ - Test that an orphaned approval_request message in the MIDDLE of the message list - is properly filtered before sending to LLM. - - This is the key scenario for the race condition bug: - - approval_request gets persisted - - User sends another message before responding to approval - - Due to race conditions, the orphaned approval_request ends up in context without tool_return - """ - system_msg = Message(role=MessageRole.system, content=[TextContent(text="You are a helpful assistant.")]) - user_msg1 = Message(role=MessageRole.user, content=[TextContent(text="Do something")]) - # Orphaned approval_request with tool_calls but no response - orphaned_approval = Message( - role=MessageRole.approval, - tool_calls=[_create_tool_call("tc_orphan", "dangerous_tool")], - ) - # User sends another message (race condition scenario) - user_msg2 = Message(role=MessageRole.user, content=[TextContent(text="Actually do something else")]) - - messages = [system_msg, user_msg1, orphaned_approval, user_msg2] - filtered = Message.filter_messages_for_llm_api(messages) - - # The orphaned approval_request should be removed, but user messages stay - assert len(filtered) == 3 - assert all(msg.role != MessageRole.approval for msg in filtered) - # Verify the messages are in correct order - assert filtered[0].role == MessageRole.system - assert filtered[1].role == MessageRole.user - assert filtered[2].role == MessageRole.user - - -def test_filter_preserves_approval_with_tool_return(): - """ - Test that approval_request messages WITH corresponding tool_return are NOT filtered. - This is the normal happy path that should not be affected by the fix. - """ - system_msg = Message(role=MessageRole.system, content=[TextContent(text="You are a helpful assistant.")]) - # Approval request - approval_request = Message( - role=MessageRole.approval, - tool_calls=[_create_tool_call("tc_valid", "approved_tool")], - ) - # Tool return with matching tool_call_id - tool_return_msg = Message( - role=MessageRole.tool, - tool_returns=[ToolReturn(tool_call_id="tc_valid", status="success", func_response="done")], - ) - - messages = [system_msg, approval_request, tool_return_msg] - filtered = Message.filter_messages_for_llm_api(messages) - - # All messages should be preserved - assert len(filtered) == 3 - assert filtered[1].role == MessageRole.approval - assert filtered[2].role == MessageRole.tool - - -def test_filter_orphaned_approval_removes_related_assistant_message(): - """ - Test that when an orphaned approval_request is removed, the preceding assistant - message with overlapping tool_calls is also removed. - - In parallel tool calling scenarios, tool_calls might be split between - assistant and approval messages - both need to be removed to avoid orphaned tool_use blocks. - """ - system_msg = Message(role=MessageRole.system, content=[TextContent(text="You are a helpful assistant.")]) - # Assistant message with tool_calls - assistant_msg = Message( - role=MessageRole.assistant, - tool_calls=[_create_tool_call("tc_shared", "some_tool")], - ) - # Orphaned approval_request with same tool_call_id - orphaned_approval = Message( - role=MessageRole.approval, - tool_calls=[_create_tool_call("tc_shared", "some_tool")], - ) - user_msg = Message(role=MessageRole.user, content=[TextContent(text="Do something else")]) - - messages = [system_msg, assistant_msg, orphaned_approval, user_msg] - filtered = Message.filter_messages_for_llm_api(messages) - - # Both assistant and approval messages should be removed - assert len(filtered) == 2 - assert filtered[0].role == MessageRole.system - assert filtered[1].role == MessageRole.user - - -def test_no_orphaned_tool_use_blocks_after_filter(): - """ - Comprehensive test: After filtering, verify there are NO tool_use blocks - without corresponding tool_result blocks. - - This is the ultimate test for the Anthropic API requirement: - "tool_use ids were found without tool_result blocks immediately after" - """ - system_msg = Message(role=MessageRole.system, content=[TextContent(text="You are a helpful assistant.")]) - user_msg = Message(role=MessageRole.user, content=[TextContent(text="Do something")]) - # Multiple orphaned approval_requests scattered in the list - orphaned1 = Message( - role=MessageRole.approval, - tool_calls=[_create_tool_call("tc_orphan1", "tool1")], - ) - orphaned2 = Message( - role=MessageRole.approval, - tool_calls=[_create_tool_call("tc_orphan2", "tool2"), _create_tool_call("tc_orphan3", "tool3")], - ) - user_msg2 = Message(role=MessageRole.user, content=[TextContent(text="Another message")]) - - messages = [system_msg, user_msg, orphaned1, orphaned2, user_msg2] - filtered = Message.filter_messages_for_llm_api(messages) - - # Collect all tool_call_ids from messages with tool_calls - tool_call_ids = set() - for msg in filtered: - if msg.tool_calls: - for tc in msg.tool_calls: - tool_call_ids.add(tc.id) - - # Collect all tool_return ids - tool_return_ids = set() - for msg in filtered: - if msg.role == MessageRole.tool: - if msg.tool_returns: - for tr in msg.tool_returns: - if tr.tool_call_id: - tool_return_ids.add(tr.tool_call_id) - elif msg.tool_call_id: - tool_return_ids.add(msg.tool_call_id) - - # Every tool_call_id must have a corresponding tool_return - orphaned_tool_calls = tool_call_ids - tool_return_ids - assert len(orphaned_tool_calls) == 0, f"Found orphaned tool_use blocks without tool_result: {orphaned_tool_calls}" - - -def test_filter_approval_with_only_approval_return_is_still_orphaned(): - """ - Test that an approval_request followed by approval_response with only ApprovalReturn - (approve/deny decision without tool_result) is STILL treated as orphaned. - - ApprovalReturn only contains approve=True/False, not the actual tool_result content. - The LLM API requires tool_use blocks to have matching tool_result blocks. - """ - system_msg = Message(role=MessageRole.system, content=[TextContent(text="You are a helpful assistant.")]) - # Approval request with tool_calls - approval_request = Message( - role=MessageRole.approval, - tool_calls=[_create_tool_call("tc_approved", "some_tool")], - ) - # Approval response with ApprovalReturn (approve=True but NO tool_result) - approval_response = Message( - role=MessageRole.approval, - approve=None, # New API pattern: approve=None with approvals list - approvals=[ApprovalReturn(tool_call_id="tc_approved", approve=True)], - ) - user_msg = Message(role=MessageRole.user, content=[TextContent(text="Continue")]) - - messages = [system_msg, approval_request, approval_response, user_msg] - filtered = Message.filter_messages_for_llm_api(messages) - - # The approval_request should be removed because ApprovalReturn doesn't have tool_result - assert all(not (msg.role == MessageRole.approval and msg.tool_calls) for msg in filtered), ( - "Orphaned approval_request with only ApprovalReturn should be filtered" - ) - - -def test_filter_preserves_approval_with_tool_return_in_approvals(): - """ - Test that approval_request followed by approval_response containing ToolReturn - (actual tool result) is NOT filtered - this is a valid completed flow. - """ - system_msg = Message(role=MessageRole.system, content=[TextContent(text="You are a helpful assistant.")]) - # Approval request with tool_calls - approval_request = Message( - role=MessageRole.approval, - tool_calls=[_create_tool_call("tc_with_result", "some_tool")], - ) - # Approval response with ToolReturn (has actual func_response content) - # Using the local ToolReturn from message.py (not LettaToolReturn from letta_message.py) - approval_response = Message( - role=MessageRole.approval, - approve=None, - approvals=[ToolReturn(tool_call_id="tc_with_result", status="success", func_response="result data")], - ) - - messages = [system_msg, approval_request, approval_response] - filtered = Message.filter_messages_for_llm_api(messages) - - # All messages should be preserved - this is a valid flow - assert len(filtered) == 3 - assert filtered[1].role == MessageRole.approval - assert filtered[1].tool_calls is not None # approval_request preserved - - -def test_filter_legacy_single_tool_call_id(): - """ - Test that the legacy single tool_call_id field on tool messages is correctly - recognized as having a tool_return (not just tool_returns array). - """ - system_msg = Message(role=MessageRole.system, content=[TextContent(text="You are a helpful assistant.")]) - # Approval request - approval_request = Message( - role=MessageRole.approval, - tool_calls=[_create_tool_call("tc_legacy", "some_tool")], - ) - # Tool message with legacy single tool_call_id (not tool_returns array) - tool_msg = Message( - role=MessageRole.tool, - tool_call_id="tc_legacy", - content=[TextContent(text="Tool result")], - ) - - messages = [system_msg, approval_request, tool_msg] - filtered = Message.filter_messages_for_llm_api(messages) - - # All messages should be preserved - assert len(filtered) == 3 - assert filtered[1].role == MessageRole.approval - assert filtered[1].tool_calls is not None - - -def test_filter_tool_return_before_approval_request_still_orphaned(): - """ - Test that an approval_request is still filtered as orphaned even if a tool_return - with matching ID exists BEFORE it in the message list. - - The LLM API requires tool_result to come AFTER tool_use. A tool_return before - the approval_request doesn't satisfy this requirement. - - This catches the ordering bug where position-unaware lookups would incorrectly - keep the approval_request because a matching tool_return "exists somewhere". - """ - system_msg = Message(role=MessageRole.system, content=[TextContent(text="You are a helpful assistant.")]) - # Tool return BEFORE the approval_request (wrong order) - tool_return_msg = Message( - role=MessageRole.tool, - tool_returns=[ToolReturn(tool_call_id="tc_wrong_order", status="success", func_response="done")], - ) - # Approval request with same tool_call_id - approval_request = Message( - role=MessageRole.approval, - tool_calls=[_create_tool_call("tc_wrong_order", "some_tool")], - ) - user_msg = Message(role=MessageRole.user, content=[TextContent(text="Continue")]) - - messages = [system_msg, tool_return_msg, approval_request, user_msg] - filtered = Message.filter_messages_for_llm_api(messages) - - # The approval_request should be removed because tool_return is BEFORE it, not after - assert all(not (msg.role == MessageRole.approval and msg.tool_calls) for msg in filtered), ( - "Approval_request with tool_return BEFORE it should still be filtered as orphaned" - )