From c1435ac2010002fba57c60cbb8227e27c94bc92f Mon Sep 17 00:00:00 2001
From: Giles Odigwe <79032838+giles17@users.noreply.github.com>
Date: Wed, 25 Mar 2026 19:08:47 -0700
Subject: [PATCH 1/4] Python: Fix A2AAgent to surface message content from
 in-progress TaskStatusUpdateEvents (#4798)

* Fix A2AAgent dropping message content from in-progress TaskStatusUpdateEvents (#4783)

_updates_from_task() returned [] for working-state tasks when
background=False, silently discarding all intermediate message content
from task.status.message. Now extracts and yields message parts from
in-progress status updates during streaming.

Also fixed MockA2AClient.send_message to yield all queued responses
(enabling multi-event streaming tests) and added text parameter to
add_in_progress_task_response for tests that need status messages.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Fix: gate intermediate status updates behind emit_intermediate flag and add missing test coverage

- Add emit_intermediate parameter to _updates_from_task and _map_a2a_stream
- Thread stream flag from run() so only streaming callers see intermediate updates
- Add IN_PROGRESS_TASK_STATES guard to emit_intermediate condition
- Add role parameter to test helper add_in_progress_task_response
- Add clarifying comment on MockA2AClient.send_message batch semantics
- Add tests for user role mapping, background precedence, non-streaming behavior,
  terminal task with no artifacts, and empty parts edge case

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

---------

Co-authored-by: Copilot <copilot@github.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
---
 .../a2a/agent_framework_a2a/_agent.py         |  55 +++++-
 python/packages/a2a/tests/test_a2a_agent.py   | 157 +++++++++++++++++-
 2 files changed, 202 insertions(+), 10 deletions(-)

diff --git a/python/packages/a2a/agent_framework_a2a/_agent.py b/python/packages/a2a/agent_framework_a2a/_agent.py
index d016caae7c..4b6d9cc19b 100644
--- a/python/packages/a2a/agent_framework_a2a/_agent.py
+++ b/python/packages/a2a/agent_framework_a2a/_agent.py
@@ -313,6 +313,7 @@ class A2AAgent(AgentTelemetryLayer, BaseAgent):
             self._map_a2a_stream(
                 a2a_stream,
                 background=background,
+                emit_intermediate=stream,
                 session=provider_session,
                 session_context=session_context,
             ),
@@ -327,6 +328,7 @@ class A2AAgent(AgentTelemetryLayer, BaseAgent):
         a2a_stream: AsyncIterable[A2AStreamItem],
         *,
         background: bool = False,
+        emit_intermediate: bool = False,
         session: AgentSession | None = None,
         session_context: SessionContext | None = None,
     ) -> AsyncIterable[AgentResponseUpdate]:
@@ -339,6 +341,10 @@ class A2AAgent(AgentTelemetryLayer, BaseAgent):
             background: When False, in-progress task updates are silently
                 consumed (the stream keeps iterating until a terminal state).
                 When True, they are yielded with a continuation token.
+            emit_intermediate: When True, in-progress status updates that
+                carry message content are yielded to the caller.  Typically
+                set for streaming callers so non-streaming consumers only
+                receive terminal task outputs.
             session: The agent session for context providers.
             session_context: The session context for context providers.
         """
@@ -373,7 +379,11 @@ class A2AAgent(AgentTelemetryLayer, BaseAgent):
                 yield update
             elif isinstance(item, tuple) and len(item) == 2 and isinstance(item[0], Task):
                 task, _update_event = item
-                for update in self._updates_from_task(task, background=background):
+                for update in self._updates_from_task(
+                    task,
+                    background=background,
+                    emit_intermediate=emit_intermediate,
+                ):
                     all_updates.append(update)
                     yield update
             else:
@@ -389,15 +399,26 @@ class A2AAgent(AgentTelemetryLayer, BaseAgent):
     # Task helpers
     # ------------------------------------------------------------------
 
-    def _updates_from_task(self, task: Task, *, background: bool = False) -> list[AgentResponseUpdate]:
+    def _updates_from_task(
+        self,
+        task: Task,
+        *,
+        background: bool = False,
+        emit_intermediate: bool = False,
+    ) -> list[AgentResponseUpdate]:
         """Convert an A2A Task into AgentResponseUpdate(s).
 
         Terminal tasks produce updates from their artifacts/history.
-        In-progress tasks produce a continuation token update only when
-        ``background=True``; otherwise they are silently skipped so the
-        caller keeps consuming the stream until completion.
+        In-progress tasks produce a continuation token update when
+        ``background=True``.  When ``emit_intermediate=True`` (typically
+        set for streaming callers), any message content attached to an
+        in-progress status update is surfaced; otherwise the update is
+        silently skipped so the caller keeps consuming the stream until
+        completion.
         """
-        if task.status.state in TERMINAL_TASK_STATES:
+        status = task.status
+
+        if status.state in TERMINAL_TASK_STATES:
             task_messages = self._parse_messages_from_task(task)
             if task_messages:
                 return [
@@ -412,7 +433,7 @@ class A2AAgent(AgentTelemetryLayer, BaseAgent):
                 ]
             return [AgentResponseUpdate(contents=[], role="assistant", response_id=task.id, raw_representation=task)]
 
-        if background and task.status.state in IN_PROGRESS_TASK_STATES:
+        if background and status.state in IN_PROGRESS_TASK_STATES:
             token = self._build_continuation_token(task)
             return [
                 AgentResponseUpdate(
@@ -424,6 +445,26 @@ class A2AAgent(AgentTelemetryLayer, BaseAgent):
                 )
             ]
 
+        # Surface message content from in-progress status updates (e.g. working state)
+        # Only emitted when the caller opts in (streaming), so non-streaming
+        # consumers keep receiving only terminal task outputs.
+        if (
+            emit_intermediate
+            and status.state in IN_PROGRESS_TASK_STATES
+            and status.message is not None
+            and status.message.parts
+        ):
+            contents = self._parse_contents_from_a2a(status.message.parts)
+            if contents:
+                return [
+                    AgentResponseUpdate(
+                        contents=contents,
+                        role="assistant" if status.message.role == A2ARole.agent else "user",
+                        response_id=task.id,
+                        raw_representation=task,
+                    )
+                ]
+
         return []
 
     @staticmethod
diff --git a/python/packages/a2a/tests/test_a2a_agent.py b/python/packages/a2a/tests/test_a2a_agent.py
index b8633938fc..0d81179cd1 100644
--- a/python/packages/a2a/tests/test_a2a_agent.py
+++ b/python/packages/a2a/tests/test_a2a_agent.py
@@ -91,9 +91,18 @@ class MockA2AClient:
         task_id: str,
         context_id: str = "test-context",
         state: TaskState = TaskState.working,
+        text: str | None = None,
+        role: A2ARole = A2ARole.agent,
     ) -> None:
         """Add a mock in-progress Task response (non-terminal)."""
-        status = TaskStatus(state=state, message=None)
+        message = None
+        if text is not None:
+            message = A2AMessage(
+                message_id=str(uuid4()),
+                role=role,
+                parts=[Part(root=TextPart(text=text))],
+            )
+        status = TaskStatus(state=state, message=message)
         task = Task(id=task_id, context_id=context_id, status=status)
         client_event = (task, None)
         self.responses.append(client_event)
@@ -102,9 +111,10 @@ class MockA2AClient:
         """Mock send_message method that yields responses."""
         self.call_count += 1
 
-        if self.responses:
-            response = self.responses.pop(0)
+        # All queued responses are delivered as a single streaming batch per call.
+        for response in self.responses:
             yield response
+        self.responses.clear()
 
     async def resubscribe(self, request: Any) -> AsyncIterator[Any]:
         """Mock resubscribe method that yields responses."""
@@ -1039,3 +1049,144 @@ async def test_run_with_continuation_token_does_not_require_messages(mock_a2a_cl
 
 
 # endregion
+
+# region Streaming with in-progress message content
+
+
+async def test_streaming_working_updates_yield_message_content(
+    a2a_agent: A2AAgent, mock_a2a_client: MockA2AClient
+) -> None:
+    """Test that streaming working updates with status.message yield content."""
+    mock_a2a_client.add_in_progress_task_response("task-w", context_id="ctx-w", text="Processing step 1...")
+    mock_a2a_client.add_in_progress_task_response("task-w", context_id="ctx-w", text="Processing step 2...")
+    mock_a2a_client.add_task_response("task-w", [{"id": "art-w", "content": "Final result"}])
+
+    updates: list[AgentResponseUpdate] = []
+    async for update in a2a_agent.run("Hello", stream=True):
+        updates.append(update)
+
+    assert len(updates) == 3
+    assert updates[0].contents[0].text == "Processing step 1..."
+    assert updates[1].contents[0].text == "Processing step 2..."
+    assert updates[2].contents[0].text == "Final result"
+
+
+async def test_streaming_single_working_update_with_message(
+    a2a_agent: A2AAgent, mock_a2a_client: MockA2AClient
+) -> None:
+    """Test that a single working update with message content is not dropped."""
+    mock_a2a_client.add_in_progress_task_response("task-s", context_id="ctx-s", text="Thinking...")
+    mock_a2a_client.add_task_response("task-s", [{"id": "art-s", "content": "Done"}])
+
+    updates: list[AgentResponseUpdate] = []
+    async for update in a2a_agent.run("Hello", stream=True):
+        updates.append(update)
+
+    assert len(updates) == 2
+    assert updates[0].contents[0].text == "Thinking..."
+    assert updates[0].role == "assistant"
+    assert updates[1].contents[0].text == "Done"
+
+
+async def test_streaming_working_update_without_message_is_skipped(
+    a2a_agent: A2AAgent, mock_a2a_client: MockA2AClient
+) -> None:
+    """Test that working updates without status.message are still silently skipped."""
+    mock_a2a_client.add_in_progress_task_response("task-n", context_id="ctx-n")
+    mock_a2a_client.add_task_response("task-n", [{"id": "art-n", "content": "Result"}])
+
+    updates: list[AgentResponseUpdate] = []
+    async for update in a2a_agent.run("Hello", stream=True):
+        updates.append(update)
+
+    assert len(updates) == 1
+    assert updates[0].contents[0].text == "Result"
+
+
+async def test_streaming_working_update_user_role_mapping(a2a_agent: A2AAgent, mock_a2a_client: MockA2AClient) -> None:
+    """Test that A2ARole.user in status message maps to role='user'."""
+    mock_a2a_client.add_in_progress_task_response("task-u", context_id="ctx-u", text="User echo", role=A2ARole.user)
+    mock_a2a_client.add_task_response("task-u", [{"id": "art-u", "content": "Done"}])
+
+    updates: list[AgentResponseUpdate] = []
+    async for update in a2a_agent.run("Hello", stream=True):
+        updates.append(update)
+
+    assert len(updates) == 2
+    assert updates[0].contents[0].text == "User echo"
+    assert updates[0].role == "user"
+
+
+async def test_background_with_status_message_yields_continuation_token(
+    a2a_agent: A2AAgent, mock_a2a_client: MockA2AClient
+) -> None:
+    """Test that background=True takes precedence over status message content."""
+    mock_a2a_client.add_in_progress_task_response("task-bg", context_id="ctx-bg", text="Should be ignored")
+
+    updates: list[AgentResponseUpdate] = []
+    async for update in a2a_agent.run("Hello", stream=True, background=True):
+        updates.append(update)
+
+    assert len(updates) == 1
+    assert updates[0].continuation_token is not None
+    assert updates[0].continuation_token["task_id"] == "task-bg"
+    assert updates[0].contents == []
+
+
+async def test_non_streaming_does_not_surface_intermediate_messages(
+    a2a_agent: A2AAgent, mock_a2a_client: MockA2AClient
+) -> None:
+    """Test that run(stream=False) does not include intermediate status messages."""
+    mock_a2a_client.add_in_progress_task_response("task-ns", context_id="ctx-ns", text="Intermediate")
+    mock_a2a_client.add_task_response("task-ns", [{"id": "art-ns", "content": "Final"}])
+
+    response = await a2a_agent.run("Hello")
+
+    assert len(response.messages) == 1
+    assert response.messages[0].text == "Final"
+
+
+async def test_terminal_no_artifacts_after_working_with_content(
+    a2a_agent: A2AAgent, mock_a2a_client: MockA2AClient
+) -> None:
+    """Test that a terminal task with no artifacts after working-state messages does not re-emit the working content."""
+    mock_a2a_client.add_in_progress_task_response("task-t", context_id="ctx-t", text="Working on it...")
+    # Terminal task with no artifacts and no history
+    status = TaskStatus(state=TaskState.completed, message=None)
+    task = Task(id="task-t", context_id="ctx-t", status=status)
+    mock_a2a_client.responses.append((task, None))
+
+    updates: list[AgentResponseUpdate] = []
+    async for update in a2a_agent.run("Hello", stream=True):
+        updates.append(update)
+
+    assert len(updates) == 2
+    assert updates[0].contents[0].text == "Working on it..."
+    # Terminal task with no artifacts yields an empty-contents update
+    assert updates[1].contents == []
+
+
+async def test_streaming_working_update_with_empty_parts_is_skipped(
+    a2a_agent: A2AAgent, mock_a2a_client: MockA2AClient
+) -> None:
+    """Test that a working update with status.message but empty parts list is skipped."""
+    # Construct a message with an empty parts list (distinct from message=None)
+    message = A2AMessage(
+        message_id=str(uuid4()),
+        role=A2ARole.agent,
+        parts=[],
+    )
+    status = TaskStatus(state=TaskState.working, message=message)
+    task = Task(id="task-ep", context_id="ctx-ep", status=status)
+    mock_a2a_client.responses.append((task, None))
+    mock_a2a_client.add_task_response("task-ep", [{"id": "art-ep", "content": "Result"}])
+
+    updates: list[AgentResponseUpdate] = []
+    async for update in a2a_agent.run("Hello", stream=True):
+        updates.append(update)
+
+    assert len(updates) == 1
+    assert updates[0].contents[0].text == "Result"
+
+
+# endregion

From dc27740f1a752360426b8ef29efd8ec259325e1b Mon Sep 17 00:00:00 2001
From: Giles Odigwe <79032838+giles17@users.noreply.github.com>
Date: Wed, 25 Mar 2026 19:09:22 -0700
Subject: [PATCH 2/4] Python: Fix streaming path to emit mcp_server_tool_result
 on output_item.done instead of output_item.added (#4821)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* Fix streaming path to deliver mcp_server_tool_result content (#4814)

Remove premature mcp_server_tool_result emission from the
response.output_item.added/mcp_call handler — at that point the MCP
server has not yet responded and output is always None.

Add a handler for response.mcp_call.completed that emits
mcp_server_tool_result with the actual tool output, matching the
non-streaming path behavior.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Fix streaming path to deliver mcp_server_tool_result content (#4814)

Stop eagerly emitting mcp_server_tool_result on response.output_item.added
(when output is always None). Instead, handle response.output_item.done for
mcp_call items, which carries the full McpCall with populated output.

This matches the non-streaming path which guards with 'if item.output is not
None' before emitting the result.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Fix test docstring to match actual implementation event name

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Address review: call_id fallback and raw_representation consistency (#4814)

- Add call_id fallback in response.output_item.done mcp_call handler to
  match the output_item.added handler pattern
- Use done_item instead of event for raw_representation to keep
  consistent with other output_item branches and non-streaming path
- Add test for call_id fallback when id attribute is missing
- Add raw_representation assertions to existing done handler tests

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Address review: call_id fallback for non-streaming path and test coverage (#4814)

- Apply defensive call_id fallback (getattr with id/call_id/empty) to
  non-streaming mcp_call path for consistency with streaming path
- Add raw_representation assertion to call_id fallback test
- Add test for empty-string fallback when neither id nor call_id exist

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

---------

Co-authored-by: Copilot <copilot@github.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
---
 .../agent_framework_openai/_chat_client.py    |  39 +++---
 .../tests/openai/test_openai_chat_client.py   | 129 ++++++++++++++++--
 2 files changed, 136 insertions(+), 32 deletions(-)

diff --git a/python/packages/openai/agent_framework_openai/_chat_client.py b/python/packages/openai/agent_framework_openai/_chat_client.py
index e1e40c87c0..86af86895e 100644
--- a/python/packages/openai/agent_framework_openai/_chat_client.py
+++ b/python/packages/openai/agent_framework_openai/_chat_client.py
@@ -1728,7 +1728,7 @@ class RawOpenAIChatClient(  # type: ignore[misc]
                         )
                     )
                 case "mcp_call":
-                    call_id = item.id
+                    call_id = getattr(item, "id", None) or getattr(item, "call_id", None) or ""
                     contents.append(
                         Content.from_mcp_server_tool_call(
                             call_id=call_id,
@@ -2118,27 +2118,7 @@ class RawOpenAIChatClient(  # type: ignore[misc]
                                 raw_representation=event_item,
                             )
                         )
-                        result_output = (
-                            getattr(event_item, "result", None)
-                            or getattr(event_item, "output", None)
-                            or getattr(event_item, "outputs", None)
-                        )
-                        parsed_output: list[Content] | None = None
-                        if result_output:
-                            normalized = (  # pyright: ignore[reportUnknownVariableType]
-                                result_output
-                                if isinstance(result_output, Sequence)
-                                and not isinstance(result_output, (str, bytes, MutableMapping))
-                                else [result_output]
-                            )
-                            parsed_output = [Content.from_dict(output_item) for output_item in normalized]  # pyright: ignore[reportArgumentType,reportUnknownVariableType]
-                        contents.append(
-                            Content.from_mcp_server_tool_result(
-                                call_id=call_id,
-                                output=parsed_output,
-                                raw_representation=event_item,
-                            )
-                        )
+                        # Result deferred to response.output_item.done
                     case "code_interpreter_call":  # ResponseOutputCodeInterpreterCall
                         call_id = getattr(event_item, "call_id", None) or getattr(event_item, "id", None)
                         outputs: list[Content] = []
@@ -2408,6 +2388,21 @@ class RawOpenAIChatClient(  # type: ignore[misc]
                         )
                 else:
                     logger.debug("Unparsed annotation type in streaming: %s", ann_type)
+            case "response.output_item.done":
+                done_item = event.item
+                if getattr(done_item, "type", None) == "mcp_call":
+                    call_id = getattr(done_item, "id", None) or getattr(done_item, "call_id", None) or ""
+                    output_text = getattr(done_item, "output", None)
+                    parsed_output: list[Content] | None = (
+                        [Content.from_text(text=output_text)] if isinstance(output_text, str) else None
+                    )
+                    contents.append(
+                        Content.from_mcp_server_tool_result(
+                            call_id=call_id,
+                            output=parsed_output,
+                            raw_representation=done_item,
+                        )
+                    )
             case _:
                 logger.debug("Unparsed event of type: %s: %s", event.type, event)
 
diff --git a/python/packages/openai/tests/openai/test_openai_chat_client.py b/python/packages/openai/tests/openai/test_openai_chat_client.py
index a29fbf443a..897fe5f913 100644
--- a/python/packages/openai/tests/openai/test_openai_chat_client.py
+++ b/python/packages/openai/tests/openai/test_openai_chat_client.py
@@ -1184,11 +1184,13 @@ def test_parse_response_from_openai_with_mcp_server_tool_result() -> None:
     assert result_content.output is not None
 
 
-def test_parse_chunk_from_openai_with_mcp_call_result() -> None:
-    """Test _parse_chunk_from_openai with MCP call output."""
+def test_parse_chunk_from_openai_with_mcp_call_added_defers_result() -> None:
+    """Test that response.output_item.added for mcp_call emits only the call, not the result.
+
+    The result is deferred to response.output_item.done.
+    """
     client = OpenAIChatClient(model="test-model", api_key="test-key")
 
-    # Mock event with MCP call that has output
     mock_event = MagicMock()
     mock_event.type = "response.output_item.added"
 
@@ -1199,8 +1201,9 @@ def test_parse_chunk_from_openai_with_mcp_call_result() -> None:
     mock_item.name = "fetch_resource"
     mock_item.server_label = "ResourceServer"
     mock_item.arguments = {"resource_id": "123"}
-    # Use proper content structure that _parse_content can handle
-    mock_item.result = [{"type": "text", "text": "test result"}]
+    mock_item.result = None
+    mock_item.output = None
+    mock_item.outputs = None
 
     mock_event.item = mock_item
     mock_event.output_index = 0
@@ -1209,18 +1212,124 @@ def test_parse_chunk_from_openai_with_mcp_call_result() -> None:
 
     update = client._parse_chunk_from_openai(mock_event, options={}, function_call_ids=function_call_ids)
 
-    # Should have both call and result in contents
-    assert len(update.contents) == 2
-    call_content, result_content = update.contents
+    # Should have only the call content — result is deferred
+    assert len(update.contents) == 1
+    call_content = update.contents[0]
 
     assert call_content.type == "mcp_server_tool_call"
     assert call_content.call_id in ["mcp_call_456", "call_456"]
     assert call_content.tool_name == "fetch_resource"
 
+    # No result should be emitted at this point
+    result_contents = [c for c in update.contents if c.type == "mcp_server_tool_result"]
+    assert len(result_contents) == 0
+
+
+def test_parse_chunk_from_openai_with_mcp_output_item_done() -> None:
+    """Test that response.output_item.done for mcp_call emits mcp_server_tool_result with output."""
+    client = OpenAIChatClient(model="test-model", api_key="test-key")
+
+    mock_event = MagicMock()
+    mock_event.type = "response.output_item.done"
+
+    mock_item = MagicMock()
+    mock_item.type = "mcp_call"
+    mock_item.id = "mcp_call_456"
+    mock_item.output = "The weather in Seattle is 72F and sunny."
+    mock_event.item = mock_item
+
+    function_call_ids: dict[int, tuple[str, str]] = {}
+
+    update = client._parse_chunk_from_openai(mock_event, options={}, function_call_ids=function_call_ids)
+
+    assert len(update.contents) == 1
+    result_content = update.contents[0]
+
     assert result_content.type == "mcp_server_tool_result"
-    assert result_content.call_id in ["mcp_call_456", "call_456"]
-    # Verify the output was parsed
+    assert result_content.call_id == "mcp_call_456"
     assert result_content.output is not None
+    assert len(result_content.output) == 1
+    assert result_content.output[0].text == "The weather in Seattle is 72F and sunny."
+    assert result_content.raw_representation is mock_item
+
+
+def test_parse_chunk_from_openai_with_mcp_output_item_done_no_output() -> None:
+    """Test that response.output_item.done for mcp_call with no output emits result with None output."""
+    client = OpenAIChatClient(model="test-model", api_key="test-key")
+
+    mock_event = MagicMock()
+    mock_event.type = "response.output_item.done"
+
+    mock_item = MagicMock()
+    mock_item.type = "mcp_call"
+    mock_item.id = "mcp_call_789"
+    mock_item.output = None
+    mock_event.item = mock_item
+
+    function_call_ids: dict[int, tuple[str, str]] = {}
+
+    update = client._parse_chunk_from_openai(mock_event, options={}, function_call_ids=function_call_ids)
+
+    assert len(update.contents) == 1
+    result_content = update.contents[0]
+
+    assert result_content.type == "mcp_server_tool_result"
+    assert result_content.call_id == "mcp_call_789"
+    assert result_content.output is None
+    assert result_content.raw_representation is mock_item
+
+
+def test_parse_chunk_from_openai_with_mcp_output_item_done_call_id_fallback() -> None:
+    """Test that response.output_item.done for mcp_call falls back to call_id when id is missing."""
+    client = OpenAIChatClient(model="test-model", api_key="test-key")
+
+    mock_event = MagicMock()
+    mock_event.type = "response.output_item.done"
+
+    mock_item = MagicMock(spec=[])
+    mock_item.type = "mcp_call"
+    mock_item.call_id = "mcp_fallback_123"
+    mock_item.output = "fallback result"
+    mock_event.item = mock_item
+
+    function_call_ids: dict[int, tuple[str, str]] = {}
+
+    update = client._parse_chunk_from_openai(mock_event, options={}, function_call_ids=function_call_ids)
+
+    assert len(update.contents) == 1
+    result_content = update.contents[0]
+
+    assert result_content.type == "mcp_server_tool_result"
+    assert result_content.call_id == "mcp_fallback_123"
+    assert result_content.output is not None
+    assert result_content.output[0].text == "fallback result"
+    assert result_content.raw_representation is mock_item
+
+
+def test_parse_chunk_from_openai_with_mcp_output_item_done_no_id_fallback() -> None:
+    """Test that response.output_item.done for mcp_call falls back to empty string when neither id nor call_id exist."""
+    client = OpenAIChatClient(model="test-model", api_key="test-key")
+
+    mock_event = MagicMock()
+    mock_event.type = "response.output_item.done"
+
+    mock_item = MagicMock(spec=[])
+    mock_item.type = "mcp_call"
+    mock_item.output = "some result"
+    mock_event.item = mock_item
+
+    function_call_ids: dict[int, tuple[str, str]] = {}
+
+    update = client._parse_chunk_from_openai(mock_event, options={}, function_call_ids=function_call_ids)
+
+    assert len(update.contents) == 1
+    result_content = update.contents[0]
+
+    assert result_content.type == "mcp_server_tool_result"
+    assert result_content.call_id == ""
+    assert result_content.output is not None
+    assert result_content.output[0].text == "some result"
+    assert result_content.raw_representation is mock_item
 
 
 def test_prepare_message_for_openai_with_function_approval_response() -> None:

From dd3d085539375c38927e6101ac60688b69fa08c1 Mon Sep 17 00:00:00 2001
From: Evan Mattson <35585003+moonbox3@users.noreply.github.com>
Date: Thu, 26 Mar 2026 14:56:10 +0900
Subject: [PATCH 3/4] Python: Include reasoning messages in MESSAGES_SNAPSHOT
 events (#4844)

* Include reasoning messages in MESSAGES_SNAPSHOT (#4843)

FlowState now tracks reasoning messages emitted during a run.
_emit_text_reasoning() persists reasoning (including encrypted_value)
into flow.reasoning_messages, and _build_messages_snapshot() appends
them to the final MESSAGES_SNAPSHOT event.

Changes:
- Add reasoning_messages field to FlowState
- Update _emit_text_reasoning() to accept optional flow parameter
- Include reasoning_messages in _build_messages_snapshot()
- Add 'reasoning' to ALLOWED_AGUI_ROLES so normalize_agui_role()
  preserves the role through snapshot round-trips
- Skip reasoning messages in agui_messages_to_agent_framework() since
  they are UI-only state and should not be forwarded to LLM providers
- Add regression tests for snapshot emission, encrypted value
  preservation, and multi-turn round-trip with reasoning

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Python: Include reasoning messages in MESSAGES_SNAPSHOT events

Fixes #4843

* Fix PR review feedback for reasoning persistence (#4843)

- Accumulate reasoning text per message_id (append deltas) instead of
  storing only the current chunk, matching flow.accumulated_text pattern
- Use camelCase encryptedValue in snapshot JSON to match AG-UI protocol
  conventions (toolCallId, encryptedValue)
- Normalize snake_case encrypted_value to encryptedValue in
  agui_messages_to_snapshot_format for input compatibility
- Update normalize_agui_role docstring to include reasoning role
- Add tests for incremental reasoning accumulation and key normalization

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Address review feedback for #4843: Python: agent-framework-ag-ui: include reasoning messages in MESSAGES_SNAPSHOT

---------

Co-authored-by: Copilot <copilot@github.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
---
 .../ag-ui/agent_framework_ag_ui/_agent_run.py |   8 +-
 .../_message_adapters.py                      |   9 +
 .../agent_framework_ag_ui/_run_common.py      |  40 ++++-
 .../ag-ui/agent_framework_ag_ui/_utils.py     |   4 +-
 .../tests/ag_ui/test_message_adapters.py      |  91 ++++++++++
 python/packages/ag-ui/tests/ag_ui/test_run.py | 155 ++++++++++++++++++
 .../packages/ag-ui/tests/ag_ui/test_utils.py  |   1 +
 7 files changed, 303 insertions(+), 5 deletions(-)

diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py b/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py
index 4b00330283..e9ce610b10 100644
--- a/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py
+++ b/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py
@@ -684,6 +684,10 @@ def _build_messages_snapshot(
             }
         )
 
+    # Add reasoning messages so frontends that reconcile state from
+    # MESSAGES_SNAPSHOT retain reasoning content after streaming ends.
+    all_messages.extend(flow.reasoning_messages)
+
     return MessagesSnapshotEvent(messages=all_messages)  # type: ignore[arg-type]
 
 
@@ -1061,7 +1065,9 @@ async def run_agent_stream(
 
     # Emit MessagesSnapshotEvent if we have tool calls or results
     # Feature #5: Suppress intermediate snapshots for predictive tools without confirmation
-    should_emit_snapshot = flow.pending_tool_calls or flow.tool_results or flow.accumulated_text
+    should_emit_snapshot = (
+        flow.pending_tool_calls or flow.tool_results or flow.accumulated_text or flow.reasoning_messages
+    )
     if should_emit_snapshot:
         # Check if we should suppress for predictive tool
         last_tool_name = None
diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_message_adapters.py b/python/packages/ag-ui/agent_framework_ag_ui/_message_adapters.py
index 2e5294a6b6..5e4fced97c 100644
--- a/python/packages/ag-ui/agent_framework_ag_ui/_message_adapters.py
+++ b/python/packages/ag-ui/agent_framework_ag_ui/_message_adapters.py
@@ -604,6 +604,10 @@ def agui_messages_to_agent_framework(messages: list[dict[str, Any]]) -> list[Mes
         # Handle standard tool result messages early (role="tool") to preserve provider invariants
         # This path maps AG‑UI tool messages to function_result content with the correct tool_call_id
         role_str = normalize_agui_role(msg.get("role", "user"))
+        if role_str == "reasoning":
+            # Reasoning messages are UI-only state carried in MESSAGES_SNAPSHOT.
+            # They should not be forwarded to the LLM provider.
+            continue
         if role_str == "tool":
             # Prefer explicit tool_call_id fields; fall back to backend fields only if necessary
             tool_call_id = msg.get("tool_call_id") or msg.get("toolCallId")
@@ -1020,6 +1024,11 @@ def agui_messages_to_snapshot_format(messages: list[dict[str, Any]]) -> list[dic
             elif "toolCallId" not in normalized_msg:
                 normalized_msg["toolCallId"] = ""
 
+        # Normalize encrypted_value to encryptedValue for reasoning messages
+        if normalized_msg.get("role") == "reasoning" and "encrypted_value" in normalized_msg:
+            normalized_msg["encryptedValue"] = normalized_msg["encrypted_value"]
+            del normalized_msg["encrypted_value"]
+
         result.append(normalized_msg)
 
     return result
diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_run_common.py b/python/packages/ag-ui/agent_framework_ag_ui/_run_common.py
index 0a9f4cea9c..155f559a94 100644
--- a/python/packages/ag-ui/agent_framework_ag_ui/_run_common.py
+++ b/python/packages/ag-ui/agent_framework_ag_ui/_run_common.py
@@ -126,6 +126,8 @@ class FlowState:
     tool_results: list[dict[str, Any]] = field(default_factory=list)  # pyright: ignore[reportUnknownVariableType]
     tool_calls_ended: set[str] = field(default_factory=set)  # pyright: ignore[reportUnknownVariableType]
     interrupts: list[dict[str, Any]] = field(default_factory=list)  # pyright: ignore[reportUnknownVariableType]
+    reasoning_messages: list[dict[str, Any]] = field(default_factory=list)  # pyright: ignore[reportUnknownVariableType]
+    accumulated_reasoning: dict[str, str] = field(default_factory=dict)  # pyright: ignore[reportUnknownVariableType]
 
     def get_tool_name(self, call_id: str | None) -> str | None:
         """Get tool name by call ID."""
@@ -460,7 +462,7 @@ def _emit_mcp_tool_result(
     return _emit_tool_result_common(content.call_id, raw_output, flow, predictive_handler)
 
 
-def _emit_text_reasoning(content: Content) -> list[BaseEvent]:
+def _emit_text_reasoning(content: Content, flow: FlowState | None = None) -> list[BaseEvent]:
     """Emit AG-UI reasoning events for text_reasoning content.
 
     Uses the protocol-defined reasoning event types so that AG-UI consumers
@@ -470,6 +472,10 @@ def _emit_text_reasoning(content: Content) -> list[BaseEvent]:
     ``content.protected_data`` is present it is emitted as a
     ``ReasoningEncryptedValueEvent`` so that consumers can persist encrypted
     reasoning for state continuity without conflating it with display text.
+
+    When *flow* is provided the reasoning message is persisted into
+    ``flow.reasoning_messages`` so that ``_build_messages_snapshot`` can
+    include it in the final ``MESSAGES_SNAPSHOT``.
     """
     text = content.text or ""
     if not text and content.protected_data is None:
@@ -498,6 +504,36 @@ def _emit_text_reasoning(content: Content) -> list[BaseEvent]:
 
     events.append(ReasoningEndEvent(message_id=message_id))
 
+    # Persist reasoning into flow state for MESSAGES_SNAPSHOT.
+    # Accumulate reasoning text per message_id, similar to flow.accumulated_text,
+    # so that incremental deltas build the full reasoning string.
+    if flow is not None:
+        if text:
+            previous_text = flow.accumulated_reasoning.get(message_id, "")
+            flow.accumulated_reasoning[message_id] = previous_text + text
+        full_text = flow.accumulated_reasoning.get(message_id, text or "")
+
+        # Update existing reasoning entry for this message_id if present; otherwise append a new one.
+        existing_entry: dict[str, Any] | None = None
+        for entry in flow.reasoning_messages:
+            if isinstance(entry, dict) and entry.get("id") == message_id:
+                existing_entry = entry
+                break
+
+        if existing_entry is None:
+            reasoning_entry: dict[str, Any] = {
+                "id": message_id,
+                "role": "reasoning",
+                "content": full_text,
+            }
+            if content.protected_data is not None:
+                reasoning_entry["encryptedValue"] = content.protected_data
+            flow.reasoning_messages.append(reasoning_entry)
+        else:
+            existing_entry["content"] = full_text
+            if content.protected_data is not None:
+                existing_entry["encryptedValue"] = content.protected_data
+
     return events
 
 
@@ -527,6 +563,6 @@ def _emit_content(
     if content_type == "mcp_server_tool_result":
         return _emit_mcp_tool_result(content, flow, predictive_handler)
     if content_type == "text_reasoning":
-        return _emit_text_reasoning(content)
+        return _emit_text_reasoning(content, flow)
     logger.debug("Skipping unsupported content type in AG-UI emitter: %s", content_type)
     return []
diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_utils.py b/python/packages/ag-ui/agent_framework_ag_ui/_utils.py
index bfda3948ec..c68301f7d2 100644
--- a/python/packages/ag-ui/agent_framework_ag_ui/_utils.py
+++ b/python/packages/ag-ui/agent_framework_ag_ui/_utils.py
@@ -27,7 +27,7 @@ FRAMEWORK_TO_AGUI_ROLE: dict[str, str] = {
     "system": "system",
 }
 
-ALLOWED_AGUI_ROLES: set[str] = {"user", "assistant", "system", "tool"}
+ALLOWED_AGUI_ROLES: set[str] = {"user", "assistant", "system", "tool", "reasoning"}
 
 
 def generate_event_id() -> str:
@@ -82,7 +82,7 @@ def normalize_agui_role(raw_role: Any) -> str:
         raw_role: Raw role value from AG-UI message
 
     Returns:
-        Normalized role string (user, assistant, system, or tool)
+        Normalized role string (user, assistant, system, tool, or reasoning)
     """
     if not isinstance(raw_role, str):
         return "user"
diff --git a/python/packages/ag-ui/tests/ag_ui/test_message_adapters.py b/python/packages/ag-ui/tests/ag_ui/test_message_adapters.py
index cc4f1230df..9508b53085 100644
--- a/python/packages/ag-ui/tests/ag_ui/test_message_adapters.py
+++ b/python/packages/ag-ui/tests/ag_ui/test_message_adapters.py
@@ -1669,3 +1669,94 @@ def test_agui_fresh_approval_is_still_processed():
     assert len(approval_contents) == 1, "Fresh approval should produce function_approval_response"
     assert approval_contents[0].approved is True
     assert approval_contents[0].function_call.name == "get_datetime"
+
+
+class TestReasoningRoundTrip:
+    """Tests for reasoning message handling in inbound/outbound adapters."""
+
+    def test_reasoning_skipped_on_inbound(self):
+        """Reasoning messages from prior snapshot are not forwarded to the LLM."""
+        messages_input = [
+            {"id": "u1", "role": "user", "content": "Hello"},
+            {"id": "r1", "role": "reasoning", "content": "Thinking..."},
+            {"id": "a1", "role": "assistant", "content": "Hi there"},
+        ]
+
+        result = agui_messages_to_agent_framework(messages_input)
+
+        roles = [m.role if hasattr(m.role, "value") else str(m.role) for m in result]
+        assert "reasoning" not in roles
+        assert len(result) == 2
+
+    def test_reasoning_preserved_in_snapshot_format(self):
+        """Reasoning messages retain their role through snapshot normalization."""
+        messages_input = [
+            {"id": "u1", "role": "user", "content": "Hello"},
+            {"id": "r1", "role": "reasoning", "content": "Thinking about this..."},
+            {"id": "a1", "role": "assistant", "content": "Answer"},
+        ]
+
+        result = agui_messages_to_snapshot_format(messages_input)
+
+        reasoning_msgs = [m for m in result if m.get("role") == "reasoning"]
+        assert len(reasoning_msgs) == 1
+        assert reasoning_msgs[0]["content"] == "Thinking about this..."
+
+    def test_reasoning_with_encrypted_value_in_snapshot_format(self):
+        """Reasoning with encryptedValue passes through snapshot normalization."""
+        messages_input = [
+            {
+                "id": "r1",
+                "role": "reasoning",
+                "content": "visible",
+                "encryptedValue": "secret-data",
+            },
+        ]
+
+        result = agui_messages_to_snapshot_format(messages_input)
+
+        assert len(result) == 1
+        assert result[0]["role"] == "reasoning"
+        assert result[0]["encryptedValue"] == "secret-data"
+
+    def test_reasoning_encrypted_value_snake_case_normalized(self):
+        """Snake-case encrypted_value is normalized to encryptedValue in snapshot format."""
+        messages_input = [
+            {
+                "id": "r1",
+                "role": "reasoning",
+                "content": "visible",
+                "encrypted_value": "snake-case-data",
+            },
+        ]
+
+        result = agui_messages_to_snapshot_format(messages_input)
+
+        assert len(result) == 1
+        assert result[0]["encryptedValue"] == "snake-case-data"
+        assert "encrypted_value" not in result[0]
+
+    def test_multi_turn_with_reasoning_in_prior_snapshot(self):
+        """Second turn with reasoning from prior snapshot does not corrupt messages."""
+        messages_input = [
+            {"id": "u1", "role": "user", "content": "First question"},
+            {"id": "r1", "role": "reasoning", "content": "Prior reasoning"},
+            {"id": "a1", "role": "assistant", "content": "First answer"},
+            {"id": "u2", "role": "user", "content": "Follow-up question"},
+        ]
+
+        result = agui_messages_to_agent_framework(messages_input)
+
+        roles = [m.role if hasattr(m.role, "value") else str(m.role) for m in result]
+        # Reasoning is filtered out, other messages preserved in order
+        assert roles == ["user", "assistant", "user"]
+        # Content not corrupted
+        texts = []
+        for m in result:
+            for c in m.contents or []:
+                if hasattr(c, "text") and c.text:
+                    texts.append(c.text)
+        assert "First question" in texts
+        assert "First answer" in texts
+        assert "Follow-up question" in texts
+        assert "Prior reasoning" not in texts
diff --git a/python/packages/ag-ui/tests/ag_ui/test_run.py b/python/packages/ag-ui/tests/ag_ui/test_run.py
index ae8c5e85b0..0e5c329ce9 100644
--- a/python/packages/ag-ui/tests/ag_ui/test_run.py
+++ b/python/packages/ag-ui/tests/ag_ui/test_run.py
@@ -1346,3 +1346,158 @@ class TestEmitContentMcpRouting:
 
         assert len(events) == 5
         assert isinstance(events[0], ReasoningStartEvent)
+
+
+class TestReasoningInSnapshot:
+    """Tests for reasoning message inclusion in MESSAGES_SNAPSHOT."""
+
+    def test_reasoning_persisted_to_flow_state(self):
+        """_emit_text_reasoning with flow persists reasoning into flow.reasoning_messages."""
+        flow = FlowState()
+        content = Content.from_text_reasoning(
+            id="reason_persist",
+            text="Let me think step by step.",
+        )
+
+        _emit_text_reasoning(content, flow)
+
+        assert len(flow.reasoning_messages) == 1
+        assert flow.reasoning_messages[0]["id"] == "reason_persist"
+        assert flow.reasoning_messages[0]["role"] == "reasoning"
+        assert flow.reasoning_messages[0]["content"] == "Let me think step by step."
+        assert "encryptedValue" not in flow.reasoning_messages[0]
+
+    def test_reasoning_with_encrypted_value_persisted(self):
+        """Reasoning with protected_data preserves encryptedValue in flow state."""
+        flow = FlowState()
+        content = Content.from_text_reasoning(
+            id="reason_enc",
+            text="visible reasoning",
+            protected_data="encrypted-data-123",
+        )
+
+        _emit_text_reasoning(content, flow)
+
+        assert len(flow.reasoning_messages) == 1
+        assert flow.reasoning_messages[0]["encryptedValue"] == "encrypted-data-123"
+
+    def test_snapshot_includes_reasoning(self):
+        """_build_messages_snapshot includes reasoning messages from flow state."""
+        from agent_framework_ag_ui._agent_run import _build_messages_snapshot
+
+        flow = FlowState()
+        flow.accumulated_text = "Here is my answer."
+        flow.reasoning_messages = [
+            {"id": "r1", "role": "reasoning", "content": "Thinking..."},
+        ]
+
+        snapshot = _build_messages_snapshot(flow, [])
+
+        roles = [m.get("role") if isinstance(m, dict) else getattr(m, "role", None) for m in snapshot.messages]
+        assert "reasoning" in roles
+
+    def test_snapshot_preserves_reasoning_encrypted_value(self):
+        """Snapshot reasoning with encryptedValue is preserved end-to-end."""
+        from agent_framework_ag_ui._agent_run import _build_messages_snapshot
+
+        flow = FlowState()
+        content = Content.from_text_reasoning(
+            id="reason_e2e",
+            text="visible",
+            protected_data="secret-data",
+        )
+        _emit_text_reasoning(content, flow)
+
+        text_content = Content.from_text("Final answer.")
+        _emit_text(text_content, flow)
+
+        snapshot = _build_messages_snapshot(flow, [])
+
+        reasoning_msgs = [
+            m
+            for m in snapshot.messages
+            if (m.get("role") if isinstance(m, dict) else getattr(m, "role", None)) == "reasoning"
+        ]
+        assert len(reasoning_msgs) == 1
+        msg = reasoning_msgs[0]
+        if isinstance(msg, dict):
+            assert msg["content"] == "visible"
+            assert msg["encryptedValue"] == "secret-data"
+
+    def test_emit_content_routes_reasoning_with_flow(self):
+        """_emit_content passes flow to _emit_text_reasoning for persistence."""
+        flow = FlowState()
+        content = Content.from_text_reasoning(text="routed reasoning")
+
+        _emit_content(content, flow)
+
+        assert len(flow.reasoning_messages) == 1
+        assert flow.reasoning_messages[0]["content"] == "routed reasoning"
+
+    def test_reasoning_without_flow_does_not_error(self):
+        """Calling _emit_text_reasoning without flow still works (backward compat)."""
+        content = Content.from_text_reasoning(text="no flow")
+
+        events = _emit_text_reasoning(content)
+
+        assert len(events) == 5
+        assert isinstance(events[0], ReasoningStartEvent)
+
+    def test_snapshot_reasoning_ordering(self):
+        """Reasoning messages appear after assistant text in snapshot."""
+        from agent_framework_ag_ui._agent_run import _build_messages_snapshot
+
+        flow = FlowState()
+        reasoning_content = Content.from_text_reasoning(id="r1", text="Thinking...")
+        _emit_text_reasoning(reasoning_content, flow)
+
+        text_content = Content.from_text("Answer")
+        _emit_text(text_content, flow)
+
+        snapshot = _build_messages_snapshot(flow, [{"id": "u1", "role": "user", "content": "Hi"}])
+
+        # user -> assistant text -> reasoning
+        assert len(snapshot.messages) == 3
+        roles = [m.get("role") if isinstance(m, dict) else getattr(m, "role", None) for m in snapshot.messages]
+        assert roles == ["user", "assistant", "reasoning"]
+
+    def test_reasoning_accumulates_incremental_deltas(self):
+        """Multiple reasoning deltas with the same id accumulate into one entry."""
+        flow = FlowState()
+        content1 = Content.from_text_reasoning(id="reason_inc", text="First ")
+        content2 = Content.from_text_reasoning(id="reason_inc", text="second ")
+        content3 = Content.from_text_reasoning(id="reason_inc", text="third.")
+
+        _emit_text_reasoning(content1, flow)
+        _emit_text_reasoning(content2, flow)
+        _emit_text_reasoning(content3, flow)
+
+        assert len(flow.reasoning_messages) == 1
+        assert flow.reasoning_messages[0]["id"] == "reason_inc"
+        assert flow.reasoning_messages[0]["content"] == "First second third."
+
+    def test_reasoning_accumulates_distinct_message_ids(self):
+        """Reasoning entries with different ids are stored separately."""
+        flow = FlowState()
+        content_a = Content.from_text_reasoning(id="a", text="alpha")
+        content_b = Content.from_text_reasoning(id="b", text="beta")
+
+        _emit_text_reasoning(content_a, flow)
+        _emit_text_reasoning(content_b, flow)
+
+        assert len(flow.reasoning_messages) == 2
+        assert flow.reasoning_messages[0]["content"] == "alpha"
+        assert flow.reasoning_messages[1]["content"] == "beta"
+
+    def test_reasoning_encrypted_value_updated_on_later_delta(self):
+        """encryptedValue is set even when it arrives with a later delta."""
+        flow = FlowState()
+        content1 = Content.from_text_reasoning(id="enc_late", text="part1 ")
+        content2 = Content.from_text_reasoning(id="enc_late", text="part2", protected_data="encrypted-payload")
+
+        _emit_text_reasoning(content1, flow)
+        _emit_text_reasoning(content2, flow)
+
+        assert len(flow.reasoning_messages) == 1
+        assert flow.reasoning_messages[0]["content"] == "part1 part2"
+        assert flow.reasoning_messages[0]["encryptedValue"] == "encrypted-payload"
diff --git a/python/packages/ag-ui/tests/ag_ui/test_utils.py b/python/packages/ag-ui/tests/ag_ui/test_utils.py
index 0f453132f7..f353d2f0a7 100644
--- a/python/packages/ag-ui/tests/ag_ui/test_utils.py
+++ b/python/packages/ag-ui/tests/ag_ui/test_utils.py
@@ -450,6 +450,7 @@ def test_normalize_agui_role_valid():
     assert normalize_agui_role("assistant") == "assistant"
     assert normalize_agui_role("system") == "system"
     assert normalize_agui_role("tool") == "tool"
+    assert normalize_agui_role("reasoning") == "reasoning"
 
 
 def test_normalize_agui_role_invalid():

From efb14cedb12b8e7f7b868921ddec1713e9e21972 Mon Sep 17 00:00:00 2001
From: Eduard van Valkenburg <eavanvalkenburg@users.noreply.github.com>
Date: Thu, 26 Mar 2026 08:33:19 +0100
Subject: [PATCH 4/4] Python: Support structuredContent in MCP tool results and
 fix sampling options type (#4763)

* Support MCP sampling tools capability (#4625)

Forward systemPrompt, tools, and toolChoice from MCP sampling requests
to the chat client's get_response() call. Also advertise the
sampling.tools capability to MCP servers when a client is configured.

- Pass SamplingCapability with tools support to ClientSession
- Convert systemPrompt to instructions in options
- Convert MCP Tool objects to FunctionTool instances for options
- Map MCP ToolChoice.mode to tool_choice in options
- Add tests for all new behaviors and update existing sampling tests

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Fix #4625: Support MCP sampling tool with proper typing and structured content

- Fix mypy error by typing sampling callback options as ChatOptions[None]
  instead of dict[str, Any], and importing ChatOptions from _types
- Handle structuredContent from CallToolResult in _parse_tool_result_from_mcp,
  serializing it as JSON text Content when present
- Add tests for structuredContent parsing (with and without regular content)

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Fix lint: add author to TODO comment

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Address review feedback for #4625: remove default=str, add edge-case tests

- Remove default=str from json.dumps for structuredContent to fail fast
  on non-JSON-serializable values instead of silently converting
- Add test for non-JSON-serializable structuredContent (TypeError)
- Add tests for empty systemPrompt ('') and empty tools list ([]) edge
  cases in sampling callback
- Expand TODO comment noting list[Content] return type constraint for
  future result_type support

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Sanitize sampling callback error to avoid leaking internals (#4625)

Log exception details at DEBUG level instead of including them in the
ErrorData message returned to the MCP server, which may be untrusted.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Address review feedback for #4625: move params to options, restore error info

- Remove stale TODO comment about response_format (ChatOptions already has it)
- Restore {ex} in sampling callback error message for useful debugging info
- Set structuredContent as additional_property on Content for structured access
- Move temperature, max_tokens, stop into options dict (not top-level kwargs)
- Only set temperature when provided (not all models support it)
- Add tests for generation params in options and temperature omission

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Fix MCP sampling callback and structured content error handling (#4625)

- Guard max_tokens like temperature: only set when not None, so options
  can properly evaluate to None when all params are absent
- Wrap json.dumps of structuredContent in try/except to fall back to
  str() for non-serializable values instead of propagating TypeError
- Extract test_connect_sampling_capabilities_with_client into its own
  test function so pytest can discover it independently
- Add test for max_tokens=None omission from options
- Update structured content non-serializable test to expect fallback

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Address review feedback for #4625: review comment fixes

* Fix MCP and Azure validation regressions

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

---------

Co-authored-by: Copilot <copilot@github.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
---
 python/packages/core/agent_framework/_mcp.py  |  41 +-
 python/packages/core/tests/core/test_mcp.py   | 366 +++++++++++++++++-
 .../openai/agent_framework_openai/_shared.py  |   4 +-
 .../packages/openai/tests/openai/conftest.py  |   2 +
 ...est_openai_chat_completion_client_azure.py |   3 +-
 5 files changed, 409 insertions(+), 7 deletions(-)

diff --git a/python/packages/core/agent_framework/_mcp.py b/python/packages/core/agent_framework/_mcp.py
index c7db6177da..267e176ee8 100644
--- a/python/packages/core/agent_framework/_mcp.py
+++ b/python/packages/core/agent_framework/_mcp.py
@@ -18,7 +18,11 @@ from typing import TYPE_CHECKING, Any, Literal, TypedDict, cast
 from opentelemetry import propagate
 
 from ._tools import FunctionTool
-from ._types import Content, Message
+from ._types import (
+    ChatOptions,
+    Content,
+    Message,
+)
 from .exceptions import ToolException, ToolExecutionException
 
 if sys.version_info >= (3, 11):
@@ -640,6 +644,7 @@ class MCPTool:
                 raise ToolException(error_msg, inner_exception=ex) from ex
             try:
                 try:
+                    from mcp import types
                     from mcp.client.session import ClientSession as runtime_client_session
                 except ModuleNotFoundError as ex:
                     await self._safe_close_exit_stack()
@@ -647,6 +652,12 @@ class MCPTool:
                         "MCP support requires `mcp`. Please install `mcp`.",
                         inner_exception=ex,
                     ) from ex
+
+                sampling_capabilities = None
+                if self.client is not None:
+                    sampling_capabilities = types.SamplingCapability(
+                        tools=types.SamplingToolsCapability(),
+                    )
                 session = await self._exit_stack.enter_async_context(
                     runtime_client_session(
                         read_stream=transport[0],
@@ -657,6 +668,7 @@ class MCPTool:
                         message_handler=self.message_handler,
                         logging_callback=self.logging_callback,
                         sampling_callback=self.sampling_callback,
+                        sampling_capabilities=sampling_capabilities,
                     )
                 )
             except Exception as ex:
@@ -733,14 +745,35 @@ class MCPTool:
         messages: list[Message] = []
         for msg in params.messages:
             messages.append(self._parse_message_from_mcp(msg))
+
+        options: ChatOptions[None] = {}
+        if params.systemPrompt is not None:
+            options["instructions"] = params.systemPrompt
+        if params.tools is not None:
+            options["tools"] = [
+                FunctionTool(
+                    name=tool.name,
+                    description=tool.description or "",
+                    input_model=tool.inputSchema,
+                )
+                for tool in params.tools
+            ]
+        if params.toolChoice is not None and params.toolChoice.mode is not None:
+            options["tool_choice"] = params.toolChoice.mode
+
+        if params.temperature is not None:
+            options["temperature"] = params.temperature
+        options["max_tokens"] = params.maxTokens
+        if params.stopSequences is not None:
+            options["stop"] = params.stopSequences
+
         try:
             response = await self.client.get_response(
                 messages,
-                temperature=params.temperature,
-                max_tokens=params.maxTokens,
-                stop=params.stopSequences,
+                options=options or None,
             )
         except Exception as ex:
+            logger.debug("Sampling callback error: %s", ex, exc_info=True)
             return types.ErrorData(
                 code=types.INTERNAL_ERROR,
                 message=f"Failed to get chat message content: {ex}",
diff --git a/python/packages/core/tests/core/test_mcp.py b/python/packages/core/tests/core/test_mcp.py
index 3c340a13b8..eb233eea99 100644
--- a/python/packages/core/tests/core/test_mcp.py
+++ b/python/packages/core/tests/core/test_mcp.py
@@ -1696,12 +1696,15 @@ async def test_mcp_tool_sampling_callback_chat_client_exception():
     params.temperature = None
     params.maxTokens = None
     params.stopSequences = None
+    params.systemPrompt = None
+    params.tools = None
+    params.toolChoice = None
 
     result = await tool.sampling_callback(Mock(), params)
 
     assert isinstance(result, types.ErrorData)
     assert result.code == types.INTERNAL_ERROR
-    assert "Failed to get chat message content: Chat client error" in result.message
+    assert "Failed to get chat message content" in result.message
 
 
 async def test_mcp_tool_sampling_callback_no_valid_content():
@@ -1739,6 +1742,9 @@ async def test_mcp_tool_sampling_callback_no_valid_content():
     params.temperature = None
     params.maxTokens = None
     params.stopSequences = None
+    params.systemPrompt = None
+    params.tools = None
+    params.toolChoice = None
 
     result = await tool.sampling_callback(Mock(), params)
 
@@ -1757,6 +1763,9 @@ async def test_mcp_tool_sampling_callback_no_response_and_successful_message_cre
     params.temperature = None
     params.maxTokens = None
     params.stopSequences = None
+    params.systemPrompt = None
+    params.tools = None
+    params.toolChoice = None
 
     tool.client.get_response.return_value = None
     no_response = await tool.sampling_callback(Mock(), params)
@@ -1787,6 +1796,361 @@ async def test_mcp_tool_logging_callback_logs_at_requested_level() -> None:
     mock_log.assert_called_once_with(logging.WARNING, "be careful")
 
 
+async def test_mcp_tool_sampling_callback_forwards_system_prompt():
+    """Test sampling callback passes systemPrompt as instructions in options."""
+    from agent_framework import Message
+
+    tool = MCPStdioTool(name="test_tool", command="python")
+
+    mock_chat_client = AsyncMock()
+    mock_response = Mock()
+    mock_response.messages = [Message(role="assistant", contents=[Content.from_text("response")])]
+    mock_response.model_id = "test-model"
+    mock_chat_client.get_response.return_value = mock_response
+
+    tool.client = mock_chat_client
+
+    params = Mock()
+    mock_message = Mock()
+    mock_message.role = "user"
+    mock_message.content = Mock()
+    mock_message.content.text = "Test question"
+    params.messages = [mock_message]
+    params.temperature = None
+    params.maxTokens = None
+    params.stopSequences = None
+    params.systemPrompt = "You are a helpful assistant"
+    params.tools = None
+    params.toolChoice = None
+
+    result = await tool.sampling_callback(Mock(), params)
+
+    assert isinstance(result, types.CreateMessageResult)
+    call_kwargs = mock_chat_client.get_response.call_args
+    options = call_kwargs.kwargs.get("options") or {}
+    assert options.get("instructions") == "You are a helpful assistant"
+
+
+async def test_mcp_tool_sampling_callback_forwards_tools():
+    """Test sampling callback converts MCP tools to FunctionTools and passes them in options."""
+    from agent_framework import FunctionTool, Message
+
+    tool = MCPStdioTool(name="test_tool", command="python")
+
+    mock_chat_client = AsyncMock()
+    mock_response = Mock()
+    mock_response.messages = [Message(role="assistant", contents=[Content.from_text("response")])]
+    mock_response.model_id = "test-model"
+    mock_chat_client.get_response.return_value = mock_response
+
+    tool.client = mock_chat_client
+
+    mcp_tool = types.Tool(
+        name="get_weather",
+        description="Get weather",
+        inputSchema={"type": "object", "properties": {"city": {"type": "string"}}},
+    )
+
+    params = Mock()
+    mock_message = Mock()
+    mock_message.role = "user"
+    mock_message.content = Mock()
+    mock_message.content.text = "Test question"
+    params.messages = [mock_message]
+    params.temperature = None
+    params.maxTokens = None
+    params.stopSequences = None
+    params.systemPrompt = None
+    params.tools = [mcp_tool]
+    params.toolChoice = None
+
+    result = await tool.sampling_callback(Mock(), params)
+
+    assert isinstance(result, types.CreateMessageResult)
+    call_kwargs = mock_chat_client.get_response.call_args
+    options = call_kwargs.kwargs.get("options") or {}
+    tools = options.get("tools")
+    assert tools is not None
+    assert len(tools) == 1
+    assert isinstance(tools[0], FunctionTool)
+    assert tools[0].name == "get_weather"
+    assert tools[0].description == "Get weather"
+
+
+async def test_mcp_tool_sampling_callback_forwards_tool_choice():
+    """Test sampling callback passes toolChoice mode in options."""
+    from agent_framework import Message
+
+    tool = MCPStdioTool(name="test_tool", command="python")
+
+    mock_chat_client = AsyncMock()
+    mock_response = Mock()
+    mock_response.messages = [Message(role="assistant", contents=[Content.from_text("response")])]
+    mock_response.model_id = "test-model"
+    mock_chat_client.get_response.return_value = mock_response
+
+    tool.client = mock_chat_client
+
+    params = Mock()
+    mock_message = Mock()
+    mock_message.role = "user"
+    mock_message.content = Mock()
+    mock_message.content.text = "Test question"
+    params.messages = [mock_message]
+    params.temperature = None
+    params.maxTokens = None
+    params.stopSequences = None
+    params.systemPrompt = None
+    params.tools = None
+    params.toolChoice = types.ToolChoice(mode="required")
+
+    result = await tool.sampling_callback(Mock(), params)
+
+    assert isinstance(result, types.CreateMessageResult)
+    call_kwargs = mock_chat_client.get_response.call_args
+    options = call_kwargs.kwargs.get("options") or {}
+    assert options.get("tool_choice") == "required"
+
+
+async def test_mcp_tool_sampling_callback_forwards_empty_system_prompt():
+    """Test sampling callback forwards empty string systemPrompt as instructions."""
+    from agent_framework import Message
+
+    tool = MCPStdioTool(name="test_tool", command="python")
+
+    mock_chat_client = AsyncMock()
+    mock_response = Mock()
+    mock_response.messages = [Message(role="assistant", contents=[Content.from_text("response")])]
+    mock_response.model_id = "test-model"
+    mock_chat_client.get_response.return_value = mock_response
+
+    tool.client = mock_chat_client
+
+    params = Mock()
+    mock_message = Mock()
+    mock_message.role = "user"
+    mock_message.content = Mock()
+    mock_message.content.text = "Test question"
+    params.messages = [mock_message]
+    params.temperature = None
+    params.maxTokens = None
+    params.stopSequences = None
+    params.systemPrompt = ""
+    params.tools = None
+    params.toolChoice = None
+
+    result = await tool.sampling_callback(Mock(), params)
+
+    assert isinstance(result, types.CreateMessageResult)
+    call_kwargs = mock_chat_client.get_response.call_args
+    options = call_kwargs.kwargs.get("options") or {}
+    assert options.get("instructions") == ""
+
+
+async def test_mcp_tool_sampling_callback_forwards_empty_tools_list():
+    """Test sampling callback forwards empty tools list in options."""
+    from agent_framework import Message
+
+    tool = MCPStdioTool(name="test_tool", command="python")
+
+    mock_chat_client = AsyncMock()
+    mock_response = Mock()
+    mock_response.messages = [Message(role="assistant", contents=[Content.from_text("response")])]
+    mock_response.model_id = "test-model"
+    mock_chat_client.get_response.return_value = mock_response
+
+    tool.client = mock_chat_client
+
+    params = Mock()
+    mock_message = Mock()
+    mock_message.role = "user"
+    mock_message.content = Mock()
+    mock_message.content.text = "Test question"
+    params.messages = [mock_message]
+    params.temperature = None
+    params.maxTokens = None
+    params.stopSequences = None
+    params.systemPrompt = None
+    params.tools = []
+    params.toolChoice = None
+
+    result = await tool.sampling_callback(Mock(), params)
+
+    assert isinstance(result, types.CreateMessageResult)
+    call_kwargs = mock_chat_client.get_response.call_args
+    options = call_kwargs.kwargs.get("options") or {}
+    assert options.get("tools") == []
+
+
+async def test_mcp_tool_sampling_callback_forwards_generation_params_in_options():
+    """Test sampling callback passes temperature, max_tokens, and stop in options."""
+    from agent_framework import Message
+
+    tool = MCPStdioTool(name="test_tool", command="python")
+
+    mock_chat_client = AsyncMock()
+    mock_response = Mock()
+    mock_response.messages = [Message(role="assistant", contents=[Content.from_text("response")])]
+    mock_response.model_id = "test-model"
+    mock_chat_client.get_response.return_value = mock_response
+
+    tool.client = mock_chat_client
+
+    params = Mock()
+    mock_message = Mock()
+    mock_message.role = "user"
+    mock_message.content = Mock()
+    mock_message.content.text = "Test question"
+    params.messages = [mock_message]
+    params.temperature = 0.7
+    params.maxTokens = 256
+    params.stopSequences = ["STOP"]
+    params.systemPrompt = None
+    params.tools = None
+    params.toolChoice = None
+
+    result = await tool.sampling_callback(Mock(), params)
+
+    assert isinstance(result, types.CreateMessageResult)
+    call_kwargs = mock_chat_client.get_response.call_args
+    options = call_kwargs.kwargs.get("options") or {}
+    assert options.get("temperature") == 0.7
+    assert options.get("max_tokens") == 256
+    assert options.get("stop") == ["STOP"]
+    # These should not be passed as top-level kwargs
+    assert "temperature" not in call_kwargs.kwargs
+    assert "max_tokens" not in call_kwargs.kwargs
+    assert "stop" not in call_kwargs.kwargs
+
+
+async def test_mcp_tool_sampling_callback_omits_temperature_when_none():
+    """Test sampling callback does not set temperature in options when it is None."""
+    from agent_framework import Message
+
+    tool = MCPStdioTool(name="test_tool", command="python")
+
+    mock_chat_client = AsyncMock()
+    mock_response = Mock()
+    mock_response.messages = [Message(role="assistant", contents=[Content.from_text("response")])]
+    mock_response.model_id = "test-model"
+    mock_chat_client.get_response.return_value = mock_response
+
+    tool.client = mock_chat_client
+
+    params = Mock()
+    mock_message = Mock()
+    mock_message.role = "user"
+    mock_message.content = Mock()
+    mock_message.content.text = "Test question"
+    params.messages = [mock_message]
+    params.temperature = None
+    params.maxTokens = 100
+    params.stopSequences = None
+    params.systemPrompt = None
+    params.tools = None
+    params.toolChoice = None
+
+    result = await tool.sampling_callback(Mock(), params)
+
+    assert isinstance(result, types.CreateMessageResult)
+    call_kwargs = mock_chat_client.get_response.call_args
+    options = call_kwargs.kwargs.get("options") or {}
+    assert "temperature" not in options
+    assert options.get("max_tokens") == 100
+    assert "stop" not in options
+
+
+async def test_mcp_tool_sampling_callback_always_passes_max_tokens():
+    """Test sampling callback always sets max_tokens in options since maxTokens is a required int field."""
+    from agent_framework import Message
+
+    tool = MCPStdioTool(name="test_tool", command="python")
+
+    mock_chat_client = AsyncMock()
+    mock_response = Mock()
+    mock_response.messages = [Message(role="assistant", contents=[Content.from_text("response")])]
+    mock_response.model_id = "test-model"
+    mock_chat_client.get_response.return_value = mock_response
+
+    tool.client = mock_chat_client
+
+    params = Mock()
+    mock_message = Mock()
+    mock_message.role = "user"
+    mock_message.content = Mock()
+    mock_message.content.text = "Test question"
+    params.messages = [mock_message]
+    params.temperature = None
+    params.maxTokens = 200
+    params.stopSequences = None
+    params.systemPrompt = None
+    params.tools = None
+    params.toolChoice = None
+
+    result = await tool.sampling_callback(Mock(), params)
+
+    assert isinstance(result, types.CreateMessageResult)
+    call_kwargs = mock_chat_client.get_response.call_args
+    options = call_kwargs.kwargs.get("options") or {}
+    assert options["max_tokens"] == 200
+
+
+async def test_connect_sampling_capabilities_with_client():
+    """Test connect() passes sampling_capabilities to ClientSession when client is set."""
+    tool = MCPStdioTool(name="test", command="test-command", load_tools=False, load_prompts=False)
+    tool.client = Mock()
+
+    mock_transport = (Mock(), Mock())
+    mock_context_manager = Mock()
+    mock_context_manager.__aenter__ = AsyncMock(return_value=mock_transport)
+    mock_context_manager.__aexit__ = AsyncMock(return_value=None)
+    tool.get_mcp_client = Mock(return_value=mock_context_manager)
+
+    with patch("mcp.client.session.ClientSession") as mock_session_class:
+        mock_session = AsyncMock()
+        mock_session._request_id = 1
+
+        session_cm = AsyncMock()
+        session_cm.__aenter__ = AsyncMock(return_value=mock_session)
+        session_cm.__aexit__ = AsyncMock(return_value=None)
+        mock_session_class.return_value = session_cm
+
+        await tool.connect()
+
+        call_kwargs = mock_session_class.call_args.kwargs
+        sampling_caps = call_kwargs.get("sampling_capabilities")
+        assert sampling_caps is not None
+        assert isinstance(sampling_caps, types.SamplingCapability)
+        assert sampling_caps.tools is not None
+        assert isinstance(sampling_caps.tools, types.SamplingToolsCapability)
+
+
+async def test_connect_no_sampling_capabilities_without_client():
+    """Test connect() does not pass sampling_capabilities when no client is set."""
+    tool = MCPStdioTool(name="test", command="test-command", load_tools=False, load_prompts=False)
+    # No client set
+
+    mock_transport = (Mock(), Mock())
+    mock_context_manager = Mock()
+    mock_context_manager.__aenter__ = AsyncMock(return_value=mock_transport)
+    mock_context_manager.__aexit__ = AsyncMock(return_value=None)
+    tool.get_mcp_client = Mock(return_value=mock_context_manager)
+
+    with patch("mcp.client.session.ClientSession") as mock_session_class:
+        mock_session = AsyncMock()
+        mock_session._request_id = 1
+
+        session_cm = AsyncMock()
+        session_cm.__aenter__ = AsyncMock(return_value=mock_session)
+        session_cm.__aexit__ = AsyncMock(return_value=None)
+        mock_session_class.return_value = session_cm
+
+        await tool.connect()
+
+        call_kwargs = mock_session_class.call_args.kwargs
+        assert call_kwargs.get("sampling_capabilities") is None
+
+
 # Test error handling in connect() method
 
 
diff --git a/python/packages/openai/agent_framework_openai/_shared.py b/python/packages/openai/agent_framework_openai/_shared.py
index 900592da17..c3c280d950 100644
--- a/python/packages/openai/agent_framework_openai/_shared.py
+++ b/python/packages/openai/agent_framework_openai/_shared.py
@@ -216,7 +216,9 @@ def load_openai_service_settings(
                     openai_settings["model"] = resolved_model
                     break
 
-        if not openai_settings.get("api_version"):
+        if api_version is not None:
+            openai_settings["api_version"] = api_version
+        else:
             resolved_api_version = _get_setting_from_alias(
                 "AZURE_OPENAI_API_VERSION",
                 dotenv_values_by_name=dotenv_values_by_name,
diff --git a/python/packages/openai/tests/openai/conftest.py b/python/packages/openai/tests/openai/conftest.py
index d2a89a6742..1ef52baf81 100644
--- a/python/packages/openai/tests/openai/conftest.py
+++ b/python/packages/openai/tests/openai/conftest.py
@@ -48,6 +48,7 @@ def openai_unit_test_env(monkeypatch, exclude_list, override_env_param_dict):  #
             "OPENAI_AUDIO_TO_TEXT_MODEL_ID",
             "OPENAI_TEXT_TO_AUDIO_MODEL_ID",
             "OPENAI_REALTIME_MODEL_ID",
+            "OPENAI_API_VERSION",
             "OPENAI_BASE_URL",
             "AZURE_OPENAI_ENDPOINT",
             "AZURE_OPENAI_BASE_URL",
@@ -101,6 +102,7 @@ def azure_openai_unit_test_env(monkeypatch, exclude_list, override_env_param_dic
             "OPENAI_AUDIO_TO_TEXT_MODEL_ID",
             "OPENAI_TEXT_TO_AUDIO_MODEL_ID",
             "OPENAI_REALTIME_MODEL_ID",
+            "OPENAI_API_VERSION",
             "OPENAI_BASE_URL",
             "AZURE_OPENAI_ENDPOINT",
             "AZURE_OPENAI_BASE_URL",
diff --git a/python/packages/openai/tests/openai/test_openai_chat_completion_client_azure.py b/python/packages/openai/tests/openai/test_openai_chat_completion_client_azure.py
index eb27312684..148fcb68b9 100644
--- a/python/packages/openai/tests/openai/test_openai_chat_completion_client_azure.py
+++ b/python/packages/openai/tests/openai/test_openai_chat_completion_client_azure.py
@@ -79,7 +79,8 @@ def test_init_auto_detects_azure_env(azure_openai_unit_test_env: dict[str, str])
 
 
 @pytest.mark.parametrize("exclude_list", [["AZURE_OPENAI_API_VERSION"]], indirect=True)
-def test_init_uses_default_azure_api_version(azure_openai_unit_test_env: dict[str, str]) -> None:
+def test_init_uses_default_azure_api_version(monkeypatch, azure_openai_unit_test_env: dict[str, str]) -> None:
+    monkeypatch.setenv("OPENAI_API_VERSION", "preview")
     client = _create_azure_chat_completion_client()
 
     assert client.model == azure_openai_unit_test_env["AZURE_OPENAI_DEPLOYMENT_NAME"]