Python: Fix AG-UI workflow-as-agent approval resumes

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 (
@@ -879,10 +880,11 @@ 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:
-        lifecycle.claim_batch(
+        batch = lifecycle.claim_batch(
             thread_id=thread_id,
             decisions=[
                 ResumeDecision(
@@ -893,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
@@ -1382,9 +1392,10 @@ def _canonical_approval_resume_messages(
                 original_arguments=pending_arguments,
             )
         )
+        response_id = pending_entry.response_id or 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,
@@ -1422,6 +1433,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(
@@ -1532,6 +1544,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()
@@ -1619,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
@@ -1632,10 +1645,11 @@ async def _resolve_approval_responses(
                 continue
             intents_by_response_content_id[id(primary_response)] = intent
         valid_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 pending_entry.owner is ApprovalExecutionOwner.DEFERRED:
+            deferred_response_content_ids.add(id(primary_response))
+        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:
@@ -1700,7 +1714,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 []
@@ -2306,14 +2320,16 @@ 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,
             approval_thread_id,
             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,
             authorized_executions=authorized_executions,
             retained_results=retained_approval_results,
             snapshot_reconciliations=approval_snapshot_reconciliations,
@@ -2342,6 +2358,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_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_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,
+            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)
@@ -2626,7 +2656,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],

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,
@@ -542,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.",
@@ -840,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
@@ -901,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

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,
         )

python/packages/ag-ui/agent_framework_ag_ui/_workflow_run.py

@@ -342,27 +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
-        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)
-
-
 def _coerce_json_value(value: Any) -> Any:
     """Parse JSON strings when possible; otherwise return the original value."""
     if not isinstance(value, str):
@@ -1088,7 +1067,12 @@ 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,
+            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()