Python: Resolve checkpoint get_latest timestamp ties through lineage

[Shivam Sharma (ShivamSharma43)]

Motivation & Context

InMemoryCheckpointStorage.get_latest and FileCheckpointStorage.get_latest choose max() by timestamp alone. Checkpoints taken at the same superstep boundary can carry identical timestamps (the WorkflowCheckpoint docstring notes this), and datetime.now() resolution can make distinct boundaries collide too. On a tie, max() returns whichever checkpoint is iterated first, so the result depends on save order (or file listing order), and a workflow resuming from get_latest can restore a superseded checkpoint without any error.

Description & Review Guide

• What are the major changes? A private _select_latest_checkpoint helper in _workflows/_checkpoint.py, used by both storages. It keeps the checkpoints with the newest timestamp and, among those, skips any that another one names as its previous_checkpoint_id, as the WorkflowCheckpoint docstring describes for lineage ordering.
• What is the impact of these changes? With a parent, child and grandchild sharing a timestamp, get_latest now returns the grandchild in any save order, for both storages. Without a tie, the result is the same as before. If tied checkpoints are not related by lineage (for example two separate branches), the first one is returned, as before.
• What do you want reviewers to focus on? Whether the other CheckpointStorage implementations (Cosmos, Foundry hosting, the AG-UI wrapper) should get the same tie-breaking. I kept this PR to the two storages named in the issue.

test_get_latest_resolves_timestamp_ties_through_lineage in tests/workflow/test_checkpoint.py covers both storages and both save orders. Without the fix, 3 of the 4 cases fail; the in-memory reverse-order case passes by chance.

Related Issue

Fixes #8513

Contribution Checklist

• The code builds clean without any errors or warnings
• All unit tests pass, and I have added new tests where possible
• The PR follows the Contribution Guidelines
• This PR is linked to an issue and there is no other open PR for this issue (see Related Issue above).
• This is not a breaking change. If it is a breaking change, add the breaking change label (or add "[BREAKING]" to the title prefix, before or after any language prefix) — a workflow keeps the label and title prefix in sync automatically.

🤖 Generated with Claude Code

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

[Eduard van Valkenburg (eavanvalkenburg)]

/review

[Eduard van Valkenburg (eavanvalkenburg)]

/review

[github-actions]

MAF Automated Review — Iteration 1

Result: No findings
Scope: full PR (1 commit(s)): 740fca80680c
Model: gpt-5.6-sol

Overview

The PR preserves timestamp-primary ordering while using lineage to select the unsuperseded tip when the newest checkpoints tie, and applies that behavior consistently to the two changed core backends. Workflow scoping, empty-store guards, caller-owned copy semantics, and forward/reverse save-order tests across both implementations support the change. The remaining provider-parity concern is pre-existing in unchanged backends, so no publishable PR-introduced finding remains.

Reviewed the supplied pull-request change set across correctness, security/reliability, architecture, and failure behavior.
No publishable findings remained after source verification for this scope.

[Eduard van Valkenburg (eavanvalkenburg)]

Shivam Sharma (@ShivamSharma43) this fixes the in-memory and file stores, but the same stale-resume bug remains in the first-party Cosmos and Foundry checkpoint stores. python/packages/azure-cosmos/agent_framework_azure_cosmos/_checkpoint_storage.py:356 still asks the backend for one checkpoint ordered only by timestamp, and python/packages/foundry_hosting/agent_framework_foundry_hosting/_state_store.py:230 still uses max(...timestamp...); with a parent/child pair sharing the newest timestamp, either can return the superseded parent depending on backend/list order. Since this PR closes #8513 and CheckpointStorage.get_latest is the public behavior, please apply the lineage tie-breaker to these implementations as well. Cosmos will need to fetch enough newest-timestamp candidates to resolve previous_checkpoint_id ties instead of limiting to one row.

[Shivam Sharma (ShivamSharma43)]

Eduard van Valkenburg (@eavanvalkenburg) thanks, good catch. Applied the same lineage tie-breaker to both first-party stores in 83ecc69:

• Cosmos (CosmosCheckpointStorage.get_latest): no longer ORDER BY c.timestamp DESC OFFSET 0 LIMIT 1. It now runs SELECT VALUE MAX(c.timestamp) for the workflow partition, then fetches every row at that timestamp (c.workflow_name = @workflow_name AND c.timestamp = @timestamp, parameterized and partition-scoped like the other queries) and picks the candidate no other candidate names as its previous_checkpoint_id. The MAX query returns a single value, so only the tied rows are transferred rather than a page of full checkpoint documents.
• Foundry (FoundryCheckpointStore.get_latest): replaced max(...timestamp...) with the same tie-breaker over the listed checkpoints.

The helper is mirrored in each package rather than imported from core's _select_latest_checkpoint, because both packages still allow agent-framework-core>=1.19.0, which doesn't include it. If you'd rather bump the core floor and import it, I can switch to that.

Tests: added forward/reverse row-order parent→child→grandchild tie tests for both stores, and updated the existing Cosmos get_latest tests for the two-query shape. One of them asserts both queries' parameters and partition key, and another covers the tied rows disappearing between the two queries. Local runs: azure-cosmos 183 passed, foundry_hosting 867 passed, core 8131 passed (one unrelated Windows MAX_PATH failure in test_sessions.py). ruff, pyright and mypy are clean on the touched files.

[Eduard van Valkenburg (eavanvalkenburg)]

Shivam Sharma (@ShivamSharma43) The parent/child lineage tie-breaker is now applied across the first-party stores, but same-timestamp ties with multiple unsuperseded lineage tips are still backend-order dependent. _select_latest_checkpoint() (and the duplicated Cosmos/Foundry helpers) returns the first candidate not referenced as a parent; two independent leaves therefore produce different results when their input order is reversed. Please apply a stable fallback among remaining leaf candidates, such as checkpoint_id, in all three helpers and add reversed-order coverage for two unsuperseded same-timestamp leaves.

[Shivam Sharma (ShivamSharma43)]

Eduard van Valkenburg (@eavanvalkenburg) thanks, agreed. Fixed in ff52e60. All three helpers (core _select_latest_checkpoint and the Cosmos and Foundry copies) now collect every candidate not named as a parent and return the one with the greatest checkpoint_id. Among unsuperseded same-timestamp leaves, the choice no longer depends on save, list or row order. If a malformed cycle leaves no unsuperseded candidate, the same checkpoint_id rule applies to all tied candidates, replacing the old candidates[0] fallback.

Added test_get_latest_breaks_ties_between_unsuperseded_checkpoints_by_id for in-memory and file storage (core), Cosmos and Foundry. Each case has two leaves sharing a parent and a timestamp, saved/listed/returned in both orders, and expects leaf-b both times. Five of the eight cases fail without the change; with it the touched test files pass. The two Foundry agentserver-state tests that fail here hit Windows MAX_PATH in the temp dir and fail the same way without the change.