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 78e4b2a026..c39456c778 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 ( @@ -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,6 +2320,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 +2328,8 @@ 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, 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], 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 c792b5465b..485599ce4e 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, @@ -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 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 6ed72c9c04..1a5bc7b2a1 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 a01851d1dc..866086fea7 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,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() 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 64b0270e73..0e6759a5df 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, @@ -454,6 +455,475 @@ 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_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_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.""" + + 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_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.""" @@ -4376,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 593aa1e4b5..78217a8f06 100644 --- a/python/packages/core/AGENTS.md +++ b/python/packages/core/AGENTS.md @@ -192,7 +192,11 @@ 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, resumes executors whose + 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/_agent.py b/python/packages/core/agent_framework/_workflows/_agent.py index 09bd592e92..1aed8c0596 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, @@ -151,26 +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, - 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]: ... + 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[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, - 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: ... + 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[AgentResponseUpdate, AgentResponse[Any]]: ... def run( self, @@ -180,9 +183,10 @@ def run( session: AgentSession | None = None, checkpoint_id: str | None = None, checkpoint_storage: CheckpointStorage | 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]: + 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[AgentResponseUpdate, AgentResponse[Any]] | Awaitable[AgentResponse[Any]]: """Get a response from the workflow agent. Args: @@ -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.") @@ -453,12 +470,14 @@ 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, stream=True, checkpoint_storage=checkpoint_storage, + tools=tools, function_invocation_kwargs=function_invocation_kwargs, client_kwargs=client_kwargs, ): @@ -467,6 +486,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 +497,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 +506,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 +735,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 @@ -730,22 +748,40 @@ 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 + and isinstance(pending_request.data, Content) + and pending_request.data.type == "function_call" + 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 b7787736fa..512706e36f 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] = {} @@ -308,13 +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, + 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]: @@ -410,9 +440,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 = ctx.get_runtime_tools() if not self._cache: logger.warning( @@ -422,13 +452,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 and self._accepts_runtime_tools: + run_kwargs["tools"] = tools + response = await run_agent(self._cache, **run_kwargs) # Handle any user input requests if response.user_input_requests: @@ -464,9 +496,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 = ctx.get_runtime_tools() if not self._cache: logger.warning( @@ -478,13 +510,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 and self._accepts_runtime_tools: + 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 +596,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/_executor.py b/python/packages/core/agent_framework/_workflows/_executor.py index 3571df2557..39027955a7 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, ctx: WorkflowContext[Any, Any]) -> None: + """Release executor-owned state for a cancelled workflow request.""" + 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/_runner_context.py b/python/packages/core/agent_framework/_workflows/_runner_context.py index f70b9075f7..3044c80e7e 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, @@ -303,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().""" ... @@ -339,6 +355,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 +548,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. @@ -581,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 04f3f9aa87..b980d0e081 100644 --- a/python/packages/core/agent_framework/_workflows/_workflow.py +++ b/python/packages/core/agent_framework/_workflows/_workflow.py @@ -12,11 +12,12 @@ 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 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 @@ -517,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) @@ -574,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: @@ -585,8 +584,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 pending_requests: self._status = WorkflowRunState.IDLE_WITH_PENDING_REQUESTS with _framework_event_origin(): terminal_status = WorkflowEvent.status(self._status) @@ -688,6 +690,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 +705,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 +719,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 +743,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. @@ -775,6 +781,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( @@ -783,6 +790,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 +810,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 +886,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, ): @@ -909,6 +919,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( @@ -1192,6 +1203,81 @@ def output_types(self) -> list[type[Any] | types.UnionType]: return list(output_types) + 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. + + Cancellation follows requests through nested workflows and clears any executor-owned + 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_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. + """ + 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.") + + 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) + 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=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( self, name: str | None = None, diff --git a/python/packages/core/agent_framework/_workflows/_workflow_context.py b/python/packages/core/agent_framework/_workflows/_workflow_context.py index 18a53e0bfa..08128f22af 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 1a8f988d19..f3bccd1ff1 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 = ctx.get_runtime_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 ) @@ -448,6 +450,15 @@ async def handle_propagated_request_response( ctx=ctx, ) + @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], + tools=ctx.get_runtime_tools(), + ) + await self._process_workflow_result(result, ctx) + @override async def on_checkpoint_save(self) -> dict[str, Any]: """Get the current state of the WorkflowExecutor for checkpointing purposes.""" @@ -606,5 +617,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. - result = await self.workflow.run(responses={request_id: response}) + 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 17b9332a28..cee0b1570e 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,97 @@ 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() + + 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 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_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 @@ -424,6 +516,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 +526,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 +628,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.py b/python/packages/core/tests/workflow/test_workflow.py index 2f672f591d..014694b9b8 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.""" diff --git a/python/packages/core/tests/workflow/test_workflow_kwargs.py b/python/packages/core/tests/workflow/test_workflow_kwargs.py index 93c6c93d58..5a6fe8d81c 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 @@ -584,6 +584,74 @@ 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_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") @@ -646,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 @@ -695,6 +818,64 @@ 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_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.