Python: Refactor RequestInfoExecutor (#1403)

* Refactor RequestInfoExecutor

* Update AI script

* Fix formatting

* Address comments

* fix unit test

python/packages/core/agent_framework/_workflows/__init__.py

@@ -12,6 +12,7 @@ from ._checkpoint import (
     InMemoryCheckpointStorage,
     WorkflowCheckpoint,
 )
+from ._checkpoint_summary import WorkflowCheckpointSummary, get_checkpoint_summary
 from ._concurrent import ConcurrentBuilder
 from ._const import (
     DEFAULT_MAX_ITERATIONS,
@@ -74,6 +75,7 @@ from ._magentic import (
     StandardMagenticManager,
 )
 from ._request_info_executor import (
+    PendingRequestDetails,
     RequestInfoExecutor,
     RequestInfoMessage,
     RequestResponse,
@@ -146,6 +148,7 @@ __all__ = [
     "MagenticResponseMessage",
     "MagenticStartMessage",
     "Message",
+    "PendingRequestDetails",
     "RequestInfoEvent",
     "RequestInfoExecutor",
     "RequestInfoMessage",
@@ -165,6 +168,7 @@ __all__ = [
     "WorkflowAgent",
     "WorkflowBuilder",
     "WorkflowCheckpoint",
+    "WorkflowCheckpointSummary",
     "WorkflowContext",
     "WorkflowErrorDetails",
     "WorkflowEvent",
@@ -181,6 +185,7 @@ __all__ = [
     "WorkflowViz",
     "create_edge_runner",
     "executor",
+    "get_checkpoint_summary",
     "handler",
     "validate_workflow_graph",
 ]

python/packages/core/agent_framework/_workflows/__init__.pyi

@@ -12,6 +12,7 @@ from ._checkpoint import (
     InMemoryCheckpointStorage,
     WorkflowCheckpoint,
 )
+from ._checkpoint_summary import WorkflowCheckpointSummary, get_checkpoint_summary
 from ._concurrent import ConcurrentBuilder
 from ._const import DEFAULT_MAX_ITERATIONS
 from ._edge import (
@@ -72,6 +73,7 @@ from ._magentic import (
     StandardMagenticManager,
 )
 from ._request_info_executor import (
+    PendingRequestDetails,
     RequestInfoExecutor,
     RequestInfoMessage,
     RequestResponse,
@@ -144,6 +146,7 @@ __all__ = [
     "MagenticResponseMessage",
     "MagenticStartMessage",
     "Message",
+    "PendingRequestDetails",
     "RequestInfoEvent",
     "RequestInfoExecutor",
     "RequestInfoMessage",
@@ -163,6 +166,7 @@ __all__ = [
     "WorkflowAgent",
     "WorkflowBuilder",
     "WorkflowCheckpoint",
+    "WorkflowCheckpointSummary",
     "WorkflowContext",
     "WorkflowErrorDetails",
     "WorkflowEvent",
@@ -179,6 +183,7 @@ __all__ = [
     "WorkflowViz",
     "create_edge_runner",
     "executor",
+    "get_checkpoint_summary",
     "handler",
     "validate_workflow_graph",
 ]

python/packages/core/agent_framework/_workflows/_checkpoint_summary.py

@@ -0,0 +1,191 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import logging
+from collections.abc import Iterable, Mapping
+from dataclasses import dataclass
+from textwrap import shorten
+from typing import Any
+
+from ._checkpoint import WorkflowCheckpoint
+from ._request_info_executor import PendingRequestDetails, RequestInfoMessage, RequestResponse
+from ._runner_context import _decode_checkpoint_value  # type: ignore
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class WorkflowCheckpointSummary:
+    """Human-readable summary of a workflow checkpoint."""
+
+    checkpoint_id: str
+    iteration_count: int
+    targets: list[str]
+    executor_ids: list[str]
+    status: str
+    draft_preview: str | None
+    pending_requests: list[PendingRequestDetails]
+
+
+def get_checkpoint_summary(
+    checkpoint: WorkflowCheckpoint,
+    *,
+    request_executor_ids: Iterable[str] | None = None,
+    preview_width: int = 70,
+) -> WorkflowCheckpointSummary:
+    targets = sorted(checkpoint.messages.keys())
+    executor_ids = sorted(checkpoint.executor_states.keys())
+    pending = _pending_requests_from_checkpoint(checkpoint, request_executor_ids=request_executor_ids)
+
+    draft_preview: str | None = None
+    for entry in pending:
+        if entry.draft:
+            draft_preview = shorten(entry.draft, width=preview_width, placeholder="…")
+            break
+
+    status = "idle"
+    if pending:
+        status = "awaiting request response"
+    elif not checkpoint.messages and "finalise" in executor_ids:
+        status = "completed"
+    elif checkpoint.messages:
+        status = "awaiting next superstep"
+    elif request_executor_ids is not None and any(tid in targets for tid in request_executor_ids):
+        status = "awaiting request delivery"
+
+    return WorkflowCheckpointSummary(
+        checkpoint_id=checkpoint.checkpoint_id,
+        iteration_count=checkpoint.iteration_count,
+        targets=targets,
+        executor_ids=executor_ids,
+        status=status,
+        draft_preview=draft_preview,
+        pending_requests=pending,
+    )
+
+
+def _pending_requests_from_checkpoint(
+    checkpoint: WorkflowCheckpoint,
+    *,
+    request_executor_ids: Iterable[str] | None = None,
+) -> list[PendingRequestDetails]:
+    executor_filter: set[str] | None = None
+    if request_executor_ids is not None:
+        executor_filter = {str(value) for value in request_executor_ids}
+
+    pending: dict[str, PendingRequestDetails] = {}
+
+    for state in checkpoint.executor_states.values():
+        if not isinstance(state, Mapping):
+            continue
+        inner = state.get("pending_requests")
+        if isinstance(inner, Mapping):
+            for request_id, snapshot in inner.items():  # type: ignore[attr-defined]
+                _merge_snapshot(pending, str(request_id), snapshot)  # type: ignore[arg-type]
+
+    for source_id, message_list in checkpoint.messages.items():
+        if executor_filter is not None and source_id not in executor_filter:
+            continue
+        if not isinstance(message_list, list):
+            continue
+        for message in message_list:
+            if not isinstance(message, Mapping):
+                continue
+            payload = _decode_checkpoint_value(message.get("data"))
+            _merge_message_payload(pending, payload, message)
+
+    return list(pending.values())
+
+
+def _merge_snapshot(pending: dict[str, PendingRequestDetails], request_id: str, snapshot: Any) -> None:
+    if not request_id or not isinstance(snapshot, Mapping):
+        return
+
+    details = pending.setdefault(request_id, PendingRequestDetails(request_id=request_id))
+
+    _apply_update(
+        details,
+        prompt=snapshot.get("prompt"),  # type: ignore[attr-defined]
+        draft=snapshot.get("draft"),  # type: ignore[attr-defined]
+        iteration=snapshot.get("iteration"),  # type: ignore[attr-defined]
+        source_executor_id=snapshot.get("source_executor_id"),  # type: ignore[attr-defined]
+    )
+
+    extra = snapshot.get("details")  # type: ignore[attr-defined]
+    if isinstance(extra, Mapping):
+        _apply_update(
+            details,
+            prompt=extra.get("prompt"),  # type: ignore[attr-defined]
+            draft=extra.get("draft"),  # type: ignore[attr-defined]
+            iteration=extra.get("iteration"),  # type: ignore[attr-defined]
+        )
+
+
+def _merge_message_payload(
+    pending: dict[str, PendingRequestDetails],
+    payload: Any,
+    raw_message: Mapping[str, Any],
+) -> None:
+    if isinstance(payload, RequestResponse):
+        request_id = payload.request_id or _get_field(payload.original_request, "request_id")  # type: ignore[arg-type]
+        if not request_id:
+            return
+        details = pending.setdefault(request_id, PendingRequestDetails(request_id=request_id))
+        _apply_update(
+            details,
+            prompt=_get_field(payload.original_request, "prompt"),  # type: ignore[arg-type]
+            draft=_get_field(payload.original_request, "draft"),  # type: ignore[arg-type]
+            iteration=_get_field(payload.original_request, "iteration"),  # type: ignore[arg-type]
+            source_executor_id=raw_message.get("source_id"),
+            original_request=payload.original_request,  # type: ignore[arg-type]
+        )
+    elif isinstance(payload, RequestInfoMessage):
+        request_id = getattr(payload, "request_id", None)
+        if not request_id:
+            return
+        details = pending.setdefault(request_id, PendingRequestDetails(request_id=request_id))
+        _apply_update(
+            details,
+            prompt=getattr(payload, "prompt", None),
+            draft=getattr(payload, "draft", None),
+            iteration=getattr(payload, "iteration", None),
+            source_executor_id=raw_message.get("source_id"),
+            original_request=payload,
+        )
+
+
+def _apply_update(
+    details: PendingRequestDetails,
+    *,
+    prompt: Any = None,
+    draft: Any = None,
+    iteration: Any = None,
+    source_executor_id: Any = None,
+    original_request: Any = None,
+) -> None:
+    if prompt and not details.prompt:
+        details.prompt = str(prompt)
+    if draft and not details.draft:
+        details.draft = str(draft)
+    if iteration is not None and details.iteration is None:
+        coerced = _coerce_int(iteration)
+        if coerced is not None:
+            details.iteration = coerced
+    if source_executor_id and not details.source_executor_id:
+        details.source_executor_id = str(source_executor_id)
+    if original_request is not None and details.original_request is None:
+        details.original_request = original_request
+
+
+def _get_field(obj: Any, key: str) -> Any:
+    if obj is None:
+        return None
+    if isinstance(obj, Mapping):
+        return obj.get(key)  # type: ignore[attr-defined,return-value]
+    return getattr(obj, key, None)
+
+
+def _coerce_int(value: Any) -> int | None:
+    try:
+        return int(value)
+    except (TypeError, ValueError):
+        return None

python/packages/core/agent_framework/_workflows/_request_info_executor.py

@@ -5,17 +5,14 @@ import importlib
 import json
 import logging
 import uuid
-from collections.abc import Iterable, Mapping, Sequence
+from collections.abc import Mapping, Sequence
 from dataclasses import asdict, dataclass, field, fields, is_dataclass
-from textwrap import shorten
 from typing import Any, ClassVar, Generic, TypeVar, cast
 
-from ._checkpoint import WorkflowCheckpoint
 from ._events import (
     RequestInfoEvent,  # type: ignore[reportPrivateUsage]
 )
 from ._executor import Executor, handler
-from ._runner_context import _decode_checkpoint_value  # type: ignore
 from ._workflow_context import WorkflowContext
 
 logger = logging.getLogger(__name__)
@@ -34,16 +31,17 @@ class PendingRequestDetails:
 
 
 @dataclass
-class WorkflowCheckpointSummary:
-    """Human-readable summary of a workflow checkpoint."""
+class PendingRequestSnapshot:
+    """Snapshot of a pending request for internal tracking.
 
-    checkpoint_id: str
-    iteration_count: int
-    targets: list[str]
-    executor_states: list[str]
-    status: str
-    draft_preview: str | None
-    pending_requests: list[PendingRequestDetails]
+    This snapshot should be JSON-serializable and contain enough
+    information to reconstruct the original request if needed.
+    """
+
+    request_id: str
+    source_executor_id: str
+    request_type: str
+    request_as_json_safe_dict: dict[str, Any]
 
 
 @dataclass
@@ -108,8 +106,10 @@ class RequestInfoExecutor(Executor):
         super().__init__(id=id)
         self._request_events: dict[str, RequestInfoEvent] = {}
 
+    # region Public Methods
+
     @handler
-    async def run(self, message: RequestInfoMessage, ctx: WorkflowContext) -> None:
+    async def handle_request(self, message: RequestInfoMessage, ctx: WorkflowContext) -> None:
         """Run the RequestInfoExecutor with the given message."""
         # Use source_executor_id from message if available, otherwise fall back to context
         source_executor_id = message.source_executor_id or ctx.get_source_executor_id()
@@ -121,7 +121,7 @@ class RequestInfoExecutor(Executor):
             request_data=message,
         )
         self._request_events[message.request_id] = event
-        await self._record_pending_request_snapshot(message, source_executor_id, ctx)
+        await self._record_pending_request(message, source_executor_id, ctx)
         await ctx.add_event(event)
 
     async def handle_response(
@@ -139,7 +139,7 @@ class RequestInfoExecutor(Executor):
         """
         event = self._request_events.get(request_id)
         if event is None:
-            event = await self._rehydrate_request_event(request_id, ctx)
+            event = await self._rehydrate_request_event(request_id, cast(WorkflowContext, ctx))
         if event is None:
             raise ValueError(f"No request found with ID: {request_id}")
 
@@ -151,72 +151,20 @@ class RequestInfoExecutor(Executor):
         correlated_response = RequestResponse(data=response_data, original_request=event.data, request_id=request_id)
         await ctx.send_message(correlated_response, target_id=event.source_executor_id)
 
-        await self._clear_pending_request_snapshot(request_id, ctx)
-
-    async def _record_pending_request_snapshot(
-        self,
-        request: RequestInfoMessage,
-        source_executor_id: str,
-        ctx: WorkflowContext[Any],
-    ) -> None:
-        snapshot = self._build_request_snapshot(request, source_executor_id)
-
-        pending = await self._load_pending_request_state(ctx)
-        pending[request.request_id] = snapshot
-        await self._persist_pending_request_state(pending, ctx)
-        await self._write_executor_state(ctx, pending)
-
-    async def _clear_pending_request_snapshot(self, request_id: str, ctx: WorkflowContext[Any]) -> None:
-        pending = await self._load_pending_request_state(ctx)
-        if request_id in pending:
-            pending.pop(request_id, None)
-            await self._persist_pending_request_state(pending, ctx)
-        await self._write_executor_state(ctx, pending)
-
-    async def _load_pending_request_state(self, ctx: WorkflowContext[Any]) -> dict[str, Any]:
-        try:
-            existing = await ctx.get_shared_state(self._PENDING_SHARED_STATE_KEY)
-        except KeyError:
-            return {}
-        except Exception as exc:  # pragma: no cover - transport specific
-            logger.warning(f"RequestInfoExecutor {self.id} failed to read pending request state: {exc}")
-            return {}
-
-        if not isinstance(existing, dict):
-            if existing not in (None, {}):
-                logger.warning(
-                    f"RequestInfoExecutor {self.id} encountered non-dict pending state "
-                    f"({type(existing).__name__}); resetting."
-                )
-            return {}
-
-        return dict(existing)  # type: ignore[arg-type]
-
-    async def _persist_pending_request_state(self, pending: dict[str, Any], ctx: WorkflowContext[Any]) -> None:
-        await self._safe_set_shared_state(ctx, pending)
-        await self._safe_set_runner_state(ctx, pending)
-
-    async def _safe_set_shared_state(self, ctx: WorkflowContext[Any], pending: dict[str, Any]) -> None:
-        try:
-            await ctx.set_shared_state(self._PENDING_SHARED_STATE_KEY, pending)
-        except Exception as exc:  # pragma: no cover - transport specific
-            logger.warning(f"RequestInfoExecutor {self.id} failed to update shared pending state: {exc}")
-
-    async def _safe_set_runner_state(self, ctx: WorkflowContext[Any], pending: dict[str, Any]) -> None:
-        try:
-            await ctx.set_state({"pending_requests": pending})
-        except Exception as exc:  # pragma: no cover - transport specific
-            logger.warning(f"RequestInfoExecutor {self.id} failed to update runner state with pending requests: {exc}")
+        await self._erase_pending_request(request_id, cast(WorkflowContext, ctx))
 
     def snapshot_state(self) -> dict[str, Any]:
         """Serialize pending requests so checkpoint restoration can resume seamlessly."""
 
-        def _encode_event(event: RequestInfoEvent) -> dict[str, Any]:
-            request_data = event.data
-            payload: dict[str, Any]
-            data_cls = request_data.__class__ if request_data is not None else type(None)
+        def _encode_event(event: RequestInfoEvent) -> dict[str, Any] | None:
+            if event.data is None or not isinstance(event.data, RequestInfoMessage):
+                logger.warning(
+                    f"RequestInfoExecutor {self.id} encountered invalid event data for request ID {event.request_id}: "
+                    f"{type(event.data).__name__}. This request will be skipped in the checkpoint."
+                )
+                return None
 
-            payload = self._encode_request_payload(request_data, data_cls)
+            payload = self._encode_request_payload(event.data, event.data.__class__)
 
             return {
                 "source_executor_id": event.source_executor_id,
@@ -225,20 +173,125 @@ class RequestInfoExecutor(Executor):
             }
 
         return {
-            "request_events": {rid: _encode_event(event) for rid, event in self._request_events.items()},
+            "request_events": {
+                rid: encoded
+                for rid, event in self._request_events.items()
+                if (encoded := _encode_event(event)) is not None
+            },
         }
 
-    def _encode_request_payload(self, request_data: RequestInfoMessage | None, data_cls: type[Any]) -> dict[str, Any]:
-        if request_data is None or isinstance(request_data, (str, int, float, bool)):
-            return {
-                "kind": "raw",
-                "type": f"{data_cls.__module__}:{data_cls.__qualname__}",
-                "value": request_data,
-            }
+    def restore_state(self, state: dict[str, Any]) -> None:
+        """Restore pending request bookkeeping from checkpoint state."""
+        self._request_events.clear()
+        stored_events = state.get("request_events", {})
 
+        for request_id, payload in stored_events.items():
+            request_type_qual = payload.get("request_type", "")
+            try:
+                request_type = _import_qualname(request_type_qual)
+            except Exception as exc:  # pragma: no cover - defensive fallback
+                logger.debug(
+                    "RequestInfoExecutor %s failed to import %s during restore: %s",
+                    self.id,
+                    request_type_qual,
+                    exc,
+                )
+                request_type = RequestInfoMessage
+            request_data_meta = payload.get("request_data", {})
+            request_data = self._decode_request_data(request_data_meta)
+            event = RequestInfoEvent(
+                request_id=request_id,
+                source_executor_id=payload.get("source_executor_id", ""),
+                request_type=request_type,
+                request_data=request_data,
+            )
+            self._request_events[request_id] = event
+
+    async def has_pending_request(self, request_id: str, ctx: WorkflowContext) -> bool:
+        """Check if there is a pending request with the given ID.
+
+        Args:
+            request_id: The ID of the request to check.
+            ctx: The workflow context for accessing state if needed.
+
+        Returns: True if the request is pending, False otherwise.
+        """
+        if request_id in self._request_events:
+            return True
+
+        pending_requests = await self._retrieve_existing_pending_requests(ctx)
+        return request_id in pending_requests
+
+    # endregion: Public Methods
+
+    # region: Internal Methods
+
+    async def _record_pending_request(
+        self,
+        message: RequestInfoMessage,
+        source_executor_id: str,
+        ctx: WorkflowContext,
+    ) -> None:
+        """Record a pending request to the executor's state for checkpointing purposes."""
+        pending_request_snapshot = self._build_pending_request_snapshot(message, source_executor_id)
+
+        existing_pending_requests = await self._retrieve_existing_pending_requests(ctx)
+        existing_pending_requests[message.request_id] = pending_request_snapshot
+
+        await self._persist_to_executor_state(existing_pending_requests, ctx)
+
+    async def _erase_pending_request(self, request_id: str, ctx: WorkflowContext) -> None:
+        """Erase a pending request from the executor's state after it has been handled for checkpointing purposes."""
+        existing_pending_requests = await self._retrieve_existing_pending_requests(ctx)
+        if request_id in existing_pending_requests:
+            existing_pending_requests.pop(request_id)
+            await self._persist_to_executor_state(existing_pending_requests, ctx)
+
+    async def _retrieve_existing_pending_requests(self, ctx: WorkflowContext) -> dict[str, PendingRequestSnapshot]:
+        """Retrieve existing pending requests from executor state."""
+        executor_state = await ctx.get_state()
+        if executor_state is None:
+            return {}
+
+        stored_requests = executor_state.get(self._PENDING_SHARED_STATE_KEY, {})
+        if not isinstance(stored_requests, dict):
+            raise TypeError(f"Unexpected type for pending requests: {type(stored_requests).__name__}")
+
+        # Validate contents
+        for key, value in stored_requests.items():  # type: ignore
+            if not isinstance(key, str) or not isinstance(value, PendingRequestSnapshot):
+                raise TypeError(
+                    "Invalid pending request entry in executor state. "
+                    "Key must be `str` and value must be `PendingRequestSnapshot`."
+                )
+
+        return stored_requests  # type: ignore
+
+    async def _persist_to_executor_state(
+        self, pending: dict[str, PendingRequestSnapshot], ctx: WorkflowContext
+    ) -> None:
+        """Persist the current pending requests to the executor's state."""
+        executor_state = await ctx.get_state() or {}
+        executor_state[self._PENDING_SHARED_STATE_KEY] = pending
+        await ctx.set_state(executor_state)
+
+    def _build_pending_request_snapshot(
+        self, request: RequestInfoMessage, source_executor_id: str
+    ) -> PendingRequestSnapshot:
+        """Build a snapshot of the pending request for checkpointing."""
+        request_as_json_safe_dict = self._convert_request_to_json_safe_dict(request)
+
+        return PendingRequestSnapshot(
+            request_id=request.request_id,
+            source_executor_id=source_executor_id,
+            request_type=f"{type(request).__module__}:{type(request).__name__}",
+            request_as_json_safe_dict=request_as_json_safe_dict,
+        )
+
+    def _encode_request_payload(self, request_data: RequestInfoMessage, data_cls: type[Any]) -> dict[str, Any]:
         if is_dataclass(request_data) and not isinstance(request_data, type):
             dataclass_instance = cast(Any, request_data)
-            safe_value = self._make_json_safe(asdict(dataclass_instance))
+            safe_value = _make_json_safe(asdict(dataclass_instance))
             return {
                 "kind": "dataclass",
                 "type": f"{data_cls.__module__}:{data_cls.__qualname__}",
@@ -251,7 +304,7 @@ class RequestInfoExecutor(Executor):
                 dumped = to_dict_fn()
             except TypeError:
                 dumped = to_dict_fn()
-            safe_value = self._make_json_safe(dumped)
+            safe_value = _make_json_safe(dumped)
             return {
                 "kind": "dict",
                 "type": f"{data_cls.__module__}:{data_cls.__qualname__}",
@@ -278,76 +331,26 @@ class RequestInfoExecutor(Executor):
                     converted = json.loads(decoded)
                 except Exception:
                     converted = decoded
-            safe_value = self._make_json_safe(converted)
+            safe_value = _make_json_safe(converted)
             return {
                 "kind": "dict" if isinstance(converted, dict) else "json",
                 "type": f"{data_cls.__module__}:{data_cls.__qualname__}",
                 "value": safe_value,
             }
 
-        details = self._serialise_request_details(request_data)
-        if details is not None:
-            safe_value = self._make_json_safe(details)
-            return {
-                "kind": "raw",
-                "type": f"{data_cls.__module__}:{data_cls.__qualname__}",
-                "value": safe_value,
-            }
-
-        safe_value = self._make_json_safe(request_data)
         return {
             "kind": "raw",
             "type": f"{data_cls.__module__}:{data_cls.__qualname__}",
-            "value": safe_value,
+            "value": self._convert_request_to_json_safe_dict(request_data),
         }
 
-    def restore_state(self, state: dict[str, Any]) -> None:
-        """Restore pending request bookkeeping from checkpoint state."""
-        self._request_events.clear()
-        stored_events = state.get("request_events", {})
-
-        for request_id, payload in stored_events.items():
-            request_type_qual = payload.get("request_type", "")
-            try:
-                request_type = self._import_qualname(request_type_qual)
-            except Exception as exc:  # pragma: no cover - defensive fallback
-                logger.debug(
-                    "RequestInfoExecutor %s failed to import %s during restore: %s",
-                    self.id,
-                    request_type_qual,
-                    exc,
-                )
-                request_type = RequestInfoMessage
-            request_data_meta = payload.get("request_data", {})
-            request_data = self._decode_request_data(request_data_meta)
-            event = RequestInfoEvent(
-                request_id=request_id,
-                source_executor_id=payload.get("source_executor_id", ""),
-                request_type=request_type,
-                request_data=request_data,
-            )
-            self._request_events[request_id] = event
-
-    @staticmethod
-    def _import_qualname(qualname: str) -> type[Any]:
-        module_name, _, type_name = qualname.partition(":")
-        if not module_name or not type_name:
-            raise ValueError(f"Invalid qualified name: {qualname}")
-        module = importlib.import_module(module_name)
-        attr: Any = module
-        for part in type_name.split("."):
-            attr = getattr(attr, part)
-        if not isinstance(attr, type):
-            raise TypeError(f"Resolved object is not a type: {qualname}")
-        return attr
-
     def _decode_request_data(self, metadata: dict[str, Any]) -> RequestInfoMessage:
         kind = metadata.get("kind")
         type_name = metadata.get("type", "")
         value: Any = metadata.get("value", {})
         if type_name:
             try:
-                imported = self._import_qualname(type_name)
+                imported = _import_qualname(type_name)
             except Exception as exc:  # pragma: no cover - defensive fallback
                 logger.debug(
                     "RequestInfoExecutor %s failed to import %s during decode: %s",
@@ -396,114 +399,22 @@ class RequestInfoExecutor(Executor):
             return target_cls()
         return RequestInfoMessage()
 
-    async def _write_executor_state(self, ctx: WorkflowContext[Any], pending: dict[str, Any]) -> None:
-        state = self.snapshot_state()
-        state["pending_requests"] = pending
+    def _convert_request_to_json_safe_dict(self, request: RequestInfoMessage) -> dict[str, Any]:
         try:
-            await ctx.set_state(state)
-        except Exception as exc:  # pragma: no cover - transport specific
-            logger.warning(f"RequestInfoExecutor {self.id} failed to persist executor state: {exc}")
-
-    def _build_request_snapshot(
-        self,
-        request: RequestInfoMessage,
-        source_executor_id: str,
-    ) -> dict[str, Any]:
-        snapshot: dict[str, Any] = {
-            "request_id": request.request_id,
-            "source_executor_id": source_executor_id,
-            "request_type": f"{type(request).__module__}:{type(request).__name__}",
-            "summary": repr(request),
-        }
-
-        details = self._serialise_request_details(request)
-        if details:
-            snapshot["details"] = details
-            for key in ("prompt", "draft", "iteration"):
-                if key in details and key not in snapshot:
-                    snapshot[key] = details[key]
-
-        return snapshot
-
-    def _serialise_request_details(self, request: RequestInfoMessage) -> dict[str, Any] | None:
-        if is_dataclass(request):
-            data = self._make_json_safe(asdict(request))
+            data = _make_json_safe(asdict(request))
             if isinstance(data, dict):
                 return cast(dict[str, Any], data)
-            return None
+            raise ValueError(f"Failed to convert {type(request).__name__} to dict")
+        except Exception as exc:
+            logger.error(f"RequestInfoExecutor {self.id} failed to serialize request: {exc}")
+            raise RuntimeError(
+                f"Failed to serialize request `{type(request).__name__}`: {exc}\n"
+                "Make sure request is a dataclass and derive from `RequestInfoMessage`."
+            ) from exc
 
-        to_dict = getattr(request, "to_dict", None)
-        if callable(to_dict):
-            try:
-                dump = self._make_json_safe(to_dict())
-            except TypeError:
-                dump = self._make_json_safe(to_dict())
-            if isinstance(dump, dict):
-                return cast(dict[str, Any], dump)
-            return None
-
-        to_json = getattr(request, "to_json", None)
-        if callable(to_json):
-            try:
-                raw = to_json()
-            except TypeError:
-                raw = to_json()
-            converted = raw
-            if isinstance(raw, (str, bytes, bytearray)):
-                decoded: str | bytes | bytearray
-                if isinstance(raw, (bytes, bytearray)):
-                    try:
-                        decoded = raw.decode()
-                    except Exception:
-                        decoded = raw
-                else:
-                    decoded = raw
-                try:
-                    converted = json.loads(decoded)
-                except Exception:
-                    converted = decoded
-            dump = self._make_json_safe(converted)
-            if isinstance(dump, dict):
-                return cast(dict[str, Any], dump)
-            return None
-
-        attrs = getattr(request, "__dict__", None)
-        if isinstance(attrs, dict):
-            cleaned = self._make_json_safe(attrs)
-            if isinstance(cleaned, dict):
-                return cast(dict[str, Any], cleaned)
-
-        return None
-
-    def _make_json_safe(self, value: Any) -> Any:
-        if value is None or isinstance(value, (str, int, float, bool)):
-            return value
-        if isinstance(value, Mapping):
-            safe_dict: dict[str, Any] = {}
-            for key, val in value.items():  # type: ignore[attr-defined]
-                safe_dict[str(key)] = self._make_json_safe(val)  # type: ignore[arg-type]
-            return safe_dict
-        if isinstance(value, Sequence) and not isinstance(value, (str, bytes, bytearray)):
-            return [self._make_json_safe(item) for item in value]  # type: ignore[misc]
-        return repr(value)
-
-    async def has_pending_request(self, request_id: str, ctx: WorkflowContext[Any]) -> bool:
-        if request_id in self._request_events:
-            return True
-        snapshot = await self._get_pending_request_snapshot(request_id, ctx)
-        return snapshot is not None
-
-    async def _rehydrate_request_event(
-        self,
-        request_id: str,
-        ctx: WorkflowContext[Any],
-    ) -> RequestInfoEvent | None:
-        snapshot = await self._get_pending_request_snapshot(request_id, ctx)
-        if snapshot is None:
-            return None
-
-        source_executor_id = snapshot.get("source_executor_id")
-        if not isinstance(source_executor_id, str) or not source_executor_id:
+    async def _rehydrate_request_event(self, request_id: str, ctx: WorkflowContext) -> RequestInfoEvent | None:
+        pending_requests = await self._retrieve_existing_pending_requests(ctx)
+        if (snapshot := pending_requests.get(request_id)) is None:
             return None
 
         request = self._construct_request_from_snapshot(snapshot)
@@ -512,57 +423,18 @@ class RequestInfoExecutor(Executor):
 
         event = RequestInfoEvent(
             request_id=request_id,
-            source_executor_id=source_executor_id,
+            source_executor_id=snapshot.source_executor_id,
             request_type=type(request),
             request_data=request,
         )
         self._request_events[request_id] = event
         return event
 
-    async def _get_pending_request_snapshot(self, request_id: str, ctx: WorkflowContext[Any]) -> dict[str, Any] | None:
-        pending = await self._collect_pending_request_snapshots(ctx)
-        snapshot = pending.get(request_id)
-        if snapshot is None:
-            return None
-        return snapshot
-
-    async def _collect_pending_request_snapshots(self, ctx: WorkflowContext[Any]) -> dict[str, dict[str, Any]]:
-        combined: dict[str, dict[str, Any]] = {}
-
-        try:
-            shared_pending = await ctx.get_shared_state(self._PENDING_SHARED_STATE_KEY)
-        except KeyError:
-            shared_pending = None
-        except Exception as exc:  # pragma: no cover - transport specific
-            logger.warning(f"RequestInfoExecutor {self.id} failed to read shared pending state during rehydrate: {exc}")
-            shared_pending = None
-
-        if isinstance(shared_pending, dict):
-            for key, value in shared_pending.items():  # type: ignore[attr-defined]
-                if isinstance(key, str) and isinstance(value, dict):
-                    combined[key] = cast(dict[str, Any], value)
-
-        try:
-            state = await ctx.get_state()
-        except Exception as exc:  # pragma: no cover - transport specific
-            logger.warning(f"RequestInfoExecutor {self.id} failed to read runner state during rehydrate: {exc}")
-            state = None
-
-        if isinstance(state, dict):
-            state_pending = state.get("pending_requests")
-            if isinstance(state_pending, dict):
-                for key, value in state_pending.items():  # type: ignore[attr-defined]
-                    if isinstance(key, str) and isinstance(value, dict) and key not in combined:
-                        combined[key] = cast(dict[str, Any], value)
-
-        return combined
-
-    def _construct_request_from_snapshot(self, snapshot: dict[str, Any]) -> RequestInfoMessage | None:
-        details_raw = snapshot.get("details")
-        details: dict[str, Any] = cast(dict[str, Any], details_raw) if isinstance(details_raw, dict) else {}
+    def _construct_request_from_snapshot(self, snapshot: PendingRequestSnapshot) -> RequestInfoMessage | None:
+        json_safe_dict = snapshot.request_as_json_safe_dict
 
         request_cls: type[RequestInfoMessage] = RequestInfoMessage
-        request_type_str = snapshot.get("request_type")
+        request_type_str = snapshot.request_type
         if isinstance(request_type_str, str) and ":" in request_type_str:
             module_name, class_name = request_type_str.split(":", 1)
             try:
@@ -574,19 +446,19 @@ class RequestInfoExecutor(Executor):
                 logger.warning(f"RequestInfoExecutor {self.id} could not import {module_name}.{class_name}: {exc}")
                 request_cls = RequestInfoMessage
 
-        request: RequestInfoMessage | None = self._instantiate_request(request_cls, details)
+        request: RequestInfoMessage | None = self._instantiate_request(request_cls, json_safe_dict)
 
         if request is None and request_cls is not RequestInfoMessage:
-            request = self._instantiate_request(RequestInfoMessage, details)
+            request = self._instantiate_request(RequestInfoMessage, json_safe_dict)
 
         if request is None:
             logger.warning(
                 f"RequestInfoExecutor {self.id} could not reconstruct request "
-                f"{request_type_str or RequestInfoMessage.__name__} from snapshot keys {sorted(details.keys())}"
+                f"{request_type_str or RequestInfoMessage.__name__} from snapshot keys {sorted(json_safe_dict.keys())}"
             )
             return None
 
-        for key, value in details.items():
+        for key, value in json_safe_dict.items():
             if key == "request_id":
                 continue
             try:
@@ -597,7 +469,7 @@ class RequestInfoExecutor(Executor):
                 )
                 continue
 
-        snapshot_request_id = snapshot.get("request_id")
+        snapshot_request_id = snapshot.request_id
         if isinstance(snapshot_request_id, str) and snapshot_request_id:
             try:
                 request.request_id = snapshot_request_id
@@ -664,178 +536,38 @@ class RequestInfoExecutor(Executor):
 
         return instance
 
-    @staticmethod
-    def pending_requests_from_checkpoint(
-        checkpoint: WorkflowCheckpoint,
-        *,
-        request_executor_ids: Iterable[str] | None = None,
-    ) -> list[PendingRequestDetails]:
-        executor_filter: set[str] | None = None
-        if request_executor_ids is not None:
-            executor_filter = {str(value) for value in request_executor_ids}
+    # endregion: Internal Methods
 
-        pending: dict[str, PendingRequestDetails] = {}
 
-        shared_map = checkpoint.shared_state.get(RequestInfoExecutor._PENDING_SHARED_STATE_KEY)
-        if isinstance(shared_map, Mapping):
-            for request_id, snapshot in shared_map.items():  # type: ignore[attr-defined]
-                RequestInfoExecutor._merge_snapshot(pending, str(request_id), snapshot)  # type: ignore[arg-type]
+# region: Utility Functions
 
-        for state in checkpoint.executor_states.values():
-            if not isinstance(state, Mapping):
-                continue
-            inner = state.get("pending_requests")
-            if isinstance(inner, Mapping):
-                for request_id, snapshot in inner.items():  # type: ignore[attr-defined]
-                    RequestInfoExecutor._merge_snapshot(pending, str(request_id), snapshot)  # type: ignore[arg-type]
 
-        for source_id, message_list in checkpoint.messages.items():
-            if executor_filter is not None and source_id not in executor_filter:
-                continue
-            if not isinstance(message_list, list):
-                continue
-            for message in message_list:
-                if not isinstance(message, Mapping):
-                    continue
-                payload = _decode_checkpoint_value(message.get("data"))
-                RequestInfoExecutor._merge_message_payload(pending, payload, message)
+def _make_json_safe(value: Any) -> Any:
+    """Recursively convert a value to a JSON-safe representation."""
+    if value is None or isinstance(value, (str, int, float, bool)):
+        return value
+    if isinstance(value, Mapping):
+        safe_dict: dict[str, Any] = {}
+        for key, val in value.items():  # type: ignore[attr-defined]
+            safe_dict[str(key)] = _make_json_safe(val)  # type: ignore[arg-type]
+        return safe_dict
+    if isinstance(value, Sequence) and not isinstance(value, (str, bytes, bytearray)):
+        return [_make_json_safe(item) for item in value]  # type: ignore[misc]
+    return repr(value)
 
-        return list(pending.values())
 
-    @staticmethod
-    def checkpoint_summary(
-        checkpoint: WorkflowCheckpoint,
-        *,
-        request_executor_ids: Iterable[str] | None = None,
-        preview_width: int = 70,
-    ) -> WorkflowCheckpointSummary:
-        targets = sorted(checkpoint.messages.keys())
-        executor_states = sorted(checkpoint.executor_states.keys())
-        pending = RequestInfoExecutor.pending_requests_from_checkpoint(
-            checkpoint, request_executor_ids=request_executor_ids
-        )
+def _import_qualname(qualname: str) -> type[Any]:
+    """Import a type given its qualified name in the format 'module:TypeName'."""
+    module_name, _, type_name = qualname.partition(":")
+    if not module_name or not type_name:
+        raise ValueError(f"Invalid qualified name: {qualname}")
+    module = importlib.import_module(module_name)
+    attr: Any = module
+    for part in type_name.split("."):
+        attr = getattr(attr, part)
+    if not isinstance(attr, type):
+        raise TypeError(f"Resolved object is not a type: {qualname}")
+    return attr
 
-        draft_preview: str | None = None
-        for entry in pending:
-            if entry.draft:
-                draft_preview = shorten(entry.draft, width=preview_width, placeholder="…")
-                break
 
-        status = "idle"
-        if pending:
-            status = "awaiting human response"
-        elif not checkpoint.messages and "finalise" in executor_states:
-            status = "completed"
-        elif checkpoint.messages:
-            status = "awaiting next superstep"
-        elif request_executor_ids is not None and any(tid in targets for tid in request_executor_ids):
-            status = "awaiting request delivery"
-
-        return WorkflowCheckpointSummary(
-            checkpoint_id=checkpoint.checkpoint_id,
-            iteration_count=checkpoint.iteration_count,
-            targets=targets,
-            executor_states=executor_states,
-            status=status,
-            draft_preview=draft_preview,
-            pending_requests=pending,
-        )
-
-    @staticmethod
-    def _merge_snapshot(
-        pending: dict[str, PendingRequestDetails],
-        request_id: str,
-        snapshot: Any,
-    ) -> None:
-        if not request_id or not isinstance(snapshot, Mapping):
-            return
-
-        details = pending.setdefault(request_id, PendingRequestDetails(request_id=request_id))
-
-        RequestInfoExecutor._apply_update(
-            details,
-            prompt=snapshot.get("prompt"),  # type: ignore[attr-defined]
-            draft=snapshot.get("draft"),  # type: ignore[attr-defined]
-            iteration=snapshot.get("iteration"),  # type: ignore[attr-defined]
-            source_executor_id=snapshot.get("source_executor_id"),  # type: ignore[attr-defined]
-        )
-
-        extra = snapshot.get("details")  # type: ignore[attr-defined]
-        if isinstance(extra, Mapping):
-            RequestInfoExecutor._apply_update(
-                details,
-                prompt=extra.get("prompt"),  # type: ignore[attr-defined]
-                draft=extra.get("draft"),  # type: ignore[attr-defined]
-                iteration=extra.get("iteration"),  # type: ignore[attr-defined]
-            )
-
-    @staticmethod
-    def _merge_message_payload(
-        pending: dict[str, PendingRequestDetails],
-        payload: Any,
-        raw_message: Mapping[str, Any],
-    ) -> None:
-        if isinstance(payload, RequestResponse):
-            request_id = payload.request_id or RequestInfoExecutor._get_field(payload.original_request, "request_id")  # type: ignore[arg-type]
-            if not request_id:
-                return
-            details = pending.setdefault(request_id, PendingRequestDetails(request_id=request_id))
-            RequestInfoExecutor._apply_update(
-                details,
-                prompt=RequestInfoExecutor._get_field(payload.original_request, "prompt"),  # type: ignore[arg-type]
-                draft=RequestInfoExecutor._get_field(payload.original_request, "draft"),  # type: ignore[arg-type]
-                iteration=RequestInfoExecutor._get_field(payload.original_request, "iteration"),  # type: ignore[arg-type]
-                source_executor_id=raw_message.get("source_id"),
-                original_request=payload.original_request,  # type: ignore[arg-type]
-            )
-        elif isinstance(payload, RequestInfoMessage):
-            request_id = getattr(payload, "request_id", None)
-            if not request_id:
-                return
-            details = pending.setdefault(request_id, PendingRequestDetails(request_id=request_id))
-            RequestInfoExecutor._apply_update(
-                details,
-                prompt=getattr(payload, "prompt", None),
-                draft=getattr(payload, "draft", None),
-                iteration=getattr(payload, "iteration", None),
-                source_executor_id=raw_message.get("source_id"),
-                original_request=payload,
-            )
-
-    @staticmethod
-    def _apply_update(
-        details: PendingRequestDetails,
-        *,
-        prompt: Any = None,
-        draft: Any = None,
-        iteration: Any = None,
-        source_executor_id: Any = None,
-        original_request: Any = None,
-    ) -> None:
-        if prompt and not details.prompt:
-            details.prompt = str(prompt)
-        if draft and not details.draft:
-            details.draft = str(draft)
-        if iteration is not None and details.iteration is None:
-            coerced = RequestInfoExecutor._coerce_int(iteration)
-            if coerced is not None:
-                details.iteration = coerced
-        if source_executor_id and not details.source_executor_id:
-            details.source_executor_id = str(source_executor_id)
-        if original_request is not None and details.original_request is None:
-            details.original_request = original_request
-
-    @staticmethod
-    def _get_field(obj: Any, key: str) -> Any:
-        if obj is None:
-            return None
-        if isinstance(obj, Mapping):
-            return obj.get(key)  # type: ignore[attr-defined,return-value]
-        return getattr(obj, key, None)
-
-    @staticmethod
-    def _coerce_int(value: Any) -> int | None:
-        try:
-            return int(value)
-        except (TypeError, ValueError):
-            return None
+# endregion: Utility Functions

python/packages/core/agent_framework/_workflows/_runner.py

@@ -304,6 +304,8 @@ class Runner:
                     checkpoint_id,
                 )
 
+            await self._restore_executor_states(checkpoint.executor_states)
+
             state = self._checkpoint_to_state(checkpoint)
             await self._ctx.set_checkpoint_state(state)
             if checkpoint.workflow_id:
@@ -323,6 +325,27 @@ class Runner:
             logger.error(f"Failed to restore from checkpoint {checkpoint_id}: {e}")
             return False
 
+    async def _restore_executor_states(self, executor_states: dict[str, dict[str, Any]]) -> None:
+        for exec_id, state in executor_states.items():
+            executor = self._executors.get(exec_id)
+            if not executor:
+                logger.debug(f"Executor {exec_id} not found during state restoration; skipping.")
+                continue
+
+            restored = False
+            restore_method = getattr(executor, "restore_state", None)
+            try:
+                if callable(restore_method):
+                    maybe = restore_method(state)
+                    if asyncio.iscoroutine(maybe):  # type: ignore[arg-type]
+                        await maybe  # type: ignore[arg-type]
+                    restored = True
+            except Exception as ex:  # pragma: no cover - defensive
+                logger.debug(f"Executor {exec_id} restore_state failed: {ex}")
+
+            if not restored:
+                logger.debug(f"Executor {exec_id} does not support state restoration; skipping.")
+
     async def _restore_shared_state_from_context(self) -> None:
         try:
             restored_state = await self._ctx.get_checkpoint_state()

python/packages/core/agent_framework/_workflows/_workflow_context.py

@@ -441,14 +441,11 @@ class WorkflowContext(Generic[T_Out, T_W_Out]):
         Executors call this with a JSON-serializable dict capturing the minimal
         state needed to resume. It replaces any previously stored state.
         """
-        if hasattr(self._runner_context, "set_state"):
-            await self._runner_context.set_state(self._executor_id, state)  # type: ignore[arg-type]
+        await self._runner_context.set_state(self._executor_id, state)
 
     async def get_state(self) -> dict[str, Any] | None:
         """Retrieve previously persisted state for this executor, if any."""
-        if hasattr(self._runner_context, "get_state"):
-            return await self._runner_context.get_state(self._executor_id)  # type: ignore[return-value]
-        return None
+        return await self._runner_context.get_state(self._executor_id)
 
     def is_streaming(self) -> bool:
         """Check if the workflow is running in streaming mode.

python/packages/core/tests/workflow/test_request_info_executor_rehydrate.py

@@ -6,17 +6,19 @@ from datetime import datetime, timezone
 from typing import Any
 
 from agent_framework._workflows._checkpoint import CheckpointStorage, WorkflowCheckpoint
+from agent_framework._workflows._checkpoint_summary import get_checkpoint_summary
 from agent_framework._workflows._events import RequestInfoEvent, WorkflowEvent
 from agent_framework._workflows._request_info_executor import (
     PendingRequestDetails,
+    PendingRequestSnapshot,
     RequestInfoExecutor,
     RequestInfoMessage,
     RequestResponse,
 )
-from agent_framework._workflows._runner_context import (  # type: ignore
+from agent_framework._workflows._runner_context import (
     CheckpointState,
     Message,
-    _encode_checkpoint_value,
+    _encode_checkpoint_value,  # type: ignore
 )
 from agent_framework._workflows._shared_state import SharedState
 from agent_framework._workflows._workflow_context import WorkflowContext
@@ -85,6 +87,12 @@ class _StubRunnerContext:
     async def set_checkpoint_state(self, state: CheckpointState) -> None:  # pragma: no cover - unused
         pass
 
+    def set_streaming(self, streaming: bool) -> None:  # pragma: no cover - unused
+        pass
+
+    def is_streaming(self) -> bool:  # pragma: no cover - unused
+        return False
+
 
 @dataclass(kw_only=True)
 class SimpleApproval(RequestInfoMessage):
@@ -109,30 +117,18 @@ async def test_rehydrate_falls_back_when_request_type_missing() -> None:
     This simulates resuming a workflow where the HumanApprovalRequest class is unavailable
     in the current process (e.g., defined in __main__ during the original run).
     """
-
     request_id = "request-123"
-    snapshot = {
-        "request_id": request_id,
-        "source_executor_id": "review_gateway",
-        "request_type": "nonexistent.module:MissingRequest",
-        "summary": "...",
-        "details": {
+    snapshot = PendingRequestSnapshot(
+        request_id=request_id,
+        source_executor_id="review_gateway",
+        request_type="nonexistent.module:MissingRequest",
+        request_as_json_safe_dict={
             "request_id": request_id,
-            "prompt": "Review draft",
-            "draft": "Draft text",
-            "iteration": 2,
         },
-    }
+    )
 
-    shared_state = SharedState()
-    async with shared_state.hold():
-        await shared_state.set_within_hold(
-            PENDING_STATE_KEY,
-            {request_id: snapshot},
-        )
-
-    runner_ctx = _StubRunnerContext({"pending_requests": {request_id: snapshot}})
-    ctx: WorkflowContext[Any] = WorkflowContext("request_info", ["workflow"], shared_state, runner_ctx)
+    runner_ctx = _StubRunnerContext({PENDING_STATE_KEY: {request_id: snapshot}})
+    ctx: WorkflowContext[Any] = WorkflowContext("request_info", ["workflow"], SharedState(), runner_ctx)
 
     executor = RequestInfoExecutor(id="request_info")
 
@@ -141,31 +137,21 @@ async def test_rehydrate_falls_back_when_request_type_missing() -> None:
     assert event is not None
     assert event.request_id == request_id
     assert isinstance(event.data, RequestInfoMessage)
-    assert getattr(event.data, "prompt", None) == "Review draft"
-    assert getattr(event.data, "iteration", None) == 2
 
 
 async def test_has_pending_request_detects_snapshot() -> None:
-    request_id = "req-pending"
-    snapshot = {
-        "request_id": request_id,
-        "source_executor_id": "review_gateway",
-        "details": {
+    request_id = "request-123"
+    snapshot = PendingRequestSnapshot(
+        request_id=request_id,
+        source_executor_id="review_gateway",
+        request_type="nonexistent.module:MissingRequest",
+        request_as_json_safe_dict={
             "request_id": request_id,
-            "prompt": "Review",
-            "draft": "Draft",
         },
-    }
+    )
 
-    shared_state = SharedState()
-    async with shared_state.hold():
-        await shared_state.set_within_hold(
-            PENDING_STATE_KEY,
-            {request_id: snapshot},
-        )
-
-    runner_ctx = _StubRunnerContext({"pending_requests": {request_id: snapshot}})
-    ctx: WorkflowContext[Any] = WorkflowContext("request_info", ["workflow"], shared_state, runner_ctx)
+    runner_ctx = _StubRunnerContext({PENDING_STATE_KEY: {request_id: snapshot}})
+    ctx: WorkflowContext[Any] = WorkflowContext("request_info", ["workflow"], SharedState(), runner_ctx)
 
     executor = RequestInfoExecutor(id="request_info")
 
@@ -221,7 +207,12 @@ def test_pending_requests_from_checkpoint_and_summary() -> None:
         iteration_count=1,
     )
 
-    pending = RequestInfoExecutor.pending_requests_from_checkpoint(checkpoint)
+    summary = get_checkpoint_summary(checkpoint)
+    assert summary.checkpoint_id == "cp-1"
+    assert summary.status == "awaiting request response"
+    assert summary.pending_requests[0].request_id == "req-42"
+
+    pending = summary.pending_requests
     assert len(pending) == 1
     entry = pending[0]
     assert isinstance(entry, PendingRequestDetails)
@@ -231,11 +222,6 @@ def test_pending_requests_from_checkpoint_and_summary() -> None:
     assert entry.iteration == 3
     assert entry.original_request is not None
 
-    summary = RequestInfoExecutor.checkpoint_summary(checkpoint)
-    assert summary.checkpoint_id == "cp-1"
-    assert summary.status == "awaiting human response"
-    assert summary.pending_requests[0].request_id == "req-42"
-
 
 def test_snapshot_state_serializes_non_json_payloads() -> None:
     executor = RequestInfoExecutor(id="request_info")
@@ -305,13 +291,10 @@ async def test_run_persists_pending_requests_in_runner_state() -> None:
     await executor.execute(approval, ctx.source_executor_ids, shared_state, runner_ctx)
 
     # Runner state should include both pending snapshot and serialized request events
-    assert "pending_requests" in runner_ctx._state  # pyright: ignore[reportPrivateUsage]
-    assert approval.request_id in runner_ctx._state["pending_requests"]  # pyright: ignore[reportPrivateUsage]
-    assert "request_events" in runner_ctx._state  # pyright: ignore[reportPrivateUsage]
-    assert approval.request_id in runner_ctx._state["request_events"]  # pyright: ignore[reportPrivateUsage]
+    assert PENDING_STATE_KEY in runner_ctx._state  # pyright: ignore[reportPrivateUsage]
+    assert approval.request_id in runner_ctx._state[PENDING_STATE_KEY]  # pyright: ignore[reportPrivateUsage]
 
     response_ctx: WorkflowContext[None] = WorkflowContext("request_info", ["source"], shared_state, runner_ctx)
     await executor.handle_response("approved", approval.request_id, response_ctx)  # type: ignore
 
-    assert runner_ctx._state["pending_requests"] == {}  # pyright: ignore[reportPrivateUsage]
-    assert runner_ctx._state.get("request_events", {}).get(approval.request_id) is None  # pyright: ignore[reportPrivateUsage]
+    assert runner_ctx._state[PENDING_STATE_KEY] == {}  # pyright: ignore[reportPrivateUsage]

python/packages/core/tests/workflow/test_workflow_states.py

@@ -1,7 +1,6 @@
 # Copyright (c) Microsoft. All rights reserved.
 
 from dataclasses import dataclass
-from typing import Any
 
 import pytest
 from typing_extensions import Never
@@ -25,7 +24,6 @@ from agent_framework import (
     WorkflowStatusEvent,
     handler,
 )
-from agent_framework import WorkflowContext as WFContext
 
 
 class FailingExecutor(Executor):
@@ -182,39 +180,3 @@ class SnapshotRequester(Executor):
     @handler
     async def ask(self, _: str, ctx: WorkflowContext[SnapshotRequest]) -> None:  # pragma: no cover - simple helper
         await ctx.send_message(SnapshotRequest(prompt=self._prompt, draft=self._draft, iteration=1))
-
-
-async def test_request_info_executor_tracks_pending_requests_via_shared_state():
-    prompt = "Review the launch copy"
-    draft = "Limited edition grinder now $249"
-    requester = SnapshotRequester(id="snapshot_req", prompt=prompt, draft=draft)
-    request_info = RequestInfoExecutor(id="request_info")
-
-    wf = WorkflowBuilder().set_start_executor(requester).add_edge(requester, request_info).build()
-
-    events = [event async for event in wf.run_stream("start")]
-    assert any(isinstance(event, RequestInfoEvent) for event in events)
-
-    pending_map: dict[str, Any] = await wf._shared_state.get(RequestInfoExecutor._PENDING_SHARED_STATE_KEY)  # type: ignore[reportPrivateUsage]
-    assert isinstance(pending_map, dict)
-    assert len(pending_map) == 1
-    snapshot: dict[str, Any] = next(iter(pending_map.values()))
-    assert snapshot["prompt"] == prompt
-    assert snapshot["draft"] == draft
-    assert snapshot.get("iteration") == 1
-
-    request_id: str = snapshot["request_id"]
-
-    request_info_resume = RequestInfoExecutor(id="request_info_resume")
-    resume_context: WFContext[Any] = WFContext(
-        executor_id=request_info_resume.id,
-        source_executor_ids=[wf.__class__.__name__],
-        shared_state=wf._shared_state,  # type: ignore[reportPrivateUsage]
-        runner_context=wf._runner_context,  # type: ignore[reportPrivateUsage]
-    )
-
-    await request_info_resume.handle_response("approve", request_id, resume_context)
-
-    updated_pending: dict[str, Any] = await wf._shared_state.get(RequestInfoExecutor._PENDING_SHARED_STATE_KEY)  # type: ignore[reportPrivateUsage]
-    assert isinstance(updated_pending, dict)
-    assert request_id not in updated_pending