diff --git a/python/packages/declarative/agent_framework_declarative/_workflows/__init__.py b/python/packages/declarative/agent_framework_declarative/_workflows/__init__.py index 31ee990cb69..5421a48e8dd 100644 --- a/python/packages/declarative/agent_framework_declarative/_workflows/__init__.py +++ b/python/packages/declarative/agent_framework_declarative/_workflows/__init__.py @@ -12,10 +12,13 @@ - Full integration with the workflow runtime """ +from agent_framework import register_checkpoint_type + from ._declarative_base import ( DECLARATIVE_STATE_KEY, ActionComplete, ActionTrigger, + ConditionResult, ConversationData, DeclarativeActionExecutor, DeclarativeMessage, @@ -97,6 +100,36 @@ ) from ._state import WorkflowState +# Framework-owned types that cross a checkpoint boundary. +# +# Checkpoint restore runs pickle through a restricted unpickler whose default allowlist +# trusts ``agent_framework.`` and ``agent_framework_orchestrations.`` but not this +# distribution, so these types opt in by name, as the orchestrations package does. +# Without this, ``FileCheckpointStorage`` cannot save any checkpoint taken after the +# first action, and a paused workflow cannot be resumed from storage. +for _checkpoint_type in ( + # Executor-to-executor message envelopes (``DeclarativeMessage``). + ActionTrigger, + ActionComplete, + ConditionResult, + LoopIterationResult, + LoopControl, + # ``request_info`` payloads and their response types, checkpointed as pending + # requests while a workflow waits for input or approval. + ExternalInputRequest, + ExternalInputResponse, + AgentExternalInputRequest, + AgentExternalInputResponse, + ToolApprovalRequest, + ToolApprovalResponse, + MCPToolApprovalRequest, + # Stored in workflow state while an agent's external input loop is paused. + ExternalLoopState, +): + register_checkpoint_type(_checkpoint_type) + +del _checkpoint_type + __all__ = [ "AGENT_ACTION_EXECUTORS", "AGENT_REGISTRY_KEY", diff --git a/python/packages/declarative/agent_framework_declarative/_workflows/_factory.py b/python/packages/declarative/agent_framework_declarative/_workflows/_factory.py index a8329a59734..6d51abafd09 100644 --- a/python/packages/declarative/agent_framework_declarative/_workflows/_factory.py +++ b/python/packages/declarative/agent_framework_declarative/_workflows/_factory.py @@ -66,7 +66,7 @@ class WorkflowFactory: from agent_framework import FileCheckpointStorage # With checkpointing for pause/resume support - storage = FileCheckpointStorage(path="./checkpoints") + storage = FileCheckpointStorage("./checkpoints") factory = WorkflowFactory(checkpoint_storage=storage) workflow = factory.create_workflow_from_yaml_path("workflow.yaml") diff --git a/python/packages/declarative/tests/declarative/test_declarative_checkpoint_types.py b/python/packages/declarative/tests/declarative/test_declarative_checkpoint_types.py new file mode 100644 index 00000000000..feb571881e6 --- /dev/null +++ b/python/packages/declarative/tests/declarative/test_declarative_checkpoint_types.py @@ -0,0 +1,164 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Built-in declarative workflow types must survive a checkpoint round trip. + +Checkpoint restore runs pickle through a restricted unpickler. Its default allowlist +trusts ``agent_framework.`` and ``agent_framework_orchestrations.`` but not +``agent_framework_declarative``, so framework-owned declarative types have to opt in by +name. Importing the package is what registers them. + +These tests use a ``FileCheckpointStorage`` built with **no** ``allowed_checkpoint_types``, +because the bug was that users had to supply framework-internal module paths themselves. +""" + +from __future__ import annotations + +import os +from pathlib import Path +from typing import Any + +import pytest +from agent_framework import Content, FileCheckpointStorage, Message, WorkflowCheckpoint + +from agent_framework_declarative import ExternalInputResponse, WorkflowFactory +from agent_framework_declarative._workflows import ( + ActionComplete, + ActionTrigger, + AgentExternalInputRequest, + AgentExternalInputResponse, + ExternalInputRequest, + ExternalLoopState, + LoopControl, + LoopIterationResult, + MCPToolApprovalRequest, + ToolApprovalRequest, + ToolApprovalResponse, +) +from agent_framework_declarative._workflows._declarative_base import ConditionResult + +try: + import powerfx # noqa: F401 + + _powerfx_available = True +except (ImportError, RuntimeError): + _powerfx_available = False + + +def _payloads() -> list[tuple[str, Any]]: + """One populated instance of every declarative type that crosses a checkpoint.""" + message = Message(role="assistant", contents=["Which city?"]) + return [ + ("ActionTrigger", ActionTrigger(data={"input": "hello"})), + ("ActionComplete", ActionComplete(result={"value": 1})), + ("ConditionResult", ConditionResult(matched=True, branch_index=0, value=True)), + ("LoopIterationResult", LoopIterationResult(has_next=True, current_item="a", current_index=0)), + ("LoopControl", LoopControl(action="break")), + ("ExternalInputRequest", ExternalInputRequest(request_id="r1", message="Approve?", metadata={"k": "v"})), + ("ExternalInputResponse", ExternalInputResponse(user_input="yes", value={"approved": True})), + ( + "AgentExternalInputRequest", + AgentExternalInputRequest( + request_id="r2", + agent_name="planner", + agent_response="Which city?", + iteration=1, + messages=[message], + function_calls=[Content.from_function_call(call_id="c1", name="lookup", arguments={"q": "x"})], + ), + ), + ( + "AgentExternalInputResponse", + AgentExternalInputResponse( + user_input="Seattle", + messages=[Message(role="user", contents=["Seattle"])], + function_results={"c1": Content.from_function_result(call_id="c1", result="ok")}, + ), + ), + ("ToolApprovalRequest", ToolApprovalRequest(request_id="r3", function_name="send", arguments={"to": "a"})), + ("ToolApprovalResponse", ToolApprovalResponse(approved=False, reason="not now")), + ( + "MCPToolApprovalRequest", + MCPToolApprovalRequest( + request_id="r4", + tool_name="search", + server_url="https://example.com/mcp", + server_label="docs", + arguments={"query": "agents"}, + header_names=["Authorization"], + metadata={"source": "test"}, + ), + ), + ( + "ExternalLoopState", + ExternalLoopState( + agent_name="planner", + iteration=2, + external_loop_when="=true", + messages_var="Local.messages", + response_obj_var=None, + result_property=None, + auto_send=True, + ), + ), + ] + + +@pytest.mark.parametrize(("name", "payload"), _payloads(), ids=[name for name, _ in _payloads()]) +async def test_declarative_type_survives_checkpoint_round_trip(tmp_path: Path, name: str, payload: Any) -> None: + """Saving and restoring a built-in declarative payload needs no user allowlist.""" + storage = FileCheckpointStorage(tmp_path) # deliberately no allowed_checkpoint_types + await storage.save( + WorkflowCheckpoint( + workflow_name="declarative", + graph_signature_hash="test-hash", + checkpoint_id=f"ck-{name}", + state={"payload": payload}, + ) + ) + + restored = await storage.load(f"ck-{name}") + + restored_payload = restored.state["payload"] + assert type(restored_payload) is type(payload) + assert vars(restored_payload).keys() == vars(payload).keys() + + +@pytest.mark.skipif(not _powerfx_available, reason="PowerFx engine not available") +async def test_paused_workflow_resumes_from_file_checkpoint(tmp_path: Path) -> None: + """A workflow paused on external input can be restored from storage and resumed.""" + yaml_definition = """ +name: checkpoint-resume +actions: + - kind: SetValue + path: Local.greeting + value: Hello + - kind: RequestExternalInput + id: ask + prompt: Approve? + variable: Local.answer + - kind: SendActivity + activity: + text: '=Concatenate(Local.greeting, " - ", Local.answer)' +""" + storage = FileCheckpointStorage(tmp_path) # deliberately no allowed_checkpoint_types + workflow = WorkflowFactory(checkpoint_storage=storage).create_workflow_from_yaml(yaml_definition) + + result = await workflow.run({}) + requests = result.get_request_info_events() + assert len(requests) == 1 + + # Every checkpoint written to disk must also decode; undecodable ones are skipped by the listing. + files_on_disk = [name for name in os.listdir(tmp_path) if name.endswith(".json")] + listed = await storage.list_checkpoints(workflow_name=workflow.name) + assert len(listed) == len(files_on_disk) > 1 + + latest = await storage.get_latest(workflow_name=workflow.name) + assert latest is not None + + restored_workflow = WorkflowFactory(checkpoint_storage=storage).create_workflow_from_yaml(yaml_definition) + await restored_workflow.run(checkpoint_id=latest.checkpoint_id) + resumed = await restored_workflow.run( + responses={requests[0].request_id: ExternalInputResponse(user_input="approved")} + ) + + assert resumed.get_outputs() == ["Hello - approved"]