diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_workflow.py b/python/packages/ag-ui/agent_framework_ag_ui/_workflow.py index a5ccd1a9ca..73b2a1d7ec 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_workflow.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_workflow.py @@ -87,12 +87,11 @@ def __init__(self, storage: CheckpointStorage, owner: WorkflowRequestOwner) -> N async def save(self, checkpoint: WorkflowCheckpoint) -> CheckpointID: """Save a checkpoint with ownership for any pending request occurrences.""" - if checkpoint.pending_request_info_events: - checkpoint.metadata = dict(checkpoint.metadata) - checkpoint.metadata[_CHECKPOINT_REQUEST_OWNER_KEY] = { - "snapshot_scope": self._owner[0], - "thread_id": self._owner[1], - } + checkpoint.metadata = dict(checkpoint.metadata) + checkpoint.metadata[_CHECKPOINT_REQUEST_OWNER_KEY] = { + "snapshot_scope": self._owner[0], + "thread_id": self._owner[1], + } return await self._storage.save(checkpoint) async def load(self, checkpoint_id: CheckpointID) -> WorkflowCheckpoint: @@ -398,9 +397,8 @@ async def run(self, input_data: dict[str, Any]) -> AsyncGenerator[BaseEvent]: code="WORKFLOW_CHECKPOINT_LOAD_FAILED", ) return - checkpoint_pending_ids = {str(request_id) for request_id in checkpoint.pending_request_info_events} checkpoint_owner = _checkpoint_request_owner(checkpoint.metadata) - if checkpoint_pending_ids and checkpoint_owner != request_owner: + if checkpoint_owner is not None and checkpoint_owner != request_owner: yield RunStartedEvent(run_id=run_id, thread_id=thread_id) yield RunErrorEvent( message=f"No pending interrupt found for checkpointId '{checkpoint_id}'.", diff --git a/python/packages/ag-ui/tests/ag_ui/test_endpoint.py b/python/packages/ag-ui/tests/ag_ui/test_endpoint.py index 120e387817..6f38bd8f26 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 @@ SupportsAgentRun, ToolApprovalMiddleware, WorkflowBuilder, + WorkflowCheckpoint, WorkflowContext, WorkflowExecutor, executor, @@ -55,7 +56,11 @@ from agent_framework_ag_ui._agent import AgentFrameworkAgent from agent_framework_ag_ui._approval_lifecycle import ApprovalExecutionOwner, ApprovalLifecycle, ApprovalStatus from agent_framework_ag_ui._approval_state import InMemoryAGUIApprovalStateStore, approval_state_thread_id -from agent_framework_ag_ui._workflow import AgentFrameworkWorkflow +from agent_framework_ag_ui._workflow import ( + _CHECKPOINT_REQUEST_OWNER_KEY, + AgentFrameworkWorkflow, + _OwnedWorkflowCheckpointStorage, +) def _decode_sse_events(response: Any) -> list[dict[str, Any]]: @@ -5283,6 +5288,76 @@ async def test_endpoint_workflow_request_info_rejects_unowned_pending_interrupt( assert attacker_errors[0]["code"] == "WORKFLOW_RESUME_NOT_FOUND" +async def test_owned_checkpoint_storage_stamps_owner_without_pending_events() -> None: + """The request owner is stamped on every save, not only when pending request events exist.""" + storage = InMemoryCheckpointStorage() + owned_storage = _OwnedWorkflowCheckpointStorage(storage, ("scope-1", "thread-1")) + checkpoint = WorkflowCheckpoint(workflow_name="owned-workflow", graph_signature_hash="signature") + assert not checkpoint.pending_request_info_events + + checkpoint_id = await owned_storage.save(checkpoint) + + expected_owner = {"snapshot_scope": "scope-1", "thread_id": "thread-1"} + assert checkpoint.metadata[_CHECKPOINT_REQUEST_OWNER_KEY] == expected_owner + stored = await storage.load(checkpoint_id) + assert stored.metadata[_CHECKPOINT_REQUEST_OWNER_KEY] == expected_owner + + +async def test_endpoint_workflow_checkpoint_resume_rejects_foreign_clean_checkpoint(): + """A checkpoint with no pending request events is still owned and cannot be resumed by another thread.""" + storage = InMemoryCheckpointStorage() + first_app = FastAPI() + first_workflow = _build_flight_choice_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": "victim-thread", + "messages": [{"role": "user", "content": "Book me a flight"}], + }, + ) + assert pause_response.status_code == 200 + + checkpoints = await storage.list_checkpoints(workflow_name=first_workflow.name) + clean_checkpoints = [checkpoint for checkpoint in checkpoints if not checkpoint.pending_request_info_events] + assert clean_checkpoints, "expected at least one checkpoint without pending request events" + checkpoint = min(clean_checkpoints, key=lambda checkpoint: checkpoint.timestamp) + assert checkpoint.metadata.get(_CHECKPOINT_REQUEST_OWNER_KEY) is not None + + second_app = FastAPI() + add_agent_framework_fastapi_endpoint( + second_app, + _build_flight_choice_workflow(), + path="/workflow", + checkpoint_storage=storage, + ) + + with TestClient(second_app) as client: + attacker_response = client.post( + "/workflow", + json={ + "runId": "run-attacker", + "threadId": "attacker-thread", + "messages": [], + "forwardedProps": {"checkpointId": checkpoint.checkpoint_id}, + }, + ) + + attacker_events = _decode_sse_events(attacker_response) + attacker_errors = [event for event in attacker_events if event.get("type") == "RUN_ERROR"] + assert len(attacker_errors) == 1 + assert attacker_errors[0]["code"] == "WORKFLOW_RESUME_NOT_FOUND" + assert not [event for event in attacker_events if event.get("type") == "TEXT_MESSAGE_CONTENT"] + + async def test_endpoint_workflow_checkpoint_resume_rejects_threaded_resume_after_restart(): """An explicitly threaded cold checkpoint resume fails closed when ownership is unavailable.""" storage = InMemoryCheckpointStorage()