Merge branch 'main' into fix/playwrightcontroller-sleep-targetclosed

.github/ISSUE_TEMPLATE/1-bug_report.yml

@@ -90,6 +90,7 @@ body:
       multiple: false
       options:
         - "Python dev (main branch)"
+        - "Python 0.5.7"
         - "Python 0.5.6"
         - "Python 0.5.5"
         - "Python 0.5.4"

.github/workflows/docs.yml

@@ -33,7 +33,7 @@ jobs:
           [
             # For main use the workflow target
             { ref: "${{github.ref}}", dest-dir: dev, uv-version: "0.5.13", sphinx-release-override: "dev" },
-            { ref: "python-v0.5.6", dest-dir: stable, uv-version: "0.5.13", sphinx-release-override: "stable" },
+            { ref: "python-v0.5.7", dest-dir: stable, uv-version: "0.5.13", sphinx-release-override: "stable" },
             { ref: "v0.4.0.post1", dest-dir: "0.4.0", uv-version: "0.5.13", sphinx-release-override: "" },
             { ref: "v0.4.1", dest-dir: "0.4.1", uv-version: "0.5.13", sphinx-release-override: "" },
             { ref: "v0.4.2", dest-dir: "0.4.2", uv-version: "0.5.13", sphinx-release-override: "" },
@@ -50,6 +50,7 @@ jobs:
             { ref: "python-v0.5.4", dest-dir: "0.5.4", uv-version: "0.5.13", sphinx-release-override: "" },
             { ref: "python-v0.5.5", dest-dir: "0.5.5", uv-version: "0.5.13", sphinx-release-override: "" },
             { ref: "python-v0.5.6", dest-dir: "0.5.6", uv-version: "0.5.13", sphinx-release-override: "" },
+            { ref: "python-v0.5.7", dest-dir: "0.5.7", uv-version: "0.5.13", sphinx-release-override: "" },
           ]
     steps:
       - name: Checkout

docs/switcher.json

@@ -5,11 +5,16 @@
         "url": "/autogen/dev/"
     },
     {
-        "name": "0.5.6 (stable)",
+        "name": "0.5.7 (stable)",
         "version": "stable",
         "url": "/autogen/stable/",
         "preferred": true
     },
+    {
+        "name": "0.5.6",
+        "version": "0.5.6",
+        "url": "/autogen/0.5.6/"
+    },
     {
         "name": "0.5.5",
         "version": "0.5.5",
@@ -90,4 +95,4 @@
         "version": "0.2",
         "url": "/autogen/0.2/"
     }
-]
+]

dotnet/Directory.Packages.props

@@ -5,7 +5,8 @@
     <MicrosoftSemanticKernelStableVersion>1.45.0</MicrosoftSemanticKernelStableVersion>
     <MicrosoftSemanticKernelPreviewVersion>$(MicrosoftSemanticKernelStableVersion)-preview</MicrosoftSemanticKernelPreviewVersion>
     <MicrosoftSemanticKernelAlphaVersion>$(MicrosoftSemanticKernelStableVersion)-alpha</MicrosoftSemanticKernelAlphaVersion>
-    <MicrosoftExtensionsAIVersion>9.3.0-preview.1.25161.3</MicrosoftExtensionsAIVersion>
+    <MicrosoftExtensionsAIVersion>9.5.0</MicrosoftExtensionsAIVersion>
+    <MicrosoftExtensionsAIPreviewVersion>9.5.0-preview.1.25265.7</MicrosoftExtensionsAIPreviewVersion>
     <MicrosoftExtensionConfiguration>9.0.0</MicrosoftExtensionConfiguration>
     <MicrosoftExtensionDependencyInjection>9.0.3</MicrosoftExtensionDependencyInjection>
     <MicrosoftExtensionLogging>9.0.0</MicrosoftExtensionLogging>
@@ -64,9 +65,9 @@
     <PackageVersion Include="Microsoft.DotNet.Interactive.PackageManagement" Version="$(MicrosoftDotNetInteractive)" />
     <PackageVersion Include="Microsoft.Extensions.AI" Version="$(MicrosoftExtensionsAIVersion)" />
     <PackageVersion Include="Microsoft.Extensions.AI.Abstractions" Version="$(MicrosoftExtensionsAIVersion)" />
-    <PackageVersion Include="Microsoft.Extensions.AI.AzureAIInference" Version="$(MicrosoftExtensionsAIVersion)" />
-    <PackageVersion Include="Microsoft.Extensions.AI.Ollama" Version="$(MicrosoftExtensionsAIVersion)" />
-    <PackageVersion Include="Microsoft.Extensions.AI.OpenAI" Version="$(MicrosoftExtensionsAIVersion)" />
+    <PackageVersion Include="Microsoft.Extensions.AI.AzureAIInference" Version="$(MicrosoftExtensionsAIPreviewVersion)" />
+    <PackageVersion Include="Microsoft.Extensions.AI.Ollama" Version="$(MicrosoftExtensionsAIPreviewVersion)" />
+    <PackageVersion Include="Microsoft.Extensions.AI.OpenAI" Version="$(MicrosoftExtensionsAIPreviewVersion)" />
     <PackageVersion Include="Microsoft.Extensions.Azure" Version="1.8.0" />
     <PackageVersion Include="Microsoft.Extensions.Caching.Abstractions" Version="8.0.0" />
     <PackageVersion Include="Microsoft.Extensions.Configuration.EnvironmentVariables" Version="$(MicrosoftExtensionConfiguration)" />
@@ -135,4 +136,4 @@
     <PackageVersion Include="xunit.runner.console" Version="2.9.2" />
     <PackageVersion Include="xunit.runner.visualstudio" Version="2.8.2" />
   </ItemGroup>
-</Project>
+</Project>

dotnet/README.md

@@ -40,7 +40,7 @@ var assistantAgent = new AssistantAgent(
 // set human input mode to ALWAYS so that user always provide input
 var userProxyAgent = new UserProxyAgent(
     name: "user",
-    humanInputMode: ConversableAgent.HumanInputMode.ALWAYS)
+    humanInputMode: HumanInputMode.ALWAYS)
     .RegisterPrintMessage();
 
 // start the conversation

dotnet/src/AutoGen.Core/Middleware/FunctionCallMiddleware.cs

@@ -189,12 +189,12 @@ public class FunctionCallMiddleware : IStreamingMiddleware
         }
     }
 
-    private Func<string, Task<string>> AIToolInvokeWrapper(Func<IEnumerable<KeyValuePair<string, object?>>?, CancellationToken, Task<object?>> lambda)
+    private Func<string, Task<string>> AIToolInvokeWrapper(Func<AIFunctionArguments?, CancellationToken, ValueTask<object?>> lambda)
     {
         return async (string args) =>
         {
             var arguments = JsonSerializer.Deserialize<Dictionary<string, object?>>(args);
-            var result = await lambda(arguments, CancellationToken.None);
+            var result = await lambda(new(arguments), CancellationToken.None);
 
             return result switch
             {

dotnet/src/Microsoft.AutoGen/Extensions/MEAI/ServiceCollectionChatCompletionExtensions.cs

@@ -89,7 +89,7 @@ public static class ServiceCollectionChatClientExtensions
             .AddChatClient(service =>
             {
                 var openAiClient = service.GetRequiredService<OpenAIClient>();
-                return openAiClient.AsChatClient(modelOrDeploymentName);
+                return openAiClient.GetChatClient(modelOrDeploymentName).AsIChatClient();
             });
 
         return services;
@@ -112,7 +112,7 @@ public static class ServiceCollectionChatClientExtensions
         var endpoint = $"{serviceName}:Endpoint" ?? throw new InvalidOperationException($"No endpoint was specified for the Azure Inference Chat Client");
         var endpointUri = string.IsNullOrEmpty(endpoint) ? null : new Uri(endpoint);
         var token = Environment.GetEnvironmentVariable("AZURE_OPENAI_API_KEY") ?? throw new InvalidOperationException("No model access token was found in the environment variable AZURE_OPENAI_API_KEY");
-        var chatClient = new ChatCompletionsClient(endpointUri, new AzureKeyCredential(token)).AsChatClient(modelOrDeploymentName);
+        var chatClient = new ChatCompletionsClient(endpointUri, new AzureKeyCredential(token)).AsIChatClient(modelOrDeploymentName);
         hostBuilder.Services.AddChatClient(chatClient);
 
         return hostBuilder.Services;

python/packages/autogen-agentchat/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "autogen-agentchat"
-version = "0.5.6"
+version = "0.5.7"
 license = {file = "LICENSE-CODE"}
 description = "AutoGen agents and teams library"
 readme = "README.md"
@@ -15,7 +15,7 @@ classifiers = [
     "Operating System :: OS Independent",
 ]
 dependencies = [
-    "autogen-core==0.5.6",
+    "autogen-core==0.5.7",
 ]
 
 [tool.ruff]

python/packages/autogen-agentchat/src/autogen_agentchat/agents/_assistant_agent.py

@@ -131,17 +131,26 @@ class AssistantAgent(BaseChatAgent, Component[AssistantAgentConfig]):
 
     * If the model returns no tool call, then the response is immediately returned as a :class:`~autogen_agentchat.messages.TextMessage` or a :class:`~autogen_agentchat.messages.StructuredMessage` (when using structured output) in :attr:`~autogen_agentchat.base.Response.chat_message`.
     * When the model returns tool calls, they will be executed right away:
-        - When `reflect_on_tool_use` is False, the tool call results are returned as a :class:`~autogen_agentchat.messages.ToolCallSummaryMessage` in :attr:`~autogen_agentchat.base.Response.chat_message`. `tool_call_summary_format` can be used to customize the tool call summary.
+        - When `reflect_on_tool_use` is False, the tool call results are returned as a :class:`~autogen_agentchat.messages.ToolCallSummaryMessage` in :attr:`~autogen_agentchat.base.Response.chat_message`. You can customise the summary with either a static format string (`tool_call_summary_format`) **or** a callable (`tool_call_summary_formatter`); the callable is evaluated once per tool call.
         - When `reflect_on_tool_use` is True, the another model inference is made using the tool calls and results, and final response is returned as a :class:`~autogen_agentchat.messages.TextMessage` or a :class:`~autogen_agentchat.messages.StructuredMessage` (when using structured output) in :attr:`~autogen_agentchat.base.Response.chat_message`.
         - `reflect_on_tool_use` is set to `True` by default when `output_content_type` is set.
         - `reflect_on_tool_use` is set to `False` by default when `output_content_type` is not set.
     * If the model returns multiple tool calls, they will be executed concurrently. To disable parallel tool calls you need to configure the model client. For example, set `parallel_tool_calls=False` for :class:`~autogen_ext.models.openai.OpenAIChatCompletionClient` and :class:`~autogen_ext.models.openai.AzureOpenAIChatCompletionClient`.
 
     .. tip::
-        By default, the tool call results are returned as response when tool calls are made.
-        So it is recommended to pay attention to the formatting of the tools return values,
-        especially if another agent is expecting them in a specific format.
-        Use `tool_call_summary_format` to customize the tool call summary, if needed.
+
+        By default, the tool call results are returned as the response when tool
+        calls are made, so pay close attention to how the tools’ return values
+        are formatted—especially if another agent expects a specific schema.
+
+        * Use **`tool_call_summary_format`** for a simple static template.
+        * Use **`tool_call_summary_formatter`** for full programmatic control
+          (e.g., “hide large success payloads, show full details on error”).
+
+        *Note*: `tool_call_summary_formatter` is **not serializable** and will
+        be ignored when an agent is loaded from, or exported to, YAML/JSON
+        configuration files.
+
 
     **Hand off behavior:**
 
@@ -199,13 +208,22 @@ class AssistantAgent(BaseChatAgent, Component[AssistantAgentConfig]):
             If this is set, the agent will respond with a :class:`~autogen_agentchat.messages.StructuredMessage` instead of a :class:`~autogen_agentchat.messages.TextMessage`
             in the final response, unless `reflect_on_tool_use` is `False` and a tool call is made.
         output_content_type_format (str | None, optional): (Experimental) The format string used for the content of a :class:`~autogen_agentchat.messages.StructuredMessage` response.
-        tool_call_summary_format (str, optional): The format string used to create the content for a :class:`~autogen_agentchat.messages.ToolCallSummaryMessage` response.
-            The format string is used to format the tool call summary for every tool call result.
-            Defaults to "{result}".
-            When `reflect_on_tool_use` is `False`, a concatenation of all the tool call summaries, separated by a new line character ('\\n')
-            will be returned as the response.
-            Available variables: `{tool_name}`, `{arguments}`, `{result}`.
-            For example, `"{tool_name}: {result}"` will create a summary like `"tool_name: result"`.
+        tool_call_summary_format (str, optional): Static format string applied to each tool call result when composing the :class:`~autogen_agentchat.messages.ToolCallSummaryMessage`.
+            Defaults to ``"{result}"``. Ignored if `tool_call_summary_formatter` is provided. When `reflect_on_tool_use` is ``False``, the summaries for all tool
+            calls are concatenated with a newline ('\\n') and returned as the response.  Placeholders available in the template:
+            `{tool_name}`, `{arguments}`, `{result}`, `{is_error}`.
+        tool_call_summary_formatter (Callable[[FunctionCall, FunctionExecutionResult], str] | None, optional):
+            Callable that receives the ``FunctionCall`` and its ``FunctionExecutionResult`` and returns the summary string.
+            Overrides `tool_call_summary_format` when supplied and allows conditional logic — for example, emitting static string like
+            ``"Tool FooBar executed successfully."`` on success and a full payload (including all passed arguments etc.) only on failure.
+
+            **Limitation**: The callable is *not serializable*; values provided via YAML/JSON configs are ignored.
+
+    .. note::
+
+        `tool_call_summary_formatter` is intended for in-code use only. It cannot currently be saved or restored via
+        configuration files.
+
         memory (Sequence[Memory] | None, optional): The memory store to use for the agent. Defaults to `None`.
         metadata (Dict[str, str] | None, optional): Optional metadata for tracking.
 
@@ -652,6 +670,7 @@ class AssistantAgent(BaseChatAgent, Component[AssistantAgentConfig]):
         model_client_stream: bool = False,
         reflect_on_tool_use: bool | None = None,
         tool_call_summary_format: str = "{result}",
+        tool_call_summary_formatter: Callable[[FunctionCall, FunctionExecutionResult], str] | None = None,
         output_content_type: type[BaseModel] | None = None,
         output_content_type_format: str | None = None,
         memory: Sequence[Memory] | None = None,
@@ -756,6 +775,7 @@ class AssistantAgent(BaseChatAgent, Component[AssistantAgentConfig]):
                 stacklevel=2,
             )
         self._tool_call_summary_format = tool_call_summary_format
+        self._tool_call_summary_formatter = tool_call_summary_formatter
         self._is_running = False
 
     @property
@@ -803,6 +823,7 @@ class AssistantAgent(BaseChatAgent, Component[AssistantAgentConfig]):
         model_client_stream = self._model_client_stream
         reflect_on_tool_use = self._reflect_on_tool_use
         tool_call_summary_format = self._tool_call_summary_format
+        tool_call_summary_formatter = self._tool_call_summary_formatter
         output_content_type = self._output_content_type
         format_string = self._output_content_type_format
 
@@ -873,6 +894,7 @@ class AssistantAgent(BaseChatAgent, Component[AssistantAgentConfig]):
             model_client_stream=model_client_stream,
             reflect_on_tool_use=reflect_on_tool_use,
             tool_call_summary_format=tool_call_summary_format,
+            tool_call_summary_formatter=tool_call_summary_formatter,
             output_content_type=output_content_type,
             format_string=format_string,
         ):
@@ -976,6 +998,7 @@ class AssistantAgent(BaseChatAgent, Component[AssistantAgentConfig]):
         model_client_stream: bool,
         reflect_on_tool_use: bool,
         tool_call_summary_format: str,
+        tool_call_summary_formatter: Callable[[FunctionCall, FunctionExecutionResult], str] | None,
         output_content_type: type[BaseModel] | None,
         format_string: str | None = None,
     ) -> AsyncGenerator[BaseAgentEvent | BaseChatMessage | Response, None]:
@@ -1078,6 +1101,7 @@ class AssistantAgent(BaseChatAgent, Component[AssistantAgentConfig]):
                 inner_messages=inner_messages,
                 handoffs=handoffs,
                 tool_call_summary_format=tool_call_summary_format,
+                tool_call_summary_formatter=tool_call_summary_formatter,
                 agent_name=agent_name,
             )
 
@@ -1232,6 +1256,7 @@ class AssistantAgent(BaseChatAgent, Component[AssistantAgentConfig]):
         inner_messages: List[BaseAgentEvent | BaseChatMessage],
         handoffs: Dict[str, HandoffBase],
         tool_call_summary_format: str,
+        tool_call_summary_formatter: Callable[[FunctionCall, FunctionExecutionResult], str] | None,
         agent_name: str,
     ) -> Response:
         """
@@ -1239,15 +1264,22 @@ class AssistantAgent(BaseChatAgent, Component[AssistantAgentConfig]):
         """
         # Filter out calls which were actually handoffs
         normal_tool_calls = [(call, result) for call, result in executed_calls_and_results if call.name not in handoffs]
-        tool_call_summaries: List[str] = []
-        for tool_call, tool_call_result in normal_tool_calls:
-            tool_call_summaries.append(
-                tool_call_summary_format.format(
-                    tool_name=tool_call.name,
-                    arguments=tool_call.arguments,
-                    result=tool_call_result.content,
-                )
+
+        def default_tool_call_summary_formatter(call: FunctionCall, result: FunctionExecutionResult) -> str:
+            return tool_call_summary_format
+
+        summary_formatter = tool_call_summary_formatter or default_tool_call_summary_formatter
+
+        tool_call_summaries = [
+            summary_formatter(call, result).format(
+                tool_name=call.name,
+                arguments=call.arguments,
+                result=result.content,
+                is_error=result.is_error,
             )
+            for call, result in normal_tool_calls
+        ]
+
         tool_call_summary = "\n".join(tool_call_summaries)
         return Response(
             chat_message=ToolCallSummaryMessage(

python/packages/autogen-agentchat/src/autogen_agentchat/agents/_code_executor_agent.py

@@ -326,8 +326,8 @@ class CodeExecutorAgent(BaseChatAgent, Component[CodeExecutorAgentConfig]):
     """
 
     DEFAULT_TERMINAL_DESCRIPTION = "A computer terminal that performs no other action than running Python scripts (provided to it quoted in ```python code blocks), or sh shell scripts (provided to it quoted in ```sh code blocks)."
-    DEFAULT_AGENT_DESCRIPTION = "A Code Execution Agent that generates and executes Python and shell scripts based on user instructions. Python code should be provided in ```python code blocks, and sh shell scripts should be provided in ```sh code blocks for execution. It ensures correctness, efficiency, and minimal errors while gracefully handling edge cases."
-    DEFAULT_SYSTEM_MESSAGE = "You are a Code Execution Agent. Your role is to generate and execute Python code based on user instructions, ensuring correctness, efficiency, and minimal errors. Handle edge cases gracefully."
+    DEFAULT_AGENT_DESCRIPTION = "A Code Execution Agent that generates and executes Python and shell scripts based on user instructions. It ensures correctness, efficiency, and minimal errors while gracefully handling edge cases."
+    DEFAULT_SYSTEM_MESSAGE = "You are a Code Execution Agent. Your role is to generate and execute Python code and shell scripts based on user instructions, ensuring correctness, efficiency, and minimal errors. Handle edge cases gracefully. Python code should be provided in ```python code blocks, and sh shell scripts should be provided in ```sh code blocks for execution."
     NO_CODE_BLOCKS_FOUND_MESSAGE = "No code blocks found in the thread. Please provide at least one markdown-encoded code block to execute (i.e., quoting code in ```python or ```sh code blocks)."
 
     component_config_schema = CodeExecutorAgentConfig

python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_base_group_chat_manager.py

@@ -2,7 +2,7 @@ import asyncio
 from abc import ABC, abstractmethod
 from typing import Any, List, Sequence
 
-from autogen_core import DefaultTopicId, MessageContext, event, rpc
+from autogen_core import CancellationToken, DefaultTopicId, MessageContext, event, rpc
 
 from ...base import TerminationCondition
 from ...messages import BaseAgentEvent, BaseChatMessage, MessageFactory, SelectSpeakerEvent, StopMessage
@@ -79,6 +79,7 @@ class BaseGroupChatManager(SequentialRoutedAgent, ABC):
         self._current_turn = 0
         self._message_factory = message_factory
         self._emit_team_events = emit_team_events
+        self._active_speakers: List[str] = []
 
     @rpc
     async def handle_start(self, message: GroupChatStart, ctx: MessageContext) -> None:
@@ -122,64 +123,64 @@ class BaseGroupChatManager(SequentialRoutedAgent, ABC):
                 # Stop the group chat.
                 return
 
-        # Select a speaker to start/continue the conversation
-        speaker_name_future = asyncio.ensure_future(self.select_speaker(self._message_thread))
-        # Link the select speaker future to the cancellation token.
-        ctx.cancellation_token.link_future(speaker_name_future)
-        speaker_name = await speaker_name_future
-        if speaker_name not in self._participant_name_to_topic_type:
-            raise RuntimeError(f"Speaker {speaker_name} not found in participant names.")
-        await self._log_speaker_selection(speaker_name)
-
-        # Send the message to the next speaker
-        speaker_topic_type = self._participant_name_to_topic_type[speaker_name]
-        await self.publish_message(
-            GroupChatRequestPublish(),
-            topic_id=DefaultTopicId(type=speaker_topic_type),
-            cancellation_token=ctx.cancellation_token,
-        )
-
-    async def update_message_thread(self, messages: Sequence[BaseAgentEvent | BaseChatMessage]) -> None:
-        self._message_thread.extend(messages)
+        # Select speakers to start/continue the conversation
+        await self._transition_to_next_speakers(ctx.cancellation_token)
 
     @event
     async def handle_agent_response(self, message: GroupChatAgentResponse, ctx: MessageContext) -> None:
         try:
-            # Append the message to the message thread and construct the delta.
+            # Construct the detla from the agent response.
             delta: List[BaseAgentEvent | BaseChatMessage] = []
             if message.agent_response.inner_messages is not None:
                 for inner_message in message.agent_response.inner_messages:
                     delta.append(inner_message)
             delta.append(message.agent_response.chat_message)
+
+            # Append the messages to the message thread.
             await self.update_message_thread(delta)
 
+            # Remove the agent from the active speakers list.
+            self._active_speakers.remove(message.agent_name)
+            if len(self._active_speakers) > 0:
+                # If there are still active speakers, return without doing anything.
+                return
+
             # Check if the conversation should be terminated.
             if await self._apply_termination_condition(delta, increment_turn_count=True):
                 # Stop the group chat.
                 return
 
-            # Select a speaker to continue the conversation.
-            speaker_name_future = asyncio.ensure_future(self.select_speaker(self._message_thread))
-            # Link the select speaker future to the cancellation token.
-            ctx.cancellation_token.link_future(speaker_name_future)
-            speaker_name = await speaker_name_future
+            # Select speakers to continue the conversation.
+            await self._transition_to_next_speakers(ctx.cancellation_token)
+        except Exception as e:
+            # Handle the exception and signal termination with an error.
+            error = SerializableException.from_exception(e)
+            await self._signal_termination_with_error(error)
+            # Raise the exception to the runtime.
+            raise
+
+    async def _transition_to_next_speakers(self, cancellation_token: CancellationToken) -> None:
+        speaker_names_future = asyncio.ensure_future(self.select_speaker(self._message_thread))
+        # Link the select speaker future to the cancellation token.
+        cancellation_token.link_future(speaker_names_future)
+        speaker_names = await speaker_names_future
+        if isinstance(speaker_names, str):
+            # If only one speaker is selected, convert it to a list.
+            speaker_names = [speaker_names]
+        for speaker_name in speaker_names:
             if speaker_name not in self._participant_name_to_topic_type:
                 raise RuntimeError(f"Speaker {speaker_name} not found in participant names.")
-            await self._log_speaker_selection(speaker_name)
+        await self._log_speaker_selection(speaker_names)
 
-            # Send the message to the next speakers
+        # Send request to publish message to the next speakers
+        for speaker_name in speaker_names:
             speaker_topic_type = self._participant_name_to_topic_type[speaker_name]
             await self.publish_message(
                 GroupChatRequestPublish(),
                 topic_id=DefaultTopicId(type=speaker_topic_type),
-                cancellation_token=ctx.cancellation_token,
+                cancellation_token=cancellation_token,
             )
-        except Exception as e:
-            # Handle the exception and signal termination with an error.
-            error = SerializableException.from_exception(e)
-            await self._signal_termination_with_error(error)
-            # Raise the exception to the runtime.
-            raise
+            self._active_speakers.append(speaker_name)
 
     async def _apply_termination_condition(
         self, delta: Sequence[BaseAgentEvent | BaseChatMessage], increment_turn_count: bool = False
@@ -216,9 +217,9 @@ class BaseGroupChatManager(SequentialRoutedAgent, ABC):
                 return True
         return False
 
-    async def _log_speaker_selection(self, speaker_name: str) -> None:
+    async def _log_speaker_selection(self, speaker_names: List[str]) -> None:
         """Log the selected speaker to the output message queue."""
-        select_msg = SelectSpeakerEvent(content=[speaker_name], source=self._name)
+        select_msg = SelectSpeakerEvent(content=speaker_names, source=self._name)
         if self._emit_team_events:
             await self.publish_message(
                 GroupChatMessage(message=select_msg),
@@ -284,10 +285,26 @@ class BaseGroupChatManager(SequentialRoutedAgent, ABC):
         """
         ...
 
+    async def update_message_thread(self, messages: Sequence[BaseAgentEvent | BaseChatMessage]) -> None:
+        """Update the message thread with the new messages.
+        This is called when the group chat receives a GroupChatStart or GroupChatAgentResponse event,
+        before calling the select_speakers method.
+        """
+        self._message_thread.extend(messages)
+
     @abstractmethod
-    async def select_speaker(self, thread: List[BaseAgentEvent | BaseChatMessage]) -> str:
-        """Select a speaker from the participants and return the
-        topic type of the selected speaker."""
+    async def select_speaker(self, thread: Sequence[BaseAgentEvent | BaseChatMessage]) -> List[str] | str:
+        """Select speakers from the participants and return the topic types of the selected speaker.
+        This is called when the group chat manager have received all responses from the participants
+        for a turn and is ready to select the next speakers for the next turn.
+
+        Args:
+            thread: The message thread of the group chat.
+
+        Returns:
+            A list of topic types of the selected speakers.
+            If only one speaker is selected, a single string is returned instead of a list.
+        """
         ...
 
     @abstractmethod

python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_chat_agent_container.py

@@ -89,7 +89,7 @@ class ChatAgentContainer(SequentialRoutedAgent):
             # Publish the response to the group chat.
             self._message_buffer.clear()
             await self.publish_message(
-                GroupChatAgentResponse(agent_response=response),
+                GroupChatAgentResponse(agent_response=response, agent_name=self._agent.name),
                 topic_id=DefaultTopicId(type=self._parent_topic_type),
                 cancellation_token=ctx.cancellation_token,
             )

python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_events.py

@@ -48,6 +48,9 @@ class GroupChatAgentResponse(BaseModel):
     agent_response: Response
     """The response from an agent."""
 
+    agent_name: str
+    """The name of the agent that produced the response."""
+
 
 class GroupChatRequestPublish(BaseModel):
     """A request to publish a message to a group chat."""

python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_graph/_digraph_group_chat.py

@@ -1,5 +1,6 @@
 import asyncio
-from typing import Any, Callable, Dict, List, Literal, Mapping, Sequence, Set
+from collections import Counter, deque
+from typing import Any, Callable, Deque, Dict, List, Literal, Mapping, Sequence, Set
 
 from autogen_core import AgentRuntime, CancellationToken, Component, ComponentModel
 from pydantic import BaseModel
@@ -208,154 +209,70 @@ class GraphFlowManager(BaseGroupChatManager):
             max_turns=max_turns,
             message_factory=message_factory,
         )
-        self._graph = graph
-        self._graph.graph_validate()
-        self._graph_has_cycles = self._graph.get_has_cycles()
-        if self._graph_has_cycles and self._termination_condition is None and self._max_turns is None:
+        graph.graph_validate()
+        if graph.get_has_cycles() and self._termination_condition is None and self._max_turns is None:
             raise ValueError("A termination condition is required for cyclic graphs without a maximum turn limit.")
-
-        self._use_default_start = self._graph.default_start_node is not None
-        self._default_start_executed = False
-        self._start_nodes = graph.get_start_nodes()
-        self._leaf_nodes = graph.get_leaf_nodes()
-        self._parents = graph.get_parents()  # Parent node dependencies - helper dict to get all incoming edges
-        self._active_nodes: Set[str] = set()  # Currently executing nodes
-        self._active_node_count: Dict[str, int] = {
-            node: 0 for node in graph.nodes
-        }  # Number of times a node has been active
-
-        # These are nodes next in line for execution as one or more of their parent nodes have started execution.
-        # They execute when all their parent nodes have executed.
-        # Nodes are added to this dict when at least one of their parent nodes becomes active.
-        # Start nodes (no parents) are added to this dict at initialization as they are always ready to run.
-        self._pending_execution: Dict[str, List[str]] = {node: [] for node in graph.get_start_nodes()}
-
-    def _get_valid_target(self, node: DiGraphNode, content: str) -> str:
-        """Check if a condition is met in the chat history."""
-        for edge in node.edges:
-            if edge.condition and edge.condition in content:
-                return edge.target
-
-        raise RuntimeError(f"Condition not met for node {node.name}. Content: {content}")
-
-    def _is_node_ready(self, node_name: str) -> bool:
-        """Check if a node is ready to execute based on its parent nodes.
-        If activation is any then execute as soon as any parent has finished
-        If activation is all then execute only when all parents have finished
-        """
-        node = self._graph.nodes[node_name]
-        if node.activation == "any":
-            return bool(self._pending_execution[node_name])
-        return all(parent in self._pending_execution[node_name] for parent in self._parents[node_name])
-
-    async def _select_speakers(self, thread: List[BaseAgentEvent | BaseChatMessage], many: bool = True) -> List[str]:
-        """Select the next set of agents to execute based on DAG constraints."""
-        next_speakers: Set[str] = set()
-        source_node: DiGraphNode | None = None
-        source: str | None = None
-
-        if thread and isinstance(thread[-1], BaseChatMessage):
-            source = thread[-1].source  # name of the agent that just finished
-            content = thread[-1].to_model_text()
-
-            # Safety check: only an active node can send a response
-            if source != "user":
-                if source not in self._active_nodes:
-                    raise RuntimeError(f"Agent '{source}' is not currently active.")
-
-                # Mark the node as no longer active (it just finished)
-                self._active_node_count[source] -= 1
-
-                if self._active_node_count[source] <= 0:
-                    self._active_nodes.remove(source)
-
-                source_node = self._graph.nodes[source]
-
-                if source_node.edges:
-                    # Case: conditional edges — only execute if condition is met
-                    target_nodes_names: List[str] = []
-                    if source_node.edges[0].condition is not None:
-                        target_nodes_names = [self._get_valid_target(source_node, content)]
-                        other_nodes = [
-                            edge.target for edge in source_node.edges if edge.target != target_nodes_names[0]
-                        ]
-                        for other_node in other_nodes:
-                            other_active_parents = [
-                                parent
-                                for parent in self._parents[other_node]
-                                if (parent != source and parent in self._active_nodes)
-                            ]
-                            if not other_active_parents:
-                                self._pending_execution.pop(other_node)
-                            else:
-                                self._pending_execution[other_node] = other_active_parents
-
-                    else:
-                        # Case: unconditional edges — mark this source as completed for all its children
-                        target_nodes_names = [edge.target for edge in source_node.edges]
-
-                    for target in target_nodes_names:
-                        self._pending_execution[target].append(source)
+        self._graph = graph
+        # Lookup table for incoming edges for each node.
+        self._parents = graph.get_parents()
+        # Lookup table for outgoing edges for each node.
+        self._edges: Dict[str, List[DiGraphEdge]] = {n: node.edges for n, node in graph.nodes.items()}
+        # Activation lookup table for each node.
+        self._activation: Dict[str, Literal["any", "all"]] = {n: node.activation for n, node in graph.nodes.items()}
+
+        # === Mutable states for the graph execution ===
+        # Count the number of remaining parents to activate each node.
+        self._remaining: Counter[str] = Counter({n: len(p) for n, p in self._parents.items()})
+        # Lookup table for nodes that have been enqueued through an any activation.
+        # This is used to prevent re-adding the same node multiple times.
+        self._enqueued_any: Dict[str, bool] = {n: False for n in graph.nodes}
+        # Ready queue for nodes that are ready to execute, starting with the start nodes.
+        self._ready: Deque[str] = deque([n for n in graph.get_start_nodes()])
+
+    async def update_message_thread(self, messages: Sequence[BaseAgentEvent | BaseChatMessage]) -> None:
+        await super().update_message_thread(messages)
+
+        # Find the node that ran in the current turn.
+        message = messages[-1]
+        if message.source not in self._graph.nodes:
+            # Ignore messages from sources outside of the graph.
+            return
+        assert isinstance(message, BaseChatMessage)
+        source = message.source
+        content = message.to_model_text()
+
+        # Propagate the update to the children of the node.
+        for edge in self._edges[source]:
+            if edge.condition and edge.condition not in content:
+                continue
+            if self._activation[edge.target] == "all":
+                self._remaining[edge.target] -= 1
+                if self._remaining[edge.target] == 0:
+                    # If all parents are done, add to the ready queue.
+                    self._ready.append(edge.target)
             else:
-                # TODO: Check if there are any usecase where the User can decide on the next speaker
-                pass
-
-        # After updating _pending_execution, check which nodes are now unblocked
-        for node_name in list(self._pending_execution):
-            if self._use_default_start and not self._default_start_executed:
-                if node_name == self._graph.default_start_node:
-                    next_speakers.add(node_name)
-                    self._default_start_executed = True
-                    break
-
-            if self._is_node_ready(node_name):
-                next_speakers.add(node_name)
-                node = self._graph.nodes[node_name]
-                if node.activation == "all":
-                    self._pending_execution.pop(node_name)
-                else:
-                    # If activation is any, remove the parent that just finished
-                    if source is not None:
-                        self._pending_execution[node_name] = [
-                            parent for parent in self._pending_execution[node_name] if parent != source
-                        ]
-
-                    # If none of the other parents of this node are active, remove this node from pending execution
-                    node_parents = self._parents[node_name]
-                    if not any(parent in self._active_nodes for parent in node_parents):
-                        self._pending_execution.pop(node_name)
-
-                if not many:
-                    break
-
-        # Prepopulate children of next_speakers into _pending_execution
-        for node_name in next_speakers:
-            for edge in self._graph.nodes[node_name].edges:
-                if edge.target not in self._pending_execution:
-                    self._pending_execution[edge.target] = []
-
-        # Mark newly selected speakers as active
-        for speaker in next_speakers:
-            if speaker not in self._active_nodes:
-                self._active_nodes.add(speaker)
-
-            self._active_node_count[speaker] += 1
-
-        if not self._pending_execution and not next_speakers and not self._active_nodes:
-            next_speakers = set([_DIGRAPH_STOP_AGENT_NAME])  # Call the termination agent
-
-        return list(next_speakers)
-
-    async def select_speakers(self, thread: List[BaseAgentEvent | BaseChatMessage]) -> List[str]:
-        return await self._select_speakers(thread)
-
-    async def select_speaker(self, thread: List[BaseAgentEvent | BaseChatMessage]) -> str:
-        """Select a speaker from the participants and return the
-        topic type of the selected speaker."""
-        speakers = await self._select_speakers(thread, many=False)
+                # If activation is any, add to the ready queue if not already enqueued.
+                if not self._enqueued_any[edge.target]:
+                    self._ready.append(edge.target)
+                    self._enqueued_any[edge.target] = True
+
+    async def select_speaker(self, thread: Sequence[BaseAgentEvent | BaseChatMessage]) -> List[str]:
+        # Drain the ready queue for the next set of speakers.
+        speakers: List[str] = []
+        while self._ready:
+            speaker = self._ready.popleft()
+            speakers.append(speaker)
+            # Reset the bookkeeping for the node that were selected.
+            if self._activation[speaker] == "any":
+                self._enqueued_any[speaker] = False
+            else:
+                self._remaining[speaker] = len(self._parents[speaker])
+
+        # If there are no speakers, trigger the stop agent.
         if not speakers:
-            raise RuntimeError("No available speakers found.")
-        return speakers[0]
+            speakers = [_DIGRAPH_STOP_AGENT_NAME]
+
+        return speakers
 
     async def validate_group_state(self, messages: List[BaseChatMessage] | None) -> None:
         pass
@@ -365,10 +282,9 @@ class GraphFlowManager(BaseGroupChatManager):
         state = {
             "message_thread": [message.dump() for message in self._message_thread],
             "current_turn": self._current_turn,
-            "active_nodes": list(self._active_nodes),
-            "pending_execution": self._pending_execution,
-            "active_node_count": self._active_node_count,
-            "default_start_executed": self._default_start_executed,
+            "remaining": dict(self._remaining),
+            "enqueued_any": dict(self._enqueued_any),
+            "ready": list(self._ready),
         }
         return state
 
@@ -376,10 +292,9 @@ class GraphFlowManager(BaseGroupChatManager):
         """Restore execution state from saved data."""
         self._message_thread = [self._message_factory.create(msg) for msg in state["message_thread"]]
         self._current_turn = state["current_turn"]
-        self._active_nodes = set(state["active_nodes"])
-        self._pending_execution = state["pending_execution"]
-        self._active_node_count = state["active_node_count"]
-        self._default_start_executed = state.get("default_start_executed", False)
+        self._remaining = Counter(state["remaining"])
+        self._enqueued_any = state["enqueued_any"]
+        self._ready = deque(state["ready"])
 
     async def reset(self) -> None:
         """Reset execution state to the start of the graph."""
@@ -387,11 +302,9 @@ class GraphFlowManager(BaseGroupChatManager):
         self._message_thread.clear()
         if self._termination_condition:
             await self._termination_condition.reset()
-
-        self._active_nodes = set()
-        self._active_node_count = {node: 0 for node in self._graph.nodes}
-        self._pending_execution = {node: [] for node in self._start_nodes}
-        self._default_start_executed = False
+        self._remaining = Counter({n: len(p) for n, p in self._parents.items()})
+        self._enqueued_any = {n: False for n in self._graph.nodes}
+        self._ready = deque([n for n in self._graph.get_start_nodes()])
 
 
 class _StopAgent(BaseChatAgent):

python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_magentic_one/_magentic_one_orchestrator.py

@@ -2,7 +2,7 @@ import asyncio
 import json
 import logging
 import re
-from typing import Any, Dict, List, Mapping
+from typing import Any, Dict, List, Mapping, Sequence
 
 from autogen_core import AgentId, CancellationToken, DefaultTopicId, MessageContext, event, rpc
 from autogen_core.models import (
@@ -37,6 +37,7 @@ from .._events import (
     GroupChatReset,
     GroupChatStart,
     GroupChatTermination,
+    SerializableException,
 )
 from ._prompts import (
     ORCHESTRATOR_FINAL_ANSWER_PROMPT,
@@ -187,22 +188,29 @@ class MagenticOneOrchestrator(BaseGroupChatManager):
 
     @event
     async def handle_agent_response(self, message: GroupChatAgentResponse, ctx: MessageContext) -> None:  # type: ignore
-        delta: List[BaseAgentEvent | BaseChatMessage] = []
-        if message.agent_response.inner_messages is not None:
-            for inner_message in message.agent_response.inner_messages:
-                delta.append(inner_message)
-        await self.update_message_thread([message.agent_response.chat_message])
-        delta.append(message.agent_response.chat_message)
-
-        if self._termination_condition is not None:
-            stop_message = await self._termination_condition(delta)
-            if stop_message is not None:
-                # Reset the termination conditions.
-                await self._termination_condition.reset()
-                # Signal termination.
-                await self._signal_termination(stop_message)
-                return
-        await self._orchestrate_step(ctx.cancellation_token)
+        try:
+            delta: List[BaseAgentEvent | BaseChatMessage] = []
+            if message.agent_response.inner_messages is not None:
+                for inner_message in message.agent_response.inner_messages:
+                    delta.append(inner_message)
+            await self.update_message_thread([message.agent_response.chat_message])
+            delta.append(message.agent_response.chat_message)
+
+            if self._termination_condition is not None:
+                stop_message = await self._termination_condition(delta)
+                if stop_message is not None:
+                    # Reset the termination conditions.
+                    await self._termination_condition.reset()
+                    # Signal termination.
+                    await self._signal_termination(stop_message)
+                    return
+
+            await self._orchestrate_step(ctx.cancellation_token)
+        except Exception as e:
+            error = SerializableException.from_exception(e)
+            await self._signal_termination_with_error(error)
+            # Raise the error to the runtime.
+            raise
 
     async def validate_group_state(self, messages: List[BaseChatMessage] | None) -> None:
         pass
@@ -229,9 +237,9 @@ class MagenticOneOrchestrator(BaseGroupChatManager):
         self._n_rounds = orchestrator_state.n_rounds
         self._n_stalls = orchestrator_state.n_stalls
 
-    async def select_speaker(self, thread: List[BaseAgentEvent | BaseChatMessage]) -> str:
+    async def select_speaker(self, thread: Sequence[BaseAgentEvent | BaseChatMessage]) -> List[str] | str:
         """Not used in this orchestrator, we select next speaker in _orchestrate_step."""
-        return ""
+        return [""]
 
     async def reset(self) -> None:
         """Reset the group chat manager."""
@@ -275,7 +283,7 @@ class MagenticOneOrchestrator(BaseGroupChatManager):
 
         # Broadcast
         await self.publish_message(
-            GroupChatAgentResponse(agent_response=Response(chat_message=ledger_message)),
+            GroupChatAgentResponse(agent_response=Response(chat_message=ledger_message), agent_name=self._name),
             topic_id=DefaultTopicId(type=self._group_topic_type),
         )
 
@@ -389,7 +397,7 @@ class MagenticOneOrchestrator(BaseGroupChatManager):
 
         # Broadcast it
         await self.publish_message(  # Broadcast
-            GroupChatAgentResponse(agent_response=Response(chat_message=message)),
+            GroupChatAgentResponse(agent_response=Response(chat_message=message), agent_name=self._name),
             topic_id=DefaultTopicId(type=self._group_topic_type),
             cancellation_token=cancellation_token,
         )
@@ -470,7 +478,7 @@ class MagenticOneOrchestrator(BaseGroupChatManager):
 
         # Broadcast
         await self.publish_message(
-            GroupChatAgentResponse(agent_response=Response(chat_message=message)),
+            GroupChatAgentResponse(agent_response=Response(chat_message=message), agent_name=self._name),
             topic_id=DefaultTopicId(type=self._group_topic_type),
             cancellation_token=cancellation_token,
         )

python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_round_robin_group_chat.py

@@ -1,5 +1,5 @@
 import asyncio
-from typing import Any, Callable, List, Mapping
+from typing import Any, Callable, List, Mapping, Sequence
 
 from autogen_core import AgentRuntime, Component, ComponentModel
 from pydantic import BaseModel
@@ -69,8 +69,13 @@ class RoundRobinGroupChatManager(BaseGroupChatManager):
         self._current_turn = round_robin_state.current_turn
         self._next_speaker_index = round_robin_state.next_speaker_index
 
-    async def select_speaker(self, thread: List[BaseAgentEvent | BaseChatMessage]) -> str:
-        """Select a speaker from the participants in a round-robin fashion."""
+    async def select_speaker(self, thread: Sequence[BaseAgentEvent | BaseChatMessage]) -> List[str] | str:
+        """Select a speaker from the participants in a round-robin fashion.
+
+        .. note::
+
+            This method always returns a single speaker.
+        """
         current_speaker_index = self._next_speaker_index
         self._next_speaker_index = (current_speaker_index + 1) % len(self._participant_names)
         current_speaker = self._participant_names[current_speaker_index]

python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_selector_group_chat.py

@@ -150,10 +150,14 @@ class SelectorGroupChatManager(BaseGroupChatManager):
         base_chat_messages = [m for m in messages if isinstance(m, BaseChatMessage)]
         await self._add_messages_to_context(self._model_context, base_chat_messages)
 
-    async def select_speaker(self, thread: List[BaseAgentEvent | BaseChatMessage]) -> str:
+    async def select_speaker(self, thread: Sequence[BaseAgentEvent | BaseChatMessage]) -> List[str] | str:
         """Selects the next speaker in a group chat using a ChatCompletion client,
         with the selector function as override if it returns a speaker name.
 
+        .. note::
+
+            This method always returns a single speaker name.
+
         A key assumption is that the agent type is the same as the topic type, which we use as the agent name.
         """
         # Use the selector function if provided.
@@ -171,7 +175,7 @@ class SelectorGroupChatManager(BaseGroupChatManager):
                         f"Expected one of: {self._participant_names}."
                     )
                 # Skip the model based selection.
-                return speaker
+                return [speaker]
 
         # Use the candidate function to filter participants if provided
         if self._candidate_func is not None:
@@ -211,7 +215,7 @@ class SelectorGroupChatManager(BaseGroupChatManager):
             agent_name = participants[0]
         self._previous_speaker = agent_name
         trace_logger.debug(f"Selected speaker: {agent_name}")
-        return agent_name
+        return [agent_name]
 
     def construct_message_history(self, message_history: List[LLMMessage]) -> str:
         # Construct the history of the conversation.

python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_swarm_group_chat.py

@@ -1,5 +1,5 @@
 import asyncio
-from typing import Any, Callable, List, Mapping
+from typing import Any, Callable, List, Mapping, Sequence
 
 from autogen_core import AgentRuntime, Component, ComponentModel
 from pydantic import BaseModel
@@ -79,17 +79,22 @@ class SwarmGroupChatManager(BaseGroupChatManager):
             await self._termination_condition.reset()
         self._current_speaker = self._participant_names[0]
 
-    async def select_speaker(self, thread: List[BaseAgentEvent | BaseChatMessage]) -> str:
+    async def select_speaker(self, thread: Sequence[BaseAgentEvent | BaseChatMessage]) -> List[str] | str:
         """Select a speaker from the participants based on handoff message.
-        Looks for the last handoff message in the thread to determine the next speaker."""
+        Looks for the last handoff message in the thread to determine the next speaker.
+
+        .. note::
+
+            This method always returns a single speaker.
+        """
         if len(thread) == 0:
-            return self._current_speaker
+            return [self._current_speaker]
         for message in reversed(thread):
             if isinstance(message, HandoffMessage):
                 self._current_speaker = message.target
                 # The latest handoff message should always target a valid participant.
                 assert self._current_speaker in self._participant_names
-                return self._current_speaker
+                return [self._current_speaker]
         return self._current_speaker
 
     async def save_state(self) -> Mapping[str, Any]:

python/packages/autogen-agentchat/tests/test_assistant_agent.py

@@ -32,7 +32,7 @@ from autogen_core.models import (
     SystemMessage,
     UserMessage,
 )
-from autogen_core.models._model_client import ModelFamily
+from autogen_core.models._model_client import ModelFamily, ModelInfo
 from autogen_core.tools import BaseTool, FunctionTool, StaticWorkbench
 from autogen_ext.models.openai import OpenAIChatCompletionClient
 from autogen_ext.models.replay import ReplayChatCompletionClient
@@ -56,10 +56,68 @@ async def _fail_function(input: str) -> str:
     return "fail"
 
 
+async def _throw_function(input: str) -> str:
+    raise ValueError("Helpful debugging information what went wrong.")
+
+
 async def _echo_function(input: str) -> str:
     return input
 
 
+@pytest.fixture
+def model_info_all_capabilities() -> ModelInfo:
+    return {
+        "function_calling": True,
+        "vision": True,
+        "json_output": True,
+        "family": ModelFamily.GPT_4O,
+        "structured_output": True,
+    }
+
+
+@pytest.mark.asyncio
+async def test_run_with_tool_call_summary_format_function(model_info_all_capabilities: ModelInfo) -> None:
+    model_client = ReplayChatCompletionClient(
+        [
+            CreateResult(
+                finish_reason="function_calls",
+                content=[
+                    FunctionCall(id="1", arguments=json.dumps({"input": "task"}), name="_pass_function"),
+                    FunctionCall(id="2", arguments=json.dumps({"input": "task"}), name="_throw_function"),
+                ],
+                usage=RequestUsage(prompt_tokens=10, completion_tokens=5),
+                thought="Calling pass and fail function",
+                cached=False,
+            ),
+        ],
+        model_info=model_info_all_capabilities,
+    )
+
+    def conditional_string_templates(function_call: FunctionCall, function_call_result: FunctionExecutionResult) -> str:
+        if not function_call_result.is_error:
+            return "SUCCESS: {tool_name} with {arguments}"
+
+        else:
+            return "FAILURE: {result}"
+
+    agent = AssistantAgent(
+        "tool_use_agent",
+        model_client=model_client,
+        tools=[_pass_function, _throw_function],
+        tool_call_summary_formatter=conditional_string_templates,
+    )
+    result = await agent.run(task="task")
+
+    first_tool_call_summary = next((x for x in result.messages if isinstance(x, ToolCallSummaryMessage)), None)
+    if first_tool_call_summary is None:
+        raise AssertionError("Expected a ToolCallSummaryMessage but found none")
+
+    assert (
+        first_tool_call_summary.content
+        == 'SUCCESS: _pass_function with {"input": "task"}\nFAILURE: Helpful debugging information what went wrong.'
+    )
+
+
 @pytest.mark.asyncio
 async def test_run_with_tools(monkeypatch: pytest.MonkeyPatch) -> None:
     model_client = ReplayChatCompletionClient(

python/packages/autogen-agentchat/tests/test_group_chat.py

@@ -6,24 +6,6 @@ from typing import Any, AsyncGenerator, Dict, List, Mapping, Sequence
 
 import pytest
 import pytest_asyncio
-from autogen_core import AgentId, AgentRuntime, CancellationToken, FunctionCall, SingleThreadedAgentRuntime
-from autogen_core.model_context import BufferedChatCompletionContext
-from autogen_core.models import (
-    AssistantMessage,
-    CreateResult,
-    FunctionExecutionResult,
-    FunctionExecutionResultMessage,
-    LLMMessage,
-    RequestUsage,
-    UserMessage,
-)
-from autogen_core.tools import FunctionTool
-from autogen_ext.code_executors.local import LocalCommandLineCodeExecutor
-from autogen_ext.models.openai import OpenAIChatCompletionClient
-from autogen_ext.models.replay import ReplayChatCompletionClient
-from pydantic import BaseModel
-from utils import FileLogHandler
-
 from autogen_agentchat import EVENT_LOGGER_NAME
 from autogen_agentchat.agents import (
     AssistantAgent,
@@ -57,6 +39,23 @@ from autogen_agentchat.teams._group_chat._round_robin_group_chat import RoundRob
 from autogen_agentchat.teams._group_chat._selector_group_chat import SelectorGroupChatManager
 from autogen_agentchat.teams._group_chat._swarm_group_chat import SwarmGroupChatManager
 from autogen_agentchat.ui import Console
+from autogen_core import AgentId, AgentRuntime, CancellationToken, FunctionCall, SingleThreadedAgentRuntime
+from autogen_core.model_context import BufferedChatCompletionContext
+from autogen_core.models import (
+    AssistantMessage,
+    CreateResult,
+    FunctionExecutionResult,
+    FunctionExecutionResultMessage,
+    LLMMessage,
+    RequestUsage,
+    UserMessage,
+)
+from autogen_core.tools import FunctionTool
+from autogen_ext.code_executors.local import LocalCommandLineCodeExecutor
+from autogen_ext.models.openai import OpenAIChatCompletionClient
+from autogen_ext.models.replay import ReplayChatCompletionClient
+from pydantic import BaseModel
+from utils import FileLogHandler
 
 logger = logging.getLogger(EVENT_LOGGER_NAME)
 logger.setLevel(logging.DEBUG)

python/packages/autogen-agentchat/tests/test_group_chat_graph.py

@@ -1,6 +1,6 @@
 import asyncio
-from typing import Any, AsyncGenerator, Callable, Dict, List, Sequence, Set
-from unittest.mock import AsyncMock, patch
+from typing import AsyncGenerator, List, Sequence
+from unittest.mock import patch
 
 import pytest
 import pytest_asyncio
@@ -12,9 +12,8 @@ from autogen_agentchat.agents import (
     PerSourceFilter,
 )
 from autogen_agentchat.base import Response, TaskResult
-from autogen_agentchat.conditions import MaxMessageTermination
+from autogen_agentchat.conditions import MaxMessageTermination, SourceMatchTermination
 from autogen_agentchat.messages import BaseChatMessage, ChatMessage, MessageFactory, StopMessage, TextMessage
-from autogen_agentchat.messages import BaseTextChatMessage as TextChatMessage
 from autogen_agentchat.teams import (
     DiGraphBuilder,
     GraphFlow,
@@ -269,61 +268,6 @@ def test_validate_graph_mixed_conditions() -> None:
         graph.graph_validate()
 
 
-def test_get_valid_target() -> None:
-    node = DiGraphNode(
-        name="A",
-        edges=[DiGraphEdge(target="B", condition="approve"), DiGraphEdge(target="C", condition="reject")],
-    )
-    manager = GraphFlowManager.__new__(GraphFlowManager)
-
-    assert manager._get_valid_target(node, "please approve this") == "B"  # pyright: ignore[reportPrivateUsage]
-    assert manager._get_valid_target(node, "i reject this") == "C"  # pyright: ignore[reportPrivateUsage]
-    with pytest.raises(RuntimeError):
-        manager._get_valid_target(node, "unknown path")  # pyright: ignore[reportPrivateUsage]
-
-
-def test_is_node_ready_all_and_any() -> None:
-    graph = DiGraph(
-        nodes={
-            "A": DiGraphNode(name="A", edges=[DiGraphEdge(target="C")]),
-            "B": DiGraphNode(name="B", edges=[DiGraphEdge(target="C")]),
-            "C": DiGraphNode(name="C", edges=[], activation="all"),
-        }
-    )
-
-    manager = GraphFlowManager.__new__(GraphFlowManager)
-    manager._graph = graph  # pyright: ignore[reportPrivateUsage]
-    manager._parents = graph.get_parents()  # pyright: ignore[reportPrivateUsage]
-
-    # === Test "all" activation ===
-    # Case 1: No parent finished
-    manager._pending_execution = {"C": []}  # pyright: ignore[reportPrivateUsage]
-    assert not manager._is_node_ready("C")  # pyright: ignore[reportPrivateUsage]
-
-    # Case 2: One parent finished
-    manager._pending_execution = {"C": ["A"]}  # pyright: ignore[reportPrivateUsage]
-    assert not manager._is_node_ready("C")  # pyright: ignore[reportPrivateUsage]
-
-    # Case 3: All parents finished
-    manager._pending_execution = {"C": ["A", "B"]}  # pyright: ignore[reportPrivateUsage]
-    assert manager._is_node_ready("C")  # pyright: ignore[reportPrivateUsage]
-
-    # === Test "any" activation ===
-    graph.nodes["C"].activation = "any"
-
-    # Case 1: No parent finished
-    manager._pending_execution = {"C": []}  # pyright: ignore[reportPrivateUsage]
-    assert not manager._is_node_ready("C")  # pyright: ignore[reportPrivateUsage]
-
-    # Case 2: One parent finished
-    manager._pending_execution = {"C": ["B"]}  # pyright: ignore[reportPrivateUsage]
-    assert manager._is_node_ready("C")  # pyright: ignore[reportPrivateUsage]
-
-    # Case 3: All parents finished
-    manager._pending_execution = {"C": ["A", "B"]}  # pyright: ignore[reportPrivateUsage]
-    assert manager._is_node_ready("C")  # pyright: ignore[reportPrivateUsage]
-
-
 @pytest.mark.asyncio
 async def test_invalid_digraph_manager_cycle_without_termination() -> None:
     """Test GraphManager raises error for cyclic graph without termination condition."""
@@ -359,170 +303,6 @@ async def test_invalid_digraph_manager_cycle_without_termination() -> None:
             )
 
 
-@pytest.fixture
-def digraph_manager() -> Callable[..., GraphFlowManager]:
-    @patch(
-        "autogen_agentchat.teams._group_chat._base_group_chat_manager.BaseGroupChatManager.__init__", return_value=None
-    )
-    def _create(
-        _: Any,
-        graph: DiGraph,
-        active_nodes: Set[str] | None = None,
-        thread: List[BaseAgentEvent | BaseChatMessage] | None = None,
-        pending: Dict[str, List[str]] | None = None,
-    ) -> GraphFlowManager:
-        manager = GraphFlowManager.__new__(GraphFlowManager)
-        manager._graph = graph  # pyright: ignore[reportPrivateUsage]
-        manager._parents = graph.get_parents()  # pyright: ignore[reportPrivateUsage]
-        manager._start_nodes = graph.get_start_nodes()  # pyright: ignore[reportPrivateUsage]
-        manager._leaf_nodes = graph.get_leaf_nodes()  # pyright: ignore[reportPrivateUsage]
-        manager._active_nodes = set(active_nodes or [])  # pyright: ignore[reportPrivateUsage]
-        manager._active_node_count = {node: 0 for node in graph.nodes}  # pyright: ignore[reportPrivateUsage]
-        manager._message_factory = MessageFactory()  # pyright: ignore[reportPrivateUsage]
-        manager._message_thread = thread if thread is not None else []  # pyright: ignore[reportPrivateUsage]
-        manager._pending_execution = pending if pending is not None else {node: [] for node in graph.get_start_nodes()}  # pyright: ignore[reportPrivateUsage]
-        manager._name = "test_manager"  # pyright: ignore[reportPrivateUsage]
-        manager._use_default_start = False  # pyright: ignore[reportPrivateUsage]
-        return manager
-
-    return _create
-
-
-# -------------------- Test: Sequential Flow --------------------
-@pytest.mark.asyncio
-async def test_select_speakers_linear(digraph_manager: Callable[..., GraphFlowManager]) -> None:
-    graph = DiGraph(
-        nodes={
-            "A": DiGraphNode(name="A", edges=[DiGraphEdge(target="B")]),
-            "B": DiGraphNode(name="B", edges=[DiGraphEdge(target="C")]),
-            "C": DiGraphNode(name="C", edges=[]),
-        }
-    )
-    message_thread = [TextChatMessage(source="A", content="done", metadata={})]
-    manager = digraph_manager(graph=graph, active_nodes={"A"}, thread=message_thread, pending={"B": [], "C": []})
-
-    result = await manager.select_speakers(manager._message_thread)  # pyright: ignore[reportPrivateUsage]
-    assert result == ["B"]
-    assert "B" in manager._active_nodes  # pyright: ignore[reportPrivateUsage]
-
-
-# -------------------- Test: Parallel Fan-out --------------------
-
-
-@pytest.mark.asyncio
-async def test_select_speakers_parallel(digraph_manager: Callable[..., GraphFlowManager]) -> None:
-    graph = DiGraph(
-        nodes={
-            "A": DiGraphNode(name="A", edges=[DiGraphEdge(target="B"), DiGraphEdge(target="C")]),
-            "B": DiGraphNode(name="B", edges=[]),
-            "C": DiGraphNode(name="C", edges=[]),
-        }
-    )
-    message_thread = [TextChatMessage(source="A", content="done", metadata={})]
-    manager = digraph_manager(graph=graph, active_nodes={"A"}, thread=message_thread, pending={"B": [], "C": []})
-
-    result = await manager.select_speakers(manager._message_thread)  # pyright: ignore[reportPrivateUsage]
-    assert set(result) == {"B", "C"}
-    assert "B" in manager._active_nodes  # pyright: ignore[reportPrivateUsage]
-    assert "C" in manager._active_nodes  # pyright: ignore[reportPrivateUsage]
-
-
-# -------------------- Test: Conditional Path --------------------
-@pytest.mark.asyncio
-async def test_select_speakers_conditional(digraph_manager: Callable[..., GraphFlowManager]) -> None:
-    graph = DiGraph(
-        nodes={
-            "A": DiGraphNode(
-                name="A", edges=[DiGraphEdge(target="B", condition="yes"), DiGraphEdge(target="C", condition="no")]
-            ),
-            "B": DiGraphNode(name="B", edges=[]),
-            "C": DiGraphNode(name="C", edges=[]),
-        }
-    )
-    message_thread = [TextChatMessage(source="A", content="no", metadata={})]
-    manager = digraph_manager(graph=graph, active_nodes={"A"}, thread=message_thread, pending={"B": [], "C": []})
-
-    result = await manager.select_speakers(manager._message_thread)  # pyright: ignore[reportPrivateUsage]
-    assert result == ["C"]
-    assert "C" in manager._active_nodes  # pyright: ignore[reportPrivateUsage]
-
-
-@pytest.mark.asyncio
-async def test_select_speakers_from_start_nodes(digraph_manager: Callable[..., GraphFlowManager]) -> None:
-    graph = DiGraph(
-        nodes={
-            "A": DiGraphNode(name="A", edges=[]),
-            "B": DiGraphNode(name="B", edges=[]),
-        }
-    )
-    # No prior message — both are start nodes
-    manager = digraph_manager(graph=graph, active_nodes=set(), thread=[], pending={"A": [], "B": []})
-    result = await manager.select_speakers([])
-    assert set(result) == {"A", "B"}
-
-
-@pytest.mark.asyncio
-async def test_select_speakers_termination(digraph_manager: Callable[..., GraphFlowManager]) -> None:
-    graph = DiGraph(
-        nodes={
-            "A": DiGraphNode(name="A", edges=[]),
-        }
-    )
-
-    # Create the manager and manually patch _signal_termination to track calls
-    manager = digraph_manager(
-        graph=graph, active_nodes={"A"}, thread=[TextChatMessage(source="A", content="done", metadata={})], pending={}
-    )
-    manager._signal_termination = AsyncMock()  # type: ignore[assignment]
-
-    result = await manager.select_speakers(manager._message_thread)  # pyright: ignore[reportPrivateUsage]
-
-    # No speakers left to run, so result should be empty
-    assert result == [_DIGRAPH_STOP_AGENT_NAME]
-
-
-@pytest.mark.asyncio
-async def test_select_speakers_conditional_all_activation(
-    digraph_manager: Callable[..., GraphFlowManager],
-) -> None:
-    graph = DiGraph(
-        nodes={
-            "A": DiGraphNode(
-                name="A", edges=[DiGraphEdge(target="B", condition="yes"), DiGraphEdge(target="C", condition="no")]
-            ),
-            "B": DiGraphNode(name="B", edges=[], activation="all"),
-            "C": DiGraphNode(name="C", edges=[], activation="all"),
-        }
-    )
-    message_thread = [TextChatMessage(source="A", content="no", metadata={})]
-    manager = digraph_manager(graph=graph, active_nodes={"A"}, thread=message_thread, pending={"B": [], "C": []})
-
-    result = await manager.select_speakers(manager._message_thread)  # pyright: ignore[reportPrivateUsage]
-    assert result == ["C"]
-    assert "C" in manager._active_nodes  # pyright: ignore[reportPrivateUsage]
-
-
-@pytest.mark.asyncio
-async def test_select_speakers_conditional_any_activation(
-    digraph_manager: Callable[..., GraphFlowManager],
-) -> None:
-    graph = DiGraph(
-        nodes={
-            "A": DiGraphNode(
-                name="A", edges=[DiGraphEdge(target="B", condition="yes"), DiGraphEdge(target="C", condition="no")]
-            ),
-            "B": DiGraphNode(name="B", edges=[], activation="any"),
-            "C": DiGraphNode(name="C", edges=[], activation="any"),
-        }
-    )
-    message_thread = [TextChatMessage(source="A", content="yes", metadata={})]
-    manager = digraph_manager(graph=graph, active_nodes={"A"}, thread=message_thread, pending={"B": [], "C": []})
-
-    result = await manager.select_speakers(manager._message_thread)  # pyright: ignore[reportPrivateUsage]
-    assert result == ["B"]
-    assert "B" in manager._active_nodes  # pyright: ignore[reportPrivateUsage]
-
-
 class _EchoAgent(BaseChatAgent):
     def __init__(self, name: str, description: str) -> None:
         super().__init__(name, description)
@@ -1417,7 +1197,6 @@ async def test_graph_flow_serialize_deserialize() -> None:
         participants=builder.get_participants(),
         graph=builder.build(),
         runtime=None,
-        termination_condition=MaxMessageTermination(5),
     )
 
     serialized = team.dump_component()
@@ -1439,11 +1218,52 @@ async def test_graph_flow_serialize_deserialize() -> None:
     assert results.messages[1].source == "A"
     assert results.messages[1].content == "0"
     assert isinstance(results.messages[2], TextMessage)
-    assert results.messages[2].source == "A"
-    assert results.messages[2].content == "1"
-    assert isinstance(results.messages[3], TextMessage)
-    assert results.messages[3].source == "B"
-    assert results.messages[3].content == "0"
+    assert results.messages[2].source == "B"
+    assert results.messages[2].content == "0"
     assert isinstance(results.messages[-1], StopMessage)
     assert results.messages[-1].source == _DIGRAPH_STOP_AGENT_NAME
     assert results.messages[-1].content == "Digraph execution is complete"
+
+
+@pytest.mark.asyncio
+async def test_graph_flow_stateful_pause_and_resume_with_termination() -> None:
+    client_a = ReplayChatCompletionClient(["A1", "A2"])
+    client_b = ReplayChatCompletionClient(["B1"])
+
+    a = AssistantAgent("A", model_client=client_a)
+    b = AssistantAgent("B", model_client=client_b)
+
+    builder = DiGraphBuilder()
+    builder.add_node(a).add_node(b)
+    builder.add_edge(a, b)
+    builder.set_entry_point(a)
+
+    team = GraphFlow(
+        participants=builder.get_participants(),
+        graph=builder.build(),
+        runtime=None,
+        termination_condition=SourceMatchTermination(sources=["A"]),
+    )
+
+    result = await team.run(task="Start")
+    assert len(result.messages) == 2
+    assert result.messages[0].source == "user"
+    assert result.messages[1].source == "A"
+    assert result.stop_reason is not None and result.stop_reason == "'A' answered"
+
+    # Export state.
+    state = await team.save_state()
+
+    # Load state into a new team.
+    new_team = GraphFlow(
+        participants=builder.get_participants(),
+        graph=builder.build(),
+        runtime=None,
+    )
+    await new_team.load_state(state)
+
+    # Resume.
+    result = await new_team.run()
+    assert len(result.messages) == 2
+    assert result.messages[0].source == "B"
+    assert result.messages[1].source == _DIGRAPH_STOP_AGENT_NAME

python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/tutorial/models.ipynb

@@ -427,6 +427,58 @@
     "```"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Llama API (experimental)\n",
+    "\n",
+    "[Llama API](https://llama.developer.meta.com?utm_source=partner-autogen&utm_medium=readme) is the Meta's first party API offering. It currently offers an [OpenAI compatible endpoint](https://llama.developer.meta.com/docs/features/compatibility).\n",
+    "So you can use the {py:class}`~autogen_ext.models.openai.OpenAIChatCompletionClient` with the Llama API.\n",
+    "\n",
+    "This endpoint fully supports the following OpenAI client library features:\n",
+    "* Chat completions\n",
+    "* Model selection\n",
+    "* Temperature/sampling\n",
+    "* Streaming\n",
+    "* Image understanding\n",
+    "* Structured output (JSON mode)\n",
+    "* Funciton calling (tools)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from autogen_core.models import UserMessage\n",
+    "from autogen_ext.models.openai import OpenAIChatCompletionClient\n",
+    "from autogen_core import Image\n",
+    "from pathlib import Path\n",
+    "\n",
+    "# Text\n",
+    "model_client = OpenAIChatCompletionClient(\n",
+    "    model=\"Llama-4-Scout-17B-16E-Instruct-FP8\",\n",
+    "    # api_key=\"LLAMA_API_KEY\"\n",
+    ")\n",
+    "\n",
+    "response = await model_client.create([UserMessage(content=\"Write me a poem\", source=\"user\")])\n",
+    "print(response)\n",
+    "await model_client.close()\n",
+    "\n",
+    "# Image\n",
+    "model_client = OpenAIChatCompletionClient(\n",
+    "    model=\"Llama-4-Maverick-17B-128E-Instruct-FP8\",\n",
+    "    # api_key=\"LLAMA_API_KEY\"\n",
+    ")\n",
+    "image = Image.from_file(Path(\"test.png\"))\n",
+    "\n",
+    "response = await model_client.create([UserMessage(content=[\"What is in this image\", image], source=\"user\")])\n",
+    "print(response)\n",
+    "await model_client.close()"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},

python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/tutorial/state.ipynb

@@ -112,7 +112,7 @@
             "source": [
                 "```{note}\n",
                 "For {py:class}`~autogen_agentchat.agents.AssistantAgent`, its state consists of the model_context.\n",
-                "If your write your own custom agent, consider overriding the {py:meth}`~autogen_agentchat.agents.BaseChatAgent.save_state` and {py:meth}`~autogen_agentchat.agents.BaseChatAgent.load_state` methods to customize the behavior. The default implementations save and load an empty state.\n",
+                "If you write your own custom agent, consider overriding the {py:meth}`~autogen_agentchat.agents.BaseChatAgent.save_state` and {py:meth}`~autogen_agentchat.agents.BaseChatAgent.load_state` methods to customize the behavior. The default implementations save and load an empty state.\n",
                 "```"
             ]
         },

python/packages/autogen-core/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "autogen-core"
-version = "0.5.6"
+version = "0.5.7"
 license = {file = "LICENSE-CODE"}
 description = "Foundational interfaces and agent runtime implementation for AutoGen"
 readme = "README.md"
@@ -69,7 +69,7 @@ dev = [
     "pygments",
     "sphinxext-rediraffe",
 
-    "autogen_ext==0.5.6",
+    "autogen_ext==0.5.7",
 
     # Documentation tooling
     "diskcache",

python/packages/autogen-core/src/autogen_core/_agent.py

@@ -1,9 +1,13 @@
-from typing import Any, Mapping, Protocol, runtime_checkable
+from typing import TYPE_CHECKING, Any, Mapping, Protocol, runtime_checkable
 
 from ._agent_id import AgentId
 from ._agent_metadata import AgentMetadata
 from ._message_context import MessageContext
 
+# Forward declaration for type checking only
+if TYPE_CHECKING:
+    from ._agent_runtime import AgentRuntime
+
 
 @runtime_checkable
 class Agent(Protocol):
@@ -17,6 +21,15 @@ class Agent(Protocol):
         """ID of the agent."""
         ...
 
+    async def bind_id_and_runtime(self, id: AgentId, runtime: "AgentRuntime") -> None:
+        """Function used to bind an Agent instance to an `AgentRuntime`.
+
+        Args:
+            agent_id (AgentId): ID of the agent.
+            runtime (AgentRuntime): AgentRuntime instance to bind the agent to.
+        """
+        ...
+
     async def on_message(self, message: Any, ctx: MessageContext) -> Any:
         """Message handler for the agent. This should only be called by the runtime, not by other agents.
 

python/packages/autogen-core/src/autogen_core/_agent_instantiation.py

@@ -118,3 +118,9 @@ class AgentInstantiationContext:
             raise RuntimeError(
                 "AgentInstantiationContext.agent_id() must be called within an instantiation context such as when the AgentRuntime is instantiating an agent. Mostly likely this was caused by directly instantiating an agent instead of using the AgentRuntime to do so."
             ) from e
+
+    @classmethod
+    def is_in_factory_call(cls) -> bool:
+        if cls._AGENT_INSTANTIATION_CONTEXT_VAR.get(None) is None:
+            return False
+        return True

python/packages/autogen-core/src/autogen_core/_agent_runtime.py

@@ -130,6 +130,60 @@ class AgentRuntime(Protocol):
         """
         ...
 
+    async def register_agent_instance(
+        self,
+        agent_instance: Agent,
+        agent_id: AgentId,
+    ) -> AgentId:
+        """Register an agent instance with the runtime. The type may be reused, but each agent_id must be unique. All agent instances within a type must be of the same object type. This API does not add any subscriptions.
+
+        .. note::
+
+            This is a low level API and usually the agent class's `register_instance` method should be used instead, as this also handles subscriptions automatically.
+
+        Example:
+
+        .. code-block:: python
+
+            from dataclasses import dataclass
+
+            from autogen_core import AgentId, AgentRuntime, MessageContext, RoutedAgent, event
+            from autogen_core.models import UserMessage
+
+
+            @dataclass
+            class MyMessage:
+                content: str
+
+
+            class MyAgent(RoutedAgent):
+                def __init__(self) -> None:
+                    super().__init__("My core agent")
+
+                @event
+                async def handler(self, message: UserMessage, context: MessageContext) -> None:
+                    print("Event received: ", message.content)
+
+
+            async def main() -> None:
+                runtime: AgentRuntime = ...  # type: ignore
+                agent = MyAgent()
+                await runtime.register_agent_instance(
+                    agent_instance=agent, agent_id=AgentId(type="my_agent", key="default")
+                )
+
+
+            import asyncio
+
+            asyncio.run(main())
+
+
+        Args:
+            agent_instance (Agent): A concrete instance of the agent.
+            agent_id (AgentId): The agent's identifier. The agent's type is `agent_id.type`.
+        """
+        ...
+
     # TODO: uncomment out the following type ignore when this is fixed in mypy: https://github.com/python/mypy/issues/3737
     async def try_get_underlying_agent_instance(self, id: AgentId, type: Type[T] = Agent) -> T:  # type: ignore[assignment]
         """Try to get the underlying agent instance by name and namespace. This is generally discouraged (hence the long name), but can be useful in some cases.

python/packages/autogen-core/src/autogen_core/_base_agent.py

@@ -21,6 +21,7 @@ from ._subscription import Subscription, UnboundSubscription
 from ._subscription_context import SubscriptionInstantiationContext
 from ._topic import TopicId
 from ._type_prefix_subscription import TypePrefixSubscription
+from ._type_subscription import TypeSubscription
 
 T = TypeVar("T", bound=Agent)
 
@@ -82,20 +83,25 @@ class BaseAgent(ABC, Agent):
         return AgentMetadata(key=self._id.key, type=self._id.type, description=self._description)
 
     def __init__(self, description: str) -> None:
-        try:
-            runtime = AgentInstantiationContext.current_runtime()
-            id = AgentInstantiationContext.current_agent_id()
-        except LookupError as e:
-            raise RuntimeError(
-                "BaseAgent must be instantiated within the context of an AgentRuntime. It cannot be directly instantiated."
-            ) from e
-
-        self._runtime: AgentRuntime = runtime
-        self._id: AgentId = id
+        if AgentInstantiationContext.is_in_factory_call():
+            self._runtime: AgentRuntime = AgentInstantiationContext.current_runtime()
+            self._id = AgentInstantiationContext.current_agent_id()
         if not isinstance(description, str):
             raise ValueError("Agent description must be a string")
         self._description = description
 
+    async def bind_id_and_runtime(self, id: AgentId, runtime: AgentRuntime) -> None:
+        if hasattr(self, "_id"):
+            if self._id != id:
+                raise RuntimeError("Agent is already bound to a different ID")
+
+        if hasattr(self, "_runtime"):
+            if self._runtime != runtime:
+                raise RuntimeError("Agent is already bound to a different runtime")
+
+        self._id = id
+        self._runtime = runtime
+
     @property
     def type(self) -> str:
         return self.id.type
@@ -155,6 +161,56 @@ class BaseAgent(ABC, Agent):
     async def close(self) -> None:
         pass
 
+    async def register_instance(
+        self,
+        runtime: AgentRuntime,
+        agent_id: AgentId,
+        *,
+        skip_class_subscriptions: bool = True,
+        skip_direct_message_subscription: bool = False,
+    ) -> AgentId:
+        """
+        This function is similar to `register` but is used for registering an instance of an agent. A subscription based on the agent ID is created and added to the runtime.
+        """
+        agent_id = await runtime.register_agent_instance(agent_instance=self, agent_id=agent_id)
+
+        id_subscription = TypeSubscription(topic_type=agent_id.key, agent_type=agent_id.type)
+        await runtime.add_subscription(id_subscription)
+
+        if not skip_class_subscriptions:
+            with SubscriptionInstantiationContext.populate_context(AgentType(agent_id.type)):
+                subscriptions: List[Subscription] = []
+                for unbound_subscription in self._unbound_subscriptions():
+                    subscriptions_list_result = unbound_subscription()
+                    if inspect.isawaitable(subscriptions_list_result):
+                        subscriptions_list = await subscriptions_list_result
+                    else:
+                        subscriptions_list = subscriptions_list_result
+
+                    subscriptions.extend(subscriptions_list)
+            for subscription in subscriptions:
+                await runtime.add_subscription(subscription)
+
+        if not skip_direct_message_subscription:
+            # Additionally adds a special prefix subscription for this agent to receive direct messages
+            try:
+                await runtime.add_subscription(
+                    TypePrefixSubscription(
+                        # The prefix MUST include ":" to avoid collisions with other agents
+                        topic_type_prefix=agent_id.type + ":",
+                        agent_type=agent_id.type,
+                    )
+                )
+            except ValueError:
+                # We don't care if the subscription already exists
+                pass
+
+        # TODO: deduplication
+        for _message_type, serializer in self._handles_types():
+            runtime.add_message_serializer(serializer)
+
+        return agent_id
+
     @classmethod
     async def register(
         cls,

python/packages/autogen-core/src/autogen_core/_single_threaded_agent_runtime.py

@@ -2,6 +2,7 @@ from __future__ import annotations
 
 import asyncio
 import inspect
+import json
 import logging
 import sys
 import uuid
@@ -265,6 +266,7 @@ class SingleThreadedAgentRuntime(AgentRuntime):
         self._serialization_registry = SerializationRegistry()
         self._ignore_unhandled_handler_exceptions = ignore_unhandled_exceptions
         self._background_exception: BaseException | None = None
+        self._agent_instance_types: Dict[str, Type[Agent]] = {}
 
     @property
     def unprocessed_messages_count(
@@ -276,6 +278,55 @@ class SingleThreadedAgentRuntime(AgentRuntime):
     def _known_agent_names(self) -> Set[str]:
         return set(self._agent_factories.keys())
 
+    async def _create_otel_attributes(
+        self,
+        sender_agent_id: AgentId | None = None,
+        recipient_agent_id: AgentId | None = None,
+        message_context: MessageContext | None = None,
+        message: Any = None,
+    ) -> Mapping[str, str]:
+        """Create OpenTelemetry attributes for the given agent and message.
+
+        Args:
+            sender_agent (Agent, optional): The sender agent instance.
+            recipient_agent (Agent, optional): The recipient agent instance.
+            message (Any): The message instance.
+
+        Returns:
+            Attributes: A dictionary of OpenTelemetry attributes.
+        """
+        if not sender_agent_id and not recipient_agent_id and not message:
+            return {}
+        attributes: Dict[str, str] = {}
+        if sender_agent_id:
+            sender_agent = await self._get_agent(sender_agent_id)
+            attributes["sender_agent_type"] = sender_agent.id.type
+            attributes["sender_agent_class"] = sender_agent.__class__.__name__
+        if recipient_agent_id:
+            recipient_agent = await self._get_agent(recipient_agent_id)
+            attributes["recipient_agent_type"] = recipient_agent.id.type
+            attributes["recipient_agent_class"] = recipient_agent.__class__.__name__
+
+        if message_context:
+            serialized_message_context = {
+                "sender": str(message_context.sender),
+                "topic_id": str(message_context.topic_id),
+                "is_rpc": message_context.is_rpc,
+                "message_id": message_context.message_id,
+            }
+            attributes["message_context"] = json.dumps(serialized_message_context)
+
+        if message:
+            try:
+                serialized_message = self._try_serialize(message)
+            except Exception as e:
+                serialized_message = str(e)
+        else:
+            serialized_message = "No Message"
+        attributes["message"] = serialized_message
+
+        return attributes
+
     # Returns the response of the message
     async def send_message(
         self,
@@ -311,6 +362,7 @@ class SingleThreadedAgentRuntime(AgentRuntime):
             future = asyncio.get_event_loop().create_future()
             if recipient.type not in self._known_agent_names:
                 future.set_exception(Exception("Recipient not found"))
+                return await future
 
             content = message.__dict__ if hasattr(message, "__dict__") else message
             logger.info(f"Sending message of type {type(message).__name__} to {recipient.type}: {content}")
@@ -440,7 +492,17 @@ class SingleThreadedAgentRuntime(AgentRuntime):
                     cancellation_token=message_envelope.cancellation_token,
                     message_id=message_envelope.message_id,
                 )
-                with self._tracer_helper.trace_block("process", recipient_agent.id, parent=message_envelope.metadata):
+                with self._tracer_helper.trace_block(
+                    "process",
+                    recipient_agent.id,
+                    parent=message_envelope.metadata,
+                    attributes=await self._create_otel_attributes(
+                        sender_agent_id=message_envelope.sender,
+                        recipient_agent_id=recipient,
+                        message_context=message_context,
+                        message=message_envelope.message,
+                    ),
+                ):
                     with MessageHandlerContext.populate_context(recipient_agent.id):
                         response = await recipient_agent.on_message(
                             message_envelope.message,
@@ -527,7 +589,17 @@ class SingleThreadedAgentRuntime(AgentRuntime):
                     agent = await self._get_agent(agent_id)
 
                     async def _on_message(agent: Agent, message_context: MessageContext) -> Any:
-                        with self._tracer_helper.trace_block("process", agent.id, parent=message_envelope.metadata):
+                        with self._tracer_helper.trace_block(
+                            "process",
+                            agent.id,
+                            parent=message_envelope.metadata,
+                            attributes=await self._create_otel_attributes(
+                                sender_agent_id=message_envelope.sender,
+                                recipient_agent_id=agent.id,
+                                message_context=message_context,
+                                message=message_envelope.message,
+                            ),
+                        ):
                             with MessageHandlerContext.populate_context(agent.id):
                                 try:
                                     return await agent.on_message(
@@ -557,7 +629,16 @@ class SingleThreadedAgentRuntime(AgentRuntime):
             # TODO if responses are given for a publish
 
     async def _process_response(self, message_envelope: ResponseMessageEnvelope) -> None:
-        with self._tracer_helper.trace_block("ack", message_envelope.recipient, parent=message_envelope.metadata):
+        with self._tracer_helper.trace_block(
+            "ack",
+            message_envelope.recipient,
+            parent=message_envelope.metadata,
+            attributes=await self._create_otel_attributes(
+                sender_agent_id=message_envelope.sender,
+                recipient_agent_id=message_envelope.recipient,
+                message=message_envelope.message,
+            ),
+        ):
             content = (
                 message_envelope.message.__dict__
                 if hasattr(message_envelope.message, "__dict__")
@@ -830,6 +911,32 @@ class SingleThreadedAgentRuntime(AgentRuntime):
 
         return type
 
+    async def register_agent_instance(
+        self,
+        agent_instance: Agent,
+        agent_id: AgentId,
+    ) -> AgentId:
+        def agent_factory() -> Agent:
+            raise RuntimeError(
+                "Agent factory was invoked for an agent instance that was not registered. This is likely due to the agent type being incorrectly subscribed to a topic. If this exception occurs when publishing a message to the DefaultTopicId, then it is likely that `skip_class_subscriptions` needs to be turned off when registering the agent."
+            )
+
+        if agent_id in self._instantiated_agents:
+            raise ValueError(f"Agent with id {agent_id} already exists.")
+
+        if agent_id.type not in self._agent_factories:
+            self._agent_factories[agent_id.type] = agent_factory
+            self._agent_instance_types[agent_id.type] = type_func_alias(agent_instance)
+        else:
+            if self._agent_factories[agent_id.type].__code__ != agent_factory.__code__:
+                raise ValueError("Agent factories and agent instances cannot be registered to the same type.")
+            if self._agent_instance_types[agent_id.type] != type_func_alias(agent_instance):
+                raise ValueError("Agent instances must be the same object type.")
+
+        await agent_instance.bind_id_and_runtime(id=agent_id, runtime=self)
+        self._instantiated_agents[agent_id] = agent_instance
+        return agent_id
+
     async def _invoke_agent_factory(
         self,
         agent_factory: Callable[[], T | Awaitable[T]] | Callable[[AgentRuntime, AgentId], T | Awaitable[T]],
@@ -851,8 +958,7 @@ class SingleThreadedAgentRuntime(AgentRuntime):
                     raise ValueError("Agent factory must take 0 or 2 arguments.")
 
                 if inspect.isawaitable(agent):
-                    return cast(T, await agent)
-
+                    agent = cast(T, await agent)
                 return agent
 
             except BaseException as e:

python/packages/autogen-core/src/autogen_core/models/_model_client.py

@@ -37,9 +37,19 @@ class ModelFamily:
     CLAUDE_3_5_HAIKU = "claude-3-5-haiku"
     CLAUDE_3_5_SONNET = "claude-3-5-sonnet"
     CLAUDE_3_7_SONNET = "claude-3-7-sonnet"
+    LLAMA_3_3_8B = "llama-3.3-8b"
+    LLAMA_3_3_70B = "llama-3.3-70b"
+    LLAMA_4_SCOUT = "llama-4-scout"
+    LLAMA_4_MAVERICK = "llama-4-maverick"
+    CODESRAL = "codestral"
+    OPEN_CODESRAL_MAMBA = "open-codestral-mamba"
+    MISTRAL = "mistral"
+    MINISTRAL = "ministral"
+    PIXTRAL = "pixtral"
     UNKNOWN = "unknown"
 
     ANY: TypeAlias = Literal[
+        # openai_models
         "gpt-41",
         "gpt-45",
         "gpt-4o",
@@ -49,16 +59,30 @@ class ModelFamily:
         "gpt-4",
         "gpt-35",
         "r1",
+        # google_models
         "gemini-1.5-flash",
         "gemini-1.5-pro",
         "gemini-2.0-flash",
         "gemini-2.5-pro",
+        # anthropic_models
         "claude-3-haiku",
         "claude-3-sonnet",
         "claude-3-opus",
         "claude-3-5-haiku",
         "claude-3-5-sonnet",
         "claude-3-7-sonnet",
+        # llama_models
+        "llama-3.3-8b",
+        "llama-3.3-70b",
+        "llama-4-scout",
+        "llama-4-maverick",
+        # mistral_models
+        "codestral",
+        "open-codestral-mamba",
+        "mistral",
+        "ministral",
+        "pixtral",
+        # unknown
         "unknown",
     ]
 
@@ -98,6 +122,25 @@ class ModelFamily:
             ModelFamily.GPT_35,
         )
 
+    @staticmethod
+    def is_llama(family: str) -> bool:
+        return family in (
+            ModelFamily.LLAMA_3_3_8B,
+            ModelFamily.LLAMA_3_3_70B,
+            ModelFamily.LLAMA_4_SCOUT,
+            ModelFamily.LLAMA_4_MAVERICK,
+        )
+
+    @staticmethod
+    def is_mistral(family: str) -> bool:
+        return family in (
+            ModelFamily.CODESRAL,
+            ModelFamily.OPEN_CODESRAL_MAMBA,
+            ModelFamily.MISTRAL,
+            ModelFamily.MINISTRAL,
+            ModelFamily.PIXTRAL,
+        )
+
 
 @deprecated("Use the ModelInfo class instead ModelCapabilities.")
 class ModelCapabilities(TypedDict, total=False):

python/packages/autogen-core/src/autogen_core/tools/_static_workbench.py

@@ -1,4 +1,5 @@
 import asyncio
+import builtins
 from typing import Any, Dict, List, Literal, Mapping
 
 from pydantic import BaseModel
@@ -55,12 +56,12 @@ class StaticWorkbench(Workbench, Component[StaticWorkbenchConfig]):
         try:
             result_future = asyncio.ensure_future(tool.run_json(arguments, cancellation_token))
             cancellation_token.link_future(result_future)
-            result = await result_future
+            actual_tool_output = await result_future
             is_error = False
+            result_str = tool.return_value_as_string(actual_tool_output)
         except Exception as e:
-            result = str(e)
+            result_str = self._format_errors(e)
             is_error = True
-        result_str = tool.return_value_as_string(result)
         return ToolResult(name=tool.name, result=[TextResultContent(content=result_str)], is_error=is_error)
 
     async def start(self) -> None:
@@ -90,3 +91,16 @@ class StaticWorkbench(Workbench, Component[StaticWorkbenchConfig]):
     @classmethod
     def _from_config(cls, config: StaticWorkbenchConfig) -> Self:
         return cls(tools=[BaseTool.load_component(tool) for tool in config.tools])
+
+    def _format_errors(self, error: Exception) -> str:
+        """Recursively format errors into a string."""
+
+        error_message = ""
+        if hasattr(builtins, "ExceptionGroup") and isinstance(error, builtins.ExceptionGroup):
+            # ExceptionGroup is available in Python 3.11+.
+            # TODO: how to make this compatible with Python 3.10?
+            for sub_exception in error.exceptions:  # type: ignore
+                error_message += self._format_errors(sub_exception)  # type: ignore
+        else:
+            error_message += f"{str(error)}\n"
+        return error_message.strip()

python/packages/autogen-core/tests/test_base_agent.py

@@ -9,7 +9,7 @@ async def test_base_agent_create(mocker: MockerFixture) -> None:
     runtime = mocker.Mock(spec=AgentRuntime)
 
     # Shows how to set the context for the agent instantiation in a test context
-    with AgentInstantiationContext.populate_context((runtime, AgentId("name", "namespace"))):
-        agent = NoopAgent()
-        assert agent.runtime == runtime
-        assert agent.id == AgentId("name", "namespace")
+    with AgentInstantiationContext.populate_context((runtime, AgentId("name2", "namespace2"))):
+        agent2 = NoopAgent()
+        assert agent2.runtime == runtime
+        assert agent2.id == AgentId("name2", "namespace2")

python/packages/autogen-core/tests/test_runtime.py

@@ -82,6 +82,60 @@ async def test_agent_type_must_be_unique() -> None:
     await runtime.register_factory(type=AgentType("name2"), agent_factory=agent_factory, expected_class=NoopAgent)
 
 
+@pytest.mark.asyncio
+async def test_agent_type_register_instance() -> None:
+    runtime = SingleThreadedAgentRuntime()
+    agent1_id = AgentId(type="name", key="default")
+    agent2_id = AgentId(type="name", key="notdefault")
+    agent1 = NoopAgent()
+    agent1_dup = NoopAgent()
+    agent2 = NoopAgent()
+    await agent1.register_instance(runtime=runtime, agent_id=agent1_id)
+    await agent2.register_instance(runtime=runtime, agent_id=agent2_id)
+
+    assert await runtime.try_get_underlying_agent_instance(agent1_id, type=NoopAgent) == agent1
+    assert await runtime.try_get_underlying_agent_instance(agent2_id, type=NoopAgent) == agent2
+    with pytest.raises(ValueError):
+        await agent1_dup.register_instance(runtime=runtime, agent_id=agent1_id)
+
+
+@pytest.mark.asyncio
+async def test_agent_type_register_instance_different_types() -> None:
+    runtime = SingleThreadedAgentRuntime()
+    agent_id1 = AgentId(type="name", key="noop")
+    agent_id2 = AgentId(type="name", key="loopback")
+    agent1 = NoopAgent()
+    agent2 = LoopbackAgent()
+    await agent1.register_instance(runtime=runtime, agent_id=agent_id1)
+    with pytest.raises(ValueError):
+        await agent2.register_instance(runtime=runtime, agent_id=agent_id2)
+
+
+@pytest.mark.asyncio
+async def test_agent_type_register_instance_publish_new_source() -> None:
+    runtime = SingleThreadedAgentRuntime(ignore_unhandled_exceptions=False)
+    agent_id = AgentId(type="name", key="default")
+    agent1 = LoopbackAgent()
+    await agent1.register_instance(runtime=runtime, agent_id=agent_id)
+    await runtime.add_subscription(TypeSubscription("notdefault", "name"))
+
+    runtime.start()
+    with pytest.raises(RuntimeError):
+        await runtime.publish_message(MessageType(), TopicId("notdefault", "notdefault"))
+        await runtime.stop_when_idle()
+    await runtime.close()
+
+
+@pytest.mark.asyncio
+async def test_register_instance_factory() -> None:
+    runtime = SingleThreadedAgentRuntime()
+    agent1_id = AgentId(type="name", key="default")
+    agent1 = NoopAgent()
+    await agent1.register_instance(runtime=runtime, agent_id=agent1_id)
+    with pytest.raises(ValueError):
+        await NoopAgent.register(runtime, "name", lambda: NoopAgent())
+
+
 @pytest.mark.asyncio
 async def test_register_receives_publish(tracer_provider: TracerProvider) -> None:
     runtime = SingleThreadedAgentRuntime(tracer_provider=tracer_provider)

python/packages/autogen-ext/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "autogen-ext"
-version = "0.5.6"
+version = "0.5.7"
 license = {file = "LICENSE-CODE"}
 description = "AutoGen extensions library"
 readme = "README.md"
@@ -15,7 +15,7 @@ classifiers = [
     "Operating System :: OS Independent",
 ]
 dependencies = [
-    "autogen-core==0.5.6",
+    "autogen-core==0.5.7",
 ]
 
 [project.optional-dependencies]
@@ -32,7 +32,7 @@ docker = ["docker~=7.0", "asyncio_atexit>=1.0.1"]
 ollama = ["ollama>=0.4.7", "tiktoken>=0.8.0"]
 openai = ["openai>=1.66.5", "tiktoken>=0.8.0", "aiofiles"]
 file-surfer = [
-    "autogen-agentchat==0.5.6",
+    "autogen-agentchat==0.5.7",
     "magika>=0.6.1rc2",
     "markitdown[all]~=0.1.0a3",
 ]
@@ -44,21 +44,21 @@ llama-cpp = [
 graphrag = ["graphrag>=1.0.1"]
 chromadb = ["chromadb>=1.0.0"]
 web-surfer = [
-    "autogen-agentchat==0.5.6",
+    "autogen-agentchat==0.5.7",
     "playwright>=1.48.0",
     "pillow>=11.0.0",
     "magika>=0.6.1rc2",
     "markitdown[all]~=0.1.0a3",
 ]
 magentic-one = [
-    "autogen-agentchat==0.5.6",
+    "autogen-agentchat==0.5.7",
     "magika>=0.6.1rc2",
     "markitdown[all]~=0.1.0a3",
     "playwright>=1.48.0",
     "pillow>=11.0.0",
 ]
 video-surfer = [
-    "autogen-agentchat==0.5.6",
+    "autogen-agentchat==0.5.7",
     "opencv-python>=4.5",
     "ffmpeg-python",
     "openai-whisper",

python/packages/autogen-ext/src/autogen_ext/models/anthropic/__init__.py

@@ -1,4 +1,8 @@
-from ._anthropic_client import AnthropicChatCompletionClient, BaseAnthropicChatCompletionClient
+from ._anthropic_client import (
+    AnthropicBedrockChatCompletionClient,
+    AnthropicChatCompletionClient,
+    BaseAnthropicChatCompletionClient,
+)
 from .config import (
     AnthropicBedrockClientConfiguration,
     AnthropicBedrockClientConfigurationConfigModel,
@@ -10,6 +14,7 @@ from .config import (
 
 __all__ = [
     "AnthropicChatCompletionClient",
+    "AnthropicBedrockChatCompletionClient",
     "BaseAnthropicChatCompletionClient",
     "AnthropicClientConfiguration",
     "AnthropicBedrockClientConfiguration",

python/packages/autogen-ext/src/autogen_ext/models/anthropic/_anthropic_client.py

@@ -5,8 +5,6 @@ import json
 import logging
 import re
 import warnings
-
-# from asyncio import Task
 from typing import (
     Any,
     AsyncGenerator,
@@ -25,7 +23,7 @@ from typing import (
 )
 
 import tiktoken
-from anthropic import AnthropicBedrock, AsyncAnthropic, AsyncStream
+from anthropic import AsyncAnthropic, AsyncAnthropicBedrock, AsyncStream
 from anthropic.types import (
     Base64ImageSourceParam,
     ContentBlock,
@@ -1121,7 +1119,7 @@ class AnthropicBedrockChatCompletionClient(
     BaseAnthropicChatCompletionClient, Component[AnthropicBedrockClientConfigurationConfigModel]
 ):
     """
-    Chat completion client for Anthropic's Claude models.
+    Chat completion client for Anthropic's Claude models on AWS Bedrock.
 
     Args:
         model (str): The Claude model to use (e.g., "claude-3-sonnet-20240229", "claude-3-opus-20240229")
@@ -1145,28 +1143,24 @@ class AnthropicBedrockChatCompletionClient(
     .. code-block:: python
 
         import asyncio
-        from autogen_ext.models.anthropic import AnthropicBedrockChatCompletionClient
-        from autogen_core.models import UserMessage
+        from autogen_ext.models.anthropic import AnthropicBedrockChatCompletionClient, BedrockInfo
+        from autogen_core.models import UserMessage, ModelInfo
 
 
         async def main():
-            config = {
-                "model": "anthropic.claude-3-5-sonnet-20240620-v1:0",
-                "temperature": 0.1,
-                "model_info": {
-                    "vision": True,
-                    "function_calling": True,
-                    "json_output": True,
-                    "family": ModelFamily.CLAUDE_3_5_SONNET,
-                },
-                "bedrock_info": {
-                    "aws_access_key": "<aws_access_key>",
-                    "aws_secret_key": "<aws_secret_key>",
-                    "aws_session_token": "<aws_session_token>",
-                    "aws_region": "<aws_region>",
-                },
-            }
-            anthropic_client = AnthropicBedrockChatCompletionClient(**config)
+            anthropic_client = AnthropicBedrockChatCompletionClient(
+                model="anthropic.claude-3-5-sonnet-20240620-v1:0",
+                temperature=0.1,
+                model_info=ModelInfo(
+                    vision=False, function_calling=True, json_output=False, family="unknown", structured_output=True
+                ),
+                bedrock_info=BedrockInfo(
+                    aws_access_key="<aws_access_key>",
+                    aws_secret_key="<aws_secret_key>",
+                    aws_session_token="<aws_session_token>",
+                    aws_region="<aws_region>",
+                ),
+            )
 
             result = await anthropic_client.create([UserMessage(content="What is the capital of France?", source="user")])  # type: ignore
             print(result)
@@ -1178,11 +1172,11 @@ class AnthropicBedrockChatCompletionClient(
 
     component_type = "model"
     component_config_schema = AnthropicBedrockClientConfigurationConfigModel
-    component_provider_override = "autogen_ext.models.anthropic.AnthropicChatCompletionClient"
+    component_provider_override = "autogen_ext.models.anthropic.AnthropicBedrockChatCompletionClient"
 
     def __init__(self, **kwargs: Unpack[AnthropicBedrockClientConfiguration]):
         if "model" not in kwargs:
-            raise ValueError("model is required for AnthropicChatCompletionClient")
+            raise ValueError("model is required for  AnthropicBedrockChatCompletionClient")
 
         self._raw_config: Dict[str, Any] = dict(kwargs).copy()
         copied_args = dict(kwargs).copy()
@@ -1201,11 +1195,11 @@ class AnthropicBedrockChatCompletionClient(
 
         # Handle bedrock_info as secretestr
         aws_region = bedrock_info["aws_region"]
-        aws_access_key = bedrock_info["aws_access_key"].get_secret_value()
-        aws_secret_key = bedrock_info["aws_secret_key"].get_secret_value()
-        aws_session_token = bedrock_info["aws_session_token"].get_secret_value()
+        aws_access_key = bedrock_info["aws_access_key"]
+        aws_secret_key = bedrock_info["aws_secret_key"]
+        aws_session_token = bedrock_info["aws_session_token"]
 
-        client = AnthropicBedrock(
+        client = AsyncAnthropicBedrock(
             aws_access_key=aws_access_key,
             aws_secret_key=aws_secret_key,
             aws_session_token=aws_session_token,

python/packages/autogen-ext/src/autogen_ext/models/anthropic/config/__init__.py

@@ -26,11 +26,11 @@ class BedrockInfo(TypedDict):
 
     """
 
-    aws_access_key: Required[SecretStr]
+    aws_access_key: Required[str]
     """Access key for the aws account to gain bedrock model access"""
-    aws_secret_key: Required[SecretStr]
+    aws_secret_key: Required[str]
     """Access secret key for the aws account to gain bedrock model access"""
-    aws_session_token: Required[SecretStr]
+    aws_session_token: Required[str]
     """aws session token for the aws account to gain bedrock model access"""
     aws_region: Required[str]
     """aws region for the aws account to gain bedrock model access"""
@@ -83,5 +83,15 @@ class AnthropicClientConfigurationConfigModel(BaseAnthropicClientConfigurationCo
     tool_choice: Union[Literal["auto", "any", "none"], Dict[str, Any]] | None = None
 
 
+class BedrockInfoConfigModel(TypedDict):
+    aws_access_key: Required[SecretStr]
+    """Access key for the aws account to gain bedrock model access"""
+    aws_session_token: Required[SecretStr]
+    """aws session token for the aws account to gain bedrock model access"""
+    aws_region: Required[str]
+    """aws region for the aws account to gain bedrock model access"""
+    aws_secret_key: Required[SecretStr]
+
+
 class AnthropicBedrockClientConfigurationConfigModel(AnthropicClientConfigurationConfigModel):
-    bedrock_info: BedrockInfo | None = None
+    bedrock_info: BedrockInfoConfigModel | None = None

python/packages/autogen-ext/src/autogen_ext/models/ollama/_model_info.py

@@ -258,6 +258,13 @@ _MODEL_INFO: Dict[str, ModelInfo] = {
         "family": ModelFamily.UNKNOWN,
         "structured_output": True,
     },
+    "qwen3": {
+        "vision": False,
+        "function_calling": True,
+        "json_output": True,
+        "family": ModelFamily.UNKNOWN,
+        "structured_output": True,
+    },
     "snowflake-arctic-embed": {
         "vision": False,
         "function_calling": False,
@@ -351,6 +358,15 @@ _MODEL_TOKEN_LIMITS: Dict[str, int] = {
     "qwen2.5-coder:0.5b": 32768,
     "qwen2.5-coder:1.5b": 32768,
     "qwen2.5-coder:3b": 32768,
+    "qwen3": 40960,
+    "qwen3:0.6b": 40960,
+    "qwen3:1.7b": 40960,
+    "qwen3:4b": 40960,
+    "qwen3:8b": 40960,
+    "qwen3:14b": 40960,
+    "qwen3:30b": 40960,
+    "qwen3:32b": 40960,
+    "qwen3:235b": 40960,
     "snowflake-arctic-embed": 512,
     "starcoder2": 16384,
     "tinyllama": 2048,

python/packages/autogen-ext/src/autogen_ext/models/openai/_message_transform.py

@@ -275,7 +275,6 @@ base_system_message_transformers: List[Callable[[LLMMessage, Dict[str, Any]], Di
 
 base_user_transformer_funcs: List[Callable[[LLMMessage, Dict[str, Any]], Dict[str, Any]]] = [
     _assert_valid_name,
-    _set_name,
     _set_role("user"),
 ]
 
@@ -293,6 +292,7 @@ system_message_transformers: List[Callable[[LLMMessage, Dict[str, Any]], Dict[st
 single_user_transformer_funcs: List[Callable[[LLMMessage, Dict[str, Any]], Dict[str, Any]]] = (
     base_user_transformer_funcs
     + [
+        _set_name,
         _set_prepend_text_content,
     ]
 )
@@ -300,6 +300,7 @@ single_user_transformer_funcs: List[Callable[[LLMMessage, Dict[str, Any]], Dict[
 multimodal_user_transformer_funcs: List[Callable[[LLMMessage, Dict[str, Any]], Dict[str, Any]]] = (
     base_user_transformer_funcs
     + [
+        _set_name,
         _set_multimodal_content,
     ]
 )
@@ -334,6 +335,19 @@ thought_assistant_transformer_funcs_gemini: List[Callable[[LLMMessage, Dict[str,
 
 
 # === Specific message param functions ===
+single_user_transformer_funcs_mistral: List[Callable[[LLMMessage, Dict[str, Any]], Dict[str, Any]]] = (
+    base_user_transformer_funcs
+    + [
+        _set_prepend_text_content,
+    ]
+)
+
+multimodal_user_transformer_funcs_mistral: List[Callable[[LLMMessage, Dict[str, Any]], Dict[str, Any]]] = (
+    base_user_transformer_funcs
+    + [
+        _set_multimodal_content,
+    ]
+)
 
 
 # === Transformer maps ===
@@ -359,6 +373,8 @@ assistant_transformer_funcs: Dict[str, List[Callable[[LLMMessage, Dict[str, Any]
     "tools": tools_assistant_transformer_funcs,
     "thought": thought_assistant_transformer_funcs,
 }
+
+
 assistant_transformer_constructors: Dict[str, Callable[..., Any]] = {
     "text": ChatCompletionAssistantMessageParam,
     "tools": ChatCompletionAssistantMessageParam,
@@ -403,6 +419,12 @@ assistant_transformer_funcs_claude: Dict[str, List[Callable[[LLMMessage, Dict[st
 }
 
 
+user_transformer_funcs_mistral: Dict[str, List[Callable[[LLMMessage, Dict[str, Any]], Dict[str, Any]]]] = {
+    "text": single_user_transformer_funcs_mistral,
+    "multimodal": multimodal_user_transformer_funcs_mistral,
+}
+
+
 def function_execution_result_message(message: LLMMessage, context: Dict[str, Any]) -> TrasformerReturnType:
     assert isinstance(message, FunctionExecutionResultMessage)
     return [
@@ -466,6 +488,24 @@ __CLAUDE_TRANSFORMER_MAP: TransformerMap = {
     FunctionExecutionResultMessage: function_execution_result_message,
 }
 
+__MISTRAL_TRANSFORMER_MAP: TransformerMap = {
+    SystemMessage: build_transformer_func(
+        funcs=system_message_transformers + [_set_empty_to_whitespace],
+        message_param_func=ChatCompletionSystemMessageParam,
+    ),
+    UserMessage: build_conditional_transformer_func(
+        funcs_map=user_transformer_funcs_mistral,
+        message_param_func_map=user_transformer_constructors,
+        condition_func=user_condition,
+    ),
+    AssistantMessage: build_conditional_transformer_func(
+        funcs_map=assistant_transformer_funcs,
+        message_param_func_map=assistant_transformer_constructors,
+        condition_func=assistant_condition,
+    ),
+    FunctionExecutionResultMessage: function_execution_result_message,
+}
+
 
 # set openai models to use the transformer map
 total_models = get_args(ModelFamily.ANY)
@@ -475,7 +515,16 @@ __claude_models = [model for model in total_models if ModelFamily.is_claude(mode
 
 __gemini_models = [model for model in total_models if ModelFamily.is_gemini(model)]
 
-__unknown_models = list(set(total_models) - set(__openai_models) - set(__claude_models) - set(__gemini_models))
+__llama_models = [model for model in total_models if ModelFamily.is_llama(model)]
+
+__unknown_models = list(
+    set(total_models) - set(__openai_models) - set(__claude_models) - set(__gemini_models) - set(__llama_models)
+)
+__mistral_models = [model for model in total_models if ModelFamily.is_mistral(model)]
+
+__unknown_models = list(
+    set(total_models) - set(__openai_models) - set(__claude_models) - set(__gemini_models) - set(__mistral_models)
+)
 
 for model in __openai_models:
     register_transformer("openai", model, __BASE_TRANSFORMER_MAP)
@@ -486,6 +535,12 @@ for model in __claude_models:
 for model in __gemini_models:
     register_transformer("openai", model, __GEMINI_TRANSFORMER_MAP)
 
+for model in __llama_models:
+    register_transformer("openai", model, __BASE_TRANSFORMER_MAP)
+
+for model in __mistral_models:
+    register_transformer("openai", model, __MISTRAL_TRANSFORMER_MAP)
+
 for model in __unknown_models:
     register_transformer("openai", model, __BASE_TRANSFORMER_MAP)
 

python/packages/autogen-ext/src/autogen_ext/models/openai/_model_info.py

@@ -36,6 +36,11 @@ _MODEL_POINTERS = {
     "claude-3-5-haiku": "claude-3-5-haiku-20241022",
     "claude-3-5-sonnet": "claude-3-5-sonnet-20241022",
     "claude-3-7-sonnet": "claude-3-7-sonnet-20250219",
+    # Llama models
+    "llama-3.3-8b": "Llama-3.3-8B-Instruct",
+    "llama-3.3-70b": "Llama-3.3-70B-Instruct",
+    "llama-4-scout": "Llama-4-Scout-17B-16E-Instruct-FP8",
+    "llama-4-maverick": "Llama-4-Maverick-17B-128E-Instruct-FP8",
 }
 
 _MODEL_INFO: Dict[str, ModelInfo] = {
@@ -351,6 +356,38 @@ _MODEL_INFO: Dict[str, ModelInfo] = {
         "structured_output": False,
         "multiple_system_messages": True,
     },
+    "Llama-3.3-8B-Instruct": {
+        "vision": False,
+        "function_calling": True,
+        "json_output": True,
+        "family": ModelFamily.LLAMA_3_3_8B,
+        "structured_output": False,
+        "multiple_system_messages": True,
+    },
+    "Llama-3.3-70B-Instruct": {
+        "vision": False,
+        "function_calling": True,
+        "json_output": True,
+        "family": ModelFamily.LLAMA_3_3_70B,
+        "structured_output": False,
+        "multiple_system_messages": True,
+    },
+    "Llama-4-Scout-17B-16E-Instruct-FP8": {
+        "vision": True,
+        "function_calling": True,
+        "json_output": True,
+        "family": ModelFamily.LLAMA_4_SCOUT,
+        "structured_output": True,
+        "multiple_system_messages": True,
+    },
+    "Llama-4-Maverick-17B-128E-Instruct-FP8": {
+        "vision": True,
+        "function_calling": True,
+        "json_output": True,
+        "family": ModelFamily.LLAMA_4_MAVERICK,
+        "structured_output": True,
+        "multiple_system_messages": True,
+    },
 }
 
 _MODEL_TOKEN_LIMITS: Dict[str, int] = {
@@ -391,10 +428,15 @@ _MODEL_TOKEN_LIMITS: Dict[str, int] = {
     "claude-3-5-haiku-20241022": 50000,
     "claude-3-5-sonnet-20241022": 40000,
     "claude-3-7-sonnet-20250219": 20000,
+    "Llama-3.3-8B-Instruct": 128000,
+    "Llama-3.3-70B-Instruct": 128000,
+    "Llama-4-Scout-17B-16E-Instruct-FP8": 128000,
+    "Llama-4-Maverick-17B-128E-Instruct-FP8": 128000,
 }
 
 GEMINI_OPENAI_BASE_URL = "https://generativelanguage.googleapis.com/v1beta/openai/"
 ANTHROPIC_OPENAI_BASE_URL = "https://api.anthropic.com/v1/"
+LLAMA_API_BASE_URL = "https://api.llama.com/compat/v1/"
 
 
 def resolve_model(model: str) -> str:

python/packages/autogen-ext/src/autogen_ext/models/openai/_openai_client.py

@@ -637,6 +637,7 @@ class BaseOpenAIChatCompletionClient(ChatCompletionClient):
                 response=result.model_dump(),
                 prompt_tokens=usage.prompt_tokens,
                 completion_tokens=usage.completion_tokens,
+                tools=create_params.tools,
             )
         )
 
@@ -1369,6 +1370,11 @@ class OpenAIChatCompletionClient(BaseOpenAIChatCompletionClient, Component[OpenA
                 copied_args["base_url"] = _model_info.ANTHROPIC_OPENAI_BASE_URL
             if "api_key" not in copied_args and "ANTHROPIC_API_KEY" in os.environ:
                 copied_args["api_key"] = os.environ["ANTHROPIC_API_KEY"]
+        if copied_args["model"].startswith("Llama-"):
+            if "base_url" not in copied_args:
+                copied_args["base_url"] = _model_info.LLAMA_API_BASE_URL
+            if "api_key" not in copied_args and "LLAMA_API_KEY" in os.environ:
+                copied_args["api_key"] = os.environ["LLAMA_API_KEY"]
 
         client = _openai_client_from_config(copied_args)
         create_args = _create_args_from_config(copied_args)

python/packages/autogen-ext/src/autogen_ext/runtimes/grpc/_worker_runtime.py

@@ -251,6 +251,7 @@ class GrpcWorkerAgentRuntime(AgentRuntime):
         self._subscription_manager = SubscriptionManager()
         self._serialization_registry = SerializationRegistry()
         self._extra_grpc_config = extra_grpc_config or []
+        self._agent_instance_types: Dict[str, Type[Agent]] = {}
 
         if payload_serialization_format not in {JSON_DATA_CONTENT_TYPE, PROTOBUF_DATA_CONTENT_TYPE}:
             raise ValueError(f"Unsupported payload serialization format: {payload_serialization_format}")
@@ -701,6 +702,14 @@ class GrpcWorkerAgentRuntime(AgentRuntime):
         except BaseException as e:
             logger.error("Error handling event", exc_info=e)
 
+    async def _register_agent_type(self, agent_type: str) -> None:
+        if self._host_connection is None:
+            raise RuntimeError("Host connection is not set.")
+        message = agent_worker_pb2.RegisterAgentTypeRequest(type=agent_type)
+        _response: agent_worker_pb2.RegisterAgentTypeResponse = await self._host_connection.stub.RegisterAgent(
+            message, metadata=self._host_connection.metadata
+        )
+
     async def register_factory(
         self,
         type: str | AgentType,
@@ -729,14 +738,38 @@ class GrpcWorkerAgentRuntime(AgentRuntime):
             return agent_instance
 
         self._agent_factories[type.type] = factory_wrapper
-
         # Send the registration request message to the host.
-        message = agent_worker_pb2.RegisterAgentTypeRequest(type=type.type)
-        _response: agent_worker_pb2.RegisterAgentTypeResponse = await self._host_connection.stub.RegisterAgent(
-            message, metadata=self._host_connection.metadata
-        )
+        await self._register_agent_type(type.type)
+
         return type
 
+    async def register_agent_instance(
+        self,
+        agent_instance: Agent,
+        agent_id: AgentId,
+    ) -> AgentId:
+        def agent_factory() -> Agent:
+            raise RuntimeError(
+                "Agent factory was invoked for an agent instance that was not registered. This is likely due to the agent type being incorrectly subscribed to a topic. If this exception occurs when publishing a message to the DefaultTopicId, then it is likely that `skip_class_subscriptions` needs to be turned off when registering the agent."
+            )
+
+        if agent_id in self._instantiated_agents:
+            raise ValueError(f"Agent with id {agent_id} already exists.")
+
+        if agent_id.type not in self._agent_factories:
+            self._agent_factories[agent_id.type] = agent_factory
+            await self._register_agent_type(agent_id.type)
+            self._agent_instance_types[agent_id.type] = type_func_alias(agent_instance)
+        else:
+            if self._agent_factories[agent_id.type].__code__ != agent_factory.__code__:
+                raise ValueError("Agent factories and agent instances cannot be registered to the same type.")
+            if self._agent_instance_types[agent_id.type] != type_func_alias(agent_instance):
+                raise ValueError("Agent instances must be the same object type.")
+
+        await agent_instance.bind_id_and_runtime(id=agent_id, runtime=self)
+        self._instantiated_agents[agent_id] = agent_instance
+        return agent_id
+
     async def _invoke_agent_factory(
         self,
         agent_factory: Callable[[], T | Awaitable[T]] | Callable[[AgentRuntime, AgentId], T | Awaitable[T]],
@@ -757,7 +790,7 @@ class GrpcWorkerAgentRuntime(AgentRuntime):
                 raise ValueError("Agent factory must take 0 or 2 arguments.")
 
             if inspect.isawaitable(agent):
-                return cast(T, await agent)
+                agent = cast(T, await agent)
 
         return agent
 

python/packages/autogen-ext/src/autogen_ext/tools/azure/__init__.py

@@ -4,6 +4,7 @@ from ._ai_search import (
     SearchQuery,
     SearchResult,
     SearchResults,
+    VectorizableTextQuery,
 )
 from ._config import AzureAISearchConfig
 
@@ -14,4 +15,5 @@ __all__ = [
     "SearchResult",
     "SearchResults",
     "AzureAISearchConfig",
+    "VectorizableTextQuery",
 ]

python/packages/autogen-ext/src/autogen_ext/tools/azure/_config.py

@@ -6,173 +6,180 @@ settings for authentication, search behavior, retry policies, and caching.
 
 import logging
 from typing import (
-    Any,
-    Dict,
     List,
     Literal,
     Optional,
-    Type,
     TypeVar,
     Union,
 )
 
-from azure.core.credentials import AzureKeyCredential, TokenCredential
-from pydantic import BaseModel, Field, model_validator
-
-# Add explicit ignore for the specific model validator error
-# pyright: reportArgumentType=false
-# pyright: reportUnknownArgumentType=false
-# pyright: reportUnknownVariableType=false
+from azure.core.credentials import AzureKeyCredential
+from azure.core.credentials_async import AsyncTokenCredential
+from pydantic import BaseModel, Field, field_validator, model_validator
 
 T = TypeVar("T", bound="AzureAISearchConfig")
 
 logger = logging.getLogger(__name__)
 
+QueryTypeLiteral = Literal["simple", "full", "semantic", "vector"]
+DEFAULT_API_VERSION = "2023-10-01-preview"
+
 
 class AzureAISearchConfig(BaseModel):
-    """Configuration for Azure AI Search tool.
+    """Configuration for Azure AI Search with validation.
 
-    This class defines the configuration parameters for :class:`AzureAISearchTool`.
-    It provides options for customizing search behavior including query types,
-    field selection, authentication, retry policies, and caching strategies.
+    This class defines the configuration parameters for Azure AI Search tools, including
+    authentication, search behavior, caching, and embedding settings.
 
     .. note::
-
-        This class requires the :code:`azure` extra for the :code:`autogen-ext` package.
+        This class requires the ``azure`` extra for the ``autogen-ext`` package.
 
         .. code-block:: bash
 
             pip install -U "autogen-ext[azure]"
 
-    Example:
+    .. note::
+        **Prerequisites:**
+
+        1. An Azure AI Search service must be created in your Azure subscription.
+        2. The search index must be properly configured for your use case:
+
+           - For vector search: Index must have vector fields
+           - For semantic search: Index must have semantic configuration
+           - For hybrid search: Both vector fields and text fields must be configured
+        3. Required packages:
+
+           - Base functionality: ``azure-search-documents>=11.4.0``
+           - For Azure OpenAI embeddings: ``openai azure-identity``
+           - For OpenAI embeddings: ``openai``
+
+    Example Usage:
         .. code-block:: python
 
             from azure.core.credentials import AzureKeyCredential
             from autogen_ext.tools.azure import AzureAISearchConfig
 
+            # Basic configuration for full-text search
             config = AzureAISearchConfig(
-                name="doc_search",
-                endpoint="https://my-search.search.windows.net",
-                index_name="my-index",
+                name="doc-search",
+                endpoint="https://your-search.search.windows.net",  # Your Azure AI Search endpoint
+                index_name="<your-index>",  # Name of your search index
+                credential=AzureKeyCredential("<your-key>"),  # Your Azure AI Search admin key
+                query_type="simple",
+                search_fields=["content", "title"],  # Update with your searchable fields
+                top=5,
+            )
+
+            # Configuration for vector search with Azure OpenAI embeddings
+            vector_config = AzureAISearchConfig(
+                name="vector-search",
+                endpoint="https://your-search.search.windows.net",
+                index_name="<your-index>",
                 credential=AzureKeyCredential("<your-key>"),
                 query_type="vector",
-                vector_fields=["embedding"],
+                vector_fields=["embedding"],  # Update with your vector field name
+                embedding_provider="azure_openai",
+                embedding_model="text-embedding-ada-002",
+                openai_endpoint="https://your-openai.openai.azure.com",  # Your Azure OpenAI endpoint
+                openai_api_key="<your-openai-key>",  # Your Azure OpenAI key
+                top=5,
             )
 
-    For more details, see:
-        * `Azure AI Search Overview <https://learn.microsoft.com/azure/search/search-what-is-azure-search>`_
-        * `Vector Search <https://learn.microsoft.com/azure/search/vector-search-overview>`_
-
-    Args:
-        name (str): Name for the tool instance, used to identify it in the agent's toolkit.
-        description (Optional[str]): Human-readable description of what this tool does and how to use it.
-        endpoint (str): The full URL of your Azure AI Search service, in the format
-            'https://<service-name>.search.windows.net'.
-        index_name (str): Name of the target search index in your Azure AI Search service.
-            The index must be pre-created and properly configured.
-        api_version (str): Azure AI Search REST API version to use. Defaults to '2023-11-01'.
-            Only change if you need specific features from a different API version.
-        credential (Union[AzureKeyCredential, TokenCredential]): Azure authentication credential:
-            - AzureKeyCredential: For API key authentication (admin/query key)
-            - TokenCredential: For Azure AD authentication (e.g., DefaultAzureCredential)
-        query_type (Literal["keyword", "fulltext", "vector", "semantic"]): The search query mode to use:
-            - 'keyword': Basic keyword search (default)
-            - 'fulltext': Full Lucene query syntax
-            - 'vector': Vector similarity search
-            - 'semantic': Semantic search using semantic configuration
-        search_fields (Optional[List[str]]): List of index fields to search within. If not specified,
-            searches all searchable fields. Example: ['title', 'content'].
-        select_fields (Optional[List[str]]): Fields to return in search results. If not specified,
-            returns all fields. Use to optimize response size.
-        vector_fields (Optional[List[str]]): Vector field names for vector search. Must be configured
-            in your search index as vector fields. Required for vector search.
-        top (Optional[int]): Maximum number of documents to return in search results.
-            Helps control response size and processing time.
-        retry_enabled (bool): Whether to enable retry policy for transient errors. Defaults to True.
-        retry_max_attempts (Optional[int]): Maximum number of retry attempts for failed requests. Defaults to 3.
-        retry_mode (Literal["fixed", "exponential"]): Retry backoff strategy: fixed or exponential. Defaults to "exponential".
-        enable_caching (bool): Whether to enable client-side caching of search results. Defaults to False.
-        cache_ttl_seconds (int): Time-to-live for cached search results in seconds. Defaults to 300 (5 minutes).
-        filter (Optional[str]): OData filter expression to refine search results.
+            # Configuration for hybrid search with semantic ranking
+            hybrid_config = AzureAISearchConfig(
+                name="hybrid-search",
+                endpoint="https://your-search.search.windows.net",
+                index_name="<your-index>",
+                credential=AzureKeyCredential("<your-key>"),
+                query_type="semantic",
+                semantic_config_name="<your-semantic-config>",  # Name of your semantic configuration
+                search_fields=["content", "title"],  # Update with your search fields
+                vector_fields=["embedding"],  # Update with your vector field name
+                embedding_provider="openai",
+                embedding_model="text-embedding-ada-002",
+                openai_api_key="<your-openai-key>",  # Your OpenAI API key
+                top=5,
+            )
     """
 
-    name: str = Field(description="The name of the tool")
-    description: Optional[str] = Field(default=None, description="A description of the tool")
-    endpoint: str = Field(description="The endpoint URL for your Azure AI Search service")
-    index_name: str = Field(description="The name of the search index to query")
-    api_version: str = Field(default="2023-11-01", description="API version to use")
-    credential: Union[AzureKeyCredential, TokenCredential] = Field(
-        description="The credential to use for authentication"
+    name: str = Field(description="The name of this tool instance")
+    description: Optional[str] = Field(default=None, description="Description explaining the tool's purpose")
+    endpoint: str = Field(description="The full URL of your Azure AI Search service")
+    index_name: str = Field(description="Name of the search index to query")
+    credential: Union[AzureKeyCredential, AsyncTokenCredential] = Field(
+        description="Azure credential for authentication (API key or token)"
     )
-    query_type: Literal["keyword", "fulltext", "vector", "semantic"] = Field(
-        default="keyword",
-        description="Type of query to perform (keyword for classic, fulltext for Lucene, vector for embedding, semantic for semantic/AI search)",
+    api_version: str = Field(
+        default=DEFAULT_API_VERSION,
+        description=f"Azure AI Search API version to use. Defaults to {DEFAULT_API_VERSION}.",
     )
-    search_fields: Optional[List[str]] = Field(default=None, description="Optional list of fields to search in")
-    select_fields: Optional[List[str]] = Field(default=None, description="Optional list of fields to return in results")
-    vector_fields: Optional[List[str]] = Field(
-        default=None, description="Optional list of vector fields for vector search"
+    query_type: QueryTypeLiteral = Field(
+        default="simple", description="Type of search to perform: simple, full, semantic, or vector"
     )
-    top: Optional[int] = Field(default=None, description="Optional number of results to return")
-    filter: Optional[str] = Field(default=None, description="Optional OData filter expression to refine search results")
-
-    retry_enabled: bool = Field(default=True, description="Whether to enable retry policy for transient errors")
-    retry_max_attempts: Optional[int] = Field(
-        default=3, description="Maximum number of retry attempts for failed requests"
+    search_fields: Optional[List[str]] = Field(default=None, description="Fields to search within documents")
+    select_fields: Optional[List[str]] = Field(default=None, description="Fields to return in search results")
+    vector_fields: Optional[List[str]] = Field(default=None, description="Fields to use for vector search")
+    top: Optional[int] = Field(
+        default=None, description="Maximum number of results to return. For vector searches, acts as k in k-NN."
     )
-    retry_mode: Literal["fixed", "exponential"] = Field(
-        default="exponential",
-        description="Retry backoff strategy: fixed or exponential",
+    filter: Optional[str] = Field(default=None, description="OData filter expression to refine search results")
+    semantic_config_name: Optional[str] = Field(
+        default=None, description="Semantic configuration name for enhanced results"
     )
 
-    enable_caching: bool = Field(
-        default=False,
-        description="Whether to enable client-side caching of search results",
-    )
-    cache_ttl_seconds: int = Field(
-        default=300,  # 5 minutes
-        description="Time-to-live for cached search results in seconds",
-    )
+    enable_caching: bool = Field(default=False, description="Whether to cache search results")
+    cache_ttl_seconds: int = Field(default=300, description="How long to cache results in seconds")
 
     embedding_provider: Optional[str] = Field(
-        default=None,
-        description="Name of embedding provider to use (e.g., 'azure_openai', 'openai')",
-    )
-    embedding_model: Optional[str] = Field(default=None, description="Model name to use for generating embeddings")
-    embedding_dimension: Optional[int] = Field(
-        default=None, description="Dimension of embedding vectors produced by the model"
+        default=None, description="Name of embedding provider for client-side embeddings"
     )
+    embedding_model: Optional[str] = Field(default=None, description="Model name for client-side embeddings")
+    openai_api_key: Optional[str] = Field(default=None, description="API key for OpenAI/Azure OpenAI embeddings")
+    openai_api_version: Optional[str] = Field(default=None, description="API version for Azure OpenAI embeddings")
+    openai_endpoint: Optional[str] = Field(default=None, description="Endpoint URL for Azure OpenAI embeddings")
 
     model_config = {"arbitrary_types_allowed": True}
 
-    @classmethod
-    @model_validator(mode="before")
-    def validate_credentials(cls: Type[T], data: Any) -> Any:
-        """Validate and convert credential data."""
-        if not isinstance(data, dict):
-            return data
-
-        result = {}
-
-        for key, value in data.items():
-            result[str(key)] = value
-
-        if "credential" in result:
-            credential = result["credential"]
-
-            if isinstance(credential, dict) and "api_key" in credential:
-                api_key = str(credential["api_key"])
-                result["credential"] = AzureKeyCredential(api_key)
-
-        return result
-
-    def model_dump(self, **kwargs: Any) -> Dict[str, Any]:
-        """Custom model_dump to handle credentials."""
-        result: Dict[str, Any] = super().model_dump(**kwargs)
-
-        if isinstance(self.credential, AzureKeyCredential):
-            result["credential"] = {"type": "AzureKeyCredential"}
-        elif isinstance(self.credential, TokenCredential):
-            result["credential"] = {"type": "TokenCredential"}
-
-        return result
+    @field_validator("endpoint")
+    def validate_endpoint(cls, v: str) -> str:
+        """Validate that the endpoint is a valid URL."""
+        if not v.startswith(("http://", "https://")):
+            raise ValueError("endpoint must be a valid URL starting with http:// or https://")
+        return v
+
+    @field_validator("query_type")
+    def normalize_query_type(cls, v: QueryTypeLiteral) -> QueryTypeLiteral:
+        """Normalize query type to standard values."""
+        if not v:
+            return "simple"
+
+        if isinstance(v, str) and v.lower() == "fulltext":
+            return "full"
+
+        return v
+
+    @field_validator("top")
+    def validate_top(cls, v: Optional[int]) -> Optional[int]:
+        """Ensure top is a positive integer if provided."""
+        if v is not None and v <= 0:
+            raise ValueError("top must be a positive integer")
+        return v
+
+    @model_validator(mode="after")
+    def validate_interdependent_fields(self) -> "AzureAISearchConfig":
+        """Validate interdependent fields after all fields have been parsed."""
+        if self.query_type == "semantic" and not self.semantic_config_name:
+            raise ValueError("semantic_config_name must be provided when query_type is 'semantic'")
+
+        if self.query_type == "vector" and not self.vector_fields:
+            raise ValueError("vector_fields must be provided for vector search")
+
+        if (
+            self.embedding_provider
+            and self.embedding_provider.lower() == "azure_openai"
+            and self.embedding_model
+            and not self.openai_endpoint
+        ):
+            raise ValueError("openai_endpoint must be provided for azure_openai embedding provider")
+
+        return self

python/packages/autogen-ext/src/autogen_ext/tools/mcp/_base.py

@@ -74,21 +74,52 @@ class McpToolAdapter(BaseTool[BaseModel, Any], ABC, Generic[TServerParams]):
             await session.initialize()
             return await self._run(args=kwargs, cancellation_token=cancellation_token, session=session)
 
+    def _normalize_payload_to_content_list(
+        self, payload: list[TextContent | ImageContent | EmbeddedResource]
+    ) -> list[TextContent | ImageContent | EmbeddedResource]:
+        """
+        Normalizes a raw tool output payload into a list of content items.
+        - If payload is already a list of (TextContent, ImageContent, EmbeddedResource), it's returned as is.
+        - If payload is a single TextContent, ImageContent, or EmbeddedResource, it's wrapped in a list.
+        - If payload is a string, it's wrapped in [TextContent(text=payload)].
+        - Otherwise, the payload is stringified and wrapped in [TextContent(text=str(payload))].
+        """
+        if isinstance(payload, list) and all(
+            isinstance(item, (TextContent, ImageContent, EmbeddedResource)) for item in payload
+        ):
+            return payload
+        elif isinstance(payload, (TextContent, ImageContent, EmbeddedResource)):
+            return [payload]
+        elif isinstance(payload, str):
+            return [TextContent(text=payload, type="text")]
+        else:
+            return [TextContent(text=str(payload), type="text")]
+
     async def _run(self, args: Dict[str, Any], cancellation_token: CancellationToken, session: ClientSession) -> Any:
+        exceptions_to_catch: tuple[Type[BaseException], ...]
+        if hasattr(builtins, "ExceptionGroup"):
+            exceptions_to_catch = (asyncio.CancelledError, builtins.ExceptionGroup)
+        else:
+            exceptions_to_catch = (asyncio.CancelledError,)
+
         try:
             if cancellation_token.is_cancelled():
-                raise Exception("Operation cancelled")
+                raise asyncio.CancelledError("Operation cancelled")
 
             result_future = asyncio.ensure_future(session.call_tool(name=self._tool.name, arguments=args))
             cancellation_token.link_future(result_future)
             result = await result_future
 
+            normalized_content_list = self._normalize_payload_to_content_list(result.content)
+
             if result.isError:
-                raise Exception(f"MCP tool execution failed: {result.content}")
-            return result.content
-        except Exception as e:
-            error_message = self._format_errors(e)
-            raise Exception(error_message) from e
+                serialized_error_message = self.return_value_as_string(normalized_content_list)
+                raise Exception(serialized_error_message)
+            return normalized_content_list
+
+        except exceptions_to_catch:
+            # Re-raise these specific exception types directly.
+            raise
 
     @classmethod
     async def from_server_params(cls, server_params: TServerParams, tool_name: str) -> "McpToolAdapter[TServerParams]":
@@ -138,16 +169,3 @@ class McpToolAdapter(BaseTool[BaseModel, Any], ABC, Generic[TServerParams]):
                 return {}
 
         return json.dumps([serialize_item(item) for item in value])
-
-    def _format_errors(self, error: Exception) -> str:
-        """Recursively format errors into a string."""
-
-        error_message = ""
-        if hasattr(builtins, "ExceptionGroup") and isinstance(error, builtins.ExceptionGroup):
-            # ExceptionGroup is available in Python 3.11+.
-            # TODO: how to make this compatible with Python 3.10?
-            for sub_exception in error.exceptions:  # type: ignore
-                error_message += self._format_errors(sub_exception)  # type: ignore
-        else:
-            error_message += f"{str(error)}\n"
-        return error_message

python/packages/autogen-ext/src/autogen_ext/tools/mcp/_workbench.py

@@ -1,3 +1,4 @@
+import asyncio
 import builtins
 import warnings
 from typing import Any, List, Literal, Mapping
@@ -152,6 +153,7 @@ class McpWorkbench(Workbench, Component[McpWorkbenchConfig]):
         self._server_params = server_params
         # self._session: ClientSession | None = None
         self._actor: McpSessionActor | None = None
+        self._actor_loop: asyncio.AbstractEventLoop | None = None
         self._read = None
         self._write = None
 
@@ -253,6 +255,7 @@ class McpWorkbench(Workbench, Component[McpWorkbenchConfig]):
         if isinstance(self._server_params, (StdioServerParams, SseServerParams)):
             self._actor = McpSessionActor(self._server_params)
             await self._actor.initialize()
+            self._actor_loop = asyncio.get_event_loop()
         else:
             raise ValueError(f"Unsupported server params type: {type(self._server_params)}")
 
@@ -282,4 +285,10 @@ class McpWorkbench(Workbench, Component[McpWorkbenchConfig]):
 
     def __del__(self) -> None:
         # Ensure the actor is stopped when the workbench is deleted
-        pass
+        if self._actor and self._actor_loop:
+            loop = self._actor_loop
+            if loop.is_running() and not loop.is_closed():
+                loop.call_soon_threadsafe(lambda: asyncio.create_task(self.stop()))
+            else:
+                msg = "Cannot safely stop actor at [McpWorkbench.__del__]: loop is closed or not running"
+                warnings.warn(msg, RuntimeWarning, stacklevel=2)

python/packages/autogen-ext/tests/code_executors/test_docker_commandline_code_executor.py

@@ -1,8 +1,8 @@
 # mypy: disable-error-code="no-any-unimported"
 import asyncio
 import os
-import sys
 import shutil
+import sys
 import tempfile
 from pathlib import Path
 from typing import AsyncGenerator, TypeAlias

python/packages/autogen-ext/tests/models/test_anthropic_model_client.py

@@ -10,12 +10,17 @@ from autogen_core.models import (
     CreateResult,
     FunctionExecutionResult,
     FunctionExecutionResultMessage,
+    ModelInfo,
     SystemMessage,
     UserMessage,
 )
 from autogen_core.models._types import LLMMessage
 from autogen_core.tools import FunctionTool
-from autogen_ext.models.anthropic import AnthropicChatCompletionClient
+from autogen_ext.models.anthropic import (
+    AnthropicBedrockChatCompletionClient,
+    AnthropicChatCompletionClient,
+    BedrockInfo,
+)
 
 
 def _pass_function(input: str) -> str:
@@ -46,6 +51,29 @@ async def test_anthropic_serialization_api_key() -> None:
     client2 = AnthropicChatCompletionClient.load_component(config)
     assert client2
 
+    bedrock_client = AnthropicBedrockChatCompletionClient(
+        model="claude-3-haiku-20240307",  # Use haiku for faster/cheaper testing
+        api_key="sk-password",
+        model_info=ModelInfo(
+            vision=False, function_calling=True, json_output=False, family="unknown", structured_output=True
+        ),
+        bedrock_info=BedrockInfo(
+            aws_access_key="<aws_access_key>",
+            aws_secret_key="<aws_secret_key>",
+            aws_session_token="<aws_session_token>",
+            aws_region="<aws_region>",
+        ),
+    )
+    assert bedrock_client
+    bedrock_config = bedrock_client.dump_component()
+    assert bedrock_config
+    assert "sk-password" not in str(bedrock_config)
+    serialized_bedrock_config = bedrock_config.model_dump_json()
+    assert serialized_bedrock_config
+    assert "sk-password" not in serialized_bedrock_config
+    bedrock_client2 = AnthropicBedrockChatCompletionClient.load_component(bedrock_config)
+    assert bedrock_client2
+
 
 @pytest.mark.asyncio
 async def test_anthropic_basic_completion(caplog: pytest.LogCaptureFixture) -> None:

python/packages/autogen-ext/tests/models/test_azure_ai_model_client.py

@@ -3,7 +3,7 @@ import logging
 import os
 from datetime import datetime
 from typing import Any, AsyncGenerator, List, Type, Union
-from unittest.mock import MagicMock
+from unittest.mock import AsyncMock, MagicMock
 
 import pytest
 from autogen_core import CancellationToken, FunctionCall, Image
@@ -570,7 +570,7 @@ def thought_with_tool_call_stream_client(monkeypatch: pytest.MonkeyPatch) -> Azu
         )
 
     mock_client = MagicMock()
-    mock_client.close = MagicMock()
+    mock_client.close = AsyncMock()
 
     async def mock_complete(*args: Any, **kwargs: Any) -> Any:
         if kwargs.get("stream", False):

python/packages/autogen-ext/tests/models/test_ollama_chat_completion_client.py

@@ -570,7 +570,7 @@ async def test_ollama_create_structured_output(model: str, ollama_client: Ollama
 
 
 @pytest.mark.asyncio
-@pytest.mark.parametrize("model", ["qwen2.5:0.5b", "llama3.2:1b"])
+@pytest.mark.parametrize("model", ["qwen2.5:0.5b", "llama3.2:1b", "qwen3:0.6b"])
 async def test_ollama_create_tools(model: str, ollama_client: OllamaChatCompletionClient) -> None:
     def add(x: int, y: int) -> str:
         return str(x + y)
@@ -653,7 +653,7 @@ async def test_ollama_create_structured_output_with_tools(
 
 
 @pytest.mark.asyncio
-@pytest.mark.parametrize("model", ["qwen2.5:0.5b", "llama3.2:1b"])
+@pytest.mark.parametrize("model", ["qwen2.5:0.5b", "llama3.2:1b", "qwen3:0.6b"])
 async def test_ollama_create_stream_tools(model: str, ollama_client: OllamaChatCompletionClient) -> None:
     def add(x: int, y: int) -> str:
         return str(x + y)

python/packages/autogen-ext/tests/models/test_openai_model_client.py

@@ -2485,4 +2485,17 @@ async def test_multimodal_message_test(
     _ = await ocr_agent.run(task=multi_modal_message)
 
 
+@pytest.mark.asyncio
+async def test_mistral_remove_name() -> None:
+    # Test that the name pramaeter is removed from the message
+    # when the model is Mistral
+    message = UserMessage(content="foo", source="user")
+    params = to_oai_type(message, prepend_name=False, model="mistral-7b", model_family=ModelFamily.MISTRAL)
+    assert ("name" in params[0]) is False
+
+    # when the model is gpt-4o, the name parameter is not removed
+    params = to_oai_type(message, prepend_name=False, model="gpt-4o", model_family=ModelFamily.GPT_4O)
+    assert ("name" in params[0]) is True
+
+
 # TODO: add integration tests for Azure OpenAI using AAD token.

python/packages/autogen-ext/tests/test_worker_runtime.py

@@ -577,6 +577,139 @@ async def test_grpc_max_message_size() -> None:
         await host.stop()
 
 
+@pytest.mark.grpc
+@pytest.mark.asyncio
+async def test_agent_type_register_instance() -> None:
+    host_address = "localhost:50051"
+    agent1_id = AgentId(type="name", key="default")
+    agentdup_id = AgentId(type="name", key="default")
+    agent2_id = AgentId(type="name", key="notdefault")
+    host = GrpcWorkerAgentRuntimeHost(address=host_address)
+    host.start()
+
+    worker = GrpcWorkerAgentRuntime(host_address=host_address)
+    agent1 = NoopAgent()
+    agent2 = NoopAgent()
+    agentdup = NoopAgent()
+    await worker.start()
+
+    await worker.register_agent_instance(agent1, agent_id=agent1_id)
+    await worker.register_agent_instance(agent2, agent_id=agent2_id)
+
+    with pytest.raises(ValueError):
+        await worker.register_agent_instance(agentdup, agent_id=agentdup_id)
+
+    assert await worker.try_get_underlying_agent_instance(agent1_id, type=NoopAgent) == agent1
+    assert await worker.try_get_underlying_agent_instance(agent2_id, type=NoopAgent) == agent2
+
+    await worker.stop()
+    await host.stop()
+
+
+@pytest.mark.grpc
+@pytest.mark.asyncio
+async def test_agent_type_register_instance_different_types() -> None:
+    host_address = "localhost:50051"
+    agent1_id = AgentId(type="name", key="noop")
+    agent2_id = AgentId(type="name", key="loopback")
+    host = GrpcWorkerAgentRuntimeHost(address=host_address)
+    host.start()
+
+    worker = GrpcWorkerAgentRuntime(host_address=host_address)
+    agent1 = NoopAgent()
+    agent2 = LoopbackAgent()
+    await worker.start()
+
+    await worker.register_agent_instance(agent1, agent_id=agent1_id)
+    with pytest.raises(ValueError):
+        await worker.register_agent_instance(agent2, agent_id=agent2_id)
+
+    await worker.stop()
+    await host.stop()
+
+
+@pytest.mark.grpc
+@pytest.mark.asyncio
+async def test_register_instance_factory() -> None:
+    host_address = "localhost:50051"
+    agent1_id = AgentId(type="name", key="default")
+    host = GrpcWorkerAgentRuntimeHost(address=host_address)
+    host.start()
+
+    worker = GrpcWorkerAgentRuntime(host_address=host_address)
+    agent1 = NoopAgent()
+    await worker.start()
+
+    await agent1.register_instance(runtime=worker, agent_id=agent1_id)
+
+    with pytest.raises(ValueError):
+        await NoopAgent.register(runtime=worker, type="name", factory=lambda: NoopAgent())
+
+    await worker.stop()
+    await host.stop()
+
+
+@pytest.mark.grpc
+@pytest.mark.asyncio
+async def test_instance_factory_messaging() -> None:
+    host_address = "localhost:50051"
+    loopback_agent_id = AgentId(type="dm_agent", key="dm_agent")
+    cascading_agent_id = AgentId(type="instance_agent", key="instance_agent")
+    host = GrpcWorkerAgentRuntimeHost(address=host_address)
+    host.start()
+
+    worker = GrpcWorkerAgentRuntime(host_address=host_address)
+    cascading_agent = CascadingAgent(max_rounds=5)
+    loopback_agent = LoopbackAgent()
+    await worker.start()
+
+    await loopback_agent.register_instance(worker, agent_id=loopback_agent_id)
+    resp = await worker.send_message(message=ContentMessage(content="Hello!"), recipient=loopback_agent_id)
+    assert resp == ContentMessage(content="Hello!")
+
+    await cascading_agent.register_instance(worker, agent_id=cascading_agent_id)
+    await CascadingAgent.register(worker, "factory_agent", lambda: CascadingAgent(max_rounds=5))
+
+    # instance_agent will publish a message that factory_agent will pick up
+    for i in range(5):
+        await worker.publish_message(
+            CascadingMessageType(round=i + 1), TopicId(type="instance_agent", source="instance_agent")
+        )
+    await asyncio.sleep(2)
+
+    agent = await worker.try_get_underlying_agent_instance(AgentId("factory_agent", "default"), CascadingAgent)
+    assert agent.num_calls == 4
+    assert cascading_agent.num_calls == 5
+
+    await worker.stop()
+    await host.stop()
+
+
+# GrpcWorkerAgentRuntimeHost eats exceptions in the main loop
+# @pytest.mark.grpc
+# @pytest.mark.asyncio
+# async def test_agent_type_register_instance_publish_new_source() -> None:
+#     host_address = "localhost:50056"
+#     agent_id = AgentId(type="name", key="default")
+#     agent1 = LoopbackAgent()
+#     host = GrpcWorkerAgentRuntimeHost(address=host_address)
+#     host.start()
+#     worker = GrpcWorkerAgentRuntime(host_address=host_address)
+#     await worker.start()
+#     publisher = GrpcWorkerAgentRuntime(host_address=host_address)
+#     publisher.add_message_serializer(try_get_known_serializers_for_type(MessageType))
+#     await publisher.start()
+
+#     await agent1.register_instance(worker, agent_id=agent_id)
+#     await worker.add_subscription(TypeSubscription("notdefault", "name"))
+
+#     with pytest.raises(RuntimeError):
+#         await worker.publish_message(MessageType(), TopicId("notdefault", "notdefault"))
+#         await asyncio.sleep(2)
+
+#     await worker.stop()
+#     await host.stop()
+
 if __name__ == "__main__":
     os.environ["GRPC_VERBOSITY"] = "DEBUG"
     os.environ["GRPC_TRACE"] = "all"

python/packages/autogen-ext/tests/tools/azure/conftest.py

@@ -1,7 +1,7 @@
 """Test fixtures for Azure AI Search tool tests."""
 
 import warnings
-from typing import Any, Dict, Generator, List, Protocol, Type, TypeVar, Union
+from typing import Any, Dict, Iterator, List, Protocol, TypeVar, Union
 from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
@@ -9,6 +9,17 @@ from autogen_core import ComponentModel
 
 T = TypeVar("T")
 
+try:
+    from azure.core.credentials import AzureKeyCredential, TokenCredential
+
+    azure_sdk_available = True
+except ImportError:
+    azure_sdk_available = False
+
+skip_if_no_azure_sdk = pytest.mark.skipif(
+    not azure_sdk_available, reason="Azure SDK components (azure-search-documents, azure-identity) not available"
+)
+
 
 class AccessTokenProtocol(Protocol):
     """Protocol matching Azure AccessToken."""
@@ -47,18 +58,13 @@ class MockTokenCredential:
         return MockAccessToken("mock-token", 12345)
 
 
-try:
-    from azure.core.credentials import AccessToken, AzureKeyCredential, TokenCredential
-
-    _access_token_type: Type[AccessToken] = AccessToken
-    azure_sdk_available = True
-except ImportError:
-    AzureKeyCredential = MockAzureKeyCredential  # type: ignore
-    TokenCredential = MockTokenCredential  # type: ignore
-    _access_token_type = MockAccessToken  # type: ignore
-    azure_sdk_available = False
-
-CredentialType = Union[AzureKeyCredential, TokenCredential, MockAzureKeyCredential, MockTokenCredential, Any]
+CredentialType = Union[
+    AzureKeyCredential,  # pyright: ignore [reportPossiblyUnboundVariable]
+    TokenCredential,  # pyright: ignore [reportPossiblyUnboundVariable]
+    MockAzureKeyCredential,
+    MockTokenCredential,
+    Any,
+]
 
 needs_azure_sdk = pytest.mark.skipif(not azure_sdk_available, reason="Azure SDK not available")
 
@@ -70,10 +76,14 @@ warnings.filterwarnings(
 
 
 @pytest.fixture
-def mock_vectorized_query() -> Generator[MagicMock, None, None]:
+def mock_vectorized_query() -> MagicMock:
     """Create a mock VectorizedQuery for testing."""
-    with patch("azure.search.documents.models.VectorizedQuery") as mock:
-        yield mock
+    if azure_sdk_available:
+        from azure.search.documents.models import VectorizedQuery
+
+        return MagicMock(spec=VectorizedQuery)
+    else:
+        return MagicMock()
 
 
 @pytest.fixture
@@ -87,7 +97,7 @@ def test_config() -> ComponentModel:
             "endpoint": "https://test-search-service.search.windows.net",
             "index_name": "test-index",
             "api_version": "2023-10-01-Preview",
-            "credential": AzureKeyCredential("test-key") if azure_sdk_available else {"api_key": "test-key"},
+            "credential": AzureKeyCredential("test-key") if azure_sdk_available else {"api_key": "test-key"},  # pyright: ignore [reportPossiblyUnboundVariable]
             "query_type": "keyword",
             "search_fields": ["content", "title"],
             "select_fields": ["id", "content", "title", "source"],
@@ -106,7 +116,7 @@ def keyword_config() -> ComponentModel:
             "description": "Keyword search tool",
             "endpoint": "https://test-search-service.search.windows.net",
             "index_name": "test-index",
-            "credential": AzureKeyCredential("test-key") if azure_sdk_available else {"api_key": "test-key"},
+            "credential": AzureKeyCredential("test-key") if azure_sdk_available else {"api_key": "test-key"},  # pyright: ignore [reportPossiblyUnboundVariable]
             "query_type": "keyword",
             "search_fields": ["content", "title"],
             "select_fields": ["id", "content", "title", "source"],
@@ -125,7 +135,7 @@ def vector_config() -> ComponentModel:
             "endpoint": "https://test-search-service.search.windows.net",
             "index_name": "test-index",
             "api_version": "2023-10-01-Preview",
-            "credential": AzureKeyCredential("test-key") if azure_sdk_available else {"api_key": "test-key"},
+            "credential": AzureKeyCredential("test-key") if azure_sdk_available else {"api_key": "test-key"},  # pyright: ignore [reportPossiblyUnboundVariable]
             "query_type": "vector",
             "vector_fields": ["embedding"],
             "select_fields": ["id", "content", "title", "source"],
@@ -145,7 +155,7 @@ def hybrid_config() -> ComponentModel:
             "endpoint": "https://test-search-service.search.windows.net",
             "index_name": "test-index",
             "api_version": "2023-10-01-Preview",
-            "credential": AzureKeyCredential("test-key") if azure_sdk_available else {"api_key": "test-key"},
+            "credential": AzureKeyCredential("test-key") if azure_sdk_available else {"api_key": "test-key"},  # pyright: ignore [reportPossiblyUnboundVariable]
             "query_type": "keyword",
             "search_fields": ["content", "title"],
             "vector_fields": ["embedding"],
@@ -196,108 +206,18 @@ class AsyncIterator:
 
 
 @pytest.fixture
-def mock_search_client(mock_search_response: List[Dict[str, Any]]) -> tuple[MagicMock, Any]:
-    """Create a mock search client for testing."""
-    mock_client = MagicMock()
-    mock_client.__aenter__ = AsyncMock(return_value=mock_client)
-    mock_client.__aexit__ = AsyncMock(return_value=None)
-
-    search_results = AsyncIterator(mock_search_response)
-    mock_client.search = MagicMock(return_value=search_results)
-
-    patcher = patch("azure.search.documents.aio.SearchClient", return_value=mock_client)
-
-    return mock_client, patcher
-
-
-def test_validate_credentials_scenarios() -> None:
-    """Test all validate_credentials scenarios to ensure full code coverage."""
-    import sys
-
-    from autogen_ext.tools.azure._config import AzureAISearchConfig
-
-    module_path = sys.modules[AzureAISearchConfig.__module__].__file__
-    if module_path is not None:
-        assert "autogen-ext" in module_path
-
-    data: Any = "not a dict"
-    result: Any = AzureAISearchConfig.validate_credentials(data)  # type: ignore
-    assert result == data
-
-    data_empty: Dict[str, Any] = {}
-    result_empty: Dict[str, Any] = AzureAISearchConfig.validate_credentials(data_empty)  # type: ignore
-    assert isinstance(result_empty, dict)
-
-    data_items: Dict[str, Any] = {"key1": "value1", "key2": "value2"}
-    result_items: Dict[str, Any] = AzureAISearchConfig.validate_credentials(data_items)  # type: ignore
-    assert result_items["key1"] == "value1"
-    assert result_items["key2"] == "value2"
-
-    data_with_api_key: Dict[str, Any] = {
-        "name": "test",
-        "endpoint": "https://test.search.windows.net",
-        "index_name": "test-index",
-        "credential": {"api_key": "test-key"},
-    }
-    result_with_api_key: Dict[str, Any] = AzureAISearchConfig.validate_credentials(data_with_api_key)  # type: ignore
-
-    cred = result_with_api_key["credential"]  # type: ignore
-    assert isinstance(cred, (AzureKeyCredential, MockAzureKeyCredential))
-    assert hasattr(cred, "key")
-    assert cred.key == "test-key"  # type: ignore
-
-    credential: Any = AzureKeyCredential("test-key")
-    data_with_credential: Dict[str, Any] = {
-        "name": "test",
-        "endpoint": "https://test.search.windows.net",
-        "index_name": "test-index",
-        "credential": credential,
-    }
-    result_with_credential: Dict[str, Any] = AzureAISearchConfig.validate_credentials(data_with_credential)  # type: ignore
-    assert result_with_credential["credential"] is credential
-
-    data_without_api_key: Dict[str, Any] = {
-        "name": "test",
-        "endpoint": "https://test.search.windows.net",
-        "index_name": "test-index",
-        "credential": {"username": "test-user", "password": "test-pass"},
-    }
-    result_without_api_key: Dict[str, Any] = AzureAISearchConfig.validate_credentials(data_without_api_key)  # type: ignore
-    assert result_without_api_key["credential"] == {"username": "test-user", "password": "test-pass"}
-
-
-def test_model_dump_scenarios() -> None:
-    """Test all model_dump scenarios to ensure full code coverage."""
-    import sys
-
-    from autogen_ext.tools.azure._config import AzureAISearchConfig
-
-    module_path = sys.modules[AzureAISearchConfig.__module__].__file__
-    if module_path is not None:
-        assert "autogen-ext" in module_path
-
-    config = AzureAISearchConfig(
-        name="test",
-        endpoint="https://endpoint",
-        index_name="index",
-        credential=AzureKeyCredential("key"),  # type: ignore
-    )
-    result = config.model_dump()
-    assert result["credential"] == {"type": "AzureKeyCredential"}
+def mock_search_client(mock_search_response: List[Dict[str, Any]]) -> Iterator[MagicMock]:
+    """Create a mock search client for testing, with the patch active."""
+    mock_client_instance = MagicMock()
+    mock_client_instance.__aenter__ = AsyncMock(return_value=mock_client_instance)
+    mock_client_instance.__aexit__ = AsyncMock(return_value=None)
 
-    if azure_sdk_available:
-        from azure.core.credentials import AccessToken
-        from azure.core.credentials import TokenCredential as RealTokenCredential
-
-        class TestTokenCredential(RealTokenCredential):
-            def get_token(self, *args: Any, **kwargs: Any) -> AccessToken:
-                """Override of get_token method that returns proper type."""
-                return AccessToken("test-token", 12345)
-
-        config = AzureAISearchConfig(
-            name="test", endpoint="https://endpoint", index_name="index", credential=TestTokenCredential()
-        )
-        result = config.model_dump()
-        assert result["credential"] == {"type": "TokenCredential"}
-    else:
-        pytest.skip("Skipping TokenCredential test - Azure SDK not available")
+    search_results_iterator = AsyncIterator(mock_search_response)
+    mock_client_instance.search = MagicMock(return_value=search_results_iterator)
+
+    patch_target = "autogen_ext.tools.azure._ai_search.SearchClient"
+    patcher = patch(patch_target, return_value=mock_client_instance)
+
+    patcher.start()
+    yield mock_client_instance
+    patcher.stop()

python/packages/autogen-ext/tests/tools/azure/test_ai_search_config.py

@@ -0,0 +1,297 @@
+from typing import Any, Dict, cast
+
+import pytest
+from autogen_ext.tools.azure._config import AzureAISearchConfig, QueryTypeLiteral
+from azure.core.credentials import AzureKeyCredential
+from pydantic import ValidationError
+
+from tests.tools.azure.conftest import azure_sdk_available
+
+skip_if_no_azure_sdk = pytest.mark.skipif(
+    not azure_sdk_available, reason="Azure SDK components (azure-search-documents, azure-identity) not available"
+)
+
+# =====================================
+# Basic Configuration Tests
+# =====================================
+
+
+def test_basic_config_creation() -> None:
+    """Test that a basic valid configuration can be created."""
+    config = AzureAISearchConfig(
+        name="test_tool",
+        endpoint="https://test-search.search.windows.net",
+        index_name="test-index",
+        credential=AzureKeyCredential("test-key"),
+    )
+
+    assert config.name == "test_tool"
+    assert config.endpoint == "https://test-search.search.windows.net"
+    assert config.index_name == "test-index"
+    assert isinstance(config.credential, AzureKeyCredential)
+    assert config.query_type == "simple"  # default value
+
+
+def test_endpoint_validation() -> None:
+    """Test that endpoint validation works correctly."""
+    valid_endpoints = ["https://test.search.windows.net", "http://localhost:8080"]
+
+    for endpoint in valid_endpoints:
+        config = AzureAISearchConfig(
+            name="test_tool",
+            endpoint=endpoint,
+            index_name="test-index",
+            credential=AzureKeyCredential("test-key"),
+        )
+        assert config.endpoint == endpoint
+
+    invalid_endpoints = [
+        "test.search.windows.net",
+        "ftp://test.search.windows.net",
+        "",
+    ]
+
+    for endpoint in invalid_endpoints:
+        with pytest.raises(ValidationError) as exc:
+            AzureAISearchConfig(
+                name="test_tool",
+                endpoint=endpoint,
+                index_name="test-index",
+                credential=AzureKeyCredential("test-key"),
+            )
+        assert "endpoint must be a valid URL" in str(exc.value)
+
+
+def test_top_validation() -> None:
+    """Test validation of top parameter."""
+    valid_tops = [1, 5, 10, 100]
+
+    for top in valid_tops:
+        config = AzureAISearchConfig(
+            name="test_tool",
+            endpoint="https://test.search.windows.net",
+            index_name="test-index",
+            credential=AzureKeyCredential("test-key"),
+            top=top,
+        )
+        assert config.top == top
+
+    invalid_tops = [0, -1, -10]
+
+    for top in invalid_tops:
+        with pytest.raises(ValidationError) as exc:
+            AzureAISearchConfig(
+                name="test_tool",
+                endpoint="https://test.search.windows.net",
+                index_name="test-index",
+                credential=AzureKeyCredential("test-key"),
+                top=top,
+            )
+        assert "top must be a positive integer" in str(exc.value)
+
+
+# =====================================
+# Query Type Tests
+# =====================================
+
+
+def test_query_type_normalization() -> None:
+    """Test that query_type normalization works correctly."""
+    standard_query_types = {
+        "simple": "simple",
+        "full": "full",
+        "semantic": "semantic",
+        "vector": "vector",
+    }
+
+    for input_type, expected_type in standard_query_types.items():
+        config_args: Dict[str, Any] = {
+            "name": "test_tool",
+            "endpoint": "https://test.search.windows.net",
+            "index_name": "test-index",
+            "credential": AzureKeyCredential("test-key"),
+            "query_type": cast(QueryTypeLiteral, input_type),
+        }
+
+        if input_type == "semantic":
+            config_args["semantic_config_name"] = "my-semantic-config"
+        elif input_type == "vector":
+            config_args["vector_fields"] = ["content_vector"]
+
+        config = AzureAISearchConfig(**config_args)
+        assert config.query_type == expected_type
+
+    with pytest.raises(ValidationError) as exc:
+        AzureAISearchConfig(
+            name="test_tool",
+            endpoint="https://test.search.windows.net",
+            index_name="test-index",
+            credential=AzureKeyCredential("test-key"),
+            query_type=cast(Any, "invalid_type"),
+        )
+    assert "Input should be" in str(exc.value)
+
+
+def test_semantic_config_validation() -> None:
+    """Test validation of semantic configuration."""
+    config = AzureAISearchConfig(
+        name="test_tool",
+        endpoint="https://test.search.windows.net",
+        index_name="test-index",
+        credential=AzureKeyCredential("test-key"),
+        query_type=cast(QueryTypeLiteral, "semantic"),
+        semantic_config_name="my-semantic-config",
+    )
+    assert config.query_type == "semantic"
+    assert config.semantic_config_name == "my-semantic-config"
+
+    with pytest.raises(ValidationError) as exc:
+        AzureAISearchConfig(
+            name="test_tool",
+            endpoint="https://test.search.windows.net",
+            index_name="test-index",
+            credential=AzureKeyCredential("test-key"),
+            query_type=cast(QueryTypeLiteral, "semantic"),
+        )
+    assert "semantic_config_name must be provided" in str(exc.value)
+
+
+def test_vector_fields_validation() -> None:
+    """Test validation of vector fields for vector search."""
+    config = AzureAISearchConfig(
+        name="test_tool",
+        endpoint="https://test.search.windows.net",
+        index_name="test-index",
+        credential=AzureKeyCredential("test-key"),
+        query_type=cast(QueryTypeLiteral, "vector"),
+        vector_fields=["content_vector"],
+    )
+    assert config.query_type == "vector"
+    assert config.vector_fields == ["content_vector"]
+
+
+# =====================================
+# Embedding Configuration Tests
+# =====================================
+
+
+def test_azure_openai_endpoint_validation() -> None:
+    """Test validation of Azure OpenAI endpoint for client-side embeddings."""
+    config = AzureAISearchConfig(
+        name="test_tool",
+        endpoint="https://test.search.windows.net",
+        index_name="test-index",
+        credential=AzureKeyCredential("test-key"),
+        embedding_provider="azure_openai",
+        embedding_model="text-embedding-ada-002",
+        openai_endpoint="https://test.openai.azure.com",
+    )
+    assert config.embedding_provider == "azure_openai"
+    assert config.embedding_model == "text-embedding-ada-002"
+    assert config.openai_endpoint == "https://test.openai.azure.com"
+
+    with pytest.raises(ValidationError) as exc:
+        AzureAISearchConfig(
+            name="test_tool",
+            endpoint="https://test.search.windows.net",
+            index_name="test-index",
+            credential=AzureKeyCredential("test-key"),
+            embedding_provider="azure_openai",
+            embedding_model="text-embedding-ada-002",
+        )
+    assert "openai_endpoint must be provided for azure_openai" in str(exc.value)
+
+    config = AzureAISearchConfig(
+        name="test_tool",
+        endpoint="https://test.search.windows.net",
+        index_name="test-index",
+        credential=AzureKeyCredential("test-key"),
+        embedding_provider="openai",
+        embedding_model="text-embedding-ada-002",
+    )
+    assert config.embedding_provider == "openai"
+    assert config.embedding_model == "text-embedding-ada-002"
+    assert config.openai_endpoint is None
+
+
+# =====================================
+# Credential and Serialization Tests
+# =====================================
+
+
+def test_credential_validation() -> None:
+    """Test credential validation scenarios."""
+    config = AzureAISearchConfig(
+        name="test_tool",
+        endpoint="https://test.search.windows.net",
+        index_name="test-index",
+        credential=AzureKeyCredential("test-key"),
+    )
+    assert isinstance(config.credential, AzureKeyCredential)
+    assert config.credential.key == "test-key"
+
+    if azure_sdk_available:
+        from azure.core.credentials import AccessToken
+        from azure.core.credentials_async import AsyncTokenCredential
+
+        class TestTokenCredential(AsyncTokenCredential):
+            async def get_token(self, *scopes: str, **kwargs: Any) -> AccessToken:
+                return AccessToken("test-token", 12345)
+
+            async def close(self) -> None:
+                pass
+
+            async def __aenter__(self) -> "TestTokenCredential":
+                return self
+
+            async def __aexit__(self, *args: Any) -> None:
+                await self.close()
+
+        config = AzureAISearchConfig(
+            name="test",
+            endpoint="https://endpoint",
+            index_name="index",
+            credential=TestTokenCredential(),
+        )
+        assert isinstance(config.credential, AsyncTokenCredential)
+
+
+def test_model_dump_scenarios() -> None:
+    """Test all model_dump scenarios to ensure full code coverage."""
+    config = AzureAISearchConfig(
+        name="test",
+        endpoint="https://endpoint",
+        index_name="index",
+        credential=AzureKeyCredential("key"),
+    )
+    result = config.model_dump()
+    assert isinstance(result["credential"], AzureKeyCredential)
+    assert result["credential"].key == "key"
+
+    if azure_sdk_available:
+        from azure.core.credentials import AccessToken
+        from azure.core.credentials_async import AsyncTokenCredential
+
+        class TestTokenCredential(AsyncTokenCredential):
+            async def get_token(self, *scopes: str, **kwargs: Any) -> AccessToken:
+                return AccessToken("test-token", 12345)
+
+            async def close(self) -> None:
+                pass
+
+            async def __aenter__(self) -> "TestTokenCredential":
+                return self
+
+            async def __aexit__(self, *args: Any) -> None:
+                await self.close()
+
+        config = AzureAISearchConfig(
+            name="test",
+            endpoint="https://endpoint",
+            index_name="index",
+            credential=TestTokenCredential(),
+        )
+        result = config.model_dump()
+        assert isinstance(result["credential"], AsyncTokenCredential)
+    else:
+        pytest.skip("Skipping TokenCredential test - Azure SDK not available")

python/packages/autogen-ext/tests/tools/http/test_http_tool.py

@@ -176,7 +176,7 @@ async def test_invalid_request(test_config: ComponentModel, test_server: None) -
     config.config["host"] = "fake"
     tool = HttpTool.load_component(config)
 
-    with pytest.raises(httpx.ConnectError):
+    with pytest.raises((httpx.ConnectError, httpx.ConnectTimeout)):
         await tool.run_json({"query": "test query", "value": 42}, CancellationToken())
 
 

python/packages/autogen-ext/tests/tools/test_mcp_tools.py

@@ -1,12 +1,17 @@
+import asyncio
 import logging
 import os
+import threading
+from typing import cast
 from unittest.mock import AsyncMock, MagicMock
 
 import pytest
+from _pytest.logging import LogCaptureFixture  # type: ignore[import]
 from autogen_core import CancellationToken
 from autogen_core.tools import Workbench
 from autogen_core.utils import schema_to_pydantic_model
 from autogen_ext.tools.mcp import (
+    McpSessionActor,
     McpWorkbench,
     SseMcpToolAdapter,
     SseServerParams,
@@ -94,6 +99,14 @@ def cancellation_token() -> CancellationToken:
     return CancellationToken()
 
 
+@pytest.fixture
+def mock_error_tool_response() -> MagicMock:
+    response = MagicMock()
+    response.isError = True
+    response.content = [TextContent(text="error output", type="text")]
+    return response
+
+
 def test_adapter_config_serialization(sample_tool: Tool, sample_server_params: StdioServerParams) -> None:
     """Test that adapter can be saved to and loaded from config."""
     original_adapter = StdioMcpToolAdapter(server_params=sample_server_params, tool=sample_tool)
@@ -594,4 +607,165 @@ async def test_lazy_init_and_finalize_cleanup() -> None:
     assert workbench._actor is not None  # type: ignore[reportPrivateUsage]
     assert workbench._actor._active is True  # type: ignore[reportPrivateUsage]
 
+    actor = workbench._actor  # type: ignore[reportPrivateUsage]
     del workbench
+    await asyncio.sleep(0.1)
+    assert actor._active is False
+
+
+@pytest.mark.asyncio
+async def test_del_to_new_event_loop_when_get_event_loop_fails() -> None:
+    params = StdioServerParams(
+        command="npx",
+        args=[
+            "-y",
+            "@modelcontextprotocol/server-filesystem",
+            ".",
+        ],
+        read_timeout_seconds=60,
+    )
+    workbench = McpWorkbench(server_params=params)
+
+    await workbench.list_tools()
+    assert workbench._actor is not None  # type: ignore[reportPrivateUsage]
+    assert workbench._actor._active is True  # type: ignore[reportPrivateUsage]
+
+    actor = workbench._actor  # type: ignore[reportPrivateUsage]
+
+    def cleanup() -> None:
+        nonlocal workbench
+        del workbench
+
+    t = threading.Thread(target=cleanup)
+    t.start()
+    t.join()
+
+    await asyncio.sleep(0.1)
+    assert actor._active is False  # type: ignore[reportPrivateUsage]
+
+
+def test_del_raises_when_loop_closed() -> None:
+    loop = asyncio.new_event_loop()
+    asyncio.set_event_loop(loop)
+
+    params = StdioServerParams(command="echo", args=["ok"])
+    workbench = McpWorkbench(server_params=params)
+
+    workbench._actor_loop = loop  # type: ignore[reportPrivateUsage]
+    workbench._actor = cast(McpSessionActor, object())  # type: ignore[reportPrivateUsage]
+
+    loop.close()
+
+    with pytest.warns(RuntimeWarning, match="loop is closed or not running"):
+        del workbench
+
+
+def test_mcp_tool_adapter_normalize_payload(sample_tool: Tool, sample_server_params: StdioServerParams) -> None:
+    """Test the _normalize_payload_to_content_list method of McpToolAdapter."""
+    adapter = StdioMcpToolAdapter(server_params=sample_server_params, tool=sample_tool)
+
+    # Case 1: Payload is already a list of valid content items
+    valid_content_list: list[TextContent | ImageContent | EmbeddedResource] = [
+        TextContent(text="hello", type="text"),
+        ImageContent(data="base64data", mimeType="image/png", type="image"),
+        EmbeddedResource(
+            type="resource",
+            resource=TextResourceContents(text="embedded text", uri=AnyUrl(url="http://example.com/resource")),
+        ),
+    ]
+    assert adapter._normalize_payload_to_content_list(valid_content_list) == valid_content_list  # type: ignore[reportPrivateUsage]
+
+    # Case 2: Payload is a single TextContent
+    single_text_content = TextContent(text="single text", type="text")
+    assert adapter._normalize_payload_to_content_list(single_text_content) == [single_text_content]  # type: ignore[reportPrivateUsage, arg-type]
+
+    # Case 3: Payload is a single ImageContent
+    single_image_content = ImageContent(data="imagedata", mimeType="image/jpeg", type="image")
+    assert adapter._normalize_payload_to_content_list(single_image_content) == [single_image_content]  # type: ignore[reportPrivateUsage, arg-type]
+
+    # Case 4: Payload is a single EmbeddedResource
+    single_embedded_resource = EmbeddedResource(
+        type="resource",
+        resource=TextResourceContents(text="other embedded", uri=AnyUrl(url="http://example.com/other")),
+    )
+    assert adapter._normalize_payload_to_content_list(single_embedded_resource) == [single_embedded_resource]  # type: ignore[reportPrivateUsage, arg-type]
+
+    # Case 5: Payload is a string
+    string_payload = "This is a string payload."
+    expected_from_string = [TextContent(text=string_payload, type="text")]
+    assert adapter._normalize_payload_to_content_list(string_payload) == expected_from_string  # type: ignore[reportPrivateUsage, arg-type]
+
+    # Case 6: Payload is an integer
+    int_payload = 12345
+    expected_from_int = [TextContent(text=str(int_payload), type="text")]
+    assert adapter._normalize_payload_to_content_list(int_payload) == expected_from_int  # type: ignore[reportPrivateUsage, arg-type]
+
+    # Case 7: Payload is a dictionary
+    dict_payload = {"key": "value", "number": 42}
+    expected_from_dict = [TextContent(text=str(dict_payload), type="text")]
+    assert adapter._normalize_payload_to_content_list(dict_payload) == expected_from_dict  # type: ignore[reportPrivateUsage, arg-type]
+
+    # Case 8: Payload is an empty list (should still be a list of valid items, so returns as is)
+    empty_list_payload: list[TextContent | ImageContent | EmbeddedResource] = []
+    assert adapter._normalize_payload_to_content_list(empty_list_payload) == empty_list_payload  # type: ignore[reportPrivateUsage]
+
+    # Case 9: Payload is None (should be stringified)
+    none_payload = None
+    expected_from_none = [TextContent(text=str(none_payload), type="text")]
+    assert adapter._normalize_payload_to_content_list(none_payload) == expected_from_none  # type: ignore[reportPrivateUsage, arg-type]
+
+
+@pytest.mark.asyncio
+async def test_mcp_tool_adapter_run_error(
+    sample_tool: Tool,
+    sample_server_params: StdioServerParams,
+    mock_session: AsyncMock,
+    mock_error_tool_response: MagicMock,
+    cancellation_token: CancellationToken,
+) -> None:
+    """Test McpToolAdapter._run when tool returns an error."""
+    adapter = StdioMcpToolAdapter(server_params=sample_server_params, tool=sample_tool, session=mock_session)
+    mock_session.call_tool.return_value = mock_error_tool_response
+
+    args = {"test_param": "test_value"}
+    with pytest.raises(Exception) as excinfo:
+        await adapter._run(args=args, cancellation_token=cancellation_token, session=mock_session)  # type: ignore[reportPrivateUsage]
+
+    mock_session.call_tool.assert_called_once_with(name=sample_tool.name, arguments=args)
+    assert adapter.return_value_as_string([TextContent(text="error output", type="text")]) in str(excinfo.value)
+
+
+@pytest.mark.asyncio
+async def test_mcp_tool_adapter_run_cancelled_before_call(
+    sample_tool: Tool,
+    sample_server_params: StdioServerParams,
+    mock_session: AsyncMock,
+    cancellation_token: CancellationToken,
+) -> None:
+    """Test McpToolAdapter._run when operation is cancelled before tool call."""
+    adapter = StdioMcpToolAdapter(server_params=sample_server_params, tool=sample_tool, session=mock_session)
+    cancellation_token.cancel()  # Cancel before the call
+
+    args = {"test_param": "test_value"}
+    with pytest.raises(asyncio.CancelledError):
+        await adapter._run(args=args, cancellation_token=cancellation_token, session=mock_session)  # type: ignore[reportPrivateUsage]
+
+    mock_session.call_tool.assert_not_called()
+
+
+@pytest.mark.asyncio
+async def test_mcp_tool_adapter_run_cancelled_during_call(
+    sample_tool: Tool,
+    sample_server_params: StdioServerParams,
+    mock_session: AsyncMock,
+    cancellation_token: CancellationToken,
+) -> None:
+    """Test McpToolAdapter._run when operation is cancelled during tool call."""
+    adapter = StdioMcpToolAdapter(server_params=sample_server_params, tool=sample_tool, session=mock_session)
+    mock_session.call_tool.side_effect = asyncio.CancelledError("Tool call cancelled")
+
+    args = {"test_param": "test_value"}
+    with pytest.raises(asyncio.CancelledError):
+        await adapter._run(args=args, cancellation_token=cancellation_token, session=mock_session)  # type: ignore[reportPrivateUsage]
+
+    mock_session.call_tool.assert_called_once_with(name=sample_tool.name, arguments=args)

python/uv.lock

@@ -455,7 +455,7 @@ wheels = [
 
 [[package]]
 name = "autogen-agentchat"
-version = "0.5.6"
+version = "0.5.7"
 source = { editable = "packages/autogen-agentchat" }
 dependencies = [
     { name = "autogen-core" },
@@ -466,7 +466,7 @@ requires-dist = [{ name = "autogen-core", editable = "packages/autogen-core" }]
 
 [[package]]
 name = "autogen-core"
-version = "0.5.6"
+version = "0.5.7"
 source = { editable = "packages/autogen-core" }
 dependencies = [
     { name = "jsonref" },
@@ -585,7 +585,7 @@ dev = [
 
 [[package]]
 name = "autogen-ext"
-version = "0.5.6"
+version = "0.5.7"
 source = { editable = "packages/autogen-ext" }
 dependencies = [
     { name = "autogen-core" },