Reframe Wool's context model around stdlib contextvars — Closes #223, #229, #232

[conradbzura]

Summary

Re-found Wool's context model directly on Python's stdlib contextvars. There is one context system, not two: Wool chain state rides in a single Wool-owned stdlib contextvars.ContextVar (wool.__chain__) as an immutable Chain dataclass, so it propagates with stdlib visibility across every conformant event loop (including uvloop) and across every cooperative asyncio scheduling edge — call_soon / call_later / call_at, add_reader / add_writer / add_signal_handler, Future.add_done_callback. Child task creation propagates as a fork — the child inherits the parent's variable values on a freshly minted chain id, not the parent's chain itself. OS-thread offload is the one deliberate exception: an armed chain entered from a second OS thread fails loud rather than propagating silently. The chain-contention guard arms only when a wool.ContextVar has been set; a process that touches neither wool.ContextVar nor the dispatch path pays nothing.

Reshape the public API surface. Remove wool.Context, wool.current_context, wool.copy_context, wool.create_task, and wool.ContextAlreadyBound. Alias wool.Token to stdlib contextvars.Token (wool.Token is contextvars.Token); user code reuses the stdlib implementation directly and the previous home-grown Token class is deleted. Add wool.to_thread, wool.install_task_factory, and an exception hierarchy rooted at two umbrellas — wool.WoolError and wool.WoolWarning — with wool.ChainContention, wool.ContextVarCollision, wool.TaskFactoryDisplaced, wool.SerializationError, wool.ChainSerializationError, and wool.SerializationWarning underneath. wool.ContextVar and wool.RuntimeContext carry over.

Introduce a Frame/ChainManifest wire-frame abstraction in wool.runtime.worker.frame: every request and response on the worker wire wraps in a Frame subclass (TaskRequestFrame, SendRequestFrame, ThrowRequestFrame, NextRequestFrame, ResultResponseFrame, ExceptionResponseFrame, AckResponseFrame, NackResponseFrame) that carries its own decoded sender-side state alongside the payload. Receive sites mount that state into the active chain as a single canonical step, replacing the ad-hoc wire-context plumbing previously scattered across worker/session.py, worker/connection.py, and worker/service.py.

This is a breaking change to the public API; backward compatibility is not preserved.

Closes #223, #229, #232

The diagram below traces one wool.ContextVar value across a cross-process routine dispatch — arm, encode, apply, mutate, and apply back:

sequenceDiagram
    participant Caller
    participant Connection
    participant DispatchSession
    participant Routine

    rect rgb(0, 0, 0, 0)
        Note over Caller,Routine: Arm & dispatch
        Caller->>Caller: wool.ContextVar.set(v) — mint chain id, arm the Chain
        Caller->>Connection: dispatch routine
        Connection->>DispatchSession: TaskRequestFrame — pure dispatch metadata, no chain manifest
        Connection->>Connection: NextRequestFrame.for_send — auto-capture the armed chain into a ChainManifest
        Connection->>DispatchSession: NextRequestFrame
    end

    rect rgb(0, 0, 0, 0)
        Note over Caller,Routine: Worker apply & run
        DispatchSession->>DispatchSession: Frame.from_protobuf — decode, then mount the ChainManifest into the worker's cached context
        DispatchSession->>Routine: drive (await / asend)
        Routine->>Routine: wool.ContextVar.get() reads the caller's v
        Routine->>Routine: wool.ContextVar.set(v') mutates on the chain
    end

    rect rgb(0, 0, 0, 0)
        Note over Caller,Routine: Apply back
        Routine-->>DispatchSession: result
        DispatchSession->>DispatchSession: ResultResponseFrame.for_send — auto-capture the worker chain into a ChainManifest
        DispatchSession-->>Connection: ResultResponseFrame
        Connection->>Connection: Frame.from_protobuf — caller adopts worker mutations
        Connection-->>Caller: result — wool.ContextVar now reads v'
    end

    Note over Caller,Routine: The worker caches one contextvars.Context per chain id — async-generator routines repeat apply/run and apply-back once per __anext__ / asend via NextRequestFrame / SendRequestFrame / ResultResponseFrame, reusing that cached context across yields

Loading

Proposed changes

Chain-backed storage

Wool chain state — the bound-wool.ContextVar index, the logical-chain UUID, the owner-thread / owning-task ownership stamps, the reset-variable signals, and the stub-pin set — rides in a single immutable Chain dataclass (runtime/context/chain.py, frozen, identity-equality) held in one Wool-owned stdlib contextvars.ContextVar (wool.__chain__). The Chain is an index, not a value store: each wool.ContextVar's value lives in its own backing contextvars.ContextVar, and the chain's vars field records which variables are bound. Because the values live in the stdlib contextvars.Context rather than in the chain, callback write-isolation rests on native contextvars copy-on-write: a callback runs in a copy_context() copy, so its backing-variable writes (and the new chain it installs) stay in that copy.

Chain.mount is the single owner-stamping and arming boundary: it re-stamps the owning thread / task from the calling scope, rewinds reset signals, sets wool.__chain__, and self-installs Wool's task factory on the running loop (idempotent per loop via a WeakSet). Wire ingress goes through ChainManifest.mount(*, stamp_owner, install_factory, merge_with=None): it drains the decoded values into their backing variables, then either seeds a fresh chain from the manifest alone (merge_with=None) or unions the manifest onto a live receiver chain that keeps its chain id (merge_with set), and delegates the arming to Chain.mount. Incoming wire state wins on overlap; reset-only keys drop out so the reset propagates.

Reorganise the context module around the new model. The live Chain lives in chain.py; the stdlib-compatible ContextVar (backed by a per-var contextvars.ContextVar) in var.py; the chain-ownership guard in guard.py; a shared, cycle-free ContextVar registry in registry.py; install_task_factory and the task-factory composition path in factory.py; the stdlib-aligned to_thread in threading.py; and the RuntimeContext overlay for dispatch-time fields in runtime.py. Deferred wire-side state is decoupled from the live chain: ChainManifest (the decoded-but-unmounted snapshot) lives in manifest.py, and the context error hierarchy is centralised in runtime/context/exceptions.py.

Reshape the public API and root the exception hierarchy

Remove wool.Context, wool.current_context, wool.copy_context, wool.create_task, and wool.ContextAlreadyBound. Add wool.to_thread and wool.install_task_factory. wool.ContextVar and wool.RuntimeContext carry over unchanged.

Root the exception surface at two umbrellas: wool.WoolError(Exception) and wool.WoolWarning(Warning), both in src/wool/exceptions.py. Under them: wool.SerializationError, wool.ChainSerializationError (a SerializationError), wool.SerializationWarning, wool.ChainContention, wool.ContextVarCollision, and wool.TaskFactoryDisplaced, all defined in runtime/context/exceptions.py. The earlier ContextSerializationError / ContextDecodeError / ContextDecodeWarning names are removed outright — no back-compat aliases.

Alias wool.Token to stdlib contextvars.Token — wool.Token is contextvars.Token, with no Wool subclass or shim. A token is therefore a plain stdlib token: it carries no Wool-internal state and does not transport across the wire (vanilla pickle / cloudpickle / copy raise TypeError), so value-state propagation rides the chain manifest rather than the token. The dispatch-serialisation machinery lives on wool.ContextVar (guarded against vanilla pickling), not on the token.

Armed-gating, copy-on-fork, and the chain-contention guard

Enforcement is armed-gated. While unset, wool.__chain__ defaults to None; the surrounding context is unarmed and behaves as a plain contextvars.Context — no chain UUID and no guard. A context is armed by either of two events: the first wool.ContextVar.set() on it, or a ChainManifest.mount(...) applying incoming wire state.

The chain-contention guard enforces one invariant — a single chain is never run by two runners at once — across both trigger paths:

• No shared context across concurrent tasks. The task factory forks every child task onto a freshly minted chain UUID (copy-on-fork), so sibling tasks created the ordinary way never share a chain. The guard bites when a contextvars.Context is explicitly reused: handing a context that is already driving a live task to a second create_task raises ChainContention. Once the first task completes, that context is free to drive another.
• No cross-thread entry. Entering an armed chain from a thread other than its owner raises ChainContention — so a plain asyncio.to_thread from an armed context fails loud, and wool.to_thread (a detached fork onto a fresh chain) is the supported alternative.

ChainContention carries structured diagnostic fields (chain_id, kind: Literal["thread", "task"], plus the offending thread id or task) and interpolates them into the error message so a stack trace surfaces who collided with whom. wool.TaskFactoryDisplaced is raised when a third-party factory installed after Wool's silently drops copy-on-fork for subsequently-created tasks; displacement is detected at three converging points — the factory's weakref.finalize, the per-task done-callback backstop, and the set-time factory inspection.

Cooperatively-scheduled work is deliberately not arbitrated. call_soon / call_later / call_at, add_reader / add_writer / add_signal_handler, and Future.add_done_callback callbacks inherit the scheduling scope's context via copy_context() and run on its chain UUID unchanged. They execute synchronously on the owning thread and never overlap with their scheduler, so sharing the chain is safe.

install_task_factory composes with an existing factory via a __wool_inner__ chain it stamps on its wrapper. The composition-chain walk detects Wool buried under a later third-party install, so wool → third-party → wool is recognised and skipped (memoised per layer). A module-level threading.Lock guards both the install read-modify-write and the per-task registry against the (narrow) cross-thread install race.

Frame/ChainManifest wire-frame abstraction

Introduce a Frame hierarchy in wool.runtime.worker.frame that wraps every request and response on the worker wire:

• RequestFrame: TaskRequestFrame, NextRequestFrame, SendRequestFrame, ThrowRequestFrame
• ResponseFrame: ResultResponseFrame, ExceptionResponseFrame, AckResponseFrame, NackResponseFrame

Mid-stream frames mix in a _ChainManifestFrame base that captures or accepts a manifest; boundary frames (TaskRequestFrame, AckResponseFrame, NackResponseFrame) carry none. Frame.chain_manifest is a union — ChainManifest | ChainSerializationError | None — read as "the decoded manifest, the reason it couldn't be parsed, or absent." Frame.from_protobuf decodes the optional wire context field through ChainManifest.from_protobuf and, on a strict-mode decode failure, captures the ChainSerializationError as the union's error arm rather than raising at decode time. The failure is deferred to Frame.mount, which raises it for a value frame or, when the frame's payload is itself a BaseException, walks the payload's __context__ chain and attaches the decode error at the tail — so a malformed manifest never preempts the routine's primary signal. Each frame owns its own (de)serialisation, giving senders and receivers a single canonical place to encode and decode chain state.

ChainManifest.from_protobuf is the pure decoder: it resolves variable identities and deserialises values without ever touching a contextvars.Context, emitting SerializationWarning per bad entry and aggregating them into one ChainSerializationError under strict mode.

Reorganise the protocol package so the codegen has a clean home. Extract the sys.path hack, the wire_pb2 / wire_pb2_grpc import block (wrapped as ProtobufImportError), and the AddServicerToServerProtocol plus add_to_server registry into _wire.py; the package __init__ re-exports from it. The wire message is ChainManifest (value-only ContextVar entries), riding on every Request / Response as the optional context field.

Per-frame context architecture on the worker

Route every request and response on the worker wire through the new RequestFrame and ResponseFrame types. Stop threading wire-context application through ad-hoc helpers in session.py, service.py, and connection.py; let each frame carry its own manifest and mount it into the active chain at receive time as a single canonical step. Build outbound traffic in connection.py from RequestFrame subclasses and consume ResponseFrame subclasses on _read_next; emit the service.py terminal responses (result, exception, ack, nack) through ResponseFrame; reframe session.py's request and response paths around the Frame hierarchy. The chain rides only on the mid-stream payload frames (Next / Send / Throw / Result / Exception), so the first drive frame — not the dispatch metadata — propagates the armed chain.

The dispatch driver caches the worker-side step context per wool chain id in a WeakValueDictionary[UUID, contextvars.Context]. Frames sharing a chain id reuse one contextvars.Context, so a routine's set before an async-generator yield and its reset after land in the same context — stdlib async-generator parity, since asyncio drives every resumption of one generator in a single context. Frames carrying distinct chain ids resolve to distinct contexts and stay isolated. An unarmed dispatch — one whose caller propagates no chain — runs its frames in a single working context left unarmed, so a stateless routine's plain asyncio.to_thread offload copies a bare context and never trips the chain-contention guard; if the routine's own set arms that working context, the driver indexes it under the freshly minted chain id, so the back-propagated next frame (which now carries that id) reuses it.

Routine, discovery, loadbalancer

Update the @routine decorator and the underlying Task wrapper to read dispatch-time fields off RuntimeContext rather than the removed Context aliases, and surface RuntimeContext on the Task itself so call sites can inspect what was actually dispatched. Switch discovery/local.py and the matching discovery test to the wire-namespace import (from wool import protocol) so they refer to the protobuf WorkerMetadata without colliding with the wool dataclass of the same name. Update the loadbalancer exception-handling docstrings to reference the current ChainSerializationError shape rather than the removed decode-alias names.

Build, documentation, and test coverage

Add uvloop as a dev dependency and register the stdlib_parity pytest marker so the new parity suite runs under both the default asyncio loop and uvloop. Measure coverage from spawned worker processes (concurrency = multiprocessing,thread, sigterm flush) so worker-side dispatch code is visible to the gate. Rewrite the context subsystem README, the worker README, and the package README to match the shipped surface — the stdlib-aligned public API, the ChainManifest.mount wire-apply pattern, and the Frame-based dispatch vocabulary.

Test coverage spans three families: unit coverage of the new context module in tests/runtime/context/ (test_chain.py, test_manifest.py, test_runtime.py, test_var.py, test_token.py, test_factory.py, test_guard.py, test_exceptions.py); a stdlib_parity suite in tests/stdlib_parity/ (test_context_parity.py, test_task_creation.py, test_executor_offload.py, test_loop_callbacks.py) pinning Wool's behaviour against the analogous stdlib contextvars / asyncio surfaces under both loops; and Frame round-trip coverage in tests/runtime/worker/test_frame.py, rewritten worker tests across tests/runtime/worker/, plus end-to-end coverage in tests/integration/test_context_var_propagation.py. The unit suite holds the 98% coverage gate: reachable defensive branches are tested directly, and the genuinely-unreachable guards are marked # pragma: no cover.

Test cases

#	Test Suite	Given	When	Then	Coverage Target
1	TestContextVar	A ContextVar with and without a constructor default	get is called on an unarmed then armed context	Returns the fallback ladder — supplied default, constructor default, then LookupError	ContextVar.get, armed-gating
2	TestContextVar	An armed context	set then reset are called	Value rebinds then restores; the first set mints the chain and returns a single-use Token	ContextVar.set / reset
3	TestContextVar	A Token from a different chain or var, or an already-used token	reset is called with it	Stdlib contextvars.Token.reset raises ValueError	ContextVar.reset rejection via stdlib delegation
4	TestContextVar	Two distinct declarations of one (namespace, name) key	The second is constructed	Raises ContextVarCollision	Identity model
5	TestToken	A wool.Token after a first and a nested set	old_value is read, then the token is pickled, cloudpickled, and copied	old_value reports MISSING then the prior value; vanilla pickle, cloudpickle, and copy are rejected with TypeError	Token old-value semantics and serialization rejection
6	TestToken	A wool.Token (alias of contextvars.Token)	Its repr is inspected	References the backing contextvars.ContextVar	Stdlib Token alias surface
7	TestChain	A Chain dataclass with chain id, vars, reset signals, stub pins	Field access and mount field-replacement are exercised	Fields are immutable, equality is identity-based, and mount returns an armed copy	Chain dataclass shape
8	TestChain	A chain with bound and reset variables	to_protobuf emits the wire manifest	One entry per bound var, reset-var entries, and unserializable values skipped with a SerializationWarning (aggregated under strict mode)	Chain wire emission
9	TestChainManifest	A wire ChainManifest with values, resets, and an undeclared var	ChainManifest.from_protobuf decodes it	Round-trips values and resets, registers a stub for the undeclared var, and raises ChainSerializationError on strict-mode failure	Manifest decode
10	TestChainMount	A decoded manifest and a live receiver chain	ChainManifest.mount(merge_with=...) is applied	Drains values into the backings, restamps the owner, unions onto the receiver with incoming winning on overlap, propagates resets, and arms the result	Wire-apply / mount
11	TestRuntimeContext	A RuntimeContext entered with a dispatch_timeout override	__enter__ / __exit__ and a protobuf round-trip are exercised	The overlay applies on enter and restores on exit; the round-trip preserves an explicit timeout and elides an unset default; double-enter raises	RuntimeContext lifecycle and codec
12	test_factory.py	An armed parent task with the factory installed	A child task is created	The child inherits the parent's values on a freshly minted chain UUID	Copy-on-fork
13	test_factory.py	A loop with Wool's factory buried under a later third-party install, or no running loop, or a finalizer firing on an idle loop	install_task_factory is re-invoked / called outside a loop / collected	Detects buried Wool via the __wool_inner__ walk and skips; raises a Wool-flavoured RuntimeError outside a loop; treats an idle-loop finalizer as teardown, not displacement	Install entry and idempotency
14	test_factory.py	A Wool factory displaced by a third-party install, by stash, or at task completion	The next wool.ContextVar set or task completion runs	Raises TaskFactoryDisplaced so the caller knows copy-on-fork is no longer in effect	Displacement detection
15	TestChainContention	An armed chain accessed from a foreign thread	A wool.ContextVar get / set / reset runs	Raises ChainContention with kind="thread" interpolating the chain id, owner thread, and offender; no false positive when the owner task is pending-off-loop or done	Cross-thread entry guard
16	TestChainContentionGuardBoundary	An armed contextvars.Context already driving a live task	A second concurrent task is created on that same context	Raises ChainContention with kind="task"; a callback on the owning thread does not	No shared context across concurrent tasks
17	TestChainContention (exceptions)	A ChainContention instance with structured fields	It is pickled and unpickled	The structured fields and message survive the round-trip; the hierarchy roots under WoolError / WoolWarning	Exception hierarchy and serialization
18	TestWoolToThreadParity	An armed context	wool.to_thread runs a function touching a wool.ContextVar	The worker thread runs on a fresh detached chain with no apply-back and no guard trip; a no-op when unarmed	wool.to_thread
19	TestAsyncioToThreadParity	An armed context	A plain asyncio.to_thread touches a wool.ContextVar	Raises ChainContention	Cross-thread guard at the stdlib boundary
20	TestRunInExecutorParity	A wool.ContextVar set in a scope	loop.run_in_executor runs on a thread and a process pool	The executor function carries neither stdlib nor wool context — matching stdlib semantics	Executor offload parity
21	TestTaskCreationValuePropagationParity	A wool.ContextVar set in a scope	asyncio.create_task / loop.create_task / TaskGroup create child tasks	The child observes the parent's value on a freshly minted chain UUID (copy-on-fork parity)	Parity matrix (default loop + uvloop)
22	TestCallbackValuePropagationParity	A wool.ContextVar set in a scope	Work is scheduled via call_soon / call_later / call_at / Future.add_done_callback	Callbacks observe the scheduling scope's value on its shared chain, untouched by the guard	Cooperative scheduling parity
23	TestAsyncGeneratorValuePropagation	An async generator suspended on yield in an armed context	The generator advances via asend / athrow / aclose	The generator's chain id is preserved across resume and the value mutation propagates	Async generator parity
24	TestFrameToProtobuf	A RequestFrame or ResponseFrame subclass	for_send is called with payload and serializer	Returns a Frame with payload encoded into the matching protobuf field and the armed chain captured into a ChainManifest	Frame construction and emission
25	TestFrameFromProtobuf	A wire envelope, including one whose chain manifest fails to decode under strict mode	Frame.from_protobuf is invoked	Returns the matching leaf with payload decoded; a decode failure is captured as the ChainSerializationError arm of chain_manifest, deferred rather than raised	Frame ingress and deferred decode
26	TestChainsDecodeErrorOntoPayload	An exception, nack, or throw frame carrying a deferred ChainSerializationError	Frame.mount runs	The decode error chains onto the payload exception's __context__ tail rather than preempting it	Strict-mode decode chaining
27	TestContextVarPropagation (integration)	A cross-process routine dispatch, including async-generator routines interleaved on one chain	The routine reads and mutates a wool.ContextVar	The caller's value propagates in, the worker's mutation propagates back, and serialized interleaving never trips the guard	End-to-end propagation