Skip to content

Latest commit

 

History

History
251 lines (218 loc) · 13.7 KB

File metadata and controls

251 lines (218 loc) · 13.7 KB

Python Workflow Checkpoint Store implementation

This document describes implementation-map slice 17. Normative requirements are persistence, architecture, interfaces, resilience, security, testing, packages, and samples. ADRs 0012, 0018, and 0009 record rationale; their proposed status does not override the specifications.

Public contract and scope

agent_framework_mongodb.MongoDBCheckpointStorage explicitly derives from the public Agent Framework Core 1.13 CheckpointStorage protocol and implements its exact asynchronous seam:

  • save(checkpoint) -> CheckpointID
  • load(checkpoint_id) -> WorkflowCheckpoint
  • list_checkpoints(*, workflow_name) -> list[WorkflowCheckpoint]
  • delete(checkpoint_id) -> bool
  • get_latest(*, workflow_name) -> WorkflowCheckpoint | None
  • list_checkpoint_ids(*, workflow_name) -> list[CheckpointID]

Agent Framework's protocol has no tenant or run parameters. Therefore MongoDBCheckpointStorageOptions immutably binds a required tenant_id, workflow_name, and session_id at construction. application_id is optional. Every read, write, sort, limit, and delete includes the discriminator and all raw scope fields. workflow_name arguments and checkpoint payloads must equal the bound value. A checkpoint ID alone is never an authorization filter.

The inherited list methods enumerate the complete authorized run in deterministic order by repeatedly fetching configured bounded pages. Cancellation propagates from every page request. The additional list_checkpoint_page(..., cursor=None, limit=None) API lets callers consume one bounded MongoDBCheckpointPage(checkpoints, next_cursor) at a time. The default page size is 100, the configurable hard maximum defaults to 1000, and invalid or unknown-version cursors fail closed. Ordering is (sequence, checkpoint_id), not timestamp or iteration_count.

Serialization and immutable records

The implementation calls only the public WorkflowCheckpoint.to_dict() and WorkflowCheckpoint.from_dict() serialization methods. It does not import or reflect over framework internals. The public dictionary is encoded as BSON binary with Python pickle so public framework message/event objects and application executor state remain lossless. Loading uses a restricted unpickler: safe built-ins and concrete agent_framework types are permitted, while application types must be explicitly listed as module:qualname values in allowed_checkpoint_types. That allowlist permits class resolution only; it does not permit custom __reduce__, __reduce_ex__, __getstate__, __setstate__, or new-argument pickle hooks. A controlled pickler ignores the process-global copyreg dispatch table and rejects custom hooks before they execute. Application dataclasses and public to_dict types must use default object serialization. The implementation retains a narrow exact-type allowlist for standard values whose Python reducers are part of the documented codec: None, bool, int, float, str, bytes, bytearray, exact dict, list, tuple, set, frozenset, date, datetime, time, timedelta, timezone, Decimal, and UUID. The default Enum reduction is allowed; enum classes that override pickle hooks are rejected.

Python's separate copyreg extension registry can replace global class references with process-local EXT codes even when the dispatch table is empty. Save recursively walks the actual serialized object graph, including nested framework objects, dataclass fields, enum values, instance dictionaries, and slots. If any value's exact (module, qualname) or (module, name) is present in that registry, save raises MongoDBSerializationError with guidance to remove the registration or migrate the value. This happens before pickle, sequence allocation, or MongoDB I/O. Restricted-load ValueError failures, including unregistered EXT codes, are translated to the same stable category.

Before allocating a sequence or contacting MongoDB, save creates the exact payload bytes, decodes them through the same restricted load path, recursively canonicalizes the decoded graph, and compares it with the original canonical graph. Unsupported values or any mismatch raise MongoDBSerializationError. This validates the bytes that will actually be stored rather than assuming that logical projection and pickle behavior agree.

Pickle is appropriate only for application-owned, access-controlled checkpoint storage. It is not a boundary against an attacker who can modify the collection. Never load checkpoint documents from untrusted input. Use TLS, MongoDB access controls, encryption at rest, and deployment-owned client-side field-level encryption where required.

Each immutable checkpoint document is:

{
  "_id": "<deterministic scoped sha-256>",
  "_kind": "workflow_checkpoint",
  "schema_version": 1,
  "framework_version": "agent-framework-core/1:WorkflowCheckpoint.to_dict/v1",
  "payload_version": "1.0",
  "idempotency_hash_version": 2,
  "scope_discriminator": "<scope sha-256>",
  "tenant_id": "tenant-1",
  "application_id": "application-1",
  "workflow_name": "approval-workflow",
  "session_id": "run-1",
  "checkpoint_id": "framework checkpoint ID",
  "parent_checkpoint_id": "optional exact lineage edge",
  "sequence": 12,
  "created_at": "<UTC BSON datetime>",
  "expires_at": "<optional UTC BSON datetime>",
  "checkpoint": "<BSON binary of the public dictionary>",
  "payload_hash": "<sha-256>"
}

The framework checkpoint ID is preserved exactly, while _id is deterministic for the complete scope and ID. Idempotency hashes use a versioned canonical logical representation of the public checkpoint dictionary rather than pickle bytes. Exact built-in dict is the only supported mapping type, and its values are explicitly order-insensitive because their public checkpoint meaning is key/value state. The mapping guard runs before scalar, enum, dataclass, to_dict, and every other type handler, then exact-dict keys and values recurse through the same validation. OrderedDict, mapping/enum multiple inheritance, and every mapping subclass are rejected before sequence allocation with MongoDBSerializationError, even when allowlisted. Callers must migrate them to plain dict/list structures. This intentionally eliminates instance reducer, iterator-state, and concrete mapping semantics from the canonical contract. The controlled pickler excludes attempted copyreg.dispatch_table registrations; unit coverage installs a divergent reducer and proves the stored pickle bytes and canonical hash remain unchanged. Sets are stably ordered, scalar and collection types carry explicit tags, and framework/application dataclasses or public to_dict values carry stable type identities. The same logical checkpoint therefore hashes identically across processes and PYTHONHASHSEED values without discarding lossless concrete-type or order semantics. Cycles, non-finite floats, and unsupported objects fail before sequence allocation with a stable MongoDBMappingError. Pickle remains only the lossless storage encoding and is not part of identity. Hash version 2 rejects version 1 records with migration guidance. An identical retry returns the same ID. Reusing the ID with different public state raises MongoDBConcurrencyError. schema_version, framework_version, the checkpoint's public version, and idempotency_hash_version are independent compatibility gates. Unknown values raise MongoDBMappingError with migration guidance rather than best-effort loading. Python/.NET physical checkpoint interoperability is not claimed.

Sequence allocation, lineage, and retention

A separate, scoped counter document uses an atomic aggregation-pipeline upsert. Concurrent saves therefore receive unique, positive, monotonic sequences. When allocating, storage first reads the greatest retained checkpoint sequence in the exact authorized scope. The atomic upsert computes the allocation from the maximum of that observed value and the current counter before adding the batch count. Concurrent recovery after a missing counter therefore cannot reset or collide with retained sequences. When TTL is configured, the same update extends the counter's expires_at to the maximum of its current and requested values; an out-of-order shorter TTL can never move it backward. A non-expiring checkpoint atomically changes retention_mode to permanent and removes expires_at. That mode is dominant, so a later concurrent TTL write cannot restore expiration. Legacy counters with a sequence but no expiration are treated as permanent. Counter metadata therefore cannot expire while any retained permanent checkpoint could still need its sequence. Retries and failed inserts may leave sequence gaps; ordering never assumes contiguity. get_latest() sorts by descending sequence and checkpoint ID.

previous_checkpoint_id is copied unchanged to parent_checkpoint_id. Parents are not required to exist at save or load time. This permits branched framework lineage and is required because MongoDB TTL removal is asynchronous and may expire a parent before a child. Restoring a child does not traverse or rewrite its parent edge.

ttl computes an optional UTC BSON-millisecond expires_at independently from Session Store, Chat History, and Memory retention. The TTL monitor provides eventual deletion; applications must not use expiration timing as workflow coordination.

clear_run() is the explicit authorized lifecycle operation for a completed run. It applies the complete constructor-bound scope to a checkpoint delete_many followed by the exact deterministic counter delete_one, returning MongoDBCheckpointClearResult with acknowledged checkpoint and counter counts. It is retry-safe best-effort cleanup rather than a cross-deployment transaction; callers must quiesce writers before clearing and may retry after a partial driver failure. It never issues an empty or ID-only delete.

Explicit regular indexes

Construction, save, load, and workflow hooks never mutate indexes. ensure_indexes() is the explicit provisioning operation and validate_indexes() is read-only.

Name Keys after scoped prefix Options
checkpoint_scope_identity checkpoint_id unique, simple collation
checkpoint_scope_sequence sequence unique, simple collation
checkpoint_scope_lineage parent_checkpoint_id simple collation
checkpoint_expiration expires_at expireAfterSeconds: 0, checkpoints
checkpoint_counter_expiration expires_at expireAfterSeconds: 0, counters

The scoped prefix is scope_discriminator, workflow_name, and session_id. Identity, sequence, lineage, and checkpoint TTL indexes have a checkpoint-only partial filter, so the internal counter cannot collide with checkpoint uniqueness. The separate counter TTL index has a counter-only partial filter. MongoDB may process the two TTL indexes in either order; correctness never relies on the counter outliving checkpoints because allocation recovers from the indexed retained maximum.

Runtime privileges are find, insert, atomic update/upsert for the sequence counter, and targeted delete on the checkpoint collection. Provisioning also requires createIndex; validation requires index-list access.

Errors, cancellation, ownership, and logs

Missing authorized records raise MongoDBCheckpointNotFoundError, which is both an integration retrieval error and the framework's WorkflowCheckpointException. Noncanonical lossless serialization raises MongoDBSerializationError, a mapping-error subtype. Configuration, mapping, concurrency, authorization, transient retrieval, transient persistence, and other MongoDB failures use the package's stable categories while preserving driver exceptions as __cause__. asyncio.CancelledError is never caught.

Injected clients and collections remain caller-owned. A storage created from a connection string owns its PyMongo AsyncMongoClient; close() and the async context manager close it once. Ownership never changes after an error.

Completion logs contain only feature, operation, outcome, bounded result count, duration, and stable error category. They exclude IDs, scopes, payloads, collection/database names, filters, driver messages, hosts, and credentials.

Verification

Public serialization, restricted round-trip equivalence, custom-reducer and copyreg-extension rejection, mapping/enum precedence, actual workflow pause/resume, cross-process canonical idempotency, stateful-mapping rejection, conflict, lineage, concurrent sequence, missing-counter recovery, complete inherited listing, bounded pagination, latest, scope cleanup, counter TTL, TTL-gap, compatibility, index, cancellation, error, and ownership tests are in python/tests/unit/test_checkpoint_storage.py. Language-neutral outcomes are in python/tests/contracts/fixtures/checkpoint_storage_contract.json. Credential-gated real-deployment coverage is in python/tests/integration_persistence/test_checkpoint_storage_integration.py and uses a unique test-checkpoint- collection with cleanup in finally.

From python, run:

uv run pytest tests\unit\test_checkpoint_storage.py tests\contracts\test_checkpoint_storage_contract.py
uv run pytest tests\integration_persistence -m integration_persistence
uv run ruff check src tests samples
uv run ruff format --check src tests samples
uv run mypy
uv run pyright

The integration command skips cleanly without MONGODB_URI.