Python: executor on_checkpoint_save return value is validated only at restore time: non-dict state checkpoints fine then fails to restore on any storage

[ptimizeroracle]

description

_save_executor_states in _runner.py calls executor.on_checkpoint_save() and stores whatever comes back with no shape check. _restore_executor_states then validates isinstance(state, dict) and raises WorkflowCheckpointException("Executor state for {id} is not a dict[str, Any]. Unable to restore.").

so an executor whose on_checkpoint_save returns a list (or any non-dict) checkpoints successfully, the run completes, and every later restore of those checkpoints fails. this is the same save-accepts-what-restore-rejects principle as #8181, but a different locus (the executor state hooks, _runner.py, not the file storage) and it reproduces on InMemoryCheckpointStorage too, so it is not pickle- or storage-specific.

the contract is documented on the Executor base class (on_checkpoint_save() -> dict[str, Any], _executor.py:525), so a returning a non-dict is user error, but the failure surfaces only at restore time, potentially in a different process. a save-side isinstance check (raise WorkflowCheckpointException from _save_executor_states) would make it loud at capture time.

one related observation i am not claiming as a bug: when checkpoint encoding fails mid-run, the run continues by design ("note that this does not fail the workflow run" in the log message). combined with the asymmetry above, an executor returning an unencodable value means the run finishes with no checkpoints at all and only log lines saying so. that may be intended resilience, flagging it only as context.

checked tests (test_checkpoint.py, test_agent_executor.py) and the tracker: nothing pins the non-dict path and no duplicate found. found via a probe of the executor state hooks. ai assistance disclosed. standalone offline repro below.

reproduction steps

import asyncio, sys
sys.path.insert(0, "python/packages/core")
from agent_framework import WorkflowBuilder, WorkflowContext, handler
from agent_framework._workflows._checkpoint import InMemoryCheckpointStorage
from agent_framework._workflows._executor import Executor

class BadStateExecutor(Executor):
    @handler
    async def run(self, message: str, ctx: WorkflowContext) -> None:
        await ctx.yield_output(message + "-done")

    async def on_checkpoint_save(self):
        return ["not", "a", "dict"]  # type: ignore[override]

async def main():
    store = InMemoryCheckpointStorage()
    ex = BadStateExecutor(id="only")
    wf = WorkflowBuilder(max_iterations=5, start_executor=ex, checkpoint_storage=store).build()
    _ = [e async for e in wf.run("hi", stream=True)]          # run completes
    ckpts = await store.list_checkpoints(workflow_name=wf.name)
    print(f"{len(ckpts)} checkpoints saved")                    # checkpoints exist

    ex2 = BadStateExecutor(id="only")
    wf2 = WorkflowBuilder(max_iterations=5, start_executor=ex2, checkpoint_storage=store).build()
    _ = [e async for e in wf2.run(checkpoint_id=ckpts[-1].checkpoint_id, stream=True)]
    # WorkflowCheckpointException: Executor state for only is not a dict[str, Any]. Unable to restore.

asyncio.run(main())

environment

• agent-framework-core @ main $SHA (editable, python/packages/core)
• python 3.12, macOS, offline repro

[manjunathshiva]

I'd like to take this one.

Confirmed the asymmetry on main. _save_executor_states (_runner.py:398-406) passes whatever on_checkpoint_save() returns straight into _set_executor_state, which only validates the container and never the state itself, while _restore_executor_states (_runner.py:423-427) checks isinstance(state, dict) and the key types and raises WorkflowCheckpointException. So a hook returning the wrong shape is accepted at capture time and only fails later at restore, against the dict[str, Any] contract documented on Executor.on_checkpoint_save (_executor.py:525). Storage-agnostic, as you say — the check lives in the runner, so InMemoryCheckpointStorage behaves the same.

Planning the same guard on the save side so the failure surfaces where the bad value came from. Worth noting it lands softly rather than as a new crash: create_checkpoint_if_enabled (_runner.py:268-276) already wraps the call in try/except Exception and logs a warning without failing the run, and the other caller, build_checkpoint(), is reached through sub-workflow checkpointing at _workflow_executor.py:471, which is also exception-wrapped. So the practical effect is a loud warning at capture time instead of a silent bad checkpoint.

Two things I'd rather confirm than assume before opening a PR:

1. Should the save-side guard raise WorkflowCheckpointException (matching restore) or log and skip that executor's state? Raising is symmetric, but given both callers swallow exceptions it effectively becomes "warn and drop the whole checkpoint", where skipping would drop one executor's state and keep the rest.
2. Should the message name the offending type, so an author sees what their hook actually returned?

Happy to go with your preference on both; I'll pick raising-with-the-type unless you say otherwise.

[github-actions]

🤖 Automated triage reproduction notes (agent-authored — trust but verify)

Agent analysis

The bug reproduces in python/packages/core/agent_framework/_workflows/_runner.py::_save_executor_states at lines 398-404 when Executor.on_checkpoint_save() returns a non-dict; _restore_executor_states later rejects that persisted value at lines 423-426. Minimal repro: run a checkpoint-enabled single-executor workflow backed by InMemoryCheckpointStorage, return a list from the save hook, then restore either generated checkpoint.

• Failing test: python/packages/core/tests/workflow/test_checkpoint.py::test_workflow_rejects_non_dict_executor_state_when_saving
• Files examined: python/AGENTS.md, python/packages/core/pyproject.toml, python/packages/core/agent_framework/_workflows/_executor.py, python/packages/core/agent_framework/_workflows/_runner.py, python/packages/core/tests/workflow/test_checkpoint.py
• Tests run: python/packages/core/tests/workflow/test_checkpoint.py (46 passed baseline), python/packages/core/tests/workflow/test_checkpoint.py::test_workflow_rejects_non_dict_executor_state_when_saving (failed)
• Current version: 1.17.0