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 325b56ce730..921e9f3d11f 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 @@ -226,6 +226,19 @@ def _resume_to_workflow_responses(resume_payload: Any) -> dict[str, Any]: return responses +def _merge_workflow_response_sources( + resume_responses: dict[str, Any], + message_responses: dict[str, Any], +) -> dict[str, Any]: + """Merge workflow response sources with explicit resume payloads taking precedence.""" + if not resume_responses: + return dict(message_responses) + + responses = dict(message_responses) + responses.update(resume_responses) + return responses + + def _resume_entries_to_workflow_responses(entries: list[dict[str, Any]]) -> dict[str, Any]: """Convert validated resume entries into workflow responses.""" responses: dict[str, Any] = {} @@ -773,12 +786,12 @@ async def run_workflow_stream( yield resume_error return - responses = ( + resume_responses = ( _resume_entries_to_workflow_responses(resume_entries) if pending_interrupt_ids else _resume_to_workflow_responses(resume_payload) ) - responses.update(_extract_responses_from_messages(messages)) + responses = _merge_workflow_response_sources(resume_responses, _extract_responses_from_messages(messages)) responses, response_error = _coerce_responses_for_pending_requests_strict(responses, pending_before_run) if response_error is not None: yield RunStartedEvent(run_id=run_id, thread_id=thread_id) diff --git a/python/packages/ag-ui/tests/ag_ui/test_workflow_run.py b/python/packages/ag-ui/tests/ag_ui/test_workflow_run.py index aa1c05b3936..0eb50938169 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_workflow_run.py +++ b/python/packages/ag-ui/tests/ag_ui/test_workflow_run.py @@ -1403,6 +1403,83 @@ async def handle_approval(self, original_request: Content, response: Content, ct assert "outcome" not in resumed_finished +async def test_workflow_run_explicit_resume_overrides_stale_message_approval() -> None: + """Explicit resume payloads should not be overwritten by stale function_approvals in messages.""" + + class ApprovalExecutor(Executor): + def __init__(self) -> None: + super().__init__(id="approval_executor") + + @handler + async def start(self, message: Any, ctx: WorkflowContext) -> None: + del message + function_call = Content.from_function_call( + call_id="refund-call", + name="submit_refund", + arguments={"order_id": "12345", "amount": "$89.99"}, + ) + approval_request = Content.from_function_approval_request(id="approval-1", function_call=function_call) + await ctx.request_info(approval_request, Content, request_id="approval-1") + + @response_handler + async def handle_approval( + self, original_request: Content, response: Content, ctx: WorkflowContext[Any, str] + ) -> None: + del original_request + status = "approved" if bool(response.approved) else "rejected" + await ctx.yield_output(f"Refund {status}.") + + workflow = WorkflowBuilder(start_executor=ApprovalExecutor()).build() + first_events: list[Any] = [ + event async for event in run_workflow_stream({"messages": [{"role": "user", "content": "go"}]}, workflow) + ] + first_finished_events = [event for event in first_events if event.type == "RUN_FINISHED"] + assert len(first_finished_events) == 1 + interrupt_payload = _interrupts_from_run_finished(first_finished_events[0]) + assert len(interrupt_payload) == 1 + interrupt_value = _interrupt_metadata_value(interrupt_payload[0]) + + resumed_events: list[Any] = [ + event + async for event in run_workflow_stream( + { + "messages": [ + { + "role": "user", + "content": "", + "function_approvals": [ + { + "approved": True, + "id": "approval-1", + "call_id": "refund-call", + "name": "submit_refund", + "arguments": {"order_id": "12345", "amount": "$89.99"}, + } + ], + } + ], + "resume": [ + { + "interruptId": "approval-1", + "status": "resolved", + "payload": { + "type": "function_approval_response", + "approved": False, + "id": interrupt_value.get("id", "approval-1"), + "function_call": interrupt_value.get("function_call"), + }, + } + ], + }, + workflow, + ) + ] + + assistant_text = "".join(event.delta for event in resumed_events if event.type == "TEXT_MESSAGE_CONTENT") + assert "rejected" in assistant_text + assert "approved" not in assistant_text + + async def test_workflow_run_approval_argument_mismatch_emits_run_error() -> None: """Workflow approval responses must fail when function arguments change."""