diff --git a/letta/schemas/message.py b/letta/schemas/message.py index 0af6d2f0..9027c0ea 100644 --- a/letta/schemas/message.py +++ b/letta/schemas/message.py @@ -2139,6 +2139,69 @@ 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 c0124a25..2d6fc80b 100644 --- a/tests/test_message_serialization.py +++ b/tests/test_message_serialization.py @@ -1,7 +1,10 @@ +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 +from letta.schemas.message import Message, ToolReturn def _user_message_with_image_first(text: str) -> Message: @@ -30,3 +33,288 @@ 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" + )