Python: Handle agent user input request in AgentExecutor (#2022)

* Handle agent user input request in AgentExecutor * fix test * Address comments * Fix tests * Fix tests * Address comments * Address comments
2026-06-16 21:04:09 +08:00 · 2025-11-11 10:59:57 -08:00
parent 5fd2a0c287
commit cd9073aa11
6 changed files with 727 additions and 45 deletions
@@ -2,11 +2,14 @@

 import logging
 from dataclasses import dataclass
-from typing import Any
+from typing import Any, cast
+
+from agent_framework import FunctionApprovalRequestContent, FunctionApprovalResponseContent

 from .._agents import AgentProtocol, ChatAgent
 from .._threads import AgentThread
 from .._types import AgentRunResponse, AgentRunResponseUpdate, ChatMessage
+from ._checkpoint_encoding import decode_checkpoint_value, encode_checkpoint_value
 from ._conversation_state import encode_chat_messages
 from ._events import (
    AgentRunEvent,
@@ -14,6 +17,7 @@ from ._events import (
 )
 from ._executor import Executor, handler
 from ._message_utils import normalize_messages_input
+from ._request_info_mixin import response_handler
 from ._workflow_context import WorkflowContext

 logger = logging.getLogger(__name__)
@@ -83,6 +87,8 @@ class AgentExecutor(Executor):
        super().__init__(exec_id)
        self._agent = agent
        self._agent_thread = agent_thread or self._agent.get_new_thread()
+        self._pending_agent_requests: dict[str, FunctionApprovalRequestContent] = {}
+        self._pending_responses_to_agent: list[FunctionApprovalResponseContent] = []
        self._output_response = output_response
        self._cache: list[ChatMessage] = []

@@ -93,50 +99,6 @@ class AgentExecutor(Executor):
            return [AgentRunResponse]
        return []

-    async def _run_agent_and_emit(self, ctx: WorkflowContext[AgentExecutorResponse, AgentRunResponse]) -> None:
-        """Execute the underlying agent, emit events, and enqueue response.
-
-        Checks ctx.is_streaming() to determine whether to emit incremental AgentRunUpdateEvent
-        events (streaming mode) or a single AgentRunEvent (non-streaming mode).
-        """
-        if ctx.is_streaming():
-            # Streaming mode: emit incremental updates
-            updates: list[AgentRunResponseUpdate] = []
-            async for update in self._agent.run_stream(
-                self._cache,
-                thread=self._agent_thread,
-            ):
-                updates.append(update)
-                await ctx.add_event(AgentRunUpdateEvent(self.id, update))
-
-            if isinstance(self._agent, ChatAgent):
-                response_format = self._agent.chat_options.response_format
-                response = AgentRunResponse.from_agent_run_response_updates(
-                    updates,
-                    output_format_type=response_format,
-                )
-            else:
-                response = AgentRunResponse.from_agent_run_response_updates(updates)
-        else:
-            # Non-streaming mode: use run() and emit single event
-            response = await self._agent.run(
-                self._cache,
-                thread=self._agent_thread,
-            )
-            await ctx.add_event(AgentRunEvent(self.id, response))
-
-        if self._output_response:
-            await ctx.yield_output(response)
-
-        # Always construct a full conversation snapshot from inputs (cache)
-        # plus agent outputs (agent_run_response.messages). Do not mutate
-        # response.messages so AgentRunEvent remains faithful to the raw output.
-        full_conversation: list[ChatMessage] = list(self._cache) + list(response.messages)
-
-        agent_response = AgentExecutorResponse(self.id, response, full_conversation=full_conversation)
-        await ctx.send_message(agent_response)
-        self._cache.clear()
-
    @handler
    async def run(
        self, request: AgentExecutorRequest, ctx: WorkflowContext[AgentExecutorResponse, AgentRunResponse]
@@ -192,6 +154,31 @@ class AgentExecutor(Executor):
        self._cache = normalize_messages_input(messages)
        await self._run_agent_and_emit(ctx)

+    @response_handler
+    async def handle_user_input_response(
+        self,
+        original_request: FunctionApprovalRequestContent,
+        response: FunctionApprovalResponseContent,
+        ctx: WorkflowContext[AgentExecutorResponse, AgentRunResponse],
+    ) -> None:
+        """Handle user input responses for function approvals during agent execution.
+
+        This will hold the executor's execution until all pending user input requests are resolved.
+
+        Args:
+            original_request: The original function approval request sent by the agent.
+            response: The user's response to the function approval request.
+            ctx: The workflow context for emitting events and outputs.
+        """
+        self._pending_responses_to_agent.append(response)
+        self._pending_agent_requests.pop(original_request.id, None)
+
+        if not self._pending_agent_requests:
+            # All pending requests have been resolved; resume agent execution
+            self._cache = normalize_messages_input(ChatMessage(role="user", contents=self._pending_responses_to_agent))
+            self._pending_responses_to_agent.clear()
+            await self._run_agent_and_emit(ctx)
+
    async def snapshot_state(self) -> dict[str, Any]:
        """Capture current executor state for checkpointing.

@@ -226,6 +213,8 @@ class AgentExecutor(Executor):
        return {
            "cache": encode_chat_messages(self._cache),
            "agent_thread": serialized_thread,
+            "pending_agent_requests": encode_checkpoint_value(self._pending_agent_requests),
+            "pending_responses_to_agent": encode_checkpoint_value(self._pending_responses_to_agent),
        }

    async def restore_state(self, state: dict[str, Any]) -> None:
@@ -258,7 +247,109 @@ class AgentExecutor(Executor):
        else:
            self._agent_thread = self._agent.get_new_thread()

+        pending_requests_payload = state.get("pending_agent_requests")
+        if pending_requests_payload:
+            self._pending_agent_requests = decode_checkpoint_value(pending_requests_payload)
+
+        pending_responses_payload = state.get("pending_responses_to_agent")
+        if pending_responses_payload:
+            self._pending_responses_to_agent = decode_checkpoint_value(pending_responses_payload)
+
    def reset(self) -> None:
        """Reset the internal cache of the executor."""
        logger.debug("AgentExecutor %s: Resetting cache", self.id)
        self._cache.clear()
+
+    async def _run_agent_and_emit(self, ctx: WorkflowContext[AgentExecutorResponse, AgentRunResponse]) -> None:
+        """Execute the underlying agent, emit events, and enqueue response.
+
+        Checks ctx.is_streaming() to determine whether to emit incremental AgentRunUpdateEvent
+        events (streaming mode) or a single AgentRunEvent (non-streaming mode).
+        """
+        if ctx.is_streaming():
+            # Streaming mode: emit incremental updates
+            response = await self._run_agent_streaming(cast(WorkflowContext, ctx))
+        else:
+            # Non-streaming mode: use run() and emit single event
+            response = await self._run_agent(cast(WorkflowContext, ctx))
+
+        if response is None:
+            # Agent did not complete (e.g., waiting for user input); do not emit response
+            logger.info("AgentExecutor %s: Agent did not complete, awaiting user input", self.id)
+            return
+
+        if self._output_response:
+            await ctx.yield_output(response)
+
+        # Always construct a full conversation snapshot from inputs (cache)
+        # plus agent outputs (agent_run_response.messages). Do not mutate
+        # response.messages so AgentRunEvent remains faithful to the raw output.
+        full_conversation: list[ChatMessage] = list(self._cache) + list(response.messages)
+
+        agent_response = AgentExecutorResponse(self.id, response, full_conversation=full_conversation)
+        await ctx.send_message(agent_response)
+        self._cache.clear()
+
+    async def _run_agent(self, ctx: WorkflowContext) -> AgentRunResponse | None:
+        """Execute the underlying agent in non-streaming mode.
+
+        Args:
+            ctx: The workflow context for emitting events.
+
+        Returns:
+            The complete AgentRunResponse, or None if waiting for user input.
+        """
+        response = await self._agent.run(
+            self._cache,
+            thread=self._agent_thread,
+        )
+        await ctx.add_event(AgentRunEvent(self.id, response))
+
+        # Handle any user input requests
+        if response.user_input_requests:
+            for user_input_request in response.user_input_requests:
+                self._pending_agent_requests[user_input_request.id] = user_input_request
+                await ctx.request_info(user_input_request, FunctionApprovalResponseContent)
+            return None
+
+        return response
+
+    async def _run_agent_streaming(self, ctx: WorkflowContext) -> AgentRunResponse | None:
+        """Execute the underlying agent in streaming mode and collect the full response.
+
+        Args:
+            ctx: The workflow context for emitting events.
+
+        Returns:
+            The complete AgentRunResponse, or None if waiting for user input.
+        """
+        updates: list[AgentRunResponseUpdate] = []
+        user_input_requests: list[FunctionApprovalRequestContent] = []
+        async for update in self._agent.run_stream(
+            self._cache,
+            thread=self._agent_thread,
+        ):
+            updates.append(update)
+            await ctx.add_event(AgentRunUpdateEvent(self.id, update))
+
+            if update.user_input_requests:
+                user_input_requests.extend(update.user_input_requests)
+
+        # Build the final AgentRunResponse from the collected updates
+        if isinstance(self._agent, ChatAgent):
+            response_format = self._agent.chat_options.response_format
+            response = AgentRunResponse.from_agent_run_response_updates(
+                updates,
+                output_format_type=response_format,
+            )
+        else:
+            response = AgentRunResponse.from_agent_run_response_updates(updates)
+
+        # Handle any user input requests after the streaming completes
+        if user_input_requests:
+            for user_input_request in user_input_requests:
+                self._pending_agent_requests[user_input_request.id] = user_input_request
+                await ctx.request_info(user_input_request, FunctionApprovalResponseContent)
+            return None
+
+        return response
@@ -111,6 +111,10 @@ async def test_agent_executor_checkpoint_stores_and_restores_state() -> None:
    chat_store_state = thread_state["chat_message_store_state"]  # type: ignore[index]
    assert "messages" in chat_store_state, "Message store state should include messages"

+    # Verify checkpoint contains pending requests from agents and responses to be sent
+    assert "pending_agent_requests" in executor_state
+    assert "pending_responses_to_agent" in executor_state
+
    # Create a new agent and executor for restoration
    # This simulates starting from a fresh state and restoring from checkpoint
    restored_agent = _CountingAgent(id="test_agent", name="TestAgent")
@@ -5,19 +5,32 @@
 from collections.abc import AsyncIterable
 from typing import Any

+from typing_extensions import Never
+
 from agent_framework import (
    AgentExecutor,
+    AgentExecutorResponse,
    AgentRunResponse,
    AgentRunResponseUpdate,
    AgentRunUpdateEvent,
    AgentThread,
    BaseAgent,
+    ChatAgent,
    ChatMessage,
+    ChatResponse,
+    ChatResponseUpdate,
+    FunctionApprovalRequestContent,
    FunctionCallContent,
    FunctionResultContent,
+    RequestInfoEvent,
    Role,
    TextContent,
    WorkflowBuilder,
+    WorkflowContext,
+    WorkflowOutputEvent,
+    ai_function,
+    executor,
+    use_function_invocation,
 )


@@ -120,3 +133,235 @@ async def test_agent_executor_emits_tool_calls_in_streaming_mode() -> None:
    assert events[3].data is not None
    assert isinstance(events[3].data.contents[0], TextContent)
    assert "sunny" in events[3].data.contents[0].text
+
+
+@ai_function(approval_mode="always_require")
+def mock_tool_requiring_approval(query: str) -> str:
+    """Mock tool that requires approval before execution."""
+    return f"Executed tool with query: {query}"
+
+
+@use_function_invocation
+class MockChatClient:
+    """Simple implementation of a chat client."""
+
+    def __init__(self, parallel_request: bool = False) -> None:
+        self.additional_properties: dict[str, Any] = {}
+        self._iteration: int = 0
+        self._parallel_request: bool = parallel_request
+
+    async def get_response(
+        self,
+        messages: str | ChatMessage | list[str] | list[ChatMessage],
+        **kwargs: Any,
+    ) -> ChatResponse:
+        if self._iteration == 0:
+            if self._parallel_request:
+                response = ChatResponse(
+                    messages=ChatMessage(
+                        role="assistant",
+                        contents=[
+                            FunctionCallContent(
+                                call_id="1", name="mock_tool_requiring_approval", arguments='{"query": "test"}'
+                            ),
+                            FunctionCallContent(
+                                call_id="2", name="mock_tool_requiring_approval", arguments='{"query": "test"}'
+                            ),
+                        ],
+                    )
+                )
+            else:
+                response = ChatResponse(
+                    messages=ChatMessage(
+                        role="assistant",
+                        contents=[
+                            FunctionCallContent(
+                                call_id="1", name="mock_tool_requiring_approval", arguments='{"query": "test"}'
+                            )
+                        ],
+                    )
+                )
+        else:
+            response = ChatResponse(messages=ChatMessage(role="assistant", text="Tool executed successfully."))
+
+        self._iteration += 1
+        return response
+
+    async def get_streaming_response(
+        self,
+        messages: str | ChatMessage | list[str] | list[ChatMessage],
+        **kwargs: Any,
+    ) -> AsyncIterable[ChatResponseUpdate]:
+        if self._iteration == 0:
+            if self._parallel_request:
+                yield ChatResponseUpdate(
+                    contents=[
+                        FunctionCallContent(
+                            call_id="1", name="mock_tool_requiring_approval", arguments='{"query": "test"}'
+                        ),
+                        FunctionCallContent(
+                            call_id="2", name="mock_tool_requiring_approval", arguments='{"query": "test"}'
+                        ),
+                    ],
+                    role="assistant",
+                )
+            else:
+                yield ChatResponseUpdate(
+                    contents=[
+                        FunctionCallContent(
+                            call_id="1", name="mock_tool_requiring_approval", arguments='{"query": "test"}'
+                        )
+                    ],
+                    role="assistant",
+                )
+        else:
+            yield ChatResponseUpdate(text=TextContent(text="Tool executed "), role="assistant")
+            yield ChatResponseUpdate(contents=[TextContent(text="successfully.")], role="assistant")
+
+        self._iteration += 1
+
+
+@executor(id="test_executor")
+async def test_executor(agent_executor_response: AgentExecutorResponse, ctx: WorkflowContext[Never, str]) -> None:
+    await ctx.yield_output(agent_executor_response.agent_run_response.text)
+
+
+async def test_agent_executor_tool_call_with_approval() -> None:
+    """Test that AgentExecutor handles tool calls requiring approval."""
+    # Arrange
+    agent = ChatAgent(
+        chat_client=MockChatClient(),
+        name="ApprovalAgent",
+        tools=[mock_tool_requiring_approval],
+    )
+
+    workflow = WorkflowBuilder().set_start_executor(agent).add_edge(agent, test_executor).build()
+
+    # Act
+    events = await workflow.run("Invoke tool requiring approval")
+
+    # Assert
+    assert len(events.get_request_info_events()) == 1
+    approval_request = events.get_request_info_events()[0]
+    assert isinstance(approval_request.data, FunctionApprovalRequestContent)
+    assert approval_request.data.function_call.name == "mock_tool_requiring_approval"
+    assert approval_request.data.function_call.arguments == '{"query": "test"}'
+
+    # Act
+    events = await workflow.send_responses({approval_request.request_id: approval_request.data.create_response(True)})
+
+    # Assert
+    final_response = events.get_outputs()
+    assert len(final_response) == 1
+    assert final_response[0] == "Tool executed successfully."
+
+
+async def test_agent_executor_tool_call_with_approval_streaming() -> None:
+    """Test that AgentExecutor handles tool calls requiring approval in streaming mode."""
+    # Arrange
+    agent = ChatAgent(
+        chat_client=MockChatClient(),
+        name="ApprovalAgent",
+        tools=[mock_tool_requiring_approval],
+    )
+
+    workflow = WorkflowBuilder().set_start_executor(agent).add_edge(agent, test_executor).build()
+
+    # Act
+    request_info_events: list[RequestInfoEvent] = []
+    async for event in workflow.run_stream("Invoke tool requiring approval"):
+        if isinstance(event, RequestInfoEvent):
+            request_info_events.append(event)
+
+    # Assert
+    assert len(request_info_events) == 1
+    approval_request = request_info_events[0]
+    assert isinstance(approval_request.data, FunctionApprovalRequestContent)
+    assert approval_request.data.function_call.name == "mock_tool_requiring_approval"
+    assert approval_request.data.function_call.arguments == '{"query": "test"}'
+
+    # Act
+    output: str | None = None
+    async for event in workflow.send_responses_streaming({
+        approval_request.request_id: approval_request.data.create_response(True)
+    }):
+        if isinstance(event, WorkflowOutputEvent):
+            output = event.data
+
+    # Assert
+    assert output is not None
+    assert output == "Tool executed successfully."
+
+
+async def test_agent_executor_parallel_tool_call_with_approval() -> None:
+    """Test that AgentExecutor handles parallel tool calls requiring approval."""
+    # Arrange
+    agent = ChatAgent(
+        chat_client=MockChatClient(parallel_request=True),
+        name="ApprovalAgent",
+        tools=[mock_tool_requiring_approval],
+    )
+
+    workflow = WorkflowBuilder().set_start_executor(agent).add_edge(agent, test_executor).build()
+
+    # Act
+    events = await workflow.run("Invoke tool requiring approval")
+
+    # Assert
+    assert len(events.get_request_info_events()) == 2
+    for approval_request in events.get_request_info_events():
+        assert isinstance(approval_request.data, FunctionApprovalRequestContent)
+        assert approval_request.data.function_call.name == "mock_tool_requiring_approval"
+        assert approval_request.data.function_call.arguments == '{"query": "test"}'
+
+    # Act
+    responses = {
+        approval_request.request_id: approval_request.data.create_response(True)  # type: ignore
+        for approval_request in events.get_request_info_events()
+    }
+    events = await workflow.send_responses(responses)
+
+    # Assert
+    final_response = events.get_outputs()
+    assert len(final_response) == 1
+    assert final_response[0] == "Tool executed successfully."
+
+
+async def test_agent_executor_parallel_tool_call_with_approval_streaming() -> None:
+    """Test that AgentExecutor handles parallel tool calls requiring approval in streaming mode."""
+    # Arrange
+    agent = ChatAgent(
+        chat_client=MockChatClient(parallel_request=True),
+        name="ApprovalAgent",
+        tools=[mock_tool_requiring_approval],
+    )
+
+    workflow = WorkflowBuilder().set_start_executor(agent).add_edge(agent, test_executor).build()
+
+    # Act
+    request_info_events: list[RequestInfoEvent] = []
+    async for event in workflow.run_stream("Invoke tool requiring approval"):
+        if isinstance(event, RequestInfoEvent):
+            request_info_events.append(event)
+
+    # Assert
+    assert len(request_info_events) == 2
+    for approval_request in request_info_events:
+        assert isinstance(approval_request.data, FunctionApprovalRequestContent)
+        assert approval_request.data.function_call.name == "mock_tool_requiring_approval"
+        assert approval_request.data.function_call.arguments == '{"query": "test"}'
+
+    # Act
+    responses = {
+        approval_request.request_id: approval_request.data.create_response(True)  # type: ignore
+        for approval_request in request_info_events
+    }
+
+    output: str | None = None
+    async for event in workflow.send_responses_streaming(responses):
+        if isinstance(event, WorkflowOutputEvent):
+            output = event.data
+
+    # Assert
+    assert output is not None
+    assert output == "Tool executed successfully."