From 7ec56e01c46a09d5e4d1fb7ca632c877d80b4a28 Mon Sep 17 00:00:00 2001 From: yibie Date: Sat, 6 Jun 2026 16:00:46 +0800 Subject: [PATCH] fix: ContextVar cross-Context token error in ForkTapeStore.fork() Replace ContextVar.reset(token) with save/restore pattern to avoid 'Token was created in a different Context' ValueError when the async generator cleanup runs in a different asyncio Task (e.g. cancellation, TaskGroup exception propagation). --- src/bub/builtin/store.py | 18 ++++++++++++------ 1 file changed, 12 insertions(+), 6 deletions(-) diff --git a/src/bub/builtin/store.py b/src/bub/builtin/store.py index 714cfc88..b30dd751 100644 --- a/src/bub/builtin/store.py +++ b/src/bub/builtin/store.py @@ -101,16 +101,22 @@ async def append(self, tape: str, entry: TapeEntry) -> None: @contextlib.asynccontextmanager async def fork(self, tape: str, merge_back: bool = True) -> AsyncGenerator[None, None]: store = InMemoryTapeStore() - token = current_store.set(store) - tape_token = current_fork_tape.set(tape) - reset_token = current_tape_was_reset.set(False) + # Save/restore instead of ContextVar.reset(token) to avoid + # "Token was created in a different Context" when cleanup + # runs in a different asyncio Task (e.g. cancellation, TaskGroup). + prev_store = current_store.get(_empty_store) + prev_fork_tape = current_fork_tape.get() + prev_was_reset = current_tape_was_reset.get() + current_store.set(store) + current_fork_tape.set(tape) + current_tape_was_reset.set(False) try: yield finally: was_reset = current_tape_was_reset.get() - current_store.reset(token) - current_fork_tape.reset(tape_token) - current_tape_was_reset.reset(reset_token) + current_store.set(prev_store) + current_fork_tape.set(prev_fork_tape) + current_tape_was_reset.set(prev_was_reset) if merge_back: if was_reset: await self._parent.reset(tape)