microsoft / microsoft/agent-framework

Python: [Feature]: Make workflows stateless with caller-owned checkpoints and msgspec snapshots

Open
#8,522 0 comments 0 reactions 0 assignees View on GitHub
python vnext
Dominant language
Python
Stars
13.6k
Forks
2.3k
Avg merge
2d 45m
Merged PRs (30d)
358

Description

## Description

### Summary

For v2, align workflow state ownership with the agent/session model without introducing a separate `WorkflowSession` abstraction:

- `agent.create_session()` creates caller-owned state.
- `agent.run(..., session=session)` uses and updates that state.
- `workflow.create_checkpoint()` should create caller-owned workflow state.
- `workflow.run(..., checkpoint=checkpoint)` should use and update that state.

A `Checkpoint` represents one logical workflow execution and its lifecycle. It contains snapshots, and each snapshot contains the complete state required to resume from that superstep.

The developer owns the checkpoint and its lifecycle. The `Workflow` becomes a stateless, reusable definition.

This is a breaking change and should only be implemented for v2, hence the `vnext` label.

### Current behavior

Today, mutable execution state is retained by `Workflow`, its runner, runner context, executors, and edge runners. Checkpoint storage can be attached to `WorkflowBuilder` or supplied to `Workflow.run()`:

```python
storage = FileCheckpointStorage("./checkpoints")

builder = (
WorkflowBuilder(
start_executor=start,
checkpoint_storage=storage,
)
.add_edge(start, worker)
.add_edge(worker, worker)
)

workflow = builder.build()
await workflow.run(message=10)

latest = await storage.get_latest(workflow_name=workflow.name)
if latest is None:
raise RuntimeError("No checkpoint was created")

resumed_workflow = builder.build()
await resumed_workflow.run(checkpoint_id=latest.checkpoint_id)
```

This has several consequences:

- A `Workflow` instance owns mutable execution state.
- One workflow instance cannot safely execute independent runs concurrently.
- Checkpoint lifecycle is split between the workflow, storage, and application.
- Applications must query storage to discover the latest checkpoint.
- `get_latest(workflow_name=...)` is ambiguous when multiple executions use the same workflow definition.
- Resume is expressed through `checkpoint_id` and `checkpoint_storage` arguments rather than a caller-owned state object.
- `FileCheckpointStorage` uses JSON containing pickle/base64 payloads, while agent sessions use typed msgspec serialization.

## Proposed developer experience

The workflow creates a checkpoint, and the caller passes that checkpoint to related runs:

```python
storage = FileCheckpointStorage("./checkpoints")

workflow = (
WorkflowBuilder(start_executor=start)
.add_edge(start, worker)
.add_edge(worker, worker)
.build()
)

checkpoint = workflow.create_checkpoint(
checkpoint_id="factor-run-123",
checkpoint_storage=storage,
)

await workflow.run(
message=10,
checkpoint=checkpoint,
)

# Continue the same logical execution.
await workflow.run(
responses={"approval-request": True},
checkpoint=checkpoint,
)

latest_snapshot = checkpoint.head_snapshot
```

A one-shot invocation without a checkpoint remains stateless:

```python
result = await workflow.run(message=10)
```

To resume after a process restart, load the checkpoint and pass it directly to the workflow:

```python
storage = FileCheckpointStorage("./checkpoints")
checkpoint = await storage.load("factor-run-123")

workflow = build_workflow()
result = await workflow.run(checkpoint=checkpoint)
```

A checkpoint returned by a storage implementation should remain attached to that storage so subsequent snapshots continue to be persisted automatically.

The exact property names should be finalized in an ADR, but the intended public shape is:

```python
class Checkpoint:
checkpoint_id: str
workflow_name: str
graph_signature_hash: str
snapshots: list[CheckpointSnapshot]
head_snapshot_id: str | None
checkpoint_storage: CheckpointStorage | None # Runtime-only attachment

@property
def head_snapshot(self) -> CheckpointSnapshot | None: ...

def get_snapshot(self, snapshot_id: str) -> CheckpointSnapshot: ...

def restart_from(self, snapshot_id: str) -> None: ...

async def commit_snapshot(
self,
*,
state: dict[str, Any],
metadata: dict[str, Any] | None = None,
) -> CheckpointSnapshot: ...

class CheckpointSnapshot:
snapshot_id: str
parent_snapshot_id: str | None
timestamp: str
state: dict[str, Any]
metadata: dict[str, Any]

class Workflow:
def create_checkpoint(
self,
*,
checkpoint_id: str | None = None,
checkpoint_storage: CheckpointStorage | None = None,
) -> Checkpoint: ...

def run(
self,
message: Any | None = None,
*,
checkpoint: Checkpoint | None = None,
responses: Mapping[str, Any] | None = None,
stream: bool = False,
...,
) -> ...: ...
```

`checkpoint_storage` is attached to the checkpoint for automatic persistence but is not part of its serialized representation.

## Checkpoint and snapshot model

A `Checkpoint` represents the developer-owned lifecycle and snapshot history of one logical workflow execution.

A `CheckpointSnapshot` represents the complete workflow state at one resumable boundary:

- initial input has been accepted;
- a superstep has completed;
- the workflow has paused; or
- the workflow has completed.

`head_snapshot_id` identifies the snapshot from which the next run will continue. For a newly created checkpoint it is `None`. After a snapshot is committed, the head advances to that snapshot.

All mutable execution state must be stored under four framework-owned keys in `CheckpointSnapshot.state`:

```python
snapshot.state = {
"workflow_state": {
# Iteration count, run status, workflow-level shared state,
# continuation cursor, and other workflow-wide state.
},
"executor_states": {
# State keyed by executor ID.
#
# An AgentExecutor stores its AgentSession here rather than
# retaining the session on the workflow-owned executor instance.
},
"edge_states": {
# Delivery, fan-in, buffering, and other edge-runner state,
# keyed by stable edge identity.
},
"messages": {
"input": ...,
"pending": ...,
"output": ...,
},
}
```

These are the only top-level state buckets:

- `workflow_state` contains workflow-wide state, including iteration count and run status.
- `executor_states` contains state owned by individual executors. For an `AgentExecutor`, this includes its `AgentSession`.
- `edge_states` contains state required to restore message delivery and edge behavior.
- `messages.input` contains the input that entered the workflow.
- `messages.pending` contains messages or requests that have not yet been fully processed, including pending external-input requests.
- `messages.output` contains outputs produced by the workflow up to that snapshot.

Framework additions should be made inside one of these buckets rather than adding more top-level snapshot-state fields.

Snapshot identity, lineage, timestamps, format version, and diagnostic metadata remain typed snapshot fields because they describe the snapshot rather than the workflow execution state.

## Stateless workflow execution

`Workflow` should contain only the reusable workflow definition:

- graph topology and graph signature;
- executor definitions or factories;
- input/output metadata;
- execution configuration such as `max_iterations`.

Every call to `run()` should create an ephemeral runner hydrated from the checkpoint's head snapshot.

After the run returns, no mutable execution state may remain on the workflow. The supplied checkpoint is the only object through which execution state survives across calls.

A single workflow must support concurrent runs with different checkpoints. Concurrent use of the same checkpoint must be rejected to prevent lost snapshots and conflicting state updates.

Stateful executor instances are the main complication. Moving only `Workflow._runner` into a local variable is insufficient if executor objects continue to retain mutable fields. The implementation must either:

- create run-scoped executor instances; or
- move all executor state into `CheckpointSnapshot.state["executor_states"]`.

`AgentExecutor` must similarly stop retaining its `AgentSession` on a workflow-owned executor instance. Nested agent sessions belong in the executor's snapshot state.

## Restarting from an earlier snapshot

A checkpoint must make it easy to restart from any retained snapshot, not only the latest one:

```python
checkpoint = await storage.load("factor-run-123")
checkpoint.restart_from(snapshot_id="snapshot-5")

await workflow.run(checkpoint=checkpoint)
```

`restart_from()` changes the checkpoint head but does not delete snapshots created after the selected snapshot.

The next committed snapshot uses the selected snapshot as its parent:

```text
snapshot-4
|
snapshot-5 -------- snapshot-6b
|
snapshot-6 -------- snapshot-7
```

This preserves the original history while allowing a new execution branch to be created from `snapshot-5`.

Required behavior:

- `restart_from()` validates that the snapshot belongs to the checkpoint.
- Workflow graph compatibility is validated before execution resumes.
- The selected snapshot becomes `head_snapshot`.
- The next run restores exclusively from the selected snapshot's state.
- New snapshots set `parent_snapshot_id` to the current head.
- A successful snapshot commit advances `head_snapshot_id`.
- Existing descendant snapshots are retained until the developer explicitly deletes or prunes them.
- Storage persists both snapshot lineage and the checkpoint head.
- Loading a checkpoint restores its persisted head while still allowing another retained snapshot to be selected.

This makes the checkpoint a small snapshot graph rather than assuming that its history is always linear.

## Committing and persisting snapshots

The workflow runner should only export its complete state. Snapshot creation, lineage, persistence, and head advancement belong to `Checkpoint.commit_snapshot()`:

```python
snapshot = await checkpoint.commit_snapshot(
state=export_complete_run_state(),
metadata={"boundary": "superstep_completed"},
)
```

Conceptually, `commit_snapshot()` performs:

```python
class Checkpoint:
async def commit_snapshot(
self,
*,
state: dict[str, Any],
metadata: dict[str, Any] | None = None,
) -> CheckpointSnapshot:
previous_head = self.head_snapshot_id

snapshot = CheckpointSnapshot(
parent_snapshot_id=previous_head,
state=state,
metadata=metadata or {},
)

if self.checkpoint_storage is not None:
await self.checkpoint_storage.commit_snapshot(
self,
snapshot,
expected_head_snapshot_id=previous_head,
)

self.snapshots.append(snapshot)
self.head_snapshot_id = snapshot.snapshot_id
return snapshot
```

The actual implementation must synchronize commits so two callers cannot concurrently advance the same checkpoint.

`commit_snapshot()` is responsible for:

1. Capturing the current head.
2. Creating the snapshot and parent relationship.
3. Asking attached storage to persist the snapshot and new head atomically.
4. Appending the snapshot to in-memory history only after persistence succeeds.
5. Advancing the in-memory head.
6. Returning the committed snapshot.

The runner must not construct lineage, call storage directly, or update `head_snapshot_id` itself.

The lower-level storage contract can be:

```python
class CheckpointStorage(Protocol):
async def commit_snapshot(
self,
checkpoint: Checkpoint,
snapshot: CheckpointSnapshot,
*,
expected_head_snapshot_id: str | None,
) -> None: ...

async def load(self, checkpoint_id: str) -> Checkpoint: ...

async def delete(self, checkpoint_id: str) -> None: ...
```

Persisting a snapshot must also persist the checkpoint's new head in the same logical operation. There must not be separate `save_snapshot()` and `set_head()` calls.

`CheckpointStorage.commit_snapshot()` must atomically:

1. Persist the new snapshot and its parent relationship.
2. Set the persisted checkpoint head to the new snapshot.
3. Validate that the persisted head still matches `expected_head_snapshot_id`.
4. Fail with a checkpoint-conflict error if another writer advanced the head.

For file storage, this can be implemented through atomic replacement of the checkpoint representation or an atomically replaced manifest containing the head. Database-backed implementations should update the snapshot and head in one transaction.

If persistence fails:

- the persisted head remains unchanged;
- the in-memory head remains unchanged;
- the candidate snapshot is not visible as a committed snapshot;
- the failure is surfaced to the caller.

A backend may temporarily leave an unreachable snapshot record after an interrupted write, but loading the checkpoint must only expose snapshots reachable through committed checkpoint metadata.

When no storage is attached, `Checkpoint.commit_snapshot()` appends the snapshot and advances the head in memory.

Checkpoint recovery remains at-least-once. The framework cannot transactionally roll back external side effects completed during a superstep, so this limitation must be documented.

## Snapshot history and lifecycle ownership

The checkpoint has one stable `checkpoint_id` for its complete lifecycle. Individual supersteps create `snapshot_id` values rather than unrelated top-level checkpoint IDs.

This removes the need for:

```python
await storage.get_latest(workflow_name=workflow.name)
```

The caller already owns the checkpoint and can inspect its head directly:

```python
latest = checkpoint.head_snapshot
```

The developer owns:

- checkpoint identity;
- checkpoint retention;
- snapshot retention or pruning;
- head selection;
- deletion;
- persistence backend selection;
- restoration and branching.

Storage implementations may persist snapshots incrementally or load older snapshots lazily, but those implementation details must not change the public checkpoint model.

## Adopt session-style msgspec serialization

Replace the current JSON plus pickle/base64 checkpoint encoding with the same approach used by `FileSessionStore`:

- A typed, versioned internal `msgspec.Struct` for `Checkpoint`.
- A typed, versioned internal `msgspec.Struct` for `CheckpointSnapshot`.
- msgspec JSON as the readable default.
- Optional msgspec MessagePack using the same model.
- Atomic file writes or transactional backend writes.
- Strict envelope and version validation.
- Explicit custom-state type registration.
- Clear errors for unsupported values and unknown versions.

The dynamic `CheckpointSnapshot.state` dictionary should use the same recursive codec and type registry as `AgentSession.state`. We should not maintain separate serialization systems for agent and workflow state.

Types registered through `register_state_type` should round-trip in both agent sessions and workflow snapshots. If that name is too agent-specific, it can be renamed to a neutral state-serialization registration API as part of the v2 break.

Remove the pickle-based checkpoint path, including:

- `register_checkpoint_type`;
- `allowed_checkpoint_types`;
- restricted-unpickler configuration;
- pickle/base64 marker handling.

New v2 writes must never silently fall back to pickle. If v1 checkpoints must survive the major-version upgrade, provide an explicit migration reader or offline migration path rather than enabling automatic legacy deserialization during normal v2 loads.

## Other required updates

- Remove checkpoint storage from `WorkflowBuilder`.
- Replace `checkpoint_id` and `checkpoint_storage` arguments on `Workflow.run()` with `checkpoint`.
- Move the workflow-wide active-run guard to the checkpoint.
- Rework `RunnerImpl`, `InProcRunnerContext`, executor state, and edge-runner state so they are run-scoped.
- Move `AgentExecutor`'s nested `AgentSession` into snapshot executor state.
- Update functional workflows and pipelines to use the same checkpoint model.
- Update nested workflows so child workflow state is represented in the parent's snapshot state.
- Update `WorkflowAgent` to store its workflow checkpoint through the surrounding `AgentSession.state`.
- Update hosting `WorkflowState` to store and retrieve checkpoints instead of maintaining a separate `session_id -> checkpoint_id` cursor.
- Update AG-UI, A2A, Foundry hosting, and other adapters that currently forward `checkpoint_id` or `checkpoint_storage`.
- Update in-memory, file, Cosmos, and other checkpoint storage implementations.
- Include checkpoint and snapshot identifiers in telemetry.
- Update checkpoint, human-in-the-loop, hosting, and orchestration samples.
- Add an ADR covering ownership, concurrency, storage ordering, snapshot retention, serialization, and v1 migration.

## Acceptance criteria

- `workflow.create_checkpoint()` returns a caller-owned `Checkpoint`.
- `workflow.run(..., checkpoint=checkpoint)` restores from and updates that checkpoint.
- A checkpoint contains snapshots, and every snapshot contains the complete resumable state dictionary.
- Snapshot state has exactly four top-level buckets: `workflow_state`, `executor_states`, `edge_states`, and `messages`.
- Messages are divided into `input`, `pending`, and `output`.
- Workflow-wide values such as iteration count and run status are stored inside `workflow_state`.
- An `AgentExecutor` stores its `AgentSession` inside its entry in `executor_states`.
- A `Workflow` retains no mutable execution state between calls.
- Separate checkpoints can run concurrently against one workflow without state leakage.
- Concurrent use of the same checkpoint is rejected.
- Omitting a checkpoint produces an independent one-shot run.
- Reusing a checkpoint preserves all workflow, executor, edge, and message state.
- An entry snapshot and a snapshot after each completed superstep are created.
- Snapshot creation, parent assignment, persistence, and head advancement are encapsulated by `Checkpoint.commit_snapshot()`.
- Workflow runners do not manipulate checkpoint lineage or storage directly.
- Attached storage persists every snapshot and its new checkpoint head as one atomic, conflict-checked operation.
- Storage failures are surfaced and leave the checkpoint at its last durable head.
- A stored checkpoint can resume against a rebuilt workflow in a new process.
- A caller can restart from any retained snapshot without deleting newer snapshots.
- Restarting from an earlier snapshot creates a new lineage branch.
- Each snapshot records its parent, and the checkpoint records its current head.
- Graph-signature compatibility is validated before restoration.
- JSON and MessagePack checkpoints round-trip registered framework and application state through msgspec.
- New checkpoint serialization contains no pickle payloads.
- Unsupported versions, incompatible graph signatures, corrupt data, and unregistered custom values fail clearly.
- Existing workflow integrations and samples use the new checkpoint-based API.
- The change is introduced only in v2; v1 does not gain a parallel lifecycle model.

## Non-goals

- Introducing a separate `WorkflowSession` abstraction.
- Exactly-once execution of external side effects.
- Automatic snapshot retention or garbage collection.
- Treating checkpoint IDs as authorization boundaries.
- Silently loading legacy pickle-based checkpoints in the normal v2 code path.

### Language/SDK

Python

Contributor guide

Open the contributing guide

Research direction

Start with the Workflow, WorkflowBuilder, runner, AgentExecutor, Checkpoint, and CheckpointStorage entry points described in the issue. Trace where mutable execution state, executor sessions, edge state, and snapshots are currently retained. Done means caller-owned checkpoints support stateless concurrent workflow runs, restartable snapshot branches, atomic persistence, and the proposed v2 behavior.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
ai, backend-api-design
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Active
Clarity
Mostly clear
Newbie friendliness
25/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.