From 4990d4b673c5fe3c94e07b905b8b0fea73cd0c8d Mon Sep 17 00:00:00 2001 From: Wasim Date: Wed, 2 Sep 2026 10:04:02 +0530 Subject: [PATCH] fix: remove orphaned FunctionExecutionResultMessage after truncation Previously, TokenLimitedChatCompletionContext only checked index 0 for an orphaned FunctionExecutionResultMessage after truncation. When the middle-index pop removed an AssistantMessage while leaving its paired FunctionExecutionResultMessage elsewhere in the list, it was never removed. Now _remove_orphaned_function_results checks the entire message list, removing any FunctionExecutionResultMessage that doesn't have a matching call_id in an AssistantMessage. Fixes #7955 --- .../_token_limited_chat_completion_context.py | 37 ++++++++++++-- .../autogen-core/tests/test_model_context.py | 51 +++++++++++++++++++ 2 files changed, 83 insertions(+), 5 deletions(-) diff --git a/python/packages/autogen-core/src/autogen_core/model_context/_token_limited_chat_completion_context.py b/python/packages/autogen-core/src/autogen_core/model_context/_token_limited_chat_completion_context.py index b8a0258a4d62..d0442b1be589 100644 --- a/python/packages/autogen-core/src/autogen_core/model_context/_token_limited_chat_completion_context.py +++ b/python/packages/autogen-core/src/autogen_core/model_context/_token_limited_chat_completion_context.py @@ -4,7 +4,8 @@ from typing_extensions import Self from .._component_config import Component, ComponentModel -from ..models import ChatCompletionClient, FunctionExecutionResultMessage, LLMMessage +from .. import FunctionCall +from ..models import AssistantMessage, ChatCompletionClient, FunctionExecutionResultMessage, LLMMessage from ..tools import ToolSchema from ._chat_completion_context import ChatCompletionContext @@ -70,12 +71,38 @@ async def get_messages(self) -> List[LLMMessage]: middle_index = len(messages) // 2 messages.pop(middle_index) token_count = self._model_client.count_tokens(messages, tools=self._tool_schema) - if messages and isinstance(messages[0], FunctionExecutionResultMessage): - # Handle the first message is a function call result message. - # Remove the first message from the list. - messages = messages[1:] + + # Remove orphaned FunctionExecutionResultMessages that have no matching + # AssistantMessage with a corresponding FunctionCall anywhere in the list. + # This can happen when truncation removes an AssistantMessage from the + # middle of the list while leaving its paired FunctionExecutionResultMessage. + messages = self._remove_orphaned_function_results(messages) + return messages + def _remove_orphaned_function_results(self, messages: List[LLMMessage]) -> List[LLMMessage]: + """Remove FunctionExecutionResultMessage instances that have no matching + FunctionCall in any AssistantMessage within the provided message list.""" + # Collect all call_ids from AssistantMessages + call_ids: set[str] = set() + for msg in messages: + if isinstance(msg, AssistantMessage) and isinstance(msg.content, list): + for item in msg.content: + if isinstance(item, FunctionCall) and item.id: + call_ids.add(item.id) + + # Filter out orphaned FunctionExecutionResultMessages + result: List[LLMMessage] = [] + for msg in messages: + if isinstance(msg, FunctionExecutionResultMessage): + # Keep only if at least one result has a matching call_id + has_match = any(result.call_id in call_ids for result in msg.content if result.call_id) + if has_match: + result.append(msg) + else: + result.append(msg) + return result + def _to_config(self) -> TokenLimitedChatCompletionContextConfig: return TokenLimitedChatCompletionContextConfig( model_client=self._model_client.dump_component(), diff --git a/python/packages/autogen-core/tests/test_model_context.py b/python/packages/autogen-core/tests/test_model_context.py index bcd2a9c87d9e..772a5db528d9 100644 --- a/python/packages/autogen-core/tests/test_model_context.py +++ b/python/packages/autogen-core/tests/test_model_context.py @@ -208,3 +208,54 @@ async def test_token_limited_model_context_openai_with_function_result( assert type(retrieved[0]) == UserMessage # Function result should be removed assert type(retrieved[1]) == AssistantMessage assert type(retrieved[2]) == UserMessage + + +@pytest.mark.asyncio +async def test_token_limited_context_removes_orphaned_function_result() -> None: + """Test that orphaned FunctionExecutionResultMessages are removed when + the paired AssistantMessage is truncated from the middle of the list. + See https://github.com/microsoft/autogen/issues/7955""" + + class FakeCounter: + def count_tokens(self, messages, tools=None): + return len(messages) # one "token" per message for deterministic test + + def remaining_tokens(self, messages, tools=None): + return 100 - len(messages) + + model_context = TokenLimitedChatCompletionContext(model_client=FakeCounter(), token_limit=6) + + from autogen_core import FunctionCall + from autogen_core.models import FunctionExecutionResult + + messages: List[LLMMessage] = [ + UserMessage(content="m0", source="user"), + UserMessage(content="m1", source="user"), + UserMessage(content="m2", source="user"), + AssistantMessage(content=[FunctionCall(id="call_1", arguments="{}", name="tool")], source="assistant"), + FunctionExecutionResultMessage( + content=[FunctionExecutionResult(content="ok", name="tool", call_id="call_1")] + ), + UserMessage(content="m5", source="user"), + UserMessage(content="m6", source="user"), + ] + for msg in messages: + await model_context.add_message(msg) + + retrieved = await model_context.get_messages() + + # The AssistantMessage at index 3 should be popped (middle of 7), + # leaving the FunctionExecutionResultMessage orphaned + # It should be removed, not left in the list + for msg in retrieved: + if isinstance(msg, FunctionExecutionResultMessage): + # If a FunctionExecutionResultMessage exists, it must have a matching call_id + # in an AssistantMessage somewhere in the list + call_ids = set() + for m in retrieved: + if isinstance(m, AssistantMessage) and isinstance(m.content, list): + for item in m.content: + if isinstance(item, FunctionCall): + call_ids.add(item.id) + assert any(r.call_id in call_ids for r in msg.content), \ + "Orphaned FunctionExecutionResultMessage found without matching AssistantMessage"