aiotrace 0.1.0__tar.gz

aiotrace-0.1.0/.gitignore

@@ -0,0 +1,14 @@
+__pycache__/
+*.pyc
+*.pyo
+.pytest_cache/
+dist/
+*.egg-info/
+.venv/
+venv/
+.coverage
+coverage/
+htmlcov/
+.mypy_cache/
+.ruff_cache/
+.pyright/

aiotrace-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 aiotrace contributors
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

aiotrace-0.1.0/PKG-INFO

@@ -0,0 +1,148 @@
+Metadata-Version: 2.4
+Name: aiotrace
+Version: 0.1.0
+Summary: OpenTelemetry-native async context propagation for asyncio
+Author: aiotrace contributors
+License: MIT
+License-File: LICENSE
+Keywords: asyncio,contextvars,observability,opentelemetry,tracing
+Classifier: Development Status :: 3 - Alpha
+Classifier: Framework :: AsyncIO
+Classifier: Intended Audience :: Developers
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3.8
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: System :: Monitoring
+Requires-Dist: opentelemetry-api>=1.20.0
+Provides-Extra: dev
+Requires-Dist: coverage>=7.0; extra == 'dev'
+Requires-Dist: opentelemetry-sdk>=1.20.0; extra == 'dev'
+Requires-Dist: opentelemetry-test-utils>=0.40b0; extra == 'dev'
+Requires-Dist: pyright>=1.1; extra == 'dev'
+Requires-Dist: pytest-asyncio>=0.21; extra == 'dev'
+Requires-Dist: pytest>=7.0; extra == 'dev'
+Requires-Dist: ruff>=0.1; extra == 'dev'
+Provides-Extra: sdk
+Requires-Dist: opentelemetry-sdk>=1.20.0; extra == 'sdk'
+Provides-Extra: test
+Requires-Dist: coverage>=7.0; extra == 'test'
+Requires-Dist: opentelemetry-sdk>=1.20.0; extra == 'test'
+Requires-Dist: opentelemetry-test-utils>=0.40b0; extra == 'test'
+Requires-Dist: pytest-asyncio>=0.21; extra == 'test'
+Requires-Dist: pytest>=7.0; extra == 'test'
+Description-Content-Type: text/markdown
+
+# aiotrace – Async context propagation for OpenTelemetry
+
+[![PyPI](https://img.shields.io/pypi/v/aiotrace)](https://pypi.org/project/aiotrace/)
+[![Python Versions](https://img.shields.io/pypi/pyversions/aiotrace)](https://pypi.org/project/aiotrace/)
+[![License](https://img.shields.io/pypi/l/aiotrace)](https://github.com/Kubenew/aiotrace/blob/main/LICENSE)
+[![GitHub](https://img.shields.io/github/stars/Kubenew/aiotrace?style=social)](https://github.com/Kubenew/aiotrace)
+[![Tests](https://img.shields.io/github/actions/workflow/status/Kubenew/aiotrace/ci.yml?label=tests)](https://github.com/Kubenew/aiotrace/actions)
+
+Fixes OpenTelemetry context propagation across asyncio boundaries:
+`create_task`, `Queue`, `Lock`, `Event`, `Semaphore`, `TaskGroup`, and
+`run_in_executor` (thread pool).
+
+## Problem
+
+OpenTelemetry stores the current span in a `contextvars.ContextVar`.
+When asyncio tasks are created or resumed, Python copies/shallow-copies
+these contextvars.  On Python < 3.9.17 / 3.10.7 / 3.11.1, `Task.__step`
+does **not** restore the task's own context, causing:
+
+* Spans created in child tasks appearing under the wrong parent
+* Context leaking between producer/consumer queues
+* Lost context when using `run_in_executor` (threads)
+
+## Quick Start
+
+```python
+import asyncio
+from opentelemetry import trace
+from aiotrace import install
+
+install()
+
+tracer = trace.get_tracer(__name__)
+
+async def child():
+    with tracer.start_as_current_span("child"):
+        pass
+
+async def main():
+    with tracer.start_as_current_span("parent"):
+        await asyncio.create_task(child())
+
+asyncio.run(main())
+```
+
+## Manual Usage (no monkey-patching)
+
+```python
+from aiotrace import create_task_with_context, PropagatingQueue
+
+# Use explicit wrapper instead of monkey-patch
+task = create_task_with_context(some_coro())
+
+# Explicit propagating queue
+queue = PropagatingQueue()
+await queue.put(item)
+item = await queue.get()
+```
+
+## Installation
+
+```bash
+pip install aiotrace
+```
+
+## What Gets Patched
+
+| Primitive | Replacement | Description |
+|-----------|-------------|-------------|
+| `asyncio.create_task` | Wrapped | Captures OTEL context at task creation, restores inside child |
+| `asyncio.Queue` | `PropagatingQueue` | Re-attaches context after `get()`/`put()` resume |
+| `asyncio.Lock` | `PropagatingLock` | Re-attaches context after `acquire()` |
+| `asyncio.Event` | `PropagatingEvent` | Re-attaches context after `wait()` |
+| `asyncio.Semaphore` | `PropagatingSemaphore` | Re-attaches context after `acquire()` |
+| `asyncio.Condition` | `PropagatingCondition` | Re-attaches context after `wait()` |
+| `asyncio.TaskGroup` (3.11+) | `PropagatingTaskGroup` | Captures context at `create_task()` |
+
+## API
+
+### `install(patch_queue=True, patch_locks=True)`
+
+Monkey-patches asyncio primitives. Safe to call multiple times (idempotent).
+
+### `uninstall()`
+
+Restores original asyncio primitives.
+
+### `PropagatingQueue(maxsize=0)`
+
+Drop-in replacement for `asyncio.Queue` with context propagation.
+
+### `PropagatingLock`, `PropagatingEvent`, `PropagatingSemaphore`, `PropagatingCondition`
+
+Drop-in replacements for synchronization primitives.
+
+### `create_task_with_context(coro, *, name=None, otel_ctx=None)`
+
+Create a task with explicit OTEL context propagation.
+
+### `run_in_executor_with_context(executor, func, *args)`
+
+Schedule a function in a thread pool with OTEL context.
+
+## Requirements
+
+* Python 3.8–3.12
+* `opentelemetry-api >= 1.20.0`
+
+## License
+
+MIT

aiotrace-0.1.0/README.md

@@ -0,0 +1,111 @@
+# aiotrace – Async context propagation for OpenTelemetry
+
+[![PyPI](https://img.shields.io/pypi/v/aiotrace)](https://pypi.org/project/aiotrace/)
+[![Python Versions](https://img.shields.io/pypi/pyversions/aiotrace)](https://pypi.org/project/aiotrace/)
+[![License](https://img.shields.io/pypi/l/aiotrace)](https://github.com/Kubenew/aiotrace/blob/main/LICENSE)
+[![GitHub](https://img.shields.io/github/stars/Kubenew/aiotrace?style=social)](https://github.com/Kubenew/aiotrace)
+[![Tests](https://img.shields.io/github/actions/workflow/status/Kubenew/aiotrace/ci.yml?label=tests)](https://github.com/Kubenew/aiotrace/actions)
+
+Fixes OpenTelemetry context propagation across asyncio boundaries:
+`create_task`, `Queue`, `Lock`, `Event`, `Semaphore`, `TaskGroup`, and
+`run_in_executor` (thread pool).
+
+## Problem
+
+OpenTelemetry stores the current span in a `contextvars.ContextVar`.
+When asyncio tasks are created or resumed, Python copies/shallow-copies
+these contextvars.  On Python < 3.9.17 / 3.10.7 / 3.11.1, `Task.__step`
+does **not** restore the task's own context, causing:
+
+* Spans created in child tasks appearing under the wrong parent
+* Context leaking between producer/consumer queues
+* Lost context when using `run_in_executor` (threads)
+
+## Quick Start
+
+```python
+import asyncio
+from opentelemetry import trace
+from aiotrace import install
+
+install()
+
+tracer = trace.get_tracer(__name__)
+
+async def child():
+    with tracer.start_as_current_span("child"):
+        pass
+
+async def main():
+    with tracer.start_as_current_span("parent"):
+        await asyncio.create_task(child())
+
+asyncio.run(main())
+```
+
+## Manual Usage (no monkey-patching)
+
+```python
+from aiotrace import create_task_with_context, PropagatingQueue
+
+# Use explicit wrapper instead of monkey-patch
+task = create_task_with_context(some_coro())
+
+# Explicit propagating queue
+queue = PropagatingQueue()
+await queue.put(item)
+item = await queue.get()
+```
+
+## Installation
+
+```bash
+pip install aiotrace
+```
+
+## What Gets Patched
+
+| Primitive | Replacement | Description |
+|-----------|-------------|-------------|
+| `asyncio.create_task` | Wrapped | Captures OTEL context at task creation, restores inside child |
+| `asyncio.Queue` | `PropagatingQueue` | Re-attaches context after `get()`/`put()` resume |
+| `asyncio.Lock` | `PropagatingLock` | Re-attaches context after `acquire()` |
+| `asyncio.Event` | `PropagatingEvent` | Re-attaches context after `wait()` |
+| `asyncio.Semaphore` | `PropagatingSemaphore` | Re-attaches context after `acquire()` |
+| `asyncio.Condition` | `PropagatingCondition` | Re-attaches context after `wait()` |
+| `asyncio.TaskGroup` (3.11+) | `PropagatingTaskGroup` | Captures context at `create_task()` |
+
+## API
+
+### `install(patch_queue=True, patch_locks=True)`
+
+Monkey-patches asyncio primitives. Safe to call multiple times (idempotent).
+
+### `uninstall()`
+
+Restores original asyncio primitives.
+
+### `PropagatingQueue(maxsize=0)`
+
+Drop-in replacement for `asyncio.Queue` with context propagation.
+
+### `PropagatingLock`, `PropagatingEvent`, `PropagatingSemaphore`, `PropagatingCondition`
+
+Drop-in replacements for synchronization primitives.
+
+### `create_task_with_context(coro, *, name=None, otel_ctx=None)`
+
+Create a task with explicit OTEL context propagation.
+
+### `run_in_executor_with_context(executor, func, *args)`
+
+Schedule a function in a thread pool with OTEL context.
+
+## Requirements
+
+* Python 3.8–3.12
+* `opentelemetry-api >= 1.20.0`
+
+## License
+
+MIT

aiotrace-0.1.0/pyproject.toml

@@ -0,0 +1,59 @@
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[project]
+name = "aiotrace"
+version = "0.1.0"
+description = "OpenTelemetry-native async context propagation for asyncio"
+readme = "README.md"
+license = { text = "MIT" }
+authors = [
+    { name = "aiotrace contributors" },
+]
+keywords = ["opentelemetry", "asyncio", "contextvars", "tracing", "observability"]
+classifiers = [
+    "Development Status :: 3 - Alpha",
+    "Framework :: AsyncIO",
+    "Intended Audience :: Developers",
+    "License :: OSI Approved :: MIT License",
+    "Programming Language :: Python :: 3.8",
+    "Programming Language :: Python :: 3.9",
+    "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+    "Topic :: System :: Monitoring",
+]
+dependencies = [
+    "opentelemetry-api>=1.20.0",
+]
+
+[project.optional-dependencies]
+sdk = ["opentelemetry-sdk>=1.20.0"]
+test = [
+    "pytest>=7.0",
+    "pytest-asyncio>=0.21",
+    "opentelemetry-sdk>=1.20.0",
+    "opentelemetry-test-utils>=0.40b0",
+    "coverage>=7.0",
+]
+dev = ["aiotrace[test]", "ruff>=0.1", "pyright>=1.1"]
+
+[tool.ruff]
+line-length = 100
+target-version = "py38"
+
+[tool.ruff.lint]
+select = ["E", "F", "I", "W"]
+
+[tool.pytest.ini_options]
+asyncio_mode = "auto"
+testpaths = ["tests"]
+filterwarnings = ["ignore::DeprecationWarning"]
+
+[tool.hatch.build.targets.wheel]
+packages = ["src/aiotrace"]
+
+[tool.pyright]
+strict = true
+include = ["src/aiotrace"]

aiotrace-0.1.0/src/aiotrace/__init__.py

@@ -0,0 +1,19 @@
+"""aiotrace — OpenTelemetry-native async context propagation for asyncio."""
+
+from aiotrace._install import install, is_installed, uninstall
+from aiotrace._lock import PropagatingLock
+from aiotrace._queue import PropagatingQueue
+from aiotrace._task import create_task_with_context, unwrap_create_task, wrap_create_task
+
+__all__ = [
+    "install",
+    "uninstall",
+    "is_installed",
+    "wrap_create_task",
+    "unwrap_create_task",
+    "create_task_with_context",
+    "PropagatingQueue",
+    "PropagatingLock",
+]
+
+__version__ = "0.1.0"

aiotrace-0.1.0/src/aiotrace/_executor.py

@@ -0,0 +1,64 @@
+"""Executor (thread-pool) context propagation for OpenTelemetry.
+
+``run_in_executor`` runs the callable in a thread-pool thread, which
+does **not** have access to the asyncio task's ``contextvars``.
+OTEL context must be explicitly copied to the worker thread.
+"""
+
+from __future__ import annotations
+
+import asyncio
+from typing import Any, Callable, Optional, TypeVar
+
+from opentelemetry import context as otel_context
+
+T = TypeVar("T")
+
+
+def run_in_executor_with_context(
+    executor: Optional[Any],
+    func: Callable[..., T],
+    *args: Any,
+) -> asyncio.Future[T]:
+    """Schedule ``func(*args, **kwargs)`` in an executor, preserving OTEL context.
+
+    Usage::
+
+        result = await run_in_executor_with_context(None, blocking_io, path)
+    """
+    ctx = otel_context.get_current()
+
+    def wrapper() -> T:
+        token = otel_context.attach(ctx)
+        try:
+            return func(*args)
+        finally:
+            otel_context.detach(token)
+
+    loop = asyncio.get_running_loop()
+    return loop.run_in_executor(executor, wrapper)
+
+
+def wrap_run_in_executor(loop: Optional[asyncio.AbstractEventLoop] = None) -> None:
+    """Monkey-patch ``loop.run_in_executor`` to preserve OTEL context."""
+    if loop is None:
+        loop = asyncio.get_running_loop()
+    original = loop.run_in_executor
+
+    def _run_in_executor_with_context(
+        executor: Optional[Any],
+        func: Callable[..., T],
+        *args: Any,
+    ) -> asyncio.Future[T]:
+        ctx = otel_context.get_current()
+
+        def wrapper() -> T:
+            token = otel_context.attach(ctx)
+            try:
+                return func(*args)
+            finally:
+                otel_context.detach(token)
+
+        return original(executor, wrapper)
+
+    loop.run_in_executor = _run_in_executor_with_context

aiotrace-0.1.0/src/aiotrace/_install.py

@@ -0,0 +1,61 @@
+"""Central ``install()`` / ``uninstall()`` entry point.
+
+Calling ``install()`` monkey-patches ``asyncio.create_task``,
+``asyncio.Queue``, ``asyncio.Lock``, etc., with context-propagating
+wrappers.  Use ``uninstall()`` to restore the originals.
+"""
+
+from __future__ import annotations
+
+from aiotrace._lock import patch_locks
+from aiotrace._queue import patch_queue
+from aiotrace._task import unwrap_create_task, wrap_create_task
+from aiotrace._taskgroup import patch_taskgroup
+
+_installed = False
+
+
+def install(*, queue: bool = True, locks: bool = True) -> None:
+    """Monkey-patch asyncio primitives to propagate OTEL context.
+
+    Patches applied:
+
+    * ``asyncio.create_task`` — wraps child coroutines to restore
+      the parent's OTEL context at creation time.
+    * ``asyncio.Queue`` → :class:`aiotrace.PropagatingQueue` (opt-in).
+    * ``asyncio.{Lock,Event,Semaphore,Condition}`` → propagating
+      versions (opt-in).
+    * ``asyncio.TaskGroup`` → :class:`aiotrace.PropagatingTaskGroup`
+      (Python 3.11+).
+
+    Args:
+        queue: Whether to replace ``asyncio.Queue``.
+        locks: Whether to replace lock primitives.
+    """
+    global _installed
+    if _installed:
+        return
+
+    wrap_create_task()
+
+    if queue:
+        patch_queue()
+
+    if locks:
+        patch_locks()
+
+    patch_taskgroup()
+
+    _installed = True
+
+
+def uninstall() -> None:
+    """Restore all original asyncio primitives."""
+    global _installed
+    unwrap_create_task()
+    _installed = False
+
+
+def is_installed() -> bool:
+    """Return ``True`` if ``install()`` has been called."""
+    return _installed

aiotrace-0.1.0/src/aiotrace/_lock.py

@@ -0,0 +1,60 @@
+"""Lock/Event/Semaphore context propagation for OpenTelemetry.
+
+Same problem as Queue: when a coroutine is resumed after acquiring
+a lock (or waiting on an event/semaphore), it may wake in the wrong
+context on older Python versions.
+"""
+
+from __future__ import annotations
+
+import asyncio
+
+from opentelemetry import context as otel_context
+
+
+class PropagatingLock(asyncio.Lock):
+    """An ``asyncio.Lock`` that preserves OTEL context across ``acquire``."""
+
+    async def acquire(self) -> bool:
+        ctx = otel_context.get_current()
+        result = await super().acquire()
+        otel_context.attach(ctx)
+        return result
+
+
+class PropagatingEvent(asyncio.Event):
+    """An ``asyncio.Event`` that preserves OTEL context across ``wait``."""
+
+    async def wait(self) -> bool:
+        ctx = otel_context.get_current()
+        result = await super().wait()
+        otel_context.attach(ctx)
+        return result
+
+
+class PropagatingSemaphore(asyncio.Semaphore):
+    """An ``asyncio.Semaphore`` that preserves OTEL context across ``acquire``."""
+
+    async def acquire(self) -> bool:
+        ctx = otel_context.get_current()
+        result = await super().acquire()
+        otel_context.attach(ctx)
+        return result
+
+
+class PropagatingCondition(asyncio.Condition):
+    """An ``asyncio.Condition`` that preserves OTEL context across ``wait``."""
+
+    async def wait(self) -> bool:
+        ctx = otel_context.get_current()
+        result = await super().wait()
+        otel_context.attach(ctx)
+        return result
+
+
+def patch_locks() -> None:
+    """Replace stdlib synchronization primitives with propagating versions."""
+    asyncio.Lock = PropagatingLock  # type: ignore[misc]
+    asyncio.Event = PropagatingEvent  # type: ignore[misc]
+    asyncio.Semaphore = PropagatingSemaphore  # type: ignore[misc]
+    asyncio.Condition = PropagatingCondition  # type: ignore[misc]

aiotrace-0.1.0/src/aiotrace/_queue.py

@@ -0,0 +1,41 @@
+"""Queue context propagation for OpenTelemetry.
+
+On Python < 3.9.17 / 3.10.7 / 3.11.1, ``Task.__step`` does not restore
+the task's own context, so ``await queue.get()`` can resume the consumer
+in the **producer's** context.  ``PropagatingQueue`` re-attaches the
+consumer's OTEL context after every resume, fixing span nesting.
+"""
+
+from __future__ import annotations
+
+import asyncio
+from typing import Any
+
+from opentelemetry import context as otel_context
+
+
+class PropagatingQueue(asyncio.Queue):
+    """An ``asyncio.Queue`` that preserves OTEL context across ``get``/``put``.
+
+    Usage is identical to ``asyncio.Queue``::
+
+        queue = PropagatingQueue(maxsize=10)
+        await queue.put(item)
+        item = await queue.get()
+    """
+
+    async def get(self) -> Any:
+        ctx = otel_context.get_current()
+        item = await super().get()
+        otel_context.attach(ctx)
+        return item
+
+    async def put(self, item: Any) -> None:
+        ctx = otel_context.get_current()
+        await super().put(item)
+        otel_context.attach(ctx)
+
+
+def patch_queue() -> None:
+    """Replace ``asyncio.Queue`` with ``PropagatingQueue`` globally."""
+    asyncio.Queue = PropagatingQueue  # type: ignore[misc]

aiotrace-0.1.0/src/aiotrace/_task.py

@@ -0,0 +1,122 @@
+"""Task context propagation for OpenTelemetry.
+
+Captures the OTEL context at task-creation time and restores it
+inside the child task so that spans created in the child correctly
+nest under the parent span.
+"""
+
+import asyncio
+from functools import wraps
+from typing import Any, Callable, Coroutine, Optional, TypeVar
+
+from opentelemetry import context as otel_context
+
+T = TypeVar("T")
+
+_original_create_task: Optional[Callable[..., asyncio.Task]] = None
+
+
+def create_task_with_context(
+    coro: Coroutine[Any, Any, T],
+    *,
+    name: Optional[str] = None,
+    otel_ctx: Optional[otel_context.Context] = None,
+) -> asyncio.Task[T]:
+    """Create an asyncio Task that preserves the current OTEL context.
+
+    Unlike ``asyncio.create_task``, which relies on ``contextvars.copy_context()``
+    at ``Task.__init__`` time, this wrapper explicitly captures the **OTEL**
+    context and restores it inside the coroutine wrapper.  This guarantees
+    that the child task sees the same OTEL context even if the parent's
+    context changes between creation and first ``await``.
+
+    Args:
+        coro: The coroutine to run in the new task.
+        name: Optional task name (passed through to ``asyncio.create_task``).
+        otel_ctx: An explicit OTEL context to use.  Defaults to
+            ``otel_context.get_current()`` at call time.
+
+    Returns:
+        A new ``asyncio.Task`` instance.
+    """
+    ctx = otel_ctx if otel_ctx is not None else otel_context.get_current()
+
+    async def _wrapper() -> T:
+        token = otel_context.attach(ctx)
+        try:
+            return await coro
+        finally:
+            otel_context.detach(token)
+
+    return asyncio.create_task(_wrapper(), name=name)
+
+
+def wrap_create_task() -> None:
+    """Replace ``asyncio.create_task`` with a context-propagating version.
+
+    This is a lighter alternative to ``install()`` – it only patches
+    the task-creation boundary, not queues or locks.
+    """
+    global _original_create_task
+    if _original_create_task is not None:
+        return  # already wrapped
+
+    _original_create_task = asyncio.create_task
+
+    @wraps(_original_create_task)
+    def _create_task_with_context(
+        coro: Coroutine[Any, Any, T],
+        *,
+        name: Optional[str] = None,
+        context: Optional[otel_context.Context] = None,
+    ) -> asyncio.Task[T]:
+        ctx = context or otel_context.get_current()
+
+        async def _wrapper() -> T:
+            token = otel_context.attach(ctx)
+            try:
+                return await coro
+            finally:
+                otel_context.detach(token)
+
+        return _original_create_task(_wrapper(), name=name)
+
+    asyncio.create_task = _create_task_with_context
+
+
+def unwrap_create_task() -> None:
+    """Restore the original ``asyncio.create_task``."""
+    global _original_create_task
+    if _original_create_task is not None:
+        asyncio.create_task = _original_create_task
+        _original_create_task = None
+
+
+def _patch_loop_create_task(loop: asyncio.AbstractEventLoop) -> None:
+    """Patch ``loop.create_task`` for the given event loop.
+
+    Some libraries call ``loop.create_task`` directly instead of
+    ``asyncio.create_task``.  This function patches the loop-level
+    method for completeness.
+    """
+    original = loop.create_task
+
+    @wraps(original)
+    def _loop_create_task(
+        coro: Coroutine[Any, Any, T],
+        *,
+        name: Optional[str] = None,
+        context: Optional[otel_context.Context] = None,
+    ) -> asyncio.Task[T]:
+        ctx = context or otel_context.get_current()
+
+        async def _wrapper() -> T:
+            token = otel_context.attach(ctx)
+            try:
+                return await coro
+            finally:
+                otel_context.detach(token)
+
+        return original(_wrapper(), name=name)
+
+    loop.create_task = _loop_create_task

aiotrace-0.1.0/src/aiotrace/_taskgroup.py

@@ -0,0 +1,63 @@
+"""TaskGroup context propagation for OpenTelemetry.
+
+:func:`asyncio.TaskGroup` (Python 3.11+)  creates child tasks via
+``loop.create_task``.  If the loop-level method has not been patched,
+the child tasks may inherit a stale context.
+
+This module provides ``PropagatingTaskGroup`` which explicitly captures
+the OTEL context at ``create_task`` time and wraps the child coroutine.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import sys
+from typing import Any, Coroutine, Optional, TypeVar
+
+from opentelemetry import context as otel_context
+
+T = TypeVar("T")
+
+if sys.version_info >= (3, 11):
+
+    class PropagatingTaskGroup(asyncio.TaskGroup):
+        """An :class:`asyncio.TaskGroup` that preserves OTEL context.
+
+        Usage::
+
+            async with PropagatingTaskGroup() as tg:
+                tg.create_task(some_coro())
+        """
+
+        def create_task(
+            self,
+            coro: Coroutine[Any, Any, T],
+            *,
+            name: Optional[str] = None,
+            context: Optional[otel_context.Context] = None,
+        ) -> asyncio.Task[T]:
+            ctx = context or otel_context.get_current()
+
+            async def _wrapper() -> T:
+                token = otel_context.attach(ctx)
+                try:
+                    return await coro
+                finally:
+                    otel_context.detach(token)
+
+            return super().create_task(_wrapper(), name=name)
+
+    def patch_taskgroup() -> None:
+        """Replace ``asyncio.TaskGroup`` with ``PropagatingTaskGroup``."""
+        asyncio.TaskGroup = PropagatingTaskGroup  # type: ignore[misc]
+
+else:
+
+    class PropagatingTaskGroup:  # type: ignore[no-redef]
+        """Placeholder — ``asyncio.TaskGroup`` requires Python 3.11+."""
+
+        def __init__(self, *args: Any, **kwargs: Any):
+            raise RuntimeError("TaskGroup requires Python 3.11+")
+
+    def patch_taskgroup() -> None:
+        """No-op on Python < 3.11."""

aiotrace-0.1.0/tests/conftest.py

@@ -0,0 +1,48 @@
+"""pytest fixtures for aiotrace tests.
+
+WARNING: ``set_tracer_provider`` uses an internal ``Once`` guard — calling
+``get_tracer_provider()`` during import (e.g. via module-level code in any
+test file) consumes it, making all subsequent ``set_tracer_provider()``
+calls no-ops.  We work around this by directly assigning
+``trace._TRACER_PROVIDER`` instead.
+"""
+
+from __future__ import annotations
+
+import pytest
+from opentelemetry import trace
+from opentelemetry.sdk.trace import TracerProvider
+from opentelemetry.sdk.trace.export import SimpleSpanProcessor, Span, SpanExporter
+
+
+class CapturingExporter(SpanExporter):
+    """Collects exported spans in memory for test assertions."""
+
+    def __init__(self):
+        self.spans: list[Span] = []
+
+    def export(self, spans: list[Span]) -> None:
+        self.spans.extend(spans)
+
+    def shutdown(self) -> None:
+        self.spans.clear()
+
+
+@pytest.fixture
+def exporter() -> CapturingExporter:
+    return CapturingExporter()
+
+
+@pytest.fixture(autouse=True)
+def setup_otel(exporter: CapturingExporter):
+    provider = TracerProvider()
+    provider.add_span_processor(SimpleSpanProcessor(exporter))
+    # Direct assignment to bypass OTEL's Once guard
+    trace._TRACER_PROVIDER = provider
+    yield
+    trace._TRACER_PROVIDER = None
+
+
+@pytest.fixture
+def tracer():
+    return trace.get_tracer("aiotrace.test")

aiotrace-0.1.0/tests/test_install.py

@@ -0,0 +1,41 @@
+"""Tests for the install/uninstall mechanism."""
+
+from __future__ import annotations
+
+import asyncio
+
+import pytest
+
+from aiotrace import install, is_installed, uninstall
+
+
+@pytest.mark.asyncio
+async def test_install_and_uninstall():
+    assert not is_installed()
+    install()
+    assert is_installed()
+    uninstall()
+    assert not is_installed()
+
+
+@pytest.mark.asyncio
+async def test_install_idempotent():
+    install()
+    assert is_installed()
+    # Second call should be a no-op
+    install()
+    assert is_installed()
+    uninstall()
+
+
+@pytest.mark.asyncio
+async def test_create_task_still_works_after_install():
+    install()
+    try:
+        async def dummy():
+            return 42
+
+        result = await asyncio.create_task(dummy())
+        assert result == 42
+    finally:
+        uninstall()

aiotrace-0.1.0/tests/test_queue.py

@@ -0,0 +1,67 @@
+"""Tests for queue context propagation."""
+
+from __future__ import annotations
+
+import asyncio
+
+import pytest
+
+from aiotrace import PropagatingQueue
+
+
+@pytest.mark.asyncio
+async def test_queue_producer_consumer_spans(tracer, exporter):
+    async def consumer(q: asyncio.Queue):
+        item = await q.get()
+        with tracer.start_as_current_span("process"):
+            pass
+        q.task_done()
+        return item
+
+    async def producer(q: asyncio.Queue):
+        with tracer.start_as_current_span("produce"):
+            await q.put("item")
+
+    q = PropagatingQueue()
+    prod = asyncio.create_task(producer(q))
+    cons = asyncio.create_task(consumer(q))
+    await asyncio.gather(prod, cons)
+    await q.join()
+
+    names = {s.name for s in exporter.spans}
+    assert "produce" in names
+    assert "process" in names
+
+
+@pytest.mark.asyncio
+async def test_queue_multiple_consumers(tracer, exporter):
+    async def consumer(q: asyncio.Queue, name: str):
+        await q.get()
+        with tracer.start_as_current_span(f"consume_{name}"):
+            pass
+        q.task_done()
+
+    async def producer(q: asyncio.Queue):
+        with tracer.start_as_current_span("produce"):
+            for i in range(3):
+                await q.put(i)
+
+    q = PropagatingQueue(maxsize=1)
+    prod = asyncio.create_task(producer(q))
+    cons = [asyncio.create_task(consumer(q, str(i))) for i in range(3)]
+    await asyncio.gather(prod, *cons)
+    await q.join()
+
+    assert len(exporter.spans) >= 4
+
+
+@pytest.mark.asyncio
+async def test_queue_standard_behavior(tracer, exporter):
+    q = PropagatingQueue()
+    await q.put(1)
+    await q.put(2)
+    result = await q.get()
+    assert result == 1
+    result = await q.get()
+    assert result == 2
+    assert q.empty()

aiotrace-0.1.0/tests/test_task.py

@@ -0,0 +1,131 @@
+"""Tests for task context propagation."""
+
+from __future__ import annotations
+
+import asyncio
+import sys
+
+import pytest
+
+from aiotrace import unwrap_create_task, wrap_create_task
+
+
+@pytest.mark.asyncio
+async def test_child_task_inherits_parent_context(tracer, exporter):
+    wrap_create_task()
+    try:
+        async def child():
+            with tracer.start_as_current_span("child"):
+                pass
+
+        async def parent():
+            with tracer.start_as_current_span("parent"):
+                await asyncio.create_task(child())
+
+        await parent()
+        assert len(exporter.spans) == 2
+        names = {s.name for s in exporter.spans}
+        assert names == {"parent", "child"}
+    finally:
+        unwrap_create_task()
+
+
+@pytest.mark.asyncio
+async def test_gather_preserves_context(tracer, exporter):
+    wrap_create_task()
+    try:
+        async def child(n):
+            with tracer.start_as_current_span(f"child_{n}"):
+                pass
+
+        async def parent():
+            with tracer.start_as_current_span("parent"):
+                await asyncio.gather(child(1), child(2), child(3))
+
+        await parent()
+        assert len(exporter.spans) == 4
+    finally:
+        unwrap_create_task()
+
+
+@pytest.mark.asyncio
+async def test_nested_create_task(tracer, exporter):
+    wrap_create_task()
+    try:
+        async def inner():
+            with tracer.start_as_current_span("inner"):
+                pass
+
+        async def outer():
+            with tracer.start_as_current_span("outer"):
+                await asyncio.create_task(inner())
+
+        await outer()
+        assert len(exporter.spans) == 2
+        names = {s.name for s in exporter.spans}
+        assert names == {"outer", "inner"}
+    finally:
+        unwrap_create_task()
+
+
+@pytest.mark.skipif(sys.version_info < (3, 11), reason="TaskGroup requires 3.11+")
+@pytest.mark.asyncio
+async def test_create_task_in_taskgroup(tracer, exporter):
+    wrap_create_task()
+    try:
+        async def child():
+            with tracer.start_as_current_span("child"):
+                pass
+
+        async def parent():
+            with tracer.start_as_current_span("parent"):
+                async with asyncio.TaskGroup() as tg:
+                    tg.create_task(child())
+                    tg.create_task(child())
+
+        await parent()
+        assert len(exporter.spans) == 3
+    finally:
+        unwrap_create_task()
+
+
+@pytest.mark.asyncio
+async def test_context_maintained_after_task_cancel(tracer, exporter):
+    wrap_create_task()
+    try:
+        async def cancellable():
+            try:
+                await asyncio.sleep(10)
+            except asyncio.CancelledError:
+                raise
+
+        async def parent():
+            with tracer.start_as_current_span("parent"):
+                task = asyncio.create_task(cancellable())
+                await asyncio.sleep(0.01)
+                task.cancel()
+                try:
+                    await task
+                except asyncio.CancelledError:
+                    pass
+
+        await parent()
+        assert len(exporter.spans) == 1
+        assert exporter.spans[0].name == "parent"
+    finally:
+        unwrap_create_task()
+
+
+@pytest.mark.asyncio
+async def test_create_task_without_patch(tracer, exporter):
+    """Sanity check: standard create_task on modern Python."""
+    async def child():
+        with tracer.start_as_current_span("child"):
+            pass
+
+    async def parent():
+        with tracer.start_as_current_span("parent"):
+            await asyncio.create_task(child())
+
+    await parent()
+    assert len(exporter.spans) == 2