Python: Stamp AG-UI checkpoint owner on every save

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}'.",

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()