`at_most_once` raises in simple usage

[rileysdev]

We have a test for at_most_once which passes without raising. When I created the simplest at_most_once example I could think of in a tool call, I got this stack trace thought the test still passes:

Traceback (most recent call last):
  File "/workspaces/durable-mcp-python/reboot/mcp/server.py", line 946, in wrapper
    return await fn(**dict(bound.arguments))
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/workspaces/durable-mcp-python/tests/test_at_most_once_simple.py", line 38, in add
    result = await at_most_once(
             ^^^^^^^^^^^^^^^^^^^
  File "/workspaces/durable-mcp-python/.venv/lib/python3.12/site-packages/rebootdev/aio/workflows.py", line 109, in at_most_once
    return await memoize(
           ^^^^^^^^^^^^^^
  File "/workspaces/durable-mcp-python/.venv/lib/python3.12/site-packages/rebootdev/aio/memoize.py", line 109, in memoize
    raise AtMostOnceFailedBeforeCompleting(
rebootdev.aio.workflows.AtMostOnceFailedBeforeCompleting: ... it looks like an external failure occurred (e.g., the machine failed, your container was rescheduled, etc) while your code was executing

Here is a branch with a simple at_most_once test that exhibits the same behavior: https://github.com/reboot-dev/durable-mcp-python/tree/riley/at_most_once_simple_test

File: https://github.com/reboot-dev/durable-mcp-python/blob/riley/at_most_once_simple_test/tests/test_at_most_once_simple.py

[skaar]

I think this is the at_most_once issue I ran into in previous tests/examples, where I temporarily worked around it by disabling effect validation - I patched my local memoize.py to work around it:

--- .venv/lib/python3.12/site-packages/rebootdev/aio/memoize.py.orig    2025-11-03 15:58:38.807487171 -0500
+++ .venv/lib/python3.12/site-packages/rebootdev/aio/memoize.py 2025-11-03 16:02:42.028862772 -0500
@@ -106,13 +106,28 @@
     status = await memoize.always().Status(context)
 
     if at_most_once and status.started and not status.stored:
-        raise AtMostOnceFailedBeforeCompleting(
-            status.failure if status.failed else (
+        # We have `started=True` but `stored=False`. This could be:
+        # 1. Effect validation rollback (`Store()` succeeded but was
+        #    rolled back by `context.restore()`).
+        # 2. Real failure (callable threw exception, process crashed,
+        #    etc).
+
+        # If the callable explicitly failed, always raise the error.
+        if status.failed:
+            raise AtMostOnceFailedBeforeCompleting(status.failure)
+
+        # During effect validation, `Store()` from the first execution
+        # gets rolled back by checkpoint/restore in `server.py`,
+        # leaving `started=True` and `stored=False`. This is expected
+        # during validation - we should re-execute rather than raise.
+        # However, if we're NOT in effect validation, this is a real
+        # failure.
+        if context._effect_validation != EffectValidation.ENABLED:
+            raise AtMostOnceFailedBeforeCompleting(
                 '... it looks like an external failure occurred (e.g., '
                 'the machine failed, your container was rescheduled, etc) '
                 'while your code was executing'
             )
-        )
     elif status.stored:
         t = pickle.loads(status.data)
         if type(t) is not type_t:
@@ -124,8 +139,9 @@
         return t
 
     # Only need to call `Start` for "at most once" semantics.
-    if at_most_once:
-        assert not status.started
+    # Skip `Start()` if already started (happens during effect
+    # validation).
+    if at_most_once and not status.started:
         await memoize.always().Start(context)
 
     try:
@@ -236,15 +252,32 @@
         # TODO(benh): retry just this part in the event of retryable
         # errors so that we aren't the cause of raising
         # `AtMostOnceFailedBeforeCompleting`.
-        await memoize.idempotently(
-            f'{idempotency.alias} store',
-            # Don't mangle idempotency alias; any loop iterations
-            # should already be accounted for in `idempotency`.
-            each_iteration=False,
-        ).Store(
-            context,
-            data=pickle.dumps(t),
-        )
+        #
+        # NOTE: Use `.always()` for `at_most_once` to persist `Store()`
+        # through effect validation checkpoint/restore cycles. Without
+        # this, the `Store()` gets rolled back by `context.restore()`,
+        # leaving `started=True` and `stored=False` on the validation
+        # re-run.
+        if at_most_once:
+            # During effect validation re-runs (when `status.started`
+            # was `True`), the `Store()` already happened in the first
+            # execution and persisted. Skip the `Store()` call to avoid
+            # assertion error.
+            if not status.started:
+                await memoize.always().Store(
+                    context,
+                    data=pickle.dumps(t),
+                )
+        else:
+            await memoize.idempotently(
+                f'{idempotency.alias} store',
+                # Don't mangle idempotency alias; any loop iterations
+                # should already be accounted for in `idempotency`.
+                each_iteration=False,
+            ).Store(
+                context,
+                data=pickle.dumps(t),
+            )

[benh]

Will be fixed in next Reboot release.