Python: [BREAKING] Python: Make executor ID required, improvements around handling rehydrating checkpoints (#832)

* Make executor ID required, improvements around handling rehydrating checkpoints.

* Duplicate executor validation added

* fix remaining issues

---------

Co-authored-by: Eric Zhu <ekzhu@users.noreply.github.com>
This commit is contained in:
Evan Mattson
2025-09-20 03:57:09 +09:00
committed by GitHub
Unverified
parent 7cd45e313b
commit aba094b5cf
33 changed files with 1967 additions and 275 deletions
@@ -91,6 +91,7 @@ from ._shared_state import SharedState
from ._telemetry import EdgeGroupDeliveryStatus, WorkflowTracer, workflow_tracer
from ._validation import (
EdgeDuplicationError,
ExecutorDuplicationError,
GraphConnectivityError,
HandlerOutputAnnotationError,
TypeCompatibilityError,
@@ -118,6 +119,7 @@ __all__ = [
"EdgeGroupDeliveryStatus",
"Executor",
"ExecutorCompletedEvent",
"ExecutorDuplicationError",
"ExecutorEvent",
"ExecutorFailedEvent",
"ExecutorInvokedEvent",
@@ -87,6 +87,7 @@ from ._shared_state import SharedState
from ._telemetry import EdgeGroupDeliveryStatus, WorkflowTracer, workflow_tracer
from ._validation import (
EdgeDuplicationError,
ExecutorDuplicationError,
GraphConnectivityError,
HandlerOutputAnnotationError,
TypeCompatibilityError,
@@ -114,6 +115,7 @@ __all__ = [
"EdgeGroupDeliveryStatus",
"Executor",
"ExecutorCompletedEvent",
"ExecutorDuplicationError",
"ExecutorEvent",
"ExecutorFailedEvent",
"ExecutorInvokedEvent",
@@ -145,7 +145,10 @@ class _CallbackAggregator(Executor):
"""
def __init__(self, callback: Callable[..., Any], id: str | None = None) -> None:
super().__init__(id)
derived_id = getattr(callback, "__name__", "") or ""
if not derived_id or derived_id == "<lambda>":
derived_id = f"{type(self).__name__}_unnamed"
super().__init__(id or derived_id)
self._callback = callback
self._param_count = len(inspect.signature(callback).parameters)
@@ -2,12 +2,15 @@
import contextlib
import functools
import importlib
import inspect
import logging
import uuid
from collections.abc import Awaitable, Callable
from dataclasses import dataclass, field
from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence
from dataclasses import asdict, dataclass, field, fields, is_dataclass
from textwrap import shorten
from types import UnionType
from typing import TYPE_CHECKING, Any, Generic, TypeVar, Union, get_args, get_origin, overload
from typing import TYPE_CHECKING, Any, ClassVar, Generic, TypeVar, Union, cast, get_args, get_origin, overload
if TYPE_CHECKING:
from ._workflow import Workflow
@@ -17,6 +20,7 @@ from pydantic import Field
from agent_framework import AgentProtocol, AgentRunResponse, AgentRunResponseUpdate, AgentThread, ChatMessage
from agent_framework._pydantic import AFBaseModel
from ._checkpoint import WorkflowCheckpoint
from ._events import (
AgentRunEvent,
AgentRunUpdateEvent,
@@ -25,35 +29,63 @@ from ._events import (
RequestInfoEvent,
_framework_event_origin, # pyright: ignore[reportPrivateUsage]
)
from ._runner_context import _decode_checkpoint_value
from ._typing_utils import is_instance_of
from ._workflow_context import WorkflowContext
logger = logging.getLogger(__name__)
# region Executor
@dataclass
class PendingRequestDetails:
"""Lightweight information about a pending request captured in a checkpoint."""
request_id: str
prompt: str | None = None
draft: str | None = None
iteration: int | None = None
source_executor_id: str | None = None
original_request: "RequestInfoMessage | dict[str, Any] | None" = None
@dataclass
class WorkflowCheckpointSummary:
"""Human-readable summary of a workflow checkpoint."""
checkpoint_id: str
iteration_count: int
targets: list[str]
executor_states: list[str]
status: str
draft_preview: str | None
pending_requests: list[PendingRequestDetails]
class Executor(AFBaseModel):
"""An executor is a component that processes messages in a workflow."""
# Provide a default so static analyzers (e.g., pyright) don't require passing `id`.
# Runtime still sets a concrete value in __init__.
id: str = Field(
default_factory=lambda: str(uuid.uuid4()),
...,
min_length=1,
description="Unique identifier for the executor",
)
type_: str = Field(default="", alias="type", description="The type of executor, corresponding to the class name")
def __init__(self, id: str | None = None, **kwargs: Any) -> None:
def __init__(self, id: str, **kwargs: Any) -> None:
"""Initialize the executor with a unique identifier.
Args:
id: A unique identifier for the executor. If None, a new ID will be generated
following the format <class_name>/<uuid>.
id: A unique identifier for the executor.
kwargs: Additional keyword arguments. Unused in this implementation.
"""
executor_id = f"{self.__class__.__name__}/{uuid.uuid4()}" if id is None else id
if not id:
raise ValueError("Executor ID must be a non-empty string.")
kwargs.update({"id": executor_id})
kwargs.update({"id": id})
if "type" not in kwargs and "type_" not in kwargs:
kwargs["type_"] = self.__class__.__name__
@@ -677,17 +709,15 @@ class RequestInfoExecutor(Executor):
a response is provided externally, it emits the response as a message.
"""
def __init__(self, id: str | None = None):
"""Initialize the RequestInfoExecutor with an optional custom ID.
_PENDING_SHARED_STATE_KEY: ClassVar[str] = "_af_pending_request_info"
def __init__(self, id: str):
"""Initialize the RequestInfoExecutor with a unique ID.
Args:
id: Optional custom ID for this RequestInfoExecutor. If not provided,
a unique ID will be generated.
id: Unique ID for this RequestInfoExecutor.
"""
import uuid
executor_id = id or f"request_info_{uuid.uuid4().hex[:8]}"
super().__init__(id=executor_id)
super().__init__(id=id)
self._request_events: dict[str, RequestInfoEvent] = {}
self._sub_workflow_contexts: dict[str, dict[str, str]] = {}
@@ -703,6 +733,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 ctx.add_event(event)
@handler
@@ -748,10 +779,13 @@ class RequestInfoExecutor(Executor):
response_data: The data returned in the response.
ctx: The workflow context for sending the response.
"""
if request_id not in self._request_events:
event = self._request_events.get(request_id)
if event is None:
event = await self._rehydrate_request_event(request_id, ctx)
if event is None:
raise ValueError(f"No request found with ID: {request_id}")
event = self._request_events.pop(request_id)
self._request_events.pop(request_id, None)
# Check if this was a forwarded sub-workflow request
if request_id in self._sub_workflow_contexts:
@@ -779,6 +813,472 @@ class RequestInfoExecutor(Executor):
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)
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 not in pending:
return
pending.pop(request_id, None)
await self._persist_pending_request_state(pending, ctx)
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 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) or existing is None:
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)
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}")
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))
if isinstance(data, dict):
return cast(dict[str, Any], data)
return None
model_dump = getattr(request, "model_dump", None)
if callable(model_dump):
try:
dump = self._make_json_safe(model_dump(mode="json"))
except TypeError:
dump = self._make_json_safe(model_dump())
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():
safe_dict[str(key)] = self._make_json_safe(val)
return safe_dict
if isinstance(value, Sequence) and not isinstance(value, (str, bytes, bytearray)):
return [self._make_json_safe(item) for item in value]
return repr(value)
async def has_pending_request(self, request_id: str, ctx: WorkflowContext[Any]) -> bool:
if request_id in self._request_events or request_id in self._sub_workflow_contexts:
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:
return None
request = self._construct_request_from_snapshot(snapshot)
if request is None:
return None
event = RequestInfoEvent(
request_id=request_id,
source_executor_id=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 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():
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():
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 {}
request_cls: type[RequestInfoMessage] = RequestInfoMessage
request_type_str = snapshot.get("request_type")
if isinstance(request_type_str, str) and ":" in request_type_str:
module_name, class_name = request_type_str.split(":", 1)
try:
module = importlib.import_module(module_name)
candidate = getattr(module, class_name)
if isinstance(candidate, type) and issubclass(candidate, RequestInfoMessage):
request_cls = candidate
except Exception as exc:
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)
if request is None and request_cls is not RequestInfoMessage:
request = self._instantiate_request(RequestInfoMessage, details)
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())}"
)
return None
for key, value in details.items():
if key == "request_id":
continue
try:
setattr(request, key, value)
except Exception as exc:
logger.debug(
f"RequestInfoExecutor {self.id} could not set attribute {key} on {type(request).__name__}: {exc}"
)
continue
snapshot_request_id = snapshot.get("request_id")
if isinstance(snapshot_request_id, str) and snapshot_request_id:
try:
request.request_id = snapshot_request_id
except Exception as exc:
logger.debug(
f"RequestInfoExecutor {self.id} could not apply snapshot "
f"request_id to {type(request).__name__}: {exc}"
)
return request
def _instantiate_request(
self,
request_cls: type[RequestInfoMessage],
details: dict[str, Any],
) -> RequestInfoMessage | None:
try:
model_validate = getattr(request_cls, "model_validate", None)
if callable(model_validate):
return cast(RequestInfoMessage, model_validate(details))
except (TypeError, ValueError) as exc:
logger.debug(
f"RequestInfoExecutor {self.id} validation failed for {request_cls.__name__} via model_validate: {exc}"
)
except Exception as exc:
logger.warning(
f"RequestInfoExecutor {self.id} encountered unexpected error during "
f"{request_cls.__name__}.model_validate: {exc}"
)
if is_dataclass(request_cls):
try:
field_names = {f.name for f in fields(request_cls)}
ctor_kwargs = {name: details[name] for name in field_names if name in details}
return request_cls(**ctor_kwargs) # type: ignore[call-arg]
except (TypeError, ValueError) as exc:
logger.debug(
f"RequestInfoExecutor {self.id} could not instantiate dataclass "
f"{request_cls.__name__} with snapshot data: {exc}"
)
except Exception as exc:
logger.warning(
f"RequestInfoExecutor {self.id} encountered unexpected error "
f"constructing dataclass {request_cls.__name__}: {exc}"
)
try:
instance = request_cls() # type: ignore[call-arg]
except Exception as exc:
logger.warning(
f"RequestInfoExecutor {self.id} could not instantiate {request_cls.__name__} without arguments: {exc}"
)
return None
for key, value in details.items():
if key == "request_id":
continue
try:
setattr(instance, key, value)
except Exception as exc:
logger.debug(
f"RequestInfoExecutor {self.id} could not set attribute {key} on "
f"{request_cls.__name__} during instantiation: {exc}"
)
continue
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}
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():
RequestInfoExecutor._merge_snapshot(pending, str(request_id), snapshot)
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():
RequestInfoExecutor._merge_snapshot(pending, str(request_id), snapshot)
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)
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
)
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"),
draft=snapshot.get("draft"),
iteration=snapshot.get("iteration"),
source_executor_id=snapshot.get("source_executor_id"),
)
extra = snapshot.get("details")
if isinstance(extra, Mapping):
RequestInfoExecutor._apply_update(
details,
prompt=extra.get("prompt"),
draft=extra.get("draft"),
iteration=extra.get("iteration"),
)
@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")
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"),
draft=RequestInfoExecutor._get_field(payload.original_request, "draft"),
iteration=RequestInfoExecutor._get_field(payload.original_request, "iteration"),
source_executor_id=raw_message.get("source_id"),
original_request=payload.original_request,
)
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)
return getattr(obj, key, None)
@staticmethod
def _coerce_int(value: Any) -> int | None:
try:
return int(value)
except (TypeError, ValueError):
return None
# endregion: Request Info Executor
@@ -840,7 +1340,11 @@ class AgentExecutor(Executor):
exec_id = id
else:
agent_name = agent.name
exec_id = str(agent_name) if agent_name else f"executor_{uuid.uuid4()}"
if agent_name:
exec_id = str(agent_name)
else:
logger.warning("Agent has no name, using fallback ID 'executor_unnamed'")
exec_id = "executor_unnamed"
super().__init__(exec_id)
self._agent = agent
self._agent_thread = agent_thread or self._agent.get_new_thread()
@@ -952,12 +1456,12 @@ class WorkflowExecutor(Executor):
workflow: "Workflow" = Field(description="The workflow to execute as a sub-workflow")
def __init__(self, workflow: "Workflow", id: str | None = None, **kwargs: Any):
def __init__(self, workflow: "Workflow", id: str, **kwargs: Any):
"""Initialize the WorkflowExecutor.
Args:
workflow: The workflow to execute as a sub-workflow.
id: Optional unique identifier for this executor.
id: Unique identifier for this executor.
**kwargs: Additional keyword arguments passed to the parent constructor.
"""
kwargs.update({"workflow": workflow})
@@ -1635,7 +1635,7 @@ class MagenticBuilder:
if self._enable_plan_review:
from ._executor import RequestInfoExecutor
request_info = RequestInfoExecutor()
request_info = RequestInfoExecutor(id="request_info")
workflow_builder = (
workflow_builder
# Only route plan review asks to request_info
@@ -13,7 +13,14 @@ from ._edge import EdgeGroup
from ._edge_runner import EdgeRunner, create_edge_runner
from ._events import WorkflowCompletedEvent, WorkflowEvent, _framework_event_origin
from ._executor import Executor
from ._runner_context import Message, RunnerContext
from ._runner_context import (
_DATACLASS_MARKER,
_PYDANTIC_MARKER,
CheckpointState,
Message,
RunnerContext,
_decode_checkpoint_value,
)
from ._shared_state import SharedState
from ._typing_utils import is_instance_of
from ._workflow_context import WorkflowContext
@@ -53,6 +60,7 @@ class Runner:
self._workflow_id = workflow_id
self._running = False
self._resumed_from_checkpoint = False # Track whether we resumed
self.graph_signature_hash: str | None = None
# Set workflow ID in context if provided
if workflow_id:
@@ -244,6 +252,19 @@ class Runner:
"""Inner loop to deliver a single message through an edge runner."""
return await edge_runner.send_message(message, self._shared_state, self._ctx)
def _normalize_message_payload(message: Message) -> None:
data = message.data
if not isinstance(data, dict):
return
if _PYDANTIC_MARKER not in data and _DATACLASS_MARKER not in data:
return
try:
decoded = _decode_checkpoint_value(data)
except Exception as exc: # pragma: no cover - defensive
logger.debug("Failed to decode checkpoint payload during delivery: %s", exc)
return
message.data = decoded
# Handle SubWorkflowRequestInfo messages specially
await _deliver_sub_workflow_requests(messages)
@@ -266,6 +287,7 @@ class Runner:
associated_edge_runners = self._edge_runner_map.get(source_executor_id, [])
for message in non_sub_workflow_messages:
_normalize_message_payload(message)
# Deliver a message through all edge runners associated with the source executor concurrently.
tasks = [_deliver_message_inner(edge_runner, message) for edge_runner in associated_edge_runners]
if not tasks:
@@ -332,6 +354,8 @@ class Runner:
"superstep": self._iteration,
"checkpoint_type": checkpoint_category,
}
if self.graph_signature_hash:
metadata["graph_signature"] = self.graph_signature_hash
checkpoint_id = await self._ctx.create_checkpoint(metadata=metadata)
logger.info(f"Created {checkpoint_type} checkpoint: {checkpoint_id}")
return checkpoint_id
@@ -403,14 +427,45 @@ class Runner:
return False
try:
success = await self._ctx.restore_from_checkpoint(checkpoint_id)
if not success:
checkpoint = await self._ctx.load_checkpoint(checkpoint_id)
if not checkpoint:
logger.error(f"Checkpoint {checkpoint_id} not found")
return False
graph_hash = getattr(self, "graph_signature_hash", None)
checkpoint_hash = (checkpoint.metadata or {}).get("graph_signature")
if graph_hash and checkpoint_hash and graph_hash != checkpoint_hash:
raise ValueError(
"Workflow graph has changed since the checkpoint was created. "
"Please rebuild the original workflow before resuming."
)
if graph_hash and not checkpoint_hash:
logger.warning(
"Checkpoint %s does not include graph signature metadata; skipping topology validation.",
checkpoint_id,
)
state: CheckpointState = {
"messages": checkpoint.messages,
"shared_state": checkpoint.shared_state,
"executor_states": checkpoint.executor_states,
"iteration_count": checkpoint.iteration_count,
"max_iterations": checkpoint.max_iterations,
}
await self._ctx.set_checkpoint_state(state)
if checkpoint.workflow_id:
self._ctx.set_workflow_id(checkpoint.workflow_id)
self._workflow_id = checkpoint.workflow_id
await self._restore_shared_state_from_context()
self.mark_resumed() # mark resumed; iteration/max already restored from context
self.mark_resumed(
iteration=checkpoint.iteration_count,
max_iterations=checkpoint.max_iterations,
)
logger.info(f"Successfully restored workflow from checkpoint: {checkpoint_id}")
return True
except ValueError:
raise
except Exception as e:
logger.error(f"Failed to restore from checkpoint {checkpoint_id}: {e}")
return False
@@ -3,6 +3,7 @@
import asyncio
import importlib
import logging
import sys
import uuid
from collections import defaultdict
from dataclasses import dataclass, fields, is_dataclass
@@ -60,6 +61,39 @@ _MAX_ENCODE_DEPTH = 100
_CYCLE_SENTINEL = "<cycle>"
def _instantiate_checkpoint_dataclass(cls: type[Any], payload: Any) -> Any | None:
if not isinstance(cls, type):
logger.debug(f"Checkpoint decoder received non-type dataclass reference: {cls!r}")
return None
if isinstance(payload, dict):
try:
return cls(**payload) # type: ignore[arg-type]
except TypeError as exc:
logger.debug(f"Checkpoint decoder could not call {cls.__name__}(**payload): {exc}")
except Exception as exc:
logger.warning(f"Checkpoint decoder encountered unexpected error calling {cls.__name__}(**payload): {exc}")
try:
instance = object.__new__(cls)
except Exception as exc:
logger.debug(f"Checkpoint decoder could not allocate {cls.__name__} without __init__: {exc}")
return None
for key, val in payload.items():
try:
setattr(instance, key, val)
except Exception as exc:
logger.debug(f"Checkpoint decoder could not set attribute {key} on {cls.__name__}: {exc}")
return instance
try:
return cls(payload) # type: ignore[call-arg]
except TypeError as exc:
logger.debug(f"Checkpoint decoder could not call {cls.__name__}({payload!r}): {exc}")
except Exception as exc:
logger.warning(f"Checkpoint decoder encountered unexpected error calling {cls.__name__}({payload!r}): {exc}")
return None
def _is_pydantic_model(obj: object) -> bool:
"""Best-effort check for Pydantic models (e.g., AFBaseModel).
@@ -99,7 +133,7 @@ def _encode_checkpoint_value(value: Any) -> Any:
"value": v.model_dump(mode="json"),
}
except Exception as exc: # best-effort fallback
logger.debug("Pydantic model_dump failed for %s: %s", cls, exc)
logger.debug(f"Pydantic model_dump failed for {cls}: {exc}")
return str(v)
# Dataclasses (instances only)
@@ -178,36 +212,32 @@ def _decode_checkpoint_value(value: Any) -> Any:
if isinstance(type_key, str):
try:
module_name, class_name = type_key.split(":", 1)
module = importlib.import_module(module_name)
module = sys.modules.get(module_name)
if module is None:
module = importlib.import_module(module_name)
cls: Any = getattr(module, class_name)
if hasattr(cls, "model_validate"):
return cls.model_validate(raw)
except Exception as exc:
logger.debug(
"Failed to decode pydantic model %s: %s; returning raw value",
type_key,
exc,
)
logger.debug(f"Failed to decode pydantic model {type_key}: {exc}; returning raw value")
# Dataclass marker handling
if _DATACLASS_MARKER in value_dict and "value" in value_dict:
type_key_dc: str | None = value_dict.get(_DATACLASS_MARKER) # type: ignore[assignment]
raw_dc: Any = value_dict.get("value")
decoded_raw = _decode_checkpoint_value(raw_dc)
if isinstance(type_key_dc, str):
try:
module_name, class_name = type_key_dc.split(":", 1)
module = importlib.import_module(module_name)
module = sys.modules.get(module_name)
if module is None:
module = importlib.import_module(module_name)
cls_dc: Any = getattr(module, class_name)
decoded_raw = _decode_checkpoint_value(raw_dc)
if isinstance(decoded_raw, dict):
return cls_dc(**decoded_raw)
constructed = _instantiate_checkpoint_dataclass(cls_dc, decoded_raw)
if constructed is not None:
return constructed
except Exception as exc:
logger.debug(
"Failed to decode dataclass %s: %s; returning raw value",
type_key_dc,
exc,
)
# Fallback to decoded raw value
return _decode_checkpoint_value(raw_dc)
logger.debug(f"Failed to decode dataclass {type_key_dc}: {exc}; returning raw value")
return decoded_raw
# Regular dict: decode recursively
decoded: dict[str, Any] = {}
@@ -338,6 +368,10 @@ class RunnerContext(Protocol):
"""
...
async def load_checkpoint(self, checkpoint_id: str) -> WorkflowCheckpoint | None:
"""Load a checkpoint without mutating the current context state."""
...
async def get_checkpoint_state(self) -> CheckpointState:
"""Get the current state of the context suitable for checkpointing."""
...
@@ -409,7 +443,7 @@ class InProcRunnerContext:
return
except Exception as exc: # pragma: no cover - defensive logging path
# Best-effort filtering only; never block event delivery on filtering errors
logger.debug("Error while filtering event %r: %s", event, exc, exc_info=True)
logger.debug(f"Error while filtering event {event!r}: {exc}", exc_info=True)
await self._event_queue.put(event)
@@ -497,6 +531,11 @@ class InProcRunnerContext:
logger.info(f"Restored state from checkpoint {checkpoint_id}'")
return True
async def load_checkpoint(self, checkpoint_id: str) -> WorkflowCheckpoint | None:
if not self._checkpoint_storage:
raise ValueError("Checkpoint storage not configured")
return await self._checkpoint_storage.load_checkpoint(checkpoint_id)
async def get_checkpoint_state(self) -> CheckpointState:
serializable_messages: dict[str, list[dict[str, Any]]] = {}
for source_id, message_list in self._messages.items():
@@ -1,7 +1,58 @@
# Copyright (c) Microsoft. All rights reserved.
import logging
from dataclasses import fields, is_dataclass
from typing import Any, Union, get_args, get_origin
logger = logging.getLogger(__name__)
def _coerce_to_type(value: Any, target_type: type) -> Any | None:
"""Best-effort conversion of value into target_type."""
if isinstance(value, target_type):
return value
# Convert dataclass instances or objects with __dict__ into dict first
if not isinstance(value, dict):
if is_dataclass(value):
value = {f.name: getattr(value, f.name) for f in fields(value)}
else:
value_dict = getattr(value, "__dict__", None)
if isinstance(value_dict, dict):
value = dict(value_dict)
if isinstance(value, dict):
ctor_kwargs: dict[str, Any] = dict(value)
if is_dataclass(target_type):
field_names = {f.name for f in fields(target_type)}
ctor_kwargs = {k: v for k, v in value.items() if k in field_names}
try:
return target_type(**ctor_kwargs) # type: ignore[arg-type]
except TypeError as exc:
logger.debug(f"_coerce_to_type could not call {target_type.__name__}(**..): {exc}")
except Exception as exc: # pragma: no cover - unexpected constructor failure
logger.warning(
f"_coerce_to_type encountered unexpected error calling {target_type.__name__} constructor: {exc}"
)
try:
instance: Any = object.__new__(target_type)
except Exception as exc: # pragma: no cover - pathological type
logger.debug(f"_coerce_to_type could not allocate {target_type.__name__} without __init__: {exc}")
return None
for key, val in value.items():
try:
setattr(instance, key, val)
except Exception as exc:
logger.debug(
f"_coerce_to_type could not set {target_type.__name__}.{key} during fallback assignment: {exc}"
)
continue
return instance
return None
def is_instance_of(data: Any, target_type: type) -> bool:
"""Check if the data is an instance of the target type.
@@ -63,7 +114,10 @@ def is_instance_of(data: Any, target_type: type) -> bool:
and data.original_request is not None
and not is_instance_of(data.original_request, request_type)
):
return False
coerced = _coerce_to_type(data.original_request, request_type)
if coerced is None:
return False
data.original_request = coerced
if hasattr(data, "data") and data.data is not None and not is_instance_of(data.data, response_type):
return False
return True
@@ -34,6 +34,7 @@ class ValidationTypeEnum(Enum):
"""Enumeration of workflow validation types."""
EDGE_DUPLICATION = "EDGE_DUPLICATION"
EXECUTOR_DUPLICATION = "EXECUTOR_DUPLICATION"
TYPE_COMPATIBILITY = "TYPE_COMPATIBILITY"
GRAPH_CONNECTIVITY = "GRAPH_CONNECTIVITY"
HANDLER_OUTPUT_ANNOTATION = "HANDLER_OUTPUT_ANNOTATION"
@@ -63,6 +64,20 @@ class EdgeDuplicationError(WorkflowValidationError):
self.edge_id = edge_id
class ExecutorDuplicationError(WorkflowValidationError):
"""Exception raised when duplicate executor identifiers are detected."""
def __init__(self, executor_id: str):
super().__init__(
message=(
f"Duplicate executor id detected: '{executor_id}'. Executor ids must be globally unique within a "
"workflow."
),
validation_type=ValidationTypeEnum.EXECUTOR_DUPLICATION,
)
self.executor_id = executor_id
class TypeCompatibilityError(WorkflowValidationError):
"""Exception raised when type incompatibility is detected between connected executors."""
@@ -133,10 +148,17 @@ class WorkflowGraphValidator:
def __init__(self) -> None:
self._edges: list[Edge] = []
self._executors: dict[str, Executor] = {}
self._duplicate_executor_ids: set[str] = set()
self._start_executor_ref: Executor | str | None = None
# region Core Validation Methods
def validate_workflow(
self, edge_groups: Sequence[EdgeGroup], executors: dict[str, Executor], start_executor: Executor | str
self,
edge_groups: Sequence[EdgeGroup],
executors: dict[str, Executor],
start_executor: Executor | str,
*,
duplicate_executor_ids: Sequence[str] | None = None,
) -> None:
"""Validate the entire workflow graph.
@@ -144,6 +166,7 @@ class WorkflowGraphValidator:
edge_groups: list of edge groups in the workflow
executors: Map of executor IDs to executor instances
start_executor: The starting executor (can be instance or ID)
duplicate_executor_ids: Optional list of known duplicate executor IDs to pre-populate
Raises:
WorkflowValidationError: If any validation fails
@@ -151,6 +174,8 @@ class WorkflowGraphValidator:
self._executors = executors
self._edges = [edge for group in edge_groups for edge in group.edges]
self._edge_groups = edge_groups
self._duplicate_executor_ids = set(duplicate_executor_ids or [])
self._start_executor_ref = start_executor
# If only the start executor exists, add it to the executor map
# Handle the special case where the workflow consists of only a single executor and no edges.
@@ -185,6 +210,7 @@ class WorkflowGraphValidator:
)
# Run all checks
self._validate_executor_id_uniqueness(start_executor_id)
self._validate_edge_duplication()
self._validate_handler_output_annotations()
self._validate_type_compatibility()
@@ -353,6 +379,26 @@ class WorkflowGraphValidator:
# endregion
def _validate_executor_id_uniqueness(self, start_executor_id: str) -> None:
"""Ensure executor identifiers are unique throughout the workflow graph."""
duplicates: set[str] = set(self._duplicate_executor_ids)
id_counts: defaultdict[str, int] = defaultdict(int)
for key, executor in self._executors.items():
id_counts[executor.id] += 1
if key != executor.id:
duplicates.add(executor.id)
duplicates.update({executor_id for executor_id, count in id_counts.items() if count > 1})
if isinstance(self._start_executor_ref, Executor):
mapped = self._executors.get(start_executor_id)
if mapped is not None and mapped is not self._start_executor_ref:
duplicates.add(start_executor_id)
if duplicates:
raise ExecutorDuplicationError(sorted(duplicates)[0])
# region Edge and Type Validation
def _validate_edge_duplication(self) -> None:
"""Validate that there are no duplicate edges in the workflow.
@@ -793,7 +839,11 @@ class WorkflowGraphValidator:
def validate_workflow_graph(
edge_groups: Sequence[EdgeGroup], executors: dict[str, Executor], start_executor: Executor | str
edge_groups: Sequence[EdgeGroup],
executors: dict[str, Executor],
start_executor: Executor | str,
*,
duplicate_executor_ids: Sequence[str] | None = None,
) -> None:
"""Convenience function to validate a workflow graph.
@@ -801,9 +851,15 @@ def validate_workflow_graph(
edge_groups: list of edge groups in the workflow
executors: Map of executor IDs to executor instances
start_executor: The starting executor (can be instance or ID)
duplicate_executor_ids: Optional list of known duplicate executor IDs to pre-populate
Raises:
WorkflowValidationError: If any validation fails
"""
validator = WorkflowGraphValidator()
validator.validate_workflow(edge_groups, executors, start_executor)
validator.validate_workflow(
edge_groups,
executors,
start_executor,
duplicate_executor_ids=duplicate_executor_ids,
)
@@ -1,6 +1,8 @@
# Copyright (c) Microsoft. All rights reserved.
import asyncio
import hashlib
import json
import logging
import sys
import uuid
@@ -177,6 +179,12 @@ class Workflow(AFBaseModel):
workflow_id=id,
)
# Capture a canonical fingerprint of the workflow graph so checkpoints
# can assert they are resumed with an equivalent topology.
self._graph_signature = self._compute_graph_signature()
self._graph_signature_hash = self._hash_graph_signature(self._graph_signature)
self._runner.graph_signature_hash = self._graph_signature_hash
def model_dump(self, **kwargs: Any) -> dict[str, Any]:
"""Custom serialization that properly handles WorkflowExecutor nested workflows."""
data = super().model_dump(**kwargs)
@@ -377,17 +385,26 @@ class Workflow(AFBaseModel):
request_info_executor = self._find_request_info_executor()
if request_info_executor:
for request_id, response_data in responses.items():
ctx: WorkflowContext[Any] = WorkflowContext(
request_info_executor.id,
[self.__class__.__name__],
self._shared_state,
self._runner.context,
trace_contexts=None, # No parent trace context for new workflow span
source_span_ids=None, # No source span for response handling
)
if not await request_info_executor.has_pending_request(request_id, ctx):
logger.debug(
f"Skipping pre-supplied response for request {request_id}; no pending request found "
f"after checkpoint restoration."
)
continue
await request_info_executor.handle_response(
response_data,
request_id,
WorkflowContext(
request_info_executor.id,
[self.__class__.__name__],
self._shared_state,
self._runner.context,
trace_contexts=None, # No parent trace context for new workflow span
source_span_ids=None, # No source span for response handling
),
ctx,
)
async for event in self._run_workflow_with_tracing(
@@ -590,6 +607,19 @@ class Workflow(AFBaseModel):
if not checkpoint:
return False
graph_hash = getattr(self._runner, "graph_signature_hash", None)
checkpoint_hash = (checkpoint.metadata or {}).get("graph_signature")
if graph_hash and checkpoint_hash and graph_hash != checkpoint_hash:
raise ValueError(
"Workflow graph has changed since the checkpoint was created. "
"Please rebuild the original workflow before resuming."
)
if graph_hash and not checkpoint_hash:
logger.warning(
f"Checkpoint {checkpoint_id} does not include graph signature metadata; "
f"skipping topology validation."
)
temp_context = InProcRunnerContext(checkpoint_storage)
state: CheckpointState = {
"messages": checkpoint.messages,
@@ -608,10 +638,9 @@ class Workflow(AFBaseModel):
return True
except ValueError:
raise
except Exception as e:
import logging
logger = logging.getLogger(__name__)
logger.error(f"Failed to restore from external checkpoint {checkpoint_id}: {e}")
return False
@@ -632,7 +661,7 @@ class Workflow(AFBaseModel):
self._shared_state._state.clear() # type: ignore[attr-defined]
self._shared_state._state.update(shared_state_data) # type: ignore[attr-defined]
except Exception as exc: # pragma: no cover
logger.debug("Failed to restore shared_state during external restore: %s", exc)
logger.debug(f"Failed to restore shared_state during external restore: {exc}")
# Restore executor states into the context so ctx.get_state() calls after resume succeed
try:
@@ -641,9 +670,9 @@ class Workflow(AFBaseModel):
try:
await self._runner.context.set_state(exec_id, state)
except Exception as exc: # pragma: no cover - ignore per-executor failures
logger.debug("Failed to restore executor state for %s during external restore: %s", exec_id, exc)
logger.debug(f"Failed to restore executor state for {exec_id} during external restore: {exc}")
except Exception as exc: # pragma: no cover
logger.debug("Failed to iterate executor_states during external restore: %s", exc)
logger.debug(f"Failed to iterate executor_states during external restore: {exc}")
# Transfer pending messages into the context for delivery in the next superstep
messages_data = restored_state["messages"]
@@ -671,6 +700,71 @@ class Workflow(AFBaseModel):
)
)
# Graph signature helpers
def _compute_graph_signature(self) -> dict[str, Any]:
"""Build a canonical fingerprint of the workflow graph topology for checkpoint validation.
This creates a minimal, stable representation that captures only the structural
elements of the workflow (executor types, edge relationships, topology) while
ignoring data/state changes. Used to verify that a workflow's structure hasn't
changed when resuming from checkpoints.
"""
executors_signature = {
executor_id: f"{executor.__class__.__module__}.{executor.__class__.__name__}"
for executor_id, executor in self.executors.items()
}
edge_groups_signature: list[dict[str, Any]] = []
for group in self.edge_groups:
edges = [
{
"source": edge.source_id,
"target": edge.target_id,
"condition": getattr(edge, "condition_name", None),
}
for edge in group.edges
]
edges.sort(key=lambda e: (e["source"], e["target"], e["condition"] or ""))
group_info: dict[str, Any] = {
"group_type": group.__class__.__name__,
"sources": sorted(group.source_executor_ids),
"targets": sorted(group.target_executor_ids),
"edges": edges,
}
if isinstance(group, FanOutEdgeGroup):
group_info["selection_func"] = group.selection_func_name
edge_groups_signature.append(group_info)
edge_groups_signature.sort(
key=lambda info: (
info["group_type"],
tuple(info["sources"]),
tuple(info["targets"]),
json.dumps(info["edges"], sort_keys=True),
json.dumps(info.get("selection_func")),
)
)
return {
"start_executor": self.start_executor_id,
"executors": executors_signature,
"edge_groups": edge_groups_signature,
"max_iterations": self.max_iterations,
}
@staticmethod
def _hash_graph_signature(signature: dict[str, Any]) -> str:
canonical = json.dumps(signature, sort_keys=True, separators=(",", ":"))
return hashlib.sha256(canonical.encode("utf-8")).hexdigest()
@property
def graph_signature_hash(self) -> str:
return self._graph_signature_hash
def as_agent(self, name: str | None = None) -> WorkflowAgent:
"""Create a WorkflowAgent that wraps this workflow.
@@ -699,6 +793,7 @@ class WorkflowBuilder:
"""Initialize the WorkflowBuilder with an empty list of edges and no starting executor."""
self._edge_groups: list[EdgeGroup] = []
self._executors: dict[str, Executor] = {}
self._duplicate_executor_ids: set[str] = set()
self._start_executor: Executor | str | None = None
self._checkpoint_storage: CheckpointStorage | None = None
self._max_iterations: int = max_iterations
@@ -712,7 +807,11 @@ class WorkflowBuilder:
def _add_executor(self, executor: Executor) -> str:
"""Add an executor to the map and return its ID."""
self._executors[executor.id] = executor
existing = self._executors.get(executor.id)
if existing is not None and existing is not executor:
self._duplicate_executor_ids.add(executor.id)
else:
self._executors[executor.id] = executor
return executor.id
def _maybe_wrap_agent(self, candidate: Executor | AgentProtocol) -> Executor:
@@ -739,7 +838,10 @@ class WorkflowBuilder:
if name:
proposed_id = str(name)
if proposed_id in self._executors:
proposed_id = f"{proposed_id}-{uuid.uuid4().hex[:8]}"
raise ValueError(
f"Duplicate executor ID '{proposed_id}' from agent name. "
"Agent names must be unique within a workflow."
)
wrapper = AgentExecutor(candidate, id=proposed_id, streaming=True)
self._agent_wrappers[id(candidate)] = wrapper
return wrapper
@@ -934,8 +1036,9 @@ class WorkflowBuilder:
self._start_executor = wrapped
# Ensure the start executor is present in the executor map so validation succeeds
# even if no edges are added yet, or before edges wrap the same agent again.
if wrapped.id not in self._executors:
self._executors[wrapped.id] = wrapped
existing = self._executors.get(wrapped.id)
if existing is not wrapped:
self._add_executor(wrapped)
return self
def set_max_iterations(self, max_iterations: int) -> Self:
@@ -986,7 +1089,12 @@ class WorkflowBuilder:
)
# Perform validation before creating the workflow
validate_workflow_graph(self._edge_groups, self._executors, self._start_executor)
validate_workflow_graph(
self._edge_groups,
self._executors,
self._start_executor,
duplicate_executor_ids=tuple(self._duplicate_executor_ids),
)
# Add validation completed event
workflow_tracer.add_build_event("build.validation_completed")