Python: stream tool results for approval-resolution execution

python/packages/core/agent_framework/_tools.py

@@ -2362,14 +2362,18 @@ async def _process_function_requests(
                         )
             _replace_approval_contents_with_results(prepped_messages, fcc_todo, approved_function_results)
             executed_count = sum(1 for r in approved_function_results if r.type == "function_result")
+            has_user_input_request = any(
+                r.type in {"function_approval_request", "function_call"} or r.user_input_request
+                for r in approved_function_results
+            )
             # Continue to call chat client with updated messages (containing function results)
             # so it can generate the final response
             return {
                 "action": "return" if should_terminate else "continue",
                 "errors_in_a_row": errors_in_a_row,
                 "result_message": None,
-                "update_role": None,
-                "function_call_results": None,
+                "update_role": "assistant" if has_user_input_request else ("tool" if executed_count else None),
+                "function_call_results": approved_function_results or None,
                 "function_call_count": executed_count,
             }
 
@@ -2777,6 +2781,13 @@ async def _stream() -> AsyncIterable[ChatResponseUpdate]:
                 errors_in_a_row = approval_result.get("errors_in_a_row", errors_in_a_row)
                 total_function_calls += approval_result.get("function_call_count", 0)
                 budget_state["total_function_calls"] = total_function_calls
+                if role := approval_result.get("update_role"):
+                    # Stream the results of tools executed while resolving approvals,
+                    # mirroring the gate after the second _process_function_requests call.
+                    yield ChatResponseUpdate(
+                        contents=approval_result.get("function_call_results") or [],
+                        role=role,
+                    )
                 if max_function_calls is not None and total_function_calls >= max_function_calls:
                     logger.info(
                         "Maximum function calls reached (%d/%d). Stopping further function calls for this request.",

python/packages/core/tests/core/test_harness_tool_approval.py

@@ -553,6 +553,128 @@ def second_streamed_tool() -> str:
     assert calls == 2
 
 
+async def test_tool_approval_resolution_streams_tool_results(
+    chat_client_base: MockBaseChatClient,
+) -> None:
+    """A tool executed while resolving an approval must have its result streamed."""
+    calls = 0
+
+    @tool(name="guarded_tool", approval_mode="always_require")
+    def guarded_tool() -> str:
+        nonlocal calls
+        calls += 1
+        return "guarded result"
+
+    agent = Agent(
+        client=chat_client_base,
+        tools=[guarded_tool],
+        middleware=[ToolApprovalMiddleware()],
+    )
+    session = AgentSession(session_id="approval-result-stream")
+    chat_client_base.streaming_responses = [
+        [
+            ChatResponseUpdate(
+                contents=[Content.from_function_call(call_id="call_guarded", name="guarded_tool", arguments="{}")],
+                role="assistant",
+            )
+        ]
+    ]
+
+    first_updates = [update async for update in agent.run("run guarded", stream=True, session=session)]
+    requests = [content for update in first_updates for content in update.user_input_requests]
+    assert len(requests) == 1
+    assert calls == 0
+
+    chat_client_base.streaming_responses = [
+        [ChatResponseUpdate(contents=[Content.from_text("done")], role="assistant")]
+    ]
+    second_updates = [
+        update
+        async for update in agent.run(
+            requests[0].to_function_approval_response(approved=True), stream=True, session=session
+        )
+    ]
+
+    tool_results = [
+        content
+        for update in second_updates
+        if update.role == "tool"
+        for content in update.contents
+        if content.type == "function_result"
+    ]
+    assert len(tool_results) == 1
+    assert tool_results[0].call_id == "call_guarded"
+    assert tool_results[0].result == "guarded result"
+    assert calls == 1
+
+
+async def test_tool_approval_resolution_streams_user_input_request_as_assistant(
+    chat_client_base: MockBaseChatClient,
+) -> None:
+    """A user-input request raised while resolving an approval must stream as an assistant update."""
+    from agent_framework.exceptions import UserInputRequiredException
+
+    calls = 0
+
+    @tool(name="oauth_tool", approval_mode="always_require")
+    def oauth_tool() -> str:
+        nonlocal calls
+        calls += 1
+        raise UserInputRequiredException(
+            contents=[Content.from_oauth_consent_request(consent_link="https://example.com/consent")]
+        )
+
+    agent = Agent(
+        client=chat_client_base,
+        tools=[oauth_tool],
+        middleware=[ToolApprovalMiddleware()],
+    )
+    session = AgentSession(session_id="approval-user-input-stream")
+    chat_client_base.streaming_responses = [
+        [
+            ChatResponseUpdate(
+                contents=[Content.from_function_call(call_id="call_oauth", name="oauth_tool", arguments="{}")],
+                role="assistant",
+            )
+        ]
+    ]
+
+    first_updates = [update async for update in agent.run("run oauth", stream=True, session=session)]
+    requests = [content for update in first_updates for content in update.user_input_requests]
+    assert len(requests) == 1
+    assert calls == 0
+
+    chat_client_base.streaming_responses = [
+        [ChatResponseUpdate(contents=[Content.from_text("done")], role="assistant")]
+    ]
+    second_updates = [
+        update
+        async for update in agent.run(
+            requests[0].to_function_approval_response(approved=True), stream=True, session=session
+        )
+    ]
+
+    assistant_requests = [
+        content
+        for update in second_updates
+        if update.role == "assistant"
+        for content in update.contents
+        if content.user_input_request
+    ]
+    assert len(assistant_requests) == 1
+    assert assistant_requests[0].type == "oauth_consent_request"
+    assert calls == 1
+
+    # it must not be mislabeled as a tool update
+    assert not [
+        content
+        for update in second_updates
+        if update.role == "tool"
+        for content in update.contents
+        if content.user_input_request
+    ]
+
+
 async def test_tool_approval_middleware_always_approve_tool_rule(
     chat_client_base: MockBaseChatClient,
 ) -> None: