From 6ee7a03bee9552b9068259586cc2a57527e223a0 Mon Sep 17 00:00:00 2001 From: Shivani767 <101629653+Shivani767@users.noreply.github.com> Date: Thu, 20 Aug 2026 23:54:18 +0530 Subject: [PATCH 1/2] Python: Preserve structured output on streaming harness tool-approval path Forward response_format through ToolApprovalMiddleware's stream finalizer so AgentResponse.value is parsed the same way as the non-streaming path. --- .../_harness/_tool_approval.py | 27 +++- .../tests/core/test_harness_tool_approval.py | 135 ++++++++++++++++++ 2 files changed, 161 insertions(+), 1 deletion(-) diff --git a/python/packages/core/agent_framework/_harness/_tool_approval.py b/python/packages/core/agent_framework/_harness/_tool_approval.py index 5c3bf7a2c9..6ac36a2b55 100644 --- a/python/packages/core/agent_framework/_harness/_tool_approval.py +++ b/python/packages/core/agent_framework/_harness/_tool_approval.py @@ -7,6 +7,7 @@ import json from asyncio import sleep from collections.abc import AsyncIterable, Awaitable, Callable, Iterable, Mapping, MutableMapping, Sequence +from functools import partial from typing import Any, Literal, cast from .._middleware import AgentContext, AgentMiddleware @@ -79,6 +80,24 @@ def _contents_from_state(values: Any) -> list[Content]: return [_content_from_state(value) for value in state_items] +def _structured_response_format(context: AgentContext) -> Any | None: + """Return the structured-output schema for this invocation, if any. + + Run ``options`` take precedence over the agent's ``default_options``. The + streaming re-wrap in :meth:`ToolApprovalMiddleware._process_stream` must + forward this to ``AgentResponse.from_updates`` or ``response.value`` is + dropped even when the inner stream already parsed it. + """ + if context.options is not None: + response_format = context.options.get("response_format") + if response_format is not None: + return response_format + default_options = getattr(context.agent, "default_options", None) + if isinstance(default_options, Mapping): + return default_options.get("response_format") + return None + + def _content_to_state(content: Content) -> dict[str, Any]: return content.to_dict() @@ -499,7 +518,13 @@ async def _stream() -> AsyncIterable[AgentResponseUpdate]: context.messages = [] context.result = None - return ResponseStream(_stream(), finalizer=AgentResponse.from_updates) + return ResponseStream( + _stream(), + finalizer=partial( + AgentResponse.from_updates, + output_format_type=_structured_response_format(context), + ), + ) def _prepare_inbound_messages( self, diff --git a/python/packages/core/tests/core/test_harness_tool_approval.py b/python/packages/core/tests/core/test_harness_tool_approval.py index b637bc751e..06688153b3 100644 --- a/python/packages/core/tests/core/test_harness_tool_approval.py +++ b/python/packages/core/tests/core/test_harness_tool_approval.py @@ -1426,3 +1426,138 @@ def optional_args_tool(value: str = "default") -> str: requests = _approval_requests(second_response.messages) assert [_function_call(request).arguments for request in requests] == ['{"value": "custom"}'] assert calls == 1 + + +@pytest.mark.parametrize("via", ["run_options", "default_options"], ids=["run-options", "default-options"]) +async def test_streaming_tool_approval_preserves_structured_value( + chat_client_base: MockBaseChatClient, + via: str, +) -> None: + """Streaming ToolApprovalMiddleware must forward response_format to the outer finalizer. + + Regression for https://github.com/microsoft/agent-framework/issues/7418: + re-wrapping with ``AgentResponse.from_updates`` dropped ``output_format_type``, + so ``response.value`` was None even when the inner stream had parsed it. + """ + from pydantic import BaseModel + + class Answer(BaseModel): + answer: str + + json_text = '{"answer": "42"}' + + @tool(name="echo", approval_mode="always_require") + def echo(text: str) -> str: + return text + + default_options = {"response_format": Answer} if via == "default_options" else None + run_options = {"response_format": Answer} if via == "run_options" else None + agent = Agent( + client=chat_client_base, + tools=[echo], + middleware=[ToolApprovalMiddleware()], + default_options=default_options, # type: ignore[arg-type] + ) + session = AgentSession(session_id=f"structured-stream-{via}") + chat_client_base.streaming_responses = [ + [ + ChatResponseUpdate( + role="assistant", + contents=[Content.from_text(json_text)], + finish_reason="stop", + ) + ] + ] + + stream = agent.run("return an Answer", stream=True, session=session, options=run_options) + async for _update in stream: + pass + response = await stream.get_final_response() + + assert response.text == json_text + assert isinstance(response.value, Answer) + assert response.value.answer == "42" + + +async def test_streaming_auto_approved_tool_preserves_structured_value( + chat_client_base: MockBaseChatClient, +) -> None: + """Auto-approved tool calls must still parse structured output on the streaming path.""" + from pydantic import BaseModel + + class Answer(BaseModel): + answer: str + + json_text = '{"answer": "42"}' + calls = 0 + + @tool(name="echo", approval_mode="always_require") + def echo(text: str) -> str: + nonlocal calls + calls += 1 + return text + + agent = Agent( + client=chat_client_base, + tools=[echo], + middleware=[ToolApprovalMiddleware(auto_approval_rules=[lambda function_call: True])], + ) + session = AgentSession(session_id="structured-stream-auto-approve") + function_call = Content.from_function_call(call_id="call_echo", name="echo", arguments='{"text": "hi"}') + chat_client_base.streaming_responses = [ + [ChatResponseUpdate(role="assistant", contents=[function_call])], + [ + ChatResponseUpdate( + role="assistant", + contents=[Content.from_text(json_text)], + finish_reason="stop", + ) + ], + ] + + stream = agent.run( + "Call echo and return an Answer.", + stream=True, + session=session, + options={"response_format": Answer}, + ) + async for _update in stream: + pass + response = await stream.get_final_response() + + assert calls == 1 + assert isinstance(response.value, Answer) + assert response.value.answer == "42" + + +async def test_non_streaming_tool_approval_preserves_structured_value( + chat_client_base: MockBaseChatClient, +) -> None: + """Non-streaming ToolApprovalMiddleware already returns the inner AgentResponse.""" + from pydantic import BaseModel + + class Answer(BaseModel): + answer: str + + json_text = '{"answer": "42"}' + + @tool(name="echo", approval_mode="always_require") + def echo(text: str) -> str: + return text + + agent = Agent( + client=chat_client_base, + tools=[echo], + middleware=[ToolApprovalMiddleware()], + ) + session = AgentSession(session_id="structured-non-stream") + chat_client_base.run_responses = [ChatResponse(messages=Message(role="assistant", contents=[json_text]))] + + response = await agent.run( + "return an Answer", + session=session, + options={"response_format": Answer}, + ) + + assert isinstance(response.value, Answer) + assert response.value.answer == "42" From 305164af72afe4941983f87cbcfbfefe30e8befa Mon Sep 17 00:00:00 2001 From: Shivani767 <101629653+Shivani767@users.noreply.github.com> Date: Fri, 21 Aug 2026 00:08:23 +0530 Subject: [PATCH 2/2] Python: Keep terminal structured value after auto-approved tool preamble Return the last inner AgentResponse from the streaming tool-approval finalizer so coalesced preamble text cannot clobber response.value. --- .../_harness/_tool_approval.py | 23 ++++--- .../tests/core/test_harness_tool_approval.py | 61 ++++++++++++++++++- 2 files changed, 74 insertions(+), 10 deletions(-) diff --git a/python/packages/core/agent_framework/_harness/_tool_approval.py b/python/packages/core/agent_framework/_harness/_tool_approval.py index 6ac36a2b55..8ab6a2caa3 100644 --- a/python/packages/core/agent_framework/_harness/_tool_approval.py +++ b/python/packages/core/agent_framework/_harness/_tool_approval.py @@ -7,7 +7,6 @@ import json from asyncio import sleep from collections.abc import AsyncIterable, Awaitable, Callable, Iterable, Mapping, MutableMapping, Sequence -from functools import partial from typing import Any, Literal, cast from .._middleware import AgentContext, AgentMiddleware @@ -457,6 +456,13 @@ def _process_stream( call_next: Callable[[], Awaitable[None]], state: ToolApprovalState, ) -> ResponseStream[AgentResponseUpdate, AgentResponse]: + # Last inner AgentResponse. Returning it from the outer finalizer matches + # the non-streaming path and preserves an already-parsed structured value + # even when earlier auto-approved turns yielded assistant text that + # AgentResponse.from_updates would coalesce into the JSON message. + holder: dict[str, AgentResponse | None] = {"final": None} + response_format = _structured_response_format(context) + async def _stream() -> AsyncIterable[AgentResponseUpdate]: if context.session is None: raise RuntimeError("ToolApprovalMiddleware requires an AgentSession.") @@ -503,7 +509,7 @@ async def _stream() -> AsyncIterable[AgentResponseUpdate]: additional_properties=update.additional_properties, raw_representation=update.raw_representation, ) - await context.result.get_final_response() + holder["final"] = await context.result.get_final_response() if not approval_requests: return @@ -518,13 +524,12 @@ async def _stream() -> AsyncIterable[AgentResponseUpdate]: context.messages = [] context.result = None - return ResponseStream( - _stream(), - finalizer=partial( - AgentResponse.from_updates, - output_format_type=_structured_response_format(context), - ), - ) + def _finalize(updates: Sequence[AgentResponseUpdate]) -> AgentResponse: + if holder["final"] is not None: + return holder["final"] + return AgentResponse.from_updates(updates, output_format_type=response_format) + + return ResponseStream(_stream(), finalizer=_finalize) def _prepare_inbound_messages( self, diff --git a/python/packages/core/tests/core/test_harness_tool_approval.py b/python/packages/core/tests/core/test_harness_tool_approval.py index 06688153b3..16fa04628c 100644 --- a/python/packages/core/tests/core/test_harness_tool_approval.py +++ b/python/packages/core/tests/core/test_harness_tool_approval.py @@ -1456,7 +1456,7 @@ def echo(text: str) -> str: client=chat_client_base, tools=[echo], middleware=[ToolApprovalMiddleware()], - default_options=default_options, # type: ignore[arg-type] + default_options=default_options, # type: ignore[arg-type, typeddict-item] # pyrefly: ignore[bad-argument-type] # ty: ignore[invalid-argument-type] ) session = AgentSession(session_id=f"structured-stream-{via}") chat_client_base.streaming_responses = [ @@ -1530,6 +1530,65 @@ def echo(text: str) -> str: assert response.value.answer == "42" +async def test_streaming_auto_approved_tool_preserves_value_when_preamble_text_coalesces( + chat_client_base: MockBaseChatClient, +) -> None: + """Preamble assistant text plus a tool call must not clobber the parsed structured value. + + Updates without ``message_id`` are coalesced by ``AgentResponse.from_updates``. If the + outer finalizer rebuilt from every yielded update, ``Calling echo…{"answer":"42"}`` + would fail to parse. The terminal inner response's value must be kept instead. + """ + from pydantic import BaseModel + + class Answer(BaseModel): + answer: str + + json_text = '{"answer": "42"}' + calls = 0 + + @tool(name="echo", approval_mode="always_require") + def echo(text: str) -> str: + nonlocal calls + calls += 1 + return text + + agent = Agent( + client=chat_client_base, + tools=[echo], + middleware=[ToolApprovalMiddleware(auto_approval_rules=[lambda function_call: True])], + ) + session = AgentSession(session_id="structured-stream-auto-approve-preamble") + function_call = Content.from_function_call(call_id="call_echo", name="echo", arguments='{"text": "hi"}') + chat_client_base.streaming_responses = [ + [ + ChatResponseUpdate(role="assistant", contents=[Content.from_text("Calling echo…")]), + ChatResponseUpdate(role="assistant", contents=[function_call]), + ], + [ + ChatResponseUpdate( + role="assistant", + contents=[Content.from_text(json_text)], + finish_reason="stop", + ) + ], + ] + + stream = agent.run( + "Call echo and return an Answer.", + stream=True, + session=session, + options={"response_format": Answer}, + ) + updates = [update async for update in stream] + response = await stream.get_final_response() + + assert any("Calling echo" in (content.text or "") for update in updates for content in update.contents) + assert calls == 1 + assert isinstance(response.value, Answer) + assert response.value.answer == "42" + + async def test_non_streaming_tool_approval_preserves_structured_value( chat_client_base: MockBaseChatClient, ) -> None: