From 3c56b2666d8c358c6a13fde1e2a078e98191d4c2 Mon Sep 17 00:00:00 2001 From: Choppaaahh Date: Wed, 2 Sep 2026 08:05:12 -0400 Subject: [PATCH 1/2] Stamp AG-UI checkpoint owner on every save --- .../ag-ui/agent_framework_ag_ui/_workflow.py | 11 +++++----- .../ag-ui/tests/ag_ui/test_endpoint.py | 22 ++++++++++++++++++- 2 files changed, 26 insertions(+), 7 deletions(-) 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..552f844393 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: 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 dc4efc6802..2b2e2cedea 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_endpoint.py +++ b/python/packages/ag-ui/tests/ag_ui/test_endpoint.py @@ -31,6 +31,7 @@ SupportsAgentRun, ToolApprovalMiddleware, WorkflowBuilder, + WorkflowCheckpoint, WorkflowContext, WorkflowExecutor, executor, @@ -54,7 +55,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]]: @@ -4725,6 +4730,21 @@ 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_threaded_resume_after_restart(): """An explicitly threaded cold checkpoint resume fails closed when ownership is unavailable.""" storage = InMemoryCheckpointStorage() From 60c2e96796a16c922a53a3f3689770df132abca3 Mon Sep 17 00:00:00 2001 From: Choppaaahh Date: Thu, 3 Sep 2026 04:11:18 -0400 Subject: [PATCH 2/2] Validate checkpoint ownership whenever a checkpoint carries an owner The resume path only compared ownership when the checkpoint had pending request info events, so a checkpoint with an owner but no pending events skipped the check entirely. Combined with stamping ownership on every save, a caller could resume another Snapshot Scope's clean checkpoint and have subsequent saves re-stamped to themselves. Validate on owner presence instead. Checkpoints saved before ownership was recorded have no owner and continue to resume unchanged. --- .../ag-ui/agent_framework_ag_ui/_workflow.py | 3 +- .../ag-ui/tests/ag_ui/test_endpoint.py | 55 +++++++++++++++++++ 2 files changed, 56 insertions(+), 2 deletions(-) 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 552f844393..73b2a1d7ec 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_workflow.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_workflow.py @@ -397,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 8347c3f093..6f38bd8f26 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_endpoint.py +++ b/python/packages/ag-ui/tests/ag_ui/test_endpoint.py @@ -5303,6 +5303,61 @@ async def test_owned_checkpoint_storage_stamps_owner_without_pending_events() -> 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()