diff --git a/python/packages/core/agent_framework/_tools.py b/python/packages/core/agent_framework/_tools.py index 588ae93d5c0..0c316250605 100644 --- a/python/packages/core/agent_framework/_tools.py +++ b/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.", 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 f530483db65..26d3a3badc1 100644 --- a/python/packages/core/tests/core/test_harness_tool_approval.py +++ b/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: