Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -520,8 +520,17 @@ def _peek_next_checkpoint(self) -> CheckpointedResult:
)

def _next_operation_exists(self) -> bool:
"""True if a checkpoint exists for this context's next operation."""
return self._peek_next_checkpoint().is_existent()
"""True if the next operation or its first inner operation has a checkpoint."""
if self._peek_next_checkpoint().is_existent():
return True
# A virtual child never checkpoints itself, so its history starts at its
# first inner operation.
first_child_id: str = OperationIdNamespace(
self._peek_next_operation_id()
).create_id_for_step(1)
return self.state.get_checkpoint_result(
first_child_id, claim_end_hook=False
).is_existent()

@contextmanager
def _replay_aware(
Expand All @@ -542,14 +551,20 @@ def _replay_aware(

- Existence flip (before the op): if we are replaying and the next
operation has no checkpoint at all, it is brand-new code, so flip to
NEW immediately so the operation and its logs count as new.
NEW immediately so the operation and its logs count as new. A virtual
child never checkpoints itself, so a missing checkpoint at its id is
not this flip: with `is_virtual=True` the context stays in replay and
the child refines against its inner operation ids.
- Deferred status flip (after the op): a non-terminal next operation is
a pure resume point with no user body. We keep replay status through the
operation and flip to NEW afterwards, so the resuming operation's own
logs stay de-duplicated but subsequent code counts as new.
- Post-op existence flip (after the op): if we are still replaying once
the operation completes and the *following* operation does not exist
yet, we have reached the replay boundary.
yet, we have reached the replay boundary. Existence includes a
checkpoint at the following operation's own id or, when that is
missing, at its first inner operation id. Virtual children never
checkpoint themselves, so their history starts at that inner id.
"""
was_replaying: bool = self.is_replaying()
# Only look up a checkpoint when replaying. Operation call sites pass
Expand Down Expand Up @@ -606,10 +621,17 @@ def _replay_aware(
was_replaying and next_exists and not next_terminal and not next_is_step
)
# Before-the-op flips:
# - brand-new next op (no checkpoint): always flip to NEW.
# - brand-new next op (no checkpoint): flip to NEW, unless the operation
# is virtual. A virtual container never checkpoints itself, so the
# absence of a checkpoint at its id is by design and is not evidence
# of new work; its inner operations may still have history to replay.
# Keep replay status and let the child refine to NEW via its own
# look-ahead against its inner step ids.
# - non-terminal STEP-type op (brand-new or retrying): the user function
# is about to run real work, so flip to NEW before it.
if was_replaying and (not next_exists or (next_is_step and not next_terminal)):
if was_replaying and (
(not next_exists and not is_virtual) or (next_is_step and not next_terminal)
):
Comment on lines +632 to +634

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This fixes a virtual child that is the first operation in its context, but not one that comes after a replayed operation. After each operation, context.py:626, the look-ahead switches the context to NEW when context.py:522-524, _next_operation_exists() finds no checkpoint at the next operation's id. A virtual child never has a checkpoint at its own id. So when a replayed operation comes before it, the context is already NEW by the time run_in_child_context reaches this condition, and the child runs as NEW. Can _next_operation_exists() also check the first operation inside the next one, as the JS SDK does?

Two things I checked support this:

  1. With DurableFunctionTestRunner and a plugin, I ran a handler that calls a step and then a virtual child that waits. On the resumed invocation it behaves as on main: is_replaying() returns False inside the child, and the plugin gets no on_operation_end for the wait. The second of two virtual children in a row fails in the same two ways. When the virtual child is the first operation, this branch makes is_replaying() return True and delivers the on_operation_end.
  2. The JS SDK fixed the same problem in aws-durable-execution-sdk-js #697: keep replay mode active after a virtual context. Before a context leaves replay mode, durable-context.ts:316-333, checkAndUpdateReplayMode also checks the first child id, ${pendingStepId}-1.

This is the change I tested on top of this branch:

def _next_operation_exists(self) -> bool:
    """True if the next operation or its first inner operation has a checkpoint."""
    if self._peek_next_checkpoint().is_existent():
        return True
    # A virtual child never checkpoints itself, so its history starts at its
    # first inner operation.
    first_child_id: str = OperationIdNamespace(
        self._peek_next_operation_id()
    ).create_id_for_step(1)
    return self.state.get_checkpoint_result(
        first_child_id, claim_end_hook=False
    ).is_existent()

With it, the core unit and e2e tests pass, and both handlers above get the right is_replaying() and on_operation_end. The second read is an in-memory lookup. It runs only while replaying, when the next operation has no checkpoint at its own id. One case is still not covered: after a replayed operation, the context still switches to NEW before a virtual child that has no checkpoint at its first inner id, such as an empty one. JS has the same limit, so I think that can stay out of this PR.

Can the context.py:543-552, _replay_aware docstring also describe the virtual case? It still says that a next operation without a checkpoint switches the context to NEW, both before and after the operation.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Addressed in 84b7b61. _next_operation_exists() now also checks the first inner id via OperationIdNamespace(...).create_id_for_step(1), matching the JS SDK. The _replay_aware docstring now describes that virtual case.

Kept the is_virtual before-op guard: it is still needed when the virtual child is the first operation in its context. _next_operation_exists() only runs after a previous operation, so it cannot cover that first-op path.

self._set_replay_status_new()
try:
yield
Expand Down
Loading