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

python/samples/_run_all_samples.py

@@ -0,0 +1,235 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""
+Script to run all Python samples in the samples directory concurrently.
+This script will run all samples and report results at the end.
+
+Note: This script is AI generated. This is for internal validation purposes only.
+
+Samples that require human interaction are known to fail.
+
+Usage:
+    python run_all_samples.py                          # Run all samples using uv run (concurrent)
+    python run_all_samples.py --direct                 # Run all samples directly (concurrent,
+                                                       # assumes environment is set up)
+    python run_all_samples.py --subdir <directory>     # Run samples only in specific subdirectory
+    python run_all_samples.py --subdir getting_started/workflows  # Example: run only workflow samples
+"""
+
+import argparse
+import os
+import subprocess
+import sys
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from pathlib import Path
+
+
+def find_python_samples(samples_dir: Path, subdir: str | None = None) -> list[Path]:
+    """Find all Python sample files in the samples directory or a subdirectory."""
+    python_files: list[Path] = []
+
+    # Determine the search directory
+    if subdir:
+        search_dir = samples_dir / subdir
+        if not search_dir.exists():
+            print(f"Warning: Subdirectory '{subdir}' does not exist in {samples_dir}")
+            return []
+        print(f"Searching in subdirectory: {search_dir}")
+    else:
+        search_dir = samples_dir
+        print(f"Searching in all samples: {search_dir}")
+
+    # Walk through all subdirectories and find .py files
+    for root, dirs, files in os.walk(search_dir):
+        # Skip __pycache__ directories
+        dirs[:] = [d for d in dirs if d != "__pycache__"]
+
+        for file in files:
+            if file.endswith(".py") and not file.startswith("_") and file != "_run_all_samples.py":
+                python_files.append(Path(root) / file)
+
+    # Sort files for consistent execution order
+    return sorted(python_files)
+
+
+def run_sample(
+    sample_path: Path,
+    use_uv: bool = True,
+    python_root: Path | None = None,
+) -> tuple[bool, str, str]:
+    """
+    Run a single sample file using subprocess and return (success, output, error_info).
+
+    Args:
+        sample_path: Path to the sample file
+        use_uv: Whether to use uv run
+        python_root: Root directory for uv run
+
+    Returns:
+        Tuple of (success, output, error_info)
+    """
+    if use_uv and python_root:
+        cmd = ["uv", "run", "python", str(sample_path)]
+        cwd = python_root
+    else:
+        cmd = [sys.executable, sample_path.name]
+        cwd = sample_path.parent
+
+    try:
+        result = subprocess.run(
+            cmd,
+            cwd=cwd,
+            capture_output=True,
+            text=True,
+            timeout=60,  # 60 second timeout
+        )
+
+        if result.returncode == 0:
+            output = result.stdout.strip() if result.stdout.strip() else "No output"
+            return True, output, ""
+
+        error_info = f"Exit code: {result.returncode}"
+        if result.stderr.strip():
+            error_info += f"\nSTDERR: {result.stderr}"
+
+        return False, result.stdout.strip() if result.stdout.strip() else "", error_info
+
+    except subprocess.TimeoutExpired:
+        return False, "", f"TIMEOUT: {sample_path.name} (exceeded 60 seconds)"
+    except Exception as e:
+        return False, "", f"ERROR: {sample_path.name} - Exception: {str(e)}"
+
+
+def parse_arguments() -> argparse.Namespace:
+    """Parse command line arguments."""
+    parser = argparse.ArgumentParser(
+        description="Run Python samples concurrently",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog="""
+Examples:
+  python run_all_samples.py                                    # Run all samples
+  python run_all_samples.py --direct                           # Run all samples directly
+  python run_all_samples.py --subdir getting_started           # Run only getting_started samples
+  python run_all_samples.py --subdir getting_started/workflows # Run only workflow samples
+  python run_all_samples.py --subdir semantic-kernel-migration # Run only SK migration samples
+        """,
+    )
+
+    parser.add_argument(
+        "--direct", action="store_true", help="Run samples directly with python instead of using uv run"
+    )
+
+    parser.add_argument(
+        "--subdir", type=str, help="Run samples only in the specified subdirectory (relative to samples/)"
+    )
+
+    parser.add_argument(
+        "--max-workers", type=int, default=16, help="Maximum number of concurrent workers (default: 16)"
+    )
+
+    return parser.parse_args()
+
+
+def main() -> None:
+    """Main function to run all samples concurrently."""
+    args = parse_arguments()
+
+    # Get the samples directory (assuming this script is in the samples directory)
+    samples_dir = Path(__file__).parent
+    python_root = samples_dir.parent  # Go up to the python/ directory
+
+    print("Python samples runner")
+    print(f"Samples directory: {samples_dir}")
+
+    if args.direct:
+        print("Running samples directly (assuming environment is set up)")
+    else:
+        print(f"Using uv run from: {python_root}")
+
+    if args.subdir:
+        print(f"Filtering to subdirectory: {args.subdir}")
+
+    print("🚀 Running samples concurrently...")
+
+    # Find all Python sample files
+    sample_files = find_python_samples(samples_dir, args.subdir)
+
+    if not sample_files:
+        print("No Python sample files found!")
+        return
+
+    print(f"Found {len(sample_files)} Python sample files")
+
+    # Run samples concurrently
+    results: list[tuple[Path, bool, str, str]] = []
+
+    with ThreadPoolExecutor(max_workers=args.max_workers) as executor:
+        # Submit all tasks
+        future_to_sample = {
+            executor.submit(run_sample, sample_path, not args.direct, python_root): sample_path
+            for sample_path in sample_files
+        }
+
+        # Collect results as they complete
+        for future in as_completed(future_to_sample):
+            sample_path = future_to_sample[future]
+            try:
+                success, output, error_info = future.result()
+                results.append((sample_path, success, output, error_info))
+
+                # Print progress - show relative path from samples directory
+                relative_path = sample_path.relative_to(samples_dir)
+                if success:
+                    print(f"✅ {relative_path}")
+                else:
+                    print(f"❌ {relative_path} - {error_info.split(':', 1)[0]}")
+
+            except Exception as e:
+                error_info = f"Future exception: {str(e)}"
+                results.append((sample_path, False, "", error_info))
+                relative_path = sample_path.relative_to(samples_dir)
+                print(f"❌ {relative_path} - {error_info}")
+
+    # Sort results by original file order for consistent reporting
+    sample_to_index = {path: i for i, path in enumerate(sample_files)}
+    results.sort(key=lambda x: sample_to_index[x[0]])
+
+    successful_runs = sum(1 for _, success, _, _ in results if success)
+    failed_runs = len(results) - successful_runs
+
+    # Print detailed results
+    print(f"\n{'=' * 80}")
+    print("DETAILED RESULTS:")
+    print(f"{'=' * 80}")
+
+    for sample_path, success, output, error_info in results:
+        relative_path = sample_path.relative_to(samples_dir)
+        if success:
+            print(f"✅ {relative_path}")
+            if output and output != "No output":
+                print(f"   Output preview: {output[:100]}{'...' if len(output) > 100 else ''}")
+        else:
+            print(f"❌ {relative_path}")
+            print(f"   Error: {error_info}")
+
+    # Print summary
+    print(f"\n{'=' * 80}")
+    if failed_runs == 0:
+        print("🎉 ALL SAMPLES COMPLETED SUCCESSFULLY!")
+    else:
+        print(f"❌ {failed_runs} SAMPLE(S) FAILED!")
+    print(f"Successful runs: {successful_runs}")
+    print(f"Failed runs: {failed_runs}")
+
+    if args.subdir:
+        print(f"Subdirectory filter: {args.subdir}")
+
+    print(f"{'=' * 80}")
+
+    # Exit with error code if any samples failed
+    if failed_runs > 0:
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()

python/samples/getting_started/workflows/agents/workflow_as_agent_human_in_the_loop.py

@@ -5,7 +5,7 @@ import sys
 from collections.abc import Mapping
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Any, cast
+from typing import Any
 
 # Ensure local getting_started package can be imported when running as a script.
 _SAMPLES_ROOT = Path(__file__).resolve().parents[3]
@@ -150,23 +150,17 @@ async def main() -> None:
         else:
             raise TypeError("Unexpected argument type for human review function call.")
 
-        request_payload_obj: Any = request.data
-        if not isinstance(request_payload_obj, Mapping):
-            raise ValueError("Human review request payload must be a mapping.")
-        request_payload = cast(Mapping[str, Any], request_payload_obj)
+        request_payload: Any = request.data
+        if not isinstance(request_payload, HumanReviewRequest):
+            raise ValueError("Human review request payload must be a HumanReviewRequest.")
 
-        agent_request_obj = request_payload.get("agent_request")
-        if not isinstance(agent_request_obj, Mapping):
-            raise ValueError("Human review request must include agent_request mapping data.")
-        agent_request_data = cast(Mapping[str, Any], agent_request_obj)
-
-        request_id_obj = agent_request_data.get("request_id")
-        if not isinstance(request_id_obj, str):
-            raise ValueError("Human review request_id must be a string.")
-        request_id_value = request_id_obj
+        agent_request = request_payload.agent_request
+        if agent_request is None:
+            raise ValueError("Human review request must include agent_request.")
 
+        request_id = agent_request.request_id
         # Mock a human response approval for demonstration purposes.
-        human_response = ReviewResponse(request_id=request_id_value, feedback="Approved", approved=True)
+        human_response = ReviewResponse(request_id=request_id, feedback="Approved", approved=True)
 
         # Create the function call result object to send back to the agent.
         human_review_function_result = FunctionResultContent(

python/samples/getting_started/workflows/checkpoint/checkpoint_with_human_in_the_loop.py

@@ -23,6 +23,7 @@ from agent_framework import (
     WorkflowOutputEvent,
     WorkflowRunState,
     WorkflowStatusEvent,
+    get_checkpoint_summary,
     handler,
 )
 from agent_framework.azure import AzureOpenAIChatClient
@@ -246,14 +247,12 @@ def _render_checkpoint_summary(checkpoints: list["WorkflowCheckpoint"]) -> None:
     """Pretty-print saved checkpoints with the new framework summaries."""
 
     print("\nCheckpoint summary:")
-    for summary in [
-        RequestInfoExecutor.checkpoint_summary(cp) for cp in sorted(checkpoints, key=lambda c: c.timestamp)
-    ]:
+    for summary in [get_checkpoint_summary(cp) for cp in sorted(checkpoints, key=lambda c: c.timestamp)]:
         # Compose a single line per checkpoint so the user can scan the output
         # and pick the resume point that still has outstanding human work.
         line = (
             f"- {summary.checkpoint_id} | iter={summary.iteration_count} "
-            f"| targets={summary.targets} | states={summary.executor_states}"
+            f"| targets={summary.targets} | states={summary.executor_ids}"
         )
         if summary.status:
             line += f" | status={summary.status}"
@@ -312,7 +311,7 @@ def _prompt_for_responses(requests: list[tuple[str, HumanApprovalRequest]]) -> d
 def _maybe_pre_supply_responses(cp: "WorkflowCheckpoint") -> dict[str, str] | None:
     """Offer to collect responses before resuming a checkpoint."""
 
-    pending = RequestInfoExecutor.pending_requests_from_checkpoint(cp)
+    pending = get_checkpoint_summary(cp).pending_requests
     if not pending:
         return None
 
@@ -468,7 +467,7 @@ async def main() -> None:
         return
 
     chosen = sorted_cps[idx]
-    summary = RequestInfoExecutor.checkpoint_summary(chosen)
+    summary = get_checkpoint_summary(chosen)
     if summary.status == "completed":
         print("Selected checkpoint already reflects a completed workflow; nothing to resume.")
         return

python/samples/getting_started/workflows/checkpoint/checkpoint_with_resume.py

@@ -12,10 +12,10 @@ from agent_framework import (
     ChatMessage,
     Executor,
     FileCheckpointStorage,
-    RequestInfoExecutor,
     Role,
     WorkflowBuilder,
     WorkflowContext,
+    get_checkpoint_summary,
     handler,
 )
 from agent_framework.azure import AzureOpenAIChatClient
@@ -194,7 +194,7 @@ def _render_checkpoint_summary(checkpoints: list["WorkflowCheckpoint"]) -> None:
 
     print("\nCheckpoint summary:")
     for cp in sorted(checkpoints, key=lambda c: c.timestamp):
-        summary = RequestInfoExecutor.checkpoint_summary(cp)
+        summary = get_checkpoint_summary(cp)
         msg_count = sum(len(v) for v in cp.messages.values())
         state_keys = sorted(cp.executor_states.keys())
         orig = cp.shared_state.get("original_input")
@@ -241,7 +241,7 @@ async def main():
 
     print("\nAvailable checkpoints to resume from:")
     for idx, cp in enumerate(sorted_cps):
-        summary = RequestInfoExecutor.checkpoint_summary(cp)
+        summary = get_checkpoint_summary(cp)
         line = f"  [{idx}] id={summary.checkpoint_id} iter={summary.iteration_count}"
         if summary.status:
             line += f" status={summary.status}"