diff --git a/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/MessageMergerTests.cs b/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/MessageMergerTests.cs index 704e25b14a..4181dad409 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/MessageMergerTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/MessageMergerTests.cs @@ -17,7 +17,7 @@ public class MessageMergerTests [Fact] public void Test_MessageMerger_AssemblesMessage() { - DateTimeOffset creationTime = DateTimeOffset.UtcNow; + DateTimeOffset creationTime = DateTimeOffset.UtcNow.Subtract(TimeSpan.FromSeconds(1)); string responseId = Guid.NewGuid().ToString("N"); string messageId = Guid.NewGuid().ToString("N"); diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_message_adapters.py b/python/packages/ag-ui/agent_framework_ag_ui/_message_adapters.py index 4a846ea41d..2e5294a6b6 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_message_adapters.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_message_adapters.py @@ -242,8 +242,16 @@ def _deduplicate_messages(messages: list[Message]) -> list[Message]: unique_messages.append(msg) else: - content_str = str([str(c) for c in msg.contents]) if msg.contents else "" - key = (role_value, hash(content_str)) + # Use message_id for deduplication when available — two messages with the + # same id are definitively the same message (e.g. upstream replays), while + # different messages that happen to share identical content (e.g. repeated + # "yes" confirmations) will have distinct ids and be preserved. + # Fall back to content-hash when message_id is absent or empty. + if msg.message_id: + key = ("id", msg.message_id) + else: + content_str = str([str(c) for c in msg.contents]) if msg.contents else "" + key = ("content", role_value, hash(content_str)) if key in seen_keys: logger.info(f"Skipping duplicate message at index {idx}: role={role_value}") diff --git a/python/packages/ag-ui/tests/ag_ui/test_message_adapters.py b/python/packages/ag-ui/tests/ag_ui/test_message_adapters.py index 5227d376bb..cc4f1230df 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_message_adapters.py +++ b/python/packages/ag-ui/tests/ag_ui/test_message_adapters.py @@ -1015,15 +1015,111 @@ def test_deduplicate_assistant_tool_calls(): assert len(result) == 1 -def test_deduplicate_general_messages(): - """Duplicate general user messages are deduplicated.""" +def test_deduplicate_by_message_id(): + """Messages with the same message_id are deduplicated.""" + from agent_framework_ag_ui._message_adapters import _deduplicate_messages + + msg1 = Message(role="user", contents=[Content.from_text(text="Hello")]) + msg1.message_id = "msg-1" + msg2 = Message(role="user", contents=[Content.from_text(text="Hello")]) + msg2.message_id = "msg-1" + + result = _deduplicate_messages([msg1, msg2]) + assert len(result) == 1 + assert result == [msg1] + + +def test_deduplicate_preserves_repeated_confirmations_with_distinct_ids(): + """Identical content with different message_ids is preserved.""" + from agent_framework_ag_ui._message_adapters import _deduplicate_messages + + assistant = Message(role="assistant", contents=[Content.from_text(text="Are you sure?")]) + assistant.message_id = "msg-1" + confirm1 = Message(role="user", contents=[Content.from_text(text="yes")]) + confirm1.message_id = "msg-2" + confirm2 = Message(role="user", contents=[Content.from_text(text="yes")]) + confirm2.message_id = "msg-3" + + result = _deduplicate_messages([confirm1, assistant, confirm2]) + assert result == [confirm1, assistant, confirm2] + + +def test_deduplicate_preserves_repeated_system_messages_with_distinct_ids(): + """Non-consecutive identical system messages with different ids are preserved.""" + from agent_framework_ag_ui._message_adapters import _deduplicate_messages + + sys1 = Message(role="system", contents=[Content.from_text(text="You are a helpful assistant.")]) + sys1.message_id = "msg-1" + user_msg = Message(role="user", contents=[Content.from_text(text="Hi")]) + user_msg.message_id = "msg-2" + sys2 = Message(role="system", contents=[Content.from_text(text="You are a helpful assistant.")]) + sys2.message_id = "msg-3" + + result = _deduplicate_messages([sys1, user_msg, sys2]) + assert result == [sys1, user_msg, sys2] + + +def test_deduplicate_skips_replayed_system_messages_with_same_id(): + """System messages replayed with the same message_id are deduplicated.""" + from agent_framework_ag_ui._message_adapters import _deduplicate_messages + + msgs = [] + for _ in range(3): + m = Message(role="system", contents=[Content.from_text(text="You are a helpful assistant.")]) + m.message_id = "msg-1" + msgs.append(m) + + result = _deduplicate_messages(msgs) + assert len(result) == 1 + + +def test_deduplicate_without_message_id_uses_content_hash(): + """Messages without message_id are deduplicated by content hash.""" from agent_framework_ag_ui._message_adapters import _deduplicate_messages msg1 = Message(role="user", contents=[Content.from_text(text="Hello")]) msg2 = Message(role="user", contents=[Content.from_text(text="Hello")]) result = _deduplicate_messages([msg1, msg2]) - assert len(result) == 1 + assert result == [msg1] + + +def test_deduplicate_without_message_id_preserves_different_content(): + """Messages without message_id but different content are preserved.""" + from agent_framework_ag_ui._message_adapters import _deduplicate_messages + + msg1 = Message(role="user", contents=[Content.from_text(text="Hello")]) + msg2 = Message(role="user", contents=[Content.from_text(text="World")]) + + result = _deduplicate_messages([msg1, msg2]) + assert result == [msg1, msg2] + + +def test_deduplicate_handles_none_contents(): + """Messages with contents=None pass through without errors; duplicates are deduped.""" + from agent_framework_ag_ui._message_adapters import _deduplicate_messages + + msg1 = Message(role="user", contents=None) + msg2 = Message(role="assistant", contents=[Content.from_text(text="Hello")]) + msg3 = Message(role="user", contents=None) + + result = _deduplicate_messages([msg1, msg2, msg3]) + assert result == [msg1, msg2] + + +def test_deduplicate_mixed_id_and_no_id(): + """Messages with and without message_id coexist correctly.""" + from agent_framework_ag_ui._message_adapters import _deduplicate_messages + + msg1 = Message(role="user", contents=[Content.from_text(text="Hello")]) + msg1.message_id = "msg-1" + msg2 = Message(role="user", contents=[Content.from_text(text="Hello")]) # no id + msg3 = Message(role="user", contents=[Content.from_text(text="Hello")]) + msg3.message_id = "msg-1" # duplicate of msg1 + + result = _deduplicate_messages([msg1, msg2, msg3]) + assert len(result) == 2 + assert result == [msg1, msg2] def test_deduplicate_replaces_empty_tool_result(): @@ -1038,7 +1134,30 @@ def test_deduplicate_replaces_empty_tool_result(): assert result[0].contents[0].result == "actual result" -# ── Multimodal & content conversion edge cases ── +def test_deduplicate_empty_string_message_id_falls_back_to_content_hash(): + """Empty-string message_id is treated as missing; content-hash dedup is used.""" + from agent_framework_ag_ui._message_adapters import _deduplicate_messages + + msg1 = Message(role="user", contents=[Content.from_text(text="Hello")]) + msg1.message_id = "" + msg2 = Message(role="user", contents=[Content.from_text(text="World")]) + msg2.message_id = "" + + result = _deduplicate_messages([msg1, msg2]) + assert result == [msg1, msg2], "Different content with empty IDs should both be preserved" + + +def test_deduplicate_empty_string_message_id_deduplicates_same_content(): + """Empty-string message_id with identical content should be deduplicated.""" + from agent_framework_ag_ui._message_adapters import _deduplicate_messages + + msg1 = Message(role="user", contents=[Content.from_text(text="Hello")]) + msg1.message_id = "" + msg2 = Message(role="user", contents=[Content.from_text(text="Hello")]) + msg2.message_id = "" + + result = _deduplicate_messages([msg1, msg2]) + assert result == [msg1], "Same content with empty IDs should be deduplicated" def test_convert_agui_content_unknown_source_type_fallback(): diff --git a/python/packages/core/agent_framework/_mcp.py b/python/packages/core/agent_framework/_mcp.py index 81227c7e73..5901e34dd9 100644 --- a/python/packages/core/agent_framework/_mcp.py +++ b/python/packages/core/agent_framework/_mcp.py @@ -480,6 +480,10 @@ class MCPTool: self.load_prompts_flag = load_prompts self.parse_prompt_results = parse_prompt_results self._exit_stack = AsyncExitStack() + self._lifecycle_lock = asyncio.Lock() + self._lifecycle_request_lock = asyncio.Lock() + self._lifecycle_queue: asyncio.Queue[tuple[str, bool, asyncio.Future[None]]] | None = None + self._lifecycle_owner_task: asyncio.Task[None] | None = None self.session = session self.request_timeout = request_timeout self.client = client @@ -510,39 +514,113 @@ class MCPTool: filtered_functions.append(func) return filtered_functions + async def _ensure_lifecycle_owner(self) -> None: + async with self._lifecycle_lock: + if self._lifecycle_owner_task is not None and not self._lifecycle_owner_task.done(): + return + + self._lifecycle_queue = asyncio.Queue() + self._lifecycle_owner_task = asyncio.create_task( + self._run_lifecycle_owner(), + name=f"mcp-lifecycle:{self.name}", + ) + + async def _run_lifecycle_owner(self) -> None: + queue = self._lifecycle_queue + if queue is None: + return + + stop_error: BaseException | None = None + try: + while True: + action, reset, future = await queue.get() + + try: + if action == "connect": + await self._connect_on_owner(reset=reset) + elif action == "close": + await self._close_on_owner() + else: + raise RuntimeError(f"Unknown MCP lifecycle action: {action}") + except asyncio.CancelledError as ex: + stop_error = ex + if not future.done(): + future.set_exception(ex) + raise + except Exception as ex: + if not future.done(): + future.set_exception(ex) + else: + if not future.done(): + future.set_result(None) + + if action == "close": + return + except asyncio.CancelledError as ex: + stop_error = ex + raise + finally: + while True: + try: + _, _, future = queue.get_nowait() + except asyncio.QueueEmpty: + break + if not future.done(): + future.set_exception(stop_error or RuntimeError("MCP lifecycle owner stopped unexpectedly.")) + + self._lifecycle_queue = None + self._lifecycle_owner_task = None + + def _is_lifecycle_owner_task(self) -> bool: + owner_task = self._lifecycle_owner_task + return owner_task is not None and asyncio.current_task() is owner_task + + async def _run_on_lifecycle_owner(self, action: str, *, reset: bool = False) -> None: + await self._ensure_lifecycle_owner() + + if self._is_lifecycle_owner_task(): + if action == "connect": + await self._connect_on_owner(reset=reset) + elif action == "close": + await self._close_on_owner() + else: + raise RuntimeError(f"Unknown MCP lifecycle action: {action}") + return + + queue = self._lifecycle_queue + if queue is None: + raise RuntimeError("MCP lifecycle owner is not available.") + + future = asyncio.get_running_loop().create_future() + await queue.put((action, reset, future)) + await future + async def _safe_close_exit_stack(self) -> None: - """Safely close the exit stack, handling cross-task boundary errors. - - anyio's cancel scopes are bound to the task they were created in. - If aclose() is called from a different task (e.g., during streaming reconnection), - anyio will raise a RuntimeError or CancelledError. In this case, we log a warning - and allow garbage collection to clean up the resources. - - Known error variants: - - "Attempted to exit cancel scope in a different task than it was entered in" - - "Attempted to exit a cancel scope that isn't the current task's current cancel scope" - - CancelledError from anyio cancel scope cleanup - """ + """Safely close the exit stack, handling unexpected cleanup failures.""" try: await self._exit_stack.aclose() except RuntimeError as e: error_msg = str(e).lower() - # Check for anyio cancel scope errors (multiple variants exist) if "cancel scope" in error_msg: logger.warning( "Could not cleanly close MCP exit stack due to cancel scope error. " - "Old resources will be garbage collected. Error: %s", + "This indicates MCP lifecycle ownership was lost. Error: %s", e, ) else: raise except asyncio.CancelledError: - # CancelledError can occur during cleanup when cancel scopes are involved - logger.warning( - "Could not cleanly close MCP exit stack due to cancellation. Old resources will be garbage collected." - ) + logger.warning("Could not cleanly close MCP exit stack because the lifecycle owner task was cancelled.") async def connect(self, *, reset: bool = False) -> None: + if self._is_lifecycle_owner_task(): + await self._connect_on_owner(reset=reset) + return + + async with self._lifecycle_request_lock: + await self._run_on_lifecycle_owner("connect", reset=reset) + + async def _connect_on_owner(self, *, reset: bool = False) -> None: """Connect to the MCP server. Establishes a connection to the MCP server, initializes the session, @@ -844,14 +922,23 @@ class MCPTool: break params = types.PaginatedRequestParams(cursor=tool_list.nextCursor) + async def _close_on_owner(self) -> None: + await self._safe_close_exit_stack() + self._exit_stack = AsyncExitStack() + self.session = None + self.is_connected = False + async def close(self) -> None: """Disconnect from the MCP server. Closes the connection and cleans up resources. """ - await self._safe_close_exit_stack() - self.session = None - self.is_connected = False + if self._is_lifecycle_owner_task(): + await self._close_on_owner() + return + + async with self._lifecycle_request_lock: + await self._run_on_lifecycle_owner("close") @abstractmethod def get_mcp_client(self) -> _AsyncGeneratorContextManager[Any, None]: @@ -1043,7 +1130,7 @@ class MCPTool: except ToolException: raise except Exception as ex: - await self._safe_close_exit_stack() + await self.close() raise ToolExecutionException("Failed to enter context manager.", inner_exception=ex) from ex async def __aexit__( diff --git a/python/packages/core/tests/core/test_mcp.py b/python/packages/core/tests/core/test_mcp.py index 70aff972fe..b29ec1a794 100644 --- a/python/packages/core/tests/core/test_mcp.py +++ b/python/packages/core/tests/core/test_mcp.py @@ -2525,67 +2525,169 @@ async def test_mcp_tool_get_prompt_reconnection_on_closed_resource_error(): assert "failed to reconnect" in str(exc_info.value).lower() -async def test_mcp_tool_reconnection_handles_cross_task_cancel_scope_error(): - """Test that reconnection gracefully handles anyio cancel scope errors. +async def test_mcp_tool_close_cleans_up_in_original_task(caplog): + """Closing an MCP tool from another task should still unwind contexts in the owner task.""" + import asyncio - This tests the fix for the bug where calling connect(reset=True) from a - different task than where the connection was originally established would - cause: RuntimeError: Attempted to exit cancel scope in a different task - than it was entered in + class TaskBoundTransportContext: + def __init__(self) -> None: + self.enter_task = None + self.exit_task = None + self.closed_cleanly = False - This happens when using multiple MCP tools with AG-UI streaming - the first - tool call succeeds, but when the connection closes, the second tool call - triggers a reconnection from within the streaming loop (a different task). - """ - from contextlib import AsyncExitStack + async def __aenter__(self): + self.enter_task = asyncio.current_task() + return (Mock(), Mock()) - from agent_framework._mcp import MCPStdioTool + async def __aexit__(self, exc_type, exc, tb): + self.exit_task = asyncio.current_task() + if self.exit_task is not self.enter_task: + raise RuntimeError("Attempted to exit cancel scope in a different task than it was entered in") + self.closed_cleanly = True + return - # Use load_tools=False and load_prompts=False to avoid triggering them during connect() - tool = MCPStdioTool( + tool = MCPStreamableHTTPTool( name="test_server", - command="test_command", - args=["arg1"], + url="https://example.com/mcp", load_tools=False, load_prompts=False, ) - # Mock the exit stack to raise the cross-task cancel scope error - mock_exit_stack = AsyncMock(spec=AsyncExitStack) - mock_exit_stack.aclose = AsyncMock( - side_effect=RuntimeError("Attempted to exit cancel scope in a different task than it was entered in") - ) - tool._exit_stack = mock_exit_stack - tool.session = Mock() - tool.is_connected = True + transport_context = TaskBoundTransportContext() + mock_session = Mock() + mock_session._request_id = 1 + mock_session.initialize = AsyncMock() - # Mock get_mcp_client to return a mock transport - mock_transport = (Mock(), Mock()) - mock_context = AsyncMock() - mock_context.__aenter__ = AsyncMock(return_value=mock_transport) - mock_context.__aexit__ = AsyncMock() + mock_session_context = AsyncMock() + mock_session_context.__aenter__ = AsyncMock(return_value=mock_session) + mock_session_context.__aexit__ = AsyncMock(return_value=None) with ( - patch.object(tool, "get_mcp_client", return_value=mock_context), - patch("agent_framework._mcp.ClientSession") as mock_session_class, + patch.object(tool, "get_mcp_client", return_value=transport_context), + patch("agent_framework._mcp.ClientSession", return_value=mock_session_context), ): - mock_session = Mock() - mock_session._request_id = 1 - mock_session.initialize = AsyncMock() - mock_session.set_logging_level = AsyncMock() - mock_session_context = AsyncMock() - mock_session_context.__aenter__ = AsyncMock(return_value=mock_session) - mock_session_context.__aexit__ = AsyncMock() - mock_session_class.return_value = mock_session_context + await asyncio.create_task(tool.connect()) - # This should NOT raise even though aclose() raised the cancel scope error - # The _safe_close_exit_stack method should catch and log the error - await tool.connect(reset=True) + caplog.clear() + with caplog.at_level(logging.WARNING, logger=logger.name): + await tool.close() - # Verify a new exit stack was created (the old mock was replaced) - assert tool._exit_stack is not mock_exit_stack - assert tool.session is not None + assert transport_context.closed_cleanly is True + assert transport_context.exit_task is transport_context.enter_task + assert not any("cancel scope" in record.getMessage().lower() for record in caplog.records) + + +async def test_mcp_tool_connect_reset_cleans_up_in_original_task(caplog): + """Resetting an MCP tool from another task should unwind and reconnect on the owner task.""" + import asyncio + + class TaskBoundTransportContext: + def __init__(self) -> None: + self.enter_task = None + self.exit_task = None + self.closed_cleanly = False + + async def __aenter__(self): + self.enter_task = asyncio.current_task() + return (Mock(), Mock()) + + async def __aexit__(self, exc_type, exc, tb): + self.exit_task = asyncio.current_task() + if self.exit_task is not self.enter_task: + raise RuntimeError("Attempted to exit cancel scope in a different task than it was entered in") + self.closed_cleanly = True + return + + tool = MCPStreamableHTTPTool( + name="test_server", + url="https://example.com/mcp", + load_tools=False, + load_prompts=False, + ) + + transport_contexts = [TaskBoundTransportContext(), TaskBoundTransportContext()] + sessions = [] + session_contexts = [] + for _ in range(2): + session = Mock() + session._request_id = 1 + session.initialize = AsyncMock() + session.set_logging_level = AsyncMock() + sessions.append(session) + + session_context = AsyncMock() + session_context.__aenter__ = AsyncMock(return_value=session) + session_context.__aexit__ = AsyncMock(return_value=None) + session_contexts.append(session_context) + + with ( + patch.object(tool, "get_mcp_client", side_effect=transport_contexts), + patch("agent_framework._mcp.ClientSession", side_effect=session_contexts), + ): + await tool.connect() + + caplog.clear() + with caplog.at_level(logging.WARNING, logger=logger.name): + await asyncio.create_task(tool.connect(reset=True)) + + assert transport_contexts[0].closed_cleanly is True + assert transport_contexts[0].exit_task is transport_contexts[0].enter_task + assert transport_contexts[1].enter_task is transport_contexts[0].enter_task + assert tool.session is sessions[1] assert tool.is_connected is True + assert not any("cancel scope" in record.getMessage().lower() for record in caplog.records) + + await tool.close() + + +async def test_mcp_tool_connect_from_lifecycle_owner_bypasses_request_lock() -> None: + """connect(reset=True) should bypass the request queue when already on the owner task.""" + import asyncio + + tool = MCPStreamableHTTPTool( + name="test_server", + url="https://example.com/mcp", + load_tools=False, + load_prompts=False, + ) + + async def connect_from_owner_task() -> None: + tool._lifecycle_owner_task = asyncio.current_task() + try: + async with tool._lifecycle_request_lock: + await tool.connect(reset=True) + finally: + tool._lifecycle_owner_task = None + + with patch.object(tool, "_connect_on_owner", AsyncMock()) as mock_connect_on_owner: + await asyncio.wait_for(connect_from_owner_task(), timeout=0.1) + + mock_connect_on_owner.assert_awaited_once_with(reset=True) + + +async def test_mcp_tool_close_from_lifecycle_owner_bypasses_request_lock() -> None: + """close() should bypass the request queue when already on the owner task.""" + import asyncio + + tool = MCPStreamableHTTPTool( + name="test_server", + url="https://example.com/mcp", + load_tools=False, + load_prompts=False, + ) + + async def close_from_owner_task() -> None: + tool._lifecycle_owner_task = asyncio.current_task() + try: + async with tool._lifecycle_request_lock: + await tool.close() + finally: + tool._lifecycle_owner_task = None + + with patch.object(tool, "_close_on_owner", AsyncMock()) as mock_close_on_owner: + await asyncio.wait_for(close_from_owner_task(), timeout=0.1) + + mock_close_on_owner.assert_awaited_once_with() async def test_mcp_tool_safe_close_reraises_other_runtime_errors(): diff --git a/python/samples/04-hosting/azure_functions/09_workflow_shared_state/requirements.txt b/python/samples/04-hosting/azure_functions/09_workflow_shared_state/requirements.txt index 5739f93aa3..39ad8a124f 100644 --- a/python/samples/04-hosting/azure_functions/09_workflow_shared_state/requirements.txt +++ b/python/samples/04-hosting/azure_functions/09_workflow_shared_state/requirements.txt @@ -1,3 +1,2 @@ agent-framework-azurefunctions azure-identity -agents-maf \ No newline at end of file diff --git a/python/samples/04-hosting/azure_functions/12_workflow_hitl/requirements.txt b/python/samples/04-hosting/azure_functions/12_workflow_hitl/requirements.txt index 85e158b8d4..39ad8a124f 100644 --- a/python/samples/04-hosting/azure_functions/12_workflow_hitl/requirements.txt +++ b/python/samples/04-hosting/azure_functions/12_workflow_hitl/requirements.txt @@ -1,3 +1,2 @@ agent-framework-azurefunctions azure-identity -agents-maf