Merge pull request #488 from goodboy/wkt/moc_teardown_completion

Wait for ctx exit in `maybe_open_context()`

ai/prompt-io/opencode/20260804T030309Z_65bf9df5_prompt_io.md

@@ -0,0 +1,39 @@
+---
+model: openai/gpt-5.6-sol
+service: opencode
+session: moc-teardown-completion-20260804
+timestamp: 2026-08-04T03:03:09Z
+git_ref: 65bf9df5
+scope: code
+substantive: true
+raw_file: 20260804T030309Z_65bf9df5_prompt_io.raw.md
+---
+
+## Prompt
+
+Patch Tractor's `maybe_open_context()` so the final consumer waits for
+resource `__aexit__()` completion and receives cleanup errors. Reuse
+`outcome.Outcome` for the exit result; use the smaller mutable
+`_CtxExit` holder if that makes the implementation simpler. Add and run
+the relevant existing unit tests, but do not commit or push the patch.
+After reviewing the result, simplify `_CtxExit` back to an optional
+exception because the success outcome carries no useful value.
+
+## Response summary
+
+Added an exception-backed completion handshake between
+`_Cache.run_ctx()` and the final `maybe_open_context()` consumer.
+Serialized consumer registration and final teardown under the per-key
+lock, preserving that lock for queued entrants. Added deterministic
+regressions for normal exit, cleanup errors, cancellation interactions,
+service-nursery cancellation, and teardown re-entry.
+
+## Files changed
+
+- `tractor/trionics/_mngrs.py` - publish and unwrap cached exit outcomes.
+- `tests/test_resource_cache.py` - cover completion and cancellation.
+
+## Human edits
+
+The user directed the final simplification from `outcome.Outcome` to an
+optional exception field. The patch remains uncommitted.

ai/prompt-io/opencode/20260804T030309Z_65bf9df5_prompt_io.raw.md

@@ -0,0 +1,36 @@
+---
+model: openai/gpt-5.6-sol
+service: opencode
+timestamp: 2026-08-04T03:03:09Z
+git_ref: 65bf9df5
+diff_cmd: git diff HEAD~1..HEAD
+---
+
+Implemented cached-context exit completion in
+`tractor.trionics.maybe_open_context()`.
+
+> `git diff HEAD~1..HEAD -- tractor/trionics/_mngrs.py`
+
+The generated implementation adds `_CtxExit`, whose `done` event
+publishes an `outcome.Outcome[None]`. `_Cache.run_ctx()` records either
+`Value(None)` or `Error(exc)` after the resource exit attempt. The final
+MOC consumer signals `no_more_users`, waits for completion under a
+shielded cancel scope, removes the per-key lock, and unwraps the outcome
+so ordinary cleanup failures are raised at the consumer boundary.
+`trio.Cancelled`, `KeyboardInterrupt`, and `SystemExit` continue through
+the service task rather than being converted into regular cleanup
+errors.
+
+> `git diff HEAD~1..HEAD -- tests/test_resource_cache.py`
+
+The generated regressions cover successful exit blocking, cleanup-error
+delivery, final-user cancellation, cancellation combined with a cleanup
+error, and service-nursery cancellation. The existing teardown re-entry
+test now uses explicit events instead of a ten-second cleanup sleep and
+asserts that the replacement resource is a fresh cache miss.
+
+Verification:
+
+`env PYTHONPATH="$PWD" /home/goodboy/repos/tractor/py313/bin/python -m pytest tests/test_resource_cache.py`
+
+Result: `16 passed in 8.37s`.

ai/prompt-io/opencode/README.md

@@ -0,0 +1,22 @@
+# AI Prompt I/O Log - OpenCode
+
+This directory tracks prompt inputs and model outputs for AI-assisted
+development using `opencode`.
+
+## Policy
+
+Prompt logging follows the [NLNet generative AI policy][nlnet-ai]. All
+substantive AI contributions are logged with:
+
+- Model name and version
+- Timestamps
+- The prompts that produced the output
+- Unedited model output (`.raw.md` files)
+
+[nlnet-ai]: https://nlnet.nl/foundation/policies/generativeAI/
+
+## Usage
+
+Entries are created by the prompt-io workflow. Human contributors remain
+accountable for all decisions. AI-generated content is never presented as
+human-authored work.

tests/test_resource_cache.py

@@ -9,6 +9,7 @@ from typing import Awaitable
 
 import pytest
 import trio
+from trio.testing import wait_all_tasks_blocked
 import tractor
 from tractor.trionics import (
     maybe_open_context,
@@ -94,6 +95,232 @@ def test_resource_only_entered_once(key_on):
     trio.run(main)
 
 
+def test_last_moc_user_waits_for_resource_exit():
+    '''
+    Verify the final user cannot return before resource teardown.
+
+    Previously the final `maybe_open_context()` user only signalled
+    `_Cache.run_ctx()` through its `no_more_users` event. The user
+    then returned while the service task was still running the
+    resource's `__aexit__()`, so callers could observe stale external
+    state immediately after their `async with` block.
+
+    The resource sets `exit_started` before blocking on
+    `allow_exit`. The user task must remain inside MOC until the test
+    releases that deterministic checkpoint and `__aexit__()` sets
+    `exit_finished`.
+
+    '''
+    async def main():
+        exit_started = trio.Event()
+        allow_exit = trio.Event()
+        exit_finished = trio.Event()
+        user_returned = trio.Event()
+
+        @acm
+        async def open_resource():
+            try:
+                yield
+            finally:
+                exit_started.set()
+                await allow_exit.wait()
+                exit_finished.set()
+
+        async def use_resource():
+            async with maybe_open_context(open_resource):
+                pass
+
+            assert exit_finished.is_set()
+            user_returned.set()
+
+        async with (
+            tractor.open_root_actor(),
+            trio.open_nursery() as tn,
+        ):
+            tn.start_soon(use_resource)
+            await exit_started.wait()
+            assert not user_returned.is_set()
+            allow_exit.set()
+            await user_returned.wait()
+
+    trio.run(main)
+
+
+def test_moc_delivers_resource_exit_error():
+    '''
+    Verify a resource exit error reaches the final MOC user.
+
+    Previously `_Cache.run_ctx()` executed the cached resource's
+    `__aexit__()` after the final user had returned. An exit failure
+    therefore surfaced later through the actor service nursery rather
+    than at the user's `async with maybe_open_context()` boundary.
+
+    This resource raises a unique `ResourceExitError` during exit.
+    Catching that exact instance around MOC proves the service task
+    delivered the failure to the final user without replacing it.
+
+    '''
+    class ResourceExitError(Exception):
+        pass
+
+    exit_error = ResourceExitError('resource exit failed')
+
+    async def main():
+        @acm
+        async def open_resource():
+            yield
+            raise exit_error
+
+        async with tractor.open_root_actor():
+            with pytest.raises(ResourceExitError) as exc_info:
+                async with maybe_open_context(open_resource):
+                    pass
+
+            assert exc_info.value is exit_error
+
+    trio.run(main)
+
+
+def test_moc_final_user_cancellation_waits_for_exit():
+    '''
+    Verify final-user cancellation still waits for successful exit.
+
+    Previously cancellation escaped the final MOC user immediately
+    after it signalled `_Cache.run_ctx()`, leaving resource exit to
+    finish later in the actor service task. This violated the context
+    manager boundary even when cleanup itself succeeded.
+
+    The consumer cancels its own scope while holding the sole cached
+    resource. The resource sets `exit_finished` from its `finally`
+    block, and the consumer checks that event immediately after its
+    cancel scope catches `trio.Cancelled`. This proves MOC's
+    completion wait is shielded without suppressing the original
+    cancellation.
+
+    '''
+    async def main():
+        exit_finished = trio.Event()
+
+        @acm
+        async def open_resource():
+            try:
+                yield
+            finally:
+                exit_finished.set()
+
+        async with tractor.open_root_actor():
+            with trio.CancelScope() as cs:
+                async with maybe_open_context(open_resource):
+                    cs.cancel()
+                    await trio.sleep_forever()
+
+            assert cs.cancelled_caught
+            assert exit_finished.is_set()
+
+    trio.run(main)
+
+
+def test_moc_exit_error_masks_final_user_cancellation():
+    '''
+    Verify cleanup errors survive final-user cancellation.
+
+    A cancelled final user previously signalled `no_more_users` and
+    propagated `trio.Cancelled` before `_Cache.run_ctx()` completed
+    resource exit. If `__aexit__()` then failed, its error was
+    detached from the API call which caused teardown.
+
+    The consumer cancels its own scope at a deterministic checkpoint
+    inside MOC. Resource exit raises `ResourceExitError`; observing
+    that exact error outside the cancel scope proves MOC shields the
+    completion wait and applies normal context-manager masking, where
+    a cleanup failure replaces the active cancellation.
+
+    '''
+    class ResourceExitError(Exception):
+        pass
+
+    exit_error = ResourceExitError('resource exit failed')
+
+    async def main():
+        @acm
+        async def open_resource():
+            yield
+            raise exit_error
+
+        async with tractor.open_root_actor():
+            with pytest.raises(ResourceExitError) as exc_info:
+                with trio.CancelScope() as cs:
+                    async with maybe_open_context(open_resource):
+                        cs.cancel()
+                        await trio.sleep_forever()
+
+            assert exc_info.value is exit_error
+
+    trio.run(main)
+
+
+def test_moc_service_nursery_cancellation_completes_exit():
+    '''
+    Verify service-nursery cancellation cannot strand a final user.
+
+    `_Cache.run_ctx()` and an MOC consumer may share a
+    caller-provided service nursery. Cancelling that nursery
+    interrupts the service task's `no_more_users` wait and the
+    consumer body together. A shielded final-user wait would deadlock
+    if `run_ctx()` failed to publish completion while propagating its
+    own `trio.Cancelled`.
+
+    The outer task waits for resource entry, then cancels the exact
+    nursery containing both tasks. The resource shields one cleanup
+    checkpoint and sets `exit_finished`; observing both it and
+    `service_finished` proves cancellation propagated normally while
+    MOC's completion handshake terminated deterministically.
+
+    '''
+    async def main():
+        resource_entered = trio.Event()
+        exit_finished = trio.Event()
+        service_finished = trio.Event()
+        service_tn: trio.Nursery|None = None
+
+        @acm
+        async def open_resource():
+            try:
+                resource_entered.set()
+                yield
+            finally:
+                with trio.CancelScope(shield=True):
+                    await trio.lowlevel.checkpoint()
+                    exit_finished.set()
+
+        async def use_resource(tn: trio.Nursery):
+            async with maybe_open_context(
+                open_resource,
+                tn=tn,
+            ):
+                await trio.sleep_forever()
+
+        async def run_service():
+            nonlocal service_tn
+
+            async with trio.open_nursery() as tn:
+                service_tn = tn
+                tn.start_soon(use_resource, tn)
+
+            service_finished.set()
+
+        async with trio.open_nursery() as outer_tn:
+            outer_tn.start_soon(run_service)
+            await resource_entered.wait()
+            assert service_tn is not None
+            service_tn.cancel_scope.cancel()
+            await service_finished.wait()
+
+        assert exit_finished.is_set()
+
+    trio.run(main)
+
+
 @tractor.context
 async def streamer(
     ctx: tractor.Context,
@@ -548,43 +775,52 @@ def test_moc_reentry_during_teardown(
     loglevel: str,
 ):
     '''
-    Reproduce the piker `open_cached_client('kraken')` race:
+    Reproduce re-entry while an identical cached context exits.
 
-    - same `acm_func`, NO kwargs (identical `ctx_key`)
-    - multiple tasks share the cached resource
-    - all users exit -> teardown starts
-    - a NEW task enters during `_Cache.run_ctx.__aexit__`
-    - `values[ctx_key]` is gone (popped in inner finally)
-      but `resources[ctx_key]` still exists (outer finally
-      hasn't run yet bc the acm cleanup has checkpoints)
-    - old code: `assert not resources.get(ctx_key)` FIRES
+    - multiple tasks use the same `acm_func` with no kwargs,
+      producing an identical `ctx_key`;
+    - all users leave and the final user starts resource teardown;
+    - `_Cache.run_ctx()` removes the cached value and resource entry
+      before entering the resource's blocking `__aexit__()` body;
+    - a new task attempts to enter that same `ctx_key` during exit;
+    - the per-key lock keeps that entrant queued until exit
+      completes;
+    - the entrant then receives a fresh cache miss and resource.
 
-    This models the real-world scenario where `brokerd.kraken`
-    tasks concurrently call `open_cached_client('kraken')`
-    (same `acm_func`, empty kwargs, shared `ctx_key`) and
-    the teardown/re-entry race triggers intermittently.
+    Without teardown sharing the registration lock, re-entry could
+    race resource replacement while the prior generation was still
+    exiting. The final user could also return before that exit
+    completed.
+
+    The first resource generation signals `in_aexit` and waits on
+    `allow_aexit`. The re-entry task signals `reentry_started` and
+    blocks inside MOC; only after `wait_all_tasks_blocked()` confirms
+    that ordering does the coordinator release cleanup. The entrant
+    must then receive a fresh cache miss. `first_done` additionally
+    proves the first MOC user observed completed teardown before
+    returning.
 
     '''
     async def main():
         in_aexit = trio.Event()
+        allow_aexit = trio.Event()
+        reentry_started = trio.Event()
+        generation: int = 0
 
         @acm
         async def cached_client():
             '''
-            Simulates `kraken.api.get_client()`:
-            - no params (all callers share one `ctx_key`)
-            - slow-ish cleanup to widen the race window
-              between `values.pop()` and `resources.pop()`
-              inside `_Cache.run_ctx`.
+            Simulate a no-argument `kraken.api.get_client()`.
 
             '''
+            nonlocal generation
+
+            generation += 1
+            resource_generation: int = generation
             yield 'the-client'
-            # Signal that we're in __aexit__ — at this
-            # point `values` has already been popped by
-            # `run_ctx`'s inner finally, but `resources`
-            # is still alive (outer finally hasn't run).
-            in_aexit.set()
-            await trio.sleep(10)
+            if resource_generation == 1:
+                in_aexit.set()
+                await allow_aexit.wait()
 
         first_done = trio.Event()
 
@@ -598,16 +834,25 @@ def test_moc_reentry_during_teardown(
         async def reenter_during_teardown():
             '''
             Wait for the acm's `__aexit__` to start (meaning
-            `values` is popped but `resources` still exists),
-            then re-enter — triggering the assert.
+            the cached value is no longer available), then re-enter.
 
             '''
             await in_aexit.wait()
+
+            # Tell the coordinator this task is about to enter MOC.
+            # `Event.set()` is not a checkpoint. Though `async with`
+            # awaits MOC's `__aenter__()`, its async generator runs
+            # synchronously until the held per-key `lock.acquire()`
+            # actually suspends this task.
+            reentry_started.set()
             async with maybe_open_context(
                 cached_client,
             ) as (cache_hit, value):
+                assert not cache_hit
                 assert value == 'the-client'
 
+            await first_done.wait()
+
         with trio.fail_after(5):
             async with (
                 tractor.open_root_actor(
@@ -619,5 +864,15 @@ def test_moc_reentry_during_teardown(
             ):
                 tn.start_soon(use_and_exit)
                 tn.start_soon(reenter_during_teardown)
+                await reentry_started.wait()
+
+                # Wait until the re-entry task is queued on MOC's
+                # per-key lock while `_Cache.run_ctx()` remains
+                # blocked in the first generation's `__aexit__()`.
+                # Only then release cleanup, making the intended
+                # enter-during-sibling-exit ordering deterministic.
+                await wait_all_tasks_blocked()
+                assert not first_done.is_set()
+                allow_aexit.set()
 
     trio.run(main)

tractor/trionics/_mngrs.py

@@ -198,6 +198,16 @@ async def gather_contexts(
 # Further potential examples of interest:
 # https://gist.github.com/njsmith/cf6fc0a97f53865f2c671659c88c1798#file-cache-py-L8
 
+class _CtxExit:
+    '''
+    Completion state for a cached context's shared exit.
+
+    '''
+    def __init__(self) -> None:
+        self.done = trio.Event()
+        self.error: Exception|None = None
+
+
 class _Cache:
     '''
     Globally (actor-processs scoped) cached, task access to
@@ -213,7 +223,11 @@ class _Cache:
     values: dict[Any,  Any] = {}
     resources: dict[
         Hashable,
-        tuple[trio.Nursery, trio.Event]
+        tuple[
+            trio.Nursery,
+            trio.Event,
+            _CtxExit,
+        ],
     ] = {}
     # nurseries: dict[int, trio.Nursery] = {}
     no_more_users: trio.Event|None = None
@@ -223,18 +237,38 @@ class _Cache:
         cls,
         mng,
         ctx_key: tuple,
+        ctx_exit: _CtxExit,
         task_status: trio.TaskStatus[T] = trio.TASK_STATUS_IGNORED,
 
     ) -> None:
-        async with mng as value:
-            _, no_more_users = cls.resources[ctx_key]
-            cls.values[ctx_key] = value
-            task_status.started(value)
-            try:
-                await no_more_users.wait()
-            finally:
-                value = cls.values.pop(ctx_key)
-                cls.resources.pop(ctx_key)
+        entered: bool = False
+        try:
+            async with mng as value:
+                entered = True
+                (
+                    _,
+                    no_more_users,
+                    _,
+                ) = cls.resources[ctx_key]
+                cls.values[ctx_key] = value
+                task_status.started(value)
+                try:
+                    await no_more_users.wait()
+                finally:
+                    cls.values.pop(ctx_key)
+                    cls.resources.pop(ctx_key)
+
+        except Exception as exc:
+            if not entered:
+                raise
+
+            # Deliver regular `__aexit__()` failures to the final
+            # consumer instead of raising into the service nursery.
+            ctx_exit.error = exc
+
+        finally:
+            if entered:
+                ctx_exit.done.set()
 
 
 class _UnresolvedCtx:
@@ -281,9 +315,10 @@ async def maybe_open_context(
         )
 
     # yielded output
-    # sentinel = object()
     yielded: Any = _UnresolvedCtx
     user_registered: bool = False
+    ctx_exit: _CtxExit|None = None
+    exit_error: Exception|None = None
 
     # Lock resource acquisition around task racing  / ``trio``'s
     # scheduler protocol.
@@ -300,7 +335,6 @@ async def maybe_open_context(
         ] = trio.StrictFIFOLock()
         header: str = 'Allocated NEW lock for @acm_func,\n'
     else:
-        await trio.lowlevel.checkpoint()
         header: str = 'Reusing OLD lock for @acm_func,\n'
 
     log.debug(
@@ -368,7 +402,6 @@ async def maybe_open_context(
             resources = _Cache.resources
             entry: tuple|None = resources.get(ctx_key)
             if entry:
-                service_tn, ev = entry
                 raise RuntimeError(
                     f'Caching resources ALREADY exist?!\n'
                     f'ctx_key={ctx_key!r}\n'
@@ -376,12 +409,18 @@ async def maybe_open_context(
                     f'task: {task}\n'
                 )
 
-            resources[ctx_key] = (service_tn, trio.Event())
+            ctx_exit = _CtxExit()
+            resources[ctx_key] = (
+                service_tn,
+                trio.Event(),
+                ctx_exit,
+            )
             try:
                 yielded: Any = await service_tn.start(
                     _Cache.run_ctx,
                     mngr,
                     ctx_key,
+                    ctx_exit,
                 )
             except BaseException:
                 # If `run_ctx` (wrapping the acm's `__aenter__`)
@@ -427,6 +466,11 @@ async def maybe_open_context(
             raise taskc
     else:
         # XXX, cached-entry-path
+        (
+            _,
+            _,
+            ctx_exit,
+        ) = _Cache.resources[ctx_key]
         _Cache.users[ctx_key] += 1
         user_registered = True
         log.debug(
@@ -445,44 +489,57 @@ async def maybe_open_context(
         )
 
     finally:
-        if lock.locked():
-            stats: trio.LockStatistics = lock.statistics()
-            owner: trio.Task|None = stats.owner
-            log.error(
-                f'Lock never released by last owner={owner!r} !?\n'
-                f'{stats}\n'
-                f'\n'
-                f'task={task!r}\n'
-                f'ctx_key={ctx_key!r}\n'
-                f'acm_func={acm_func}\n'
-
-            )
-
         if user_registered:
-            _Cache.users[ctx_key] -= 1
+            # Serialize user registration and teardown under the same
+            # per-key lock so no entrant can acquire a resource after
+            # its final user has committed to exiting it.
+            with trio.CancelScope(shield=True):
+                await lock.acquire()
+                try:
+                    _Cache.users[ctx_key] -= 1
 
-        if yielded is not _UnresolvedCtx:
-            # if no more consumers, teardown the client
-            if _Cache.users[ctx_key] <= 0:
-                log.debug(
-                    f'De-allocating @acm-func entry\n'
-                    f'ctx_key={ctx_key!r}\n'
-                    f'acm_func={acm_func!r}\n'
-                )
+                    # If no consumers remain, keep entrants queued
+                    # until the cached context has completely exited.
+                    if _Cache.users[ctx_key] <= 0:
+                        log.debug(
+                            f'De-allocating @acm-func entry\n'
+                            f'ctx_key={ctx_key!r}\n'
+                            f'acm_func={acm_func!r}\n'
+                        )
 
-                # XXX: if we're cancelled we the entry may have never
-                # been entered since the nursery task was killed.
-                # _, no_more_users = _Cache.resources[ctx_key]
-                entry = _Cache.resources.get(ctx_key)
-                if entry:
-                    _, no_more_users = entry
-                    no_more_users.set()
+                        # XXX: if we're cancelled, the entry may
+                        # have never been entered since the nursery
+                        # task was killed.
+                        entry = _Cache.resources.get(ctx_key)
+                        if entry:
+                            (
+                                _,
+                                no_more_users,
+                                ctx_exit,
+                            ) = entry
+                            no_more_users.set()
 
-                maybe_lock = _Cache.locks.pop(
-                    ctx_key,
-                    None,
-                )
-                if maybe_lock is None:
-                    log.error(
-                        f'Resource lock for {ctx_key} ALREADY POPPED?'
-                    )
+                        assert ctx_exit is not None
+                        await ctx_exit.done.wait()
+                        exit_error = ctx_exit.error
+
+                        # A queued entrant already holds a reference
+                        # to this lock. Keep it registered until that
+                        # task has acquired and released it.
+                        stats = lock.statistics()
+                        if not stats.tasks_waiting:
+                            maybe_lock = _Cache.locks.get(ctx_key)
+                            if maybe_lock is lock:
+                                _Cache.locks.pop(ctx_key)
+                            else:
+                                log.error(
+                                    f'Resource lock for {ctx_key} '
+                                    f'was replaced before teardown?'
+                                )
+                finally:
+                    lock.release()
+
+        if exit_error is not None:
+            # Always re-raise a regular `__aexit__()` error at the
+            # final consumer's context boundary.
+            raise exit_error