Picklable wool.Token: enable cross-wire token transport with stdlib parity

[conradbzura]

Description

Reintroduce wool.Token as a concrete wrapper class (currently aliased to contextvars.Token per the per-frame-context architecture) so that tokens minted by wool.ContextVar.set() cross the wire when passed as routine args, kwargs, or return values. Receiver-side reset uses stdlib's own tok_ctx / tok_used machinery for identity validation via an internal anchor ContextVar, preserving stdlib's reset semantics verbatim (including the canonical ValueError error strings) without introducing any Wool-level chain-id parity check.

The full design is captured in docs/design/picklable-tokens.md and reproduced below. Depends on the per-frame-context architecture (chain-registry as load-bearing primitive), which is part of PR #224's design materials and lands as docs/design/per-frame-context.md.

Mechanism summary

• wool.Token wrapper carrying _var, _old_value, _native: contextvars.Token | None, _id, _chain_id, _used. Locally-minted tokens have _native bound to the var's backing; wire-reconstituted tokens have _native bound to a Wool-internal anchor ContextVar.
• Sender-side _live_tokens: ContextVar[WeakSet[contextvars.Token]] with copy-on-write on set() (not on reset — the token's _used flag carries the consumed state). Membership at pickle time determines the wire do_mount flag.
• Anchor-token mechanism on receiver: wire reconstitution mints _token_anchor.set(token_id) inside the receiving task's stdlib.Context. The resulting stdlib token's tok_ctx is the receiver's Context; reset validates via _token_anchor.reset(native) (stdlib raises with canonical error strings on cross-Context / already-used).
• Reintroduce external_used as cross-process consumed-token tracking. Ship the full set per frame in v1; prune-on-re-set bounds growth structurally; diff encoding deferred to the broader wire-diff future extension.
• Pickle dispatch: __reduce__ raises TypeError; __wool_reduce__ is the real reducer. Matches existing wool.ContextVar pattern.

Stdlib parity model

Tokens in stdlib have three lifecycle states relevant to transport:

State	tok_used	tok_ctx validity	var.reset(token) outcome
Fresh, live	False	Valid (the Context that minted it)	Restores tok_oldval; sets tok_used=True
Consumed	True	Was-valid	Raises ValueError("<Token> has already been used once")
Orphan	False	Context is gone / different from current	Raises ValueError("<Token> was created in a different Context")

A token in any of these states can be passed to any function — stdlib's var.reset just fails appropriately on attempt.

The picklable-token wire format preserves all three states. The wire format ships enough metadata for the receiver to reconstruct a wool.Token in the equivalent state, and the receiver's wool.ContextVar.reset matches stdlib's check order (used first, then context validity).

Failure modes and stdlib-parallel equivalents

Pattern	In-process behavior	Wool behavior
Reset a live token, same Context	OK	OK
Reset a used token	ValueError("<Token> has already been used once")	Same error message
Reset an orphan token (different Context)	ValueError("<Token> was created in a different Context")	Same error message
Reset a token across a create_task fork	Stdlib raises "different Context" (the fork's Context is a copy, but tok_ctx is the parent)	Wool raises "different Context" — the fork lands in a different chain id, different cached Context, do_mount=False on the wire token
Reset a token across processes (caller → worker)	n/a	Worker resets in the cached Context for the token's chain_id
Same token reset twice across processes	n/a	Second reset raises "already used" via external_used
Worker creates token, returns it to caller; caller resets it	n/a	On wire reconstitution, the token's anchor is minted in whichever caller task received the response. Reset works only from that same task (stdlib tok_ctx semantics)

Detailed design choices

1. Identity validation for wire-reconstituted tokens: anchor-token mechanism — wire-reconstituted tokens mint a stdlib token in _token_anchor at reconstitution. Reset validates via _token_anchor.reset(native), then writes _old_value to var directly. No Wool-level chain_id check; tokens are bound to the task that produced/received them, matching stdlib's tok_ctx semantic.
2. wool.ContextVar.reset accepts only wool.Token — rejects stdlib contextvars.Token with TypeError. Mixing would bypass Wool's wire-aware bookkeeping (_live_tokens, external_used).
3. _live_tokens long-running cost: ignore for v1 — WeakSet auto-cleanup handles the common case; per-var.set CoW cost is O(n) where typical n ∈ [0, 5]. Future mitigation if profiling shows it matters: soft-bound prune at dispatch time.
4. external_used propagation: full set per frame, with prune-on-re-set — full set per frame in v1; finding C26's prune-on-re-set (var re-set() prunes stale external_used entries for that var_key) bounds growth structurally. Diff encoding deferred to the broader wire-diff extension.
5. Pickle dispatch: __reduce__ raises, __wool_reduce__ is the real reducer — vanilla pickle.dumps(token) raises TypeError; Wool's CloudpickleSerializer dispatches via __wool_reduce__. Matches existing wool.ContextVar pattern.

The full design including code sketches lives in docs/design/picklable-tokens.md.

Motivation

Without this feature, the following pattern fails:

```python
@wool.routine
async def undo_after(token):
await asyncio.sleep(1)
wool.var.reset(token)

token = wool.var.set("a")
await undo_after(token) # would fail: token's tok_ctx is on the caller process
```

Tokens become functionally untransportable — users either avoid passing them through @routine boundaries (forcing awkward refactors to value-only protocols) or hit confusing reset failures with no Wool-aware error message. The pattern is stdlib-natural; Wool's "distributed-but-stdlib-equivalent" tenet means it should just work.

The per-frame-context architecture (PR #224's design) makes this implementable cleanly for the first time: the per-chain stdlib.Context registry on the worker is the substrate the anchor-token mechanism builds on. Without that foundation, wire-token reconstitution had no stable receiver-side Context to bind to — which is why commit 7bf3ae3 dropped Token wire-serialization in the first place.

Expected outcome

When this is done:

• wool.ContextVar.set() returns a wool.Token instance (concrete wrapper, not the stdlib alias).
• Tokens passed as routine args, kwargs, or return values pickle cleanly through wool.__serializer__ and reconstruct on the receiver as reset-able tokens.
• Vanilla pickle.dumps(token) raises TypeError with a message pointing at Wool's serializer.
• wool.ContextVar.reset(token) accepts only wool.Token; passing contextvars.Token raises TypeError.
• All three stdlib reset failure modes (live / consumed / orphan) reproduce with stdlib's canonical error strings.
• Cross-process double-reset of the same token raises "<Token> has already been used once" on the second process to try.
• The stdlib_parity test suite gains coverage for token-across-wire patterns — @routine versions of stdlib Token tests assert byte-for-byte error parity.
• wool.runtime.context.token is reinstated as a real module (currently deleted under the per-frame-context design).
• external_used returns to the wire schema with prune-on-re-set semantics.

References

• Depends on docs/design/per-frame-context.md (PR Reframe Wool's context model around stdlib contextvars — Closes #223, #229, #232 #224 design materials) — chain-registry primitive.
• Full design: docs/design/picklable-tokens.md.
• REVIEW-PLAN-224-3.md finding C26 — external_used prune-on-re-set.
• Commit 7bf3ae3 — drop Token wire-serialization (this design's reintroduction).
• CPython Lib/contextvars.py and Modules/_contextvarsmodule.c — stdlib Token semantics reference.
• PEP 567 — Context Variables.