From 2f684c4af26f97eb111de8b487d245730e916f77 Mon Sep 17 00:00:00 2001 From: Evan Mattson Date: Thu, 20 Aug 2026 15:56:47 +0900 Subject: [PATCH 1/7] Python: Fix AG-UI workflow-as-agent approval resumes (#7707) --- .../ag-ui/agent_framework_ag_ui/_agent_run.py | 18 +++- .../ag-ui/tests/ag_ui/test_endpoint.py | 84 +++++++++++++++++++ .../core/agent_framework/_workflows/_agent.py | 29 +++++-- .../_workflows/_agent_executor.py | 45 +++++----- .../agent_framework/_workflows/_workflow.py | 14 +++- .../_workflows/_workflow_executor.py | 5 +- .../tests/workflow/test_workflow_kwargs.py | 31 +++++++ 7 files changed, 194 insertions(+), 32 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 78e4b2a026b..0292e9a3620 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 @@ -36,6 +36,7 @@ MESSAGE_INJECTION_PENDING_MESSAGES_STATE_KEY, Message, SupportsAgentRun, + WorkflowAgent, ) from agent_framework._middleware import FunctionMiddlewarePipeline from agent_framework._tools import ( @@ -1192,6 +1193,7 @@ def _canonical_approval_resume_messages( lifecycle: ApprovalLifecycle, tools: list[Any] | None = None, has_deferred_owner: bool = False, + workflow_agent_owns_approval: bool = False, authorized_executions: dict[ApprovalOccurrenceIdentity, AuthorizedExecution] | None = None, retained_results: list[Content] | None = None, snapshot_reconciliations: list[ApprovalSnapshotReconciliation] | None = None, @@ -1382,9 +1384,15 @@ def _canonical_approval_resume_messages( original_arguments=pending_arguments, ) ) + response_id = interrupt_id + if workflow_agent_owns_approval: + response_id = next( + (alias for alias in pending_entry.aliases if alias != interrupt_id), + interrupt_id, + ) function_approvals = [ { - "id": interrupt_id, + "id": response_id, "call_id": pending_entry.identity.call_id, "name": _pending_approval_name(pending_entry) or "", "approved": accepted, @@ -2306,6 +2314,7 @@ async def run_agent_stream( client_tools = convert_agui_tools_to_agent_framework(input_data.get("tools")) server_tools = collect_server_tools(agent) tools = merge_tools(server_tools, client_tools) + workflow_agent_owns_approval = isinstance(agent, WorkflowAgent) approval_resume_messages, handled_resume_ids, cancelled_resume_ids, resume_error = ( _canonical_approval_resume_messages( resume_payload, @@ -2313,7 +2322,9 @@ async def run_agent_stream( expected_interrupt_ids=stored_pending_approval_interrupt_ids or None, lifecycle=approval_state_store.lifecycle, tools=tools, - has_deferred_owner=approval_state_store.has_tool_approval_state(approval_thread_id), + has_deferred_owner=approval_state_store.has_tool_approval_state(approval_thread_id) + or workflow_agent_owns_approval, + workflow_agent_owns_approval=workflow_agent_owns_approval, authorized_executions=authorized_executions, retained_results=retained_approval_results, snapshot_reconciliations=approval_snapshot_reconciliations, @@ -2626,7 +2637,8 @@ async def run_agent_stream( execution_owner = _function_call_execution_owner( content.function_call, tools, - has_deferred_owner=_TOOL_APPROVAL_STATE_KEY in session.state, + has_deferred_owner=_TOOL_APPROVAL_STATE_KEY in session.state + or workflow_agent_owns_approval, ) registration_kwargs = { "thread_ids": [approval_thread_id, provider_approval_thread_id], diff --git a/python/packages/ag-ui/tests/ag_ui/test_endpoint.py b/python/packages/ag-ui/tests/ag_ui/test_endpoint.py index 64b0270e73c..663e6e123b6 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_endpoint.py +++ b/python/packages/ag-ui/tests/ag_ui/test_endpoint.py @@ -454,6 +454,90 @@ async def approve( ) +async def test_endpoint_workflow_as_agent_resumes_with_client_tools() -> None: + """A workflow exposed through AgentFrameworkAgent accepts client tools across approval resume.""" + + class ApprovalExecutor(Executor): + def __init__(self) -> None: + super().__init__(id="approval") + + @handler + async def start(self, message: Any, ctx: WorkflowContext[Any, Any]) -> None: + del message + function_call = Content.from_function_call( + call_id="refund-call", + name="submit_refund", + arguments={"order_id": "12345", "amount": 89.99}, + ) + await ctx.request_info( + Content.from_function_approval_request(id="approval-1", function_call=function_call), + Content, + request_id="approval-1", + ) + + @response_handler + async def approve( + self, + original_request: Content, + response: Content, + ctx: WorkflowContext[Any, Any], + ) -> None: + del original_request + arguments = response.function_call.parse_arguments() if response.function_call is not None else None + await ctx.yield_output(json.dumps(arguments, sort_keys=True)) # type: ignore[arg-type] + + client_tool = { + "name": "submit_refund", + "description": "Submit a refund", + "parameters": { + "type": "object", + "properties": {"order_id": {"type": "string"}, "amount": {"type": "number"}}, + }, + } + app = FastAPI() + workflow = WorkflowBuilder(start_executor=ApprovalExecutor()).build() + wrapped_agent = AgentFrameworkAgent(agent=workflow.as_agent()) + add_agent_framework_fastapi_endpoint(app, wrapped_agent, path="/workflow-as-agent") + + with TestClient(app) as client: + pause_response = client.post( + "/workflow-as-agent", + json={ + "runId": "run-pause", + "threadId": "thread-client-tools", + "messages": [{"role": "user", "content": "Refund the order"}], + "tools": [client_tool], + }, + ) + pause_events = _decode_sse_events(pause_response) + assert not [event for event in pause_events if event.get("type") == "RUN_ERROR"] + pause_finished = [event for event in pause_events if event.get("type") == "RUN_FINISHED"] + assert _run_finished_interrupts(pause_finished[-1])[0]["id"] == "refund-call" + + resume_response = client.post( + "/workflow-as-agent", + json={ + "runId": "run-resume", + "threadId": "thread-client-tools", + "messages": [], + "tools": [client_tool], + "resume": [ + { + "interruptId": "refund-call", + "status": "resolved", + "payload": {"accepted": True, "amount": 49.5}, + } + ], + }, + ) + resume_events = _decode_sse_events(resume_response) + + assert not [event for event in resume_events if event.get("type") == "RUN_ERROR"] + assert '{"amount": 49.5, "order_id": "12345"}' == "".join( + str(event.get("delta", "")) for event in resume_events if event.get("type") == "TEXT_MESSAGE_CONTENT" + ) + + async def test_workflow_endpoint_applies_canonical_approval_edited_args() -> None: """Workflow approvals apply standard editedArgs as a full replacement.""" diff --git a/python/packages/core/agent_framework/_workflows/_agent.py b/python/packages/core/agent_framework/_workflows/_agent.py index 09bd592e922..345295adcb7 100644 --- a/python/packages/core/agent_framework/_workflows/_agent.py +++ b/python/packages/core/agent_framework/_workflows/_agent.py @@ -5,7 +5,7 @@ import logging import sys import uuid -from collections.abc import AsyncIterable, Awaitable, Mapping, Sequence +from collections.abc import AsyncIterable, Awaitable, Callable, Mapping, Sequence from dataclasses import dataclass from datetime import datetime, timezone from typing import TYPE_CHECKING, Any, ClassVar, Literal, cast, overload @@ -18,6 +18,7 @@ InMemoryHistoryProvider, SessionContext, ) +from .._tools import ToolTypes from .._types import ( AgentResponse, AgentResponseUpdate, @@ -155,6 +156,7 @@ def run( session: AgentSession | None = None, checkpoint_id: str | None = None, checkpoint_storage: CheckpointStorage | None = None, + tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None = None, function_invocation_kwargs: Mapping[str, Mapping[str, Any]] | Mapping[str, Any] | None = None, client_kwargs: Mapping[str, Mapping[str, Any]] | Mapping[str, Any] | None = None, ) -> ResponseStream[AgentResponseUpdate, AgentResponse]: ... @@ -168,6 +170,7 @@ async def run( session: AgentSession | None = None, checkpoint_id: str | None = None, checkpoint_storage: CheckpointStorage | None = None, + tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None = None, function_invocation_kwargs: Mapping[str, Mapping[str, Any]] | Mapping[str, Any] | None = None, client_kwargs: Mapping[str, Mapping[str, Any]] | Mapping[str, Any] | None = None, ) -> AgentResponse: ... @@ -180,6 +183,7 @@ def run( session: AgentSession | None = None, checkpoint_id: str | None = None, checkpoint_storage: CheckpointStorage | None = None, + tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None = None, function_invocation_kwargs: Mapping[str, Mapping[str, Any]] | Mapping[str, Any] | None = None, client_kwargs: Mapping[str, Mapping[str, Any]] | Mapping[str, Any] | None = None, ) -> ResponseStream[AgentResponseUpdate, AgentResponse] | Awaitable[AgentResponse]: @@ -198,6 +202,7 @@ def run( checkpoint_storage: Runtime checkpoint storage. When provided with checkpoint_id, used to load and restore the checkpoint. When provided without checkpoint_id, enables checkpointing for this run. + tools: Tools available to agents inside the workflow for this run. function_invocation_kwargs: Keyword arguments forwarded to tool invocations in subagents. Either a mapping of agent name/executor id to kwargs, or a flat mapping of kwargs for all tool invocations. @@ -224,6 +229,7 @@ def run( session, checkpoint_id, checkpoint_storage, + tools=tools, function_invocation_kwargs=function_invocation_kwargs, client_kwargs=client_kwargs, ), @@ -235,6 +241,7 @@ def run( session, checkpoint_id, checkpoint_storage, + tools=tools, function_invocation_kwargs=function_invocation_kwargs, client_kwargs=client_kwargs, ) @@ -246,6 +253,7 @@ async def _run_impl( session: AgentSession | None, checkpoint_id: str | None = None, checkpoint_storage: CheckpointStorage | None = None, + tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None = None, function_invocation_kwargs: Mapping[str, Mapping[str, Any]] | Mapping[str, Any] | None = None, client_kwargs: Mapping[str, Mapping[str, Any]] | Mapping[str, Any] | None = None, ) -> AgentResponse: @@ -257,6 +265,7 @@ async def _run_impl( session: The agent session for conversation context. checkpoint_id: ID of checkpoint to restore from. checkpoint_storage: Runtime checkpoint storage. + tools: Tools available to agents inside the workflow for this run. function_invocation_kwargs: Optional kwargs for tool invocations. client_kwargs: Optional kwargs for chat client calls. @@ -304,6 +313,7 @@ async def _run_impl( checkpoint_id, checkpoint_storage, streaming=False, + tools=tools, function_invocation_kwargs=function_invocation_kwargs, client_kwargs=client_kwargs, ): @@ -326,6 +336,7 @@ async def _run_stream_impl( session: AgentSession | None, checkpoint_id: str | None = None, checkpoint_storage: CheckpointStorage | None = None, + tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None = None, function_invocation_kwargs: Mapping[str, Mapping[str, Any]] | Mapping[str, Any] | None = None, client_kwargs: Mapping[str, Mapping[str, Any]] | Mapping[str, Any] | None = None, ) -> AsyncIterable[AgentResponseUpdate]: @@ -337,6 +348,7 @@ async def _run_stream_impl( session: The agent session for conversation context. checkpoint_id: ID of checkpoint to restore from. checkpoint_storage: Runtime checkpoint storage. + tools: Tools available to agents inside the workflow for this run. function_invocation_kwargs: Optional kwargs for tool invocations. client_kwargs: Optional kwargs for chat client calls. @@ -384,6 +396,7 @@ async def _run_stream_impl( checkpoint_id, checkpoint_storage, streaming=True, + tools=tools, function_invocation_kwargs=function_invocation_kwargs, client_kwargs=client_kwargs, ): @@ -405,6 +418,7 @@ async def _run_core( checkpoint_id: str | None, checkpoint_storage: CheckpointStorage | None, streaming: bool, + tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None = None, function_invocation_kwargs: Mapping[str, Mapping[str, Any]] | Mapping[str, Any] | None = None, client_kwargs: Mapping[str, Mapping[str, Any]] | Mapping[str, Any] | None = None, ) -> AsyncIterable[WorkflowEvent]: @@ -415,6 +429,7 @@ async def _run_core( checkpoint_id: ID of checkpoint to restore from. checkpoint_storage: Runtime checkpoint storage. streaming: Whether to use streaming workflow methods. + tools: Tools available to agents inside the workflow for this run. function_invocation_kwargs: Optional kwargs for tool invocations. client_kwargs: Optional kwargs for chat client calls. @@ -432,12 +447,14 @@ async def _run_core( stream=True, checkpoint_id=checkpoint_id, checkpoint_storage=checkpoint_storage, + tools=tools, ): pass else: _ = await self.workflow.run( checkpoint_id=checkpoint_id, checkpoint_storage=checkpoint_storage, + tools=tools, ) if not input_messages: logger.info("No input messages provided; the workflow has been restored to the checkpoint state.") @@ -459,6 +476,7 @@ async def _run_core( responses=function_responses, stream=True, checkpoint_storage=checkpoint_storage, + tools=tools, function_invocation_kwargs=function_invocation_kwargs, client_kwargs=client_kwargs, ): @@ -467,6 +485,7 @@ async def _run_core( for event in await self.workflow.run( responses=function_responses, checkpoint_storage=checkpoint_storage, + tools=tools, function_invocation_kwargs=function_invocation_kwargs, client_kwargs=client_kwargs, ): @@ -477,6 +496,7 @@ async def _run_core( message=input_messages, stream=True, checkpoint_storage=checkpoint_storage, + tools=tools, function_invocation_kwargs=function_invocation_kwargs, client_kwargs=client_kwargs, ): @@ -485,6 +505,7 @@ async def _run_core( for event in await self.workflow.run( message=input_messages, checkpoint_storage=checkpoint_storage, + tools=tools, function_invocation_kwargs=function_invocation_kwargs, client_kwargs=client_kwargs, ): @@ -713,11 +734,7 @@ def _process_request_info_event( Note: Text requests use the function-call envelope so callers can reply with a matching function result. """ - if ( - isinstance(event.data, Content) - and event.data.user_input_request - and event.data.type != "text" - ): + if isinstance(event.data, Content) and event.data.user_input_request and event.data.type != "text": # Preserve specialized requests that callers already understand how to present. return event.data diff --git a/python/packages/core/agent_framework/_workflows/_agent_executor.py b/python/packages/core/agent_framework/_workflows/_agent_executor.py index b7787736fa5..0627073ef58 100644 --- a/python/packages/core/agent_framework/_workflows/_agent_executor.py +++ b/python/packages/core/agent_framework/_workflows/_agent_executor.py @@ -410,9 +410,9 @@ async def _run_agent(self, ctx: WorkflowContext[Never, AgentResponse]) -> AgentR Returns: The complete AgentResponse, or None if waiting for user input. """ - function_invocation_kwargs, client_kwargs = self._prepare_agent_run_args( - ctx.get_state(WORKFLOW_RUN_KWARGS_KEY, {}) - ) + raw_run_kwargs = ctx.get_state(WORKFLOW_RUN_KWARGS_KEY, {}) + function_invocation_kwargs, client_kwargs = self._prepare_agent_run_args(raw_run_kwargs) + tools = raw_run_kwargs.get("tools") if not self._cache: logger.warning( @@ -422,13 +422,15 @@ async def _run_agent(self, ctx: WorkflowContext[Never, AgentResponse]) -> AgentR ) run_agent = cast(Callable[..., Awaitable[AgentResponse[Any]]], self._agent.run) - response = await run_agent( - self._cache, - stream=False, - session=self._session, - function_invocation_kwargs=function_invocation_kwargs, - client_kwargs=client_kwargs, - ) + run_kwargs: dict[str, Any] = { + "stream": False, + "session": self._session, + "function_invocation_kwargs": function_invocation_kwargs, + "client_kwargs": client_kwargs, + } + if tools is not None: + run_kwargs["tools"] = tools + response = await run_agent(self._cache, **run_kwargs) # Handle any user input requests if response.user_input_requests: @@ -464,9 +466,9 @@ async def _run_agent_streaming(self, ctx: WorkflowContext[Never, AgentResponseUp Returns: The complete AgentResponse, or None if waiting for user input. """ - function_invocation_kwargs, client_kwargs = self._prepare_agent_run_args( - ctx.get_state(WORKFLOW_RUN_KWARGS_KEY, {}) - ) + raw_run_kwargs = ctx.get_state(WORKFLOW_RUN_KWARGS_KEY, {}) + function_invocation_kwargs, client_kwargs = self._prepare_agent_run_args(raw_run_kwargs) + tools = raw_run_kwargs.get("tools") if not self._cache: logger.warning( @@ -478,13 +480,15 @@ async def _run_agent_streaming(self, ctx: WorkflowContext[Never, AgentResponseUp updates: list[AgentResponseUpdate] = [] streamed_user_input_requests: list[Content] = [] run_agent_stream = cast(Callable[..., ResponseStream[AgentResponseUpdate, AgentResponse[Any]]], self._agent.run) - stream = run_agent_stream( - self._cache, - stream=True, - session=self._session, - function_invocation_kwargs=function_invocation_kwargs, - client_kwargs=client_kwargs, - ) + run_kwargs: dict[str, Any] = { + "stream": True, + "session": self._session, + "function_invocation_kwargs": function_invocation_kwargs, + "client_kwargs": client_kwargs, + } + if tools is not None: + run_kwargs["tools"] = tools + stream = run_agent_stream(self._cache, **run_kwargs) async for update in stream: updates.append(update) if update.user_input_requests: @@ -562,7 +566,6 @@ def _prepare_agent_run_args( """ fi_resolved = raw_run_kwargs.get("function_invocation_kwargs") ci_resolved = raw_run_kwargs.get("client_kwargs") - function_invocation_kwargs = self._resolve_executor_kwargs(fi_resolved) client_kwargs = self._resolve_executor_kwargs(ci_resolved) diff --git a/python/packages/core/agent_framework/_workflows/_workflow.py b/python/packages/core/agent_framework/_workflows/_workflow.py index 04f3f9aa875..92eef4b1ffd 100644 --- a/python/packages/core/agent_framework/_workflows/_workflow.py +++ b/python/packages/core/agent_framework/_workflows/_workflow.py @@ -17,6 +17,7 @@ from typing import TYPE_CHECKING, Any, Literal, overload from .._sessions import ContextProvider +from .._tools import ToolTypes from .._types import Content, ResponseStream from ..exceptions import WorkflowException from ..observability import OtelAttr, capture_exception, create_workflow_span @@ -480,6 +481,7 @@ async def _run_workflow_with_tracing( initial_executor_fn: Callable[[], Awaitable[None]] | None = None, is_continuation: bool = False, streaming: bool = False, + tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None = None, function_invocation_kwargs: Mapping[str, Mapping[str, Any]] | Mapping[str, Any] | None = None, client_kwargs: Mapping[str, Mapping[str, Any]] | Mapping[str, Any] | None = None, ) -> AsyncIterable[WorkflowEvent]: @@ -498,6 +500,7 @@ async def _run_workflow_with_tracing( fresh-message runs reset them. Shared workflow state is preserved in both cases. streaming: Whether to enable streaming mode for agents. + tools: Runtime tools available to agent executors. function_invocation_kwargs: Optional kwargs to store in State for function invocations in subagents. client_kwargs: Optional kwargs to store in State for chat client @@ -550,8 +553,10 @@ async def _run_workflow_with_tracing( # - On a continuation (checkpoint restore or responses), the # prior run's kwargs are preserved unless the caller # explicitly provides new kwargs. - if function_invocation_kwargs is not None or client_kwargs is not None: + if function_invocation_kwargs is not None or client_kwargs is not None or tools is not None: combined_kwargs: dict[str, Any] = {} + if tools is not None: + combined_kwargs["tools"] = tools if function_invocation_kwargs is not None: combined_kwargs["function_invocation_kwargs"] = self._resolve_invocation_kwargs( function_invocation_kwargs, "function_invocation_kwargs" @@ -688,6 +693,7 @@ def run( responses: Mapping[str, Any] | None = None, checkpoint_id: str | None = None, checkpoint_storage: CheckpointStorage | None = None, + tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None = None, function_invocation_kwargs: Mapping[str, Any] | None = None, client_kwargs: Mapping[str, Any] | None = None, ) -> ResponseStream[WorkflowEvent, WorkflowRunResult]: ... @@ -702,6 +708,7 @@ def run( checkpoint_id: str | None = None, checkpoint_storage: CheckpointStorage | None = None, include_status_events: bool = False, + tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None = None, function_invocation_kwargs: Mapping[str, Any] | None = None, client_kwargs: Mapping[str, Any] | None = None, ) -> Awaitable[WorkflowRunResult]: ... @@ -715,6 +722,7 @@ def run( checkpoint_id: str | None = None, checkpoint_storage: CheckpointStorage | None = None, include_status_events: bool = False, + tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None = None, function_invocation_kwargs: Mapping[str, Mapping[str, Any]] | Mapping[str, Any] | None = None, client_kwargs: Mapping[str, Mapping[str, Any]] | Mapping[str, Any] | None = None, ) -> ResponseStream[WorkflowEvent, WorkflowRunResult] | Awaitable[WorkflowRunResult]: @@ -738,6 +746,7 @@ def run( (restore then send responses). checkpoint_storage: Runtime checkpoint storage. include_status_events: Whether to include status events (non-streaming only). + tools: Runtime tools available to agent executors. function_invocation_kwargs: Keyword arguments forwarded to tool invocations in subagents. Either a mapping for agent name or agent executor id to kwargs, or a flat mapping of kwargs for all tool invocations. @@ -783,6 +792,7 @@ def run( checkpoint_id=checkpoint_id, checkpoint_storage=checkpoint_storage, streaming=stream, + tools=tools, function_invocation_kwargs=function_invocation_kwargs, client_kwargs=client_kwargs, ), @@ -802,6 +812,7 @@ async def _run_core( checkpoint_id: str | None = None, checkpoint_storage: CheckpointStorage | None = None, streaming: bool = False, + tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None = None, function_invocation_kwargs: Mapping[str, Mapping[str, Any]] | Mapping[str, Any] | None = None, client_kwargs: Mapping[str, Mapping[str, Any]] | Mapping[str, Any] | None = None, ) -> AsyncIterable[WorkflowEvent]: @@ -877,6 +888,7 @@ async def _run_core( initial_executor_fn=initial_executor_fn, is_continuation=(message is None), streaming=streaming, + tools=tools, function_invocation_kwargs=function_invocation_kwargs, client_kwargs=client_kwargs, ): diff --git a/python/packages/core/agent_framework/_workflows/_workflow_executor.py b/python/packages/core/agent_framework/_workflows/_workflow_executor.py index 1a8f988d19b..f9b24cb9daa 100644 --- a/python/packages/core/agent_framework/_workflows/_workflow_executor.py +++ b/python/packages/core/agent_framework/_workflows/_workflow_executor.py @@ -381,6 +381,7 @@ async def process_workflow(self, input_data: object, ctx: WorkflowContext[Any, A # against the subworkflow's own executor IDs. fi_kwargs: dict[str, Any] | None = None ci_kwargs: dict[str, Any] | None = None + tools = parent_kwargs.get("tools") for key in ("function_invocation_kwargs", "client_kwargs"): resolved = parent_kwargs.get(key) if isinstance(resolved, dict): @@ -394,6 +395,7 @@ async def process_workflow(self, input_data: object, ctx: WorkflowContext[Any, A # Run the sub-workflow and collect all events, passing parent kwargs result = await self.workflow.run( input_data, + tools=tools, function_invocation_kwargs=fi_kwargs, # type: ignore client_kwargs=ci_kwargs, # type: ignore ) @@ -606,5 +608,6 @@ async def _handle_response( # Forward the response to the sub-workflow, which resumes and validates it against its own # pending requests, then process whatever the sub-workflow produces. - result = await self.workflow.run(responses={request_id: response}) + parent_kwargs: dict[str, Any] = ctx.get_state(WORKFLOW_RUN_KWARGS_KEY, {}) + result = await self.workflow.run(responses={request_id: response}, tools=parent_kwargs.get("tools")) await self._process_workflow_result(result, ctx) diff --git a/python/packages/core/tests/workflow/test_workflow_kwargs.py b/python/packages/core/tests/workflow/test_workflow_kwargs.py index 93c6c93d580..dc2d28b72c8 100644 --- a/python/packages/core/tests/workflow/test_workflow_kwargs.py +++ b/python/packages/core/tests/workflow/test_workflow_kwargs.py @@ -584,6 +584,19 @@ async def test_workflow_as_agent_run_propagates_kwargs_to_underlying_agent() -> assert received.get("function_invocation_kwargs") == fi_kwargs +async def test_workflow_as_agent_run_propagates_tools_to_underlying_agent() -> None: + """Tools passed to workflow.as_agent().run() reach the underlying agent executor.""" + agent = _KwargsCapturingAgent(name="inner_agent") + workflow = SequentialBuilder(participants=[agent]).build() + workflow_agent = workflow.as_agent(name="TestWorkflowAgent") + + client_tools = [object()] + _ = await workflow_agent.run("test message", tools=client_tools) + + assert len(agent.captured_kwargs) >= 1, "Inner agent should have been invoked at least once" + assert agent.captured_kwargs[0].get("tools") is client_tools + + async def test_workflow_as_agent_run_stream_propagates_kwargs_to_underlying_agent() -> None: """Test that function_invocation_kwargs passed to workflow_agent.run(stream=True) flow through.""" agent = _KwargsCapturingAgent(name="inner_agent") @@ -695,6 +708,24 @@ async def test_subworkflow_kwargs_propagation() -> None: ) +async def test_subworkflow_tools_propagation() -> None: + """Tools passed to a parent workflow reach agents inside a nested workflow.""" + from agent_framework._workflows._workflow_executor import WorkflowExecutor + + inner_agent = _KwargsCapturingAgent(name="inner_agent") + inner_workflow = SequentialBuilder(participants=[inner_agent]).build() + subworkflow_executor = WorkflowExecutor(workflow=inner_workflow, id="subworkflow_executor") + outer_workflow = SequentialBuilder(participants=[subworkflow_executor]).build() + + client_tools = [object()] + async for event in outer_workflow.run("test message for subworkflow", stream=True, tools=client_tools): + if event.type == "status" and event.state == WorkflowRunState.IDLE: + break + + assert len(inner_agent.captured_kwargs) >= 1, "Inner agent in subworkflow should have been invoked" + assert inner_agent.captured_kwargs[0].get("tools") is client_tools + + async def test_subworkflow_kwargs_accessible_via_state() -> None: """Test that kwargs are accessible via State within subworkflow. From 260dc00c40b48a03cce78c098fdc85b93decf185 Mon Sep 17 00:00:00 2001 From: Evan Mattson Date: Thu, 20 Aug 2026 16:53:43 +0900 Subject: [PATCH 2/7] fix workflow agent review issues --- .../ag-ui/agent_framework_ag_ui/_agent_run.py | 31 +++- .../_approval_lifecycle.py | 6 + .../agent_framework_ag_ui/_approval_state.py | 1 + .../agent_framework_ag_ui/_workflow_run.py | 3 + .../ag-ui/tests/ag_ui/test_endpoint.py | 170 ++++++++++++++++++ .../core/agent_framework/_workflows/_agent.py | 50 ++++-- .../_workflows/_agent_executor.py | 19 +- .../_workflows/_runner_context.py | 25 +++ .../agent_framework/_workflows/_workflow.py | 6 +- .../_workflows/_workflow_context.py | 4 + .../_workflows/_workflow_executor.py | 5 +- .../test_agent_executor_tool_calls.py | 32 ++++ .../tests/workflow/test_workflow_kwargs.py | 152 +++++++++++++++- 13 files changed, 467 insertions(+), 37 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 0292e9a3620..6539925f6bf 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 @@ -74,6 +74,7 @@ from ._message_adapters import normalize_agui_input_messages from ._predictive_state import PredictiveStateHandler from ._tooling import collect_server_tools, merge_tools +from ._workflow_run import _consume_cancelled_workflow_requests # pyright: ignore[reportPrivateUsage] from ._run_common import ( FlowState, _approval_interrupt_for_function_call, # type: ignore @@ -880,6 +881,7 @@ def _register_server_generated_approval_response( name=response.function_call.name, arguments=arguments, aliases=[str(response.function_call.call_id)] if response.function_call.call_id else None, + response_id=str(response_id), server_label=_function_call_server_label(response.function_call), ) if response.approved is not True: @@ -1193,7 +1195,6 @@ def _canonical_approval_resume_messages( lifecycle: ApprovalLifecycle, tools: list[Any] | None = None, has_deferred_owner: bool = False, - workflow_agent_owns_approval: bool = False, authorized_executions: dict[ApprovalOccurrenceIdentity, AuthorizedExecution] | None = None, retained_results: list[Content] | None = None, snapshot_reconciliations: list[ApprovalSnapshotReconciliation] | None = None, @@ -1384,12 +1385,7 @@ def _canonical_approval_resume_messages( original_arguments=pending_arguments, ) ) - response_id = interrupt_id - if workflow_agent_owns_approval: - response_id = next( - (alias for alias in pending_entry.aliases if alias != interrupt_id), - interrupt_id, - ) + response_id = pending_entry.response_id or interrupt_id function_approvals = [ { "id": response_id, @@ -1430,6 +1426,7 @@ def _canonical_approval_resume_messages( name=function_call.name, arguments=sibling_arguments, aliases=[sibling_call_id], + response_id=str(response_id), server_label=_function_call_server_label(function_call), ) lifecycle_decisions.append( @@ -1540,6 +1537,7 @@ async def _resolve_approval_responses( valid_response_content_ids: set[int] | None = None pending_local_response_content_ids: set[int] | None = None validated_forwarded_approvals: list[Content] = [] + deferred_response_content_ids: set[int] = set() response_content_ids_to_strip: set[int] = set() valid_response_content_ids = set() pending_local_response_content_ids = set() @@ -1640,6 +1638,8 @@ async def _resolve_approval_responses( continue intents_by_response_content_id[id(primary_response)] = intent valid_response_content_ids.add(id(primary_response)) + if pending_entry.owner is ApprovalExecutionOwner.DEFERRED: + deferred_response_content_ids.add(id(primary_response)) if ( primary_response.approved is True and intent is not None @@ -1708,7 +1708,7 @@ async def forward_hosted_decision(approval: Content = approval) -> list[Content] fcc_todo = { response_id: response for response_id, response in fcc_todo.items() - if id(response) in valid_response_content_ids + if id(response) in valid_response_content_ids and id(response) not in deferred_response_content_ids } if not fcc_todo: return [] @@ -2324,7 +2324,6 @@ async def run_agent_stream( tools=tools, has_deferred_owner=approval_state_store.has_tool_approval_state(approval_thread_id) or workflow_agent_owns_approval, - workflow_agent_owns_approval=workflow_agent_owns_approval, authorized_executions=authorized_executions, retained_results=retained_approval_results, snapshot_reconciliations=approval_snapshot_reconciliations, @@ -2353,6 +2352,20 @@ async def run_agent_stream( await snapshot_session.clear_interrupts(interrupt_ids=retired_interrupt_ids or cancelled_resume_ids or None) yield resume_error return + if cancelled_resume_ids and isinstance(agent, WorkflowAgent): + cancelled_workflow_entries: list[dict[str, Any]] = [] + for interrupt_id in cancelled_resume_ids: + occurrence = approval_state_store.lifecycle.occurrence_for_alias( + thread_id=approval_thread_id, + interrupt_id=interrupt_id, + ) + cancelled_workflow_entries.append( + { + "interrupt_id": occurrence.response_id if occurrence and occurrence.response_id else interrupt_id, + "status": "cancelled", + } + ) + _consume_cancelled_workflow_requests(agent.workflow, cancelled_workflow_entries) if cancelled_resume_ids and handled_resume_ids == cancelled_resume_ids: yield RunStartedEvent(run_id=run_id, thread_id=thread_id) _clear_tool_approval_state(approval_state_store, approval_thread_id) diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_approval_lifecycle.py b/python/packages/ag-ui/agent_framework_ag_ui/_approval_lifecycle.py index c792b5465bd..7defdef5f70 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_approval_lifecycle.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_approval_lifecycle.py @@ -181,6 +181,7 @@ class ApprovalOccurrence: owner: ApprovalExecutionOwner scope: str | None = None aliases: tuple[str, ...] = () + response_id: str | None = None already_approved_requests: tuple[dict[str, Any], ...] = () server_label: str | None = None idempotency_key: str | None = None @@ -276,6 +277,7 @@ def register( name: str, arguments: str, aliases: list[str] | None = None, + response_id: str | None = None, already_approved_requests: list[dict[str, Any]] | None = None, server_label: str | None = None, idempotency_key: str | None = None, @@ -293,6 +295,7 @@ def register( owner=owner, scope=scope, aliases=aliases, + response_id=response_id, already_approved_requests=already_approved_requests, server_label=server_label, idempotency_key=idempotency_key, @@ -310,6 +313,7 @@ def _register_aliases( owner: ApprovalExecutionOwner, scope: str | None = None, aliases: list[str] | None = None, + response_id: str | None = None, already_approved_requests: list[dict[str, Any]] | None = None, server_label: str | None = None, idempotency_key: str | None = None, @@ -342,6 +346,7 @@ def _register_aliases( or occurrence.arguments != arguments or occurrence.owner is not owner or occurrence.scope != scope + or occurrence.response_id != response_id or occurrence.idempotency_key != idempotency_key or occurrence.already_approved_requests != tuple(already_approved_requests or ()) or occurrence.server_label != server_label @@ -371,6 +376,7 @@ def _register_aliases( owner=owner, scope=scope, aliases=occurrence_aliases, + response_id=response_id, already_approved_requests=tuple(already_approved_requests or ()), server_label=server_label, idempotency_key=idempotency_key, diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_approval_state.py b/python/packages/ag-ui/agent_framework_ag_ui/_approval_state.py index 6ed72c9c049..1a5bc7b2a11 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_approval_state.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_approval_state.py @@ -99,6 +99,7 @@ def register( name=name, arguments=arguments, aliases=[request_id], + response_id=request_id, already_approved_requests=already_approved_requests, server_label=server_label, ) diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_workflow_run.py b/python/packages/ag-ui/agent_framework_ag_ui/_workflow_run.py index a01851d1dce..888ee2637ea 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_workflow_run.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_workflow_run.py @@ -362,6 +362,9 @@ def _consume_cancelled_workflow_requests(workflow: Workflow, resume_entries: lis if isinstance(pending_agent_requests, dict): cast(dict[str, Any], pending_agent_requests).pop(interrupt_id, None) + if not pending_events: + workflow._status = WorkflowRunState.IDLE # pyright: ignore[reportPrivateUsage] + def _coerce_json_value(value: Any) -> Any: """Parse JSON strings when possible; otherwise return the original value.""" diff --git a/python/packages/ag-ui/tests/ag_ui/test_endpoint.py b/python/packages/ag-ui/tests/ag_ui/test_endpoint.py index 663e6e123b6..53207cb30de 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_endpoint.py +++ b/python/packages/ag-ui/tests/ag_ui/test_endpoint.py @@ -538,6 +538,176 @@ async def approve( ) +async def test_endpoint_workflow_as_agent_cancellation_allows_next_turn() -> None: + """Cancelling a wrapped workflow approval consumes its pending request correlation.""" + + class ApprovalExecutor(Executor): + def __init__(self) -> None: + super().__init__(id="approval") + self.run_count = 0 + + @handler + async def start(self, message: Any, ctx: WorkflowContext[Any, Any]) -> None: + del message + self.run_count += 1 + if self.run_count > 1: + await ctx.yield_output("Follow-up completed.") # type: ignore[arg-type] + return + function_call = Content.from_function_call( + call_id="refund-call", + name="submit_refund", + arguments={"order_id": "12345"}, + ) + await ctx.request_info( + Content.from_function_approval_request(id="approval-1", function_call=function_call), + Content, + request_id="approval-1", + ) + + @response_handler + async def approve( + self, + original_request: Content, + response: Content, + ctx: WorkflowContext[Any, Any], + ) -> None: + del original_request, response + await ctx.yield_output("Approval resolved.") # type: ignore[arg-type] + + client_tool = { + "name": "submit_refund", + "description": "Submit a refund", + "parameters": { + "type": "object", + "properties": {"order_id": {"type": "string"}}, + }, + } + app = FastAPI() + workflow = WorkflowBuilder(start_executor=ApprovalExecutor()).build() + add_agent_framework_fastapi_endpoint(app, AgentFrameworkAgent(agent=workflow.as_agent()), path="/workflow-agent") + + with TestClient(app) as client: + pause_response = client.post( + "/workflow-agent", + json={ + "runId": "run-pause", + "threadId": "thread-cancel", + "messages": [{"role": "user", "content": "Refund the order"}], + "tools": [client_tool], + }, + ) + pause_events = _decode_sse_events(pause_response) + assert not [event for event in pause_events if event.get("type") == "RUN_ERROR"] + + cancel_response = client.post( + "/workflow-agent", + json={ + "runId": "run-cancel", + "threadId": "thread-cancel", + "messages": [], + "tools": [client_tool], + "resume": [{"interruptId": "refund-call", "status": "cancelled"}], + }, + ) + cancel_events = _decode_sse_events(cancel_response) + assert not [event for event in cancel_events if event.get("type") == "RUN_ERROR"] + + follow_up_response = client.post( + "/workflow-agent", + json={ + "runId": "run-follow-up", + "threadId": "thread-cancel", + "messages": [{"role": "user", "content": "Continue without the refund"}], + "tools": [client_tool], + }, + ) + follow_up_events = _decode_sse_events(follow_up_response) + + assert not [event for event in follow_up_events if event.get("type") == "RUN_ERROR"] + assert "Follow-up completed." == "".join( + str(event.get("delta", "")) for event in follow_up_events if event.get("type") == "TEXT_MESSAGE_CONTENT" + ) + + +async def test_endpoint_workflow_as_agent_rejection_reaches_response_handler() -> None: + """A rejected deferred approval remains typed until the wrapped workflow consumes it.""" + + class ApprovalExecutor(Executor): + def __init__(self) -> None: + super().__init__(id="approval") + + @handler + async def start(self, message: Any, ctx: WorkflowContext[Any, Any]) -> None: + del message + function_call = Content.from_function_call( + call_id="refund-call", + name="submit_refund", + arguments={"order_id": "12345"}, + ) + await ctx.request_info( + Content.from_function_approval_request(id="approval-1", function_call=function_call), + Content, + request_id="approval-1", + ) + + @response_handler + async def approve( + self, + original_request: Content, + response: Content, + ctx: WorkflowContext[Any, Any], + ) -> None: + del original_request + await ctx.yield_output(f"{response.type}:{response.approved}") # type: ignore[arg-type] + + client_tool = { + "name": "submit_refund", + "description": "Submit a refund", + "parameters": { + "type": "object", + "properties": {"order_id": {"type": "string"}}, + }, + } + app = FastAPI() + workflow = WorkflowBuilder(start_executor=ApprovalExecutor()).build() + add_agent_framework_fastapi_endpoint(app, AgentFrameworkAgent(agent=workflow.as_agent()), path="/workflow-agent") + + with TestClient(app) as client: + pause_response = client.post( + "/workflow-agent", + json={ + "runId": "run-pause", + "threadId": "thread-reject", + "messages": [{"role": "user", "content": "Refund the order"}], + "tools": [client_tool], + }, + ) + assert not [event for event in _decode_sse_events(pause_response) if event.get("type") == "RUN_ERROR"] + + reject_response = client.post( + "/workflow-agent", + json={ + "runId": "run-reject", + "threadId": "thread-reject", + "messages": [], + "tools": [client_tool], + "resume": [ + { + "interruptId": "refund-call", + "status": "resolved", + "payload": {"accepted": False}, + } + ], + }, + ) + reject_events = _decode_sse_events(reject_response) + + assert not [event for event in reject_events if event.get("type") == "RUN_ERROR"] + assert "function_approval_response:False" == "".join( + str(event.get("delta", "")) for event in reject_events if event.get("type") == "TEXT_MESSAGE_CONTENT" + ) + + async def test_workflow_endpoint_applies_canonical_approval_edited_args() -> None: """Workflow approvals apply standard editedArgs as a full replacement.""" diff --git a/python/packages/core/agent_framework/_workflows/_agent.py b/python/packages/core/agent_framework/_workflows/_agent.py index 345295adcb7..f913667c108 100644 --- a/python/packages/core/agent_framework/_workflows/_agent.py +++ b/python/packages/core/agent_framework/_workflows/_agent.py @@ -152,28 +152,28 @@ def run( self, messages: AgentRunInputs | None = None, *, - stream: Literal[True], + stream: Literal[False] = ..., session: AgentSession | None = None, checkpoint_id: str | None = None, checkpoint_storage: CheckpointStorage | None = None, tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None = None, - function_invocation_kwargs: Mapping[str, Mapping[str, Any]] | Mapping[str, Any] | None = None, - client_kwargs: Mapping[str, Mapping[str, Any]] | Mapping[str, Any] | None = None, - ) -> ResponseStream[AgentResponseUpdate, AgentResponse]: ... + function_invocation_kwargs: Mapping[str, Any] | None = None, + client_kwargs: Mapping[str, Any] | None = None, + ) -> Awaitable[AgentResponse[Any]]: ... @overload - async def run( + def run( self, messages: AgentRunInputs | None = None, *, - stream: Literal[False] = ..., + stream: Literal[True], session: AgentSession | None = None, checkpoint_id: str | None = None, checkpoint_storage: CheckpointStorage | None = None, tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None = None, - function_invocation_kwargs: Mapping[str, Mapping[str, Any]] | Mapping[str, Any] | None = None, - client_kwargs: Mapping[str, Mapping[str, Any]] | Mapping[str, Any] | None = None, - ) -> AgentResponse: ... + function_invocation_kwargs: Mapping[str, Any] | None = None, + client_kwargs: Mapping[str, Any] | None = None, + ) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ... def run( self, @@ -184,9 +184,9 @@ def run( checkpoint_id: str | None = None, checkpoint_storage: CheckpointStorage | None = None, tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None = None, - function_invocation_kwargs: Mapping[str, Mapping[str, Any]] | Mapping[str, Any] | None = None, - client_kwargs: Mapping[str, Mapping[str, Any]] | Mapping[str, Any] | None = None, - ) -> ResponseStream[AgentResponseUpdate, AgentResponse] | Awaitable[AgentResponse]: + function_invocation_kwargs: Mapping[str, Any] | None = None, + client_kwargs: Mapping[str, Any] | None = None, + ) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]] | Awaitable[AgentResponse[Any]]: """Get a response from the workflow agent. Args: @@ -470,7 +470,8 @@ async def _run_core( # NOTE: It is possible that some pending requests are not fulfilled, # and we will let the workflow to handle this -- the agent does not # have an opinion on this. - function_responses = self._extract_function_responses(input_messages) + pending_requests = await self.workflow._runner_context.get_pending_request_info_events() # pyright: ignore[reportPrivateUsage] + function_responses = self._extract_function_responses(input_messages, pending_requests) if streaming: async for event in self.workflow.run( responses=function_responses, @@ -747,22 +748,37 @@ def _process_request_info_event( arguments=args, ) - def _extract_function_responses(self, input_messages: Sequence[Message]) -> dict[str, Any]: + def _extract_function_responses( + self, + input_messages: Sequence[Message], + pending_requests: Mapping[str, WorkflowEvent[Any]] | None = None, + ) -> dict[str, Any]: """Extract function responses from input messages. The responses are for pending requests that the workflow is waiting on, and will be passed to the workflow. The pending requests are processed to either `function_approval_request` or `function_call` content by `_process_request_info_event`. """ + pending_requests = pending_requests or {} function_responses: dict[str, Any] = {} for message in input_messages: for content in message.contents: if content.type == "function_approval_response": - request_id: str = content.id # type: ignore[assignment] + request_id = content.id + if request_id is None: + raise AgentInvalidResponseException("Function approval response is missing its request ID.") function_responses[request_id] = content elif content.type == "function_result": - response_data = content.result if hasattr(content, "result") else str(content) - function_responses[content.call_id] = response_data # type: ignore + request_id = content.call_id + if request_id is None: + raise AgentInvalidResponseException("Function result is missing its call ID.") + pending_request = pending_requests.get(request_id) + response_data = ( + content + if pending_request is not None and pending_request.response_type is Content + else content.result + ) + function_responses[request_id] = response_data else: raise AgentInvalidResponseException( "Unexpected content type while awaiting request info responses." diff --git a/python/packages/core/agent_framework/_workflows/_agent_executor.py b/python/packages/core/agent_framework/_workflows/_agent_executor.py index 0627073ef58..3bb92682f92 100644 --- a/python/packages/core/agent_framework/_workflows/_agent_executor.py +++ b/python/packages/core/agent_framework/_workflows/_agent_executor.py @@ -1,5 +1,6 @@ # Copyright (c) Microsoft. All rights reserved. +import inspect import logging import sys from collections.abc import Awaitable, Callable @@ -29,6 +30,15 @@ logger = logging.getLogger(__name__) +def _accepts_runtime_tools(agent: SupportsAgentRun) -> bool: + """Return whether the agent run surface accepts a tools keyword.""" + try: + parameters = inspect.signature(agent.run).parameters.values() + except (TypeError, ValueError): + return False + return any(parameter.name == "tools" or parameter.kind is inspect.Parameter.VAR_KEYWORD for parameter in parameters) + + @dataclass class AgentExecutorRequest: """A request to an agent executor. @@ -166,6 +176,7 @@ def __init__( raise ValueError("Agent must have a non-empty name or id or an explicit id must be provided.") super().__init__(exec_id) self._agent = agent + self._accepts_runtime_tools = _accepts_runtime_tools(agent) self._session = session or self._agent.create_session() self._pending_agent_requests: dict[str, Content] = {} @@ -412,7 +423,7 @@ async def _run_agent(self, ctx: WorkflowContext[Never, AgentResponse]) -> AgentR """ raw_run_kwargs = ctx.get_state(WORKFLOW_RUN_KWARGS_KEY, {}) function_invocation_kwargs, client_kwargs = self._prepare_agent_run_args(raw_run_kwargs) - tools = raw_run_kwargs.get("tools") + tools = ctx.get_runtime_tools() if not self._cache: logger.warning( @@ -428,7 +439,7 @@ async def _run_agent(self, ctx: WorkflowContext[Never, AgentResponse]) -> AgentR "function_invocation_kwargs": function_invocation_kwargs, "client_kwargs": client_kwargs, } - if tools is not None: + if tools is not None and self._accepts_runtime_tools: run_kwargs["tools"] = tools response = await run_agent(self._cache, **run_kwargs) @@ -468,7 +479,7 @@ async def _run_agent_streaming(self, ctx: WorkflowContext[Never, AgentResponseUp """ raw_run_kwargs = ctx.get_state(WORKFLOW_RUN_KWARGS_KEY, {}) function_invocation_kwargs, client_kwargs = self._prepare_agent_run_args(raw_run_kwargs) - tools = raw_run_kwargs.get("tools") + tools = ctx.get_runtime_tools() if not self._cache: logger.warning( @@ -486,7 +497,7 @@ async def _run_agent_streaming(self, ctx: WorkflowContext[Never, AgentResponseUp "function_invocation_kwargs": function_invocation_kwargs, "client_kwargs": client_kwargs, } - if tools is not None: + if tools is not None and self._accepts_runtime_tools: run_kwargs["tools"] = tools stream = run_agent_stream(self._cache, **run_kwargs) async for update in stream: diff --git a/python/packages/core/agent_framework/_workflows/_runner_context.py b/python/packages/core/agent_framework/_workflows/_runner_context.py index f70b9075f77..f16e74b9ca4 100644 --- a/python/packages/core/agent_framework/_workflows/_runner_context.py +++ b/python/packages/core/agent_framework/_workflows/_runner_context.py @@ -201,6 +201,18 @@ def is_streaming(self) -> bool: """ ... + def set_runtime_tools(self, tools: Any | None) -> None: + """Set request-scoped tools for the active workflow run.""" + ... + + def get_runtime_tools(self) -> Any | None: + """Get request-scoped tools for the active workflow run.""" + ... + + def clear_runtime_tools(self) -> None: + """Clear request-scoped tools after the active workflow run.""" + ... + async def build_checkpoint( self, workflow_name: str, @@ -339,6 +351,7 @@ def __init__(self, checkpoint_storage: CheckpointStorage | None = None): # Streaming flag - set by workflow's run(..., stream=True) vs run(..., stream=False) self._streaming: bool = False + self._runtime_tools: Any | None = None self._yield_output_classifier: YieldOutputClassifier = lambda _executor_id: "output" # region Messaging and Events @@ -531,6 +544,18 @@ def is_streaming(self) -> bool: """ return self._streaming + def set_runtime_tools(self, tools: Any | None) -> None: + """Set request-scoped tools for the active workflow run.""" + self._runtime_tools = tools + + def get_runtime_tools(self) -> Any | None: + """Get request-scoped tools for the active workflow run.""" + return self._runtime_tools + + def clear_runtime_tools(self) -> None: + """Clear request-scoped tools after the active workflow run.""" + self._runtime_tools = None + async def add_request_info_event(self, event: WorkflowEvent[Any]) -> None: """Add a request_info event to the context and track it for correlation. diff --git a/python/packages/core/agent_framework/_workflows/_workflow.py b/python/packages/core/agent_framework/_workflows/_workflow.py index 92eef4b1ffd..1e933644c26 100644 --- a/python/packages/core/agent_framework/_workflows/_workflow.py +++ b/python/packages/core/agent_framework/_workflows/_workflow.py @@ -553,10 +553,8 @@ async def _run_workflow_with_tracing( # - On a continuation (checkpoint restore or responses), the # prior run's kwargs are preserved unless the caller # explicitly provides new kwargs. - if function_invocation_kwargs is not None or client_kwargs is not None or tools is not None: + if function_invocation_kwargs is not None or client_kwargs is not None: combined_kwargs: dict[str, Any] = {} - if tools is not None: - combined_kwargs["tools"] = tools if function_invocation_kwargs is not None: combined_kwargs["function_invocation_kwargs"] = self._resolve_invocation_kwargs( function_invocation_kwargs, "function_invocation_kwargs" @@ -784,6 +782,7 @@ def run( # its async-generator finalizer ran. Clear it so this run starts clean and does # not silently inherit the prior run's runtime checkpoint storage. self._runner.context.clear_runtime_checkpoint_storage() + self._runner.context.set_runtime_tools(tools) response_stream = ResponseStream[WorkflowEvent, WorkflowRunResult]( self._run_core( @@ -921,6 +920,7 @@ async def _run_core( # deferred finalizer can't clear a successor's storage. if checkpoint_storage is not None: self._runner.context.clear_runtime_checkpoint_storage() + self._runner.context.clear_runtime_tools() @staticmethod def _finalize_events( diff --git a/python/packages/core/agent_framework/_workflows/_workflow_context.py b/python/packages/core/agent_framework/_workflows/_workflow_context.py index 18a53e0bfa6..08128f22afd 100644 --- a/python/packages/core/agent_framework/_workflows/_workflow_context.py +++ b/python/packages/core/agent_framework/_workflows/_workflow_context.py @@ -487,3 +487,7 @@ def is_streaming(self) -> bool: True if the workflow was started with stream=True, False otherwise. """ return self._runner_context.is_streaming() + + def get_runtime_tools(self) -> Any | None: + """Get request-scoped tools supplied to the active workflow run.""" + return self._runner_context.get_runtime_tools() diff --git a/python/packages/core/agent_framework/_workflows/_workflow_executor.py b/python/packages/core/agent_framework/_workflows/_workflow_executor.py index f9b24cb9daa..8ab9d6d6ab9 100644 --- a/python/packages/core/agent_framework/_workflows/_workflow_executor.py +++ b/python/packages/core/agent_framework/_workflows/_workflow_executor.py @@ -381,7 +381,7 @@ async def process_workflow(self, input_data: object, ctx: WorkflowContext[Any, A # against the subworkflow's own executor IDs. fi_kwargs: dict[str, Any] | None = None ci_kwargs: dict[str, Any] | None = None - tools = parent_kwargs.get("tools") + tools = ctx.get_runtime_tools() for key in ("function_invocation_kwargs", "client_kwargs"): resolved = parent_kwargs.get(key) if isinstance(resolved, dict): @@ -608,6 +608,5 @@ async def _handle_response( # Forward the response to the sub-workflow, which resumes and validates it against its own # pending requests, then process whatever the sub-workflow produces. - parent_kwargs: dict[str, Any] = ctx.get_state(WORKFLOW_RUN_KWARGS_KEY, {}) - result = await self.workflow.run(responses={request_id: response}, tools=parent_kwargs.get("tools")) + result = await self.workflow.run(responses={request_id: response}, tools=ctx.get_runtime_tools()) await self._process_workflow_result(result, ctx) diff --git a/python/packages/core/tests/workflow/test_agent_executor_tool_calls.py b/python/packages/core/tests/workflow/test_agent_executor_tool_calls.py index 17b9332a28f..717f28019d6 100644 --- a/python/packages/core/tests/workflow/test_agent_executor_tool_calls.py +++ b/python/packages/core/tests/workflow/test_agent_executor_tool_calls.py @@ -424,6 +424,7 @@ def __init__(self, parallel_request: bool = False) -> None: BaseChatClient.__init__(self) self._iteration: int = 0 self._parallel_request: bool = parallel_request + self.received_messages: list[list[Message]] = [] def _inner_get_response( self, @@ -433,6 +434,7 @@ def _inner_get_response( options: Mapping[str, Any], **kwargs: Any, ) -> Awaitable[ChatResponse] | ResponseStream[ChatResponseUpdate, ChatResponse]: + self.received_messages.append(list(messages)) if stream: return self._build_response_stream(self._stream_response()) @@ -534,6 +536,36 @@ async def test_agent_executor_declaration_only_tool_emits_request_info() -> None assert final_response[0] == "Tool executed successfully." +async def test_workflow_agent_preserves_structured_declaration_only_tool_result() -> None: + """WorkflowAgent keeps a client-tool result typed and tool-role for the owning AgentExecutor.""" + client = DeclarationOnlyMockChatClient() + agent = Agent( + client=client, + name="DeclarationOnlyAgent", + tools=[declaration_only_tool], + ) + workflow = WorkflowBuilder(start_executor=agent).build() + workflow_agent = workflow.as_agent() + + paused = await workflow_agent.run("Use the client side tool") + [request] = paused.user_input_requests + assert request.call_id is not None + + resumed = await workflow_agent.run( + Message( + role="tool", + contents=[Content.from_function_result(call_id=request.call_id, result={"answer": 42})], + ) + ) + + assert resumed.text == "Tool executed successfully." + tool_messages = [message for message in client.received_messages[-1] if message.role == "tool"] + assert len(tool_messages) == 1 + [result] = tool_messages[0].contents + assert result.type == "function_result" + assert result.result == '{"answer": 42}' + + async def test_agent_executor_declaration_only_tool_emits_request_info_streaming() -> None: """Test that AgentExecutor emits request_info for declaration-only tools in streaming mode.""" agent = Agent( diff --git a/python/packages/core/tests/workflow/test_workflow_kwargs.py b/python/packages/core/tests/workflow/test_workflow_kwargs.py index dc2d28b72c8..5a6fe8d81cd 100644 --- a/python/packages/core/tests/workflow/test_workflow_kwargs.py +++ b/python/packages/core/tests/workflow/test_workflow_kwargs.py @@ -1,6 +1,6 @@ # Copyright (c) Microsoft. All rights reserved. -from collections.abc import AsyncIterable, Awaitable +from collections.abc import AsyncIterable, Awaitable, Mapping from typing import TYPE_CHECKING, Any, Literal, overload import pytest @@ -597,6 +597,61 @@ async def test_workflow_as_agent_run_propagates_tools_to_underlying_agent() -> N assert agent.captured_kwargs[0].get("tools") is client_tools +async def test_workflow_tools_do_not_break_exact_supports_agent_run_signature() -> None: + """Runtime tools do not reach an agent whose documented run contract has no tools keyword.""" + + class ExactSignatureAgent(BaseAgent): + @overload + def run( + self, + messages: AgentRunInputs | None = ..., + *, + stream: Literal[False] = ..., + session: AgentSession | None = ..., + function_invocation_kwargs: Mapping[str, Any] | None = ..., + client_kwargs: Mapping[str, Any] | None = ..., + ) -> Awaitable[AgentResponse[Any]]: ... + + @overload + def run( + self, + messages: AgentRunInputs | None = ..., + *, + stream: Literal[True], + session: AgentSession | None = ..., + function_invocation_kwargs: Mapping[str, Any] | None = ..., + client_kwargs: Mapping[str, Any] | None = ..., + ) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ... + + def run( + self, + messages: AgentRunInputs | None = None, + *, + stream: bool = False, + session: AgentSession | None = None, + function_invocation_kwargs: Mapping[str, Any] | None = None, + client_kwargs: Mapping[str, Any] | None = None, + ) -> Awaitable[AgentResponse[Any]] | ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: + del messages, session, function_invocation_kwargs, client_kwargs + if stream: + + async def _stream() -> AsyncIterable[AgentResponseUpdate]: + yield AgentResponseUpdate(contents=[Content.from_text(text="exact response")]) + + return ResponseStream(_stream(), finalizer=AgentResponse.from_updates) + + async def _run() -> AgentResponse[Any]: + return AgentResponse(messages=[Message(role="assistant", contents=["exact response"])]) + + return _run() + + workflow = SequentialBuilder(participants=[ExactSignatureAgent(name="exact")]).build() + + result = await workflow.run("test message", tools=[object()]) + + assert "exact response" in str(result.get_outputs()[0]) + + async def test_workflow_as_agent_run_stream_propagates_kwargs_to_underlying_agent() -> None: """Test that function_invocation_kwargs passed to workflow_agent.run(stream=True) flow through.""" agent = _KwargsCapturingAgent(name="inner_agent") @@ -659,6 +714,61 @@ async def test_workflow_as_agent_kwargs_with_complex_nested_data() -> None: assert received.get("function_invocation_kwargs") == complex_data +async def test_continuation_tools_preserve_existing_invocation_kwargs() -> None: + """Supplying request-scoped tools on resume does not replace stored invocation kwargs.""" + from agent_framework import Executor, WorkflowBuilder, WorkflowContext, handler, response_handler + + captured_run_kwargs: list[dict[str, Any]] = [] + + class RequestingExecutor(Executor): + @handler + async def start(self, messages: list[Message], ctx: WorkflowContext[Any, Any]) -> None: + del messages + await ctx.request_info("Continue?", str, request_id="continue-request") + + @response_handler + async def resume(self, request: str, response: str, ctx: WorkflowContext[Any, Any]) -> None: + del request, response + captured_run_kwargs.append(ctx.get_state(WORKFLOW_RUN_KWARGS_KEY, {})) + await ctx.yield_output("resumed") # type: ignore[arg-type] + + workflow = WorkflowBuilder(start_executor=RequestingExecutor(id="requester")).build() + function_kwargs = {"api_key": "secret"} + client_kwargs = {"model": "test-model"} + + _ = await workflow.run( + [Message(role="user", contents=["start"])], + function_invocation_kwargs=function_kwargs, + client_kwargs=client_kwargs, + ) + _ = await workflow.run(responses={"continue-request": "yes"}, tools=[object()]) + + assert captured_run_kwargs == [ + { + "function_invocation_kwargs": {"__global__": function_kwargs}, + "client_kwargs": {"__global__": client_kwargs}, + } + ] + + +async def test_runtime_tools_are_not_written_to_checkpoints(tmp_path) -> None: + """Request-scoped tools do not enter serialized workflow state.""" + from agent_framework import FileCheckpointStorage + + agent = _KwargsCapturingAgent(name="checkpoint_agent") + workflow = SequentialBuilder(participants=[agent]).build() + storage = FileCheckpointStorage(tmp_path) + runtime_tool = lambda: None # noqa: E731 + + _ = await workflow.run("checkpoint tools", tools=[runtime_tool], checkpoint_storage=storage) + + checkpoints = await storage.list_checkpoints(workflow_name=workflow.name) + assert checkpoints + for checkpoint in checkpoints: + run_kwargs = checkpoint.state.get(WORKFLOW_RUN_KWARGS_KEY, {}) + assert "tools" not in run_kwargs + + # endregion @@ -726,6 +836,46 @@ async def test_subworkflow_tools_propagation() -> None: assert inner_agent.captured_kwargs[0].get("tools") is client_tools +async def test_subworkflow_resume_tools_preserve_child_invocation_kwargs() -> None: + """Nested workflow resume tools do not replace the child's stored invocation kwargs.""" + from agent_framework import Executor, WorkflowBuilder, WorkflowContext, handler, response_handler + from agent_framework._workflows._workflow_executor import WorkflowExecutor + + captured_child_kwargs: list[dict[str, Any]] = [] + + class RequestingExecutor(Executor): + @handler + async def start(self, messages: list[Message], ctx: WorkflowContext[Any, Any]) -> None: + del messages + await ctx.request_info("Continue?", str, request_id="child-request") + + @response_handler + async def resume(self, request: str, response: str, ctx: WorkflowContext[Any, Any]) -> None: + del request, response + captured_child_kwargs.append(ctx.get_state(WORKFLOW_RUN_KWARGS_KEY, {})) + await ctx.yield_output("child resumed") # type: ignore[arg-type] + + child = WorkflowBuilder(start_executor=RequestingExecutor(id="child-requester")).build() + child_executor = WorkflowExecutor(child, id="child", propagate_request=True) + parent = WorkflowBuilder(start_executor=child_executor).build() + function_kwargs = {"api_key": "secret"} + client_kwargs = {"model": "test-model"} + + _ = await parent.run( + [Message(role="user", contents=["start"])], + function_invocation_kwargs=function_kwargs, + client_kwargs=client_kwargs, + ) + _ = await parent.run(responses={"child-request": "yes"}, tools=[object()]) + + assert captured_child_kwargs == [ + { + "function_invocation_kwargs": {"__global__": function_kwargs}, + "client_kwargs": {"__global__": client_kwargs}, + } + ] + + async def test_subworkflow_kwargs_accessible_via_state() -> None: """Test that kwargs are accessible via State within subworkflow. From c7c402d3780af40a4372eb24297f5e5c6aa995be Mon Sep 17 00:00:00 2001 From: Evan Mattson Date: Thu, 20 Aug 2026 17:11:10 +0900 Subject: [PATCH 3/7] fix workflow agent generic input responses --- python/packages/core/agent_framework/_workflows/_agent.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/python/packages/core/agent_framework/_workflows/_agent.py b/python/packages/core/agent_framework/_workflows/_agent.py index f913667c108..1aed8c0596c 100644 --- a/python/packages/core/agent_framework/_workflows/_agent.py +++ b/python/packages/core/agent_framework/_workflows/_agent.py @@ -775,7 +775,10 @@ def _extract_function_responses( pending_request = pending_requests.get(request_id) response_data = ( content - if pending_request is not None and pending_request.response_type is Content + if pending_request is not None + and pending_request.response_type is Content + and isinstance(pending_request.data, Content) + and pending_request.data.type == "function_call" else content.result ) function_responses[request_id] = response_data From 11ea19d2cafac22cf8cd610fae9df1c47cafc36d Mon Sep 17 00:00:00 2001 From: Evan Mattson Date: Thu, 20 Aug 2026 17:44:37 +0900 Subject: [PATCH 4/7] fix nested workflow approval cancellation --- .../agent_framework_ag_ui/_workflow_run.py | 6 ++ .../ag-ui/tests/ag_ui/test_endpoint.py | 99 +++++++++++++++++++ 2 files changed, 105 insertions(+) diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_workflow_run.py b/python/packages/ag-ui/agent_framework_ag_ui/_workflow_run.py index 888ee2637ea..4702b0d613d 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_workflow_run.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_workflow_run.py @@ -358,6 +358,12 @@ def _consume_cancelled_workflow_requests(workflow: Workflow, resume_entries: lis request_event = pending_events.pop(interrupt_id, None) source_executor_id = getattr(request_event, "source_executor_id", None) executor = workflow.executors.get(source_executor_id) if source_executor_id else None + nested_workflow = getattr(executor, "workflow", None) + if isinstance(nested_workflow, Workflow): + _consume_cancelled_workflow_requests( + nested_workflow, + [{"interrupt_id": interrupt_id, "status": "cancelled"}], + ) pending_agent_requests = getattr(executor, "_pending_agent_requests", None) if isinstance(pending_agent_requests, dict): cast(dict[str, Any], pending_agent_requests).pop(interrupt_id, None) diff --git a/python/packages/ag-ui/tests/ag_ui/test_endpoint.py b/python/packages/ag-ui/tests/ag_ui/test_endpoint.py index 53207cb30de..c712e577dac 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_endpoint.py +++ b/python/packages/ag-ui/tests/ag_ui/test_endpoint.py @@ -32,6 +32,7 @@ ToolApprovalMiddleware, WorkflowBuilder, WorkflowContext, + WorkflowExecutor, executor, handler, response_handler, @@ -629,6 +630,104 @@ async def approve( ) +async def test_workflow_endpoint_nested_mixed_approval_resume(streaming_chat_client_stub) -> None: + """Cancelling one nested approval does not block its approved sibling.""" + + def function_call(order_id: str, call_id: str) -> Content: + return Content.from_function_call( + call_id=call_id, + name="submit_refund", + arguments={"order_id": order_id}, + ) + + call_count = 0 + executed_orders: list[str] = [] + + def submit_refund(order_id: str) -> str: + executed_orders.append(order_id) + return f"Refunded {order_id}" + + async def stream_fn(messages: Any, options: Any, **kwargs: Any) -> AsyncIterator[ChatResponseUpdate]: + nonlocal call_count + del messages, options, kwargs + call_count += 1 + if call_count == 1: + yield ChatResponseUpdate( + contents=[ + function_call("order-1", "refund-call-1"), + function_call("order-2", "refund-call-2"), + ] + ) + return + yield ChatResponseUpdate(contents=[Content.from_text(text="Approved sibling completed.")]) + + child_agent = Agent( + name="nested-agent", + client=streaming_chat_client_stub(stream_fn), + tools=[ + FunctionTool( + name="submit_refund", + description="Submit a refund", + func=submit_refund, + approval_mode="always_require", + ) + ], + ) + child_workflow = WorkflowBuilder(start_executor=child_agent).build() + parent_workflow = WorkflowBuilder( + start_executor=WorkflowExecutor( + child_workflow, + id="nested-workflow", + propagate_request=True, + allow_direct_output=True, + ) + ).build() + app = FastAPI() + add_agent_framework_fastapi_endpoint( + app, + parent_workflow, + path="/nested-workflow", + ) + + with TestClient(app) as client: + pause_response = client.post( + "/nested-workflow", + json={ + "runId": "run-pause", + "threadId": "thread-nested-mixed", + "messages": [{"role": "user", "content": "Refund both orders"}], + }, + ) + pause_events = _decode_sse_events(pause_response) + pause_finished = [event for event in pause_events if event.get("type") == "RUN_FINISHED"] + pause_interrupts = _run_finished_interrupts(pause_finished[-1]) + assert {interrupt["id"] for interrupt in pause_interrupts} == { + "refund-call-1", + "refund-call-2", + }, pause_events + + resume_response = client.post( + "/nested-workflow", + json={ + "runId": "run-resume", + "threadId": "thread-nested-mixed", + "messages": [], + "resume": [ + {"interruptId": "refund-call-1", "status": "cancelled"}, + {"interruptId": "refund-call-2", "status": "resolved", "payload": {"approved": True}}, + ], + }, + ) + + resume_events = _decode_sse_events(resume_response) + assert not [event for event in resume_events if event.get("type") == "RUN_ERROR"] + assert "Approved sibling completed." == "".join( + str(event.get("delta", "")) for event in resume_events if event.get("type") == "TEXT_MESSAGE_CONTENT" + ) + assert call_count == 2 + assert executed_orders == ["order-2"] + + async def test_endpoint_workflow_as_agent_rejection_reaches_response_handler() -> None: """A rejected deferred approval remains typed until the wrapped workflow consumes it.""" From 80a4851b2401d7fda2a6b560540812591008e24e Mon Sep 17 00:00:00 2001 From: Evan Mattson Date: Thu, 20 Aug 2026 18:50:51 +0900 Subject: [PATCH 5/7] fix approval recovery and workflow cancellation --- .../ag-ui/agent_framework_ag_ui/_agent_run.py | 31 ++--- .../_approval_lifecycle.py | 35 +++++- .../agent_framework_ag_ui/_workflow_run.py | 32 +---- .../ag-ui/tests/ag_ui/test_endpoint.py | 117 ++++++++++++++++++ python/packages/core/AGENTS.md | 4 +- .../_workflows/_agent_executor.py | 5 + .../agent_framework/_workflows/_executor.py | 4 + .../_workflows/_runner_context.py | 12 ++ .../agent_framework/_workflows/_workflow.py | 34 ++++- .../_workflows/_workflow_executor.py | 5 + .../test_agent_executor_tool_calls.py | 31 +++++ 11 files changed, 260 insertions(+), 50 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 6539925f6bf..d9fb48912ca 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 @@ -74,7 +74,6 @@ from ._message_adapters import normalize_agui_input_messages from ._predictive_state import PredictiveStateHandler from ._tooling import collect_server_tools, merge_tools -from ._workflow_run import _consume_cancelled_workflow_requests # pyright: ignore[reportPrivateUsage] from ._run_common import ( FlowState, _approval_interrupt_for_function_call, # type: ignore @@ -885,7 +884,7 @@ def _register_server_generated_approval_response( server_label=_function_call_server_label(response.function_call), ) if response.approved is not True: - lifecycle.claim_batch( + batch = lifecycle.claim_batch( thread_id=thread_id, decisions=[ ResumeDecision( @@ -896,6 +895,14 @@ def _register_server_generated_approval_response( ) ], ) + if batch.authorized_executions: + intent = batch.authorized_executions[0] + lifecycle.begin_execution(intent, owner=intent.owner) + lifecycle.settle_forwarded( + intent, + [response], + owner=intent.owner, + ) return None if execution_owner is ApprovalExecutionOwner.UNAVAILABLE: return None @@ -1625,7 +1632,7 @@ async def _resolve_approval_responses( else: primary_response.function_call.additional_properties.pop("server_label", None) if ( - primary_response.approved is True + (primary_response.approved is True or pending_entry.owner is ApprovalExecutionOwner.DEFERRED) and lifecycle is not None and authorized_executions is not None and primary_response.function_call is not None @@ -1640,10 +1647,9 @@ async def _resolve_approval_responses( valid_response_content_ids.add(id(primary_response)) if pending_entry.owner is ApprovalExecutionOwner.DEFERRED: deferred_response_content_ids.add(id(primary_response)) - if ( - primary_response.approved is True - and intent is not None - and intent.owner in {ApprovalExecutionOwner.HOSTED, ApprovalExecutionOwner.DEFERRED} + if intent is not None and ( + intent.owner is ApprovalExecutionOwner.DEFERRED + or (primary_response.approved is True and intent.owner is ApprovalExecutionOwner.HOSTED) ): validated_forwarded_approvals.append(primary_response) if not server_label: @@ -2353,19 +2359,16 @@ async def run_agent_stream( yield resume_error return if cancelled_resume_ids and isinstance(agent, WorkflowAgent): - cancelled_workflow_entries: list[dict[str, Any]] = [] + cancelled_workflow_request_ids: set[str] = set() for interrupt_id in cancelled_resume_ids: occurrence = approval_state_store.lifecycle.occurrence_for_alias( thread_id=approval_thread_id, interrupt_id=interrupt_id, ) - cancelled_workflow_entries.append( - { - "interrupt_id": occurrence.response_id if occurrence and occurrence.response_id else interrupt_id, - "status": "cancelled", - } + cancelled_workflow_request_ids.add( + occurrence.response_id if occurrence and occurrence.response_id else interrupt_id ) - _consume_cancelled_workflow_requests(agent.workflow, cancelled_workflow_entries) + await agent.workflow.cancel_pending_requests(cancelled_workflow_request_ids) if cancelled_resume_ids and handled_resume_ids == cancelled_resume_ids: yield RunStartedEvent(run_id=run_id, thread_id=thread_id) _clear_tool_approval_state(approval_state_store, approval_thread_id) diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_approval_lifecycle.py b/python/packages/ag-ui/agent_framework_ag_ui/_approval_lifecycle.py index 7defdef5f70..485599ce4ed 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_approval_lifecycle.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_approval_lifecycle.py @@ -548,6 +548,19 @@ def claim_batch( continue occurrence.decision = decision if not decision.accepted: + if occurrence.owner is ApprovalExecutionOwner.DEFERRED: + occurrence.status = ApprovalStatus.CLAIMED + self._emit_event("claim", occurrence) + intents.append( + AuthorizedExecution( + identity=occurrence.identity, + name=occurrence.name, + arguments=occurrence.arguments, + owner=occurrence.owner, + idempotency_key=occurrence.idempotency_key, + ) + ) + continue result = Content.from_function_result( call_id=occurrence.identity.call_id, result="Error: Tool call invocation was rejected by user.", @@ -846,6 +859,11 @@ def recover_execution( ) if occurrence.status is not ApprovalStatus.EXECUTING: raise ApprovalSettlementConflictError(f"Approval occurrence is not executing: {occurrence.status}.") + if occurrence.decision is not None and not occurrence.decision.accepted: + occurrence.status = ApprovalStatus.PENDING + occurrence.pending_since = self._clock() + self._emit_event("rejection_recovery", occurrence) + return None if intent.idempotency_key is not None and intent.idempotency_key == occurrence.idempotency_key: occurrence.status = ApprovalStatus.CLAIMED return intent @@ -907,19 +925,30 @@ def settle_forwarded( result for result in results if result.type == "function_approval_response" - and result.approved is True + and isinstance(result.approved, bool) and result.function_call is not None and result.function_call.call_id == occurrence.identity.call_id ] if len(replayable_results) + len(forwarded_responses) != 1: raise ValueError("A hosted approval must record exactly one outcome for its original call.") + is_rejection = len(forwarded_responses) == 1 and forwarded_responses[0].approved is False + if is_rejection: + rejection_result = Content.from_function_result( + call_id=occurrence.identity.call_id, + result="Error: Tool call invocation was rejected by user.", + ) + replayable_results = [ReplayableToolResult(content=rejection_result)] + result_group = (rejection_result,) + occurrence.status = ApprovalStatus.REJECTED + else: + result_group = tuple(results) + occurrence.status = ApprovalStatus.SETTLED occurrence.replayable_results = replayable_results - occurrence.status = ApprovalStatus.SETTLED self._remove_pending_aliases(occurrence) outcome = ApprovalOutcome( identity=occurrence.identity, replayable_results=tuple(replayable_results), - result_group=tuple(results), + result_group=result_group, snapshot_reconciliation=self._snapshot_reconciliation(occurrence), ) occurrence.outcome = outcome diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_workflow_run.py b/python/packages/ag-ui/agent_framework_ag_ui/_workflow_run.py index 4702b0d613d..074276ec0cc 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_workflow_run.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_workflow_run.py @@ -342,36 +342,6 @@ def _resume_error_for_pending_workflow_requests( return None -def _consume_cancelled_workflow_requests(workflow: Workflow, resume_entries: list[dict[str, Any]]) -> None: - """Remove cancelled workflow requests from runner and owning agent-executor state.""" - cancelled_ids = {str(entry["interrupt_id"]) for entry in resume_entries if entry.get("status") == "cancelled"} - if not cancelled_ids: - return - - runner_context = getattr(workflow, "_runner_context", None) - pending_events = getattr(runner_context, "_pending_request_info_events", None) - if not isinstance(pending_events, dict): - return - pending_events = cast(dict[str, Any], pending_events) - - for interrupt_id in cancelled_ids: - request_event = pending_events.pop(interrupt_id, None) - source_executor_id = getattr(request_event, "source_executor_id", None) - executor = workflow.executors.get(source_executor_id) if source_executor_id else None - nested_workflow = getattr(executor, "workflow", None) - if isinstance(nested_workflow, Workflow): - _consume_cancelled_workflow_requests( - nested_workflow, - [{"interrupt_id": interrupt_id, "status": "cancelled"}], - ) - pending_agent_requests = getattr(executor, "_pending_agent_requests", None) - if isinstance(pending_agent_requests, dict): - cast(dict[str, Any], pending_agent_requests).pop(interrupt_id, None) - - if not pending_events: - workflow._status = WorkflowRunState.IDLE # pyright: ignore[reportPrivateUsage] - - def _coerce_json_value(value: Any) -> Any: """Parse JSON strings when possible; otherwise return the original value.""" if not isinstance(value, str): @@ -1097,7 +1067,7 @@ async def run_workflow_stream( yield response_error return if cancelled_request_ids: - _consume_cancelled_workflow_requests(workflow, resume_entries) + await workflow.cancel_pending_requests(cancelled_request_ids) pending_before_run = { request_id: request_event for request_id, request_event in pending_before_run.items() diff --git a/python/packages/ag-ui/tests/ag_ui/test_endpoint.py b/python/packages/ag-ui/tests/ag_ui/test_endpoint.py index c712e577dac..a9aea56ce7e 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_endpoint.py +++ b/python/packages/ag-ui/tests/ag_ui/test_endpoint.py @@ -807,6 +807,123 @@ async def approve( ) +async def test_endpoint_workflow_as_agent_rejection_retries_after_transient_failure() -> None: + """A deferred rejection remains retryable until the wrapped workflow consumes it.""" + + class FailFirstResumeProvider(ContextProvider): + def __init__(self) -> None: + super().__init__("fail-first-resume") + self.call_count = 0 + + async def before_run( + self, + *, + agent: SupportsAgentRun, + session: AgentSession, + context: SessionContext, + state: dict[str, Any], + ) -> None: + del agent, session, context, state + self.call_count += 1 + if self.call_count == 2: + raise RuntimeError("transient workflow failure") + + class ApprovalExecutor(Executor): + def __init__(self) -> None: + super().__init__(id="approval") + + @handler + async def start(self, message: Any, ctx: WorkflowContext[Any, Any]) -> None: + del message + function_call = Content.from_function_call( + call_id="refund-call", + name="submit_refund", + arguments={"order_id": "12345"}, + ) + await ctx.request_info( + Content.from_function_approval_request(id="approval-1", function_call=function_call), + Content, + request_id="approval-1", + ) + + @response_handler + async def approve( + self, + original_request: Content, + response: Content, + ctx: WorkflowContext[Any, Any], + ) -> None: + del original_request + await ctx.yield_output(f"{response.type}:{response.approved}") # type: ignore[arg-type] + + provider = FailFirstResumeProvider() + workflow = WorkflowBuilder(start_executor=ApprovalExecutor()).build() + workflow_agent = workflow.as_agent(context_providers=[provider]) + app = FastAPI() + add_agent_framework_fastapi_endpoint( + app, + AgentFrameworkAgent(agent=workflow_agent), + path="/workflow-agent-retry", + ) + resume = [ + { + "interruptId": "refund-call", + "status": "resolved", + "payload": {"accepted": False}, + } + ] + client_tool = { + "name": "submit_refund", + "description": "Submit a refund", + "parameters": { + "type": "object", + "properties": {"order_id": {"type": "string"}}, + }, + } + + with TestClient(app) as client: + pause_response = client.post( + "/workflow-agent-retry", + json={ + "runId": "run-pause", + "threadId": "thread-reject-retry", + "messages": [{"role": "user", "content": "Refund the order"}], + "tools": [client_tool], + }, + ) + assert not [event for event in _decode_sse_events(pause_response) if event.get("type") == "RUN_ERROR"] + + failed_response = client.post( + "/workflow-agent-retry", + json={ + "runId": "run-failed", + "threadId": "thread-reject-retry", + "messages": [], + "tools": [client_tool], + "resume": resume, + }, + ) + failed_errors = [event for event in _decode_sse_events(failed_response) if event.get("type") == "RUN_ERROR"] + assert len(failed_errors) == 1 + + retry_response = client.post( + "/workflow-agent-retry", + json={ + "runId": "run-retry", + "threadId": "thread-reject-retry", + "messages": [], + "tools": [client_tool], + "resume": resume, + }, + ) + + retry_events = _decode_sse_events(retry_response) + assert not [event for event in retry_events if event.get("type") == "RUN_ERROR"] + assert "function_approval_response:False" == "".join( + str(event.get("delta", "")) for event in retry_events if event.get("type") == "TEXT_MESSAGE_CONTENT" + ) + + async def test_workflow_endpoint_applies_canonical_approval_edited_args() -> None: """Workflow approvals apply standard editedArgs as a full replacement.""" diff --git a/python/packages/core/AGENTS.md b/python/packages/core/AGENTS.md index 593aa1e4b54..3d856d4a0a1 100644 --- a/python/packages/core/AGENTS.md +++ b/python/packages/core/AGENTS.md @@ -192,7 +192,9 @@ agent_framework/ ### Workflows (`_workflows/`) -- **`Workflow`** - Graph-based workflow definition +- **`Workflow`** - Graph-based workflow definition. `cancel_pending_requests(request_ids)` cancels selected external + requests without synthesizing responses, recursively releases nested executor correlation, and returns the IDs + that were still pending. - **`WorkflowBuilder`** - Fluent API for building workflows, including explicit `output_from` / `intermediate_output_from` selection for caller-facing emissions. `output_from` is an allow-list for **Workflow Output**; unselected executor payloads are hidden unless diff --git a/python/packages/core/agent_framework/_workflows/_agent_executor.py b/python/packages/core/agent_framework/_workflows/_agent_executor.py index 3bb92682f92..3411f8e85c2 100644 --- a/python/packages/core/agent_framework/_workflows/_agent_executor.py +++ b/python/packages/core/agent_framework/_workflows/_agent_executor.py @@ -327,6 +327,11 @@ async def handle_user_input_response( self._pending_responses_to_agent.clear() await self._run_agent_and_emit(ctx) + @override + async def _cancel_pending_request(self, request_id: str) -> None: + """Release an agent-owned user-input request after workflow cancellation.""" + self._pending_agent_requests.pop(request_id, None) + @override async def on_checkpoint_save(self) -> dict[str, Any]: """Capture current executor state for checkpointing. diff --git a/python/packages/core/agent_framework/_workflows/_executor.py b/python/packages/core/agent_framework/_workflows/_executor.py index 3571df2557c..ab46b7047c9 100644 --- a/python/packages/core/agent_framework/_workflows/_executor.py +++ b/python/packages/core/agent_framework/_workflows/_executor.py @@ -518,6 +518,10 @@ def _find_handler(self, message: Any) -> Callable[[Any, WorkflowContext[Any, Any return self._handlers[message_type] raise RuntimeError(f"Executor {self.__class__.__name__} cannot handle message of type {type(message)}.") + async def _cancel_pending_request(self, request_id: str) -> None: + """Release executor-owned state for a cancelled workflow request.""" + del request_id + async def on_checkpoint_save(self) -> dict[str, Any]: """Hook called when the workflow is being saved to a checkpoint. diff --git a/python/packages/core/agent_framework/_workflows/_runner_context.py b/python/packages/core/agent_framework/_workflows/_runner_context.py index f16e74b9ca4..3044c80e7e4 100644 --- a/python/packages/core/agent_framework/_workflows/_runner_context.py +++ b/python/packages/core/agent_framework/_workflows/_runner_context.py @@ -315,6 +315,10 @@ async def get_pending_request_info_events(self) -> dict[str, WorkflowEvent[Any]] """ ... + async def cancel_request_info_events(self, request_ids: set[str]) -> dict[str, WorkflowEvent[Any]]: + """Remove and return pending request_info events selected for cancellation.""" + ... + def set_yield_output_classifier(self, classifier: YieldOutputClassifier) -> None: """Set the classifier used by WorkflowContext.yield_output().""" ... @@ -606,6 +610,14 @@ async def get_pending_request_info_events(self) -> dict[str, WorkflowEvent[Any]] """ return dict(self._pending_request_info_events) + async def cancel_request_info_events(self, request_ids: set[str]) -> dict[str, WorkflowEvent[Any]]: + """Remove and return pending request_info events selected for cancellation.""" + return { + request_id: event + for request_id in request_ids + if (event := self._pending_request_info_events.pop(request_id, None)) is not None + } + def set_yield_output_classifier(self, classifier: YieldOutputClassifier) -> None: """Set the classifier used by WorkflowContext.yield_output().""" self._yield_output_classifier = classifier diff --git a/python/packages/core/agent_framework/_workflows/_workflow.py b/python/packages/core/agent_framework/_workflows/_workflow.py index 1e933644c26..bfb5899d725 100644 --- a/python/packages/core/agent_framework/_workflows/_workflow.py +++ b/python/packages/core/agent_framework/_workflows/_workflow.py @@ -12,7 +12,7 @@ import uuid import warnings import weakref -from collections.abc import AsyncIterable, Awaitable, Callable, Mapping, Sequence +from collections.abc import AsyncIterable, Awaitable, Callable, Collection, Mapping, Sequence from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any, Literal, overload @@ -1204,6 +1204,38 @@ def output_types(self) -> list[type[Any] | types.UnionType]: return list(output_types) + async def cancel_pending_requests(self, request_ids: Collection[str]) -> set[str]: + """Cancel pending external requests and release their owning executor state. + + Cancellation follows requests through nested workflows and clears any executor-owned + correlation without synthesizing a response. Unknown or already-handled request IDs + are ignored. + + Args: + request_ids: Request identifiers to cancel. + + Returns: + The set of request identifiers that were pending and are now cancelled. + """ + selected_ids = set(request_ids) + if not all(isinstance(request_id, str) and request_id for request_id in selected_ids): + raise ValueError("Pending workflow request IDs must be non-empty strings.") + cancelled_events = await self._runner.context.cancel_request_info_events(selected_ids) + for request_id, request_event in cancelled_events.items(): + source_executor_id = request_event.source_executor_id + executor = self.executors.get(source_executor_id) if source_executor_id else None + if executor is not None: + await executor._cancel_pending_request(request_id) # pyright: ignore[reportPrivateUsage] + + if ( + cancelled_events + and not await self._runner.context.get_pending_request_info_events() + and self._status + in {WorkflowRunState.IDLE_WITH_PENDING_REQUESTS, WorkflowRunState.IN_PROGRESS_PENDING_REQUESTS} + ): + self._status = WorkflowRunState.IDLE + return set(cancelled_events) + def as_agent( self, name: str | None = None, diff --git a/python/packages/core/agent_framework/_workflows/_workflow_executor.py b/python/packages/core/agent_framework/_workflows/_workflow_executor.py index 8ab9d6d6ab9..47627f4f55d 100644 --- a/python/packages/core/agent_framework/_workflows/_workflow_executor.py +++ b/python/packages/core/agent_framework/_workflows/_workflow_executor.py @@ -450,6 +450,11 @@ async def handle_propagated_request_response( ctx=ctx, ) + @override + async def _cancel_pending_request(self, request_id: str) -> None: + """Propagate cancellation into the wrapped workflow.""" + await self.workflow.cancel_pending_requests([request_id]) + @override async def on_checkpoint_save(self) -> dict[str, Any]: """Get the current state of the WorkflowExecutor for checkpointing purposes.""" diff --git a/python/packages/core/tests/workflow/test_agent_executor_tool_calls.py b/python/packages/core/tests/workflow/test_agent_executor_tool_calls.py index 717f28019d6..e2dfe48dcce 100644 --- a/python/packages/core/tests/workflow/test_agent_executor_tool_calls.py +++ b/python/packages/core/tests/workflow/test_agent_executor_tool_calls.py @@ -25,6 +25,7 @@ WorkflowBuilder, WorkflowContext, WorkflowEvent, + WorkflowExecutor, executor, tool, ) @@ -366,6 +367,36 @@ async def test_agent_executor_parallel_tool_call_with_approval() -> None: assert final_response[0] == "Tool executed successfully." +async def test_workflow_cancels_nested_pending_request_without_blocking_sibling() -> None: + """Cancelling one nested request lets its resolved sibling complete the workflow.""" + agent = Agent( + client=MockChatClient(parallel_request=True), + name="ApprovalAgent", + tools=[mock_tool_requiring_approval], + ) + child = WorkflowBuilder(start_executor=agent, output_from=[test_executor]).add_edge(agent, test_executor).build() + nested = WorkflowExecutor( + child, + id="nested-workflow", + propagate_request=True, + allow_direct_output=True, + ) + parent = WorkflowBuilder(start_executor=nested).build() + + paused = await parent.run([Message(role="user", contents=["Invoke tool requiring approval"])]) + first_request, second_request = paused.get_request_info_events() + + cancelled = await parent.cancel_pending_requests([first_request.request_id]) + resumed = await parent.run( + responses={ + second_request.request_id: second_request.data.to_function_approval_response(True), + } + ) + + assert cancelled == {first_request.request_id} + assert resumed.get_outputs() == ["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 From 750717dfdc12b563d43d23eed6b810aed86be31d Mon Sep 17 00:00:00 2001 From: Evan Mattson Date: Thu, 20 Aug 2026 19:25:13 +0900 Subject: [PATCH 6/7] fix workflow cancellation ordering --- .../agent_framework_ag_ui/_workflow_run.py | 11 ++- .../ag-ui/tests/ag_ui/test_endpoint.py | 91 +++++++++++++++++++ python/packages/core/AGENTS.md | 4 +- .../_workflows/_agent_executor.py | 30 ++++-- .../agent_framework/_workflows/_executor.py | 4 +- .../agent_framework/_workflows/_workflow.py | 74 ++++++++++----- .../_workflows/_workflow_executor.py | 5 +- .../test_agent_executor_tool_calls.py | 26 +++++- 8 files changed, 207 insertions(+), 38 deletions(-) diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_workflow_run.py b/python/packages/ag-ui/agent_framework_ag_ui/_workflow_run.py index 074276ec0cc..7c3a4ceedd8 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_workflow_run.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_workflow_run.py @@ -1067,7 +1067,16 @@ async def run_workflow_stream( yield response_error return if cancelled_request_ids: - await workflow.cancel_pending_requests(cancelled_request_ids) + if checkpoint_id is not None: + _ = await workflow.run( + checkpoint_id=checkpoint_id, + checkpoint_storage=checkpoint_storage, + ) + checkpoint_id = None + await workflow.cancel_pending_requests( + cancelled_request_ids, + checkpoint_storage=checkpoint_storage, + ) pending_before_run = { request_id: request_event for request_id, request_event in pending_before_run.items() diff --git a/python/packages/ag-ui/tests/ag_ui/test_endpoint.py b/python/packages/ag-ui/tests/ag_ui/test_endpoint.py index a9aea56ce7e..0e6759a5df0 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_endpoint.py +++ b/python/packages/ag-ui/tests/ag_ui/test_endpoint.py @@ -4846,6 +4846,97 @@ async def test_endpoint_workflow_checkpoint_resume_same_owner_after_restart(): assert "Booked KLM" in text_deltas +async def test_endpoint_workflow_checkpoint_cancellation_survives_cold_restore() -> None: + """Cold restore applies cancellation before resolving the remaining sibling.""" + + class BatchApprovalExecutor(Executor): + def __init__(self) -> None: + super().__init__(id="batch-approval") + + @handler + async def start(self, messages: list[Message], ctx: WorkflowContext[Any, Any]) -> None: + del messages + await ctx.request_info({"order_id": "order-1"}, dict, request_id="approval-1") + await ctx.request_info({"order_id": "order-2"}, dict, request_id="approval-2") + + @response_handler + async def approve( + self, + original_request: dict[str, Any], + response: dict[str, Any], + ctx: WorkflowContext[Any, Any], + ) -> None: + assert response == {"approved": True} + await ctx.yield_output(f"Approved {original_request['order_id']}") # type: ignore[arg-type] + + def build_workflow() -> Any: + return WorkflowBuilder( + name="cold-checkpoint-cancellation", + start_executor=BatchApprovalExecutor(), + ).build() + + storage = InMemoryCheckpointStorage() + first_app = FastAPI() + first_workflow = build_workflow() + add_agent_framework_fastapi_endpoint( + first_app, + first_workflow, + path="/workflow", + checkpoint_storage=storage, + ) + + with TestClient(first_app) as client: + pause_response = client.post( + "/workflow", + json={ + "runId": "run-pause", + "threadId": "owner-thread", + "messages": [{"role": "user", "content": "Approve both orders"}], + }, + ) + assert pause_response.status_code == 200 + + checkpoints = await storage.list_checkpoints(workflow_name=first_workflow.name) + pending_checkpoints = [checkpoint for checkpoint in checkpoints if checkpoint.pending_request_info_events] + assert pending_checkpoints + checkpoint_id = max(pending_checkpoints, key=lambda checkpoint: checkpoint.timestamp).checkpoint_id + + second_app = FastAPI() + add_agent_framework_fastapi_endpoint( + second_app, + build_workflow(), + path="/workflow", + checkpoint_storage=storage, + ) + + with TestClient(second_app) as client: + cancel_response = client.post( + "/workflow", + json={ + "runId": "run-cancel", + "threadId": "owner-thread", + "messages": [], + "forwardedProps": {"checkpointId": checkpoint_id}, + "resume": [ + {"interruptId": "approval-1", "status": "cancelled"}, + { + "interruptId": "approval-2", + "status": "resolved", + "payload": {"approved": True}, + }, + ], + }, + ) + + cancel_events = _decode_sse_events(cancel_response) + assert not [event for event in cancel_events if event.get("type") == "RUN_ERROR"] + finished = [event for event in cancel_events if event.get("type") == "RUN_FINISHED"] + assert finished[-1].get("outcome") is None + assert "Approved order-2" == "".join( + str(event.get("delta", "")) for event in cancel_events if event.get("type") == "TEXT_MESSAGE_CONTENT" + ) + + async def test_endpoint_workflow_checkpoint_resume_uses_checkpoint_owner_not_live_reused_id(): """A live reused interrupt ID cannot authorize a different checkpoint occurrence.""" storage = InMemoryCheckpointStorage() diff --git a/python/packages/core/AGENTS.md b/python/packages/core/AGENTS.md index 3d856d4a0a1..06a77c42117 100644 --- a/python/packages/core/AGENTS.md +++ b/python/packages/core/AGENTS.md @@ -193,8 +193,8 @@ agent_framework/ ### Workflows (`_workflows/`) - **`Workflow`** - Graph-based workflow definition. `cancel_pending_requests(request_ids)` cancels selected external - requests without synthesizing responses, recursively releases nested executor correlation, and returns the IDs - that were still pending. + requests without synthesizing responses, recursively releases nested executor correlation, resumes executors whose + remaining requests were already answered, and returns the resulting `WorkflowRunResult`. - **`WorkflowBuilder`** - Fluent API for building workflows, including explicit `output_from` / `intermediate_output_from` selection for caller-facing emissions. `output_from` is an allow-list for **Workflow Output**; unselected executor payloads are hidden unless diff --git a/python/packages/core/agent_framework/_workflows/_agent_executor.py b/python/packages/core/agent_framework/_workflows/_agent_executor.py index 3411f8e85c2..512706e36f4 100644 --- a/python/packages/core/agent_framework/_workflows/_agent_executor.py +++ b/python/packages/core/agent_framework/_workflows/_agent_executor.py @@ -319,18 +319,32 @@ async def handle_user_input_response( self._pending_agent_requests.pop(original_request.id, None) # type: ignore[arg-type] if not self._pending_agent_requests: - # All pending requests have been resolved; resume agent execution. - # Use role="tool" for function_result responses (from declaration-only tools) - # so the LLM receives proper tool results instead of orphaned tool_calls. - role = "tool" if all(r.type == "function_result" for r in self._pending_responses_to_agent) else "user" - self._cache = normalize_messages_input(Message(role=role, contents=self._pending_responses_to_agent)) - self._pending_responses_to_agent.clear() - await self._run_agent_and_emit(ctx) + await self._resume_with_pending_responses(ctx) + + async def _resume_with_pending_responses( + self, + ctx: WorkflowContext[AgentExecutorResponse, AgentResponse | AgentResponseUpdate], + ) -> None: + """Resume agent execution after every pending request has reached an outcome.""" + if not self._pending_responses_to_agent: + return + # Use role="tool" for function_result responses (from declaration-only tools) + # so the LLM receives proper tool results instead of orphaned tool_calls. + role = "tool" if all(r.type == "function_result" for r in self._pending_responses_to_agent) else "user" + self._cache = normalize_messages_input(Message(role=role, contents=self._pending_responses_to_agent)) + self._pending_responses_to_agent.clear() + await self._run_agent_and_emit(ctx) @override - async def _cancel_pending_request(self, request_id: str) -> None: + async def _cancel_pending_request( + self, + request_id: str, + ctx: WorkflowContext[AgentExecutorResponse, AgentResponse | AgentResponseUpdate], + ) -> None: """Release an agent-owned user-input request after workflow cancellation.""" self._pending_agent_requests.pop(request_id, None) + if not self._pending_agent_requests: + await self._resume_with_pending_responses(ctx) @override async def on_checkpoint_save(self) -> dict[str, Any]: diff --git a/python/packages/core/agent_framework/_workflows/_executor.py b/python/packages/core/agent_framework/_workflows/_executor.py index ab46b7047c9..39027955a7f 100644 --- a/python/packages/core/agent_framework/_workflows/_executor.py +++ b/python/packages/core/agent_framework/_workflows/_executor.py @@ -518,9 +518,9 @@ def _find_handler(self, message: Any) -> Callable[[Any, WorkflowContext[Any, Any return self._handlers[message_type] raise RuntimeError(f"Executor {self.__class__.__name__} cannot handle message of type {type(message)}.") - async def _cancel_pending_request(self, request_id: str) -> None: + async def _cancel_pending_request(self, request_id: str, ctx: WorkflowContext[Any, Any]) -> None: """Release executor-owned state for a cancelled workflow request.""" - del request_id + del request_id, ctx async def on_checkpoint_save(self) -> dict[str, Any]: """Hook called when the workflow is being saved to a checkpoint. diff --git a/python/packages/core/agent_framework/_workflows/_workflow.py b/python/packages/core/agent_framework/_workflows/_workflow.py index bfb5899d725..fc013f3fe58 100644 --- a/python/packages/core/agent_framework/_workflows/_workflow.py +++ b/python/packages/core/agent_framework/_workflows/_workflow.py @@ -588,8 +588,11 @@ async def _run_workflow_with_tracing( with _framework_event_origin(): pending_status = WorkflowEvent.status(self._status) yield pending_status - # Workflow runs until idle - emit final status based on whether requests are pending - if saw_request: + # Workflow runs until idle - emit final status based on whether requests are pending. + # Continuations such as cancellation may retain an existing sibling request without + # re-emitting its request_info event during this run. + pending_requests = await self._runner.context.get_pending_request_info_events() + if saw_request or pending_requests: self._status = WorkflowRunState.IDLE_WITH_PENDING_REQUESTS with _framework_event_origin(): terminal_status = WorkflowEvent.status(self._status) @@ -1204,37 +1207,66 @@ def output_types(self) -> list[type[Any] | types.UnionType]: return list(output_types) - async def cancel_pending_requests(self, request_ids: Collection[str]) -> set[str]: + async def cancel_pending_requests( + self, + request_ids: Collection[str], + *, + checkpoint_storage: CheckpointStorage | None = None, + ) -> WorkflowRunResult: """Cancel pending external requests and release their owning executor state. Cancellation follows requests through nested workflows and clears any executor-owned - correlation without synthesizing a response. Unknown or already-handled request IDs - are ignored. + correlation without synthesizing a response. If cancellation drains an executor's pending + set after sibling responses were already accepted, the executor resumes through its normal + continuation path. Unknown or already-handled request IDs are ignored. Args: request_ids: Request identifiers to cancel. + Keyword Args: + checkpoint_storage: Runtime checkpoint storage for the cancellation continuation. + Returns: - The set of request identifiers that were pending and are now cancelled. + Events produced while applying cancellation and any resulting continuation. """ selected_ids = set(request_ids) if not all(isinstance(request_id, str) and request_id for request_id in selected_ids): raise ValueError("Pending workflow request IDs must be non-empty strings.") - cancelled_events = await self._runner.context.cancel_request_info_events(selected_ids) - for request_id, request_event in cancelled_events.items(): - source_executor_id = request_event.source_executor_id - executor = self.executors.get(source_executor_id) if source_executor_id else None - if executor is not None: - await executor._cancel_pending_request(request_id) # pyright: ignore[reportPrivateUsage] - - if ( - cancelled_events - and not await self._runner.context.get_pending_request_info_events() - and self._status - in {WorkflowRunState.IDLE_WITH_PENDING_REQUESTS, WorkflowRunState.IN_PROGRESS_PENDING_REQUESTS} - ): - self._status = WorkflowRunState.IDLE - return set(cancelled_events) + + async def apply_cancellations() -> None: + cancelled_events = await self._runner.context.cancel_request_info_events(selected_ids) + for request_id, request_event in cancelled_events.items(): + source_executor_id = request_event.source_executor_id + executor = self.executors.get(source_executor_id) if source_executor_id else None + if executor is None: + continue + context = executor._create_context_for_handler( # pyright: ignore[reportPrivateUsage] + source_executor_ids=[INTERNAL_SOURCE_ID(executor.id)], + state=self._runner.state, + runner_context=self._runner.context, + ) + await executor._cancel_pending_request( # pyright: ignore[reportPrivateUsage] + request_id, + context, + ) + + if checkpoint_storage is not None: + self._runner.context.set_runtime_checkpoint_storage(checkpoint_storage) + events: list[WorkflowEvent[Any]] = [] + try: + async for event in self._run_workflow_with_tracing( + initial_executor_fn=apply_cancellations, + is_continuation=True, + streaming=False, + tools=self._runner.context.get_runtime_tools(), + function_invocation_kwargs=None, + client_kwargs=None, + ): + events.append(event) + finally: + if checkpoint_storage is not None: + self._runner.context.clear_runtime_checkpoint_storage() + return self._finalize_events(events) def as_agent( self, diff --git a/python/packages/core/agent_framework/_workflows/_workflow_executor.py b/python/packages/core/agent_framework/_workflows/_workflow_executor.py index 47627f4f55d..1c4f2f8ad1c 100644 --- a/python/packages/core/agent_framework/_workflows/_workflow_executor.py +++ b/python/packages/core/agent_framework/_workflows/_workflow_executor.py @@ -451,9 +451,10 @@ async def handle_propagated_request_response( ) @override - async def _cancel_pending_request(self, request_id: str) -> None: + async def _cancel_pending_request(self, request_id: str, ctx: WorkflowContext[Any, Any]) -> None: """Propagate cancellation into the wrapped workflow.""" - await self.workflow.cancel_pending_requests([request_id]) + result = await self.workflow.cancel_pending_requests([request_id]) + await self._process_workflow_result(result, ctx) @override async def on_checkpoint_save(self) -> dict[str, Any]: diff --git a/python/packages/core/tests/workflow/test_agent_executor_tool_calls.py b/python/packages/core/tests/workflow/test_agent_executor_tool_calls.py index e2dfe48dcce..3d4e9c4921a 100644 --- a/python/packages/core/tests/workflow/test_agent_executor_tool_calls.py +++ b/python/packages/core/tests/workflow/test_agent_executor_tool_calls.py @@ -386,17 +386,39 @@ async def test_workflow_cancels_nested_pending_request_without_blocking_sibling( paused = await parent.run([Message(role="user", contents=["Invoke tool requiring approval"])]) first_request, second_request = paused.get_request_info_events() - cancelled = await parent.cancel_pending_requests([first_request.request_id]) + await parent.cancel_pending_requests([first_request.request_id]) resumed = await parent.run( responses={ second_request.request_id: second_request.data.to_function_approval_response(True), } ) - assert cancelled == {first_request.request_id} assert resumed.get_outputs() == ["Tool executed successfully."] +async def test_workflow_final_cancellation_resumes_accumulated_sibling_response() -> None: + """Cancelling the final request resumes an agent that already received its sibling response.""" + agent = Agent( + client=MockChatClient(parallel_request=True), + name="ApprovalAgent", + tools=[mock_tool_requiring_approval], + ) + workflow = WorkflowBuilder(start_executor=agent, output_from=[test_executor]).add_edge(agent, test_executor).build() + + paused = await workflow.run("Invoke tool requiring approval") + first_request, second_request = paused.get_request_info_events() + partial = await workflow.run( + responses={ + first_request.request_id: first_request.data.to_function_approval_response(True), + } + ) + + cancelled = await workflow.cancel_pending_requests([second_request.request_id]) + + assert partial.get_outputs() == [] + assert cancelled.get_outputs() == ["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 From 84354b66bbf5b25a7df0d72f211223e87b032c10 Mon Sep 17 00:00:00 2001 From: Evan Mattson Date: Fri, 21 Aug 2026 16:02:26 +0900 Subject: [PATCH 7/7] fix cancellation continuation context --- .../ag-ui/agent_framework_ag_ui/_agent_run.py | 5 ++- .../agent_framework_ag_ui/_workflow_run.py | 8 +--- python/packages/core/AGENTS.md | 4 +- .../agent_framework/_workflows/_workflow.py | 26 +++++++++---- .../_workflows/_workflow_executor.py | 5 ++- .../test_agent_executor_tool_calls.py | 39 +++++++++++++++++++ .../core/tests/workflow/test_workflow.py | 31 +++++++++++++++ 7 files changed, 101 insertions(+), 17 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 d9fb48912ca..c39456c7782 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 @@ -2368,7 +2368,10 @@ async def run_agent_stream( cancelled_workflow_request_ids.add( occurrence.response_id if occurrence and occurrence.response_id else interrupt_id ) - await agent.workflow.cancel_pending_requests(cancelled_workflow_request_ids) + await agent.workflow.cancel_pending_requests( + cancelled_workflow_request_ids, + tools=tools, + ) if cancelled_resume_ids and handled_resume_ids == cancelled_resume_ids: yield RunStartedEvent(run_id=run_id, thread_id=thread_id) _clear_tool_approval_state(approval_state_store, approval_thread_id) diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_workflow_run.py b/python/packages/ag-ui/agent_framework_ag_ui/_workflow_run.py index 7c3a4ceedd8..866086fea77 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_workflow_run.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_workflow_run.py @@ -1067,16 +1067,12 @@ async def run_workflow_stream( yield response_error return if cancelled_request_ids: - if checkpoint_id is not None: - _ = await workflow.run( - checkpoint_id=checkpoint_id, - checkpoint_storage=checkpoint_storage, - ) - checkpoint_id = None await workflow.cancel_pending_requests( cancelled_request_ids, + checkpoint_id=checkpoint_id, checkpoint_storage=checkpoint_storage, ) + checkpoint_id = None pending_before_run = { request_id: request_event for request_id, request_event in pending_before_run.items() diff --git a/python/packages/core/AGENTS.md b/python/packages/core/AGENTS.md index 06a77c42117..78217a8f06e 100644 --- a/python/packages/core/AGENTS.md +++ b/python/packages/core/AGENTS.md @@ -194,7 +194,9 @@ agent_framework/ - **`Workflow`** - Graph-based workflow definition. `cancel_pending_requests(request_ids)` cancels selected external requests without synthesizing responses, recursively releases nested executor correlation, resumes executors whose - remaining requests were already answered, and returns the resulting `WorkflowRunResult`. + remaining requests were already answered, accepts the same request-scoped tools and invocation/client kwargs needed + by that continuation, can atomically restore a supplied checkpoint before cancellation, and returns the resulting + `WorkflowRunResult`. - **`WorkflowBuilder`** - Fluent API for building workflows, including explicit `output_from` / `intermediate_output_from` selection for caller-facing emissions. `output_from` is an allow-list for **Workflow Output**; unselected executor payloads are hidden unless diff --git a/python/packages/core/agent_framework/_workflows/_workflow.py b/python/packages/core/agent_framework/_workflows/_workflow.py index fc013f3fe58..b980d0e0810 100644 --- a/python/packages/core/agent_framework/_workflows/_workflow.py +++ b/python/packages/core/agent_framework/_workflows/_workflow.py @@ -520,7 +520,6 @@ async def _run_workflow_with_tracing( OtelAttr.WORKFLOW_RUN_SPAN, attributes, ) as span: - saw_request = False emitted_in_progress_pending = False try: # Add workflow started event (telemetry + surface state to consumers) @@ -577,9 +576,6 @@ async def _run_workflow_with_tracing( # All executor executions happen within workflow span async for event in self._runner.run_until_convergence(): - # Track request events for final status determination - if event.type == "request_info": - saw_request = True yield event if event.type == "request_info" and not emitted_in_progress_pending: @@ -592,7 +588,7 @@ async def _run_workflow_with_tracing( # Continuations such as cancellation may retain an existing sibling request without # re-emitting its request_info event during this run. pending_requests = await self._runner.context.get_pending_request_info_events() - if saw_request or pending_requests: + if pending_requests: self._status = WorkflowRunState.IDLE_WITH_PENDING_REQUESTS with _framework_event_origin(): terminal_status = WorkflowEvent.status(self._status) @@ -1211,7 +1207,11 @@ async def cancel_pending_requests( self, request_ids: Collection[str], *, + checkpoint_id: str | None = None, checkpoint_storage: CheckpointStorage | None = None, + tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None = None, + function_invocation_kwargs: Mapping[str, Mapping[str, Any]] | Mapping[str, Any] | None = None, + client_kwargs: Mapping[str, Mapping[str, Any]] | Mapping[str, Any] | None = None, ) -> WorkflowRunResult: """Cancel pending external requests and release their owning executor state. @@ -1224,7 +1224,11 @@ async def cancel_pending_requests( request_ids: Request identifiers to cancel. Keyword Args: + checkpoint_id: Checkpoint to restore before applying cancellation. checkpoint_storage: Runtime checkpoint storage for the cancellation continuation. + tools: Request-scoped tools available while cancellation resumes executors. + function_invocation_kwargs: Keyword arguments forwarded to resumed tool invocations. + client_kwargs: Keyword arguments forwarded to resumed chat client calls. Returns: Events produced while applying cancellation and any resulting continuation. @@ -1252,20 +1256,26 @@ async def apply_cancellations() -> None: if checkpoint_storage is not None: self._runner.context.set_runtime_checkpoint_storage(checkpoint_storage) + self._runner.context.set_runtime_tools(tools) events: list[WorkflowEvent[Any]] = [] try: + if checkpoint_id is not None: + await self._runner.restore_from_checkpoint(checkpoint_id, checkpoint_storage) async for event in self._run_workflow_with_tracing( initial_executor_fn=apply_cancellations, is_continuation=True, streaming=False, - tools=self._runner.context.get_runtime_tools(), - function_invocation_kwargs=None, - client_kwargs=None, + tools=tools, + function_invocation_kwargs=function_invocation_kwargs, + client_kwargs=client_kwargs, ): + if event.type == "request_info" and event.request_id in selected_ids: + continue events.append(event) finally: if checkpoint_storage is not None: self._runner.context.clear_runtime_checkpoint_storage() + self._runner.context.clear_runtime_tools() return self._finalize_events(events) def as_agent( diff --git a/python/packages/core/agent_framework/_workflows/_workflow_executor.py b/python/packages/core/agent_framework/_workflows/_workflow_executor.py index 1c4f2f8ad1c..f3bccd1ff15 100644 --- a/python/packages/core/agent_framework/_workflows/_workflow_executor.py +++ b/python/packages/core/agent_framework/_workflows/_workflow_executor.py @@ -453,7 +453,10 @@ async def handle_propagated_request_response( @override async def _cancel_pending_request(self, request_id: str, ctx: WorkflowContext[Any, Any]) -> None: """Propagate cancellation into the wrapped workflow.""" - result = await self.workflow.cancel_pending_requests([request_id]) + result = await self.workflow.cancel_pending_requests( + [request_id], + tools=ctx.get_runtime_tools(), + ) await self._process_workflow_result(result, ctx) @override diff --git a/python/packages/core/tests/workflow/test_agent_executor_tool_calls.py b/python/packages/core/tests/workflow/test_agent_executor_tool_calls.py index 3d4e9c4921a..cee0b1570e8 100644 --- a/python/packages/core/tests/workflow/test_agent_executor_tool_calls.py +++ b/python/packages/core/tests/workflow/test_agent_executor_tool_calls.py @@ -419,6 +419,45 @@ async def test_workflow_final_cancellation_resumes_accumulated_sibling_response( assert cancelled.get_outputs() == ["Tool executed successfully."] +async def test_workflow_final_cancellation_preserves_runtime_tool_for_approved_sibling() -> None: + """Cancellation continuation keeps request-scoped tools needed by an accepted sibling.""" + executed_queries: list[str] = [] + + def execute_runtime_tool(query: str) -> str: + executed_queries.append(query) + return f"Executed runtime tool with query: {query}" + + runtime_tool = FunctionTool( + name="mock_tool_requiring_approval", + description="Request-scoped approval tool", + func=execute_runtime_tool, + approval_mode="always_require", + ) + agent = Agent( + client=MockChatClient(parallel_request=True), + name="ApprovalAgent", + ) + workflow = WorkflowBuilder(start_executor=agent, output_from=[test_executor]).add_edge(agent, test_executor).build() + + paused = await workflow.run("Invoke tool requiring approval", tools=[runtime_tool]) + first_request, second_request = paused.get_request_info_events() + partial = await workflow.run( + responses={ + first_request.request_id: first_request.data.to_function_approval_response(True), + }, + tools=[runtime_tool], + ) + + cancelled = await workflow.cancel_pending_requests( + [second_request.request_id], + tools=[runtime_tool], + ) + + assert partial.get_outputs() == [] + assert cancelled.get_outputs() == ["Tool executed successfully."] + assert executed_queries == ["test"] + + 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 diff --git a/python/packages/core/tests/workflow/test_workflow.py b/python/packages/core/tests/workflow/test_workflow.py index 2f672f591d3..014694b9b80 100644 --- a/python/packages/core/tests/workflow/test_workflow.py +++ b/python/packages/core/tests/workflow/test_workflow.py @@ -21,6 +21,7 @@ Content, Executor, FileCheckpointStorage, + InMemoryCheckpointStorage, InProcRunnerContext, Message, ResponseStream, @@ -516,6 +517,36 @@ async def test_workflow_run_stream_from_checkpoint_with_responses( assert len(events) > 0 # Just ensure we processed some events +async def test_cancel_pending_requests_restores_checkpoint_before_cancellation() -> None: + """Cold cancellation restores persisted pending state before removing the request.""" + storage = InMemoryCheckpointStorage() + + def build_workflow() -> Any: + return WorkflowBuilder( + name="cold-core-cancellation", + start_executor=MockExecutorRequestApproval(id="approver"), + checkpoint_storage=storage, + ).build() + + first_workflow = build_workflow() + paused = await first_workflow.run(NumberMessage(data=7)) + [request] = paused.get_request_info_events() + checkpoints = await storage.list_checkpoints(workflow_name=first_workflow.name) + pending_checkpoints = [checkpoint for checkpoint in checkpoints if checkpoint.pending_request_info_events] + assert pending_checkpoints + checkpoint_id = max(pending_checkpoints, key=lambda checkpoint: checkpoint.timestamp).checkpoint_id + + fresh_workflow = build_workflow() + cancelled = await fresh_workflow.cancel_pending_requests( + [request.request_id], + checkpoint_id=checkpoint_id, + checkpoint_storage=storage, + ) + + assert cancelled.get_request_info_events() == [] + assert cancelled.get_final_state() is WorkflowRunState.IDLE + + @dataclass class StateTrackingMessage: """A message that tracks state for testing context reset behavior."""