Merge branch 'main' into benchmark

Author: jy-tan

.cursor/BUGBOT.md

@@ -0,0 +1,5 @@
+# BUGBOT Notes
+
+## Instrumentation Guidelines
+
+- When adding a new instrumentation, the README must be updated to document the new instrumentation.

.github/workflows/ci.yml

@@ -20,7 +20,7 @@ jobs:
           version: "latest"
 
       - name: Setup Python
-        run: uv python install 3.12
+        run: uv python install 3.9
 
       - name: Install dependencies
         run: uv sync --all-extras
@@ -44,7 +44,7 @@ jobs:
           version: "latest"
 
       - name: Setup Python
-        run: uv python install 3.12
+        run: uv python install 3.9
 
       - name: Install dependencies
         run: uv sync --all-extras
@@ -53,8 +53,13 @@ jobs:
         run: uv run ty check drift/ tests/
 
   test:
-    name: Unit Tests
+    name: Unit Tests (Python ${{ matrix.python-version }})
     runs-on: ubuntu-latest
+    timeout-minutes: 30
+    strategy:
+      fail-fast: false
+      matrix:
+        python-version: ["3.9", "3.14"]
     steps:
       - name: Checkout
         uses: actions/checkout@v4
@@ -65,13 +70,33 @@ jobs:
           version: "latest"
 
       - name: Setup Python
-        run: uv python install 3.12
+        run: uv python install ${{ matrix.python-version }}
+
+      - name: Cache uv + Python installs + venv
+        uses: actions/cache@v4
+        with:
+          path: |
+            ~/.cache/uv
+            ~/.local/share/uv/python
+            .venv
+          key: ${{ runner.os }}-uv-${{ matrix.python-version }}-${{ hashFiles('uv.lock') }}
 
       - name: Install dependencies
         run: uv sync --all-extras
 
       - name: Run unit tests
-        run: uv run pytest tests/unit/ -v
+        if: matrix.python-version != '3.9'
+        run: uv run pytest tests/unit/ -q
+
+      - name: Run unit tests with coverage
+        if: matrix.python-version == '3.9'
+        run: uv run pytest tests/unit/ -q --cov=drift/core --cov-report=xml
+
+      - name: Upload coverage to Coveralls
+        if: matrix.python-version == '3.9' && github.event_name == 'push' && github.ref == 'refs/heads/main'
+        uses: coverallsapp/github-action@v2
+        with:
+          file: coverage.xml
 
   build:
     name: Build Package
@@ -86,7 +111,7 @@ jobs:
           version: "latest"
 
       - name: Setup Python
-        run: uv python install 3.12
+        run: uv python install 3.9
 
       - name: Install dependencies
         run: uv sync --all-extras
@@ -99,4 +124,3 @@ jobs:
           uv venv /tmp/test-install
           uv pip install dist/*.whl --python /tmp/test-install/bin/python
           /tmp/test-install/bin/python -c "import drift; print('Package imported successfully')"
-

.github/workflows/e2e.yml

@@ -54,7 +54,7 @@ jobs:
           version: "latest"
 
       - name: Setup Python
-        run: uv python install 3.12
+        run: uv python install 3.9
 
       - name: Setup Docker Buildx
         uses: docker/setup-buildx-action@v3

.gitignore

@@ -224,4 +224,9 @@ __marimo__/
 .DS_Store
 
 # Bug tracking
-**/BUG_TRACKING.md
+**/BUG_TRACKING.md
+
+# Coverage
+coverage.lcov
+coverage.xml
+.coverage

CONTRIBUTING.md

@@ -31,18 +31,25 @@ uv run ty check drift/ tests/          # Type check
 ### Unit Tests
 
 ```bash
-uv run python -m unittest discover -s tests/unit -v
+uv run pytest tests/unit/ -v
+
+# Run with coverage
+uv run pytest tests/unit/ -v --cov=drift --cov-report=term-missing
 
 # Run a specific test file
-uv run python -m unittest tests.unit.test_json_schema_helper -v
-uv run python -m unittest tests.unit.test_adapters -v
+uv run pytest tests/unit/test_json_schema_helper.py -v
+uv run pytest tests/unit/test_adapters.py -v
+
+# Run a specific test class or function
+uv run pytest tests/unit/test_metrics.py::TestMetricsCollector -v
+uv run pytest tests/unit/test_metrics.py::TestMetricsCollector::test_record_spans_exported -v
 ```
 
 ### Integration Tests
 
 ```bash
 # Flask/FastAPI integration tests
-timeout 30 uv run python -m unittest discover -s tests/integration -v
+timeout 30 uv run pytest tests/integration/ -v
 ```
 
 ### E2E Tests

README.md

@@ -39,43 +39,64 @@ For comprehensive guides and API reference, visit our [full documentation](https
 
 ## Requirements
 
-- Python 3.12+
+- Python 3.9+
 
 Tusk Drift currently supports the following packages and versions:
 
-- **Flask**: `flask>=2.0.0`
-- **FastAPI**: `fastapi>=0.68.0`
-- **Django**: `django>=3.2.0`
-- **Requests**: `requests` (all versions)
-- **HTTPX**: `httpx` (all versions)
-- **psycopg**: `psycopg>=3.0.0`, `psycopg2>=2.8.0`
-- **Redis**: `redis` (all versions)
+| Package | Supported Versions |
+|---------|-------------------|
+| Flask | `>=2.0.0` |
+| FastAPI | `>=0.68.0` |
+| Django | `>=3.2.0` |
+| Requests | all versions |
+| HTTPX | all versions |
+| aiohttp | all versions |
+| urllib3 | all versions |
+| grpcio (client-side only) | all versions |
+| psycopg | `>=3.1.12` |
+| psycopg2 | all versions |
+| Redis | `>=4.0.0` |
+| Kinde | `>=2.0.1` |
+| PyJWT | all versions |
+| urllib.request | all versions |
 
 If you're using packages or versions not listed above, please create an issue with the package + version you'd like an instrumentation for.
 
 ## Installation
 
 ### Step 1: Install the CLI
 
-First, install and configure the Tusk Drift CLI by following our [CLI installation guide](https://github.com/Use-Tusk/tusk-drift-cli?tab=readme-ov-file#install). The CLI helps set up your Tusk configuration file and replays tests.
+First, install the Tusk Drift CLI by following our [CLI installation guide](https://github.com/Use-Tusk/tusk-drift-cli?tab=readme-ov-file#install).
 
-The wizard will eventually direct you back here when it's time to set up the SDK.
+### Step 2: Set up Tusk Drift
 
-### Step 2: Install the SDK
+#### AI-powered setup (recommended)
 
-After completing the CLI wizard, install the SDK:
+Use our AI agent to automatically set up Tusk Drift for your service:
 
 ```bash
-pip install tusk-drift-python-sdk
+cd path/to/your/service
+export ANTHROPIC_API_KEY=your-api-key
+tusk setup
 ```
 
-### Step 3: Initialize the SDK for your service
+The agent will analyze your codebase, install the SDK, instrument it into your application, create configuration files, and test the setup with recording and replay.
 
-Refer to our [initialization guide](docs/initialization.md) to set up the SDK for your service.
+#### Manual setup
 
-### Step 4: Run Your First Test
+Alternatively, you can set up Tusk Drift manually:
 
-Follow along our [quick start guide](docs/quickstart.md) to record and replay your first test!
+1. Install the SDK:
+
+   ```bash
+   pip install tusk-drift-python-sdk
+   ```
+
+2. Create configuration: Run `tusk init` to create your `.tusk/config.yaml` config file interactively, or create it manually per the [configuration docs](https://github.com/Use-Tusk/tusk-drift-cli/blob/main/docs/configuration.md).
+
+3. Initialize the SDK: Refer to the [initialization guide](docs/initialization.md) to instrument the SDK in your service.
+
+4. Record and replay: Follow the [quick start guide](docs/quickstart.md) to record and replay your first test!
 
 ## Troubleshooting
 
@@ -87,11 +108,3 @@ Having issues?
 ## Community
 
 Join our open source community on [Slack](https://join.slack.com/t/tusk-community/shared_invite/zt-3fve1s7ie-NAAUn~UpHsf1m_2tdoGjsQ).
-
-## Contributing
-
-We appreciate feedback and contributions. See [CONTRIBUTING.md](/CONTRIBUTING.md).
-
-## License
-
-This project is licensed under the Apache License 2.0 - see the [LICENSE](LICENSE) file for details.

docs/initialization.md

@@ -4,9 +4,13 @@
 
 Before setting up the SDK, ensure you have:
 
-- Completed the [CLI wizard](https://github.com/Use-Tusk/tusk-drift-cli?tab=readme-ov-file#quick-start)
+- Python 3.9 or later installed
+- Installed the [Tusk Drift CLI](https://github.com/Use-Tusk/tusk-drift-cli?tab=readme-ov-file#install)
 - Obtained an API key from the [Tusk Drift dashboard](https://usetusk.ai/app/settings/api-keys) (only required if using Tusk Cloud)
-- Python 3.12 or later installed
+
+> [!TIP]
+> For automated setup, use `tusk setup` which handles SDK installation and initialization for you.
+> The steps below are for manual setup.
 
 ## Step 1: Install the SDK
 

drift/core/batch_processor.py

@@ -63,6 +63,7 @@ def __init__(
         self._condition = threading.Condition(self._lock)
         self._shutdown_event = threading.Event()
         self._export_thread: threading.Thread | None = None
+        self._thread_loop: asyncio.AbstractEventLoop | None = None
         self._started = False
         self._dropped_spans = 0
 
@@ -158,16 +159,23 @@ def add_span(self, span: CleanSpanData) -> bool:
 
     def _export_loop(self) -> None:
         """Background thread that periodically exports spans."""
-        while not self._shutdown_event.is_set():
-            # Wait for either: batch size reached, scheduled delay, or shutdown
-            with self._condition:
-                # Wait until batch is ready or timeout
-                self._condition.wait(timeout=self._config.scheduled_delay_seconds)
+        # Create a single long-lived event loop for this thread
+        self._thread_loop = asyncio.new_event_loop()
+        asyncio.set_event_loop(self._thread_loop)
 
-            if self._shutdown_event.is_set():
-                break
+        try:
+            while not self._shutdown_event.is_set():
+                # Wait for either: batch size reached, scheduled delay, or shutdown
+                with self._condition:
+                    self._condition.wait(timeout=self._config.scheduled_delay_seconds)
 
-            self._export_batch()
+                if self._shutdown_event.is_set():
+                    break
+
+                self._export_batch()
+        finally:
+            self._thread_loop.close()
+            self._thread_loop = None
 
     def _export_batch(self) -> None:
         """Export a batch of spans from the queue."""
@@ -188,16 +196,26 @@ def _export_batch(self) -> None:
         for adapter in adapters:
             start_time = time.monotonic()
             try:
-                # Handle async adapters (create new event loop for this thread)
+                # Handle async adapters
                 if asyncio.iscoroutinefunction(adapter.export_spans):
-                    loop = asyncio.new_event_loop()
-                    asyncio.set_event_loop(loop)
-                    try:
-                        loop.run_until_complete(adapter.export_spans(batch))
-                    finally:
-                        loop.close()
+                    # Only reuse the thread's event loop if we're on the export thread.
+                    # Using it from another thread (e.g., force_flush after join timeout)
+                    # would cause RuntimeError since event loops are not thread-safe.
+                    is_export_thread = threading.current_thread() is self._export_thread
+                    can_reuse_loop = (
+                        is_export_thread and self._thread_loop is not None and not self._thread_loop.is_closed()
+                    )
+                    if can_reuse_loop and self._thread_loop is not None:
+                        self._thread_loop.run_until_complete(adapter.export_spans(batch))
+                    else:
+                        loop = asyncio.new_event_loop()
+                        asyncio.set_event_loop(loop)
+                        try:
+                            loop.run_until_complete(adapter.export_spans(batch))
+                        finally:
+                            loop.close()
                 else:
-                    adapter.export_spans(batch)  # type: ignore
+                    adapter.export_spans(batch)
 
                 latency_ms = (time.monotonic() - start_time) * 1000
                 self._metrics.record_spans_exported(len(batch))