chore(deps): bump sqlparse from 0.5.3 to 0.5.4

Bumps [sqlparse](https://github.com/andialbrecht/sqlparse) from 0.5.3 to 0.5.4.
- [Changelog](https://github.com/andialbrecht/sqlparse/blob/master/CHANGELOG)
- [Commits](https://github.com/andialbrecht/sqlparse/compare/0.5.3...0.5.4)

---
updated-dependencies:
- dependency-name: sqlparse
  dependency-version: 0.5.4
  dependency-type: indirect
...

Signed-off-by: dependabot[bot] <support@github.com>

.github/workflows/ci.yml

@@ -49,7 +49,7 @@ jobs:
         runs-on: ubuntu-latest
         strategy:
             matrix:
-                python-version: ['3.10', '3.11', '3.12', '3.13']
+                python-version: ['3.10', '3.11', '3.12', '3.13', '3.14']
 
         steps:
             - uses: actions/checkout@85e6279cec87321a52edac9c87bce653a07cf6c2
@@ -76,6 +76,29 @@ jobs:
               run: |
                   pytest --verbose --timeout=30
 
+    import-check:
+        name: Python ${{ matrix.python-version }} import check
+        runs-on: ubuntu-latest
+        strategy:
+            matrix:
+                python-version: ['3.10', '3.11', '3.12', '3.13', '3.14']
+
+        steps:
+            - uses: actions/checkout@85e6279cec87321a52edac9c87bce653a07cf6c2
+              with:
+                  fetch-depth: 1
+
+            - name: Set up Python ${{ matrix.python-version }}
+              uses: actions/setup-python@8d9ed9ac5c53483de85588cdf95a591a75ab9f55
+              with:
+                  python-version: ${{ matrix.python-version }}
+
+            - name: Install posthog
+              run: pip install .
+
+            - name: Check import produces no warnings
+              run: python -W error -c "import posthog"
+
     django5-integration:
         name: Django 5 integration tests
         runs-on: ubuntu-latest

.github/workflows/codeql.yml

@@ -0,0 +1,45 @@
+name: 'CodeQL Advanced'
+
+on:
+  push:
+    branches: ['master']
+  pull_request:
+    branches: ['master']
+  schedule:
+    - cron: '32 13 * * 1'
+
+jobs:
+  analyze:
+    name: Analyze (${{ matrix.language }})
+    runs-on: 'ubuntu-latest'
+    permissions:
+      security-events: write
+      # required to fetch internal or private CodeQL packs
+      packages: read
+
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
+          - language: actions
+            build-mode: none
+          - language: python
+            build-mode: none
+    steps:
+      - name: Checkout repository
+        uses: actions/checkout@v6
+
+      - name: Initialize CodeQL
+        uses: github/codeql-action/init@5d4e8d1aca955e8d8589aabd499c5cae939e33c7 # v4.31.9
+        with:
+          languages: ${{ matrix.language }}
+          build-mode: ${{ matrix.build-mode }}
+          # Disable TRAP caching - it creates a new cache per commit SHA which
+          # is never reused, causing wasted cache space.
+          # See: https://github.com/github/codeql-action/issues/2030
+          trap-caching: false
+
+      - name: Perform CodeQL Analysis
+        uses: github/codeql-action/analyze@5d4e8d1aca955e8d8589aabd499c5cae939e33c7 # v4.31.9
+        with:
+          category: '/language:${{matrix.language}}'

.github/workflows/generate-references.yml

@@ -7,12 +7,13 @@ jobs:
   docs-generation:
     name: Generate references
     runs-on: ubuntu-latest
+    permissions:
+      contents: write
     steps:
       - name: Checkout the repository
         uses: actions/checkout@85e6279cec87321a52edac9c87bce653a07cf6c2
         with:
           fetch-depth: 0
-          token: ${{ secrets.POSTHOG_BOT_PAT }}
 
       - name: Set up Python
         uses: actions/setup-python@8d9ed9ac5c53483de85588cdf95a591a75ab9f55

.github/workflows/release.yml

@@ -1,57 +1,249 @@
 name: "Release"
 
 on:
-  push:
-    branches:
-      - master
-    paths:
-      - "posthog/version.py"
+  pull_request:
+    types: [closed]
+    branches: [master]
   workflow_dispatch:
 
+permissions:
+  contents: read
+
+# Concurrency control: only one release process can run at a time
+# This prevents race conditions if multiple PRs with 'release' label merge simultaneously
+concurrency:
+  group: release
+  cancel-in-progress: false
+
 jobs:
-  release:
-    name: Publish release
+  check-release-label:
+    name: Check for release label
     runs-on: ubuntu-latest
-    env:
-      TWINE_USERNAME: __token__
-      TWINE_PASSWORD: ${{ secrets.PYPI_API_TOKEN }}
+    # Run when PR with 'release' label is merged to master
+    if: |
+      github.event_name == 'workflow_dispatch' ||
+       (github.event_name == 'pull_request' &&
+        github.event.pull_request.merged == true &&
+        contains(github.event.pull_request.labels.*.name, 'release'))
+    outputs:
+      should-release: ${{ steps.check.outputs.should-release }}
     steps:
-      - name: Checkout the repository
-        uses: actions/checkout@85e6279cec87321a52edac9c87bce653a07cf6c2
+      - name: Checkout repository
+        uses: actions/checkout@v4
         with:
+          ref: master
           fetch-depth: 0
-          token: ${{ secrets.POSTHOG_BOT_PAT }}
+
+      - name: Check release conditions
+        id: check
+        run: |
+          changeset_count=$(find .sampo/changesets -name '*.md' 2>/dev/null | wc -l)
+          if [ "$changeset_count" -gt 0 ]; then
+            echo "should-release=true" >> "$GITHUB_OUTPUT"
+            echo "Found $changeset_count changeset(s), ready to release"
+          else
+            echo "should-release=false" >> "$GITHUB_OUTPUT"
+            echo "No changesets to release"
+          fi
+
+  notify-approval-needed:
+    name: Notify Slack - Approval Needed
+    needs: check-release-label
+    if: needs.check-release-label.outputs.should-release == 'true'
+    uses: posthog/.github/.github/workflows/notify-approval-needed.yml@main
+    with:
+      slack_channel_id: ${{ vars.SLACK_APPROVALS_CLIENT_LIBRARIES_CHANNEL_ID }}
+      slack_user_group_id: ${{ vars.GROUP_CLIENT_LIBRARIES_SLACK_GROUP_ID }}
+    secrets:
+      slack_bot_token: ${{ secrets.SLACK_CLIENT_LIBRARIES_BOT_TOKEN }}
+      posthog_project_api_key: ${{ secrets.POSTHOG_PROJECT_API_KEY }}
+
+  release:
+    name: Release and publish
+    needs: [check-release-label, notify-approval-needed]
+    runs-on: ubuntu-latest
+    # Use `always()` to ensure the job runs even if notify-approval-needed is skipped,
+    # but still depend on it to access `needs.notify-approval-needed.outputs.slack_ts`
+    if: always() && needs.check-release-label.outputs.should-release == 'true'
+    environment: "Release" # This will require an approval from a maintainer, they are notified in Slack above
+    permissions:
+      contents: write
+      actions: write
+      id-token: write
+    steps:
+      - name: Notify Slack - Approved
+        if: needs.notify-approval-needed.outputs.slack_ts != ''
+        uses: posthog/.github/.github/actions/slack-thread-reply@main
+        with:
+          slack_bot_token: ${{ secrets.SLACK_CLIENT_LIBRARIES_BOT_TOKEN }}
+          slack_channel_id: ${{ vars.SLACK_APPROVALS_CLIENT_LIBRARIES_CHANNEL_ID }}
+          thread_ts: ${{ needs.notify-approval-needed.outputs.slack_ts }}
+          message: "✅ Release approved! Version bump in progress..."
+          emoji_reaction: "white_check_mark"
+
+      - name: Get GitHub App token
+        id: releaser
+        uses: actions/create-github-app-token@v2
+        with:
+          app-id: ${{ secrets.GH_APP_POSTHOG_PYTHON_RELEASER_APP_ID }}
+          private-key: ${{ secrets.GH_APP_POSTHOG_PYTHON_RELEASER_PRIVATE_KEY }}
+
+      - name: Checkout repository
+        uses: actions/checkout@v4
+        with:
+          ref: master
+          fetch-depth: 0
+          token: ${{ steps.releaser.outputs.token }}
 
       - name: Set up Python
-        uses: actions/setup-python@8d9ed9ac5c53483de85588cdf95a591a75ab9f55
+        uses: actions/setup-python@v5
         with:
           python-version: 3.11.11
 
       - name: Install uv
-        uses: astral-sh/setup-uv@0c5e2b8115b80b4c7c5ddf6ffdd634974642d182 # v5.4.1
+        uses: astral-sh/setup-uv@v5
         with:
-            enable-cache: true
-            pyproject-file: 'pyproject.toml'
-            
-      - name: Detect version
-        run: echo "REPO_VERSION=$(python3 posthog/version.py)" >> $GITHUB_ENV
+          enable-cache: true
+          pyproject-file: "pyproject.toml"
 
-      - name: Prepare for building release
+      - name: Install Rust
+        uses: dtolnay/rust-toolchain@0b1efabc08b657293548b77fb76cc02d26091c7e
+        with:
+          toolchain: 1.91.1
+          components: cargo
+
+      - name: Cache Sampo CLI
+        id: cache-sampo
+        uses: actions/cache@v3
+        with:
+          path: ~/.cargo/bin/sampo
+          key: sampo-${{ runner.os }}-${{ runner.arch }}
+
+      - name: Install Sampo CLI
+        if: steps.cache-sampo.outputs.cache-hit != 'true'
+        run: cargo install sampo
+
+      - name: Install dependencies
         run: uv sync --extra dev
 
-      - name: Push releases to PyPI
-        run: uv run make release && uv run make release_analytics
+      - name: Configure Git
+        run: |
+          git config user.name "github-actions[bot]"
+          git config user.email "github-actions[bot]@users.noreply.github.com"
 
-      - name: Create GitHub release
-        uses: actions/create-release@0cb9c9b65d5d1901c1f53e5e66eaf4afd303e70e # v1
+      - name: Prepare release with Sampo
+        id: sampo-release
         env:
-          GITHUB_TOKEN: ${{ secrets.POSTHOG_BOT_PAT }}
-        with:
-          tag_name: v${{ env.REPO_VERSION }}
-          release_name: ${{ env.REPO_VERSION }}
-      
-      - name: Dispatch generate-references for posthog-python
+          GITHUB_TOKEN: ${{ steps.releaser.outputs.token }}
+        run: |
+          sampo release
+          new_version=$(python3 -c "import tomllib; print(tomllib.load(open('pyproject.toml', 'rb'))['project']['version'])")
+          echo "new_version=$new_version" >> "$GITHUB_OUTPUT"
+
+      - name: Sync version to posthog/version.py
+        run: |
+          echo 'VERSION = "${{ steps.sampo-release.outputs.new_version }}"' > posthog/version.py
+
+      - name: Commit release changes
+        id: commit-release
+        env:
+          GITHUB_TOKEN: ${{ steps.releaser.outputs.token }}
+        run: |
+          git add -A
+          if git diff --staged --quiet; then
+            echo "No changes to commit"
+            echo "committed=false" >> "$GITHUB_OUTPUT"
+          else
+            git commit -m "chore: Release v${{ steps.sampo-release.outputs.new_version }}"
+            git push origin master
+            echo "committed=true" >> "$GITHUB_OUTPUT"
+          fi
+
+      # Publishing is done manually (not via `sampo publish`) because we need to
+      # publish both `posthog` and `posthoganalytics` packages to PyPI.
+      # Sampo only knows about the `posthog` package, so we handle both here.
+      # Both packages use PyPI OIDC trusted publishing (no API tokens needed).
+      - name: Build posthog
+        if: steps.commit-release.outputs.committed == 'true'
+        run: uv run make build_release
+
+      - name: Publish posthog to PyPI
+        if: steps.commit-release.outputs.committed == 'true'
+        uses: pypa/gh-action-pypi-publish@release/v1
+
+      # The `posthoganalytics` package is a mirror of `posthog` published under
+      # a different name for backwards compatibility. The make target handles
+      # copying, renaming imports, and building the dist automatically.
+      - name: Build posthoganalytics
+        if: steps.commit-release.outputs.committed == 'true'
+        run: uv run make build_release_analytics
+
+      - name: Publish posthoganalytics to PyPI
+        if: steps.commit-release.outputs.committed == 'true'
+        uses: pypa/gh-action-pypi-publish@release/v1
+
+      # We skip `sampo publish` (which normally creates the tag) because we
+      # need to publish both posthog and posthoganalytics manually, so we
+      # create the tag ourselves.
+      - name: Tag release
+        if: steps.commit-release.outputs.committed == 'true'
+        run: git tag "v${{ steps.sampo-release.outputs.new_version }}"
+
+      - name: Push tags
+        if: steps.commit-release.outputs.committed == 'true'
+        run: git push origin --tags
+
+      - name: Create GitHub Release
+        if: steps.commit-release.outputs.committed == 'true'
+        env:
+          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+        run: gh release create "v${{ steps.sampo-release.outputs.new_version }}" --generate-notes
+
+      - name: Dispatch generate-references
+        if: steps.commit-release.outputs.committed == 'true'
         env:
           GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
-        run: |
-          gh workflow run generate-references.yml --ref master
+        run: gh workflow run generate-references.yml --ref master
+
+      # Notify in case of a failure
+      - name: Send failure event to PostHog
+        if: ${{ failure() }}
+        uses: PostHog/posthog-github-action@v0.1
+        with:
+          posthog-token: "${{ secrets.POSTHOG_PROJECT_API_KEY }}"
+          event: "posthog-python-github-release-workflow-failure"
+          properties: >-
+            {
+              "commitSha": "${{ github.sha }}",
+              "jobStatus": "${{ job.status }}",
+              "ref": "${{ github.ref }}",
+              "version": "v${{ steps.sampo-release.outputs.new_version }}"
+            }
+
+      - name: Notify Slack - Failed
+        if: ${{ failure() && needs.notify-approval-needed.outputs.slack_ts != '' }}
+        uses: posthog/.github/.github/actions/slack-thread-reply@main
+        with:
+          slack_bot_token: ${{ secrets.SLACK_CLIENT_LIBRARIES_BOT_TOKEN }}
+          slack_channel_id: ${{ vars.SLACK_APPROVALS_CLIENT_LIBRARIES_CHANNEL_ID }}
+          thread_ts: ${{ needs.notify-approval-needed.outputs.slack_ts }}
+          message: "❌ Failed to release `posthog-python@v${{ steps.sampo-release.outputs.new_version }}`! <https://github.com/${{ github.repository }}/actions/runs/${{ github.run_id }}|View logs>"
+          emoji_reaction: "x"
+
+  notify-released:
+    name: Notify Slack - Released
+    needs: [check-release-label, notify-approval-needed, release]
+    runs-on: ubuntu-latest
+    if: always() && needs.release.result == 'success' && needs.notify-approval-needed.outputs.slack_ts != ''
+    steps:
+      - name: Checkout repository
+        uses: actions/checkout@v4
+
+      - name: Notify Slack - Released
+        uses: posthog/.github/.github/actions/slack-thread-reply@main
+        with:
+          slack_bot_token: ${{ secrets.SLACK_CLIENT_LIBRARIES_BOT_TOKEN }}
+          slack_channel_id: ${{ vars.SLACK_APPROVALS_CLIENT_LIBRARIES_CHANNEL_ID }}
+          thread_ts: ${{ needs.notify-approval-needed.outputs.slack_ts }}
+          message: "🚀 posthog-python released successfully!"
+          emoji_reaction: "rocket"

.github/workflows/sdk-compliance.yml

@@ -0,0 +1,21 @@
+name: SDK Compliance Tests
+
+permissions:
+  contents: read
+  packages: read
+  pull-requests: write
+
+on:
+  pull_request:
+  push:
+    branches:
+      - master
+
+jobs:
+  compliance:
+    name: PostHog SDK compliance tests
+    uses: PostHog/posthog-sdk-test-harness/.github/workflows/test-sdk-action.yml@main
+    with:
+      adapter-dockerfile: "sdk_compliance_adapter/Dockerfile"
+      adapter-context: "."
+      test-harness-version: "latest"

.sampo/config.toml

@@ -0,0 +1,19 @@
+# Sampo configuration
+version = 1
+
+[git]
+default_branch = "master"
+short_tags = "posthog"    # Tag with v1.2.3 rather than posthog-v1.2.3
+
+[github]
+repository = "posthog/posthog-python"
+
+[changelog]
+# Options for release notes generation.
+# show_commit_hash = true (default)
+# show_acknowledgments = true (default)
+
+[packages]
+# Options for package discovery and filtering.
+# ignore_unpublished = false (default)
+# ignore = ["internal-*", "examples/*"]

CHANGELOG.md

@@ -1,4 +1,122 @@
-# 7.0.0 - 2025-11-11
+# posthog
+
+## 7.9.2 — 2026-02-18
+
+### Patch changes
+
+- [f1dc4d7](https://github.com/posthog/posthog-python/commit/f1dc4d73914712983a7f715ee4fe1b70e66e770a) Add sampo to the project — Thanks @rafaeelaudibert!
+
+## 7.9.1 - 2026-02-17
+
+fix(llma): make prompt fetches deterministic by requiring project_api_key and sending it as token query param
+
+## 7.9.0 - 2026-02-17
+
+feat: Support device_id as bucketing identifier for local evaluation
+
+## 7.8.6 - 2026-02-09
+
+fix: limit collections scanning in code variables
+
+## 7.8.5 - 2026-02-09
+
+fix: further optimize code variables pattern matching
+
+## 7.8.4 - 2026-02-09
+
+fix: do not pattern match long values in code variables
+
+## 7.8.3 - 2026-02-06
+
+fix: openAI input image sanitization
+
+## 7.8.2 - 2026-02-04
+
+fix(llma): fix prompts default url
+
+## 7.8.1 - 2026-02-03
+
+fix(llma): small fixes for prompt management
+
+## 7.8.0 - 2026-01-28
+
+feat(llma): add prompt management
+
+Adds the Prompt Management feature. At the time of release, this feature is in a closed alpha.
+
+## 7.7.0 - 2026-01-15
+
+feat(ai): Add OpenAI Agents SDK integration
+
+Automatic tracing for agent workflows, handoffs, tool calls, guardrails, and custom spans. Includes `$ai_total_tokens`, `$ai_error_type` categorization, and `$ai_framework` property.
+
+## 7.6.0 - 2026-01-12
+
+feat: add device_id to flags request payload
+
+Add device_id parameter to all feature flag methods, allowing the server to track device identifiers for flag evaluation. The device_id can be passed explicitly or set via context using `set_context_device_id()`.
+
+## 7.5.1 - 2026-01-07
+
+fix: avoid return from finally block to fix Python 3.14 SyntaxWarning (#361) - thanks @jodal
+
+## 7.5.0 - 2026-01-06
+
+feat: Capture Langchain, OpenAI and Anthropic errors as exceptions (if exception autocapture is enabled)
+feat: Add reference to exception in LLMA trace and span events
+
+## 7.4.3 - 2026-01-02
+
+Fixes cache creation cost for Langchain with Anthropic
+
+## 7.4.2 - 2025-12-22
+
+feat: add `in_app_modules` option to control code variables capturing
+
+## 7.4.1 - 2025-12-19
+
+fix: extract model from response for OpenAI stored prompts
+
+When using OpenAI stored prompts, the model is defined in the OpenAI dashboard rather than passed in the API request. This fix adds a fallback to extract the model from the response object when not provided in kwargs, ensuring generations show up with the correct model and enabling cost calculations.
+
+## 7.4.0 - 2025-12-16
+
+feat: Add automatic retries for feature flag requests
+
+Feature flag API requests now automatically retry on transient failures:
+
+- Network errors (connection refused, DNS failures, timeouts)
+- Server errors (500, 502, 503, 504)
+- Up to 2 retries with exponential backoff (0.5s, 1s delays)
+
+Rate limit (429) and quota (402) errors are not retried.
+
+## 7.3.1 - 2025-12-06
+
+fix: remove unused $exception_message and $exception_type
+
+## 7.3.0 - 2025-12-05
+
+feat: improve code variables capture masking
+
+## 7.2.0 - 2025-12-01
+
+feat: add $feature_flag_evaluated_at properties to $feature_flag_called events
+
+## 7.1.0 - 2025-11-26
+
+Add support for the async version of Gemini.
+
+## 7.0.2 - 2025-11-18
+
+Add support for Python 3.14.
+Projects upgrading to Python 3.14 should ensure any Pydantic models passed into the SDK use Pydantic v2, as Pydantic v1 is not compatible with Python 3.14.
+
+## 7.0.1 - 2025-11-15
+
+Try to use repr() when formatting code variables
+
+## 7.0.0 - 2025-11-11
 
 NB Python 3.9 is no longer supported
 
@@ -12,155 +130,155 @@ NB Python 3.9 is no longer supported
   - langchain-community: 0.3.29 → 0.4.1
   - langgraph: 0.6.6 → 1.0.2
 
-# 6.9.3 - 2025-11-10
+## 6.9.3 - 2025-11-10
 
 - feat(ph-ai): PostHog properties dict in GenerationMetadata
 
-# 6.9.2 - 2025-11-10
+## 6.9.2 - 2025-11-10
 
 - fix(llma): fix cache token double subtraction in Langchain for non-Anthropic providers causing negative costs
 
-# 6.9.1 - 2025-11-07
+## 6.9.1 - 2025-11-07
 
 - fix(error-tracking): pass code variables config from init to client
 
-# 6.9.0 - 2025-11-06
+## 6.9.0 - 2025-11-06
 
 - feat(error-tracking): add local variables capture
 
-# 6.8.0 - 2025-11-03
+## 6.8.0 - 2025-11-03
 
 - feat(llma): send web search calls to be used for LLM cost calculations
 
-# 6.7.14 - 2025-11-03
+## 6.7.14 - 2025-11-03
 
 - fix(django): Handle request.user access in async middleware context to prevent SynchronousOnlyOperation errors in Django 5+ (fixes #355)
 - test(django): Add Django 5 integration test suite with real ASGI application testing async middleware behavior
 
-# 6.7.13 - 2025-11-02
+## 6.7.13 - 2025-11-02
 
 - fix(llma): cache cost calculation in the LangChain callback
 
-# 6.7.12 - 2025-11-02
+## 6.7.12 - 2025-11-02
 
 - fix(django): Restore process_exception method to capture view and downstream middleware exceptions (fixes #329)
 - fix(ai/langchain): Add LangChain 1.0+ compatibility for CallbackHandler imports (fixes #362)
 
-# 6.7.11 - 2025-10-28
+## 6.7.11 - 2025-10-28
 
 - feat(ai): Add `$ai_framework` property for framework integrations (e.g. LangChain)
 
-# 6.7.10 - 2025-10-24
+## 6.7.10 - 2025-10-24
 
 - fix(django): Make middleware truly hybrid - compatible with both sync (WSGI) and async (ASGI) Django stacks without breaking sync-only deployments
 
-# 6.7.9 - 2025-10-22
+## 6.7.9 - 2025-10-22
 
 - fix(flags): multi-condition flags with static cohorts returning wrong variants
 
-# 6.7.8 - 2025-10-16
+## 6.7.8 - 2025-10-16
 
 - fix(llma): missing async for OpenAI's streaming implementation
 
-# 6.7.7 - 2025-10-14
+## 6.7.7 - 2025-10-14
 
 - fix: remove deprecated attribute $exception_personURL from exception events
 
-# 6.7.6 - 2025-09-16
+## 6.7.6 - 2025-09-16
 
 - fix: don't sort condition sets with variant overrides to the top
 - fix: Prevent core Client methods from raising exceptions
 
-# 6.7.5 - 2025-09-16
+## 6.7.5 - 2025-09-16
 
 - feat: Django middleware now supports async request handling.
 
-# 6.7.4 - 2025-09-05
+## 6.7.4 - 2025-09-05
 
 - fix: Missing system prompts for some providers
 
-# 6.7.3 - 2025-09-04
+## 6.7.3 - 2025-09-04
 
 - fix: missing usage tokens in Gemini
 
-# 6.7.2 - 2025-09-03
+## 6.7.2 - 2025-09-03
 
 - fix: tool call results in streaming providers
 
-# 6.7.1 - 2025-09-01
+## 6.7.1 - 2025-09-01
 
 - fix: Add base64 inline image sanitization
 
-# 6.7.0 - 2025-08-26
+## 6.7.0 - 2025-08-26
 
 - feat: Add support for feature flag dependencies
 
-# 6.6.1 - 2025-08-21
+## 6.6.1 - 2025-08-21
 
 - fix: Prevent `NoneType` error when `group_properties` is `None`
 
-# 6.6.0 - 2025-08-15
+## 6.6.0 - 2025-08-15
 
 - feat: Add `flag_keys_to_evaluate` parameter to optimize feature flag evaluation performance by only evaluating specified flags
 - feat: Add `flag_keys_filter` option to `send_feature_flags` for selective flag evaluation in capture events
 
-# 6.5.0 - 2025-08-08
+## 6.5.0 - 2025-08-08
 
 - feat: Add `$context_tags` to an event to know which properties were included as tags
 
-# 6.4.1 - 2025-08-06
+## 6.4.1 - 2025-08-06
 
 - fix: Always pass project API key in `remote_config` requests for deterministic project routing
 
-# 6.4.0 - 2025-08-05
+## 6.4.0 - 2025-08-05
 
 - feat: support Vertex AI for Gemini
 
-# 6.3.4 - 2025-08-04
+## 6.3.4 - 2025-08-04
 
 - fix: set `$ai_tools` for all providers and `$ai_output_choices` for all non-streaming provider flows properly
 
-# 6.3.3 - 2025-08-01
+## 6.3.3 - 2025-08-01
 
 - fix: `get_feature_flag_result` now correctly returns FeatureFlagResult when payload is empty string instead of None
 
-# 6.3.2 - 2025-07-31
+## 6.3.2 - 2025-07-31
 
 - fix: Anthropic's tool calls are now handled properly
 
-# 6.3.0 - 2025-07-22
+## 6.3.0 - 2025-07-22
 
 - feat: Enhanced `send_feature_flags` parameter to accept `SendFeatureFlagsOptions` object for declarative control over local/remote evaluation and custom properties
 
-# 6.2.1 - 2025-07-21
+## 6.2.1 - 2025-07-21
 
 - feat: make `posthog_client` an optional argument in PostHog AI providers wrappers (`posthog.ai.*`), intuitively using the default client as the default
 
-# 6.1.1 - 2025-07-16
+## 6.1.1 - 2025-07-16
 
 - fix: correctly capture exceptions processed by Django from views or middleware
 
-# 6.1.0 - 2025-07-10
+## 6.1.0 - 2025-07-10
 
 - feat: decouple feature flag local evaluation from personal API keys; support decrypting remote config payloads without relying on the feature flags poller
 
-# 6.0.4 - 2025-07-09
+## 6.0.4 - 2025-07-09
 
 - fix: add POSTHOG_MW_CLIENT setting to django middleware, to support custom clients for exception capture.
 
-# 6.0.3 - 2025-07-07
+## 6.0.3 - 2025-07-07
 
 - feat: add a feature flag evaluation cache (local storage or redis) to support returning flag evaluations when the service is down
 
-# 6.0.2 - 2025-07-02
+## 6.0.2 - 2025-07-02
 
 - fix: send_feature_flags changed to default to false in `Client::capture_exception`
 
-# 6.0.1
+## 6.0.1
 
 - fix: response `$process_person_profile` property when passed to capture
 
-# 6.0.0
+## 6.0.0
 
 This release contains a number of major breaking changes:
 
@@ -187,15 +305,15 @@ with posthog.new_context():
 
 Generally, arguments are now appropriately typed, and docstrings have been updated. If something is unclear, please open an issue, or submit a PR!
 
-# 5.4.0 - 2025-06-20
+## 5.4.0 - 2025-06-20
 
 - feat: add support to session_id context on page method
 
-# 5.3.0 - 2025-06-19
+## 5.3.0 - 2025-06-19
 
 - fix: safely handle exception values
 
-# 5.2.0 - 2025-06-19
+## 5.2.0 - 2025-06-19
 
 - feat: construct artificial stack traces if no traceback is available on a captured exception
 

Makefile

@@ -5,12 +5,26 @@ test:
 	coverage run -m pytest
 	coverage report
 
-release:
+build_release:
 	rm -rf dist/*
 	python setup.py sdist bdist_wheel
-	twine upload dist/*
 
-release_analytics:
+# Builds the `posthoganalytics` PyPI package, which is a mirror of `posthog`
+# published under a different name for internal use by posthog/posthog.
+#
+# The process works in three phases:
+#   1. posthog -> posthoganalytics: Copy the source, rewrite all imports,
+#      remove the original posthog/ dir, and build the dist.
+#   2. posthoganalytics -> posthog: Reverse the import rewrites, copy
+#      everything back into posthog/, and clean up.
+#   3. Restore pyproject.toml from backup (setup_analytics.py modifies it).
+#
+# This ensures the working tree is left in the same state it started in.
+#
+# NOTE: This target clears dist/ before building. In the release workflow,
+# `build_release` (posthog) must be published BEFORE running this target,
+# otherwise the posthog dist artifacts will be lost.
+build_release_analytics:
 	rm -rf dist
 	rm -rf build
 	rm -rf posthoganalytics
@@ -21,7 +35,6 @@ release_analytics:
 	find ./posthoganalytics -name "*.bak" -delete
 	rm -rf posthog
 	python setup_analytics.py sdist bdist_wheel
-	twine upload dist/*
 	mkdir posthog
 	find ./posthoganalytics -type f -name "*.py" -exec sed -i.bak -e 's/from posthoganalytics /from posthog /g' {} \;
 	find ./posthoganalytics -type f -name "*.py" -exec sed -i.bak -e 's/from posthoganalytics\./from posthog\./g' {} \;
@@ -54,4 +67,4 @@ prep_local:
 	@echo "Local copy created at ../posthog-python-local"
 	@echo "Install with: pip install -e ../posthog-python-local"
 
-.PHONY: test lint release e2e_test prep_local
+.PHONY: test lint build_release build_release_analytics e2e_test prep_local

README.md

@@ -12,6 +12,14 @@
 
 Please see the [Python integration docs](https://posthog.com/docs/integrations/python-integration) for details.
 
+## Python Version Support
+
+| SDK Version | Python Versions Supported | Notes |
+|-------------|---------------------------|-------|
+| 7.3.1+      | 3.10, 3.11, 3.12, 3.13, 3.14 | Added Python 3.14 support |
+| 7.0.0 - 7.0.1 | 3.10, 3.11, 3.12, 3.13  | Dropped Python 3.9 support |
+| 4.0.1 - 6.x | 3.9, 3.10, 3.11, 3.12, 3.13 | Python 3.9+ required |
+
 ## Development
 
 ### Testing Locally

README_ANALYTICS.md

@@ -0,0 +1,11 @@
+# posthoganalytics
+
+> **Do not use this package.** Use [`posthog`](https://pypi.org/project/posthog/) instead.
+
+```bash
+pip install posthog
+```
+
+This package exists solely for internal use by [posthog/posthog](https://github.com/posthog/posthog) to avoid import conflicts with the local `posthog` package in that repository. It is an automatically generated mirror of `posthog` — same code, same versions, just published under a different name.
+
+If you are not working on the PostHog main repository, you should never need this package. All documentation, issues, and development happen in [`posthog-python`](https://github.com/posthog/posthog-python).

example.py

@@ -35,54 +35,40 @@ project_key = os.getenv("POSTHOG_PROJECT_API_KEY", "")
 personal_api_key = os.getenv("POSTHOG_PERSONAL_API_KEY", "")
 host = os.getenv("POSTHOG_HOST", "http://localhost:8000")
 
-# Check if credentials are provided
-if not project_key or not personal_api_key:
-    print("❌ Missing PostHog credentials!")
-    print(
-        "   Please set POSTHOG_PROJECT_API_KEY and POSTHOG_PERSONAL_API_KEY environment variables"
-    )
+# Check if project key is provided (required)
+if not project_key:
+    print("❌ Missing PostHog project API key!")
+    print("   Please set POSTHOG_PROJECT_API_KEY environment variable")
     print("   or copy .env.example to .env and fill in your values")
     exit(1)
 
-# Test authentication before proceeding
-print("🔑 Testing PostHog authentication...")
+# Configure PostHog with credentials
+posthog.debug = False
+posthog.api_key = project_key
+posthog.project_api_key = project_key
+posthog.host = host
+posthog.poll_interval = 10
 
-try:
-    # Configure PostHog with credentials
-    posthog.debug = False  # Keep quiet during auth test
-    posthog.api_key = project_key
-    posthog.project_api_key = project_key
+# Check if personal API key is available for local evaluation
+local_eval_available = bool(personal_api_key)
+if personal_api_key:
     posthog.personal_api_key = personal_api_key
-    posthog.host = host
-    posthog.poll_interval = 10
 
-    # Test by attempting to get feature flags (this validates both keys)
-    # This will fail if credentials are invalid
-    test_flags = posthog.get_all_flags("test_user", only_evaluate_locally=True)
-
-    # If we get here without exception, credentials work
-    print("✅ Authentication successful!")
-    print(f"   Project API Key: {project_key[:9]}...")
-    print("   Personal API Key: [REDACTED]")
-    print(f"   Host: {host}\n\n")
-
-except Exception as e:
-    print("❌ Authentication failed!")
-    print(f"   Error: {e}")
-    print("\n   Please check your credentials:")
-    print("   - POSTHOG_PROJECT_API_KEY: Project API key from PostHog settings")
-    print(
-        "   - POSTHOG_PERSONAL_API_KEY: Personal API key (required for local evaluation)"
-    )
-    print("   - POSTHOG_HOST: Your PostHog instance URL")
-    exit(1)
+print("🔑 PostHog Configuration:")
+print(f"   Project API Key: {project_key[:9]}...")
+if local_eval_available:
+    print("   Personal API Key: [SET]")
+else:
+    print("   Personal API Key: [NOT SET] - Local evaluation examples will be skipped")
+print(f"   Host: {host}\n")
 
 # Display menu and get user choice
 print("🚀 PostHog Python SDK Demo - Choose an example to run:\n")
 print("1. Identify and capture examples")
-print("2. Feature flag local evaluation examples")
+local_eval_note = "" if local_eval_available else " [requires personal API key]"
+print(f"2. Feature flag local evaluation examples{local_eval_note}")
 print("3. Feature flag payload examples")
-print("4. Flag dependencies examples")
+print(f"4. Flag dependencies examples{local_eval_note}")
 print("5. Context management and tagging examples")
 print("6. Run all examples")
 print("7. Exit")
@@ -148,6 +134,14 @@ if choice == "1":
     )
 
 elif choice == "2":
+    if not local_eval_available:
+        print("\n❌ This example requires a personal API key for local evaluation.")
+        print(
+            "   Set POSTHOG_PERSONAL_API_KEY environment variable to run this example."
+        )
+        posthog.shutdown()
+        exit(1)
+
     print("\n" + "=" * 60)
     print("FEATURE FLAG LOCAL EVALUATION EXAMPLES")
     print("=" * 60)
@@ -215,6 +209,14 @@ elif choice == "3":
         print(f"Value (variant or enabled): {result.get_value()}")
 
 elif choice == "4":
+    if not local_eval_available:
+        print("\n❌ This example requires a personal API key for local evaluation.")
+        print(
+            "   Set POSTHOG_PERSONAL_API_KEY environment variable to run this example."
+        )
+        posthog.shutdown()
+        exit(1)
+
     print("\n" + "=" * 60)
     print("FLAG DEPENDENCIES EXAMPLES")
     print("=" * 60)
@@ -429,6 +431,8 @@ elif choice == "5":
 
 elif choice == "6":
     print("\n🔄 Running all examples...")
+    if not local_eval_available:
+        print("   (Skipping local evaluation examples - no personal API key set)\n")
 
     # Run example 1
     print(f"\n{'🔸' * 20} IDENTIFY AND CAPTURE {'🔸' * 20}")
@@ -447,35 +451,37 @@ elif choice == "6":
         distinct_id="new_distinct_id", properties={"email": "something@something.com"}
     )
 
-    # Run example 2
-    print(f"\n{'🔸' * 20} FEATURE FLAGS {'🔸' * 20}")
-    print("🏁 Testing basic feature flags...")
-    print(f"beta-feature: {posthog.feature_enabled('beta-feature', 'distinct_id')}")
-    print(
-        f"Sydney user: {posthog.feature_enabled('test-flag', 'random_id_12345', person_properties={'$geoip_city_name': 'Sydney'})}"
-    )
+    # Run example 2 (requires local evaluation)
+    if local_eval_available:
+        print(f"\n{'🔸' * 20} FEATURE FLAGS {'🔸' * 20}")
+        print("🏁 Testing basic feature flags...")
+        print(f"beta-feature: {posthog.feature_enabled('beta-feature', 'distinct_id')}")
+        print(
+            f"Sydney user: {posthog.feature_enabled('test-flag', 'random_id_12345', person_properties={'$geoip_city_name': 'Sydney'})}"
+        )
 
     # Run example 3
     print(f"\n{'🔸' * 20} PAYLOADS {'🔸' * 20}")
     print("📦 Testing payloads...")
     print(f"Payload: {posthog.get_feature_flag_payload('beta-feature', 'distinct_id')}")
 
-    # Run example 4
-    print(f"\n{'🔸' * 20} FLAG DEPENDENCIES {'🔸' * 20}")
-    print("🔗 Testing flag dependencies...")
-    result1 = posthog.feature_enabled(
-        "test-flag-dependency",
-        "demo_user",
-        person_properties={"email": "user@example.com"},
-        only_evaluate_locally=True,
-    )
-    result2 = posthog.feature_enabled(
-        "test-flag-dependency",
-        "demo_user2",
-        person_properties={"email": "user@other.com"},
-        only_evaluate_locally=True,
-    )
-    print(f"✅ @example.com user: {result1}, regular user: {result2}")
+    # Run example 4 (requires local evaluation)
+    if local_eval_available:
+        print(f"\n{'🔸' * 20} FLAG DEPENDENCIES {'🔸' * 20}")
+        print("🔗 Testing flag dependencies...")
+        result1 = posthog.feature_enabled(
+            "test-flag-dependency",
+            "demo_user",
+            person_properties={"email": "user@example.com"},
+            only_evaluate_locally=True,
+        )
+        result2 = posthog.feature_enabled(
+            "test-flag-dependency",
+            "demo_user2",
+            person_properties={"email": "user@other.com"},
+            only_evaluate_locally=True,
+        )
+        print(f"✅ @example.com user: {result1}, regular user: {result2}")
 
     # Run example 5
     print(f"\n{'🔸' * 20} CONTEXT MANAGEMENT {'🔸' * 20}")

examples/redis_flag_cache.py

@@ -0,0 +1,144 @@
+"""
+Redis-based distributed cache for PostHog feature flag definitions.
+
+This example demonstrates how to implement a FlagDefinitionCacheProvider
+using Redis for multi-instance deployments (leader election pattern).
+
+Usage:
+    import redis
+    from posthog import Posthog
+
+    redis_client = redis.Redis(host='localhost', port=6379, decode_responses=True)
+    cache = RedisFlagCache(redis_client, service_key="my-service")
+
+    posthog = Posthog(
+        "<project_api_key>",
+        personal_api_key="<personal_api_key>",
+        flag_definition_cache_provider=cache,
+    )
+
+Requirements:
+    pip install redis
+"""
+
+import json
+import uuid
+
+from posthog import FlagDefinitionCacheData, FlagDefinitionCacheProvider
+from redis import Redis
+from typing import Optional
+
+
+class RedisFlagCache(FlagDefinitionCacheProvider):
+    """
+    A distributed cache for PostHog feature flag definitions using Redis.
+
+    In a multi-instance deployment (e.g., multiple serverless functions or containers),
+    we want only ONE instance to poll PostHog for flag updates, while all instances
+    share the cached results. This prevents N instances from making N redundant API calls.
+
+    The implementation uses leader election:
+    - One instance "wins" and becomes responsible for fetching
+    - Other instances read from the shared cache
+    - If the leader dies, the lock expires (TTL) and another instance takes over
+
+    Uses Lua scripts for atomic operations, following Redis distributed lock best practices:
+    https://redis.io/docs/latest/develop/clients/patterns/distributed-locks/
+    """
+
+    LOCK_TTL_MS = 60 * 1000  # 60 seconds, should be longer than the flags poll interval
+    CACHE_TTL_SECONDS = 60 * 60 * 24  # 24 hours
+
+    # Lua script: acquire lock if free, or extend if we own it
+    _LUA_TRY_LEAD = """
+        local current = redis.call('GET', KEYS[1])
+        if current == false then
+            redis.call('SET', KEYS[1], ARGV[1], 'PX', ARGV[2])
+            return 1
+        elseif current == ARGV[1] then
+            redis.call('PEXPIRE', KEYS[1], ARGV[2])
+            return 1
+        end
+        return 0
+    """
+
+    # Lua script: release lock only if we own it
+    _LUA_STOP_LEAD = """
+        if redis.call('GET', KEYS[1]) == ARGV[1] then
+            return redis.call('DEL', KEYS[1])
+        end
+        return 0
+    """
+
+    def __init__(self, redis: Redis[str], service_key: str):
+        """
+        Initialize the Redis flag cache.
+
+        Args:
+            redis: A redis-py client instance. Must be configured with
+                   decode_responses=True for correct string handling.
+            service_key: A unique identifier for this service/environment.
+                         Used to scope Redis keys, allowing multiple services
+                         or environments to share the same Redis instance.
+                         Examples: "my-api-prod", "checkout-service", "staging".
+
+        Redis Keys Created:
+            - posthog:flags:{service_key} - Cached flag definitions (JSON)
+            - posthog:flags:{service_key}:lock - Leader election lock
+
+        Example:
+            redis_client = redis.Redis(
+                host='localhost',
+                port=6379,
+                decode_responses=True
+            )
+            cache = RedisFlagCache(redis_client, service_key="my-api-prod")
+        """
+        self._redis = redis
+        self._cache_key = f"posthog:flags:{service_key}"
+        self._lock_key = f"posthog:flags:{service_key}:lock"
+        self._instance_id = str(uuid.uuid4())
+        self._try_lead = self._redis.register_script(self._LUA_TRY_LEAD)
+        self._stop_lead = self._redis.register_script(self._LUA_STOP_LEAD)
+
+    def get_flag_definitions(self) -> Optional[FlagDefinitionCacheData]:
+        """
+        Retrieve cached flag definitions from Redis.
+
+        Returns:
+            Cached flag definitions if available, None otherwise.
+        """
+        cached = self._redis.get(self._cache_key)
+        return json.loads(cached) if cached else None
+
+    def should_fetch_flag_definitions(self) -> bool:
+        """
+        Determines if this instance should fetch flag definitions from PostHog.
+
+        Atomically either:
+        - Acquires the lock if no one holds it, OR
+        - Extends the lock TTL if we already hold it
+
+        Returns:
+            True if this instance is the leader and should fetch, False otherwise.
+        """
+        result = self._try_lead(
+            keys=[self._lock_key],
+            args=[self._instance_id, self.LOCK_TTL_MS],
+        )
+        return result == 1
+
+    def on_flag_definitions_received(self, data: FlagDefinitionCacheData) -> None:
+        """
+        Store fetched flag definitions in Redis.
+
+        Args:
+            data: The flag definitions to cache.
+        """
+        self._redis.set(self._cache_key, json.dumps(data), ex=self.CACHE_TTL_SECONDS)
+
+    def shutdown(self) -> None:
+        """
+        Release leadership if we hold it. Safe to call even if not the leader.
+        """
+        self._stop_lead(keys=[self._lock_key], args=[self._instance_id])

mypy-baseline.txt

@@ -26,14 +26,9 @@ posthog/client.py:0: error: Incompatible types in assignment (expression has typ
 posthog/client.py:0: error: Incompatible types in assignment (expression has type "dict[Any, Any]", variable has type "None")  [assignment]
 posthog/client.py:0: error: "None" has no attribute "__iter__" (not iterable)  [attr-defined]
 posthog/client.py:0: error: Statement is unreachable  [unreachable]
-posthog/client.py:0: error: Incompatible types in assignment (expression has type "Any | dict[Any, Any]", variable has type "None")  [assignment]
-posthog/client.py:0: error: Incompatible types in assignment (expression has type "Any | dict[Any, Any]", variable has type "None")  [assignment]
-posthog/client.py:0: error: Incompatible types in assignment (expression has type "dict[Never, Never]", variable has type "None")  [assignment]
-posthog/client.py:0: error: Incompatible types in assignment (expression has type "dict[Never, Never]", variable has type "None")  [assignment]
 posthog/client.py:0: error: Right operand of "and" is never evaluated  [unreachable]
 posthog/client.py:0: error: Incompatible types in assignment (expression has type "Poller", variable has type "None")  [assignment]
 posthog/client.py:0: error: "None" has no attribute "start"  [attr-defined]
-posthog/client.py:0: error: "None" has no attribute "get"  [attr-defined]
 posthog/client.py:0: error: Statement is unreachable  [unreachable]
 posthog/client.py:0: error: Statement is unreachable  [unreachable]
 posthog/client.py:0: error: Name "urlparse" already defined (possibly by an import)  [no-redef]

posthog/__init__.py

@@ -1,30 +1,65 @@
 import datetime  # noqa: F401
-from typing import Callable, Dict, Optional, Any  # noqa: F401
+from typing import Any, Callable, Dict, Optional  # noqa: F401
+
 from typing_extensions import Unpack
 
-from posthog.args import OptionalCaptureArgs, OptionalSetArgs, ExceptionArg
+from posthog.args import ExceptionArg, OptionalCaptureArgs, OptionalSetArgs
 from posthog.client import Client
 from posthog.contexts import (
-    new_context as inner_new_context,
-    scoped as inner_scoped,
-    tag as inner_tag,
-    set_context_session as inner_set_context_session,
     identify_context as inner_identify_context,
+)
+from posthog.contexts import (
+    new_context as inner_new_context,
+)
+from posthog.contexts import (
+    scoped as inner_scoped,
+)
+from posthog.contexts import (
     set_capture_exception_code_variables_context as inner_set_capture_exception_code_variables_context,
-    set_code_variables_mask_patterns_context as inner_set_code_variables_mask_patterns_context,
+)
+from posthog.contexts import (
     set_code_variables_ignore_patterns_context as inner_set_code_variables_ignore_patterns_context,
 )
+from posthog.contexts import (
+    set_code_variables_mask_patterns_context as inner_set_code_variables_mask_patterns_context,
+)
+from posthog.contexts import (
+    set_context_device_id as inner_set_context_device_id,
+)
+from posthog.contexts import (
+    set_context_session as inner_set_context_session,
+)
+from posthog.contexts import (
+    tag as inner_tag,
+)
+from posthog.contexts import (
+    get_tags as inner_get_tags,
+)
 from posthog.exception_utils import (
     DEFAULT_CODE_VARIABLES_IGNORE_PATTERNS,
     DEFAULT_CODE_VARIABLES_MASK_PATTERNS,
 )
 from posthog.feature_flags import (
     InconclusiveMatchError as InconclusiveMatchError,
+)
+from posthog.feature_flags import (
     RequiresServerEvaluation as RequiresServerEvaluation,
 )
+from posthog.flag_definition_cache import (
+    FlagDefinitionCacheData as FlagDefinitionCacheData,
+    FlagDefinitionCacheProvider as FlagDefinitionCacheProvider,
+)
+from posthog.request import (
+    disable_connection_reuse as disable_connection_reuse,
+    enable_keep_alive as enable_keep_alive,
+    set_socket_options as set_socket_options,
+    SocketOptions as SocketOptions,
+)
 from posthog.types import (
     FeatureFlag,
     FlagsAndPayloads,
+)
+from posthog.types import (
     FeatureFlagResult as FeatureFlagResult,
 )
 from posthog.version import VERSION
@@ -101,6 +136,26 @@ def set_context_session(session_id: str):
     return inner_set_context_session(session_id)
 
 
+def set_context_device_id(device_id: str):
+    """
+    Set the device ID for the current context, associating all feature flag requests
+    in this or child contexts with the given device ID.
+
+    Args:
+        device_id: The device ID to associate with the current context and its children
+
+    Examples:
+        ```python
+        from posthog import set_context_device_id
+        set_context_device_id("device_123")
+        ```
+
+    Category:
+        Contexts
+    """
+    return inner_set_context_device_id(device_id)
+
+
 def identify_context(distinct_id: str):
     """
     Identify the current context with a distinct ID.
@@ -161,6 +216,19 @@ def tag(name: str, value: Any):
     return inner_tag(name, value)
 
 
+def get_tags() -> Dict[str, Any]:
+    """
+    Get all tags from the current context.
+
+    Returns:
+        Dict of all tags in the current context
+
+    Category:
+        Contexts
+    """
+    return inner_get_tags()
+
+
 """Settings."""
 api_key = None  # type: Optional[str]
 host = None  # type: Optional[str]
@@ -191,6 +259,7 @@ default_client = None  # type: Optional[Client]
 capture_exception_code_variables = False
 code_variables_mask_patterns = DEFAULT_CODE_VARIABLES_MASK_PATTERNS
 code_variables_ignore_patterns = DEFAULT_CODE_VARIABLES_IGNORE_PATTERNS
+in_app_modules = None  # type: Optional[list[str]]
 
 
 # NOTE - this and following functions take unpacked kwargs because we needed to make
@@ -437,6 +506,7 @@ def feature_enabled(
     only_evaluate_locally=False,  # type: bool
     send_feature_flag_events=True,  # type: bool
     disable_geoip=None,  # type: Optional[bool]
+    device_id=None,  # type: Optional[str]
 ):
     # type: (...) -> bool
     """
@@ -476,6 +546,7 @@ def feature_enabled(
         only_evaluate_locally=only_evaluate_locally,
         send_feature_flag_events=send_feature_flag_events,
         disable_geoip=disable_geoip,
+        device_id=device_id,
     )
 
 
@@ -488,6 +559,7 @@ def get_feature_flag(
     only_evaluate_locally=False,  # type: bool
     send_feature_flag_events=True,  # type: bool
     disable_geoip=None,  # type: Optional[bool]
+    device_id=None,  # type: Optional[str]
 ) -> Optional[FeatureFlag]:
     """
     Get feature flag variant for users. Used with experiments.
@@ -526,6 +598,7 @@ def get_feature_flag(
         only_evaluate_locally=only_evaluate_locally,
         send_feature_flag_events=send_feature_flag_events,
         disable_geoip=disable_geoip,
+        device_id=device_id,
     )
 
 
@@ -536,6 +609,7 @@ def get_all_flags(
     group_properties=None,  # type: Optional[dict]
     only_evaluate_locally=False,  # type: bool
     disable_geoip=None,  # type: Optional[bool]
+    device_id=None,  # type: Optional[str]
 ) -> Optional[dict[str, FeatureFlag]]:
     """
     Get all flags for a given user.
@@ -568,6 +642,7 @@ def get_all_flags(
         group_properties=group_properties or {},
         only_evaluate_locally=only_evaluate_locally,
         disable_geoip=disable_geoip,
+        device_id=device_id,
     )
 
 
@@ -580,6 +655,7 @@ def get_feature_flag_result(
     only_evaluate_locally=False,
     send_feature_flag_events=True,
     disable_geoip=None,  # type: Optional[bool]
+    device_id=None,  # type: Optional[str]
 ):
     # type: (...) -> Optional[FeatureFlagResult]
     """
@@ -611,6 +687,7 @@ def get_feature_flag_result(
         only_evaluate_locally=only_evaluate_locally,
         send_feature_flag_events=send_feature_flag_events,
         disable_geoip=disable_geoip,
+        device_id=device_id,
     )
 
 
@@ -624,6 +701,7 @@ def get_feature_flag_payload(
     only_evaluate_locally=False,
     send_feature_flag_events=True,
     disable_geoip=None,  # type: Optional[bool]
+    device_id=None,  # type: Optional[str]
 ) -> Optional[str]:
     return _proxy(
         "get_feature_flag_payload",
@@ -636,6 +714,7 @@ def get_feature_flag_payload(
         only_evaluate_locally=only_evaluate_locally,
         send_feature_flag_events=send_feature_flag_events,
         disable_geoip=disable_geoip,
+        device_id=device_id,
     )
 
 
@@ -666,6 +745,7 @@ def get_all_flags_and_payloads(
     group_properties=None,  # type: Optional[dict]
     only_evaluate_locally=False,
     disable_geoip=None,  # type: Optional[bool]
+    device_id=None,  # type: Optional[str]
 ) -> FlagsAndPayloads:
     return _proxy(
         "get_all_flags_and_payloads",
@@ -675,6 +755,7 @@ def get_all_flags_and_payloads(
         group_properties=group_properties or {},
         only_evaluate_locally=only_evaluate_locally,
         disable_geoip=disable_geoip,
+        device_id=device_id,
     )
 
 
@@ -789,6 +870,7 @@ def setup() -> Client:
             capture_exception_code_variables=capture_exception_code_variables,
             code_variables_mask_patterns=code_variables_mask_patterns,
             code_variables_ignore_patterns=code_variables_ignore_patterns,
+            in_app_modules=in_app_modules,
         )
 
     # always set incase user changes it

posthog/ai/__init__.py

@@ -0,0 +1,3 @@
+from posthog.ai.prompts import Prompts
+
+__all__ = ["Prompts"]

posthog/ai/anthropic/anthropic_converter.py

@@ -17,6 +17,7 @@ from posthog.ai.types import (
     TokenUsage,
     ToolInProgress,
 )
+from posthog.ai.utils import serialize_raw_usage
 
 
 def format_anthropic_response(response: Any) -> List[FormattedMessage]:
@@ -221,6 +222,12 @@ def extract_anthropic_usage_from_response(response: Any) -> TokenUsage:
     if web_search_count > 0:
         result["web_search_count"] = web_search_count
 
+    # Capture raw usage metadata for backend processing
+    # Serialize to dict here in the converter (not in utils)
+    serialized = serialize_raw_usage(response.usage)
+    if serialized:
+        result["raw_usage"] = serialized
+
     return result
 
 
@@ -247,6 +254,11 @@ def extract_anthropic_usage_from_event(event: Any) -> TokenUsage:
             usage["cache_read_input_tokens"] = getattr(
                 event.message.usage, "cache_read_input_tokens", 0
             )
+            # Capture raw usage metadata for backend processing
+            # Serialize to dict here in the converter (not in utils)
+            serialized = serialize_raw_usage(event.message.usage)
+            if serialized:
+                usage["raw_usage"] = serialized
 
     # Handle usage stats from message_delta event
     if hasattr(event, "usage") and event.usage:
@@ -262,6 +274,12 @@ def extract_anthropic_usage_from_event(event: Any) -> TokenUsage:
                 if web_search_count > 0:
                     usage["web_search_count"] = web_search_count
 
+        # Capture raw usage metadata for backend processing
+        # Serialize to dict here in the converter (not in utils)
+        serialized = serialize_raw_usage(event.usage)
+        if serialized:
+            usage["raw_usage"] = serialized
+
     return usage
 
 

posthog/ai/gemini/__init__.py

@@ -1,4 +1,5 @@
 from .gemini import Client
+from .gemini_async import AsyncClient
 from .gemini_converter import (
     format_gemini_input,
     format_gemini_response,
@@ -9,12 +10,14 @@ from .gemini_converter import (
 # Create a genai-like module for perfect drop-in replacement
 class _GenAI:
     Client = Client
+    AsyncClient = AsyncClient
 
 
 genai = _GenAI()
 
 __all__ = [
     "Client",
+    "AsyncClient",
     "genai",
     "format_gemini_input",
     "format_gemini_response",

posthog/ai/gemini/gemini.py

@@ -304,7 +304,7 @@ class Models:
 
         def generator():
             nonlocal usage_stats
-            nonlocal accumulated_content  # noqa: F824
+            nonlocal accumulated_content
             try:
                 for chunk in response:
                     # Extract usage stats from chunk

posthog/ai/gemini/gemini_async.py

@@ -0,0 +1,423 @@
+import os
+import time
+import uuid
+from typing import Any, Dict, Optional
+
+from posthog.ai.types import TokenUsage, StreamingEventData
+from posthog.ai.utils import merge_system_prompt
+
+try:
+    from google import genai
+except ImportError:
+    raise ModuleNotFoundError(
+        "Please install the Google Gemini SDK to use this feature: 'pip install google-genai'"
+    )
+
+from posthog import setup
+from posthog.ai.utils import (
+    call_llm_and_track_usage_async,
+    capture_streaming_event,
+    merge_usage_stats,
+)
+from posthog.ai.gemini.gemini_converter import (
+    extract_gemini_usage_from_chunk,
+    extract_gemini_content_from_chunk,
+    format_gemini_streaming_output,
+)
+from posthog.ai.sanitization import sanitize_gemini
+from posthog.client import Client as PostHogClient
+
+
+class AsyncClient:
+    """
+    An async drop-in replacement for genai.Client that automatically sends LLM usage events to PostHog.
+
+    Usage:
+        client = AsyncClient(
+            api_key="your_api_key",
+            posthog_client=posthog_client,
+            posthog_distinct_id="default_user",  # Optional defaults
+            posthog_properties={"team": "ai"}    # Optional defaults
+        )
+        response = await client.models.generate_content(
+            model="gemini-2.0-flash",
+            contents=["Hello world"],
+            posthog_distinct_id="specific_user"  # Override default
+        )
+    """
+
+    _ph_client: PostHogClient
+
+    def __init__(
+        self,
+        api_key: Optional[str] = None,
+        vertexai: Optional[bool] = None,
+        credentials: Optional[Any] = None,
+        project: Optional[str] = None,
+        location: Optional[str] = None,
+        debug_config: Optional[Any] = None,
+        http_options: Optional[Any] = None,
+        posthog_client: Optional[PostHogClient] = None,
+        posthog_distinct_id: Optional[str] = None,
+        posthog_properties: Optional[Dict[str, Any]] = None,
+        posthog_privacy_mode: bool = False,
+        posthog_groups: Optional[Dict[str, Any]] = None,
+        **kwargs,
+    ):
+        """
+        Args:
+            api_key: Google AI API key. If not provided, will use GOOGLE_API_KEY or API_KEY environment variable (not required for Vertex AI)
+            vertexai: Whether to use Vertex AI authentication
+            credentials: Vertex AI credentials object
+            project: GCP project ID for Vertex AI
+            location: GCP location for Vertex AI
+            debug_config: Debug configuration for the client
+            http_options: HTTP options for the client
+            posthog_client: PostHog client for tracking usage
+            posthog_distinct_id: Default distinct ID for all calls (can be overridden per call)
+            posthog_properties: Default properties for all calls (can be overridden per call)
+            posthog_privacy_mode: Default privacy mode for all calls (can be overridden per call)
+            posthog_groups: Default groups for all calls (can be overridden per call)
+            **kwargs: Additional arguments (for future compatibility)
+        """
+
+        self._ph_client = posthog_client or setup()
+
+        if self._ph_client is None:
+            raise ValueError("posthog_client is required for PostHog tracking")
+
+        self.models = AsyncModels(
+            api_key=api_key,
+            vertexai=vertexai,
+            credentials=credentials,
+            project=project,
+            location=location,
+            debug_config=debug_config,
+            http_options=http_options,
+            posthog_client=self._ph_client,
+            posthog_distinct_id=posthog_distinct_id,
+            posthog_properties=posthog_properties,
+            posthog_privacy_mode=posthog_privacy_mode,
+            posthog_groups=posthog_groups,
+            **kwargs,
+        )
+
+
+class AsyncModels:
+    """
+    Async Models interface that mimics genai.Client().aio.models with PostHog tracking.
+    """
+
+    _ph_client: PostHogClient  # Not None after __init__ validation
+
+    def __init__(
+        self,
+        api_key: Optional[str] = None,
+        vertexai: Optional[bool] = None,
+        credentials: Optional[Any] = None,
+        project: Optional[str] = None,
+        location: Optional[str] = None,
+        debug_config: Optional[Any] = None,
+        http_options: Optional[Any] = None,
+        posthog_client: Optional[PostHogClient] = None,
+        posthog_distinct_id: Optional[str] = None,
+        posthog_properties: Optional[Dict[str, Any]] = None,
+        posthog_privacy_mode: bool = False,
+        posthog_groups: Optional[Dict[str, Any]] = None,
+        **kwargs,
+    ):
+        """
+        Args:
+            api_key: Google AI API key. If not provided, will use GOOGLE_API_KEY or API_KEY environment variable (not required for Vertex AI)
+            vertexai: Whether to use Vertex AI authentication
+            credentials: Vertex AI credentials object
+            project: GCP project ID for Vertex AI
+            location: GCP location for Vertex AI
+            debug_config: Debug configuration for the client
+            http_options: HTTP options for the client
+            posthog_client: PostHog client for tracking usage
+            posthog_distinct_id: Default distinct ID for all calls
+            posthog_properties: Default properties for all calls
+            posthog_privacy_mode: Default privacy mode for all calls
+            posthog_groups: Default groups for all calls
+            **kwargs: Additional arguments (for future compatibility)
+        """
+
+        self._ph_client = posthog_client or setup()
+
+        if self._ph_client is None:
+            raise ValueError("posthog_client is required for PostHog tracking")
+
+        # Store default PostHog settings
+        self._default_distinct_id = posthog_distinct_id
+        self._default_properties = posthog_properties or {}
+        self._default_privacy_mode = posthog_privacy_mode
+        self._default_groups = posthog_groups
+
+        # Build genai.Client arguments
+        client_args: Dict[str, Any] = {}
+
+        # Add Vertex AI parameters if provided
+        if vertexai is not None:
+            client_args["vertexai"] = vertexai
+
+        if credentials is not None:
+            client_args["credentials"] = credentials
+
+        if project is not None:
+            client_args["project"] = project
+
+        if location is not None:
+            client_args["location"] = location
+
+        if debug_config is not None:
+            client_args["debug_config"] = debug_config
+
+        if http_options is not None:
+            client_args["http_options"] = http_options
+
+        # Handle API key authentication
+        if vertexai:
+            # For Vertex AI, api_key is optional
+            if api_key is not None:
+                client_args["api_key"] = api_key
+        else:
+            # For non-Vertex AI mode, api_key is required (backwards compatibility)
+            if api_key is None:
+                api_key = os.environ.get("GOOGLE_API_KEY") or os.environ.get("API_KEY")
+
+            if api_key is None:
+                raise ValueError(
+                    "API key must be provided either as parameter or via GOOGLE_API_KEY/API_KEY environment variable"
+                )
+
+            client_args["api_key"] = api_key
+
+        self._client = genai.Client(**client_args)
+        self._base_url = "https://generativelanguage.googleapis.com"
+
+    def _merge_posthog_params(
+        self,
+        call_distinct_id: Optional[str],
+        call_trace_id: Optional[str],
+        call_properties: Optional[Dict[str, Any]],
+        call_privacy_mode: Optional[bool],
+        call_groups: Optional[Dict[str, Any]],
+    ):
+        """Merge call-level PostHog parameters with client defaults."""
+
+        # Use call-level values if provided, otherwise fall back to defaults
+        distinct_id = (
+            call_distinct_id
+            if call_distinct_id is not None
+            else self._default_distinct_id
+        )
+        privacy_mode = (
+            call_privacy_mode
+            if call_privacy_mode is not None
+            else self._default_privacy_mode
+        )
+        groups = call_groups if call_groups is not None else self._default_groups
+
+        # Merge properties: default properties + call properties (call properties override)
+        properties = dict(self._default_properties)
+
+        if call_properties:
+            properties.update(call_properties)
+
+        if call_trace_id is None:
+            call_trace_id = str(uuid.uuid4())
+
+        return distinct_id, call_trace_id, properties, privacy_mode, groups
+
+    async def generate_content(
+        self,
+        model: str,
+        contents,
+        posthog_distinct_id: Optional[str] = None,
+        posthog_trace_id: Optional[str] = None,
+        posthog_properties: Optional[Dict[str, Any]] = None,
+        posthog_privacy_mode: Optional[bool] = None,
+        posthog_groups: Optional[Dict[str, Any]] = None,
+        **kwargs: Any,
+    ):
+        """
+        Generate content using Gemini's API while tracking usage in PostHog.
+
+        This method signature exactly matches genai.Client().aio.models.generate_content()
+        with additional PostHog tracking parameters.
+
+        Args:
+            model: The model to use (e.g., 'gemini-2.0-flash')
+            contents: The input content for generation
+            posthog_distinct_id: ID to associate with the usage event (overrides client default)
+            posthog_trace_id: Trace UUID for linking events (auto-generated if not provided)
+            posthog_properties: Extra properties to include in the event (merged with client defaults)
+            posthog_privacy_mode: Whether to redact sensitive information (overrides client default)
+            posthog_groups: Group analytics properties (overrides client default)
+            **kwargs: Arguments passed to Gemini's generate_content
+        """
+
+        # Merge PostHog parameters
+        distinct_id, trace_id, properties, privacy_mode, groups = (
+            self._merge_posthog_params(
+                posthog_distinct_id,
+                posthog_trace_id,
+                posthog_properties,
+                posthog_privacy_mode,
+                posthog_groups,
+            )
+        )
+
+        kwargs_with_contents = {"model": model, "contents": contents, **kwargs}
+
+        return await call_llm_and_track_usage_async(
+            distinct_id,
+            self._ph_client,
+            "gemini",
+            trace_id,
+            properties,
+            privacy_mode,
+            groups,
+            self._base_url,
+            self._client.aio.models.generate_content,
+            **kwargs_with_contents,
+        )
+
+    async def _generate_content_streaming(
+        self,
+        model: str,
+        contents,
+        distinct_id: Optional[str],
+        trace_id: Optional[str],
+        properties: Optional[Dict[str, Any]],
+        privacy_mode: bool,
+        groups: Optional[Dict[str, Any]],
+        **kwargs: Any,
+    ):
+        start_time = time.time()
+        usage_stats: TokenUsage = TokenUsage(input_tokens=0, output_tokens=0)
+        accumulated_content = []
+
+        kwargs_without_stream = {"model": model, "contents": contents, **kwargs}
+        response = await self._client.aio.models.generate_content_stream(
+            **kwargs_without_stream
+        )
+
+        async def async_generator():
+            nonlocal usage_stats
+            nonlocal accumulated_content
+
+            try:
+                async for chunk in response:
+                    # Extract usage stats from chunk
+                    chunk_usage = extract_gemini_usage_from_chunk(chunk)
+
+                    if chunk_usage:
+                        # Gemini reports cumulative totals, not incremental values
+                        merge_usage_stats(usage_stats, chunk_usage, mode="cumulative")
+
+                    # Extract content from chunk (now returns content blocks)
+                    content_block = extract_gemini_content_from_chunk(chunk)
+
+                    if content_block is not None:
+                        accumulated_content.append(content_block)
+
+                    yield chunk
+
+            finally:
+                end_time = time.time()
+                latency = end_time - start_time
+
+                self._capture_streaming_event(
+                    model,
+                    contents,
+                    distinct_id,
+                    trace_id,
+                    properties,
+                    privacy_mode,
+                    groups,
+                    kwargs,
+                    usage_stats,
+                    latency,
+                    accumulated_content,
+                )
+
+        return async_generator()
+
+    def _capture_streaming_event(
+        self,
+        model: str,
+        contents,
+        distinct_id: Optional[str],
+        trace_id: Optional[str],
+        properties: Optional[Dict[str, Any]],
+        privacy_mode: bool,
+        groups: Optional[Dict[str, Any]],
+        kwargs: Dict[str, Any],
+        usage_stats: TokenUsage,
+        latency: float,
+        output: Any,
+    ):
+        # Prepare standardized event data
+        formatted_input = self._format_input(contents, **kwargs)
+        sanitized_input = sanitize_gemini(formatted_input)
+
+        event_data = StreamingEventData(
+            provider="gemini",
+            model=model,
+            base_url=self._base_url,
+            kwargs=kwargs,
+            formatted_input=sanitized_input,
+            formatted_output=format_gemini_streaming_output(output),
+            usage_stats=usage_stats,
+            latency=latency,
+            distinct_id=distinct_id,
+            trace_id=trace_id,
+            properties=properties,
+            privacy_mode=privacy_mode,
+            groups=groups,
+        )
+
+        # Use the common capture function
+        capture_streaming_event(self._ph_client, event_data)
+
+    def _format_input(self, contents, **kwargs):
+        """Format input contents for PostHog tracking"""
+
+        # Create kwargs dict with contents for merge_system_prompt
+        input_kwargs = {"contents": contents, **kwargs}
+        return merge_system_prompt(input_kwargs, "gemini")
+
+    async def generate_content_stream(
+        self,
+        model: str,
+        contents,
+        posthog_distinct_id: Optional[str] = None,
+        posthog_trace_id: Optional[str] = None,
+        posthog_properties: Optional[Dict[str, Any]] = None,
+        posthog_privacy_mode: Optional[bool] = None,
+        posthog_groups: Optional[Dict[str, Any]] = None,
+        **kwargs: Any,
+    ):
+        # Merge PostHog parameters
+        distinct_id, trace_id, properties, privacy_mode, groups = (
+            self._merge_posthog_params(
+                posthog_distinct_id,
+                posthog_trace_id,
+                posthog_properties,
+                posthog_privacy_mode,
+                posthog_groups,
+            )
+        )
+
+        return await self._generate_content_streaming(
+            model,
+            contents,
+            distinct_id,
+            trace_id,
+            properties,
+            privacy_mode,
+            groups,
+            **kwargs,
+        )

posthog/ai/gemini/gemini_converter.py

@@ -12,6 +12,7 @@ from posthog.ai.types import (
     FormattedMessage,
     TokenUsage,
 )
+from posthog.ai.utils import serialize_raw_usage
 
 
 class GeminiPart(TypedDict, total=False):
@@ -29,35 +30,76 @@ class GeminiMessage(TypedDict, total=False):
     text: str
 
 
-def _extract_text_from_parts(parts: List[Any]) -> str:
+def _format_parts_as_content_blocks(parts: List[Any]) -> List[FormattedContentItem]:
     """
-    Extract and concatenate text from a parts array.
+    Format Gemini parts array into structured content blocks.
+
+    Preserves structure for multimodal content (text + images) instead of
+    concatenating everything into a string.
 
     Args:
-        parts: List of parts that may contain text content
+        parts: List of parts that may contain text, inline_data, etc.
 
     Returns:
-        Concatenated text from all parts
+        List of formatted content blocks
     """
-
-    content_parts = []
+    content_blocks: List[FormattedContentItem] = []
 
     for part in parts:
+        # Handle dict with text field
         if isinstance(part, dict) and "text" in part:
-            content_parts.append(part["text"])
+            content_blocks.append({"type": "text", "text": part["text"]})
 
+        # Handle string parts
         elif isinstance(part, str):
-            content_parts.append(part)
+            content_blocks.append({"type": "text", "text": part})
 
+        # Handle dict with inline_data (images, documents, etc.)
+        elif isinstance(part, dict) and "inline_data" in part:
+            inline_data = part["inline_data"]
+            mime_type = inline_data.get("mime_type", "")
+            content_type = "image" if mime_type.startswith("image/") else "document"
+
+            content_blocks.append(
+                {
+                    "type": content_type,
+                    "inline_data": inline_data,
+                }
+            )
+
+        # Handle object with text attribute
         elif hasattr(part, "text"):
-            # Get the text attribute value
             text_value = getattr(part, "text", "")
-            content_parts.append(text_value if text_value else str(part))
+            if text_value:
+                content_blocks.append({"type": "text", "text": text_value})
 
-        else:
-            content_parts.append(str(part))
+        # Handle object with inline_data attribute
+        elif hasattr(part, "inline_data"):
+            inline_data = part.inline_data
+            # Convert to dict if needed
+            if hasattr(inline_data, "mime_type") and hasattr(inline_data, "data"):
+                # Determine type based on mime_type
+                mime_type = inline_data.mime_type
+                content_type = "image" if mime_type.startswith("image/") else "document"
 
-    return "".join(content_parts)
+                content_blocks.append(
+                    {
+                        "type": content_type,
+                        "inline_data": {
+                            "mime_type": mime_type,
+                            "data": inline_data.data,
+                        },
+                    }
+                )
+            else:
+                content_blocks.append(
+                    {
+                        "type": "image",
+                        "inline_data": inline_data,
+                    }
+                )
+
+    return content_blocks
 
 
 def _format_dict_message(item: Dict[str, Any]) -> FormattedMessage:
@@ -73,16 +115,17 @@ def _format_dict_message(item: Dict[str, Any]) -> FormattedMessage:
 
     # Handle dict format with parts array (Gemini-specific format)
     if "parts" in item and isinstance(item["parts"], list):
-        content = _extract_text_from_parts(item["parts"])
-        return {"role": item.get("role", "user"), "content": content}
+        content_blocks = _format_parts_as_content_blocks(item["parts"])
+        return {"role": item.get("role", "user"), "content": content_blocks}
 
     # Handle dict with content field
     if "content" in item:
         content = item["content"]
 
         if isinstance(content, list):
-            # If content is a list, extract text from it
-            content = _extract_text_from_parts(content)
+            # If content is a list, format it as content blocks
+            content_blocks = _format_parts_as_content_blocks(content)
+            return {"role": item.get("role", "user"), "content": content_blocks}
 
         elif not isinstance(content, str):
             content = str(content)
@@ -110,14 +153,14 @@ def _format_object_message(item: Any) -> FormattedMessage:
 
     # Handle object with parts attribute
     if hasattr(item, "parts") and hasattr(item.parts, "__iter__"):
-        content = _extract_text_from_parts(item.parts)
+        content_blocks = _format_parts_as_content_blocks(list(item.parts))
         role = getattr(item, "role", "user") if hasattr(item, "role") else "user"
 
         # Ensure role is a string
         if not isinstance(role, str):
             role = "user"
 
-        return {"role": role, "content": content}
+        return {"role": role, "content": content_blocks}
 
     # Handle object with text attribute
     if hasattr(item, "text"):
@@ -140,7 +183,8 @@ def _format_object_message(item: Any) -> FormattedMessage:
         content = item.content
 
         if isinstance(content, list):
-            content = _extract_text_from_parts(content)
+            content_blocks = _format_parts_as_content_blocks(content)
+            return {"role": role, "content": content_blocks}
 
         elif not isinstance(content, str):
             content = str(content)
@@ -193,6 +237,29 @@ def format_gemini_response(response: Any) -> List[FormattedMessage]:
                                 }
                             )
 
+                        elif hasattr(part, "inline_data") and part.inline_data:
+                            # Handle audio/media inline data
+                            import base64
+
+                            inline_data = part.inline_data
+                            mime_type = getattr(inline_data, "mime_type", "audio/pcm")
+                            raw_data = getattr(inline_data, "data", b"")
+
+                            # Encode binary data as base64 string for JSON serialization
+                            if isinstance(raw_data, bytes):
+                                data = base64.b64encode(raw_data).decode("utf-8")
+                            else:
+                                # Already a string (base64)
+                                data = raw_data
+
+                            content.append(
+                                {
+                                    "type": "audio",
+                                    "mime_type": mime_type,
+                                    "data": data,
+                                }
+                            )
+
                 if content:
                     output.append(
                         {
@@ -421,6 +488,12 @@ def _extract_usage_from_metadata(metadata: Any) -> TokenUsage:
         if reasoning_tokens and reasoning_tokens > 0:
             usage["reasoning_tokens"] = reasoning_tokens
 
+    # Capture raw usage metadata for backend processing
+    # Serialize to dict here in the converter (not in utils)
+    serialized = serialize_raw_usage(metadata)
+    if serialized:
+        usage["raw_usage"] = serialized
+
     return usage
 
 

posthog/ai/langchain/callbacks.py

@@ -22,8 +22,8 @@ from uuid import UUID
 
 try:
     # LangChain 1.0+ and modern 0.x with langchain-core
-    from langchain_core.callbacks.base import BaseCallbackHandler
     from langchain_core.agents import AgentAction, AgentFinish
+    from langchain_core.callbacks.base import BaseCallbackHandler
 except (ImportError, ModuleNotFoundError):
     # Fallback for older LangChain versions
     from langchain.callbacks.base import BaseCallbackHandler
@@ -35,15 +35,15 @@ from langchain_core.messages import (
     FunctionMessage,
     HumanMessage,
     SystemMessage,
-    ToolMessage,
     ToolCall,
+    ToolMessage,
 )
 from langchain_core.outputs import ChatGeneration, LLMResult
 from pydantic import BaseModel
 
 from posthog import setup
-from posthog.ai.utils import get_model_params, with_privacy_mode
 from posthog.ai.sanitization import sanitize_langchain
+from posthog.ai.utils import get_model_params, with_privacy_mode
 from posthog.client import Client
 
 log = logging.getLogger("posthog")
@@ -506,6 +506,14 @@ class CallbackHandler(BaseCallbackHandler):
         if isinstance(outputs, BaseException):
             event_properties["$ai_error"] = _stringify_exception(outputs)
             event_properties["$ai_is_error"] = True
+            event_properties = _capture_exception_and_update_properties(
+                self._ph_client,
+                outputs,
+                self._distinct_id,
+                self._groups,
+                event_properties,
+            )
+
         elif outputs is not None:
             event_properties["$ai_output_state"] = with_privacy_mode(
                 self._ph_client, self._privacy_mode, outputs
@@ -576,10 +584,24 @@ class CallbackHandler(BaseCallbackHandler):
         if run.tools:
             event_properties["$ai_tools"] = run.tools
 
+        if self._properties:
+            event_properties.update(self._properties)
+
+        if self._distinct_id is None:
+            event_properties["$process_person_profile"] = False
+
         if isinstance(output, BaseException):
             event_properties["$ai_http_status"] = _get_http_status(output)
             event_properties["$ai_error"] = _stringify_exception(output)
             event_properties["$ai_is_error"] = True
+
+            event_properties = _capture_exception_and_update_properties(
+                self._ph_client,
+                output,
+                self._distinct_id,
+                self._groups,
+                event_properties,
+            )
         else:
             # Add usage
             usage = _parse_usage(output, run.provider, run.model)
@@ -607,12 +629,6 @@ class CallbackHandler(BaseCallbackHandler):
                 self._ph_client, self._privacy_mode, completions
             )
 
-        if self._properties:
-            event_properties.update(self._properties)
-
-        if self._distinct_id is None:
-            event_properties["$process_person_profile"] = False
-
         self._ph_client.capture(
             distinct_id=self._distinct_id or trace_id,
             event="$ai_generation",
@@ -773,9 +789,11 @@ def _parse_usage_model(
             for mapped_key, dataclass_key in field_mapping.items()
         },
     )
-    # For Anthropic providers, LangChain reports input_tokens as the sum of input and cache read tokens.
+    # For Anthropic providers, LangChain reports input_tokens as the sum of all input tokens.
     # Our cost calculation expects them to be separate for Anthropic, so we subtract cache tokens.
-    # For other providers (OpenAI, etc.), input_tokens already includes cache tokens as expected.
+    # Both cache_read and cache_write tokens should be subtracted since Anthropic's raw API
+    # reports input_tokens as tokens NOT read from or used to create a cache.
+    # For other providers (OpenAI, etc.), input_tokens already excludes cache tokens as expected.
     # Match logic consistent with plugin-server: exact match on provider OR substring match on model
     is_anthropic = False
     if provider and provider.lower() == "anthropic":
@@ -783,14 +801,14 @@ def _parse_usage_model(
     elif model and "anthropic" in model.lower():
         is_anthropic = True
 
-    if (
-        is_anthropic
-        and normalized_usage.input_tokens
-        and normalized_usage.cache_read_tokens
-    ):
-        normalized_usage.input_tokens = max(
-            normalized_usage.input_tokens - normalized_usage.cache_read_tokens, 0
+    if is_anthropic and normalized_usage.input_tokens:
+        cache_tokens = (normalized_usage.cache_read_tokens or 0) + (
+            normalized_usage.cache_write_tokens or 0
         )
+        if cache_tokens > 0:
+            normalized_usage.input_tokens = max(
+                normalized_usage.input_tokens - cache_tokens, 0
+            )
     return normalized_usage
 
 
@@ -861,6 +879,27 @@ def _parse_usage(
     return llm_usage
 
 
+def _capture_exception_and_update_properties(
+    client: Client,
+    exception: BaseException,
+    distinct_id: Optional[Union[str, int, UUID]],
+    groups: Optional[Dict[str, Any]],
+    event_properties: Dict[str, Any],
+):
+    if client.enable_exception_autocapture:
+        exception_id = client.capture_exception(
+            exception,
+            distinct_id=distinct_id,
+            groups=groups,
+            properties=event_properties,
+        )
+
+        if exception_id:
+            event_properties["$exception_event_id"] = exception_id
+
+    return event_properties
+
+
 def _get_http_status(error: BaseException) -> int:
     # OpenAI: https://github.com/openai/openai-python/blob/main/src/openai/_exceptions.py
     # Anthropic: https://github.com/anthropics/anthropic-sdk-python/blob/main/src/anthropic/_exceptions.py

posthog/ai/openai/openai.py

@@ -124,14 +124,23 @@ class WrappedResponses:
         start_time = time.time()
         usage_stats: TokenUsage = TokenUsage()
         final_content = []
+        model_from_response: Optional[str] = None
         response = self._original.create(**kwargs)
 
         def generator():
             nonlocal usage_stats
             nonlocal final_content  # noqa: F824
+            nonlocal model_from_response
 
             try:
                 for chunk in response:
+                    # Extract model from response object in chunk (for stored prompts)
+                    if hasattr(chunk, "response") and chunk.response:
+                        if model_from_response is None and hasattr(
+                            chunk.response, "model"
+                        ):
+                            model_from_response = chunk.response.model
+
                     # Extract usage stats from chunk
                     chunk_usage = extract_openai_usage_from_chunk(chunk, "responses")
 
@@ -161,6 +170,7 @@ class WrappedResponses:
                     latency,
                     output,
                     None,  # Responses API doesn't have tools
+                    model_from_response,
                 )
 
         return generator()
@@ -177,6 +187,7 @@ class WrappedResponses:
         latency: float,
         output: Any,
         available_tool_calls: Optional[List[Dict[str, Any]]] = None,
+        model_from_response: Optional[str] = None,
     ):
         from posthog.ai.types import StreamingEventData
         from posthog.ai.openai.openai_converter import (
@@ -189,9 +200,12 @@ class WrappedResponses:
         formatted_input = format_openai_streaming_input(kwargs, "responses")
         sanitized_input = sanitize_openai_response(formatted_input)
 
+        # Use model from kwargs, fallback to model from response
+        model = kwargs.get("model") or model_from_response or "unknown"
+
         event_data = StreamingEventData(
             provider="openai",
-            model=kwargs.get("model", "unknown"),
+            model=model,
             base_url=str(self._client.base_url),
             kwargs=kwargs,
             formatted_input=sanitized_input,
@@ -320,6 +334,7 @@ class WrappedCompletions:
         usage_stats: TokenUsage = TokenUsage()
         accumulated_content = []
         accumulated_tool_calls: Dict[int, Dict[str, Any]] = {}
+        model_from_response: Optional[str] = None
         if "stream_options" not in kwargs:
             kwargs["stream_options"] = {}
         kwargs["stream_options"]["include_usage"] = True
@@ -329,9 +344,14 @@ class WrappedCompletions:
             nonlocal usage_stats
             nonlocal accumulated_content  # noqa: F824
             nonlocal accumulated_tool_calls
+            nonlocal model_from_response
 
             try:
                 for chunk in response:
+                    # Extract model from chunk (Chat Completions chunks have model field)
+                    if model_from_response is None and hasattr(chunk, "model"):
+                        model_from_response = chunk.model
+
                     # Extract usage stats from chunk
                     chunk_usage = extract_openai_usage_from_chunk(chunk, "chat")
 
@@ -376,6 +396,7 @@ class WrappedCompletions:
                     accumulated_content,
                     tool_calls_list,
                     extract_available_tool_calls("openai", kwargs),
+                    model_from_response,
                 )
 
         return generator()
@@ -393,6 +414,7 @@ class WrappedCompletions:
         output: Any,
         tool_calls: Optional[List[Dict[str, Any]]] = None,
         available_tool_calls: Optional[List[Dict[str, Any]]] = None,
+        model_from_response: Optional[str] = None,
     ):
         from posthog.ai.types import StreamingEventData
         from posthog.ai.openai.openai_converter import (
@@ -405,9 +427,12 @@ class WrappedCompletions:
         formatted_input = format_openai_streaming_input(kwargs, "chat")
         sanitized_input = sanitize_openai(formatted_input)
 
+        # Use model from kwargs, fallback to model from response
+        model = kwargs.get("model") or model_from_response or "unknown"
+
         event_data = StreamingEventData(
             provider="openai",
-            model=kwargs.get("model", "unknown"),
+            model=model,
             base_url=str(self._client.base_url),
             kwargs=kwargs,
             formatted_input=sanitized_input,

posthog/ai/openai/openai_async.py

@@ -128,14 +128,23 @@ class WrappedResponses:
         start_time = time.time()
         usage_stats: TokenUsage = TokenUsage()
         final_content = []
+        model_from_response: Optional[str] = None
         response = await self._original.create(**kwargs)
 
         async def async_generator():
             nonlocal usage_stats
             nonlocal final_content  # noqa: F824
+            nonlocal model_from_response
 
             try:
                 async for chunk in response:
+                    # Extract model from response object in chunk (for stored prompts)
+                    if hasattr(chunk, "response") and chunk.response:
+                        if model_from_response is None and hasattr(
+                            chunk.response, "model"
+                        ):
+                            model_from_response = chunk.response.model
+
                     # Extract usage stats from chunk
                     chunk_usage = extract_openai_usage_from_chunk(chunk, "responses")
 
@@ -166,6 +175,7 @@ class WrappedResponses:
                     latency,
                     output,
                     extract_available_tool_calls("openai", kwargs),
+                    model_from_response,
                 )
 
         return async_generator()
@@ -182,13 +192,17 @@ class WrappedResponses:
         latency: float,
         output: Any,
         available_tool_calls: Optional[List[Dict[str, Any]]] = None,
+        model_from_response: Optional[str] = None,
     ):
         if posthog_trace_id is None:
             posthog_trace_id = str(uuid.uuid4())
 
+        # Use model from kwargs, fallback to model from response
+        model = kwargs.get("model") or model_from_response or "unknown"
+
         event_properties = {
             "$ai_provider": "openai",
-            "$ai_model": kwargs.get("model"),
+            "$ai_model": model,
             "$ai_model_parameters": get_model_params(kwargs),
             "$ai_input": with_privacy_mode(
                 self._client._ph_client,
@@ -350,6 +364,7 @@ class WrappedCompletions:
         usage_stats: TokenUsage = TokenUsage()
         accumulated_content = []
         accumulated_tool_calls: Dict[int, Dict[str, Any]] = {}
+        model_from_response: Optional[str] = None
 
         if "stream_options" not in kwargs:
             kwargs["stream_options"] = {}
@@ -360,9 +375,14 @@ class WrappedCompletions:
             nonlocal usage_stats
             nonlocal accumulated_content  # noqa: F824
             nonlocal accumulated_tool_calls
+            nonlocal model_from_response
 
             try:
                 async for chunk in response:
+                    # Extract model from chunk (Chat Completions chunks have model field)
+                    if model_from_response is None and hasattr(chunk, "model"):
+                        model_from_response = chunk.model
+
                     # Extract usage stats from chunk
                     chunk_usage = extract_openai_usage_from_chunk(chunk, "chat")
                     if chunk_usage:
@@ -405,6 +425,7 @@ class WrappedCompletions:
                     accumulated_content,
                     tool_calls_list,
                     extract_available_tool_calls("openai", kwargs),
+                    model_from_response,
                 )
 
         return async_generator()
@@ -422,13 +443,17 @@ class WrappedCompletions:
         output: Any,
         tool_calls: Optional[List[Dict[str, Any]]] = None,
         available_tool_calls: Optional[List[Dict[str, Any]]] = None,
+        model_from_response: Optional[str] = None,
     ):
         if posthog_trace_id is None:
             posthog_trace_id = str(uuid.uuid4())
 
+        # Use model from kwargs, fallback to model from response
+        model = kwargs.get("model") or model_from_response or "unknown"
+
         event_properties = {
             "$ai_provider": "openai",
-            "$ai_model": kwargs.get("model"),
+            "$ai_model": model,
             "$ai_model_parameters": get_model_params(kwargs),
             "$ai_input": with_privacy_mode(
                 self._client._ph_client,

posthog/ai/openai/openai_converter.py

@@ -16,6 +16,7 @@ from posthog.ai.types import (
     FormattedTextContent,
     TokenUsage,
 )
+from posthog.ai.utils import serialize_raw_usage
 
 
 def format_openai_response(response: Any) -> List[FormattedMessage]:
@@ -67,6 +68,12 @@ def format_openai_response(response: Any) -> List[FormattedMessage]:
                             }
                         )
 
+                # Handle audio output (gpt-4o-audio-preview)
+                if hasattr(choice.message, "audio") and choice.message.audio:
+                    # Convert Pydantic model to dict to capture all fields from OpenAI
+                    audio_dict = choice.message.audio.model_dump()
+                    content.append({"type": "audio", **audio_dict})
+
         if content:
             output.append(
                 {
@@ -423,6 +430,12 @@ def extract_openai_usage_from_response(response: Any) -> TokenUsage:
     if web_search_count > 0:
         result["web_search_count"] = web_search_count
 
+    # Capture raw usage metadata for backend processing
+    # Serialize to dict here in the converter (not in utils)
+    serialized = serialize_raw_usage(response.usage)
+    if serialized:
+        result["raw_usage"] = serialized
+
     return result
 
 
@@ -476,6 +489,12 @@ def extract_openai_usage_from_chunk(
                 chunk.usage.completion_tokens_details.reasoning_tokens
             )
 
+        # Capture raw usage metadata for backend processing
+        # Serialize to dict here in the converter (not in utils)
+        serialized = serialize_raw_usage(chunk.usage)
+        if serialized:
+            usage["raw_usage"] = serialized
+
     elif provider_type == "responses":
         # For Responses API, usage is only in chunk.response.usage for completed events
         if hasattr(chunk, "type") and chunk.type == "response.completed":
@@ -510,6 +529,12 @@ def extract_openai_usage_from_chunk(
                     if web_search_count > 0:
                         usage["web_search_count"] = web_search_count
 
+                # Capture raw usage metadata for backend processing
+                # Serialize to dict here in the converter (not in utils)
+                serialized = serialize_raw_usage(response_usage)
+                if serialized:
+                    usage["raw_usage"] = serialized
+
     return usage
 
 

posthog/ai/openai_agents/__init__.py

@@ -0,0 +1,76 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Union
+
+if TYPE_CHECKING:
+    from agents.tracing import Trace
+
+    from posthog.client import Client
+
+try:
+    import agents  # noqa: F401
+except ImportError:
+    raise ModuleNotFoundError(
+        "Please install the OpenAI Agents SDK to use this feature: 'pip install openai-agents'"
+    )
+
+from posthog.ai.openai_agents.processor import PostHogTracingProcessor
+
+__all__ = ["PostHogTracingProcessor", "instrument"]
+
+
+def instrument(
+    client: Optional[Client] = None,
+    distinct_id: Optional[Union[str, Callable[[Trace], Optional[str]]]] = None,
+    privacy_mode: bool = False,
+    groups: Optional[Dict[str, Any]] = None,
+    properties: Optional[Dict[str, Any]] = None,
+) -> PostHogTracingProcessor:
+    """
+    One-liner to instrument OpenAI Agents SDK with PostHog tracing.
+
+    This registers a PostHogTracingProcessor with the OpenAI Agents SDK,
+    automatically capturing traces, spans, and LLM generations.
+
+    Args:
+        client: Optional PostHog client instance. If not provided, uses the default client.
+        distinct_id: Optional distinct ID to associate with all traces.
+            Can also be a callable that takes a trace and returns a distinct ID.
+        privacy_mode: If True, redacts input/output content from events.
+        groups: Optional PostHog groups to associate with events.
+        properties: Optional additional properties to include with all events.
+
+    Returns:
+        PostHogTracingProcessor: The registered processor instance.
+
+    Example:
+        ```python
+        from posthog.ai.openai_agents import instrument
+
+        # Simple setup
+        instrument(distinct_id="user@example.com")
+
+        # With custom properties
+        instrument(
+            distinct_id="user@example.com",
+            privacy_mode=True,
+            properties={"environment": "production"}
+        )
+
+        # Now run agents as normal - traces automatically sent to PostHog
+        from agents import Agent, Runner
+        agent = Agent(name="Assistant", instructions="You are helpful.")
+        result = Runner.run_sync(agent, "Hello!")
+        ```
+    """
+    from agents.tracing import add_trace_processor
+
+    processor = PostHogTracingProcessor(
+        client=client,
+        distinct_id=distinct_id,
+        privacy_mode=privacy_mode,
+        groups=groups,
+        properties=properties,
+    )
+    add_trace_processor(processor)
+    return processor

posthog/ai/openai_agents/processor.py

@@ -0,0 +1,863 @@
+import json
+import logging
+import time
+from datetime import datetime
+from typing import Any, Callable, Dict, Optional, Union
+
+from agents.tracing import Span, Trace
+from agents.tracing.processor_interface import TracingProcessor
+from agents.tracing.span_data import (
+    AgentSpanData,
+    CustomSpanData,
+    FunctionSpanData,
+    GenerationSpanData,
+    GuardrailSpanData,
+    HandoffSpanData,
+    MCPListToolsSpanData,
+    ResponseSpanData,
+    SpeechGroupSpanData,
+    SpeechSpanData,
+    TranscriptionSpanData,
+)
+
+from posthog import setup
+from posthog.client import Client
+
+log = logging.getLogger("posthog")
+
+
+def _ensure_serializable(obj: Any) -> Any:
+    """Ensure an object is JSON-serializable, converting to str as fallback.
+
+    Returns the original object if it's already serializable (dict, list, str,
+    int, etc.), or str(obj) for non-serializable types so that downstream
+    json.dumps() calls won't fail.
+    """
+    if obj is None:
+        return None
+    try:
+        json.dumps(obj)
+        return obj
+    except (TypeError, ValueError):
+        return str(obj)
+
+
+def _parse_iso_timestamp(iso_str: Optional[str]) -> Optional[float]:
+    """Parse ISO timestamp to Unix timestamp."""
+    if not iso_str:
+        return None
+    try:
+        dt = datetime.fromisoformat(iso_str.replace("Z", "+00:00"))
+        return dt.timestamp()
+    except (ValueError, AttributeError):
+        return None
+
+
+class PostHogTracingProcessor(TracingProcessor):
+    """
+    A tracing processor that sends OpenAI Agents SDK traces to PostHog.
+
+    This processor implements the TracingProcessor interface from the OpenAI Agents SDK
+    and maps agent traces, spans, and generations to PostHog's LLM analytics events.
+
+    Example:
+        ```python
+        from agents import Agent, Runner
+        from agents.tracing import add_trace_processor
+        from posthog.ai.openai_agents import PostHogTracingProcessor
+
+        # Create and register the processor
+        processor = PostHogTracingProcessor(
+            distinct_id="user@example.com",
+            privacy_mode=False,
+        )
+        add_trace_processor(processor)
+
+        # Run agents as normal - traces automatically sent to PostHog
+        agent = Agent(name="Assistant", instructions="You are helpful.")
+        result = Runner.run_sync(agent, "Hello!")
+        ```
+    """
+
+    def __init__(
+        self,
+        client: Optional[Client] = None,
+        distinct_id: Optional[Union[str, Callable[[Trace], Optional[str]]]] = None,
+        privacy_mode: bool = False,
+        groups: Optional[Dict[str, Any]] = None,
+        properties: Optional[Dict[str, Any]] = None,
+    ):
+        """
+        Initialize the PostHog tracing processor.
+
+        Args:
+            client: Optional PostHog client instance. If not provided, uses the default client.
+            distinct_id: Either a string distinct ID or a callable that takes a Trace
+                and returns a distinct ID. If not provided, uses the trace_id.
+            privacy_mode: If True, redacts input/output content from events.
+            groups: Optional PostHog groups to associate with all events.
+            properties: Optional additional properties to include with all events.
+        """
+        self._client = client or setup()
+        self._distinct_id = distinct_id
+        self._privacy_mode = privacy_mode
+        self._groups = groups or {}
+        self._properties = properties or {}
+
+        # Track span start times for latency calculation
+        self._span_start_times: Dict[str, float] = {}
+
+        # Track trace metadata for associating with spans
+        self._trace_metadata: Dict[str, Dict[str, Any]] = {}
+
+        # Max entries to prevent unbounded growth if on_span_end/on_trace_end
+        # is never called (e.g., due to an exception in the Agents SDK).
+        self._max_tracked_entries = 10000
+
+    def _get_distinct_id(self, trace: Optional[Trace]) -> Optional[str]:
+        """Resolve the distinct ID for a trace.
+
+        Returns the user-provided distinct ID (string or callable result),
+        or None if no user-provided ID is available. Callers should treat
+        None as a signal to use a fallback ID in personless mode.
+        """
+        if callable(self._distinct_id):
+            if trace:
+                result = self._distinct_id(trace)
+                if result:
+                    return str(result)
+            return None
+        elif self._distinct_id:
+            return str(self._distinct_id)
+        return None
+
+    def _with_privacy_mode(self, value: Any) -> Any:
+        """Apply privacy mode redaction if enabled."""
+        if self._privacy_mode or (
+            hasattr(self._client, "privacy_mode") and self._client.privacy_mode
+        ):
+            return None
+        return value
+
+    def _evict_stale_entries(self) -> None:
+        """Evict oldest entries if dicts exceed max size to prevent unbounded growth."""
+        if len(self._span_start_times) > self._max_tracked_entries:
+            # Remove oldest entries by start time
+            sorted_spans = sorted(self._span_start_times.items(), key=lambda x: x[1])
+            for span_id, _ in sorted_spans[: len(sorted_spans) // 2]:
+                del self._span_start_times[span_id]
+            log.debug(
+                "Evicted stale span start times (exceeded %d entries)",
+                self._max_tracked_entries,
+            )
+
+        if len(self._trace_metadata) > self._max_tracked_entries:
+            # Remove half the entries (oldest inserted via dict ordering in Python 3.7+)
+            keys = list(self._trace_metadata.keys())
+            for key in keys[: len(keys) // 2]:
+                del self._trace_metadata[key]
+            log.debug(
+                "Evicted stale trace metadata (exceeded %d entries)",
+                self._max_tracked_entries,
+            )
+
+    def _get_group_id(self, trace_id: str) -> Optional[str]:
+        """Get the group_id for a trace from stored metadata."""
+        if trace_id in self._trace_metadata:
+            return self._trace_metadata[trace_id].get("group_id")
+        return None
+
+    def _capture_event(
+        self,
+        event: str,
+        properties: Dict[str, Any],
+        distinct_id: Optional[str] = None,
+    ) -> None:
+        """Capture an event to PostHog with error handling.
+
+        Args:
+            distinct_id: The resolved distinct ID. When the user didn't provide
+                one, callers should pass ``user_distinct_id or fallback_id``
+                (matching the langchain/openai pattern) and separately set
+                ``$process_person_profile`` in properties.
+        """
+        try:
+            if not hasattr(self._client, "capture") or not callable(
+                self._client.capture
+            ):
+                return
+
+            final_properties = {
+                **properties,
+                **self._properties,
+            }
+
+            self._client.capture(
+                distinct_id=distinct_id or "unknown",
+                event=event,
+                properties=final_properties,
+                groups=self._groups,
+            )
+        except Exception as e:
+            log.debug(f"Failed to capture PostHog event: {e}")
+
+    def on_trace_start(self, trace: Trace) -> None:
+        """Called when a new trace begins. Stores metadata for spans; the $ai_trace event is emitted in on_trace_end."""
+        try:
+            self._evict_stale_entries()
+            trace_id = trace.trace_id
+            trace_name = trace.name
+            group_id = getattr(trace, "group_id", None)
+            metadata = getattr(trace, "metadata", None)
+
+            distinct_id = self._get_distinct_id(trace)
+
+            # Store trace metadata for later (used by spans and on_trace_end)
+            self._trace_metadata[trace_id] = {
+                "name": trace_name,
+                "group_id": group_id,
+                "metadata": metadata,
+                "distinct_id": distinct_id,
+                "start_time": time.time(),
+            }
+        except Exception as e:
+            log.debug(f"Error in on_trace_start: {e}")
+
+    def on_trace_end(self, trace: Trace) -> None:
+        """Called when a trace completes. Emits the $ai_trace event with full metadata."""
+        try:
+            trace_id = trace.trace_id
+
+            # Pop stored metadata (also cleans up)
+            trace_info = self._trace_metadata.pop(trace_id, {})
+            trace_name = trace_info.get("name") or trace.name
+            group_id = trace_info.get("group_id") or getattr(trace, "group_id", None)
+            metadata = trace_info.get("metadata") or getattr(trace, "metadata", None)
+            distinct_id = trace_info.get("distinct_id") or self._get_distinct_id(trace)
+
+            # Calculate trace-level latency
+            start_time = trace_info.get("start_time")
+            latency = (time.time() - start_time) if start_time else None
+
+            properties = {
+                "$ai_trace_id": trace_id,
+                "$ai_trace_name": trace_name,
+                "$ai_provider": "openai",
+                "$ai_framework": "openai-agents",
+            }
+
+            if latency is not None:
+                properties["$ai_latency"] = latency
+
+            # Include group_id for linking related traces (e.g., conversation threads)
+            if group_id:
+                properties["$ai_group_id"] = group_id
+
+            # Include trace metadata if present
+            if metadata:
+                properties["$ai_trace_metadata"] = _ensure_serializable(metadata)
+
+            if distinct_id is None:
+                properties["$process_person_profile"] = False
+
+            self._capture_event(
+                event="$ai_trace",
+                distinct_id=distinct_id or trace_id,
+                properties=properties,
+            )
+        except Exception as e:
+            log.debug(f"Error in on_trace_end: {e}")
+
+    def on_span_start(self, span: Span[Any]) -> None:
+        """Called when a new span begins."""
+        try:
+            self._evict_stale_entries()
+            span_id = span.span_id
+            self._span_start_times[span_id] = time.time()
+        except Exception as e:
+            log.debug(f"Error in on_span_start: {e}")
+
+    def on_span_end(self, span: Span[Any]) -> None:
+        """Called when a span completes."""
+        try:
+            span_id = span.span_id
+            trace_id = span.trace_id
+            parent_id = span.parent_id
+            span_data = span.span_data
+
+            # Calculate latency
+            start_time = self._span_start_times.pop(span_id, None)
+            if start_time:
+                latency = time.time() - start_time
+            else:
+                # Fall back to parsing timestamps
+                started = _parse_iso_timestamp(span.started_at)
+                ended = _parse_iso_timestamp(span.ended_at)
+                latency = (ended - started) if (started and ended) else 0
+
+            # Get user-provided distinct ID from trace metadata (resolved at trace start).
+            # None means no user-provided ID — use trace_id as fallback in personless mode,
+            # matching the langchain/openai pattern: `distinct_id or trace_id`.
+            trace_info = self._trace_metadata.get(trace_id, {})
+            distinct_id = trace_info.get("distinct_id") or self._get_distinct_id(None)
+
+            # Get group_id from trace metadata for linking
+            group_id = self._get_group_id(trace_id)
+
+            # Get error info if present
+            error_info = span.error
+            error_properties = {}
+            if error_info:
+                if isinstance(error_info, dict):
+                    error_message = error_info.get("message", str(error_info))
+                    error_type_raw = error_info.get("type", "")
+                else:
+                    error_message = str(error_info)
+                    error_type_raw = ""
+
+                # Categorize error type for cross-provider filtering/alerting
+                error_type = "unknown"
+                if (
+                    "ModelBehaviorError" in error_type_raw
+                    or "ModelBehaviorError" in error_message
+                ):
+                    error_type = "model_behavior_error"
+                elif "UserError" in error_type_raw or "UserError" in error_message:
+                    error_type = "user_error"
+                elif (
+                    "InputGuardrailTripwireTriggered" in error_type_raw
+                    or "InputGuardrailTripwireTriggered" in error_message
+                ):
+                    error_type = "input_guardrail_triggered"
+                elif (
+                    "OutputGuardrailTripwireTriggered" in error_type_raw
+                    or "OutputGuardrailTripwireTriggered" in error_message
+                ):
+                    error_type = "output_guardrail_triggered"
+                elif (
+                    "MaxTurnsExceeded" in error_type_raw
+                    or "MaxTurnsExceeded" in error_message
+                ):
+                    error_type = "max_turns_exceeded"
+
+                error_properties = {
+                    "$ai_is_error": True,
+                    "$ai_error": error_message,
+                    "$ai_error_type": error_type,
+                }
+
+            # Personless mode: no user-provided distinct_id, fallback to trace_id
+            if distinct_id is None:
+                error_properties["$process_person_profile"] = False
+                distinct_id = trace_id
+
+            # Dispatch based on span data type
+            if isinstance(span_data, GenerationSpanData):
+                self._handle_generation_span(
+                    span_data,
+                    trace_id,
+                    span_id,
+                    parent_id,
+                    latency,
+                    distinct_id,
+                    group_id,
+                    error_properties,
+                )
+            elif isinstance(span_data, FunctionSpanData):
+                self._handle_function_span(
+                    span_data,
+                    trace_id,
+                    span_id,
+                    parent_id,
+                    latency,
+                    distinct_id,
+                    group_id,
+                    error_properties,
+                )
+            elif isinstance(span_data, AgentSpanData):
+                self._handle_agent_span(
+                    span_data,
+                    trace_id,
+                    span_id,
+                    parent_id,
+                    latency,
+                    distinct_id,
+                    group_id,
+                    error_properties,
+                )
+            elif isinstance(span_data, HandoffSpanData):
+                self._handle_handoff_span(
+                    span_data,
+                    trace_id,
+                    span_id,
+                    parent_id,
+                    latency,
+                    distinct_id,
+                    group_id,
+                    error_properties,
+                )
+            elif isinstance(span_data, GuardrailSpanData):
+                self._handle_guardrail_span(
+                    span_data,
+                    trace_id,
+                    span_id,
+                    parent_id,
+                    latency,
+                    distinct_id,
+                    group_id,
+                    error_properties,
+                )
+            elif isinstance(span_data, ResponseSpanData):
+                self._handle_response_span(
+                    span_data,
+                    trace_id,
+                    span_id,
+                    parent_id,
+                    latency,
+                    distinct_id,
+                    group_id,
+                    error_properties,
+                )
+            elif isinstance(span_data, CustomSpanData):
+                self._handle_custom_span(
+                    span_data,
+                    trace_id,
+                    span_id,
+                    parent_id,
+                    latency,
+                    distinct_id,
+                    group_id,
+                    error_properties,
+                )
+            elif isinstance(
+                span_data, (TranscriptionSpanData, SpeechSpanData, SpeechGroupSpanData)
+            ):
+                self._handle_audio_span(
+                    span_data,
+                    trace_id,
+                    span_id,
+                    parent_id,
+                    latency,
+                    distinct_id,
+                    group_id,
+                    error_properties,
+                )
+            elif isinstance(span_data, MCPListToolsSpanData):
+                self._handle_mcp_span(
+                    span_data,
+                    trace_id,
+                    span_id,
+                    parent_id,
+                    latency,
+                    distinct_id,
+                    group_id,
+                    error_properties,
+                )
+            else:
+                # Unknown span type - capture as generic span
+                self._handle_generic_span(
+                    span_data,
+                    trace_id,
+                    span_id,
+                    parent_id,
+                    latency,
+                    distinct_id,
+                    group_id,
+                    error_properties,
+                )
+
+        except Exception as e:
+            log.debug(f"Error in on_span_end: {e}")
+
+    def _base_properties(
+        self,
+        trace_id: str,
+        span_id: str,
+        parent_id: Optional[str],
+        latency: float,
+        group_id: Optional[str],
+        error_properties: Dict[str, Any],
+    ) -> Dict[str, Any]:
+        """Build the base properties dict shared by all span handlers."""
+        properties = {
+            "$ai_trace_id": trace_id,
+            "$ai_span_id": span_id,
+            "$ai_parent_id": parent_id,
+            "$ai_provider": "openai",
+            "$ai_framework": "openai-agents",
+            "$ai_latency": latency,
+            **error_properties,
+        }
+        if group_id:
+            properties["$ai_group_id"] = group_id
+        return properties
+
+    def _handle_generation_span(
+        self,
+        span_data: GenerationSpanData,
+        trace_id: str,
+        span_id: str,
+        parent_id: Optional[str],
+        latency: float,
+        distinct_id: str,
+        group_id: Optional[str],
+        error_properties: Dict[str, Any],
+    ) -> None:
+        """Handle LLM generation spans - maps to $ai_generation event."""
+        # Extract token usage
+        usage = span_data.usage or {}
+        input_tokens = usage.get("input_tokens") or usage.get("prompt_tokens") or 0
+        output_tokens = (
+            usage.get("output_tokens") or usage.get("completion_tokens") or 0
+        )
+
+        # Extract model config parameters
+        model_config = span_data.model_config or {}
+        model_params = {}
+        for param in [
+            "temperature",
+            "max_tokens",
+            "top_p",
+            "frequency_penalty",
+            "presence_penalty",
+        ]:
+            if param in model_config:
+                model_params[param] = model_config[param]
+
+        properties = {
+            **self._base_properties(
+                trace_id, span_id, parent_id, latency, group_id, error_properties
+            ),
+            "$ai_model": span_data.model,
+            "$ai_model_parameters": model_params if model_params else None,
+            "$ai_input": self._with_privacy_mode(_ensure_serializable(span_data.input)),
+            "$ai_output_choices": self._with_privacy_mode(
+                _ensure_serializable(span_data.output)
+            ),
+            "$ai_input_tokens": input_tokens,
+            "$ai_output_tokens": output_tokens,
+            "$ai_total_tokens": (input_tokens or 0) + (output_tokens or 0),
+        }
+
+        # Add optional token fields if present
+        if usage.get("reasoning_tokens"):
+            properties["$ai_reasoning_tokens"] = usage["reasoning_tokens"]
+        if usage.get("cache_read_input_tokens"):
+            properties["$ai_cache_read_input_tokens"] = usage["cache_read_input_tokens"]
+        if usage.get("cache_creation_input_tokens"):
+            properties["$ai_cache_creation_input_tokens"] = usage[
+                "cache_creation_input_tokens"
+            ]
+
+        self._capture_event("$ai_generation", properties, distinct_id)
+
+    def _handle_function_span(
+        self,
+        span_data: FunctionSpanData,
+        trace_id: str,
+        span_id: str,
+        parent_id: Optional[str],
+        latency: float,
+        distinct_id: str,
+        group_id: Optional[str],
+        error_properties: Dict[str, Any],
+    ) -> None:
+        """Handle function/tool call spans - maps to $ai_span event."""
+        properties = {
+            **self._base_properties(
+                trace_id, span_id, parent_id, latency, group_id, error_properties
+            ),
+            "$ai_span_name": span_data.name,
+            "$ai_span_type": "tool",
+            "$ai_input_state": self._with_privacy_mode(
+                _ensure_serializable(span_data.input)
+            ),
+            "$ai_output_state": self._with_privacy_mode(
+                _ensure_serializable(span_data.output)
+            ),
+        }
+
+        if span_data.mcp_data:
+            properties["$ai_mcp_data"] = _ensure_serializable(span_data.mcp_data)
+
+        self._capture_event("$ai_span", properties, distinct_id)
+
+    def _handle_agent_span(
+        self,
+        span_data: AgentSpanData,
+        trace_id: str,
+        span_id: str,
+        parent_id: Optional[str],
+        latency: float,
+        distinct_id: str,
+        group_id: Optional[str],
+        error_properties: Dict[str, Any],
+    ) -> None:
+        """Handle agent execution spans - maps to $ai_span event."""
+        properties = {
+            **self._base_properties(
+                trace_id, span_id, parent_id, latency, group_id, error_properties
+            ),
+            "$ai_span_name": span_data.name,
+            "$ai_span_type": "agent",
+        }
+
+        if span_data.handoffs:
+            properties["$ai_agent_handoffs"] = span_data.handoffs
+        if span_data.tools:
+            properties["$ai_agent_tools"] = span_data.tools
+        if span_data.output_type:
+            properties["$ai_agent_output_type"] = span_data.output_type
+
+        self._capture_event("$ai_span", properties, distinct_id)
+
+    def _handle_handoff_span(
+        self,
+        span_data: HandoffSpanData,
+        trace_id: str,
+        span_id: str,
+        parent_id: Optional[str],
+        latency: float,
+        distinct_id: str,
+        group_id: Optional[str],
+        error_properties: Dict[str, Any],
+    ) -> None:
+        """Handle agent handoff spans - maps to $ai_span event."""
+        properties = {
+            **self._base_properties(
+                trace_id, span_id, parent_id, latency, group_id, error_properties
+            ),
+            "$ai_span_name": f"{span_data.from_agent} -> {span_data.to_agent}",
+            "$ai_span_type": "handoff",
+            "$ai_handoff_from_agent": span_data.from_agent,
+            "$ai_handoff_to_agent": span_data.to_agent,
+        }
+
+        self._capture_event("$ai_span", properties, distinct_id)
+
+    def _handle_guardrail_span(
+        self,
+        span_data: GuardrailSpanData,
+        trace_id: str,
+        span_id: str,
+        parent_id: Optional[str],
+        latency: float,
+        distinct_id: str,
+        group_id: Optional[str],
+        error_properties: Dict[str, Any],
+    ) -> None:
+        """Handle guardrail execution spans - maps to $ai_span event."""
+        properties = {
+            **self._base_properties(
+                trace_id, span_id, parent_id, latency, group_id, error_properties
+            ),
+            "$ai_span_name": span_data.name,
+            "$ai_span_type": "guardrail",
+            "$ai_guardrail_triggered": span_data.triggered,
+        }
+
+        self._capture_event("$ai_span", properties, distinct_id)
+
+    def _handle_response_span(
+        self,
+        span_data: ResponseSpanData,
+        trace_id: str,
+        span_id: str,
+        parent_id: Optional[str],
+        latency: float,
+        distinct_id: str,
+        group_id: Optional[str],
+        error_properties: Dict[str, Any],
+    ) -> None:
+        """Handle OpenAI Response API spans - maps to $ai_generation event."""
+        response = span_data.response
+        response_id = response.id if response else None
+
+        # Try to extract usage from response
+        usage = getattr(response, "usage", None) if response else None
+        input_tokens = 0
+        output_tokens = 0
+        if usage:
+            input_tokens = getattr(usage, "input_tokens", 0) or 0
+            output_tokens = getattr(usage, "output_tokens", 0) or 0
+
+        # Try to extract model from response
+        model = getattr(response, "model", None) if response else None
+
+        properties = {
+            **self._base_properties(
+                trace_id, span_id, parent_id, latency, group_id, error_properties
+            ),
+            "$ai_model": model,
+            "$ai_response_id": response_id,
+            "$ai_input": self._with_privacy_mode(_ensure_serializable(span_data.input)),
+            "$ai_input_tokens": input_tokens,
+            "$ai_output_tokens": output_tokens,
+            "$ai_total_tokens": input_tokens + output_tokens,
+        }
+
+        # Extract output content from response
+        if response:
+            output_items = getattr(response, "output", None)
+            if output_items:
+                properties["$ai_output_choices"] = self._with_privacy_mode(
+                    _ensure_serializable(output_items)
+                )
+
+        self._capture_event("$ai_generation", properties, distinct_id)
+
+    def _handle_custom_span(
+        self,
+        span_data: CustomSpanData,
+        trace_id: str,
+        span_id: str,
+        parent_id: Optional[str],
+        latency: float,
+        distinct_id: str,
+        group_id: Optional[str],
+        error_properties: Dict[str, Any],
+    ) -> None:
+        """Handle custom user-defined spans - maps to $ai_span event."""
+        properties = {
+            **self._base_properties(
+                trace_id, span_id, parent_id, latency, group_id, error_properties
+            ),
+            "$ai_span_name": span_data.name,
+            "$ai_span_type": "custom",
+            "$ai_custom_data": self._with_privacy_mode(
+                _ensure_serializable(span_data.data)
+            ),
+        }
+
+        self._capture_event("$ai_span", properties, distinct_id)
+
+    def _handle_audio_span(
+        self,
+        span_data: Union[TranscriptionSpanData, SpeechSpanData, SpeechGroupSpanData],
+        trace_id: str,
+        span_id: str,
+        parent_id: Optional[str],
+        latency: float,
+        distinct_id: str,
+        group_id: Optional[str],
+        error_properties: Dict[str, Any],
+    ) -> None:
+        """Handle audio-related spans (transcription, speech) - maps to $ai_span event."""
+        span_type = span_data.type  # "transcription", "speech", or "speech_group"
+
+        properties = {
+            **self._base_properties(
+                trace_id, span_id, parent_id, latency, group_id, error_properties
+            ),
+            "$ai_span_name": span_type,
+            "$ai_span_type": span_type,
+        }
+
+        # Add model info if available
+        if hasattr(span_data, "model") and span_data.model:
+            properties["$ai_model"] = span_data.model
+
+        # Add model config if available (pass-through property)
+        if hasattr(span_data, "model_config") and span_data.model_config:
+            properties["model_config"] = _ensure_serializable(span_data.model_config)
+
+        # Add time to first audio byte for speech spans (pass-through property)
+        if hasattr(span_data, "first_content_at") and span_data.first_content_at:
+            properties["first_content_at"] = span_data.first_content_at
+
+        # Add audio format info (pass-through properties)
+        if hasattr(span_data, "input_format"):
+            properties["audio_input_format"] = span_data.input_format
+        if hasattr(span_data, "output_format"):
+            properties["audio_output_format"] = span_data.output_format
+
+        # Add text input for TTS
+        if (
+            hasattr(span_data, "input")
+            and span_data.input
+            and isinstance(span_data.input, str)
+        ):
+            properties["$ai_input"] = self._with_privacy_mode(span_data.input)
+
+        # Don't include audio data (base64) - just metadata
+        if hasattr(span_data, "output") and isinstance(span_data.output, str):
+            # For transcription, output is the text
+            properties["$ai_output_state"] = self._with_privacy_mode(span_data.output)
+
+        self._capture_event("$ai_span", properties, distinct_id)
+
+    def _handle_mcp_span(
+        self,
+        span_data: MCPListToolsSpanData,
+        trace_id: str,
+        span_id: str,
+        parent_id: Optional[str],
+        latency: float,
+        distinct_id: str,
+        group_id: Optional[str],
+        error_properties: Dict[str, Any],
+    ) -> None:
+        """Handle MCP (Model Context Protocol) spans - maps to $ai_span event."""
+        properties = {
+            **self._base_properties(
+                trace_id, span_id, parent_id, latency, group_id, error_properties
+            ),
+            "$ai_span_name": f"mcp:{span_data.server}",
+            "$ai_span_type": "mcp_tools",
+            "$ai_mcp_server": span_data.server,
+            "$ai_mcp_tools": span_data.result,
+        }
+
+        self._capture_event("$ai_span", properties, distinct_id)
+
+    def _handle_generic_span(
+        self,
+        span_data: Any,
+        trace_id: str,
+        span_id: str,
+        parent_id: Optional[str],
+        latency: float,
+        distinct_id: str,
+        group_id: Optional[str],
+        error_properties: Dict[str, Any],
+    ) -> None:
+        """Handle unknown span types - maps to $ai_span event."""
+        span_type = getattr(span_data, "type", "unknown")
+
+        properties = {
+            **self._base_properties(
+                trace_id, span_id, parent_id, latency, group_id, error_properties
+            ),
+            "$ai_span_name": span_type,
+            "$ai_span_type": span_type,
+        }
+
+        # Try to export span data
+        if hasattr(span_data, "export"):
+            try:
+                exported = span_data.export()
+                properties["$ai_span_data"] = _ensure_serializable(exported)
+            except Exception:
+                pass
+
+        self._capture_event("$ai_span", properties, distinct_id)
+
+    def shutdown(self) -> None:
+        """Clean up resources when the application stops."""
+        try:
+            self._span_start_times.clear()
+            self._trace_metadata.clear()
+
+            # Flush the PostHog client if possible
+            if hasattr(self._client, "flush") and callable(self._client.flush):
+                self._client.flush()
+        except Exception as e:
+            log.debug(f"Error in shutdown: {e}")
+
+    def force_flush(self) -> None:
+        """Force immediate processing of any queued events."""
+        try:
+            if hasattr(self._client, "flush") and callable(self._client.flush):
+                self._client.flush()
+        except Exception as e:
+            log.debug(f"Error in force_flush: {e}")

posthog/ai/prompts.py

@@ -0,0 +1,286 @@
+"""
+Prompt management for PostHog AI SDK.
+
+Fetch and compile LLM prompts from PostHog with caching and fallback support.
+"""
+
+import logging
+import re
+import time
+import urllib.parse
+from typing import Any, Dict, Optional, Union
+
+from posthog.request import USER_AGENT, _get_session
+from posthog.utils import remove_trailing_slash
+
+log = logging.getLogger("posthog")
+
+APP_ENDPOINT = "https://us.posthog.com"
+DEFAULT_CACHE_TTL_SECONDS = 300  # 5 minutes
+
+PromptVariables = Dict[str, Union[str, int, float, bool]]
+
+
+class CachedPrompt:
+    """Cached prompt with metadata."""
+
+    def __init__(self, prompt: str, fetched_at: float):
+        self.prompt = prompt
+        self.fetched_at = fetched_at
+
+
+def _is_prompt_api_response(data: Any) -> bool:
+    """Check if the response is a valid prompt API response."""
+    return (
+        isinstance(data, dict)
+        and "prompt" in data
+        and isinstance(data.get("prompt"), str)
+    )
+
+
+class Prompts:
+    """
+    Fetch and compile LLM prompts from PostHog.
+
+    Can be initialized with a PostHog client or with direct options.
+
+    Examples:
+        ```python
+        from posthog import Posthog
+        from posthog.ai.prompts import Prompts
+
+        # With PostHog client
+        posthog = Posthog('phc_xxx', host='https://us.posthog.com', personal_api_key='phx_xxx')
+        prompts = Prompts(posthog)
+
+        # Or with direct options (no PostHog client needed)
+        prompts = Prompts(
+            personal_api_key='phx_xxx',
+            project_api_key='phc_xxx',
+            host='https://us.posthog.com',
+        )
+
+        # Fetch with caching and fallback
+        template = prompts.get('support-system-prompt', fallback='You are a helpful assistant.')
+
+        # Compile with variables
+        system_prompt = prompts.compile(template, {
+            'company': 'Acme Corp',
+            'tier': 'premium',
+        })
+        ```
+    """
+
+    def __init__(
+        self,
+        posthog: Optional[Any] = None,
+        *,
+        personal_api_key: Optional[str] = None,
+        project_api_key: Optional[str] = None,
+        host: Optional[str] = None,
+        default_cache_ttl_seconds: Optional[int] = None,
+    ):
+        """
+        Initialize Prompts.
+
+        Args:
+            posthog: PostHog client instance (optional if personal_api_key provided)
+            personal_api_key: Direct personal API key (optional if posthog provided)
+            project_api_key: Direct project API key (optional if posthog provided)
+            host: PostHog host (defaults to app endpoint)
+            default_cache_ttl_seconds: Default cache TTL (defaults to 300)
+        """
+        self._default_cache_ttl_seconds = (
+            default_cache_ttl_seconds or DEFAULT_CACHE_TTL_SECONDS
+        )
+        self._cache: Dict[str, CachedPrompt] = {}
+
+        if posthog is not None:
+            self._personal_api_key = getattr(posthog, "personal_api_key", None) or ""
+            self._project_api_key = getattr(posthog, "api_key", None) or ""
+            self._host = remove_trailing_slash(
+                getattr(posthog, "raw_host", None) or APP_ENDPOINT
+            )
+        else:
+            self._personal_api_key = personal_api_key or ""
+            self._project_api_key = project_api_key or ""
+            self._host = remove_trailing_slash(host or APP_ENDPOINT)
+
+    def get(
+        self,
+        name: str,
+        *,
+        cache_ttl_seconds: Optional[int] = None,
+        fallback: Optional[str] = None,
+    ) -> str:
+        """
+        Fetch a prompt by name from the PostHog API.
+
+        Caching behavior:
+        1. If cache is fresh, return cached value
+        2. If fetch fails and cache exists (stale), return stale cache with warning
+        3. If fetch fails and fallback provided, return fallback with warning
+        4. If fetch fails with no cache/fallback, raise exception
+
+        Args:
+            name: The name of the prompt to fetch
+            cache_ttl_seconds: Cache TTL in seconds (defaults to instance default)
+            fallback: Fallback prompt to use if fetch fails and no cache available
+
+        Returns:
+            The prompt string
+
+        Raises:
+            Exception: If the prompt cannot be fetched and no fallback is available
+        """
+        ttl = (
+            cache_ttl_seconds
+            if cache_ttl_seconds is not None
+            else self._default_cache_ttl_seconds
+        )
+
+        # Check cache first
+        cached = self._cache.get(name)
+        now = time.time()
+
+        if cached is not None:
+            is_fresh = (now - cached.fetched_at) < ttl
+
+            if is_fresh:
+                return cached.prompt
+
+        # Try to fetch from API
+        try:
+            prompt = self._fetch_prompt_from_api(name)
+            fetched_at = time.time()
+
+            # Update cache
+            self._cache[name] = CachedPrompt(prompt=prompt, fetched_at=fetched_at)
+
+            return prompt
+
+        except Exception as error:
+            # Fallback order:
+            # 1. Return stale cache (with warning)
+            if cached is not None:
+                log.warning(
+                    '[PostHog Prompts] Failed to fetch prompt "%s", using stale cache: %s',
+                    name,
+                    error,
+                )
+                return cached.prompt
+
+            # 2. Return fallback (with warning)
+            if fallback is not None:
+                log.warning(
+                    '[PostHog Prompts] Failed to fetch prompt "%s", using fallback: %s',
+                    name,
+                    error,
+                )
+                return fallback
+
+            # 3. Raise error
+            raise
+
+    def compile(self, prompt: str, variables: PromptVariables) -> str:
+        """
+        Replace {{variableName}} placeholders with values.
+
+        Unmatched variables are left unchanged.
+        Supports variable names with hyphens and dots (e.g., user-id, company.name).
+
+        Args:
+            prompt: The prompt template string
+            variables: Object containing variable values
+
+        Returns:
+            The compiled prompt string
+        """
+
+        def replace_variable(match: re.Match) -> str:
+            variable_name = match.group(1)
+
+            if variable_name in variables:
+                return str(variables[variable_name])
+
+            return match.group(0)
+
+        return re.sub(r"\{\{([\w.-]+)\}\}", replace_variable, prompt)
+
+    def clear_cache(self, name: Optional[str] = None) -> None:
+        """
+        Clear cached prompts.
+
+        Args:
+            name: Specific prompt to clear. If None, clears all cached prompts.
+        """
+        if name is not None:
+            self._cache.pop(name, None)
+        else:
+            self._cache.clear()
+
+    def _fetch_prompt_from_api(self, name: str) -> str:
+        """
+        Fetch prompt from PostHog API.
+
+        Endpoint: {host}/api/environments/@current/llm_prompts/name/{encoded_name}/?token={encoded_project_api_key}
+        Auth: Bearer {personal_api_key}
+
+        Args:
+            name: The name of the prompt to fetch
+
+        Returns:
+            The prompt string
+
+        Raises:
+            Exception: If the prompt cannot be fetched
+        """
+        if not self._personal_api_key:
+            raise Exception(
+                "[PostHog Prompts] personal_api_key is required to fetch prompts. "
+                "Please provide it when initializing the Prompts instance."
+            )
+        if not self._project_api_key:
+            raise Exception(
+                "[PostHog Prompts] project_api_key is required to fetch prompts. "
+                "Please provide it when initializing the Prompts instance."
+            )
+
+        encoded_name = urllib.parse.quote(name, safe="")
+        encoded_project_api_key = urllib.parse.quote(self._project_api_key, safe="")
+        url = f"{self._host}/api/environments/@current/llm_prompts/name/{encoded_name}/?token={encoded_project_api_key}"
+
+        headers = {
+            "Authorization": f"Bearer {self._personal_api_key}",
+            "User-Agent": USER_AGENT,
+        }
+
+        response = _get_session().get(url, headers=headers, timeout=10)
+
+        if not response.ok:
+            if response.status_code == 404:
+                raise Exception(f'[PostHog Prompts] Prompt "{name}" not found')
+
+            if response.status_code == 403:
+                raise Exception(
+                    f'[PostHog Prompts] Access denied for prompt "{name}". '
+                    "Check that your personal_api_key has the correct permissions and the LLM prompts feature is enabled."
+                )
+
+            raise Exception(
+                f'[PostHog Prompts] Failed to fetch prompt "{name}": HTTP {response.status_code}'
+            )
+
+        try:
+            data = response.json()
+        except Exception:
+            raise Exception(
+                f'[PostHog Prompts] Invalid response format for prompt "{name}"'
+            )
+
+        if not _is_prompt_api_response(data):
+            raise Exception(
+                f'[PostHog Prompts] Invalid response format for prompt "{name}"'
+            )
+
+        return data["prompt"]

posthog/ai/sanitization.py

@@ -1,3 +1,4 @@
+import os
 import re
 from typing import Any
 from urllib.parse import urlparse
@@ -5,6 +6,15 @@ from urllib.parse import urlparse
 REDACTED_IMAGE_PLACEHOLDER = "[base64 image redacted]"
 
 
+def _is_multimodal_enabled() -> bool:
+    """Check if multimodal capture is enabled via environment variable."""
+    return os.environ.get("_INTERNAL_LLMA_MULTIMODAL", "").lower() in (
+        "true",
+        "1",
+        "yes",
+    )
+
+
 def is_base64_data_url(text: str) -> bool:
     return re.match(r"^data:([^;]+);base64,", text) is not None
 
@@ -27,6 +37,9 @@ def is_raw_base64(text: str) -> bool:
 
 
 def redact_base64_data_url(value: Any) -> Any:
+    if _is_multimodal_enabled():
+        return value
+
     if not isinstance(value, str):
         return value
 
@@ -70,6 +83,12 @@ def sanitize_openai_image(item: Any) -> Any:
     if not isinstance(item, dict):
         return item
 
+    if item.get("type") == "input_image" and isinstance(item.get("image_url"), str):
+        return {
+            **item,
+            "image_url": redact_base64_data_url(item["image_url"]),
+        }
+
     if (
         item.get("type") == "image_url"
         and isinstance(item.get("image_url"), dict)
@@ -83,6 +102,11 @@ def sanitize_openai_image(item: Any) -> Any:
             },
         }
 
+    if item.get("type") == "audio" and "data" in item:
+        if _is_multimodal_enabled():
+            return item
+        return {**item, "data": REDACTED_IMAGE_PLACEHOLDER}
+
     return item
 
 
@@ -100,6 +124,9 @@ def sanitize_openai_response_image(item: Any) -> Any:
 
 
 def sanitize_anthropic_image(item: Any) -> Any:
+    if _is_multimodal_enabled():
+        return item
+
     if not isinstance(item, dict):
         return item
 
@@ -109,8 +136,6 @@ def sanitize_anthropic_image(item: Any) -> Any:
         and item["source"].get("type") == "base64"
         and "data" in item["source"]
     ):
-        # For Anthropic, if the source type is "base64", we should always redact the data
-        # The provider is explicitly telling us this is base64 data
         return {
             **item,
             "source": {
@@ -123,6 +148,9 @@ def sanitize_anthropic_image(item: Any) -> Any:
 
 
 def sanitize_gemini_part(part: Any) -> Any:
+    if _is_multimodal_enabled():
+        return part
+
     if not isinstance(part, dict):
         return part
 
@@ -131,8 +159,6 @@ def sanitize_gemini_part(part: Any) -> Any:
         and isinstance(part["inline_data"], dict)
         and "data" in part["inline_data"]
     ):
-        # For Gemini, the inline_data structure indicates base64 data
-        # We should redact any string data in this context
         return {
             **part,
             "inline_data": {
@@ -185,7 +211,9 @@ def sanitize_langchain_image(item: Any) -> Any:
         and isinstance(item.get("source"), dict)
         and "data" in item["source"]
     ):
-        # Anthropic style - raw base64 in structured format, always redact
+        if _is_multimodal_enabled():
+            return item
+
         return {
             **item,
             "source": {

posthog/ai/types.py

@@ -64,6 +64,7 @@ class TokenUsage(TypedDict, total=False):
     cache_creation_input_tokens: Optional[int]
     reasoning_tokens: Optional[int]
     web_search_count: Optional[int]
+    raw_usage: Optional[Any]  # Raw provider usage metadata for backend processing
 
 
 class ProviderResponse(TypedDict, total=False):

posthog/ai/utils.py

@@ -2,14 +2,63 @@ import time
 import uuid
 from typing import Any, Callable, Dict, List, Optional, cast
 
-from posthog.client import Client as PostHogClient
-from posthog.ai.types import FormattedMessage, StreamingEventData, TokenUsage
+from posthog import get_tags, identify_context, new_context, tag
 from posthog.ai.sanitization import (
-    sanitize_openai,
     sanitize_anthropic,
     sanitize_gemini,
     sanitize_langchain,
+    sanitize_openai,
 )
+from posthog.ai.types import FormattedMessage, StreamingEventData, TokenUsage
+from posthog.client import Client as PostHogClient
+
+
+def serialize_raw_usage(raw_usage: Any) -> Optional[Dict[str, Any]]:
+    """
+    Convert raw provider usage objects to JSON-serializable dicts.
+
+    Handles Pydantic models (OpenAI/Anthropic) and protobuf-like objects (Gemini)
+    with a fallback chain to ensure we never pass unserializable objects to PostHog.
+
+    Args:
+        raw_usage: Raw usage object from provider SDK
+
+    Returns:
+        Plain dict or None if conversion fails
+    """
+    if raw_usage is None:
+        return None
+
+    # Already a dict
+    if isinstance(raw_usage, dict):
+        return raw_usage
+
+    # Try Pydantic model_dump() (OpenAI/Anthropic)
+    if hasattr(raw_usage, "model_dump") and callable(raw_usage.model_dump):
+        try:
+            return raw_usage.model_dump()
+        except Exception:
+            pass
+
+    # Try to_dict() (some protobuf objects)
+    if hasattr(raw_usage, "to_dict") and callable(raw_usage.to_dict):
+        try:
+            return raw_usage.to_dict()
+        except Exception:
+            pass
+
+    # Try __dict__ / vars() for simple objects
+    try:
+        return vars(raw_usage)
+    except Exception:
+        pass
+
+    # Last resort: convert to string representation
+    # This ensures we always return something rather than failing
+    try:
+        return {"_raw": str(raw_usage)}
+    except Exception:
+        return None
 
 
 def merge_usage_stats(
@@ -59,6 +108,17 @@ def merge_usage_stats(
             current = target.get("web_search_count") or 0
             target["web_search_count"] = max(current, source_web_search)
 
+        # Merge raw_usage to avoid losing data from earlier events
+        # For Anthropic streaming: message_start has input tokens, message_delta has output
+        # Note: raw_usage is already serialized by converters, so it's a dict
+        source_raw_usage = source.get("raw_usage")
+        if source_raw_usage is not None and isinstance(source_raw_usage, dict):
+            current_raw_value = target.get("raw_usage")
+            current_raw: Dict[str, Any] = (
+                current_raw_value if isinstance(current_raw_value, dict) else {}
+            )
+            target["raw_usage"] = {**current_raw, **source_raw_usage}
+
     elif mode == "cumulative":
         # Replace with latest values (already cumulative)
         if source.get("input_tokens") is not None:
@@ -75,6 +135,9 @@ def merge_usage_stats(
             target["reasoning_tokens"] = source["reasoning_tokens"]
         if source.get("web_search_count") is not None:
             target["web_search_count"] = source["web_search_count"]
+        # Note: raw_usage is already serialized by converters, so it's a dict
+        if source.get("raw_usage") is not None:
+            target["raw_usage"] = source["raw_usage"]
 
     else:
         raise ValueError(f"Invalid mode: {mode}. Must be 'incremental' or 'cumulative'")
@@ -256,94 +319,113 @@ def call_llm_and_track_usage(
     usage: TokenUsage = TokenUsage()
     error_params: Dict[str, Any] = {}
 
-    try:
-        response = call_method(**kwargs)
-    except Exception as exc:
-        error = exc
-        http_status = getattr(
-            exc, "status_code", 0
-        )  # default to 0 becuase its likely an SDK error
-        error_params = {
-            "$ai_is_error": True,
-            "$ai_error": exc.__str__(),
-        }
-    finally:
-        end_time = time.time()
-        latency = end_time - start_time
+    with new_context(client=ph_client, capture_exceptions=False):
+        if posthog_distinct_id:
+            identify_context(posthog_distinct_id)
 
-        if posthog_trace_id is None:
-            posthog_trace_id = str(uuid.uuid4())
+        try:
+            response = call_method(**kwargs)
+        except Exception as exc:
+            error = exc
+            http_status = getattr(
+                exc, "status_code", 0
+            )  # default to 0 becuase its likely an SDK error
+            error_params = {
+                "$ai_is_error": True,
+                "$ai_error": exc.__str__(),
+            }
+            # TODO: Add exception capture for OpenAI/Anthropic/Gemini wrappers when
+            # enable_exception_autocapture is True, similar to LangChain callbacks.
+            # See _capture_exception_and_update_properties in langchain/callbacks.py
+        finally:
+            end_time = time.time()
+            latency = end_time - start_time
 
-        if response and (
-            hasattr(response, "usage")
-            or (provider == "gemini" and hasattr(response, "usage_metadata"))
-        ):
-            usage = get_usage(response, provider)
+            if posthog_trace_id is None:
+                posthog_trace_id = str(uuid.uuid4())
 
-        messages = merge_system_prompt(kwargs, provider)
-        sanitized_messages = sanitize_messages(messages, provider)
+            if response and (
+                hasattr(response, "usage")
+                or (provider == "gemini" and hasattr(response, "usage_metadata"))
+            ):
+                usage = get_usage(response, provider)
 
-        event_properties = {
-            "$ai_provider": provider,
-            "$ai_model": kwargs.get("model"),
-            "$ai_model_parameters": get_model_params(kwargs),
-            "$ai_input": with_privacy_mode(
-                ph_client, posthog_privacy_mode, sanitized_messages
-            ),
-            "$ai_output_choices": with_privacy_mode(
-                ph_client, posthog_privacy_mode, format_response(response, provider)
-            ),
-            "$ai_http_status": http_status,
-            "$ai_input_tokens": usage.get("input_tokens", 0),
-            "$ai_output_tokens": usage.get("output_tokens", 0),
-            "$ai_latency": latency,
-            "$ai_trace_id": posthog_trace_id,
-            "$ai_base_url": str(base_url),
-            **(posthog_properties or {}),
-            **(error_params or {}),
-        }
+            messages = merge_system_prompt(kwargs, provider)
+            sanitized_messages = sanitize_messages(messages, provider)
 
-        available_tool_calls = extract_available_tool_calls(provider, kwargs)
-
-        if available_tool_calls:
-            event_properties["$ai_tools"] = available_tool_calls
-
-        cache_read = usage.get("cache_read_input_tokens")
-        if cache_read is not None and cache_read > 0:
-            event_properties["$ai_cache_read_input_tokens"] = cache_read
-
-        cache_creation = usage.get("cache_creation_input_tokens")
-        if cache_creation is not None and cache_creation > 0:
-            event_properties["$ai_cache_creation_input_tokens"] = cache_creation
-
-        reasoning = usage.get("reasoning_tokens")
-        if reasoning is not None and reasoning > 0:
-            event_properties["$ai_reasoning_tokens"] = reasoning
-
-        web_search_count = usage.get("web_search_count")
-        if web_search_count is not None and web_search_count > 0:
-            event_properties["$ai_web_search_count"] = web_search_count
-
-        if posthog_distinct_id is None:
-            event_properties["$process_person_profile"] = False
-
-        # Process instructions for Responses API
-        if provider == "openai" and kwargs.get("instructions") is not None:
-            event_properties["$ai_instructions"] = with_privacy_mode(
-                ph_client, posthog_privacy_mode, kwargs.get("instructions")
+            tag("$ai_provider", provider)
+            tag("$ai_model", kwargs.get("model") or getattr(response, "model", None))
+            tag("$ai_model_parameters", get_model_params(kwargs))
+            tag(
+                "$ai_input",
+                with_privacy_mode(ph_client, posthog_privacy_mode, sanitized_messages),
             )
-
-        # send the event to posthog
-        if hasattr(ph_client, "capture") and callable(ph_client.capture):
-            ph_client.capture(
-                distinct_id=posthog_distinct_id or posthog_trace_id,
-                event="$ai_generation",
-                properties=event_properties,
-                groups=posthog_groups,
+            tag(
+                "$ai_output_choices",
+                with_privacy_mode(
+                    ph_client, posthog_privacy_mode, format_response(response, provider)
+                ),
             )
+            tag("$ai_http_status", http_status)
+            tag("$ai_input_tokens", usage.get("input_tokens", 0))
+            tag("$ai_output_tokens", usage.get("output_tokens", 0))
+            tag("$ai_latency", latency)
+            tag("$ai_trace_id", posthog_trace_id)
+            tag("$ai_base_url", str(base_url))
 
-    if error:
-        raise error
+            available_tool_calls = extract_available_tool_calls(provider, kwargs)
+
+            if available_tool_calls:
+                tag("$ai_tools", available_tool_calls)
+
+            cache_read = usage.get("cache_read_input_tokens")
+            if cache_read is not None and cache_read > 0:
+                tag("$ai_cache_read_input_tokens", cache_read)
+
+            cache_creation = usage.get("cache_creation_input_tokens")
+            if cache_creation is not None and cache_creation > 0:
+                tag("$ai_cache_creation_input_tokens", cache_creation)
+
+            reasoning = usage.get("reasoning_tokens")
+            if reasoning is not None and reasoning > 0:
+                tag("$ai_reasoning_tokens", reasoning)
+
+            web_search_count = usage.get("web_search_count")
+            if web_search_count is not None and web_search_count > 0:
+                tag("$ai_web_search_count", web_search_count)
+
+            raw_usage = usage.get("raw_usage")
+            if raw_usage is not None:
+                # Already serialized by converters
+                tag("$ai_usage", raw_usage)
+
+            if posthog_distinct_id is None:
+                tag("$process_person_profile", False)
+
+            # Process instructions for Responses API
+            if provider == "openai" and kwargs.get("instructions") is not None:
+                tag(
+                    "$ai_instructions",
+                    with_privacy_mode(
+                        ph_client, posthog_privacy_mode, kwargs.get("instructions")
+                    ),
+                )
+
+            # send the event to posthog
+            if hasattr(ph_client, "capture") and callable(ph_client.capture):
+                ph_client.capture(
+                    distinct_id=posthog_distinct_id or posthog_trace_id,
+                    event="$ai_generation",
+                    properties={
+                        **get_tags(),
+                        **(posthog_properties or {}),
+                        **(error_params or {}),
+                    },
+                    groups=posthog_groups,
+                )
+
+        if error:
+            raise error
 
     return response
 
@@ -367,94 +449,113 @@ async def call_llm_and_track_usage_async(
     usage: TokenUsage = TokenUsage()
     error_params: Dict[str, Any] = {}
 
-    try:
-        response = await call_async_method(**kwargs)
-    except Exception as exc:
-        error = exc
-        http_status = getattr(
-            exc, "status_code", 0
-        )  # default to 0 because its likely an SDK error
-        error_params = {
-            "$ai_is_error": True,
-            "$ai_error": exc.__str__(),
-        }
-    finally:
-        end_time = time.time()
-        latency = end_time - start_time
+    with new_context(client=ph_client, capture_exceptions=False):
+        if posthog_distinct_id:
+            identify_context(posthog_distinct_id)
 
-        if posthog_trace_id is None:
-            posthog_trace_id = str(uuid.uuid4())
+        try:
+            response = await call_async_method(**kwargs)
+        except Exception as exc:
+            error = exc
+            http_status = getattr(
+                exc, "status_code", 0
+            )  # default to 0 because its likely an SDK error
+            error_params = {
+                "$ai_is_error": True,
+                "$ai_error": exc.__str__(),
+            }
+            # TODO: Add exception capture for OpenAI/Anthropic/Gemini wrappers when
+            # enable_exception_autocapture is True, similar to LangChain callbacks.
+            # See _capture_exception_and_update_properties in langchain/callbacks.py
+        finally:
+            end_time = time.time()
+            latency = end_time - start_time
 
-        if response and (
-            hasattr(response, "usage")
-            or (provider == "gemini" and hasattr(response, "usage_metadata"))
-        ):
-            usage = get_usage(response, provider)
+            if posthog_trace_id is None:
+                posthog_trace_id = str(uuid.uuid4())
 
-        messages = merge_system_prompt(kwargs, provider)
-        sanitized_messages = sanitize_messages(messages, provider)
+            if response and (
+                hasattr(response, "usage")
+                or (provider == "gemini" and hasattr(response, "usage_metadata"))
+            ):
+                usage = get_usage(response, provider)
 
-        event_properties = {
-            "$ai_provider": provider,
-            "$ai_model": kwargs.get("model"),
-            "$ai_model_parameters": get_model_params(kwargs),
-            "$ai_input": with_privacy_mode(
-                ph_client, posthog_privacy_mode, sanitized_messages
-            ),
-            "$ai_output_choices": with_privacy_mode(
-                ph_client, posthog_privacy_mode, format_response(response, provider)
-            ),
-            "$ai_http_status": http_status,
-            "$ai_input_tokens": usage.get("input_tokens", 0),
-            "$ai_output_tokens": usage.get("output_tokens", 0),
-            "$ai_latency": latency,
-            "$ai_trace_id": posthog_trace_id,
-            "$ai_base_url": str(base_url),
-            **(posthog_properties or {}),
-            **(error_params or {}),
-        }
+            messages = merge_system_prompt(kwargs, provider)
+            sanitized_messages = sanitize_messages(messages, provider)
 
-        available_tool_calls = extract_available_tool_calls(provider, kwargs)
-
-        if available_tool_calls:
-            event_properties["$ai_tools"] = available_tool_calls
-
-        cache_read = usage.get("cache_read_input_tokens")
-        if cache_read is not None and cache_read > 0:
-            event_properties["$ai_cache_read_input_tokens"] = cache_read
-
-        cache_creation = usage.get("cache_creation_input_tokens")
-        if cache_creation is not None and cache_creation > 0:
-            event_properties["$ai_cache_creation_input_tokens"] = cache_creation
-
-        reasoning = usage.get("reasoning_tokens")
-        if reasoning is not None and reasoning > 0:
-            event_properties["$ai_reasoning_tokens"] = reasoning
-
-        web_search_count = usage.get("web_search_count")
-        if web_search_count is not None and web_search_count > 0:
-            event_properties["$ai_web_search_count"] = web_search_count
-
-        if posthog_distinct_id is None:
-            event_properties["$process_person_profile"] = False
-
-        # Process instructions for Responses API
-        if provider == "openai" and kwargs.get("instructions") is not None:
-            event_properties["$ai_instructions"] = with_privacy_mode(
-                ph_client, posthog_privacy_mode, kwargs.get("instructions")
+            tag("$ai_provider", provider)
+            tag("$ai_model", kwargs.get("model") or getattr(response, "model", None))
+            tag("$ai_model_parameters", get_model_params(kwargs))
+            tag(
+                "$ai_input",
+                with_privacy_mode(ph_client, posthog_privacy_mode, sanitized_messages),
             )
-
-        # send the event to posthog
-        if hasattr(ph_client, "capture") and callable(ph_client.capture):
-            ph_client.capture(
-                distinct_id=posthog_distinct_id or posthog_trace_id,
-                event="$ai_generation",
-                properties=event_properties,
-                groups=posthog_groups,
+            tag(
+                "$ai_output_choices",
+                with_privacy_mode(
+                    ph_client, posthog_privacy_mode, format_response(response, provider)
+                ),
             )
+            tag("$ai_http_status", http_status)
+            tag("$ai_input_tokens", usage.get("input_tokens", 0))
+            tag("$ai_output_tokens", usage.get("output_tokens", 0))
+            tag("$ai_latency", latency)
+            tag("$ai_trace_id", posthog_trace_id)
+            tag("$ai_base_url", str(base_url))
 
-    if error:
-        raise error
+            available_tool_calls = extract_available_tool_calls(provider, kwargs)
+
+            if available_tool_calls:
+                tag("$ai_tools", available_tool_calls)
+
+            cache_read = usage.get("cache_read_input_tokens")
+            if cache_read is not None and cache_read > 0:
+                tag("$ai_cache_read_input_tokens", cache_read)
+
+            cache_creation = usage.get("cache_creation_input_tokens")
+            if cache_creation is not None and cache_creation > 0:
+                tag("$ai_cache_creation_input_tokens", cache_creation)
+
+            reasoning = usage.get("reasoning_tokens")
+            if reasoning is not None and reasoning > 0:
+                tag("$ai_reasoning_tokens", reasoning)
+
+            web_search_count = usage.get("web_search_count")
+            if web_search_count is not None and web_search_count > 0:
+                tag("$ai_web_search_count", web_search_count)
+
+            raw_usage = usage.get("raw_usage")
+            if raw_usage is not None:
+                # Already serialized by converters
+                tag("$ai_usage", raw_usage)
+
+            if posthog_distinct_id is None:
+                tag("$process_person_profile", False)
+
+            # Process instructions for Responses API
+            if provider == "openai" and kwargs.get("instructions") is not None:
+                tag(
+                    "$ai_instructions",
+                    with_privacy_mode(
+                        ph_client, posthog_privacy_mode, kwargs.get("instructions")
+                    ),
+                )
+
+            # send the event to posthog
+            if hasattr(ph_client, "capture") and callable(ph_client.capture):
+                ph_client.capture(
+                    distinct_id=posthog_distinct_id or posthog_trace_id,
+                    event="$ai_generation",
+                    properties={
+                        **get_tags(),
+                        **(posthog_properties or {}),
+                        **(error_params or {}),
+                    },
+                    groups=posthog_groups,
+                )
+
+        if error:
+            raise error
 
     return response
 
@@ -565,6 +666,12 @@ def capture_streaming_event(
     ):
         event_properties["$ai_web_search_count"] = web_search_count
 
+    # Add raw usage metadata if present (all providers)
+    raw_usage = event_data["usage_stats"].get("raw_usage")
+    if raw_usage is not None:
+        # Already serialized by converters
+        event_properties["$ai_usage"] = raw_usage
+
     # Handle provider-specific fields
     if (
         event_data["provider"] == "openai"

posthog/client.py

@@ -2,53 +2,64 @@ import atexit
 import logging
 import os
 import sys
+import warnings
 from datetime import datetime, timedelta
 from typing import Any, Dict, Optional, Union
-from typing_extensions import Unpack
 from uuid import uuid4
 
 from dateutil.tz import tzutc
 from six import string_types
+from typing_extensions import Unpack
 
-from posthog.args import OptionalCaptureArgs, OptionalSetArgs, ID_TYPES, ExceptionArg
+from posthog.args import ID_TYPES, ExceptionArg, OptionalCaptureArgs, OptionalSetArgs
 from posthog.consumer import Consumer
+from posthog.contexts import (
+    _get_current_context,
+    get_capture_exception_code_variables_context,
+    get_code_variables_ignore_patterns_context,
+    get_code_variables_mask_patterns_context,
+    get_context_device_id,
+    get_context_distinct_id,
+    get_context_session_id,
+    new_context,
+)
 from posthog.exception_capture import ExceptionCapture
 from posthog.exception_utils import (
+    DEFAULT_CODE_VARIABLES_IGNORE_PATTERNS,
+    DEFAULT_CODE_VARIABLES_MASK_PATTERNS,
     exc_info_from_error,
+    exception_is_already_captured,
     exceptions_from_error_tuple,
     handle_in_app,
-    exception_is_already_captured,
     mark_exception_as_captured,
     try_attach_code_variables_to_frames,
-    DEFAULT_CODE_VARIABLES_MASK_PATTERNS,
-    DEFAULT_CODE_VARIABLES_IGNORE_PATTERNS,
 )
 from posthog.feature_flags import (
     InconclusiveMatchError,
     RequiresServerEvaluation,
     match_feature_flag_properties,
+    resolve_bucketing_value,
+)
+from posthog.flag_definition_cache import (
+    FlagDefinitionCacheData,
+    FlagDefinitionCacheProvider,
 )
 from posthog.poller import Poller
 from posthog.request import (
     DEFAULT_HOST,
     APIError,
+    QuotaLimitError,
+    RequestsConnectionError,
+    RequestsTimeout,
     batch_post,
     determine_server_host,
     flags,
     get,
     remote_config,
 )
-from posthog.contexts import (
-    _get_current_context,
-    get_context_distinct_id,
-    get_context_session_id,
-    get_capture_exception_code_variables_context,
-    get_code_variables_mask_patterns_context,
-    get_code_variables_ignore_patterns_context,
-    new_context,
-)
 from posthog.types import (
     FeatureFlag,
+    FeatureFlagError,
     FeatureFlagResult,
     FlagMetadata,
     FlagsAndPayloads,
@@ -184,9 +195,11 @@ class Client(object):
         before_send=None,
         flag_fallback_cache_url=None,
         enable_local_evaluation=True,
+        flag_definition_cache_provider: Optional[FlagDefinitionCacheProvider] = None,
         capture_exception_code_variables=False,
         code_variables_mask_patterns=None,
         code_variables_ignore_patterns=None,
+        in_app_modules: list[str] | None = None,
     ):
         """
         Initialize a new PostHog client instance.
@@ -222,8 +235,8 @@ class Client(object):
         self.timeout = timeout
         self._feature_flags = None  # private variable to store flags
         self.feature_flags_by_key = None
-        self.group_type_mapping = None
-        self.cohorts = None
+        self.group_type_mapping: Optional[dict[str, str]] = None
+        self.cohorts: Optional[dict[str, Any]] = None
         self.poll_interval = poll_interval
         self.feature_flags_request_timeout_seconds = (
             feature_flags_request_timeout_seconds
@@ -232,6 +245,8 @@ class Client(object):
         self.distinct_ids_feature_flags_reported = SizeLimitedDict(MAX_DICT_SIZE, set)
         self.flag_cache = self._initialize_flag_cache(flag_fallback_cache_url)
         self.flag_definition_version = 0
+        self._flags_etag: Optional[str] = None
+        self._flag_definition_cache_provider = flag_definition_cache_provider
         self.disabled = disabled
         self.disable_geoip = disable_geoip
         self.historical_migration = historical_migration
@@ -253,6 +268,7 @@ class Client(object):
             if code_variables_ignore_patterns is not None
             else DEFAULT_CODE_VARIABLES_IGNORE_PATTERNS
         )
+        self.in_app_modules = in_app_modules
 
         if project_root is None:
             try:
@@ -295,8 +311,9 @@ class Client(object):
             # to call flush().
             if send:
                 atexit.register(self.join)
-            for n in range(thread):
-                self.consumers = []
+
+            self.consumers = []
+            for _ in range(thread):
                 consumer = Consumer(
                     self.queue,
                     self.api_key,
@@ -367,6 +384,7 @@ class Client(object):
         group_properties=None,
         disable_geoip=None,
         flag_keys_to_evaluate: Optional[list[str]] = None,
+        device_id: Optional[str] = None,
     ) -> dict[str, Union[bool, str]]:
         """
         Get feature flag variants for a user by calling decide.
@@ -379,6 +397,7 @@ class Client(object):
             disable_geoip: Whether to disable GeoIP for this request.
             flag_keys_to_evaluate: A list of specific flag keys to evaluate. If provided,
                 only these flags will be evaluated, improving performance.
+            device_id: The device ID for this request.
 
         Category:
             Feature flags
@@ -390,6 +409,7 @@ class Client(object):
             group_properties,
             disable_geoip,
             flag_keys_to_evaluate,
+            device_id=device_id,
         )
         return to_values(resp_data) or {}
 
@@ -401,6 +421,7 @@ class Client(object):
         group_properties=None,
         disable_geoip=None,
         flag_keys_to_evaluate: Optional[list[str]] = None,
+        device_id: Optional[str] = None,
     ) -> dict[str, str]:
         """
         Get feature flag payloads for a user by calling decide.
@@ -413,6 +434,7 @@ class Client(object):
             disable_geoip: Whether to disable GeoIP for this request.
             flag_keys_to_evaluate: A list of specific flag keys to evaluate. If provided,
                 only these flags will be evaluated, improving performance.
+            device_id: The device ID for this request.
 
         Examples:
             ```python
@@ -429,6 +451,7 @@ class Client(object):
             group_properties,
             disable_geoip,
             flag_keys_to_evaluate,
+            device_id=device_id,
         )
         return to_payloads(resp_data) or {}
 
@@ -440,6 +463,7 @@ class Client(object):
         group_properties=None,
         disable_geoip=None,
         flag_keys_to_evaluate: Optional[list[str]] = None,
+        device_id: Optional[str] = None,
     ) -> FlagsAndPayloads:
         """
         Get feature flags and payloads for a user by calling decide.
@@ -452,6 +476,7 @@ class Client(object):
             disable_geoip: Whether to disable GeoIP for this request.
             flag_keys_to_evaluate: A list of specific flag keys to evaluate. If provided,
                 only these flags will be evaluated, improving performance.
+            device_id: The device ID for this request.
 
         Examples:
             ```python
@@ -468,6 +493,7 @@ class Client(object):
             group_properties,
             disable_geoip,
             flag_keys_to_evaluate,
+            device_id=device_id,
         )
         return to_flags_and_payloads(resp)
 
@@ -479,6 +505,7 @@ class Client(object):
         group_properties=None,
         disable_geoip=None,
         flag_keys_to_evaluate: Optional[list[str]] = None,
+        device_id: Optional[str] = None,
     ) -> FlagsResponse:
         """
         Get feature flags decision.
@@ -491,6 +518,7 @@ class Client(object):
             disable_geoip: Whether to disable GeoIP for this request.
             flag_keys_to_evaluate: A list of specific flag keys to evaluate. If provided,
                 only these flags will be evaluated, improving performance.
+            device_id: The device ID for this request.
 
         Examples:
             ```python
@@ -507,6 +535,9 @@ class Client(object):
         if distinct_id is None:
             distinct_id = get_context_distinct_id()
 
+        if device_id is None:
+            device_id = get_context_device_id()
+
         if disable_geoip is None:
             disable_geoip = self.disable_geoip
 
@@ -519,6 +550,7 @@ class Client(object):
             "person_properties": person_properties,
             "group_properties": group_properties,
             "geoip_disable": disable_geoip,
+            "device_id": device_id,
         }
 
         if flag_keys_to_evaluate:
@@ -621,7 +653,28 @@ class Client(object):
         if flag_options["should_send"]:
             try:
                 if flag_options["only_evaluate_locally"] is True:
-                    # Only use local evaluation
+                    # Local evaluation explicitly requested
+                    feature_variants = self.get_all_flags(
+                        distinct_id,
+                        groups=(groups or {}),
+                        person_properties=flag_options["person_properties"],
+                        group_properties=flag_options["group_properties"],
+                        disable_geoip=disable_geoip,
+                        only_evaluate_locally=True,
+                        flag_keys_to_evaluate=flag_options["flag_keys_filter"],
+                    )
+                elif flag_options["only_evaluate_locally"] is False:
+                    # Remote evaluation explicitly requested
+                    feature_variants = self.get_feature_variants(
+                        distinct_id,
+                        groups,
+                        person_properties=flag_options["person_properties"],
+                        group_properties=flag_options["group_properties"],
+                        disable_geoip=disable_geoip,
+                        flag_keys_to_evaluate=flag_options["flag_keys_filter"],
+                    )
+                elif self.feature_flags:
+                    # Local flags available, prefer local evaluation
                     feature_variants = self.get_all_flags(
                         distinct_id,
                         groups=(groups or {}),
@@ -632,7 +685,7 @@ class Client(object):
                         flag_keys_to_evaluate=flag_options["flag_keys_filter"],
                     )
                 else:
-                    # Default behavior - use remote evaluation
+                    # Fall back to remote evaluation
                     feature_variants = self.get_feature_variants(
                         distinct_id,
                         groups,
@@ -646,15 +699,6 @@ class Client(object):
                     f"[FEATURE FLAGS] Unable to get feature variants: {e}"
                 )
 
-        elif self.feature_flags and event != "$feature_flag_called":
-            # Local evaluation is enabled, flags are loaded, so try and get all flags we can without going to the server
-            feature_variants = self.get_all_flags(
-                distinct_id,
-                groups=(groups or {}),
-                disable_geoip=disable_geoip,
-                only_evaluate_locally=True,
-            )
-
         for feature, variant in (feature_variants or {}).items():
             extra_properties[f"$feature/{feature}"] = variant
 
@@ -973,15 +1017,12 @@ class Client(object):
                         "values": all_exceptions_with_trace,
                     },
                 },
+                in_app_include=self.in_app_modules,
                 project_root=self.project_root,
             )
             all_exceptions_with_trace_and_in_app = event["exception"]["values"]
 
             properties = {
-                "$exception_type": all_exceptions_with_trace_and_in_app[0].get("type"),
-                "$exception_message": all_exceptions_with_trace_and_in_app[0].get(
-                    "value"
-                ),
                 "$exception_list": all_exceptions_with_trace_and_in_app,
                 **properties,
             }
@@ -1146,17 +1187,25 @@ class Client(object):
             posthog.join()
             ```
         """
-        for consumer in self.consumers:
-            consumer.pause()
-            try:
-                consumer.join()
-            except RuntimeError:
-                # consumer thread has not started
-                pass
+        if self.consumers:
+            for consumer in self.consumers:
+                consumer.pause()
+                try:
+                    consumer.join()
+                except RuntimeError:
+                    # consumer thread has not started
+                    pass
 
         if self.poller:
             self.poller.stop()
 
+        # Shutdown the cache provider (release locks, cleanup)
+        if self._flag_definition_cache_provider:
+            try:
+                self._flag_definition_cache_provider.shutdown()
+            except Exception as e:
+                self.log.error(f"[FEATURE FLAGS] Cache provider shutdown error: {e}")
+
     def shutdown(self):
         """
         Flush all messages and cleanly shutdown the client. Call this before the process ends in serverless environments to avoid data loss.
@@ -1172,7 +1221,71 @@ class Client(object):
         if self.exception_capture:
             self.exception_capture.close()
 
+    def _update_flag_state(
+        self, data: FlagDefinitionCacheData, old_flags_by_key: Optional[dict] = None
+    ) -> None:
+        """Update internal flag state from cache data and invalidate evaluation cache if changed."""
+        self.feature_flags = data["flags"]
+        self.group_type_mapping = data["group_type_mapping"]
+        self.cohorts = data["cohorts"]
+
+        # Invalidate evaluation cache if flag definitions changed
+        if (
+            self.flag_cache
+            and old_flags_by_key is not None
+            and old_flags_by_key != (self.feature_flags_by_key or {})
+        ):
+            old_version = self.flag_definition_version
+            self.flag_definition_version += 1
+            self.flag_cache.invalidate_version(old_version)
+
     def _load_feature_flags(self):
+        should_fetch = True
+        if self._flag_definition_cache_provider:
+            try:
+                should_fetch = (
+                    self._flag_definition_cache_provider.should_fetch_flag_definitions()
+                )
+            except Exception as e:
+                self.log.error(
+                    f"[FEATURE FLAGS] Cache provider should_fetch error: {e}"
+                )
+                # Fail-safe: fetch from API if cache provider errors
+                should_fetch = True
+
+        # If not fetching, try to get from cache
+        if not should_fetch and self._flag_definition_cache_provider:
+            try:
+                cached_data = (
+                    self._flag_definition_cache_provider.get_flag_definitions()
+                )
+                if cached_data:
+                    self.log.debug(
+                        "[FEATURE FLAGS] Using cached flag definitions from external cache"
+                    )
+                    self._update_flag_state(
+                        cached_data, old_flags_by_key=self.feature_flags_by_key or {}
+                    )
+                    self._last_feature_flag_poll = datetime.now(tz=tzutc())
+                    return
+                else:
+                    # Emergency fallback: if cache is empty and we have no flags, fetch anyway.
+                    # There's really no other way of recovering in this case.
+                    if not self.feature_flags:
+                        self.log.debug(
+                            "[FEATURE FLAGS] Cache empty and no flags loaded, falling back to API fetch"
+                        )
+                        should_fetch = True
+            except Exception as e:
+                self.log.error(f"[FEATURE FLAGS] Cache provider get error: {e}")
+                # Fail-safe: fetch from API if cache provider errors
+                should_fetch = True
+
+        if should_fetch:
+            self._fetch_feature_flags_from_api()
+
+    def _fetch_feature_flags_from_api(self):
+        """Fetch feature flags from the PostHog API."""
         try:
             # Store old flags to detect changes
             old_flags_by_key: dict[str, dict] = self.feature_flags_by_key or {}
@@ -1182,25 +1295,54 @@ class Client(object):
                 f"/api/feature_flag/local_evaluation/?token={self.api_key}&send_cohorts",
                 self.host,
                 timeout=10,
+                etag=self._flags_etag,
             )
 
-            self.feature_flags = response["flags"] or []
-            self.group_type_mapping = response["group_type_mapping"] or {}
-            self.cohorts = response["cohorts"] or {}
+            # Update stored ETag (clear if server stops sending one)
+            self._flags_etag = response.etag
 
-            # Check if flag definitions changed and update version
-            if self.flag_cache and old_flags_by_key != (
-                self.feature_flags_by_key or {}
-            ):
-                old_version = self.flag_definition_version
-                self.flag_definition_version += 1
-                self.flag_cache.invalidate_version(old_version)
+            # If 304 Not Modified, flags haven't changed - skip processing
+            if response.not_modified:
+                self.log.debug(
+                    "[FEATURE FLAGS] Flags not modified (304), using cached data"
+                )
+                self._last_feature_flag_poll = datetime.now(tz=tzutc())
+                return
+
+            if response.data is None:
+                self.log.error(
+                    "[FEATURE FLAGS] Unexpected empty response data in non-304 response"
+                )
+                return
+
+            self._update_flag_state(response.data, old_flags_by_key=old_flags_by_key)
+
+            # Store in external cache if provider is configured
+            if self._flag_definition_cache_provider:
+                try:
+                    self._flag_definition_cache_provider.on_flag_definitions_received(
+                        {
+                            "flags": self.feature_flags or [],
+                            "group_type_mapping": self.group_type_mapping or {},
+                            "cohorts": self.cohorts or {},
+                        }
+                    )
+                except Exception as e:
+                    self.log.error(f"[FEATURE FLAGS] Cache provider store error: {e}")
+                    # Flags are already in memory, so continue normally
 
         except APIError as e:
             if e.status == 401:
                 self.log.error(
                     "[FEATURE FLAGS] Error loading feature flags: To use feature flags, please set a valid personal_api_key. More information: https://posthog.com/docs/api/overview"
                 )
+                self.feature_flags = []
+                self.group_type_mapping = {}
+                self.cohorts = {}
+
+                if self.flag_cache:
+                    self.flag_cache.clear()
+
                 if self.debug:
                     raise APIError(
                         status=401,
@@ -1277,6 +1419,7 @@ class Client(object):
         person_properties=None,
         group_properties=None,
         warn_on_unknown_groups=True,
+        device_id=None,
     ) -> FlagValue:
         groups = groups or {}
         person_properties = person_properties or {}
@@ -1294,7 +1437,8 @@ class Client(object):
         flag_filters = feature_flag.get("filters") or {}
         aggregation_group_type_index = flag_filters.get("aggregation_group_type_index")
         if aggregation_group_type_index is not None:
-            group_name = self.group_type_mapping.get(str(aggregation_group_type_index))
+            group_type_mapping = self.group_type_mapping or {}
+            group_name = group_type_mapping.get(str(aggregation_group_type_index))
 
             if not group_name:
                 self.log.warning(
@@ -1316,22 +1460,35 @@ class Client(object):
                     )
                 return False
 
+            if group_name not in group_properties:
+                raise InconclusiveMatchError(
+                    f"Flag has no group properties for group '{group_name}'"
+                )
             focused_group_properties = group_properties[group_name]
+            group_key = groups[group_name]
             return match_feature_flag_properties(
                 feature_flag,
-                groups[group_name],
+                group_key,
                 focused_group_properties,
-                self.feature_flags_by_key,
-                evaluation_cache,
+                cohort_properties=self.cohorts,
+                flags_by_key=self.feature_flags_by_key,
+                evaluation_cache=evaluation_cache,
+                device_id=device_id,
+                bucketing_value=group_key,
             )
         else:
+            bucketing_value = resolve_bucketing_value(
+                feature_flag, distinct_id, device_id
+            )
             return match_feature_flag_properties(
                 feature_flag,
                 distinct_id,
                 person_properties,
-                self.cohorts,
-                self.feature_flags_by_key,
-                evaluation_cache,
+                cohort_properties=self.cohorts,
+                flags_by_key=self.feature_flags_by_key,
+                evaluation_cache=evaluation_cache,
+                device_id=device_id,
+                bucketing_value=bucketing_value,
             )
 
     def feature_enabled(
@@ -1345,6 +1502,7 @@ class Client(object):
         only_evaluate_locally=False,
         send_feature_flag_events=True,
         disable_geoip=None,
+        device_id: Optional[str] = None,
     ):
         """
         Check if a feature flag is enabled for a user.
@@ -1358,6 +1516,7 @@ class Client(object):
             only_evaluate_locally: Whether to only evaluate locally.
             send_feature_flag_events: Whether to send feature flag events.
             disable_geoip: Whether to disable GeoIP for this request.
+            device_id: The device ID for this request.
 
         Examples:
             ```python
@@ -1380,12 +1539,26 @@ class Client(object):
             only_evaluate_locally=only_evaluate_locally,
             send_feature_flag_events=send_feature_flag_events,
             disable_geoip=disable_geoip,
+            device_id=device_id,
         )
 
         if response is None:
             return None
         return bool(response)
 
+    def _get_stale_flag_fallback(
+        self, distinct_id: ID_TYPES, key: str
+    ) -> Optional[FeatureFlagResult]:
+        """Returns a stale cached flag value if available, otherwise None."""
+        if self.flag_cache:
+            stale_result = self.flag_cache.get_stale_cached_flag(distinct_id, key)
+            if stale_result:
+                self.log.info(
+                    f"[FEATURE FLAGS] Using stale cached value for flag {key}"
+                )
+                return stale_result
+        return None
+
     def _get_feature_flag_result(
         self,
         key: str,
@@ -1398,6 +1571,7 @@ class Client(object):
         only_evaluate_locally=False,
         send_feature_flag_events=True,
         disable_geoip=None,
+        device_id: Optional[str] = None,
     ) -> Optional[FeatureFlagResult]:
         if self.disabled:
             return None
@@ -1418,9 +1592,15 @@ class Client(object):
         flag_result = None
         flag_details = None
         request_id = None
+        evaluated_at = None
+        feature_flag_error: Optional[str] = None
+
+        # Resolve device_id from context if not provided
+        if device_id is None:
+            device_id = get_context_device_id()
 
         flag_value = self._locally_evaluate_flag(
-            key, distinct_id, groups, person_properties, group_properties
+            key, distinct_id, groups, person_properties, group_properties, device_id
         )
         flag_was_locally_evaluated = flag_value is not None
 
@@ -1442,14 +1622,25 @@ class Client(object):
                 )
         elif not only_evaluate_locally:
             try:
-                flag_details, request_id = self._get_feature_flag_details_from_server(
-                    key,
-                    distinct_id,
-                    groups,
-                    person_properties,
-                    group_properties,
-                    disable_geoip,
+                flag_details, request_id, evaluated_at, errors_while_computing = (
+                    self._get_feature_flag_details_from_server(
+                        key,
+                        distinct_id,
+                        groups,
+                        person_properties,
+                        group_properties,
+                        disable_geoip,
+                        device_id=device_id,
+                    )
                 )
+                errors = []
+                if errors_while_computing:
+                    errors.append(FeatureFlagError.ERRORS_WHILE_COMPUTING)
+                if flag_details is None:
+                    errors.append(FeatureFlagError.FLAG_MISSING)
+                if errors:
+                    feature_flag_error = ",".join(errors)
+
                 flag_result = FeatureFlagResult.from_flag_details(
                     flag_details, override_match_value
                 )
@@ -1463,19 +1654,26 @@ class Client(object):
                 self.log.debug(
                     f"Successfully computed flag remotely: #{key} -> #{flag_result}"
                 )
+            except QuotaLimitError as e:
+                self.log.warning(f"[FEATURE FLAGS] Quota limit exceeded: {e}")
+                feature_flag_error = FeatureFlagError.QUOTA_LIMITED
+                flag_result = self._get_stale_flag_fallback(distinct_id, key)
+            except RequestsTimeout as e:
+                self.log.warning(f"[FEATURE FLAGS] Request timed out: {e}")
+                feature_flag_error = FeatureFlagError.TIMEOUT
+                flag_result = self._get_stale_flag_fallback(distinct_id, key)
+            except RequestsConnectionError as e:
+                self.log.warning(f"[FEATURE FLAGS] Connection error: {e}")
+                feature_flag_error = FeatureFlagError.CONNECTION_ERROR
+                flag_result = self._get_stale_flag_fallback(distinct_id, key)
+            except APIError as e:
+                self.log.warning(f"[FEATURE FLAGS] API error: {e}")
+                feature_flag_error = FeatureFlagError.api_error(e.status)
+                flag_result = self._get_stale_flag_fallback(distinct_id, key)
             except Exception as e:
                 self.log.exception(f"[FEATURE FLAGS] Unable to get flag remotely: {e}")
-
-                # Fallback to cached value if remote evaluation fails
-                if self.flag_cache:
-                    stale_result = self.flag_cache.get_stale_cached_flag(
-                        distinct_id, key
-                    )
-                    if stale_result:
-                        self.log.info(
-                            f"[FEATURE FLAGS] Using stale cached value for flag {key}"
-                        )
-                        flag_result = stale_result
+                feature_flag_error = FeatureFlagError.UNKNOWN_ERROR
+                flag_result = self._get_stale_flag_fallback(distinct_id, key)
 
         if send_feature_flag_events:
             self._capture_feature_flag_called(
@@ -1487,7 +1685,9 @@ class Client(object):
                 groups,
                 disable_geoip,
                 request_id,
+                evaluated_at,
                 flag_details,
+                feature_flag_error,
             )
 
         return flag_result
@@ -1503,6 +1703,7 @@ class Client(object):
         only_evaluate_locally=False,
         send_feature_flag_events=True,
         disable_geoip=None,
+        device_id: Optional[str] = None,
     ) -> Optional[FeatureFlagResult]:
         """
         Get a FeatureFlagResult object which contains the flag result and payload for a key by evaluating locally or remotely
@@ -1527,6 +1728,7 @@ class Client(object):
             only_evaluate_locally: Whether to only evaluate locally.
             send_feature_flag_events: Whether to send feature flag events.
             disable_geoip: Whether to disable GeoIP for this request.
+            device_id: The device ID for this request.
 
         Returns:
             Optional[FeatureFlagResult]: The feature flag result or None if disabled/not found.
@@ -1540,6 +1742,7 @@ class Client(object):
             only_evaluate_locally=only_evaluate_locally,
             send_feature_flag_events=send_feature_flag_events,
             disable_geoip=disable_geoip,
+            device_id=device_id,
         )
 
     def get_feature_flag(
@@ -1553,6 +1756,7 @@ class Client(object):
         only_evaluate_locally=False,
         send_feature_flag_events=True,
         disable_geoip=None,
+        device_id: Optional[str] = None,
     ) -> Optional[FlagValue]:
         """
         Get multivariate feature flag value for a user.
@@ -1566,6 +1770,7 @@ class Client(object):
             only_evaluate_locally: Whether to only evaluate locally.
             send_feature_flag_events: Whether to send feature flag events.
             disable_geoip: Whether to disable GeoIP for this request.
+            device_id: The device ID for this request.
 
         Examples:
             ```python
@@ -1588,6 +1793,7 @@ class Client(object):
             only_evaluate_locally=only_evaluate_locally,
             send_feature_flag_events=send_feature_flag_events,
             disable_geoip=disable_geoip,
+            device_id=device_id,
         )
         return feature_flag_result.get_value() if feature_flag_result else None
 
@@ -1598,6 +1804,7 @@ class Client(object):
         groups: dict[str, str],
         person_properties: dict[str, str],
         group_properties: dict[str, str],
+        device_id: Optional[str] = None,
     ) -> Optional[FlagValue]:
         if self.feature_flags is None and self.personal_api_key:
             self.load_feature_flags()
@@ -1617,6 +1824,7 @@ class Client(object):
                         groups=groups,
                         person_properties=person_properties,
                         group_properties=group_properties,
+                        device_id=device_id,
                     )
                     self.log.debug(
                         f"Successfully computed flag locally: {key} -> {response}"
@@ -1639,8 +1847,9 @@ class Client(object):
         person_properties=None,
         group_properties=None,
         only_evaluate_locally=False,
-        send_feature_flag_events=True,
+        send_feature_flag_events=False,
         disable_geoip=None,
+        device_id: Optional[str] = None,
     ):
         """
         Get the payload for a feature flag.
@@ -1653,8 +1862,9 @@ class Client(object):
             person_properties: A dictionary of person properties.
             group_properties: A dictionary of group properties.
             only_evaluate_locally: Whether to only evaluate locally.
-            send_feature_flag_events: Whether to send feature flag events.
+            send_feature_flag_events: Deprecated. Use get_feature_flag() instead if you need events.
             disable_geoip: Whether to disable GeoIP for this request.
+            device_id: The device ID for this request.
 
         Examples:
             ```python
@@ -1669,6 +1879,14 @@ class Client(object):
         Category:
             Feature flags
         """
+        if send_feature_flag_events:
+            warnings.warn(
+                "send_feature_flag_events is deprecated in get_feature_flag_payload() and will be removed "
+                "in a future version. Use get_feature_flag() if you want to send $feature_flag_called events.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+
         feature_flag_result = self._get_feature_flag_result(
             key,
             distinct_id,
@@ -1679,6 +1897,7 @@ class Client(object):
             only_evaluate_locally=only_evaluate_locally,
             send_feature_flag_events=send_feature_flag_events,
             disable_geoip=disable_geoip,
+            device_id=device_id,
         )
         return feature_flag_result.payload if feature_flag_result else None
 
@@ -1690,9 +1909,11 @@ class Client(object):
         person_properties: dict[str, str],
         group_properties: dict[str, str],
         disable_geoip: Optional[bool],
-    ) -> tuple[Optional[FeatureFlag], Optional[str]]:
+        device_id: Optional[str] = None,
+    ) -> tuple[Optional[FeatureFlag], Optional[str], Optional[int], bool]:
         """
-        Calls /flags and returns the flag details and request id
+        Calls /flags and returns the flag details, request id, evaluated at timestamp,
+        and whether there were errors while computing flags.
         """
         resp_data = self.get_flags_decision(
             distinct_id,
@@ -1701,11 +1922,14 @@ class Client(object):
             group_properties,
             disable_geoip,
             flag_keys_to_evaluate=[key],
+            device_id=device_id,
         )
         request_id = resp_data.get("requestId")
+        evaluated_at = resp_data.get("evaluatedAt")
+        errors_while_computing = resp_data.get("errorsWhileComputingFlags", False)
         flags = resp_data.get("flags")
         flag_details = flags.get(key) if flags else None
-        return flag_details, request_id
+        return flag_details, request_id, evaluated_at, errors_while_computing
 
     def _capture_feature_flag_called(
         self,
@@ -1717,16 +1941,20 @@ class Client(object):
         groups: Dict[str, str],
         disable_geoip: Optional[bool],
         request_id: Optional[str],
+        evaluated_at: Optional[int],
         flag_details: Optional[FeatureFlag],
+        feature_flag_error: Optional[str] = None,
     ):
         feature_flag_reported_key = (
             f"{key}_{'::null::' if response is None else str(response)}"
         )
 
-        if (
-            feature_flag_reported_key
-            not in self.distinct_ids_feature_flags_reported[distinct_id]
-        ):
+        reported_flags = self.distinct_ids_feature_flags_reported.get(distinct_id)
+        if reported_flags is None:
+            reported_flags = set()
+            self.distinct_ids_feature_flags_reported[distinct_id] = reported_flags
+
+        if feature_flag_reported_key not in reported_flags:
             properties: dict[str, Any] = {
                 "$feature_flag": key,
                 "$feature_flag_response": response,
@@ -1740,6 +1968,8 @@ class Client(object):
 
             if request_id:
                 properties["$feature_flag_request_id"] = request_id
+            if evaluated_at:
+                properties["$feature_flag_evaluated_at"] = evaluated_at
             if isinstance(flag_details, FeatureFlag):
                 if flag_details.reason and flag_details.reason.description:
                     properties["$feature_flag_reason"] = flag_details.reason.description
@@ -1750,6 +1980,8 @@ class Client(object):
                         )
                     if flag_details.metadata.id:
                         properties["$feature_flag_id"] = flag_details.metadata.id
+            if feature_flag_error:
+                properties["$feature_flag_error"] = feature_flag_error
 
             self.capture(
                 "$feature_flag_called",
@@ -1758,9 +1990,7 @@ class Client(object):
                 groups=groups,
                 disable_geoip=disable_geoip,
             )
-            self.distinct_ids_feature_flags_reported[distinct_id].add(
-                feature_flag_reported_key
-            )
+            reported_flags.add(feature_flag_reported_key)
 
     def get_remote_config_payload(self, key: str):
         if self.disabled:
@@ -1817,6 +2047,7 @@ class Client(object):
         only_evaluate_locally=False,
         disable_geoip=None,
         flag_keys_to_evaluate: Optional[list[str]] = None,
+        device_id: Optional[str] = None,
     ) -> Optional[dict[str, Union[bool, str]]]:
         """
         Get all feature flags for a user.
@@ -1830,6 +2061,7 @@ class Client(object):
             disable_geoip: Whether to disable GeoIP for this request.
             flag_keys_to_evaluate: A list of specific flag keys to evaluate. If provided,
                 only these flags will be evaluated, improving performance.
+            device_id: The device ID for this request.
 
         Examples:
             ```python
@@ -1847,6 +2079,7 @@ class Client(object):
             only_evaluate_locally=only_evaluate_locally,
             disable_geoip=disable_geoip,
             flag_keys_to_evaluate=flag_keys_to_evaluate,
+            device_id=device_id,
         )
 
         return response["featureFlags"]
@@ -1861,6 +2094,7 @@ class Client(object):
         only_evaluate_locally=False,
         disable_geoip=None,
         flag_keys_to_evaluate: Optional[list[str]] = None,
+        device_id: Optional[str] = None,
     ) -> FlagsAndPayloads:
         """
         Get all feature flags and their payloads for a user.
@@ -1874,6 +2108,7 @@ class Client(object):
             disable_geoip: Whether to disable GeoIP for this request.
             flag_keys_to_evaluate: A list of specific flag keys to evaluate. If provided,
                 only these flags will be evaluated, improving performance.
+            device_id: The device ID for this request.
 
         Examples:
             ```python
@@ -1892,12 +2127,17 @@ class Client(object):
             )
         )
 
+        # Resolve device_id from context if not provided
+        if device_id is None:
+            device_id = get_context_device_id()
+
         response, fallback_to_flags = self._get_all_flags_and_payloads_locally(
             distinct_id,
             groups=groups,
             person_properties=person_properties,
             group_properties=group_properties,
             flag_keys_to_evaluate=flag_keys_to_evaluate,
+            device_id=device_id,
         )
 
         if fallback_to_flags and not only_evaluate_locally:
@@ -1909,6 +2149,7 @@ class Client(object):
                     group_properties=group_properties,
                     disable_geoip=disable_geoip,
                     flag_keys_to_evaluate=flag_keys_to_evaluate,
+                    device_id=device_id,
                 )
                 return to_flags_and_payloads(decide_response)
             except Exception as e:
@@ -1927,6 +2168,7 @@ class Client(object):
         group_properties=None,
         warn_on_unknown_groups=False,
         flag_keys_to_evaluate: Optional[list[str]] = None,
+        device_id: Optional[str] = None,
     ) -> tuple[FlagsAndPayloads, bool]:
         person_properties = person_properties or {}
         group_properties = group_properties or {}
@@ -1956,6 +2198,7 @@ class Client(object):
                         person_properties=person_properties,
                         group_properties=group_properties,
                         warn_on_unknown_groups=warn_on_unknown_groups,
+                        device_id=device_id,
                     )
                     matched_payload = self._compute_payload_locally(
                         flag["key"], flags[flag["key"]]
@@ -2015,9 +2258,9 @@ class Client(object):
             return None
 
         try:
-            from urllib.parse import urlparse, parse_qs
+            from urllib.parse import parse_qs, urlparse
         except ImportError:
-            from urlparse import urlparse, parse_qs
+            from urlparse import parse_qs, urlparse
 
         try:
             parsed = urlparse(cache_url)

posthog/consumer.py

@@ -3,8 +3,6 @@ import logging
 import time
 from threading import Thread
 
-import backoff
-
 from posthog.request import APIError, DatetimeSerializer, batch_post
 
 try:
@@ -84,12 +82,16 @@ class Consumer(Thread):
             self.log.error("error uploading: %s", e)
             success = False
             if self.on_error:
-                self.on_error(e, batch)
+                try:
+                    self.on_error(e, batch)
+                except Exception as e:
+                    self.log.error("on_error handler failed: %s", e)
         finally:
             # mark items as acknowledged from queue
             for item in batch:
                 self.queue.task_done()
-            return success
+
+        return success
 
     def next(self):
         """Return the next batch of items to upload."""
@@ -124,29 +126,41 @@ class Consumer(Thread):
     def request(self, batch):
         """Attempt to upload the batch and retry before raising an error"""
 
-        def fatal_exception(exc):
+        def is_retryable(exc):
             if isinstance(exc, APIError):
                 # retry on server errors and client errors
-                # with 429 status code (rate limited),
+                # with 408 (request timeout) or 429 (rate limited),
                 # don't retry on other client errors
                 if exc.status == "N/A":
                     return False
-                return (400 <= exc.status < 500) and exc.status != 429
+                return not ((400 <= exc.status < 500) and exc.status not in (408, 429))
             else:
                 # retry on all other errors (eg. network)
-                return False
+                return True
 
-        @backoff.on_exception(
-            backoff.expo, Exception, max_tries=self.retries + 1, giveup=fatal_exception
-        )
-        def send_request():
-            batch_post(
-                self.api_key,
-                self.host,
-                gzip=self.gzip,
-                timeout=self.timeout,
-                batch=batch,
-                historical_migration=self.historical_migration,
-            )
+        last_exc = None
+        for attempt in range(self.retries + 1):
+            try:
+                batch_post(
+                    self.api_key,
+                    self.host,
+                    gzip=self.gzip,
+                    timeout=self.timeout,
+                    batch=batch,
+                    historical_migration=self.historical_migration,
+                )
+                return
+            except Exception as e:
+                last_exc = e
+                if not is_retryable(e):
+                    raise
+                if attempt < self.retries:
+                    # Respect Retry-After header if present, otherwise use exponential backoff
+                    retry_after = getattr(e, "retry_after", None)
+                    if retry_after and retry_after > 0:
+                        time.sleep(retry_after)
+                    else:
+                        time.sleep(min(2**attempt, 30))
 
-        send_request()
+        if last_exc:
+            raise last_exc

posthog/contexts.py

@@ -21,6 +21,7 @@ class ContextScope:
         self.capture_exceptions = capture_exceptions
         self.session_id: Optional[str] = None
         self.distinct_id: Optional[str] = None
+        self.device_id: Optional[str] = None
         self.tags: Dict[str, Any] = {}
         self.capture_exception_code_variables: Optional[bool] = None
         self.code_variables_mask_patterns: Optional[list] = None
@@ -32,6 +33,9 @@ class ContextScope:
     def set_distinct_id(self, distinct_id: str):
         self.distinct_id = distinct_id
 
+    def set_device_id(self, device_id: str):
+        self.device_id = device_id
+
     def add_tag(self, key: str, value: Any):
         self.tags[key] = value
 
@@ -61,15 +65,21 @@ class ContextScope:
             return self.parent.get_distinct_id()
         return None
 
+    def get_device_id(self) -> Optional[str]:
+        if self.device_id is not None:
+            return self.device_id
+        if self.parent is not None and not self.fresh:
+            return self.parent.get_device_id()
+        return None
+
     def collect_tags(self) -> Dict[str, Any]:
-        tags = self.tags.copy()
         if self.parent and not self.fresh:
             # We want child tags to take precedence over parent tags,
-            # so we can't use a simple update here, instead collecting
-            # the parent tags and then updating with the child tags.
-            new_tags = self.parent.collect_tags()
-            tags.update(new_tags)
-        return tags
+            # so collect parent tags first, then update with child tags.
+            tags = self.parent.collect_tags()
+            tags.update(self.tags)
+            return tags
+        return self.tags.copy()
 
     def get_capture_exception_code_variables(self) -> Optional[bool]:
         if self.capture_exception_code_variables is not None:
@@ -276,6 +286,39 @@ def get_context_distinct_id() -> Optional[str]:
     return None
 
 
+def set_context_device_id(device_id: str) -> None:
+    """
+    Set the device ID for the current context, associating all feature flag requests in this or
+    child contexts with the given device ID (unless set_context_device_id is called again).
+    Entering a fresh context will clear the context-level device ID.
+
+    Args:
+        device_id: The device ID to associate with the current context and its children.
+
+    Category:
+        Contexts
+    """
+    current_context = _get_current_context()
+    if current_context:
+        current_context.set_device_id(device_id)
+
+
+def get_context_device_id() -> Optional[str]:
+    """
+    Get the device ID for the current context.
+
+    Returns:
+        The device ID if set, None otherwise
+
+    Category:
+        Contexts
+    """
+    current_context = _get_current_context()
+    if current_context:
+        return current_context.get_device_id()
+    return None
+
+
 def set_capture_exception_code_variables_context(enabled: bool) -> None:
     """
     Set whether code variables are captured for the current context.

posthog/exception_utils.py

@@ -14,23 +14,23 @@ import types
 from datetime import datetime
 from types import FrameType, TracebackType  # noqa: F401
 from typing import (  # noqa: F401
+    TYPE_CHECKING,
     Any,
     Dict,
     Iterator,
     List,
     Literal,
     Optional,
+    Pattern,
     Set,
     Tuple,
     TypedDict,
     TypeVar,
     Union,
     cast,
-    TYPE_CHECKING,
-    Pattern,
 )
 
-from posthog.args import ExcInfo, ExceptionArg  # noqa: F401
+from posthog.args import ExceptionArg, ExcInfo  # noqa: F401
 
 try:
     # Python 3.11
@@ -43,22 +43,31 @@ except ImportError:
 DEFAULT_MAX_VALUE_LENGTH = 1024
 
 DEFAULT_CODE_VARIABLES_MASK_PATTERNS = [
-    r"(?i).*password.*",
-    r"(?i).*secret.*",
-    r"(?i).*passwd.*",
-    r"(?i).*pwd.*",
-    r"(?i).*api_key.*",
-    r"(?i).*apikey.*",
-    r"(?i).*auth.*",
-    r"(?i).*credentials.*",
-    r"(?i).*privatekey.*",
-    r"(?i).*private_key.*",
-    r"(?i).*token.*",
+    r"(?i)password",
+    r"(?i)secret",
+    r"(?i)passwd",
+    r"(?i)pwd",
+    r"(?i)api_key",
+    r"(?i)apikey",
+    r"(?i)auth",
+    r"(?i)credentials",
+    r"(?i)privatekey",
+    r"(?i)private_key",
+    r"(?i)token",
+    r"(?i)aws_access_key_id",
+    r"(?i)_pass",
+    r"(?i)sk_",
+    r"(?i)jwt",
 ]
 
 DEFAULT_CODE_VARIABLES_IGNORE_PATTERNS = [r"^__.*"]
 
 CODE_VARIABLES_REDACTED_VALUE = "$$_posthog_redacted_based_on_masking_rules_$$"
+CODE_VARIABLES_TOO_LONG_VALUE = "$$_posthog_value_too_long_$$"
+
+_MAX_VALUE_LENGTH_FOR_PATTERN_MATCH = 5_000
+_MAX_COLLECTION_ITEMS_TO_SCAN = 100
+_REGEX_METACHARACTERS = frozenset(r"\.^$*+?{}[]|()")
 
 DEFAULT_TOTAL_VARIABLES_SIZE_LIMIT = 20 * 1024
 
@@ -924,24 +933,95 @@ def strip_string(value, max_length=None):
     )
 
 
+def _extract_plain_substring(pattern):
+    # Matches inline flag groups like (?i), (?ai), (?ims), etc. that include the 'i' flag.
+    # Python regex flags: a=ASCII, i=IGNORECASE, L=LOCALE, m=MULTILINE, s=DOTALL, u=UNICODE, x=VERBOSE
+    inline_flags = re.match(r"^\(\?[aiLmsux]*i[aiLmsux]*\)", pattern)
+    if not inline_flags:
+        return None
+    remainder = pattern[inline_flags.end() :]
+    if not remainder or any(c in _REGEX_METACHARACTERS for c in remainder):
+        return None
+    return remainder.lower()
+
+
 def _compile_patterns(patterns):
-    compiled = []
+    if not patterns:
+        return None
+    substrings = []
+    regexes = []
     for pattern in patterns:
-        try:
-            compiled.append(re.compile(pattern))
-        except Exception:
-            pass
-    return compiled
+        simple = _extract_plain_substring(pattern)
+        if simple is not None:
+            substrings.append(simple)
+        else:
+            try:
+                regexes.append(re.compile(pattern))
+            except Exception:
+                pass
+    if not substrings and not regexes:
+        return None
+    return (substrings, regexes)
 
 
 def _pattern_matches(name, patterns):
-    for pattern in patterns:
+    if patterns is None:
+        return False
+    substrings, regexes = patterns
+    if substrings:
+        name_lower = name.lower()
+        for s in substrings:
+            if s in name_lower:
+                return True
+    for pattern in regexes:
         if pattern.search(name):
             return True
     return False
 
 
-def _serialize_variable_value(value, limiter, max_length=1024):
+def _mask_sensitive_data(value, compiled_mask, _seen=None):
+    if not compiled_mask:
+        return value
+
+    if isinstance(value, (dict, list, tuple)):
+        if _seen is None:
+            _seen = set()
+        obj_id = id(value)
+        if obj_id in _seen:
+            return "<circular ref>"
+        _seen.add(obj_id)
+
+    if isinstance(value, dict):
+        if len(value) > _MAX_COLLECTION_ITEMS_TO_SCAN:
+            return CODE_VARIABLES_TOO_LONG_VALUE
+        result = {}
+        for k, v in value.items():
+            key_str = str(k) if not isinstance(k, str) else k
+            if len(key_str) > _MAX_VALUE_LENGTH_FOR_PATTERN_MATCH:
+                result[k] = CODE_VARIABLES_TOO_LONG_VALUE
+            elif _pattern_matches(key_str, compiled_mask):
+                result[k] = CODE_VARIABLES_REDACTED_VALUE
+            else:
+                result[k] = _mask_sensitive_data(v, compiled_mask, _seen)
+        return result
+    elif isinstance(value, (list, tuple)):
+        if len(value) > _MAX_COLLECTION_ITEMS_TO_SCAN:
+            return CODE_VARIABLES_TOO_LONG_VALUE
+        masked_items = [
+            _mask_sensitive_data(item, compiled_mask, _seen) for item in value
+        ]
+        return type(value)(masked_items)
+    elif isinstance(value, str):
+        if len(value) > _MAX_VALUE_LENGTH_FOR_PATTERN_MATCH:
+            return CODE_VARIABLES_TOO_LONG_VALUE
+        if _pattern_matches(value, compiled_mask):
+            return CODE_VARIABLES_REDACTED_VALUE
+        return value
+    else:
+        return value
+
+
+def _serialize_variable_value(value, limiter, max_length=1024, compiled_mask=None):
     try:
         if value is None:
             result = "None"
@@ -954,9 +1034,15 @@ def _serialize_variable_value(value, limiter, max_length=1024):
             limiter.add(result_size)
             return value
         elif isinstance(value, str):
-            result = value
+            if len(value) > _MAX_VALUE_LENGTH_FOR_PATTERN_MATCH:
+                result = CODE_VARIABLES_TOO_LONG_VALUE
+            elif compiled_mask and _pattern_matches(value, compiled_mask):
+                result = CODE_VARIABLES_REDACTED_VALUE
+            else:
+                result = value
         else:
-            result = json.dumps(value)
+            masked_value = _mask_sensitive_data(value, compiled_mask)
+            result = json.dumps(masked_value)
 
         if len(result) > max_length:
             result = result[: max_length - 3] + "..."
@@ -969,19 +1055,30 @@ def _serialize_variable_value(value, limiter, max_length=1024):
         return result
     except Exception:
         try:
-            fallback = f"<{type(value).__name__}>"
-            fallback_size = len(fallback)
-            if not limiter.can_add(fallback_size):
+            result = repr(value)
+            if len(result) > max_length:
+                result = result[: max_length - 3] + "..."
+
+            result_size = len(result)
+            if not limiter.can_add(result_size):
                 return None
-            limiter.add(fallback_size)
-            return fallback
+            limiter.add(result_size)
+            return result
         except Exception:
-            fallback = "<unserializable object>"
-            fallback_size = len(fallback)
-            if not limiter.can_add(fallback_size):
-                return None
-            limiter.add(fallback_size)
-            return fallback
+            try:
+                fallback = f"<{type(value).__name__}>"
+                fallback_size = len(fallback)
+                if not limiter.can_add(fallback_size):
+                    return None
+                limiter.add(fallback_size)
+                return fallback
+            except Exception:
+                fallback = "<unserializable object>"
+                fallback_size = len(fallback)
+                if not limiter.can_add(fallback_size):
+                    return None
+                limiter.add(fallback_size)
+                return fallback
 
 
 def _is_simple_type(value):
@@ -1032,7 +1129,9 @@ def serialize_code_variables(
             limiter.add(redacted_size)
             result[name] = redacted_value
         else:
-            serialized = _serialize_variable_value(value, limiter, max_length)
+            serialized = _serialize_variable_value(
+                value, limiter, max_length, compiled_mask
+            )
             if serialized is None:
                 break
             result[name] = serialized
@@ -1042,6 +1141,17 @@ def serialize_code_variables(
 
 def try_attach_code_variables_to_frames(
     all_exceptions, exc_info, mask_patterns, ignore_patterns
+):
+    try:
+        attach_code_variables_to_frames(
+            all_exceptions, exc_info, mask_patterns, ignore_patterns
+        )
+    except Exception:
+        pass
+
+
+def attach_code_variables_to_frames(
+    all_exceptions, exc_info, mask_patterns, ignore_patterns
 ):
     exc_type, exc_value, traceback = exc_info
 

posthog/feature_flags.py

@@ -2,6 +2,7 @@ import datetime
 import hashlib
 import logging
 import re
+import warnings
 from typing import Optional
 
 from dateutil import parser
@@ -34,18 +35,18 @@ class RequiresServerEvaluation(Exception):
     pass
 
 
-# This function takes a distinct_id and a feature flag key and returns a float between 0 and 1.
-# Given the same distinct_id and key, it'll always return the same float. These floats are
+# This function takes a bucketing value and a feature flag key and returns a float between 0 and 1.
+# Given the same bucketing value and key, it'll always return the same float. These floats are
 # uniformly distributed between 0 and 1, so if we want to show this feature to 20% of traffic
-# we can do _hash(key, distinct_id) < 0.2
-def _hash(key: str, distinct_id: str, salt: str = "") -> float:
-    hash_key = f"{key}.{distinct_id}{salt}"
+# we can do _hash(key, bucketing_value) < 0.2
+def _hash(key: str, bucketing_value: str, salt: str = "") -> float:
+    hash_key = f"{key}.{bucketing_value}{salt}"
     hash_val = int(hashlib.sha1(hash_key.encode("utf-8")).hexdigest()[:15], 16)
     return hash_val / __LONG_SCALE__
 
 
-def get_matching_variant(flag, distinct_id):
-    hash_value = _hash(flag["key"], distinct_id, salt="variant")
+def get_matching_variant(flag, bucketing_value):
+    hash_value = _hash(flag["key"], bucketing_value, salt="variant")
     for variant in variant_lookup_table(flag):
         if hash_value >= variant["value_min"] and hash_value < variant["value_max"]:
             return variant["key"]
@@ -68,7 +69,13 @@ def variant_lookup_table(feature_flag):
 
 
 def evaluate_flag_dependency(
-    property, flags_by_key, evaluation_cache, distinct_id, properties, cohort_properties
+    property,
+    flags_by_key,
+    evaluation_cache,
+    distinct_id,
+    properties,
+    cohort_properties,
+    device_id=None,
 ):
     """
     Evaluate a flag dependency property according to the dependency chain algorithm.
@@ -80,6 +87,7 @@ def evaluate_flag_dependency(
         distinct_id: The distinct ID being evaluated
         properties: Person properties for evaluation
         cohort_properties: Cohort properties for evaluation
+        device_id: The device ID for bucketing (optional)
 
     Returns:
         bool: True if all dependencies in the chain evaluate to True, False otherwise
@@ -124,13 +132,27 @@ def evaluate_flag_dependency(
                 else:
                     # Recursively evaluate the dependency
                     try:
+                        dep_flag_filters = dep_flag.get("filters") or {}
+                        dep_aggregation_group_type_index = dep_flag_filters.get(
+                            "aggregation_group_type_index"
+                        )
+                        if dep_aggregation_group_type_index is not None:
+                            # Group flags should continue bucketing by the group key
+                            # from the current evaluation context.
+                            dep_bucketing_value = distinct_id
+                        else:
+                            dep_bucketing_value = resolve_bucketing_value(
+                                dep_flag, distinct_id, device_id
+                            )
                         dep_result = match_feature_flag_properties(
                             dep_flag,
                             distinct_id,
                             properties,
-                            cohort_properties,
-                            flags_by_key,
-                            evaluation_cache,
+                            cohort_properties=cohort_properties,
+                            flags_by_key=flags_by_key,
+                            evaluation_cache=evaluation_cache,
+                            device_id=device_id,
+                            bucketing_value=dep_bucketing_value,
                         )
                         evaluation_cache[dep_flag_key] = dep_result
                     except InconclusiveMatchError as e:
@@ -215,21 +237,54 @@ def matches_dependency_value(expected_value, actual_value):
     return False
 
 
+def resolve_bucketing_value(flag, distinct_id, device_id=None):
+    """Resolve the bucketing value for a flag based on its bucketing_identifier setting.
+
+    Returns:
+        The appropriate identifier string to use for hashing/bucketing.
+
+    Raises:
+        InconclusiveMatchError: If the flag requires device_id but none was provided.
+    """
+    flag_filters = flag.get("filters") or {}
+    bucketing_identifier = flag.get("bucketing_identifier") or flag_filters.get(
+        "bucketing_identifier"
+    )
+    if bucketing_identifier == "device_id":
+        if not device_id:
+            raise InconclusiveMatchError(
+                "Flag requires device_id for bucketing but none was provided"
+            )
+        return device_id
+    return distinct_id
+
+
 def match_feature_flag_properties(
     flag,
     distinct_id,
     properties,
+    *,
     cohort_properties=None,
     flags_by_key=None,
     evaluation_cache=None,
+    device_id=None,
+    bucketing_value=None,
 ) -> FlagValue:
-    flag_conditions = (flag.get("filters") or {}).get("groups") or []
+    if bucketing_value is None:
+        warnings.warn(
+            "Calling match_feature_flag_properties() without bucketing_value is deprecated. "
+            "Pass bucketing_value explicitly. This fallback will be removed in a future major release.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+        bucketing_value = resolve_bucketing_value(flag, distinct_id, device_id)
+
+    flag_filters = flag.get("filters") or {}
+    flag_conditions = flag_filters.get("groups") or []
     is_inconclusive = False
     cohort_properties = cohort_properties or {}
     # Some filters can be explicitly set to null, which require accessing variants like so
-    flag_variants = ((flag.get("filters") or {}).get("multivariate") or {}).get(
-        "variants"
-    ) or []
+    flag_variants = (flag_filters.get("multivariate") or {}).get("variants") or []
     valid_variant_keys = [variant["key"] for variant in flag_variants]
 
     for condition in flag_conditions:
@@ -244,12 +299,14 @@ def match_feature_flag_properties(
                 cohort_properties,
                 flags_by_key,
                 evaluation_cache,
+                bucketing_value=bucketing_value,
+                device_id=device_id,
             ):
                 variant_override = condition.get("variant")
                 if variant_override and variant_override in valid_variant_keys:
                     variant = variant_override
                 else:
-                    variant = get_matching_variant(flag, distinct_id)
+                    variant = get_matching_variant(flag, bucketing_value)
                 return variant or True
         except RequiresServerEvaluation:
             # Static cohort or other missing server-side data - must fallback to API
@@ -277,6 +334,9 @@ def is_condition_match(
     cohort_properties,
     flags_by_key=None,
     evaluation_cache=None,
+    *,
+    bucketing_value,
+    device_id=None,
 ) -> bool:
     rollout_percentage = condition.get("rollout_percentage")
     if len(condition.get("properties") or []) > 0:
@@ -290,6 +350,7 @@ def is_condition_match(
                     flags_by_key,
                     evaluation_cache,
                     distinct_id,
+                    device_id=device_id,
                 )
             elif property_type == "flag":
                 matches = evaluate_flag_dependency(
@@ -299,6 +360,7 @@ def is_condition_match(
                     distinct_id,
                     properties,
                     cohort_properties,
+                    device_id=device_id,
                 )
             else:
                 matches = match_property(prop, properties)
@@ -308,9 +370,9 @@ def is_condition_match(
         if rollout_percentage is None:
             return True
 
-    if rollout_percentage is not None and _hash(feature_flag["key"], distinct_id) > (
-        rollout_percentage / 100
-    ):
+    if rollout_percentage is not None and _hash(
+        feature_flag["key"], bucketing_value
+    ) > (rollout_percentage / 100):
         return False
 
     return True
@@ -454,6 +516,7 @@ def match_cohort(
     flags_by_key=None,
     evaluation_cache=None,
     distinct_id=None,
+    device_id=None,
 ) -> bool:
     # Cohort properties are in the form of property groups like this:
     # {
@@ -478,6 +541,7 @@ def match_cohort(
         flags_by_key,
         evaluation_cache,
         distinct_id,
+        device_id=device_id,
     )
 
 
@@ -488,6 +552,7 @@ def match_property_group(
     flags_by_key=None,
     evaluation_cache=None,
     distinct_id=None,
+    device_id=None,
 ) -> bool:
     if not property_group:
         return True
@@ -512,6 +577,7 @@ def match_property_group(
                     flags_by_key,
                     evaluation_cache,
                     distinct_id,
+                    device_id=device_id,
                 )
                 if property_group_type == "AND":
                     if not matches:
@@ -545,6 +611,7 @@ def match_property_group(
                         flags_by_key,
                         evaluation_cache,
                         distinct_id,
+                        device_id=device_id,
                     )
                 elif prop.get("type") == "flag":
                     matches = evaluate_flag_dependency(
@@ -554,6 +621,7 @@ def match_property_group(
                         distinct_id,
                         property_values,
                         cohort_properties,
+                        device_id=device_id,
                     )
                 else:
                     matches = match_property(prop, property_values)

posthog/flag_definition_cache.py

@@ -0,0 +1,127 @@
+"""
+Flag Definition Cache Provider interface for multi-worker environments.
+
+EXPERIMENTAL: This API may change in future minor version bumps.
+
+This module provides an interface for external caching of feature flag definitions,
+enabling multi-worker environments (Kubernetes, load-balanced servers, serverless
+functions) to share flag definitions and reduce API calls.
+
+Usage:
+
+    from posthog import Posthog
+    from posthog.flag_definition_cache import FlagDefinitionCacheProvider
+
+    cache = RedisFlagDefinitionCache(redis_client, "my-team")
+    posthog = Posthog(
+        "<project_api_key>",
+        personal_api_key="<personal_api_key>",
+        flag_definition_cache_provider=cache,
+    )
+"""
+
+from typing import Any, Dict, List, Optional, Protocol, runtime_checkable
+
+from typing_extensions import Required, TypedDict
+
+
+class FlagDefinitionCacheData(TypedDict):
+    """
+    Data structure for cached flag definitions.
+
+    Attributes:
+        flags: List of feature flag definition dictionaries from the API.
+        group_type_mapping: Mapping of group type indices to group names.
+        cohorts: Dictionary of cohort definitions for local evaluation.
+    """
+
+    flags: Required[List[Dict[str, Any]]]
+    group_type_mapping: Required[Dict[str, str]]
+    cohorts: Required[Dict[str, Any]]
+
+
+@runtime_checkable
+class FlagDefinitionCacheProvider(Protocol):
+    """
+    Interface for external caching of feature flag definitions.
+
+    Enables multi-worker environments to share flag definitions, reducing API
+    calls while ensuring all workers have consistent data.
+
+    EXPERIMENTAL: This API may change in future minor version bumps.
+
+    The four methods handle the complete lifecycle of flag definition caching:
+
+    1. `should_fetch_flag_definitions()` - Called before each poll to determine
+       if this worker should fetch new definitions. Use for distributed lock
+       coordination to ensure only one worker fetches at a time.
+
+    2. `get_flag_definitions()` - Called when `should_fetch_flag_definitions()`
+       returns False. Returns cached definitions if available.
+
+    3. `on_flag_definitions_received()` - Called after successfully fetching
+       new definitions from the API. Store the data in your external cache
+       and release any locks.
+
+    4. `shutdown()` - Called when the PostHog client shuts down. Release any
+       distributed locks and clean up resources.
+
+    Error Handling:
+        All methods are wrapped in try/except. Errors will be logged but will
+        never break flag evaluation. On error:
+        - `should_fetch_flag_definitions()` errors default to fetching (fail-safe)
+        - `get_flag_definitions()` errors fall back to API fetch
+        - `on_flag_definitions_received()` errors are logged but flags remain in memory
+        - `shutdown()` errors are logged but shutdown continues
+    """
+
+    def get_flag_definitions(self) -> Optional[FlagDefinitionCacheData]:
+        """
+        Retrieve cached flag definitions.
+
+        Returns:
+            Cached flag definitions if available and valid, None otherwise.
+            Returning None will trigger a fetch from the API if this worker
+            has no flags loaded yet.
+        """
+        ...
+
+    def should_fetch_flag_definitions(self) -> bool:
+        """
+        Determine whether this instance should fetch new flag definitions.
+
+        Use this for distributed lock coordination. Only one worker should
+        return True to avoid thundering herd problems. A typical implementation
+        uses a distributed lock (e.g., Redis SETNX) that expires after the
+        poll interval.
+
+        Returns:
+            True if this instance should fetch from the API, False otherwise.
+            When False, the client will call `get_flag_definitions()` to
+            retrieve cached data instead.
+        """
+        ...
+
+    def on_flag_definitions_received(self, data: FlagDefinitionCacheData) -> None:
+        """
+        Called after successfully receiving new flag definitions from PostHog.
+
+        Use this to store the data in your external cache and release any
+        distributed locks acquired in `should_fetch_flag_definitions()`.
+
+        Args:
+            data: The flag definitions to cache, containing flags,
+                  group_type_mapping, and cohorts.
+        """
+        ...
+
+    def shutdown(self) -> None:
+        """
+        Called when the PostHog client shuts down.
+
+        Use this to release any distributed locks and clean up resources.
+        This method is called even if `should_fetch_flag_definitions()`
+        returned False, so implementations should handle the case where
+        no lock was acquired.
+        """
+        ...

posthog/request.py

@@ -1,28 +1,163 @@
 import json
 import logging
-from datetime import date, datetime
+import re
+import socket
+from dataclasses import dataclass
+from datetime import date, datetime, timezone
 from gzip import GzipFile
 from io import BytesIO
-from typing import Any, Optional, Union
+from typing import Any, List, Optional, Tuple, Union
 
 import requests
 from dateutil.tz import tzutc
+from requests.adapters import HTTPAdapter  # type: ignore[import-untyped]
+from urllib3.connection import HTTPConnection
 from urllib3.util.retry import Retry
 
 from posthog.utils import remove_trailing_slash
 from posthog.version import VERSION
 
-# Retry on both connect and read errors
-# by default read errors will only retry idempotent HTTP methods (so not POST)
-adapter = requests.adapters.HTTPAdapter(
-    max_retries=Retry(
-        total=2,
-        connect=2,
-        read=2,
+SocketOptions = List[Tuple[int, int, Union[int, bytes]]]
+
+KEEPALIVE_IDLE_SECONDS = 60
+KEEPALIVE_INTERVAL_SECONDS = 60
+KEEPALIVE_PROBE_COUNT = 3
+
+# TCP keepalive probes idle connections to prevent them from being dropped.
+# SO_KEEPALIVE is cross-platform, but timing options vary:
+# - Linux: TCP_KEEPIDLE, TCP_KEEPINTVL, TCP_KEEPCNT
+# - macOS: only SO_KEEPALIVE (uses system defaults)
+# - Windows: TCP_KEEPIDLE, TCP_KEEPINTVL (since Windows 10 1709)
+KEEP_ALIVE_SOCKET_OPTIONS: SocketOptions = list(
+    HTTPConnection.default_socket_options
+) + [
+    (socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1),
+]
+for attr, value in [
+    ("TCP_KEEPIDLE", KEEPALIVE_IDLE_SECONDS),
+    ("TCP_KEEPINTVL", KEEPALIVE_INTERVAL_SECONDS),
+    ("TCP_KEEPCNT", KEEPALIVE_PROBE_COUNT),
+]:
+    if hasattr(socket, attr):
+        KEEP_ALIVE_SOCKET_OPTIONS.append((socket.SOL_TCP, getattr(socket, attr), value))
+
+# Status codes that indicate transient server errors worth retrying
+RETRY_STATUS_FORCELIST = [408, 500, 502, 503, 504]
+
+
+def _mask_tokens_in_url(url: str) -> str:
+    """Mask token values in URLs for safe logging, keeping first 10 chars visible."""
+    return re.sub(r"(token=)([^&]{10})[^&]*", r"\1\2...", url)
+
+
+@dataclass
+class GetResponse:
+    """Response from a GET request with ETag support."""
+
+    data: Any
+    etag: Optional[str] = None
+    not_modified: bool = False
+
+
+class HTTPAdapterWithSocketOptions(HTTPAdapter):
+    """HTTPAdapter with configurable socket options."""
+
+    def __init__(self, *args, socket_options: Optional[SocketOptions] = None, **kwargs):
+        self.socket_options = socket_options
+        super().__init__(*args, **kwargs)
+
+    def init_poolmanager(self, *args, **kwargs):
+        if self.socket_options is not None:
+            kwargs["socket_options"] = self.socket_options
+        super().init_poolmanager(*args, **kwargs)
+
+
+def _build_session(socket_options: Optional[SocketOptions] = None) -> requests.Session:
+    """Build a session for general requests (batch, decide, etc.)."""
+    adapter = HTTPAdapterWithSocketOptions(
+        max_retries=Retry(
+            total=2,
+            connect=2,
+            read=2,
+        ),
+        socket_options=socket_options,
     )
-)
-_session = requests.sessions.Session()
-_session.mount("https://", adapter)
+    session = requests.Session()
+    session.mount("https://", adapter)
+    return session
+
+
+def _build_flags_session(
+    socket_options: Optional[SocketOptions] = None,
+) -> requests.Session:
+    """
+    Build a session for feature flag requests with POST retries.
+
+    Feature flag requests are idempotent (read-only), so retrying POST
+    requests is safe. This session retries on transient server errors
+    (408, 5xx) and network failures with exponential backoff
+    (0.5s, 1s delays between retries).
+    """
+    adapter = HTTPAdapterWithSocketOptions(
+        max_retries=Retry(
+            total=2,
+            connect=2,
+            read=2,
+            backoff_factor=0.5,
+            status_forcelist=RETRY_STATUS_FORCELIST,
+            allowed_methods=["POST"],
+        ),
+        socket_options=socket_options,
+    )
+    session = requests.Session()
+    session.mount("https://", adapter)
+    return session
+
+
+_session = _build_session()
+_flags_session = _build_flags_session()
+_socket_options: Optional[SocketOptions] = None
+_pooling_enabled = True
+
+
+def _get_session() -> requests.Session:
+    if _pooling_enabled:
+        return _session
+    return _build_session(_socket_options)
+
+
+def _get_flags_session() -> requests.Session:
+    if _pooling_enabled:
+        return _flags_session
+    return _build_flags_session(_socket_options)
+
+
+def set_socket_options(socket_options: Optional[SocketOptions]) -> None:
+    """
+    Configure socket options for all HTTP connections.
+
+    Example:
+        from posthog import set_socket_options
+        set_socket_options([(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)])
+    """
+    global _session, _flags_session, _socket_options
+    if socket_options == _socket_options:
+        return
+    _socket_options = socket_options
+    _session = _build_session(socket_options)
+    _flags_session = _build_flags_session(socket_options)
+
+
+def enable_keep_alive() -> None:
+    """Enable TCP keepalive to prevent idle connections from being dropped."""
+    set_socket_options(KEEP_ALIVE_SOCKET_OPTIONS)
+
+
+def disable_connection_reuse() -> None:
+    """Disable connection reuse, creating a fresh connection for each request."""
+    global _pooling_enabled
+    _pooling_enabled = False
+
 
 US_INGESTION_ENDPOINT = "https://us.i.posthog.com"
 EU_INGESTION_ENDPOINT = "https://eu.i.posthog.com"
@@ -48,6 +183,7 @@ def post(
     path=None,
     gzip: bool = False,
     timeout: int = 15,
+    session: Optional[requests.Session] = None,
     **kwargs,
 ) -> requests.Response:
     """Post the `kwargs` to the API"""
@@ -68,7 +204,9 @@ def post(
             gz.write(data.encode("utf-8"))
         data = buf.getvalue()
 
-    res = _session.post(url, data=data, headers=headers, timeout=timeout)
+    res = (session or _get_session()).post(
+        url, data=data, headers=headers, timeout=timeout
+    )
 
     if res.status_code == 200:
         log.debug("data uploaded successfully")
@@ -97,12 +235,31 @@ def _process_response(
             )
             raise QuotaLimitError(res.status_code, "Feature flags quota limited")
         return response
+    retry_after = None
+    retry_after_header = res.headers.get("Retry-After")
+    if retry_after_header:
+        try:
+            retry_after = float(retry_after_header)
+        except (ValueError, TypeError):
+            try:
+                from email.utils import parsedate_to_datetime
+
+                retry_after = max(
+                    0.0,
+                    (
+                        parsedate_to_datetime(retry_after_header)
+                        - datetime.now(timezone.utc)
+                    ).total_seconds(),
+                )
+            except (ValueError, TypeError):
+                pass
+
     try:
         payload = res.json()
         log.debug("received response: %s", payload)
-        raise APIError(res.status_code, payload["detail"])
+        raise APIError(res.status_code, payload["detail"], retry_after=retry_after)
     except (KeyError, ValueError):
-        raise APIError(res.status_code, res.text)
+        raise APIError(res.status_code, res.text, retry_after=retry_after)
 
 
 def decide(
@@ -124,8 +281,16 @@ def flags(
     timeout: int = 15,
     **kwargs,
 ) -> Any:
-    """Post the `kwargs to the flags API endpoint"""
-    res = post(api_key, host, "/flags/?v=2", gzip, timeout, **kwargs)
+    """Post the kwargs to the flags API endpoint with automatic retries."""
+    res = post(
+        api_key,
+        host,
+        "/flags/?v=2",
+        gzip,
+        timeout,
+        session=_get_flags_session(),
+        **kwargs,
+    )
     return _process_response(
         res, success_message="Feature flags evaluated successfully"
     )
@@ -139,12 +304,13 @@ def remote_config(
     timeout: int = 15,
 ) -> Any:
     """Get remote config flag value from remote_config API endpoint"""
-    return get(
+    response = get(
         personal_api_key,
         f"/api/projects/@current/feature_flags/{key}/remote_config?token={project_api_key}",
         host,
         timeout,
     )
+    return response.data
 
 
 def batch_post(
@@ -162,21 +328,51 @@ def batch_post(
 
 
 def get(
-    api_key: str, url: str, host: Optional[str] = None, timeout: Optional[int] = None
-) -> requests.Response:
-    url = remove_trailing_slash(host or DEFAULT_HOST) + url
-    res = requests.get(
-        url,
-        headers={"Authorization": "Bearer %s" % api_key, "User-Agent": USER_AGENT},
-        timeout=timeout,
+    api_key: str,
+    url: str,
+    host: Optional[str] = None,
+    timeout: Optional[int] = None,
+    etag: Optional[str] = None,
+) -> GetResponse:
+    """
+    Make a GET request with optional ETag support.
+
+    If an etag is provided, sends If-None-Match header. Returns GetResponse with:
+    - not_modified=True and data=None if server returns 304
+    - not_modified=False and data=response if server returns 200
+    """
+    log = logging.getLogger("posthog")
+    full_url = remove_trailing_slash(host or DEFAULT_HOST) + url
+    headers = {"Authorization": "Bearer %s" % api_key, "User-Agent": USER_AGENT}
+
+    if etag:
+        headers["If-None-Match"] = etag
+
+    res = _get_session().get(full_url, headers=headers, timeout=timeout)
+
+    masked_url = _mask_tokens_in_url(full_url)
+
+    # Handle 304 Not Modified
+    if res.status_code == 304:
+        log.debug(f"GET {masked_url} returned 304 Not Modified")
+        response_etag = res.headers.get("ETag")
+        return GetResponse(data=None, etag=response_etag or etag, not_modified=True)
+
+    # Handle normal response
+    data = _process_response(
+        res, success_message=f"GET {masked_url} completed successfully"
     )
-    return _process_response(res, success_message=f"GET {url} completed successfully")
+    response_etag = res.headers.get("ETag")
+    return GetResponse(data=data, etag=response_etag, not_modified=False)
 
 
 class APIError(Exception):
-    def __init__(self, status: Union[int, str], message: str):
+    def __init__(
+        self, status: Union[int, str], message: str, retry_after: Optional[float] = None
+    ):
         self.message = message
         self.status = status
+        self.retry_after = retry_after
 
     def __str__(self):
         msg = "[PostHog] {0} ({1})"
@@ -187,6 +383,12 @@ class QuotaLimitError(APIError):
     pass
 
 
+# Re-export requests exceptions for use in client.py
+# This keeps all requests library imports centralized in this module
+RequestsTimeout = requests.exceptions.Timeout
+RequestsConnectionError = requests.exceptions.ConnectionError
+
+
 class DatetimeSerializer(json.JSONEncoder):
     def default(self, obj: Any):
         if isinstance(obj, (date, datetime)):

posthog/test/ai/anthropic/test_anthropic.py

@@ -1,3 +1,4 @@
+import json
 from unittest.mock import patch
 
 import pytest
@@ -306,6 +307,15 @@ def test_basic_completion(mock_client, mock_anthropic_response):
         assert props["$ai_http_status"] == 200
         assert props["foo"] == "bar"
         assert isinstance(props["$ai_latency"], float)
+        # Verify raw usage metadata is passed for backend processing
+        assert "$ai_usage" in props
+        assert props["$ai_usage"] is not None
+        # Verify it's JSON-serializable
+        json.dumps(props["$ai_usage"])
+        # Verify it has expected structure
+        assert isinstance(props["$ai_usage"], dict)
+        assert "input_tokens" in props["$ai_usage"]
+        assert "output_tokens" in props["$ai_usage"]
 
 
 def test_groups(mock_client, mock_anthropic_response):
@@ -918,6 +928,16 @@ def test_streaming_with_tool_calls(mock_client, mock_anthropic_stream_with_tools
         assert props["$ai_cache_read_input_tokens"] == 5
         assert props["$ai_cache_creation_input_tokens"] == 0
 
+        # Verify raw usage is captured in streaming mode (merged from events)
+        assert "$ai_usage" in props
+        assert props["$ai_usage"] is not None
+        # Verify it's JSON-serializable
+        json.dumps(props["$ai_usage"])
+        # Verify it has expected structure (merged from message_start and message_delta)
+        assert isinstance(props["$ai_usage"], dict)
+        assert "input_tokens" in props["$ai_usage"]
+        assert "output_tokens" in props["$ai_usage"]
+
 
 def test_async_streaming_with_tool_calls(mock_client, mock_anthropic_stream_with_tools):
     """Test that tool calls are properly captured in async streaming mode."""

posthog/test/ai/gemini/test_gemini.py

@@ -1,3 +1,4 @@
+import json
 from unittest.mock import MagicMock, patch
 
 import pytest
@@ -34,6 +35,13 @@ def mock_gemini_response():
     # Ensure cache and reasoning tokens are not present (not MagicMock)
     mock_usage.cached_content_token_count = 0
     mock_usage.thoughts_token_count = 0
+    # Make model_dump() return a proper dict for serialization
+    mock_usage.model_dump.return_value = {
+        "prompt_token_count": 20,
+        "candidates_token_count": 10,
+        "cached_content_token_count": 0,
+        "thoughts_token_count": 0,
+    }
     mock_response.usage_metadata = mock_usage
 
     mock_candidate = MagicMock()
@@ -69,6 +77,13 @@ def mock_gemini_response_with_function_calls():
     mock_usage.candidates_token_count = 15
     mock_usage.cached_content_token_count = 0
     mock_usage.thoughts_token_count = 0
+    # Make model_dump() return a proper dict for serialization
+    mock_usage.model_dump.return_value = {
+        "prompt_token_count": 25,
+        "candidates_token_count": 15,
+        "cached_content_token_count": 0,
+        "thoughts_token_count": 0,
+    }
     mock_response.usage_metadata = mock_usage
 
     # Mock function call
@@ -117,6 +132,13 @@ def mock_gemini_response_function_calls_only():
     mock_usage.candidates_token_count = 12
     mock_usage.cached_content_token_count = 0
     mock_usage.thoughts_token_count = 0
+    # Make model_dump() return a proper dict for serialization
+    mock_usage.model_dump.return_value = {
+        "prompt_token_count": 30,
+        "candidates_token_count": 12,
+        "cached_content_token_count": 0,
+        "thoughts_token_count": 0,
+    }
     mock_response.usage_metadata = mock_usage
 
     # Mock function call
@@ -174,6 +196,15 @@ def test_new_client_basic_generation(
     assert props["foo"] == "bar"
     assert "$ai_trace_id" in props
     assert props["$ai_latency"] > 0
+    # Verify raw usage metadata is passed for backend processing
+    assert "$ai_usage" in props
+    assert props["$ai_usage"] is not None
+    # Verify it's JSON-serializable
+    json.dumps(props["$ai_usage"])
+    # Verify it has expected structure
+    assert isinstance(props["$ai_usage"], dict)
+    assert "prompt_token_count" in props["$ai_usage"]
+    assert "candidates_token_count" in props["$ai_usage"]
 
 
 def test_new_client_streaming_with_generate_content_stream(
@@ -407,7 +438,9 @@ def test_new_client_different_input_formats(
     )
     call_args = mock_client.capture.call_args[1]
     props = call_args["properties"]
-    assert props["$ai_input"] == [{"role": "user", "content": "hey"}]
+    assert props["$ai_input"] == [
+        {"role": "user", "content": [{"type": "text", "text": "hey"}]}
+    ]
 
     # Test multiple parts in the parts array
     mock_client.reset_mock()
@@ -418,7 +451,15 @@ def test_new_client_different_input_formats(
     )
     call_args = mock_client.capture.call_args[1]
     props = call_args["properties"]
-    assert props["$ai_input"] == [{"role": "user", "content": "Hello world"}]
+    assert props["$ai_input"] == [
+        {
+            "role": "user",
+            "content": [
+                {"type": "text", "text": "Hello "},
+                {"type": "text", "text": "world"},
+            ],
+        }
+    ]
 
     # Test list input with string
     mock_client.capture.reset_mock()
@@ -800,6 +841,13 @@ def test_streaming_cache_and_reasoning_tokens(mock_client, mock_google_genai_cli
     chunk1_usage.candidates_token_count = 5
     chunk1_usage.cached_content_token_count = 30  # Cache tokens
     chunk1_usage.thoughts_token_count = 0
+    # Make model_dump() return a proper dict for serialization
+    chunk1_usage.model_dump.return_value = {
+        "prompt_token_count": 100,
+        "candidates_token_count": 5,
+        "cached_content_token_count": 30,
+        "thoughts_token_count": 0,
+    }
     chunk1.usage_metadata = chunk1_usage
 
     chunk2 = MagicMock()
@@ -809,6 +857,13 @@ def test_streaming_cache_and_reasoning_tokens(mock_client, mock_google_genai_cli
     chunk2_usage.candidates_token_count = 10
     chunk2_usage.cached_content_token_count = 30  # Same cache tokens
     chunk2_usage.thoughts_token_count = 5  # Reasoning tokens
+    # Make model_dump() return a proper dict for serialization
+    chunk2_usage.model_dump.return_value = {
+        "prompt_token_count": 100,
+        "candidates_token_count": 10,
+        "cached_content_token_count": 30,
+        "thoughts_token_count": 5,
+    }
     chunk2.usage_metadata = chunk2_usage
 
     mock_stream = iter([chunk1, chunk2])
@@ -838,6 +893,16 @@ def test_streaming_cache_and_reasoning_tokens(mock_client, mock_google_genai_cli
     assert props["$ai_cache_read_input_tokens"] == 30
     assert props["$ai_reasoning_tokens"] == 5
 
+    # Verify raw usage is captured in streaming mode (merged from chunks)
+    assert "$ai_usage" in props
+    assert props["$ai_usage"] is not None
+    # Verify it's JSON-serializable
+    json.dumps(props["$ai_usage"])
+    # Verify it has expected structure
+    assert isinstance(props["$ai_usage"], dict)
+    assert "prompt_token_count" in props["$ai_usage"]
+    assert "candidates_token_count" in props["$ai_usage"]
+
 
 def test_web_search_grounding(mock_client, mock_google_genai_client):
     """Test web search detection via grounding_metadata."""

posthog/test/ai/gemini/test_gemini_async.py

@@ -0,0 +1,853 @@
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import pytest
+
+try:
+    from google import genai as google_genai
+
+    from posthog.ai.gemini import AsyncClient
+
+    GEMINI_AVAILABLE = True
+except ImportError:
+    GEMINI_AVAILABLE = False
+
+pytestmark = [
+    pytest.mark.skipif(
+        not GEMINI_AVAILABLE, reason="Google Gemini package is not available"
+    ),
+    pytest.mark.asyncio,
+]
+
+
+@pytest.fixture
+def mock_client():
+    with patch("posthog.client.Client") as mock_client:
+        mock_client.privacy_mode = False
+        yield mock_client
+
+
+@pytest.fixture
+def mock_gemini_response():
+    mock_response = MagicMock()
+    mock_response.text = "Test response from Gemini"
+
+    mock_usage = MagicMock()
+    mock_usage.prompt_token_count = 20
+    mock_usage.candidates_token_count = 10
+    # Ensure cache and reasoning tokens are not present (not MagicMock)
+    mock_usage.cached_content_token_count = 0
+    mock_usage.thoughts_token_count = 0
+    mock_response.usage_metadata = mock_usage
+
+    mock_candidate = MagicMock()
+    mock_candidate.text = "Test response from Gemini"
+    mock_content = MagicMock()
+    mock_part = MagicMock()
+    mock_part.text = "Test response from Gemini"
+    mock_content.parts = [mock_part]
+    mock_candidate.content = mock_content
+    mock_response.candidates = [mock_candidate]
+
+    return mock_response
+
+
+@pytest.fixture
+def mock_google_genai_client():
+    """Mock for the google-genai Client with async support"""
+    with patch.object(google_genai, "Client") as mock_client_class:
+        mock_client_instance = MagicMock()
+        mock_models = MagicMock()
+        mock_aio = MagicMock()
+        mock_aio_models = MagicMock()
+
+        mock_client_instance.models = mock_models
+        mock_client_instance.aio = mock_aio
+        mock_aio.models = mock_aio_models
+
+        mock_client_class.return_value = mock_client_instance
+        yield mock_client_instance
+
+
+@pytest.fixture
+def mock_gemini_response_with_function_calls():
+    mock_response = MagicMock()
+
+    # Mock usage metadata
+    mock_usage = MagicMock()
+    mock_usage.prompt_token_count = 25
+    mock_usage.candidates_token_count = 15
+    mock_usage.cached_content_token_count = 0
+    mock_usage.thoughts_token_count = 0
+    mock_response.usage_metadata = mock_usage
+
+    # Mock function call
+    mock_function_call = MagicMock()
+    mock_function_call.name = "get_current_weather"
+    mock_function_call.args = {"location": "San Francisco"}
+
+    # Mock text part 1
+    mock_text_part1 = MagicMock()
+    mock_text_part1.text = "I'll check the weather for you."
+    type(mock_text_part1).text = mock_text_part1.text
+
+    # Mock text part 2
+    mock_text_part2 = MagicMock()
+    mock_text_part2.text = " Let me look that up."
+    type(mock_text_part2).text = mock_text_part2.text
+
+    # Mock function call part
+    mock_function_part = MagicMock()
+    mock_function_part.function_call = mock_function_call
+    type(mock_function_part).function_call = mock_function_part.function_call
+    del mock_function_part.text
+
+    # Mock content with 2 text parts and 1 function call part
+    mock_content = MagicMock()
+    mock_content.parts = [mock_text_part1, mock_text_part2, mock_function_part]
+
+    # Mock candidate
+    mock_candidate = MagicMock()
+    mock_candidate.content = mock_content
+    mock_response.candidates = [mock_candidate]
+
+    return mock_response
+
+
+async def test_async_client_basic_generation(
+    mock_client, mock_google_genai_client, mock_gemini_response
+):
+    """Test the async Client/AsyncModels API structure"""
+    mock_google_genai_client.aio.models.generate_content = AsyncMock(
+        return_value=mock_gemini_response
+    )
+
+    client = AsyncClient(api_key="test-key", posthog_client=mock_client)
+
+    response = await client.models.generate_content(
+        model="gemini-2.0-flash",
+        contents=["Tell me a fun fact about hedgehogs"],
+        posthog_distinct_id="test-id",
+        posthog_properties={"foo": "bar"},
+    )
+
+    assert response == mock_gemini_response
+    assert mock_client.capture.call_count == 1
+
+    call_args = mock_client.capture.call_args[1]
+    props = call_args["properties"]
+
+    assert call_args["distinct_id"] == "test-id"
+    assert call_args["event"] == "$ai_generation"
+    assert props["$ai_provider"] == "gemini"
+    assert props["$ai_model"] == "gemini-2.0-flash"
+    assert props["$ai_input_tokens"] == 20
+    assert props["$ai_output_tokens"] == 10
+    assert props["foo"] == "bar"
+    assert "$ai_trace_id" in props
+    assert props["$ai_latency"] > 0
+
+
+async def test_async_client_streaming_with_generate_content_stream(
+    mock_client, mock_google_genai_client
+):
+    """Test the async generate_content_stream method"""
+
+    async def mock_streaming_response():
+        mock_chunk1 = MagicMock()
+        mock_chunk1.text = "Hello "
+        mock_usage1 = MagicMock()
+        mock_usage1.prompt_token_count = 10
+        mock_usage1.candidates_token_count = 5
+        mock_usage1.cached_content_token_count = 0
+        mock_usage1.thoughts_token_count = 0
+        mock_chunk1.usage_metadata = mock_usage1
+        yield mock_chunk1
+
+        mock_chunk2 = MagicMock()
+        mock_chunk2.text = "world!"
+        mock_usage2 = MagicMock()
+        mock_usage2.prompt_token_count = 10
+        mock_usage2.candidates_token_count = 10
+        mock_usage2.cached_content_token_count = 0
+        mock_usage2.thoughts_token_count = 0
+        mock_chunk2.usage_metadata = mock_usage2
+        yield mock_chunk2
+
+    # Mock the async generate_content_stream method
+    mock_google_genai_client.aio.models.generate_content_stream = AsyncMock(
+        return_value=mock_streaming_response()
+    )
+
+    client = AsyncClient(api_key="test-key", posthog_client=mock_client)
+
+    response = await client.models.generate_content_stream(
+        model="gemini-2.0-flash",
+        contents=["Write a short story"],
+        posthog_distinct_id="test-id",
+        posthog_properties={"feature": "streaming"},
+    )
+
+    chunks = []
+    async for chunk in response:
+        chunks.append(chunk)
+
+    assert len(chunks) == 2
+    assert chunks[0].text == "Hello "
+    assert chunks[1].text == "world!"
+
+    # Check that the streaming event was captured
+    assert mock_client.capture.call_count == 1
+    call_args = mock_client.capture.call_args[1]
+    props = call_args["properties"]
+
+    assert call_args["distinct_id"] == "test-id"
+    assert call_args["event"] == "$ai_generation"
+    assert props["$ai_provider"] == "gemini"
+    assert props["$ai_model"] == "gemini-2.0-flash"
+    assert props["$ai_input_tokens"] == 10
+    assert props["$ai_output_tokens"] == 10
+    assert props["feature"] == "streaming"
+    assert isinstance(props["$ai_latency"], float)
+
+
+async def test_async_client_streaming_with_tools(mock_client, mock_google_genai_client):
+    """Test that tools are captured in async streaming mode"""
+
+    async def mock_streaming_response():
+        mock_chunk1 = MagicMock()
+        mock_chunk1.text = "I'll check "
+        mock_usage1 = MagicMock()
+        mock_usage1.prompt_token_count = 15
+        mock_usage1.candidates_token_count = 5
+        mock_usage1.cached_content_token_count = 0
+        mock_usage1.thoughts_token_count = 0
+        mock_chunk1.usage_metadata = mock_usage1
+        yield mock_chunk1
+
+        mock_chunk2 = MagicMock()
+        mock_chunk2.text = "the weather"
+        mock_usage2 = MagicMock()
+        mock_usage2.prompt_token_count = 15
+        mock_usage2.candidates_token_count = 10
+        mock_usage2.cached_content_token_count = 0
+        mock_usage2.thoughts_token_count = 0
+        mock_chunk2.usage_metadata = mock_usage2
+        yield mock_chunk2
+
+    # Mock the async generate_content_stream method
+    mock_google_genai_client.aio.models.generate_content_stream = AsyncMock(
+        return_value=mock_streaming_response()
+    )
+
+    client = AsyncClient(api_key="test-key", posthog_client=mock_client)
+
+    # Create mock tools configuration
+    mock_tool = MagicMock()
+    mock_tool.function_declarations = [
+        MagicMock(
+            name="get_current_weather",
+            description="Gets the current weather for a given location.",
+            parameters=MagicMock(
+                type="OBJECT",
+                properties={
+                    "location": MagicMock(
+                        type="STRING",
+                        description="The city and state, e.g. San Francisco, CA",
+                    )
+                },
+                required=["location"],
+            ),
+        )
+    ]
+
+    mock_config = MagicMock()
+    mock_config.tools = [mock_tool]
+
+    response = await client.models.generate_content_stream(
+        model="gemini-2.0-flash",
+        contents=["What's the weather in SF?"],
+        config=mock_config,
+        posthog_distinct_id="test-id",
+        posthog_properties={"feature": "streaming_with_tools"},
+    )
+
+    chunks = []
+    async for chunk in response:
+        chunks.append(chunk)
+
+    assert len(chunks) == 2
+    assert chunks[0].text == "I'll check "
+    assert chunks[1].text == "the weather"
+
+    # Check that the streaming event was captured with tools
+    assert mock_client.capture.call_count == 1
+    call_args = mock_client.capture.call_args[1]
+    props = call_args["properties"]
+
+    assert call_args["distinct_id"] == "test-id"
+    assert call_args["event"] == "$ai_generation"
+    assert props["$ai_provider"] == "gemini"
+    assert props["$ai_model"] == "gemini-2.0-flash"
+    assert props["$ai_input_tokens"] == 15
+    assert props["$ai_output_tokens"] == 10
+    assert props["feature"] == "streaming_with_tools"
+    assert isinstance(props["$ai_latency"], float)
+
+    # Verify that tools are captured in the $ai_tools property in streaming mode
+    assert props["$ai_tools"] == [mock_tool]
+
+
+async def test_async_client_groups(
+    mock_client, mock_google_genai_client, mock_gemini_response
+):
+    """Test groups functionality with async Client API"""
+    mock_google_genai_client.aio.models.generate_content = AsyncMock(
+        return_value=mock_gemini_response
+    )
+
+    client = AsyncClient(api_key="test-key", posthog_client=mock_client)
+
+    await client.models.generate_content(
+        model="gemini-2.0-flash",
+        contents=["Hello"],
+        posthog_distinct_id="test-id",
+        posthog_groups={"company": "company_123"},
+    )
+
+    call_args = mock_client.capture.call_args[1]
+    assert call_args["groups"] == {"company": "company_123"}
+
+
+async def test_async_client_privacy_mode_local(
+    mock_client, mock_google_genai_client, mock_gemini_response
+):
+    """Test local privacy mode with async Client API"""
+    mock_google_genai_client.aio.models.generate_content = AsyncMock(
+        return_value=mock_gemini_response
+    )
+
+    client = AsyncClient(api_key="test-key", posthog_client=mock_client)
+
+    await client.models.generate_content(
+        model="gemini-2.0-flash",
+        contents=["Hello"],
+        posthog_distinct_id="test-id",
+        posthog_privacy_mode=True,
+    )
+
+    call_args = mock_client.capture.call_args[1]
+    props = call_args["properties"]
+    assert props["$ai_input"] is None
+    assert props["$ai_output_choices"] is None
+
+
+async def test_async_client_privacy_mode_global(
+    mock_client, mock_google_genai_client, mock_gemini_response
+):
+    """Test global privacy mode with async Client API"""
+    mock_client.privacy_mode = True
+
+    mock_google_genai_client.aio.models.generate_content = AsyncMock(
+        return_value=mock_gemini_response
+    )
+
+    client = AsyncClient(api_key="test-key", posthog_client=mock_client)
+
+    await client.models.generate_content(
+        model="gemini-2.0-flash",
+        contents=["Hello"],
+        posthog_distinct_id="test-id",
+    )
+
+    call_args = mock_client.capture.call_args[1]
+    props = call_args["properties"]
+    assert props["$ai_input"] is None
+    assert props["$ai_output_choices"] is None
+
+
+async def test_async_client_different_input_formats(
+    mock_client, mock_google_genai_client, mock_gemini_response
+):
+    """Test different input formats with async Client API"""
+    mock_google_genai_client.aio.models.generate_content = AsyncMock(
+        return_value=mock_gemini_response
+    )
+
+    client = AsyncClient(api_key="test-key", posthog_client=mock_client)
+
+    # Test string input
+    await client.models.generate_content(
+        model="gemini-2.0-flash", contents="Hello", posthog_distinct_id="test-id"
+    )
+    call_args = mock_client.capture.call_args[1]
+    props = call_args["properties"]
+    assert props["$ai_input"] == [{"role": "user", "content": "Hello"}]
+
+    # Test Gemini-specific format with parts array
+    mock_client.reset_mock()
+    await client.models.generate_content(
+        model="gemini-2.0-flash",
+        contents=[{"role": "user", "parts": [{"text": "hey"}]}],
+        posthog_distinct_id="test-id",
+    )
+    call_args = mock_client.capture.call_args[1]
+    props = call_args["properties"]
+    assert props["$ai_input"] == [
+        {"role": "user", "content": [{"type": "text", "text": "hey"}]}
+    ]
+
+    # Test multiple parts in the parts array
+    mock_client.reset_mock()
+    await client.models.generate_content(
+        model="gemini-2.0-flash",
+        contents=[{"role": "user", "parts": [{"text": "Hello "}, {"text": "world"}]}],
+        posthog_distinct_id="test-id",
+    )
+    call_args = mock_client.capture.call_args[1]
+    props = call_args["properties"]
+    assert props["$ai_input"] == [
+        {
+            "role": "user",
+            "content": [
+                {"type": "text", "text": "Hello "},
+                {"type": "text", "text": "world"},
+            ],
+        }
+    ]
+
+    # Test list input with string
+    mock_client.capture.reset_mock()
+    await client.models.generate_content(
+        model="gemini-2.0-flash", contents=["List item"], posthog_distinct_id="test-id"
+    )
+    call_args = mock_client.capture.call_args[1]
+    props = call_args["properties"]
+    assert props["$ai_input"] == [{"role": "user", "content": "List item"}]
+
+
+async def test_async_client_model_parameters(
+    mock_client, mock_google_genai_client, mock_gemini_response
+):
+    """Test model parameters with async Client API"""
+    mock_google_genai_client.aio.models.generate_content = AsyncMock(
+        return_value=mock_gemini_response
+    )
+
+    client = AsyncClient(api_key="test-key", posthog_client=mock_client)
+
+    await client.models.generate_content(
+        model="gemini-2.0-flash",
+        contents=["Hello"],
+        posthog_distinct_id="test-id",
+        temperature=0.7,
+        max_tokens=100,
+    )
+
+    call_args = mock_client.capture.call_args[1]
+    props = call_args["properties"]
+    assert props["$ai_model_parameters"]["temperature"] == 0.7
+    assert props["$ai_model_parameters"]["max_tokens"] == 100
+
+
+async def test_async_client_default_settings(
+    mock_client, mock_google_genai_client, mock_gemini_response
+):
+    """Test async client with default PostHog settings"""
+    mock_google_genai_client.aio.models.generate_content = AsyncMock(
+        return_value=mock_gemini_response
+    )
+
+    client = AsyncClient(
+        api_key="test-key",
+        posthog_client=mock_client,
+        posthog_distinct_id="default_user",
+        posthog_properties={"team": "ai"},
+        posthog_privacy_mode=False,
+        posthog_groups={"company": "acme_corp"},
+    )
+
+    # Call without overriding defaults
+    await client.models.generate_content(model="gemini-2.0-flash", contents=["Hello"])
+
+    call_args = mock_client.capture.call_args[1]
+    props = call_args["properties"]
+
+    assert call_args["distinct_id"] == "default_user"
+    assert call_args["groups"] == {"company": "acme_corp"}
+    assert props["team"] == "ai"
+
+
+async def test_async_client_override_defaults(
+    mock_client, mock_google_genai_client, mock_gemini_response
+):
+    """Test overriding async client defaults per call"""
+    mock_google_genai_client.aio.models.generate_content = AsyncMock(
+        return_value=mock_gemini_response
+    )
+
+    client = AsyncClient(
+        api_key="test-key",
+        posthog_client=mock_client,
+        posthog_distinct_id="default_user",
+        posthog_properties={"team": "ai"},
+        posthog_privacy_mode=False,
+        posthog_groups={"company": "acme_corp"},
+    )
+
+    # Override defaults in call
+    await client.models.generate_content(
+        model="gemini-2.0-flash",
+        contents=["Hello"],
+        posthog_distinct_id="specific_user",
+        posthog_properties={"feature": "chat", "urgent": True},
+        posthog_privacy_mode=True,
+        posthog_groups={"organization": "special_org"},
+    )
+
+    call_args = mock_client.capture.call_args[1]
+    props = call_args["properties"]
+
+    # Check overrides
+    assert call_args["distinct_id"] == "specific_user"
+    assert call_args["groups"] == {"organization": "special_org"}
+    assert props["$ai_input"] is None  # privacy mode was overridden
+
+    # Check merged properties (defaults + call-specific)
+    assert props["team"] == "ai"  # from defaults
+    assert props["feature"] == "chat"  # from call
+    assert props["urgent"] is True  # from call
+
+
+async def test_async_vertex_ai_parameters_passed_through(
+    mock_client, mock_google_genai_client, mock_gemini_response
+):
+    """Test that Vertex AI parameters are properly passed to genai.Client"""
+    mock_google_genai_client.aio.models.generate_content = AsyncMock(
+        return_value=mock_gemini_response
+    )
+
+    # Mock credentials object
+    mock_credentials = MagicMock()
+    mock_debug_config = MagicMock()
+    mock_http_options = MagicMock()
+
+    # Create client with Vertex AI parameters
+    AsyncClient(
+        vertexai=True,
+        credentials=mock_credentials,
+        project="test-project",
+        location="us-central1",
+        debug_config=mock_debug_config,
+        http_options=mock_http_options,
+        posthog_client=mock_client,
+    )
+
+    # Verify genai.Client was called with correct parameters
+    google_genai.Client.assert_called_once_with(
+        vertexai=True,
+        credentials=mock_credentials,
+        project="test-project",
+        location="us-central1",
+        debug_config=mock_debug_config,
+        http_options=mock_http_options,
+    )
+
+
+async def test_async_api_key_mode(mock_client, mock_google_genai_client):
+    """Test API key authentication mode with async client"""
+
+    # Create async client with just API key (traditional mode)
+    AsyncClient(
+        api_key="test-api-key",
+        posthog_client=mock_client,
+    )
+
+    # Verify genai.Client was called with only api_key
+    google_genai.Client.assert_called_once_with(api_key="test-api-key")
+
+
+async def test_async_function_calls_in_output_choices(
+    mock_client, mock_google_genai_client, mock_gemini_response_with_function_calls
+):
+    """Test that function calls are properly included in $ai_output_choices with async"""
+    mock_google_genai_client.aio.models.generate_content = AsyncMock(
+        return_value=mock_gemini_response_with_function_calls
+    )
+
+    client = AsyncClient(api_key="test-key", posthog_client=mock_client)
+
+    response = await client.models.generate_content(
+        model="gemini-2.5-flash",
+        contents=["What's the weather in San Francisco?"],
+        posthog_distinct_id="test-id",
+    )
+
+    assert response == mock_gemini_response_with_function_calls
+    assert mock_client.capture.call_count == 1
+
+    call_args = mock_client.capture.call_args[1]
+    props = call_args["properties"]
+
+    assert call_args["distinct_id"] == "test-id"
+    assert call_args["event"] == "$ai_generation"
+    assert props["$ai_provider"] == "gemini"
+    assert props["$ai_model"] == "gemini-2.5-flash"
+    assert props["$ai_output_choices"] == [
+        {
+            "role": "assistant",
+            "content": [
+                {"type": "text", "text": "I'll check the weather for you."},
+                {"type": "text", "text": " Let me look that up."},
+                {
+                    "type": "function",
+                    "function": {
+                        "name": "get_current_weather",
+                        "arguments": {"location": "San Francisco"},
+                    },
+                },
+            ],
+        }
+    ]
+
+    # Check token usage
+    assert props["$ai_input_tokens"] == 25
+    assert props["$ai_output_tokens"] == 15
+    assert props["$ai_http_status"] == 200
+
+
+async def test_async_cache_and_reasoning_tokens(mock_client, mock_google_genai_client):
+    """Test that cache and reasoning tokens are properly extracted with async"""
+    # Create a mock response with cache and reasoning tokens
+    mock_response = MagicMock()
+    mock_response.text = "Test response with cache"
+
+    mock_usage = MagicMock()
+    mock_usage.prompt_token_count = 100
+    mock_usage.candidates_token_count = 50
+    mock_usage.cached_content_token_count = 30  # Cache tokens
+    mock_usage.thoughts_token_count = 10  # Reasoning tokens
+    mock_response.usage_metadata = mock_usage
+
+    # Mock candidates
+    mock_candidate = MagicMock()
+    mock_candidate.text = "Test response with cache"
+    mock_response.candidates = [mock_candidate]
+
+    mock_google_genai_client.aio.models.generate_content = AsyncMock(
+        return_value=mock_response
+    )
+
+    client = AsyncClient(api_key="test-key", posthog_client=mock_client)
+
+    response = await client.models.generate_content(
+        model="gemini-2.5-pro",
+        contents="Test with cache",
+        posthog_distinct_id="test-id",
+    )
+
+    assert response == mock_response
+    assert mock_client.capture.call_count == 1
+
+    call_args = mock_client.capture.call_args[1]
+    props = call_args["properties"]
+
+    # Check that all token types are present
+    assert props["$ai_input_tokens"] == 100
+    assert props["$ai_output_tokens"] == 50
+    assert props["$ai_cache_read_input_tokens"] == 30
+    assert props["$ai_reasoning_tokens"] == 10
+
+
+async def test_async_streaming_cache_and_reasoning_tokens(
+    mock_client, mock_google_genai_client
+):
+    """Test that cache and reasoning tokens are properly extracted in async streaming"""
+
+    async def mock_streaming_response():
+        # Create mock chunks with cache and reasoning tokens
+        chunk1 = MagicMock()
+        chunk1.text = "Hello "
+        chunk1_usage = MagicMock()
+        chunk1_usage.prompt_token_count = 100
+        chunk1_usage.candidates_token_count = 5
+        chunk1_usage.cached_content_token_count = 30  # Cache tokens
+        chunk1_usage.thoughts_token_count = 0
+        chunk1.usage_metadata = chunk1_usage
+        yield chunk1
+
+        chunk2 = MagicMock()
+        chunk2.text = "world!"
+        chunk2_usage = MagicMock()
+        chunk2_usage.prompt_token_count = 100
+        chunk2_usage.candidates_token_count = 10
+        chunk2_usage.cached_content_token_count = 30  # Same cache tokens
+        chunk2_usage.thoughts_token_count = 5  # Reasoning tokens
+        chunk2.usage_metadata = chunk2_usage
+        yield chunk2
+
+    mock_google_genai_client.aio.models.generate_content_stream = AsyncMock(
+        return_value=mock_streaming_response()
+    )
+
+    client = AsyncClient(api_key="test-key", posthog_client=mock_client)
+
+    response = await client.models.generate_content_stream(
+        model="gemini-2.5-pro",
+        contents="Test streaming with cache",
+        posthog_distinct_id="test-id",
+    )
+
+    # Consume the stream
+    result = []
+    async for chunk in response:
+        result.append(chunk)
+
+    assert len(result) == 2
+
+    # Check PostHog capture was called
+    assert mock_client.capture.call_count == 1
+
+    call_args = mock_client.capture.call_args[1]
+    props = call_args["properties"]
+
+    # Check that all token types are present (should use final chunk's usage)
+    assert props["$ai_input_tokens"] == 100
+    assert props["$ai_output_tokens"] == 10
+    assert props["$ai_cache_read_input_tokens"] == 30
+    assert props["$ai_reasoning_tokens"] == 5
+
+
+async def test_async_web_search_grounding(mock_client, mock_google_genai_client):
+    """Test async web search detection via grounding_metadata."""
+
+    # Create mock response with grounding metadata
+    mock_response = MagicMock()
+
+    # Mock usage metadata
+    mock_usage = MagicMock()
+    mock_usage.prompt_token_count = 60
+    mock_usage.candidates_token_count = 40
+    mock_usage.cached_content_token_count = 0
+    mock_usage.thoughts_token_count = 0
+    mock_response.usage_metadata = mock_usage
+
+    # Mock grounding metadata
+    mock_grounding_chunk = MagicMock()
+    mock_grounding_chunk.uri = "https://example.com"
+
+    mock_grounding_metadata = MagicMock()
+    mock_grounding_metadata.grounding_chunks = [mock_grounding_chunk]
+
+    # Mock text part
+    mock_text_part = MagicMock()
+    mock_text_part.text = "According to search results..."
+    type(mock_text_part).text = mock_text_part.text
+
+    # Mock content with parts
+    mock_content = MagicMock()
+    mock_content.parts = [mock_text_part]
+
+    # Mock candidate with grounding metadata
+    mock_candidate = MagicMock()
+    mock_candidate.content = mock_content
+    mock_candidate.grounding_metadata = mock_grounding_metadata
+    type(mock_candidate).grounding_metadata = mock_candidate.grounding_metadata
+
+    mock_response.candidates = [mock_candidate]
+    mock_response.text = "According to search results..."
+
+    # Mock the async generate_content method
+    mock_google_genai_client.aio.models.generate_content = AsyncMock(
+        return_value=mock_response
+    )
+
+    client = AsyncClient(api_key="test-key", posthog_client=mock_client)
+    response = await client.models.generate_content(
+        model="gemini-2.5-flash",
+        contents="What's the latest news?",
+        posthog_distinct_id="test-id",
+    )
+
+    assert response == mock_response
+    assert mock_client.capture.call_count == 1
+
+    call_args = mock_client.capture.call_args[1]
+    props = call_args["properties"]
+
+    # Verify web search count is detected (binary for grounding)
+    assert props["$ai_web_search_count"] == 1
+    assert props["$ai_input_tokens"] == 60
+    assert props["$ai_output_tokens"] == 40
+
+
+async def test_async_streaming_with_web_search(mock_client, mock_google_genai_client):
+    """Test that web search count is properly captured in async streaming mode."""
+
+    async def mock_streaming_response():
+        # Create chunk 1 with grounding metadata
+        mock_chunk1 = MagicMock()
+        mock_chunk1.text = "According to "
+
+        mock_usage1 = MagicMock()
+        mock_usage1.prompt_token_count = 30
+        mock_usage1.candidates_token_count = 5
+        mock_usage1.cached_content_token_count = 0
+        mock_usage1.thoughts_token_count = 0
+        mock_chunk1.usage_metadata = mock_usage1
+
+        # Add grounding metadata to first chunk
+        mock_grounding_chunk = MagicMock()
+        mock_grounding_chunk.uri = "https://example.com"
+
+        mock_grounding_metadata = MagicMock()
+        mock_grounding_metadata.grounding_chunks = [mock_grounding_chunk]
+
+        mock_candidate1 = MagicMock()
+        mock_candidate1.grounding_metadata = mock_grounding_metadata
+        type(mock_candidate1).grounding_metadata = mock_candidate1.grounding_metadata
+
+        mock_chunk1.candidates = [mock_candidate1]
+        yield mock_chunk1
+
+        # Create chunk 2
+        mock_chunk2 = MagicMock()
+        mock_chunk2.text = "search results..."
+
+        mock_usage2 = MagicMock()
+        mock_usage2.prompt_token_count = 30
+        mock_usage2.candidates_token_count = 15
+        mock_usage2.cached_content_token_count = 0
+        mock_usage2.thoughts_token_count = 0
+        mock_chunk2.usage_metadata = mock_usage2
+
+        mock_candidate2 = MagicMock()
+        mock_chunk2.candidates = [mock_candidate2]
+        yield mock_chunk2
+
+    # Mock the async generate_content_stream method
+    mock_google_genai_client.aio.models.generate_content_stream = AsyncMock(
+        return_value=mock_streaming_response()
+    )
+
+    client = AsyncClient(api_key="test-key", posthog_client=mock_client)
+
+    response = await client.models.generate_content_stream(
+        model="gemini-2.5-flash",
+        contents="What's the latest news?",
+        posthog_distinct_id="test-id",
+    )
+
+    chunks = []
+    async for chunk in response:
+        chunks.append(chunk)
+
+    assert len(chunks) == 2
+    assert mock_client.capture.call_count == 1
+
+    call_args = mock_client.capture.call_args[1]
+    props = call_args["properties"]
+
+    # Verify web search count is detected (binary for grounding)
+    assert props["$ai_web_search_count"] == 1
+    assert props["$ai_input_tokens"] == 30
+    assert props["$ai_output_tokens"] == 15

posthog/test/ai/langchain/test_callbacks.py

@@ -1638,6 +1638,95 @@ def test_anthropic_provider_subtracts_cache_tokens(mock_client):
     assert generation_args["properties"]["$ai_cache_read_input_tokens"] == 800
 
 
+def test_anthropic_provider_subtracts_cache_write_tokens(mock_client):
+    """Test that Anthropic provider correctly subtracts cache write tokens from input tokens."""
+    from langchain_core.outputs import LLMResult, ChatGeneration
+    from langchain_core.messages import AIMessage
+    from uuid import uuid4
+
+    cb = CallbackHandler(mock_client)
+    run_id = uuid4()
+
+    # Set up with Anthropic provider
+    cb._set_llm_metadata(
+        serialized={},
+        run_id=run_id,
+        messages=[{"role": "user", "content": "test"}],
+        metadata={"ls_provider": "anthropic", "ls_model_name": "claude-3-sonnet"},
+    )
+
+    # Response with cache creation: 1000 input (includes 800 being written to cache)
+    response = LLMResult(
+        generations=[
+            [
+                ChatGeneration(
+                    message=AIMessage(content="Response"),
+                    generation_info={
+                        "usage_metadata": {
+                            "input_tokens": 1000,
+                            "output_tokens": 50,
+                            "cache_creation_input_tokens": 800,
+                        }
+                    },
+                )
+            ]
+        ],
+        llm_output={},
+    )
+
+    cb._pop_run_and_capture_generation(run_id, None, response)
+
+    generation_args = mock_client.capture.call_args_list[0][1]
+    assert generation_args["properties"]["$ai_input_tokens"] == 200  # 1000 - 800
+    assert generation_args["properties"]["$ai_cache_creation_input_tokens"] == 800
+
+
+def test_anthropic_provider_subtracts_both_cache_read_and_write_tokens(mock_client):
+    """Test that Anthropic provider correctly subtracts both cache read and write tokens."""
+    from langchain_core.outputs import LLMResult, ChatGeneration
+    from langchain_core.messages import AIMessage
+    from uuid import uuid4
+
+    cb = CallbackHandler(mock_client)
+    run_id = uuid4()
+
+    # Set up with Anthropic provider
+    cb._set_llm_metadata(
+        serialized={},
+        run_id=run_id,
+        messages=[{"role": "user", "content": "test"}],
+        metadata={"ls_provider": "anthropic", "ls_model_name": "claude-3-sonnet"},
+    )
+
+    # Response with both cache read and creation
+    response = LLMResult(
+        generations=[
+            [
+                ChatGeneration(
+                    message=AIMessage(content="Response"),
+                    generation_info={
+                        "usage_metadata": {
+                            "input_tokens": 2000,
+                            "output_tokens": 50,
+                            "cache_read_input_tokens": 800,
+                            "cache_creation_input_tokens": 500,
+                        }
+                    },
+                )
+            ]
+        ],
+        llm_output={},
+    )
+
+    cb._pop_run_and_capture_generation(run_id, None, response)
+
+    generation_args = mock_client.capture.call_args_list[0][1]
+    # 2000 - 800 (read) - 500 (write) = 700
+    assert generation_args["properties"]["$ai_input_tokens"] == 700
+    assert generation_args["properties"]["$ai_cache_read_input_tokens"] == 800
+    assert generation_args["properties"]["$ai_cache_creation_input_tokens"] == 500
+
+
 def test_openai_cache_read_tokens(mock_client):
     """Test that OpenAI cache read tokens are captured correctly."""
     prompt = ChatPromptTemplate.from_messages(
@@ -2092,10 +2181,12 @@ def test_zero_input_tokens_with_cache_read(mock_client):
     assert generation_props["$ai_cache_read_input_tokens"] == 50
 
 
-def test_cache_write_tokens_not_subtracted_from_input(mock_client):
-    """Test that cache_creation_input_tokens (cache write) do NOT affect input_tokens.
+def test_non_anthropic_cache_write_tokens_not_subtracted_from_input(mock_client):
+    """Test that cache_creation_input_tokens do NOT affect input_tokens for non-Anthropic providers.
 
-    Only cache_read_tokens should be subtracted from input_tokens, not cache_write_tokens.
+    When no provider metadata is set (or for non-Anthropic providers), cache tokens should
+    NOT be subtracted from input_tokens. This is because different providers report tokens
+    differently - only Anthropic's LangChain integration requires subtraction.
     """
     prompt = ChatPromptTemplate.from_messages([("user", "Create cache")])
 
@@ -2350,3 +2441,206 @@ def test_billable_with_real_chain(mock_client):
     assert props["$ai_billable"] is True
     assert props["$ai_model"] == "fake-model"
     assert props["$ai_provider"] == "fake"
+
+
+# Exception Capture Integration Tests
+
+
+def test_exception_autocapture_on_span_error():
+    """Test that capture_exception is called when a span errors and autocapture is enabled."""
+    mock_client = MagicMock()
+    mock_client.privacy_mode = False
+    mock_client.enable_exception_autocapture = True
+    mock_client.capture_exception.return_value = "exception-uuid-123"
+
+    def failing_span(_):
+        raise ValueError("test error")
+
+    callbacks = [CallbackHandler(mock_client)]
+    chain = RunnableLambda(failing_span)
+
+    try:
+        chain.invoke({}, config={"callbacks": callbacks})
+    except ValueError:
+        pass
+
+    # Verify capture_exception was called
+    assert mock_client.capture_exception.call_count == 1
+    exception_call = mock_client.capture_exception.call_args
+    assert isinstance(exception_call[0][0], ValueError)
+    assert str(exception_call[0][0]) == "test error"
+
+
+def test_exception_autocapture_adds_exception_id_to_span_event():
+    """Test that $exception_event_id is added to the span event properties."""
+    mock_client = MagicMock()
+    mock_client.privacy_mode = False
+    mock_client.enable_exception_autocapture = True
+    mock_client.capture_exception.return_value = "exception-uuid-456"
+
+    def failing_span(_):
+        raise ValueError("test error")
+
+    callbacks = [CallbackHandler(mock_client)]
+    chain = RunnableLambda(failing_span)
+
+    try:
+        chain.invoke({}, config={"callbacks": callbacks})
+    except ValueError:
+        pass
+
+    # Find the span event (should have $ai_is_error=True)
+    span_calls = [
+        call
+        for call in mock_client.capture.call_args_list
+        if call[1].get("properties", {}).get("$ai_is_error") is True
+    ]
+    assert len(span_calls) >= 1
+
+    span_props = span_calls[0][1]["properties"]
+    assert span_props["$exception_event_id"] == "exception-uuid-456"
+    assert span_props["$ai_error"] == "ValueError: test error"
+
+
+def test_exception_autocapture_disabled_does_not_capture():
+    """Test that capture_exception is NOT called when autocapture is disabled."""
+    mock_client = MagicMock()
+    mock_client.privacy_mode = False
+    mock_client.enable_exception_autocapture = False
+
+    def failing_span(_):
+        raise ValueError("test error")
+
+    callbacks = [CallbackHandler(mock_client)]
+    chain = RunnableLambda(failing_span)
+
+    try:
+        chain.invoke({}, config={"callbacks": callbacks})
+    except ValueError:
+        pass
+
+    # Verify capture_exception was NOT called
+    assert mock_client.capture_exception.call_count == 0
+
+    # But the span event should still have error info
+    span_calls = [
+        call
+        for call in mock_client.capture.call_args_list
+        if call[1].get("properties", {}).get("$ai_is_error") is True
+    ]
+    assert len(span_calls) >= 1
+
+    span_props = span_calls[0][1]["properties"]
+    assert "$exception_event_id" not in span_props
+    assert span_props["$ai_error"] == "ValueError: test error"
+
+
+def test_exception_autocapture_on_llm_generation_error(mock_client):
+    """Test that capture_exception is called when an LLM generation fails."""
+    mock_client.privacy_mode = False
+    mock_client.enable_exception_autocapture = True
+    mock_client.capture_exception.return_value = "exception-uuid-789"
+
+    callbacks = CallbackHandler(mock_client)
+    run_id = uuid.uuid4()
+
+    # Simulate LLM start
+    callbacks.on_llm_start(
+        serialized={"kwargs": {"openai_api_base": "https://api.openai.com"}},
+        prompts=["Hello"],
+        run_id=run_id,
+    )
+
+    # Simulate LLM error
+    error = Exception("API rate limit exceeded")
+    callbacks.on_llm_error(error, run_id=run_id)
+
+    # Verify capture_exception was called
+    assert mock_client.capture_exception.call_count == 1
+    exception_call = mock_client.capture_exception.call_args
+    assert exception_call[0][0] is error
+
+    # Verify the generation event has $exception_event_id
+    generation_calls = [
+        call
+        for call in mock_client.capture.call_args_list
+        if call[1].get("event") == "$ai_generation"
+    ]
+    assert len(generation_calls) == 1
+
+    gen_props = generation_calls[0][1]["properties"]
+    assert gen_props["$exception_event_id"] == "exception-uuid-789"
+    assert gen_props["$ai_is_error"] is True
+
+
+def test_exception_autocapture_passes_ai_properties_to_exception():
+    """Test that AI properties are passed to the exception event."""
+    mock_client = MagicMock()
+    mock_client.privacy_mode = False
+    mock_client.enable_exception_autocapture = True
+    mock_client.capture_exception.return_value = "exception-uuid-abc"
+
+    callbacks = CallbackHandler(
+        mock_client,
+        distinct_id="user-123",
+        properties={"custom_prop": "custom_value"},
+    )
+    run_id = uuid.uuid4()
+
+    # Simulate LLM start
+    callbacks.on_llm_start(
+        serialized={"kwargs": {"openai_api_base": "https://api.openai.com"}},
+        prompts=["Hello"],
+        run_id=run_id,
+    )
+
+    # Simulate LLM error
+    error = Exception("API error")
+    callbacks.on_llm_error(error, run_id=run_id)
+
+    # Verify capture_exception received the properties
+    exception_call = mock_client.capture_exception.call_args
+    props = exception_call[1]["properties"]
+
+    # Should have AI-related properties
+    assert "$ai_trace_id" in props
+    assert "$ai_is_error" in props
+    assert props["$ai_is_error"] is True
+
+    # Should have distinct_id passed through
+    assert exception_call[1]["distinct_id"] == "user-123"
+
+
+def test_exception_autocapture_none_return_no_exception_id():
+    """Test that when capture_exception returns None, no $exception_event_id is added."""
+    mock_client = MagicMock()
+    mock_client.privacy_mode = False
+    mock_client.enable_exception_autocapture = True
+    mock_client.capture_exception.return_value = (
+        None  # e.g., exception already captured
+    )
+
+    def failing_span(_):
+        raise ValueError("test error")
+
+    callbacks = [CallbackHandler(mock_client)]
+    chain = RunnableLambda(failing_span)
+
+    try:
+        chain.invoke({}, config={"callbacks": callbacks})
+    except ValueError:
+        pass
+
+    # capture_exception was called but returned None
+    assert mock_client.capture_exception.call_count == 1
+
+    # Span event should NOT have $exception_event_id
+    span_calls = [
+        call
+        for call in mock_client.capture.call_args_list
+        if call[1].get("properties", {}).get("$ai_is_error") is True
+    ]
+    assert len(span_calls) >= 1
+
+    span_props = span_calls[0][1]["properties"]
+    assert "$exception_event_id" not in span_props

posthog/test/ai/openai/test_openai.py

@@ -1,3 +1,4 @@
+import json
 import time
 from unittest.mock import AsyncMock, patch
 
@@ -496,6 +497,15 @@ def test_basic_completion(mock_client, mock_openai_response):
         assert props["$ai_http_status"] == 200
         assert props["foo"] == "bar"
         assert isinstance(props["$ai_latency"], float)
+        # Verify raw usage metadata is passed for backend processing
+        assert "$ai_usage" in props
+        assert props["$ai_usage"] is not None
+        # Verify it's JSON-serializable
+        json.dumps(props["$ai_usage"])
+        # Verify it has expected structure
+        assert isinstance(props["$ai_usage"], dict)
+        assert "prompt_tokens" in props["$ai_usage"]
+        assert "completion_tokens" in props["$ai_usage"]
 
 
 def test_embeddings(mock_client, mock_embedding_response):
@@ -922,6 +932,16 @@ def test_streaming_with_tool_calls(mock_client, streaming_tool_call_chunks):
         assert props["$ai_input_tokens"] == 20
         assert props["$ai_output_tokens"] == 15
 
+        # Verify raw usage is captured in streaming mode
+        assert "$ai_usage" in props
+        assert props["$ai_usage"] is not None
+        # Verify it's JSON-serializable
+        json.dumps(props["$ai_usage"])
+        # Verify it has expected structure (merged from chunks)
+        assert isinstance(props["$ai_usage"], dict)
+        assert "prompt_tokens" in props["$ai_usage"]
+        assert "completion_tokens" in props["$ai_usage"]
+
 
 # test responses api
 def test_responses_api(mock_client, mock_openai_response_with_responses_api):
@@ -1676,3 +1696,459 @@ async def test_async_chat_streaming_with_web_search(
     assert props["$ai_web_search_count"] == 1
     assert props["$ai_input_tokens"] == 20
     assert props["$ai_output_tokens"] == 15
+
+
+# Tests for model extraction fallback (stored prompts support)
+
+
+def test_streaming_chat_extracts_model_from_chunk_when_not_in_kwargs(mock_client):
+    """Test that model is extracted from streaming chunks when not provided in kwargs (stored prompts)."""
+
+    # Create streaming chunks with model field but we won't pass model in kwargs
+    chunks = [
+        ChatCompletionChunk(
+            id="chunk1",
+            model="gpt-4o-stored-prompt",  # Model comes from response, not request
+            object="chat.completion.chunk",
+            created=1234567890,
+            choices=[
+                ChoiceChunk(
+                    index=0,
+                    delta=ChoiceDelta(role="assistant", content="Hello"),
+                    finish_reason=None,
+                )
+            ],
+        ),
+        ChatCompletionChunk(
+            id="chunk2",
+            model="gpt-4o-stored-prompt",
+            object="chat.completion.chunk",
+            created=1234567891,
+            choices=[
+                ChoiceChunk(
+                    index=0,
+                    delta=ChoiceDelta(content=" world"),
+                    finish_reason="stop",
+                )
+            ],
+            usage=CompletionUsage(
+                prompt_tokens=10,
+                completion_tokens=5,
+                total_tokens=15,
+            ),
+        ),
+    ]
+
+    with patch("openai.resources.chat.completions.Completions.create") as mock_create:
+        mock_create.return_value = chunks
+
+        client = OpenAI(api_key="test-key", posthog_client=mock_client)
+
+        # Note: NOT passing model in kwargs - simulates stored prompt usage
+        response_generator = client.chat.completions.create(
+            messages=[{"role": "user", "content": "Hello"}],
+            stream=True,
+            posthog_distinct_id="test-id",
+        )
+
+        # Consume the generator
+        list(response_generator)
+
+        assert mock_client.capture.call_count == 1
+        call_args = mock_client.capture.call_args[1]
+        props = call_args["properties"]
+
+        # Model should be extracted from chunk, not kwargs
+        assert props["$ai_model"] == "gpt-4o-stored-prompt"
+
+
+def test_streaming_chat_prefers_kwargs_model_over_chunk_model(mock_client):
+    """Test that model from kwargs takes precedence over model from chunk."""
+    chunks = [
+        ChatCompletionChunk(
+            id="chunk1",
+            model="gpt-4o-from-response",
+            object="chat.completion.chunk",
+            created=1234567890,
+            choices=[
+                ChoiceChunk(
+                    index=0,
+                    delta=ChoiceDelta(role="assistant", content="Hello"),
+                    finish_reason="stop",
+                )
+            ],
+            usage=CompletionUsage(
+                prompt_tokens=10,
+                completion_tokens=5,
+                total_tokens=15,
+            ),
+        ),
+    ]
+
+    with patch("openai.resources.chat.completions.Completions.create") as mock_create:
+        mock_create.return_value = chunks
+
+        client = OpenAI(api_key="test-key", posthog_client=mock_client)
+
+        response_generator = client.chat.completions.create(
+            model="gpt-4o-from-kwargs",  # Explicitly passed model
+            messages=[{"role": "user", "content": "Hello"}],
+            stream=True,
+            posthog_distinct_id="test-id",
+        )
+
+        list(response_generator)
+
+        call_args = mock_client.capture.call_args[1]
+        props = call_args["properties"]
+
+        # kwargs model should take precedence
+        assert props["$ai_model"] == "gpt-4o-from-kwargs"
+
+
+def test_streaming_responses_api_extracts_model_from_response_object(mock_client):
+    """Test that Responses API streaming extracts model from chunk.response.model (stored prompts)."""
+    from unittest.mock import MagicMock
+    from openai.types.responses import ResponseUsage
+
+    chunks = []
+
+    # Content chunk
+    chunk1 = MagicMock()
+    chunk1.type = "response.text.delta"
+    chunk1.text = "Test response"
+    # No response attribute on content chunks
+    del chunk1.response
+    chunks.append(chunk1)
+
+    # Final chunk with response object containing model
+    chunk2 = MagicMock()
+    chunk2.type = "response.completed"
+    chunk2.response = MagicMock()
+    chunk2.response.model = "gpt-4o-mini-stored"  # Model from stored prompt
+    chunk2.response.usage = ResponseUsage(
+        input_tokens=20,
+        output_tokens=10,
+        total_tokens=30,
+        input_tokens_details={"prompt_tokens": 20, "cached_tokens": 0},
+        output_tokens_details={"reasoning_tokens": 0},
+    )
+    chunk2.response.output = ["Test response"]
+    chunks.append(chunk2)
+
+    with patch("openai.resources.responses.Responses.create") as mock_create:
+        mock_create.return_value = iter(chunks)
+
+        client = OpenAI(api_key="test-key", posthog_client=mock_client)
+
+        # Note: NOT passing model - simulates stored prompt
+        response_generator = client.responses.create(
+            input=[{"role": "user", "content": "Hello"}],
+            stream=True,
+            posthog_distinct_id="test-id",
+        )
+
+        list(response_generator)
+
+        call_args = mock_client.capture.call_args[1]
+        props = call_args["properties"]
+
+        # Model should be extracted from chunk.response.model
+        assert props["$ai_model"] == "gpt-4o-mini-stored"
+
+
+def test_non_streaming_extracts_model_from_response(mock_client):
+    """Test that non-streaming calls extract model from response when not in kwargs."""
+    # Create a response with model but we won't pass model in kwargs
+    mock_response = ChatCompletion(
+        id="test",
+        model="gpt-4o-stored-prompt",
+        object="chat.completion",
+        created=int(time.time()),
+        choices=[
+            Choice(
+                finish_reason="stop",
+                index=0,
+                message=ChatCompletionMessage(
+                    content="Test response",
+                    role="assistant",
+                ),
+            )
+        ],
+        usage=CompletionUsage(
+            completion_tokens=10,
+            prompt_tokens=20,
+            total_tokens=30,
+        ),
+    )
+
+    with patch(
+        "openai.resources.chat.completions.Completions.create",
+        return_value=mock_response,
+    ):
+        client = OpenAI(api_key="test-key", posthog_client=mock_client)
+
+        # Note: NOT passing model in kwargs
+        response = client.chat.completions.create(
+            messages=[{"role": "user", "content": "Hello"}],
+            posthog_distinct_id="test-id",
+        )
+
+        assert response == mock_response
+        call_args = mock_client.capture.call_args[1]
+        props = call_args["properties"]
+
+        # Model should be extracted from response.model
+        assert props["$ai_model"] == "gpt-4o-stored-prompt"
+
+
+def test_non_streaming_responses_api_extracts_model_from_response(mock_client):
+    """Test that non-streaming Responses API extracts model from response when not in kwargs."""
+    mock_response = Response(
+        id="test",
+        model="gpt-4o-mini-stored",
+        object="response",
+        created_at=1741476542,
+        status="completed",
+        error=None,
+        incomplete_details=None,
+        instructions=None,
+        max_output_tokens=None,
+        tools=[],
+        tool_choice="auto",
+        output=[
+            ResponseOutputMessage(
+                id="msg_123",
+                type="message",
+                role="assistant",
+                status="completed",
+                content=[
+                    ResponseOutputText(
+                        type="output_text",
+                        text="Test response",
+                        annotations=[],
+                    )
+                ],
+            )
+        ],
+        parallel_tool_calls=True,
+        previous_response_id=None,
+        usage=ResponseUsage(
+            input_tokens=10,
+            output_tokens=10,
+            input_tokens_details={"prompt_tokens": 10, "cached_tokens": 0},
+            output_tokens_details={"reasoning_tokens": 0},
+            total_tokens=20,
+        ),
+        user=None,
+        metadata={},
+    )
+
+    with patch(
+        "openai.resources.responses.Responses.create",
+        return_value=mock_response,
+    ):
+        client = OpenAI(api_key="test-key", posthog_client=mock_client)
+
+        # Note: NOT passing model in kwargs
+        response = client.responses.create(
+            input="Hello",
+            posthog_distinct_id="test-id",
+        )
+
+        assert response == mock_response
+        call_args = mock_client.capture.call_args[1]
+        props = call_args["properties"]
+
+        # Model should be extracted from response.model
+        assert props["$ai_model"] == "gpt-4o-mini-stored"
+
+
+def test_non_streaming_returns_none_when_no_model(mock_client):
+    """Test that non-streaming returns None (not 'unknown') when model is not available anywhere."""
+    # Create a response without model attribute using real OpenAI types
+    mock_response = ChatCompletion(
+        id="test",
+        model="",  # Will be removed below
+        object="chat.completion",
+        created=int(time.time()),
+        choices=[
+            Choice(
+                finish_reason="stop",
+                index=0,
+                message=ChatCompletionMessage(
+                    content="Test response",
+                    role="assistant",
+                ),
+            )
+        ],
+        usage=CompletionUsage(
+            completion_tokens=5,
+            prompt_tokens=10,
+            total_tokens=15,
+        ),
+    )
+    # Remove model attribute to simulate missing model
+    object.__delattr__(mock_response, "model")
+
+    with patch(
+        "openai.resources.chat.completions.Completions.create",
+        return_value=mock_response,
+    ):
+        client = OpenAI(api_key="test-key", posthog_client=mock_client)
+
+        # Note: NOT passing model in kwargs and response has no model
+        client.chat.completions.create(
+            messages=[{"role": "user", "content": "Hello"}],
+            posthog_distinct_id="test-id",
+        )
+
+        call_args = mock_client.capture.call_args[1]
+        props = call_args["properties"]
+
+        # Should be None, NOT "unknown" (to avoid incorrect cost matching)
+        assert props["$ai_model"] is None
+
+
+def test_streaming_falls_back_to_unknown_when_no_model(mock_client):
+    """Test that streaming falls back to 'unknown' when model is not available anywhere."""
+    from unittest.mock import MagicMock
+
+    # Create a chunk without model attribute
+    chunk = MagicMock()
+    chunk.choices = [MagicMock()]
+    chunk.choices[0].delta = MagicMock()
+    chunk.choices[0].delta.content = "Hello"
+    chunk.choices[0].delta.role = "assistant"
+    chunk.choices[0].delta.tool_calls = None
+    chunk.usage = CompletionUsage(
+        prompt_tokens=10,
+        completion_tokens=5,
+        total_tokens=15,
+    )
+    # Explicitly remove model attribute
+    del chunk.model
+
+    with patch("openai.resources.chat.completions.Completions.create") as mock_create:
+        mock_create.return_value = [chunk]
+
+        client = OpenAI(api_key="test-key", posthog_client=mock_client)
+
+        response_generator = client.chat.completions.create(
+            messages=[{"role": "user", "content": "Hello"}],
+            stream=True,
+            posthog_distinct_id="test-id",
+        )
+
+        list(response_generator)
+
+        call_args = mock_client.capture.call_args[1]
+        props = call_args["properties"]
+
+        # Should fall back to "unknown"
+        assert props["$ai_model"] == "unknown"
+
+
+@pytest.mark.asyncio
+async def test_async_streaming_chat_extracts_model_from_chunk(mock_client):
+    """Test async streaming extracts model from chunk when not in kwargs."""
+    chunks = [
+        ChatCompletionChunk(
+            id="chunk1",
+            model="gpt-4o-async-stored",
+            object="chat.completion.chunk",
+            created=1234567890,
+            choices=[
+                ChoiceChunk(
+                    index=0,
+                    delta=ChoiceDelta(role="assistant", content="Hello"),
+                    finish_reason="stop",
+                )
+            ],
+            usage=CompletionUsage(
+                prompt_tokens=10,
+                completion_tokens=5,
+                total_tokens=15,
+            ),
+        ),
+    ]
+
+    async def mock_create(self, **kwargs):
+        async def chunk_iterable():
+            for chunk in chunks:
+                yield chunk
+
+        return chunk_iterable()
+
+    with patch(
+        "openai.resources.chat.completions.AsyncCompletions.create", new=mock_create
+    ):
+        client = AsyncOpenAI(api_key="test-key", posthog_client=mock_client)
+
+        # Note: NOT passing model
+        response_stream = await client.chat.completions.create(
+            messages=[{"role": "user", "content": "Hello"}],
+            stream=True,
+            posthog_distinct_id="test-id",
+        )
+
+        async for _ in response_stream:
+            pass
+
+    call_args = mock_client.capture.call_args[1]
+    props = call_args["properties"]
+
+    assert props["$ai_model"] == "gpt-4o-async-stored"
+
+
+@pytest.mark.asyncio
+async def test_async_streaming_responses_extracts_model_from_response(mock_client):
+    """Test async Responses API streaming extracts model from chunk.response.model."""
+    from unittest.mock import MagicMock
+    from openai.types.responses import ResponseUsage
+
+    chunks = []
+
+    chunk1 = MagicMock()
+    chunk1.type = "response.text.delta"
+    chunk1.text = "Test"
+    del chunk1.response
+    chunks.append(chunk1)
+
+    chunk2 = MagicMock()
+    chunk2.type = "response.completed"
+    chunk2.response = MagicMock()
+    chunk2.response.model = "gpt-4o-mini-async-stored"
+    chunk2.response.usage = ResponseUsage(
+        input_tokens=20,
+        output_tokens=10,
+        total_tokens=30,
+        input_tokens_details={"prompt_tokens": 20, "cached_tokens": 0},
+        output_tokens_details={"reasoning_tokens": 0},
+    )
+    chunk2.response.output = ["Test"]
+    chunks.append(chunk2)
+
+    async def mock_create(self, **kwargs):
+        async def chunk_iterable():
+            for chunk in chunks:
+                yield chunk
+
+        return chunk_iterable()
+
+    with patch("openai.resources.responses.AsyncResponses.create", new=mock_create):
+        client = AsyncOpenAI(api_key="test-key", posthog_client=mock_client)
+
+        response_stream = await client.responses.create(
+            input=[{"role": "user", "content": "Hello"}],
+            stream=True,
+            posthog_distinct_id="test-id",
+        )
+
+        async for _ in response_stream:
+            pass
+
+    call_args = mock_client.capture.call_args[1]
+    props = call_args["properties"]
+
+    assert props["$ai_model"] == "gpt-4o-mini-async-stored"

posthog/test/ai/openai_agents/__init__.py

@@ -0,0 +1 @@
+# Tests for OpenAI Agents SDK integration

posthog/test/ai/openai_agents/test_processor.py

@@ -0,0 +1,810 @@
+import logging
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+try:
+    from agents.tracing.span_data import (
+        AgentSpanData,
+        CustomSpanData,
+        FunctionSpanData,
+        GenerationSpanData,
+        GuardrailSpanData,
+        HandoffSpanData,
+        ResponseSpanData,
+        SpeechSpanData,
+        TranscriptionSpanData,
+    )
+
+    from posthog.ai.openai_agents import PostHogTracingProcessor, instrument
+
+    OPENAI_AGENTS_AVAILABLE = True
+except ImportError:
+    OPENAI_AGENTS_AVAILABLE = False
+
+
+# Skip all tests if OpenAI Agents SDK is not available
+pytestmark = pytest.mark.skipif(
+    not OPENAI_AGENTS_AVAILABLE, reason="OpenAI Agents SDK is not available"
+)
+
+
+@pytest.fixture(scope="function")
+def mock_client():
+    client = MagicMock()
+    client.privacy_mode = False
+    logging.getLogger("posthog").setLevel(logging.DEBUG)
+    return client
+
+
+@pytest.fixture(scope="function")
+def processor(mock_client):
+    return PostHogTracingProcessor(
+        client=mock_client,
+        distinct_id="test-user",
+        privacy_mode=False,
+    )
+
+
+@pytest.fixture
+def mock_trace():
+    trace = MagicMock()
+    trace.trace_id = "trace_123456789"
+    trace.name = "Test Workflow"
+    trace.group_id = "group_123"
+    trace.metadata = {"key": "value"}
+    return trace
+
+
+@pytest.fixture
+def mock_span():
+    span = MagicMock()
+    span.trace_id = "trace_123456789"
+    span.span_id = "span_987654321"
+    span.parent_id = None
+    span.started_at = "2024-01-01T00:00:00Z"
+    span.ended_at = "2024-01-01T00:00:01Z"
+    span.error = None
+    return span
+
+
+class TestPostHogTracingProcessor:
+    """Tests for the PostHogTracingProcessor class."""
+
+    def test_initialization(self, mock_client):
+        """Test processor initializes correctly."""
+        processor = PostHogTracingProcessor(
+            client=mock_client,
+            distinct_id="user@example.com",
+            privacy_mode=True,
+            groups={"company": "acme"},
+            properties={"env": "test"},
+        )
+
+        assert processor._client == mock_client
+        assert processor._distinct_id == "user@example.com"
+        assert processor._privacy_mode is True
+        assert processor._groups == {"company": "acme"}
+        assert processor._properties == {"env": "test"}
+
+    def test_initialization_with_callable_distinct_id(self, mock_client, mock_trace):
+        """Test processor with callable distinct_id resolver."""
+
+        def resolver(trace):
+            return trace.metadata.get("user_id", "default")
+
+        processor = PostHogTracingProcessor(
+            client=mock_client,
+            distinct_id=resolver,
+        )
+
+        mock_trace.metadata = {"user_id": "resolved-user"}
+        distinct_id = processor._get_distinct_id(mock_trace)
+        assert distinct_id == "resolved-user"
+
+    def test_on_trace_start_stores_metadata(self, processor, mock_client, mock_trace):
+        """Test that on_trace_start stores metadata but does not capture an event."""
+        processor.on_trace_start(mock_trace)
+
+        mock_client.capture.assert_not_called()
+        assert mock_trace.trace_id in processor._trace_metadata
+
+    def test_on_trace_end_captures_ai_trace(self, processor, mock_client, mock_trace):
+        """Test that on_trace_end captures $ai_trace event."""
+        processor.on_trace_start(mock_trace)
+        processor.on_trace_end(mock_trace)
+
+        mock_client.capture.assert_called_once()
+        call_kwargs = mock_client.capture.call_args[1]
+
+        assert call_kwargs["event"] == "$ai_trace"
+        assert call_kwargs["distinct_id"] == "test-user"
+        assert call_kwargs["properties"]["$ai_trace_id"] == "trace_123456789"
+        assert call_kwargs["properties"]["$ai_trace_name"] == "Test Workflow"
+        assert call_kwargs["properties"]["$ai_provider"] == "openai"
+        assert call_kwargs["properties"]["$ai_framework"] == "openai-agents"
+        assert "$ai_latency" in call_kwargs["properties"]
+
+    def test_personless_mode_when_no_distinct_id(self, mock_client, mock_trace):
+        """Test that trace events use personless mode when no distinct_id is provided."""
+        processor = PostHogTracingProcessor(
+            client=mock_client,
+        )
+
+        processor.on_trace_start(mock_trace)
+        processor.on_trace_end(mock_trace)
+
+        call_kwargs = mock_client.capture.call_args[1]
+        assert call_kwargs["properties"]["$process_person_profile"] is False
+        # Should fallback to trace_id as the distinct_id
+        assert call_kwargs["distinct_id"] == mock_trace.trace_id
+
+    def test_personless_mode_for_spans_when_no_distinct_id(
+        self, mock_client, mock_trace, mock_span
+    ):
+        """Test that span events use personless mode when no distinct_id is provided."""
+        processor = PostHogTracingProcessor(
+            client=mock_client,
+        )
+
+        processor.on_trace_start(mock_trace)
+        mock_client.capture.reset_mock()
+
+        span_data = GenerationSpanData(model="gpt-4o")
+        mock_span.span_data = span_data
+
+        processor.on_span_start(mock_span)
+        processor.on_span_end(mock_span)
+
+        call_kwargs = mock_client.capture.call_args[1]
+        assert call_kwargs["properties"]["$process_person_profile"] is False
+        assert call_kwargs["distinct_id"] == mock_span.trace_id
+
+    def test_personless_mode_when_callable_returns_none(
+        self, mock_client, mock_trace, mock_span
+    ):
+        """Test personless mode when callable distinct_id returns None."""
+
+        def resolver(trace):
+            return None  # Simulate no user ID available
+
+        processor = PostHogTracingProcessor(
+            client=mock_client,
+            distinct_id=resolver,
+        )
+
+        processor.on_trace_start(mock_trace)
+        mock_client.capture.reset_mock()
+
+        span_data = GenerationSpanData(model="gpt-4o")
+        mock_span.span_data = span_data
+
+        processor.on_span_start(mock_span)
+        processor.on_span_end(mock_span)
+
+        call_kwargs = mock_client.capture.call_args[1]
+        assert call_kwargs["properties"]["$process_person_profile"] is False
+        assert call_kwargs["distinct_id"] == mock_span.trace_id
+
+    def test_person_profile_when_distinct_id_provided(self, mock_client, mock_trace):
+        """Test that events create person profiles when distinct_id is provided."""
+        processor = PostHogTracingProcessor(
+            client=mock_client,
+            distinct_id="real-user",
+        )
+
+        processor.on_trace_start(mock_trace)
+        processor.on_trace_end(mock_trace)
+
+        call_kwargs = mock_client.capture.call_args[1]
+        assert "$process_person_profile" not in call_kwargs["properties"]
+
+    def test_on_trace_end_clears_metadata(self, processor, mock_client, mock_trace):
+        """Test that on_trace_end clears stored trace metadata."""
+        processor.on_trace_start(mock_trace)
+        assert mock_trace.trace_id in processor._trace_metadata
+
+        processor.on_trace_end(mock_trace)
+        assert mock_trace.trace_id not in processor._trace_metadata
+        # Also verify it captured the event
+        mock_client.capture.assert_called_once()
+
+    def test_on_span_start_tracks_time(self, processor, mock_span):
+        """Test that on_span_start records start time."""
+        processor.on_span_start(mock_span)
+        assert mock_span.span_id in processor._span_start_times
+
+    def test_generation_span_mapping(self, processor, mock_client, mock_span):
+        """Test GenerationSpanData maps to $ai_generation event."""
+        span_data = GenerationSpanData(
+            input=[{"role": "user", "content": "Hello"}],
+            output=[{"role": "assistant", "content": "Hi there!"}],
+            model="gpt-4o",
+            model_config={"temperature": 0.7, "max_tokens": 100},
+            usage={"input_tokens": 10, "output_tokens": 20},
+        )
+        mock_span.span_data = span_data
+
+        processor.on_span_start(mock_span)
+        processor.on_span_end(mock_span)
+
+        mock_client.capture.assert_called_once()
+        call_kwargs = mock_client.capture.call_args[1]
+
+        assert call_kwargs["event"] == "$ai_generation"
+        assert call_kwargs["properties"]["$ai_trace_id"] == "trace_123456789"
+        assert call_kwargs["properties"]["$ai_span_id"] == "span_987654321"
+        assert call_kwargs["properties"]["$ai_provider"] == "openai"
+        assert call_kwargs["properties"]["$ai_framework"] == "openai-agents"
+        assert call_kwargs["properties"]["$ai_model"] == "gpt-4o"
+        assert call_kwargs["properties"]["$ai_input_tokens"] == 10
+        assert call_kwargs["properties"]["$ai_output_tokens"] == 20
+        assert call_kwargs["properties"]["$ai_input"] == [
+            {"role": "user", "content": "Hello"}
+        ]
+        assert call_kwargs["properties"]["$ai_output_choices"] == [
+            {"role": "assistant", "content": "Hi there!"}
+        ]
+
+    def test_generation_span_with_reasoning_tokens(
+        self, processor, mock_client, mock_span
+    ):
+        """Test GenerationSpanData includes reasoning tokens when present."""
+        span_data = GenerationSpanData(
+            model="o1-preview",
+            usage={
+                "input_tokens": 100,
+                "output_tokens": 500,
+                "reasoning_tokens": 400,
+            },
+        )
+        mock_span.span_data = span_data
+
+        processor.on_span_start(mock_span)
+        processor.on_span_end(mock_span)
+
+        call_kwargs = mock_client.capture.call_args[1]
+        assert call_kwargs["properties"]["$ai_reasoning_tokens"] == 400
+
+    def test_function_span_mapping(self, processor, mock_client, mock_span):
+        """Test FunctionSpanData maps to $ai_span event with type=tool."""
+        span_data = FunctionSpanData(
+            name="get_weather",
+            input='{"city": "San Francisco"}',
+            output="Sunny, 72F",
+        )
+        mock_span.span_data = span_data
+
+        processor.on_span_start(mock_span)
+        processor.on_span_end(mock_span)
+
+        call_kwargs = mock_client.capture.call_args[1]
+
+        assert call_kwargs["event"] == "$ai_span"
+        assert call_kwargs["properties"]["$ai_span_name"] == "get_weather"
+        assert call_kwargs["properties"]["$ai_span_type"] == "tool"
+        assert (
+            call_kwargs["properties"]["$ai_input_state"] == '{"city": "San Francisco"}'
+        )
+        assert call_kwargs["properties"]["$ai_output_state"] == "Sunny, 72F"
+
+    def test_agent_span_mapping(self, processor, mock_client, mock_span):
+        """Test AgentSpanData maps to $ai_span event with type=agent."""
+        span_data = AgentSpanData(
+            name="CustomerServiceAgent",
+            handoffs=["TechnicalAgent", "BillingAgent"],
+            tools=["search", "get_order"],
+            output_type="str",
+        )
+        mock_span.span_data = span_data
+
+        processor.on_span_start(mock_span)
+        processor.on_span_end(mock_span)
+
+        call_kwargs = mock_client.capture.call_args[1]
+
+        assert call_kwargs["event"] == "$ai_span"
+        assert call_kwargs["properties"]["$ai_span_name"] == "CustomerServiceAgent"
+        assert call_kwargs["properties"]["$ai_span_type"] == "agent"
+        assert call_kwargs["properties"]["$ai_agent_handoffs"] == [
+            "TechnicalAgent",
+            "BillingAgent",
+        ]
+        assert call_kwargs["properties"]["$ai_agent_tools"] == ["search", "get_order"]
+
+    def test_handoff_span_mapping(self, processor, mock_client, mock_span):
+        """Test HandoffSpanData maps to $ai_span event with type=handoff."""
+        span_data = HandoffSpanData(
+            from_agent="TriageAgent",
+            to_agent="TechnicalAgent",
+        )
+        mock_span.span_data = span_data
+
+        processor.on_span_start(mock_span)
+        processor.on_span_end(mock_span)
+
+        call_kwargs = mock_client.capture.call_args[1]
+
+        assert call_kwargs["event"] == "$ai_span"
+        assert call_kwargs["properties"]["$ai_span_type"] == "handoff"
+        assert call_kwargs["properties"]["$ai_handoff_from_agent"] == "TriageAgent"
+        assert call_kwargs["properties"]["$ai_handoff_to_agent"] == "TechnicalAgent"
+        assert (
+            call_kwargs["properties"]["$ai_span_name"]
+            == "TriageAgent -> TechnicalAgent"
+        )
+
+    def test_guardrail_span_mapping(self, processor, mock_client, mock_span):
+        """Test GuardrailSpanData maps to $ai_span event with type=guardrail."""
+        span_data = GuardrailSpanData(
+            name="ContentFilter",
+            triggered=True,
+        )
+        mock_span.span_data = span_data
+
+        processor.on_span_start(mock_span)
+        processor.on_span_end(mock_span)
+
+        call_kwargs = mock_client.capture.call_args[1]
+
+        assert call_kwargs["event"] == "$ai_span"
+        assert call_kwargs["properties"]["$ai_span_name"] == "ContentFilter"
+        assert call_kwargs["properties"]["$ai_span_type"] == "guardrail"
+        assert call_kwargs["properties"]["$ai_guardrail_triggered"] is True
+
+    def test_custom_span_mapping(self, processor, mock_client, mock_span):
+        """Test CustomSpanData maps to $ai_span event with type=custom."""
+        span_data = CustomSpanData(
+            name="database_query",
+            data={"query": "SELECT * FROM users", "rows": 100},
+        )
+        mock_span.span_data = span_data
+
+        processor.on_span_start(mock_span)
+        processor.on_span_end(mock_span)
+
+        call_kwargs = mock_client.capture.call_args[1]
+
+        assert call_kwargs["event"] == "$ai_span"
+        assert call_kwargs["properties"]["$ai_span_name"] == "database_query"
+        assert call_kwargs["properties"]["$ai_span_type"] == "custom"
+        assert call_kwargs["properties"]["$ai_custom_data"] == {
+            "query": "SELECT * FROM users",
+            "rows": 100,
+        }
+
+    def test_privacy_mode_redacts_content(self, mock_client, mock_span):
+        """Test that privacy_mode redacts input/output content."""
+        processor = PostHogTracingProcessor(
+            client=mock_client,
+            distinct_id="test-user",
+            privacy_mode=True,
+        )
+
+        span_data = GenerationSpanData(
+            input=[{"role": "user", "content": "Secret message"}],
+            output=[{"role": "assistant", "content": "Secret response"}],
+            model="gpt-4o",
+            usage={"input_tokens": 10, "output_tokens": 20},
+        )
+        mock_span.span_data = span_data
+
+        processor.on_span_start(mock_span)
+        processor.on_span_end(mock_span)
+
+        call_kwargs = mock_client.capture.call_args[1]
+
+        # Content should be redacted
+        assert call_kwargs["properties"]["$ai_input"] is None
+        assert call_kwargs["properties"]["$ai_output_choices"] is None
+        # Token counts should still be present
+        assert call_kwargs["properties"]["$ai_input_tokens"] == 10
+        assert call_kwargs["properties"]["$ai_output_tokens"] == 20
+
+    def test_error_handling_in_span(self, processor, mock_client, mock_span):
+        """Test that span errors are captured correctly."""
+        span_data = GenerationSpanData(model="gpt-4o")
+        mock_span.span_data = span_data
+        mock_span.error = {"message": "Rate limit exceeded", "data": {"code": 429}}
+
+        processor.on_span_start(mock_span)
+        processor.on_span_end(mock_span)
+
+        call_kwargs = mock_client.capture.call_args[1]
+
+        assert call_kwargs["properties"]["$ai_is_error"] is True
+        assert call_kwargs["properties"]["$ai_error"] == "Rate limit exceeded"
+
+    def test_generation_span_includes_total_tokens(
+        self, processor, mock_client, mock_span
+    ):
+        """Test that $ai_total_tokens is calculated and included."""
+        span_data = GenerationSpanData(
+            model="gpt-4o",
+            usage={"input_tokens": 100, "output_tokens": 50},
+        )
+        mock_span.span_data = span_data
+
+        processor.on_span_start(mock_span)
+        processor.on_span_end(mock_span)
+
+        call_kwargs = mock_client.capture.call_args[1]
+        assert call_kwargs["properties"]["$ai_total_tokens"] == 150
+
+    def test_error_type_categorization_model_behavior(
+        self, processor, mock_client, mock_span
+    ):
+        """Test that ModelBehaviorError is categorized correctly."""
+        span_data = GenerationSpanData(model="gpt-4o")
+        mock_span.span_data = span_data
+        mock_span.error = {
+            "message": "ModelBehaviorError: Invalid JSON output",
+            "type": "ModelBehaviorError",
+        }
+
+        processor.on_span_start(mock_span)
+        processor.on_span_end(mock_span)
+
+        call_kwargs = mock_client.capture.call_args[1]
+        assert call_kwargs["properties"]["$ai_error_type"] == "model_behavior_error"
+
+    def test_error_type_categorization_user_error(
+        self, processor, mock_client, mock_span
+    ):
+        """Test that UserError is categorized correctly."""
+        span_data = GenerationSpanData(model="gpt-4o")
+        mock_span.span_data = span_data
+        mock_span.error = {"message": "UserError: Tool failed", "type": "UserError"}
+
+        processor.on_span_start(mock_span)
+        processor.on_span_end(mock_span)
+
+        call_kwargs = mock_client.capture.call_args[1]
+        assert call_kwargs["properties"]["$ai_error_type"] == "user_error"
+
+    def test_error_type_categorization_input_guardrail(
+        self, processor, mock_client, mock_span
+    ):
+        """Test that InputGuardrailTripwireTriggered is categorized correctly."""
+        span_data = GenerationSpanData(model="gpt-4o")
+        mock_span.span_data = span_data
+        mock_span.error = {
+            "message": "InputGuardrailTripwireTriggered: Content blocked"
+        }
+
+        processor.on_span_start(mock_span)
+        processor.on_span_end(mock_span)
+
+        call_kwargs = mock_client.capture.call_args[1]
+        assert (
+            call_kwargs["properties"]["$ai_error_type"] == "input_guardrail_triggered"
+        )
+
+    def test_error_type_categorization_output_guardrail(
+        self, processor, mock_client, mock_span
+    ):
+        """Test that OutputGuardrailTripwireTriggered is categorized correctly."""
+        span_data = GenerationSpanData(model="gpt-4o")
+        mock_span.span_data = span_data
+        mock_span.error = {
+            "message": "OutputGuardrailTripwireTriggered: Response blocked"
+        }
+
+        processor.on_span_start(mock_span)
+        processor.on_span_end(mock_span)
+
+        call_kwargs = mock_client.capture.call_args[1]
+        assert (
+            call_kwargs["properties"]["$ai_error_type"] == "output_guardrail_triggered"
+        )
+
+    def test_error_type_categorization_max_turns(
+        self, processor, mock_client, mock_span
+    ):
+        """Test that MaxTurnsExceeded is categorized correctly."""
+        span_data = GenerationSpanData(model="gpt-4o")
+        mock_span.span_data = span_data
+        mock_span.error = {"message": "MaxTurnsExceeded: Agent exceeded maximum turns"}
+
+        processor.on_span_start(mock_span)
+        processor.on_span_end(mock_span)
+
+        call_kwargs = mock_client.capture.call_args[1]
+        assert call_kwargs["properties"]["$ai_error_type"] == "max_turns_exceeded"
+
+    def test_error_type_categorization_unknown(self, processor, mock_client, mock_span):
+        """Test that unknown errors are categorized as unknown."""
+        span_data = GenerationSpanData(model="gpt-4o")
+        mock_span.span_data = span_data
+        mock_span.error = {"message": "Some random error occurred"}
+
+        processor.on_span_start(mock_span)
+        processor.on_span_end(mock_span)
+
+        call_kwargs = mock_client.capture.call_args[1]
+        assert call_kwargs["properties"]["$ai_error_type"] == "unknown"
+
+    def test_response_span_with_output_and_total_tokens(
+        self, processor, mock_client, mock_span
+    ):
+        """Test ResponseSpanData includes output choices and total tokens."""
+        # Create a mock response object
+        mock_response = MagicMock()
+        mock_response.id = "resp_123"
+        mock_response.model = "gpt-4o"
+        mock_response.output = [{"type": "message", "content": "Hello!"}]
+        mock_response.usage = MagicMock()
+        mock_response.usage.input_tokens = 25
+        mock_response.usage.output_tokens = 10
+
+        span_data = ResponseSpanData(
+            response=mock_response,
+            input="Hello, world!",
+        )
+        mock_span.span_data = span_data
+
+        processor.on_span_start(mock_span)
+        processor.on_span_end(mock_span)
+
+        call_kwargs = mock_client.capture.call_args[1]
+
+        assert call_kwargs["event"] == "$ai_generation"
+        assert call_kwargs["properties"]["$ai_total_tokens"] == 35
+        assert call_kwargs["properties"]["$ai_output_choices"] == [
+            {"type": "message", "content": "Hello!"}
+        ]
+        assert call_kwargs["properties"]["$ai_response_id"] == "resp_123"
+
+    def test_speech_span_with_pass_through_properties(
+        self, processor, mock_client, mock_span
+    ):
+        """Test SpeechSpanData includes pass-through properties."""
+        span_data = SpeechSpanData(
+            input="Hello, how can I help you?",
+            output="base64_audio_data",
+            output_format="pcm",
+            model="tts-1",
+            model_config={"voice": "alloy", "speed": 1.0},
+            first_content_at="2024-01-01T00:00:00.500Z",
+        )
+        mock_span.span_data = span_data
+
+        processor.on_span_start(mock_span)
+        processor.on_span_end(mock_span)
+
+        call_kwargs = mock_client.capture.call_args[1]
+
+        assert call_kwargs["event"] == "$ai_span"
+        assert call_kwargs["properties"]["$ai_span_type"] == "speech"
+        assert call_kwargs["properties"]["$ai_model"] == "tts-1"
+        # Pass-through properties (no $ai_ prefix)
+        assert (
+            call_kwargs["properties"]["first_content_at"] == "2024-01-01T00:00:00.500Z"
+        )
+        assert call_kwargs["properties"]["audio_output_format"] == "pcm"
+        assert call_kwargs["properties"]["model_config"] == {
+            "voice": "alloy",
+            "speed": 1.0,
+        }
+        # Text input should be captured
+        assert call_kwargs["properties"]["$ai_input"] == "Hello, how can I help you?"
+
+    def test_transcription_span_with_pass_through_properties(
+        self, processor, mock_client, mock_span
+    ):
+        """Test TranscriptionSpanData includes pass-through properties."""
+        span_data = TranscriptionSpanData(
+            input="base64_audio_data",
+            input_format="pcm",
+            output="This is the transcribed text.",
+            model="whisper-1",
+            model_config={"language": "en"},
+        )
+        mock_span.span_data = span_data
+
+        processor.on_span_start(mock_span)
+        processor.on_span_end(mock_span)
+
+        call_kwargs = mock_client.capture.call_args[1]
+
+        assert call_kwargs["event"] == "$ai_span"
+        assert call_kwargs["properties"]["$ai_span_type"] == "transcription"
+        assert call_kwargs["properties"]["$ai_model"] == "whisper-1"
+        # Pass-through properties (no $ai_ prefix)
+        assert call_kwargs["properties"]["audio_input_format"] == "pcm"
+        assert call_kwargs["properties"]["model_config"] == {"language": "en"}
+        # Transcription output should be captured
+        assert (
+            call_kwargs["properties"]["$ai_output_state"]
+            == "This is the transcribed text."
+        )
+
+    def test_latency_calculation(self, processor, mock_client, mock_span):
+        """Test that latency is calculated correctly."""
+        span_data = GenerationSpanData(model="gpt-4o")
+        mock_span.span_data = span_data
+
+        with patch("time.time") as mock_time:
+            mock_time.return_value = 1000.0
+            processor.on_span_start(mock_span)
+
+            mock_time.return_value = 1001.5  # 1.5 seconds later
+            processor.on_span_end(mock_span)
+
+        call_kwargs = mock_client.capture.call_args[1]
+        assert call_kwargs["properties"]["$ai_latency"] == pytest.approx(1.5, rel=0.01)
+
+    def test_groups_included_in_events(self, mock_client, mock_trace, mock_span):
+        """Test that groups are included in captured events."""
+        processor = PostHogTracingProcessor(
+            client=mock_client,
+            distinct_id="test-user",
+            groups={"company": "acme", "team": "engineering"},
+        )
+
+        processor.on_trace_start(mock_trace)
+        processor.on_trace_end(mock_trace)
+
+        call_kwargs = mock_client.capture.call_args[1]
+        assert call_kwargs["groups"] == {"company": "acme", "team": "engineering"}
+
+    def test_additional_properties_included(self, mock_client, mock_trace):
+        """Test that additional properties are included in events."""
+        processor = PostHogTracingProcessor(
+            client=mock_client,
+            distinct_id="test-user",
+            properties={"environment": "production", "version": "1.0"},
+        )
+
+        processor.on_trace_start(mock_trace)
+        processor.on_trace_end(mock_trace)
+
+        call_kwargs = mock_client.capture.call_args[1]
+        assert call_kwargs["properties"]["environment"] == "production"
+        assert call_kwargs["properties"]["version"] == "1.0"
+
+    def test_shutdown_clears_state(self, processor):
+        """Test that shutdown clears internal state."""
+        processor._span_start_times["span_1"] = 1000.0
+        processor._trace_metadata["trace_1"] = {"name": "test"}
+
+        processor.shutdown()
+
+        assert len(processor._span_start_times) == 0
+        assert len(processor._trace_metadata) == 0
+
+    def test_force_flush_calls_client_flush(self, processor, mock_client):
+        """Test that force_flush calls client.flush()."""
+        processor.force_flush()
+        mock_client.flush.assert_called_once()
+
+    def test_generation_span_with_no_usage(self, processor, mock_client, mock_span):
+        """Test GenerationSpanData with no usage data defaults to zero tokens."""
+        span_data = GenerationSpanData(model="gpt-4o")
+        mock_span.span_data = span_data
+
+        processor.on_span_start(mock_span)
+        processor.on_span_end(mock_span)
+
+        call_kwargs = mock_client.capture.call_args[1]
+        assert call_kwargs["properties"]["$ai_input_tokens"] == 0
+        assert call_kwargs["properties"]["$ai_output_tokens"] == 0
+        assert call_kwargs["properties"]["$ai_total_tokens"] == 0
+
+    def test_generation_span_with_partial_usage(
+        self, processor, mock_client, mock_span
+    ):
+        """Test GenerationSpanData with only input_tokens present."""
+        span_data = GenerationSpanData(
+            model="gpt-4o",
+            usage={"input_tokens": 42},
+        )
+        mock_span.span_data = span_data
+
+        processor.on_span_start(mock_span)
+        processor.on_span_end(mock_span)
+
+        call_kwargs = mock_client.capture.call_args[1]
+        assert call_kwargs["properties"]["$ai_input_tokens"] == 42
+        assert call_kwargs["properties"]["$ai_output_tokens"] == 0
+        assert call_kwargs["properties"]["$ai_total_tokens"] == 42
+
+    def test_error_type_categorization_by_type_field_only(
+        self, processor, mock_client, mock_span
+    ):
+        """Test error categorization works when only the type field matches."""
+        span_data = GenerationSpanData(model="gpt-4o")
+        mock_span.span_data = span_data
+        mock_span.error = {
+            "message": "Something went wrong",
+            "type": "ModelBehaviorError",
+        }
+
+        processor.on_span_start(mock_span)
+        processor.on_span_end(mock_span)
+
+        call_kwargs = mock_client.capture.call_args[1]
+        assert call_kwargs["properties"]["$ai_error_type"] == "model_behavior_error"
+
+    def test_distinct_id_resolved_from_trace_for_spans(
+        self, mock_client, mock_trace, mock_span
+    ):
+        """Test that spans use the distinct_id resolved at trace start."""
+
+        def resolver(trace):
+            return f"user-{trace.name}"
+
+        processor = PostHogTracingProcessor(
+            client=mock_client,
+            distinct_id=resolver,
+        )
+
+        # Start trace - this resolves and stores distinct_id
+        processor.on_trace_start(mock_trace)
+        mock_client.capture.reset_mock()
+
+        # End a span - should use the stored distinct_id from trace
+        span_data = GenerationSpanData(model="gpt-4o")
+        mock_span.span_data = span_data
+
+        processor.on_span_start(mock_span)
+        processor.on_span_end(mock_span)
+
+        call_kwargs = mock_client.capture.call_args[1]
+        assert call_kwargs["distinct_id"] == "user-Test Workflow"
+
+    def test_eviction_of_stale_entries(self, mock_client):
+        """Test that stale entries are evicted when max is exceeded."""
+        processor = PostHogTracingProcessor(
+            client=mock_client,
+            distinct_id="test-user",
+        )
+        processor._max_tracked_entries = 10
+
+        # Fill beyond max
+        for i in range(15):
+            processor._span_start_times[f"span_{i}"] = float(i)
+            processor._trace_metadata[f"trace_{i}"] = {"name": f"trace_{i}"}
+
+        processor._evict_stale_entries()
+
+        # Should have evicted half
+        assert len(processor._span_start_times) <= 10
+        assert len(processor._trace_metadata) <= 10
+
+
+class TestInstrumentHelper:
+    """Tests for the instrument() convenience function."""
+
+    def test_instrument_registers_processor(self, mock_client):
+        """Test that instrument() registers a processor."""
+        with patch("agents.tracing.add_trace_processor") as mock_add:
+            processor = instrument(
+                client=mock_client,
+                distinct_id="test-user",
+            )
+
+            mock_add.assert_called_once_with(processor)
+            assert isinstance(processor, PostHogTracingProcessor)
+
+    def test_instrument_with_privacy_mode(self, mock_client):
+        """Test instrument() respects privacy_mode."""
+        with patch("agents.tracing.add_trace_processor"):
+            processor = instrument(
+                client=mock_client,
+                privacy_mode=True,
+            )
+
+            assert processor._privacy_mode is True
+
+    def test_instrument_with_groups_and_properties(self, mock_client):
+        """Test instrument() accepts groups and properties."""
+        with patch("agents.tracing.add_trace_processor"):
+            processor = instrument(
+                client=mock_client,
+                groups={"company": "acme"},
+                properties={"env": "test"},
+            )
+
+            assert processor._groups == {"company": "acme"}
+            assert processor._properties == {"env": "test"}

posthog/test/ai/test_prompts.py

@@ -0,0 +1,607 @@
+import unittest
+from unittest.mock import MagicMock, patch
+
+from posthog.ai.prompts import Prompts
+
+
+class MockResponse:
+    """Mock HTTP response for testing."""
+
+    def __init__(self, json_data=None, status_code=200, ok=True):
+        self._json_data = json_data
+        self.status_code = status_code
+        self.ok = ok
+
+    def json(self):
+        if self._json_data is None:
+            raise ValueError("No JSON data")
+        return self._json_data
+
+
+class TestPrompts(unittest.TestCase):
+    """Tests for the Prompts class."""
+
+    mock_prompt_response = {
+        "id": 1,
+        "name": "test-prompt",
+        "prompt": "Hello, {{name}}! You are a helpful assistant for {{company}}.",
+        "version": 1,
+        "created_by": "user@example.com",
+        "created_at": "2024-01-01T00:00:00Z",
+        "updated_at": "2024-01-01T00:00:00Z",
+        "deleted": False,
+    }
+
+    def create_mock_posthog(
+        self,
+        personal_api_key="phx_test_key",
+        project_api_key="phc_test_key",
+        host="https://us.posthog.com",
+    ):
+        """Create a mock PostHog client."""
+        mock = MagicMock()
+        mock.personal_api_key = personal_api_key
+        mock.api_key = project_api_key
+        mock.raw_host = host
+        return mock
+
+
+class TestPromptsGet(TestPrompts):
+    """Tests for the Prompts.get() method."""
+
+    @patch("posthog.ai.prompts._get_session")
+    def test_successfully_fetch_a_prompt(self, mock_get_session):
+        """Should successfully fetch a prompt."""
+        mock_get = mock_get_session.return_value.get
+        mock_get.return_value = MockResponse(json_data=self.mock_prompt_response)
+
+        posthog = self.create_mock_posthog()
+        prompts = Prompts(posthog)
+
+        result = prompts.get("test-prompt")
+
+        self.assertEqual(result, self.mock_prompt_response["prompt"])
+        mock_get.assert_called_once()
+        call_args = mock_get.call_args
+        self.assertEqual(
+            call_args[0][0],
+            "https://us.posthog.com/api/environments/@current/llm_prompts/name/test-prompt/?token=phc_test_key",
+        )
+        self.assertIn("Authorization", call_args[1]["headers"])
+        self.assertEqual(
+            call_args[1]["headers"]["Authorization"], "Bearer phx_test_key"
+        )
+
+    @patch("posthog.ai.prompts._get_session")
+    @patch("posthog.ai.prompts.time.time")
+    def test_return_cached_prompt_when_fresh(self, mock_time, mock_get_session):
+        """Should return cached prompt when fresh (no API call)."""
+        mock_get = mock_get_session.return_value.get
+        mock_get.return_value = MockResponse(json_data=self.mock_prompt_response)
+        mock_time.return_value = 1000.0
+
+        posthog = self.create_mock_posthog()
+        prompts = Prompts(posthog)
+
+        # First call - fetches from API
+        result1 = prompts.get("test-prompt", cache_ttl_seconds=300)
+        self.assertEqual(result1, self.mock_prompt_response["prompt"])
+        self.assertEqual(mock_get.call_count, 1)
+
+        # Advance time by 60 seconds (still within TTL)
+        mock_time.return_value = 1060.0
+
+        # Second call - should use cache
+        result2 = prompts.get("test-prompt", cache_ttl_seconds=300)
+        self.assertEqual(result2, self.mock_prompt_response["prompt"])
+        self.assertEqual(mock_get.call_count, 1)  # No additional fetch
+
+    @patch("posthog.ai.prompts._get_session")
+    @patch("posthog.ai.prompts.time.time")
+    def test_refetch_when_cache_is_stale(self, mock_time, mock_get_session):
+        """Should refetch when cache is stale."""
+        mock_get = mock_get_session.return_value.get
+        updated_prompt_response = {
+            **self.mock_prompt_response,
+            "prompt": "Updated prompt: Hello, {{name}}!",
+        }
+
+        mock_get.side_effect = [
+            MockResponse(json_data=self.mock_prompt_response),
+            MockResponse(json_data=updated_prompt_response),
+        ]
+        mock_time.return_value = 1000.0
+
+        posthog = self.create_mock_posthog()
+        prompts = Prompts(posthog)
+
+        # First call - fetches from API
+        result1 = prompts.get("test-prompt", cache_ttl_seconds=60)
+        self.assertEqual(result1, self.mock_prompt_response["prompt"])
+        self.assertEqual(mock_get.call_count, 1)
+
+        # Advance time past TTL
+        mock_time.return_value = 1061.0
+
+        # Second call - should refetch
+        result2 = prompts.get("test-prompt", cache_ttl_seconds=60)
+        self.assertEqual(result2, updated_prompt_response["prompt"])
+        self.assertEqual(mock_get.call_count, 2)
+
+    @patch("posthog.ai.prompts._get_session")
+    @patch("posthog.ai.prompts.time.time")
+    @patch("posthog.ai.prompts.log")
+    def test_use_stale_cache_on_fetch_failure_with_warning(
+        self, mock_log, mock_time, mock_get_session
+    ):
+        """Should use stale cache on fetch failure with warning."""
+        mock_get = mock_get_session.return_value.get
+        mock_get.side_effect = [
+            MockResponse(json_data=self.mock_prompt_response),
+            Exception("Network error"),
+        ]
+        mock_time.return_value = 1000.0
+
+        posthog = self.create_mock_posthog()
+        prompts = Prompts(posthog)
+
+        # First call - populates cache
+        result1 = prompts.get("test-prompt", cache_ttl_seconds=60)
+        self.assertEqual(result1, self.mock_prompt_response["prompt"])
+
+        # Advance time past TTL
+        mock_time.return_value = 1061.0
+
+        # Second call - should use stale cache
+        result2 = prompts.get("test-prompt", cache_ttl_seconds=60)
+        self.assertEqual(result2, self.mock_prompt_response["prompt"])
+
+        # Check warning was logged
+        mock_log.warning.assert_called()
+        warning_call = mock_log.warning.call_args
+        self.assertIn("using stale cache", warning_call[0][0])
+
+    @patch("posthog.ai.prompts._get_session")
+    @patch("posthog.ai.prompts.log")
+    def test_use_fallback_when_no_cache_and_fetch_fails_with_warning(
+        self, mock_log, mock_get_session
+    ):
+        """Should use fallback when no cache and fetch fails with warning."""
+        mock_get = mock_get_session.return_value.get
+        mock_get.side_effect = Exception("Network error")
+
+        posthog = self.create_mock_posthog()
+        prompts = Prompts(posthog)
+
+        fallback = "Default system prompt."
+        result = prompts.get("test-prompt", fallback=fallback)
+
+        self.assertEqual(result, fallback)
+
+        # Check warning was logged
+        mock_log.warning.assert_called()
+        warning_call = mock_log.warning.call_args
+        self.assertIn("using fallback", warning_call[0][0])
+
+    @patch("posthog.ai.prompts._get_session")
+    def test_throw_when_no_cache_no_fallback_and_fetch_fails(self, mock_get_session):
+        """Should throw when no cache, no fallback, and fetch fails."""
+        mock_get = mock_get_session.return_value.get
+        mock_get.side_effect = Exception("Network error")
+
+        posthog = self.create_mock_posthog()
+        prompts = Prompts(posthog)
+
+        with self.assertRaises(Exception) as context:
+            prompts.get("test-prompt")
+
+        self.assertIn("Network error", str(context.exception))
+
+    @patch("posthog.ai.prompts._get_session")
+    def test_handle_404_response(self, mock_get_session):
+        """Should handle 404 response."""
+        mock_get = mock_get_session.return_value.get
+        mock_get.return_value = MockResponse(status_code=404, ok=False)
+
+        posthog = self.create_mock_posthog()
+        prompts = Prompts(posthog)
+
+        with self.assertRaises(Exception) as context:
+            prompts.get("nonexistent-prompt")
+
+        self.assertIn('Prompt "nonexistent-prompt" not found', str(context.exception))
+
+    @patch("posthog.ai.prompts._get_session")
+    def test_handle_403_response(self, mock_get_session):
+        """Should handle 403 response."""
+        mock_get = mock_get_session.return_value.get
+        mock_get.return_value = MockResponse(status_code=403, ok=False)
+
+        posthog = self.create_mock_posthog()
+        prompts = Prompts(posthog)
+
+        with self.assertRaises(Exception) as context:
+            prompts.get("restricted-prompt")
+
+        self.assertIn(
+            'Access denied for prompt "restricted-prompt"', str(context.exception)
+        )
+
+    def test_throw_when_no_personal_api_key_configured(self):
+        """Should throw when no personal_api_key is configured."""
+        posthog = self.create_mock_posthog(personal_api_key=None)
+        prompts = Prompts(posthog)
+
+        with self.assertRaises(Exception) as context:
+            prompts.get("test-prompt")
+
+        self.assertIn(
+            "personal_api_key is required to fetch prompts", str(context.exception)
+        )
+
+    def test_throw_when_no_project_api_key_configured(self):
+        """Should throw when no project_api_key is configured."""
+        posthog = self.create_mock_posthog(project_api_key=None)
+        prompts = Prompts(posthog)
+
+        with self.assertRaises(Exception) as context:
+            prompts.get("test-prompt")
+
+        self.assertIn(
+            "project_api_key is required to fetch prompts", str(context.exception)
+        )
+
+    @patch("posthog.ai.prompts._get_session")
+    def test_throw_when_api_returns_invalid_response_format(self, mock_get_session):
+        """Should throw when API returns invalid response format."""
+        mock_get = mock_get_session.return_value.get
+        mock_get.return_value = MockResponse(json_data={"invalid": "response"})
+
+        posthog = self.create_mock_posthog()
+        prompts = Prompts(posthog)
+
+        with self.assertRaises(Exception) as context:
+            prompts.get("test-prompt")
+
+        self.assertIn("Invalid response format", str(context.exception))
+
+    @patch("posthog.ai.prompts._get_session")
+    def test_use_custom_host_from_posthog_options(self, mock_get_session):
+        """Should use custom host from PostHog options."""
+        mock_get = mock_get_session.return_value.get
+        mock_get.return_value = MockResponse(json_data=self.mock_prompt_response)
+
+        posthog = self.create_mock_posthog(host="https://eu.posthog.com")
+        prompts = Prompts(posthog)
+
+        prompts.get("test-prompt")
+
+        call_args = mock_get.call_args
+        self.assertTrue(
+            call_args[0][0].startswith(
+                "https://eu.posthog.com/api/environments/@current/llm_prompts/name/test-prompt/?token=phc_test_key"
+            ),
+            f"Expected URL to start with 'https://eu.posthog.com/api/environments/@current/llm_prompts/name/test-prompt/?token=phc_test_key', got {call_args[0][0]}",
+        )
+
+    @patch("posthog.ai.prompts._get_session")
+    @patch("posthog.ai.prompts.time.time")
+    def test_use_default_cache_ttl_5_minutes(self, mock_time, mock_get_session):
+        """Should use default cache TTL (5 minutes) when not specified."""
+        mock_get = mock_get_session.return_value.get
+        mock_get.return_value = MockResponse(json_data=self.mock_prompt_response)
+        mock_time.return_value = 1000.0
+
+        posthog = self.create_mock_posthog()
+        prompts = Prompts(posthog)
+
+        # First call
+        prompts.get("test-prompt")
+        self.assertEqual(mock_get.call_count, 1)
+
+        # Advance time by 4 minutes (within default 5-minute TTL)
+        mock_time.return_value = 1000.0 + (4 * 60)
+
+        # Second call - should use cache
+        prompts.get("test-prompt")
+        self.assertEqual(mock_get.call_count, 1)
+
+        # Advance time past 5-minute TTL
+        mock_time.return_value = 1000.0 + (6 * 60)
+
+        # Third call - should refetch
+        prompts.get("test-prompt")
+        self.assertEqual(mock_get.call_count, 2)
+
+    @patch("posthog.ai.prompts._get_session")
+    @patch("posthog.ai.prompts.time.time")
+    def test_use_custom_default_cache_ttl_from_constructor(
+        self, mock_time, mock_get_session
+    ):
+        """Should use custom default cache TTL from constructor."""
+        mock_get = mock_get_session.return_value.get
+        mock_get.return_value = MockResponse(json_data=self.mock_prompt_response)
+        mock_time.return_value = 1000.0
+
+        posthog = self.create_mock_posthog()
+        prompts = Prompts(posthog, default_cache_ttl_seconds=60)
+
+        # First call
+        prompts.get("test-prompt")
+        self.assertEqual(mock_get.call_count, 1)
+
+        # Advance time past custom TTL
+        mock_time.return_value = 1061.0
+
+        # Second call - should refetch
+        prompts.get("test-prompt")
+        self.assertEqual(mock_get.call_count, 2)
+
+    @patch("posthog.ai.prompts._get_session")
+    def test_url_encode_prompt_names_with_special_characters(self, mock_get_session):
+        """Should URL-encode prompt names with special characters."""
+        mock_get = mock_get_session.return_value.get
+        mock_get.return_value = MockResponse(json_data=self.mock_prompt_response)
+
+        posthog = self.create_mock_posthog()
+        prompts = Prompts(posthog)
+
+        prompts.get("prompt with spaces/and/slashes")
+
+        call_args = mock_get.call_args
+        self.assertEqual(
+            call_args[0][0],
+            "https://us.posthog.com/api/environments/@current/llm_prompts/name/prompt%20with%20spaces%2Fand%2Fslashes/?token=phc_test_key",
+        )
+
+    @patch("posthog.ai.prompts._get_session")
+    def test_work_with_direct_options_no_posthog_client(self, mock_get_session):
+        """Should work with direct options (no PostHog client)."""
+        mock_get = mock_get_session.return_value.get
+        mock_get.return_value = MockResponse(json_data=self.mock_prompt_response)
+
+        prompts = Prompts(
+            personal_api_key="phx_direct_key", project_api_key="phc_direct_key"
+        )
+
+        result = prompts.get("test-prompt")
+
+        self.assertEqual(result, self.mock_prompt_response["prompt"])
+        call_args = mock_get.call_args
+        self.assertEqual(
+            call_args[0][0],
+            "https://us.posthog.com/api/environments/@current/llm_prompts/name/test-prompt/?token=phc_direct_key",
+        )
+        self.assertEqual(
+            call_args[1]["headers"]["Authorization"], "Bearer phx_direct_key"
+        )
+
+    @patch("posthog.ai.prompts._get_session")
+    def test_use_custom_host_from_direct_options(self, mock_get_session):
+        """Should use custom host from direct options."""
+        mock_get = mock_get_session.return_value.get
+        mock_get.return_value = MockResponse(json_data=self.mock_prompt_response)
+
+        prompts = Prompts(
+            personal_api_key="phx_direct_key",
+            project_api_key="phc_direct_key",
+            host="https://eu.posthog.com",
+        )
+
+        prompts.get("test-prompt")
+
+        call_args = mock_get.call_args
+        self.assertEqual(
+            call_args[0][0],
+            "https://eu.posthog.com/api/environments/@current/llm_prompts/name/test-prompt/?token=phc_direct_key",
+        )
+
+    @patch("posthog.ai.prompts._get_session")
+    @patch("posthog.ai.prompts.time.time")
+    def test_use_custom_default_cache_ttl_from_direct_options(
+        self, mock_time, mock_get_session
+    ):
+        """Should use custom default cache TTL from direct options."""
+        mock_get = mock_get_session.return_value.get
+        mock_get.return_value = MockResponse(json_data=self.mock_prompt_response)
+        mock_time.return_value = 1000.0
+
+        prompts = Prompts(
+            personal_api_key="phx_direct_key",
+            project_api_key="phc_direct_key",
+            default_cache_ttl_seconds=60,
+        )
+
+        # First call
+        prompts.get("test-prompt")
+        self.assertEqual(mock_get.call_count, 1)
+
+        # Advance time past custom TTL
+        mock_time.return_value = 1061.0
+
+        # Second call - should refetch
+        prompts.get("test-prompt")
+        self.assertEqual(mock_get.call_count, 2)
+
+
+class TestPromptsCompile(TestPrompts):
+    """Tests for the Prompts.compile() method."""
+
+    def test_replace_a_single_variable(self):
+        """Should replace a single variable."""
+        posthog = self.create_mock_posthog()
+        prompts = Prompts(posthog)
+
+        result = prompts.compile("Hello, {{name}}!", {"name": "World"})
+
+        self.assertEqual(result, "Hello, World!")
+
+    def test_replace_multiple_variables(self):
+        """Should replace multiple variables."""
+        posthog = self.create_mock_posthog()
+        prompts = Prompts(posthog)
+
+        result = prompts.compile(
+            "Hello, {{name}}! Welcome to {{company}}. Your tier is {{tier}}.",
+            {"name": "John", "company": "Acme Corp", "tier": "premium"},
+        )
+
+        self.assertEqual(
+            result, "Hello, John! Welcome to Acme Corp. Your tier is premium."
+        )
+
+    def test_handle_numbers(self):
+        """Should handle numbers."""
+        posthog = self.create_mock_posthog()
+        prompts = Prompts(posthog)
+
+        result = prompts.compile("You have {{count}} items.", {"count": 42})
+
+        self.assertEqual(result, "You have 42 items.")
+
+    def test_handle_booleans(self):
+        """Should handle booleans."""
+        posthog = self.create_mock_posthog()
+        prompts = Prompts(posthog)
+
+        result = prompts.compile("Feature enabled: {{enabled}}", {"enabled": True})
+
+        self.assertEqual(result, "Feature enabled: True")
+
+    def test_leave_unmatched_variables_unchanged(self):
+        """Should leave unmatched variables unchanged."""
+        posthog = self.create_mock_posthog()
+        prompts = Prompts(posthog)
+
+        result = prompts.compile(
+            "Hello, {{name}}! Your {{unknown}} is ready.", {"name": "World"}
+        )
+
+        self.assertEqual(result, "Hello, World! Your {{unknown}} is ready.")
+
+    def test_handle_prompts_with_no_variables(self):
+        """Should handle prompts with no variables."""
+        posthog = self.create_mock_posthog()
+        prompts = Prompts(posthog)
+
+        result = prompts.compile("You are a helpful assistant.", {})
+
+        self.assertEqual(result, "You are a helpful assistant.")
+
+    def test_handle_empty_variables_dict(self):
+        """Should handle empty variables dict."""
+        posthog = self.create_mock_posthog()
+        prompts = Prompts(posthog)
+
+        result = prompts.compile("Hello, {{name}}!", {})
+
+        self.assertEqual(result, "Hello, {{name}}!")
+
+    def test_handle_multiple_occurrences_of_same_variable(self):
+        """Should handle multiple occurrences of the same variable."""
+        posthog = self.create_mock_posthog()
+        prompts = Prompts(posthog)
+
+        result = prompts.compile(
+            "Hello, {{name}}! Goodbye, {{name}}!", {"name": "World"}
+        )
+
+        self.assertEqual(result, "Hello, World! Goodbye, World!")
+
+    def test_work_with_direct_options_initialization(self):
+        """Should work with direct options initialization."""
+        prompts = Prompts(
+            personal_api_key="phx_test_key", project_api_key="phc_test_key"
+        )
+
+        result = prompts.compile("Hello, {{name}}!", {"name": "World"})
+
+        self.assertEqual(result, "Hello, World!")
+
+    def test_handle_variables_with_hyphens(self):
+        """Should handle variables with hyphens."""
+        prompts = Prompts(
+            personal_api_key="phx_test_key", project_api_key="phc_test_key"
+        )
+
+        result = prompts.compile("User ID: {{user-id}}", {"user-id": "12345"})
+
+        self.assertEqual(result, "User ID: 12345")
+
+    def test_handle_variables_with_dots(self):
+        """Should handle variables with dots."""
+        prompts = Prompts(
+            personal_api_key="phx_test_key", project_api_key="phc_test_key"
+        )
+
+        result = prompts.compile("Company: {{company.name}}", {"company.name": "Acme"})
+
+        self.assertEqual(result, "Company: Acme")
+
+
+class TestPromptsClearCache(TestPrompts):
+    """Tests for the Prompts.clear_cache() method."""
+
+    @patch("posthog.ai.prompts._get_session")
+    def test_clear_a_specific_prompt_from_cache(self, mock_get_session):
+        """Should clear a specific prompt from cache."""
+        mock_get = mock_get_session.return_value.get
+        other_prompt_response = {**self.mock_prompt_response, "name": "other-prompt"}
+
+        mock_get.side_effect = [
+            MockResponse(json_data=self.mock_prompt_response),
+            MockResponse(json_data=other_prompt_response),
+            MockResponse(json_data=self.mock_prompt_response),
+        ]
+
+        posthog = self.create_mock_posthog()
+        prompts = Prompts(posthog)
+
+        # Populate cache with two prompts
+        prompts.get("test-prompt")
+        prompts.get("other-prompt")
+        self.assertEqual(mock_get.call_count, 2)
+
+        # Clear only test-prompt
+        prompts.clear_cache("test-prompt")
+
+        # test-prompt should be refetched
+        prompts.get("test-prompt")
+        self.assertEqual(mock_get.call_count, 3)
+
+        # other-prompt should still be cached
+        prompts.get("other-prompt")
+        self.assertEqual(mock_get.call_count, 3)
+
+    @patch("posthog.ai.prompts._get_session")
+    def test_clear_all_prompts_from_cache(self, mock_get_session):
+        """Should clear all prompts from cache when no name is provided."""
+        mock_get = mock_get_session.return_value.get
+        other_prompt_response = {**self.mock_prompt_response, "name": "other-prompt"}
+
+        mock_get.side_effect = [
+            MockResponse(json_data=self.mock_prompt_response),
+            MockResponse(json_data=other_prompt_response),
+            MockResponse(json_data=self.mock_prompt_response),
+            MockResponse(json_data=other_prompt_response),
+        ]
+
+        posthog = self.create_mock_posthog()
+        prompts = Prompts(posthog)
+
+        # Populate cache with two prompts
+        prompts.get("test-prompt")
+        prompts.get("other-prompt")
+        self.assertEqual(mock_get.call_count, 2)
+
+        # Clear all cache
+        prompts.clear_cache()
+
+        # Both prompts should be refetched
+        prompts.get("test-prompt")
+        prompts.get("other-prompt")
+        self.assertEqual(mock_get.call_count, 4)
+
+
+if __name__ == "__main__":
+    unittest.main()

posthog/test/ai/test_sanitization.py

@@ -1,3 +1,4 @@
+import os
 import unittest
 
 from posthog.ai.sanitization import (
@@ -68,6 +69,25 @@ class TestSanitization(unittest.TestCase):
         )
         self.assertEqual(result[0]["content"][1]["image_url"]["detail"], "high")
 
+    def test_sanitize_openai_input_image(self):
+        input_data = [
+            {
+                "role": "user",
+                "content": [
+                    {
+                        "type": "input_image",
+                        "image_url": self.sample_base64_image,
+                    }
+                ],
+            }
+        ]
+
+        result = sanitize_openai(input_data)
+
+        self.assertEqual(
+            result[0]["content"][0]["image_url"], REDACTED_IMAGE_PLACEHOLDER
+        )
+
     def test_sanitize_openai_preserves_regular_urls(self):
         input_data = [
             {
@@ -331,5 +351,191 @@ class TestSanitization(unittest.TestCase):
         )
 
 
+class TestAIMultipartRequest(unittest.TestCase):
+    """Test that _INTERNAL_LLMA_MULTIMODAL environment variable controls sanitization."""
+
+    def tearDown(self):
+        # Clean up environment variable after each test
+        if "_INTERNAL_LLMA_MULTIMODAL" in os.environ:
+            del os.environ["_INTERNAL_LLMA_MULTIMODAL"]
+
+    def test_multimodal_disabled_redacts_images(self):
+        """When _INTERNAL_LLMA_MULTIMODAL is not set, images should be redacted."""
+        if "_INTERNAL_LLMA_MULTIMODAL" in os.environ:
+            del os.environ["_INTERNAL_LLMA_MULTIMODAL"]
+
+        base64_image = "data:image/jpeg;base64,/9j/4AAQSkZJRgABAQAAAQ..."
+        result = redact_base64_data_url(base64_image)
+        self.assertEqual(result, REDACTED_IMAGE_PLACEHOLDER)
+
+    def test_multimodal_enabled_preserves_images(self):
+        """When _INTERNAL_LLMA_MULTIMODAL is true, images should be preserved."""
+        os.environ["_INTERNAL_LLMA_MULTIMODAL"] = "true"
+
+        base64_image = "data:image/jpeg;base64,/9j/4AAQSkZJRgABAQAAAQ..."
+        result = redact_base64_data_url(base64_image)
+        self.assertEqual(result, base64_image)
+
+    def test_multimodal_enabled_with_1(self):
+        """_INTERNAL_LLMA_MULTIMODAL=1 should enable multimodal."""
+        os.environ["_INTERNAL_LLMA_MULTIMODAL"] = "1"
+
+        base64_image = "data:image/jpeg;base64,/9j/4AAQSkZJRgABAQAAAQ..."
+        result = redact_base64_data_url(base64_image)
+        self.assertEqual(result, base64_image)
+
+    def test_multimodal_enabled_with_yes(self):
+        """_INTERNAL_LLMA_MULTIMODAL=yes should enable multimodal."""
+        os.environ["_INTERNAL_LLMA_MULTIMODAL"] = "yes"
+
+        base64_image = "data:image/jpeg;base64,/9j/4AAQSkZJRgABAQAAAQ..."
+        result = redact_base64_data_url(base64_image)
+        self.assertEqual(result, base64_image)
+
+    def test_multimodal_false_redacts_images(self):
+        """_INTERNAL_LLMA_MULTIMODAL=false should still redact."""
+        os.environ["_INTERNAL_LLMA_MULTIMODAL"] = "false"
+
+        base64_image = "data:image/jpeg;base64,/9j/4AAQSkZJRgABAQAAAQ..."
+        result = redact_base64_data_url(base64_image)
+        self.assertEqual(result, REDACTED_IMAGE_PLACEHOLDER)
+
+    def test_anthropic_multimodal_enabled(self):
+        """Anthropic images should be preserved when _INTERNAL_LLMA_MULTIMODAL is enabled."""
+        os.environ["_INTERNAL_LLMA_MULTIMODAL"] = "true"
+
+        input_data = [
+            {
+                "role": "user",
+                "content": [
+                    {
+                        "type": "image",
+                        "source": {
+                            "type": "base64",
+                            "media_type": "image/jpeg",
+                            "data": "base64data",
+                        },
+                    }
+                ],
+            }
+        ]
+
+        result = sanitize_anthropic(input_data)
+        self.assertEqual(result[0]["content"][0]["source"]["data"], "base64data")
+
+    def test_gemini_multimodal_enabled(self):
+        """Gemini images should be preserved when _INTERNAL_LLMA_MULTIMODAL is enabled."""
+        os.environ["_INTERNAL_LLMA_MULTIMODAL"] = "true"
+
+        input_data = [
+            {
+                "parts": [
+                    {"inline_data": {"mime_type": "image/jpeg", "data": "base64data"}}
+                ]
+            }
+        ]
+
+        result = sanitize_gemini(input_data)
+        self.assertEqual(result[0]["parts"][0]["inline_data"]["data"], "base64data")
+
+    def test_langchain_anthropic_style_multimodal_enabled(self):
+        """LangChain Anthropic-style images should be preserved when _INTERNAL_LLMA_MULTIMODAL is enabled."""
+        os.environ["_INTERNAL_LLMA_MULTIMODAL"] = "true"
+
+        input_data = [
+            {
+                "role": "user",
+                "content": [
+                    {
+                        "type": "image",
+                        "source": {"data": "base64data"},
+                    }
+                ],
+            }
+        ]
+
+        result = sanitize_langchain(input_data)
+        self.assertEqual(result[0]["content"][0]["source"]["data"], "base64data")
+
+    def test_openai_audio_redacted_by_default(self):
+        """OpenAI audio should be redacted when _INTERNAL_LLMA_MULTIMODAL is not set."""
+        if "_INTERNAL_LLMA_MULTIMODAL" in os.environ:
+            del os.environ["_INTERNAL_LLMA_MULTIMODAL"]
+
+        input_data = [
+            {
+                "role": "assistant",
+                "content": [
+                    {"type": "audio", "data": "base64audiodata", "id": "audio_123"}
+                ],
+            }
+        ]
+
+        result = sanitize_openai(input_data)
+        self.assertEqual(result[0]["content"][0]["data"], REDACTED_IMAGE_PLACEHOLDER)
+        self.assertEqual(result[0]["content"][0]["id"], "audio_123")
+
+    def test_openai_audio_preserved_with_flag(self):
+        """OpenAI audio should be preserved when _INTERNAL_LLMA_MULTIMODAL is enabled."""
+        os.environ["_INTERNAL_LLMA_MULTIMODAL"] = "true"
+
+        input_data = [
+            {
+                "role": "assistant",
+                "content": [
+                    {"type": "audio", "data": "base64audiodata", "id": "audio_123"}
+                ],
+            }
+        ]
+
+        result = sanitize_openai(input_data)
+        self.assertEqual(result[0]["content"][0]["data"], "base64audiodata")
+
+    def test_gemini_audio_redacted_by_default(self):
+        """Gemini audio should be redacted when _INTERNAL_LLMA_MULTIMODAL is not set."""
+        if "_INTERNAL_LLMA_MULTIMODAL" in os.environ:
+            del os.environ["_INTERNAL_LLMA_MULTIMODAL"]
+
+        input_data = [
+            {
+                "parts": [
+                    {
+                        "inline_data": {
+                            "mime_type": "audio/L16;codec=pcm;rate=24000",
+                            "data": "base64audiodata",
+                        }
+                    }
+                ]
+            }
+        ]
+
+        result = sanitize_gemini(input_data)
+        self.assertEqual(
+            result[0]["parts"][0]["inline_data"]["data"], REDACTED_IMAGE_PLACEHOLDER
+        )
+
+    def test_gemini_audio_preserved_with_flag(self):
+        """Gemini audio should be preserved when _INTERNAL_LLMA_MULTIMODAL is enabled."""
+        os.environ["_INTERNAL_LLMA_MULTIMODAL"] = "true"
+
+        input_data = [
+            {
+                "parts": [
+                    {
+                        "inline_data": {
+                            "mime_type": "audio/L16;codec=pcm;rate=24000",
+                            "data": "base64audiodata",
+                        }
+                    }
+                ]
+            }
+        ]
+
+        result = sanitize_gemini(input_data)
+        self.assertEqual(
+            result[0]["parts"][0]["inline_data"]["data"], "base64audiodata"
+        )
+
+
 if __name__ == "__main__":
     unittest.main()

posthog/test/ai/test_system_prompts.py

@@ -11,7 +11,10 @@ regardless of how they're passed to the providers:
 
 import time
 import unittest
-from unittest.mock import patch, MagicMock
+from unittest.mock import MagicMock, patch
+
+from posthog.client import Client
+from posthog.test.test_utils import FAKE_TEST_API_KEY
 
 
 class TestSystemPromptCapture(unittest.TestCase):
@@ -24,7 +27,8 @@ class TestSystemPromptCapture(unittest.TestCase):
         self.test_response = "I'm doing well, thank you!"
 
         # Create mock PostHog client
-        self.client = MagicMock()
+        self.client = Client(FAKE_TEST_API_KEY)
+        self.client._enqueue = MagicMock()
         self.client.privacy_mode = False
 
     def _assert_system_prompt_captured(self, captured_input):
@@ -53,10 +57,11 @@ class TestSystemPromptCapture(unittest.TestCase):
     def test_openai_messages_array_system_prompt(self):
         """Test OpenAI with system prompt in messages array."""
         try:
-            from posthog.ai.openai import OpenAI
             from openai.types.chat import ChatCompletion, ChatCompletionMessage
             from openai.types.chat.chat_completion import Choice
             from openai.types.completion_usage import CompletionUsage
+
+            from posthog.ai.openai import OpenAI
         except ImportError:
             self.skipTest("OpenAI package not available")
 
@@ -94,17 +99,18 @@ class TestSystemPromptCapture(unittest.TestCase):
                 model="gpt-4", messages=messages, posthog_distinct_id="test-user"
             )
 
-            self.assertEqual(len(self.client.capture.call_args_list), 1)
-            properties = self.client.capture.call_args_list[0][1]["properties"]
+            self.assertEqual(len(self.client._enqueue.call_args_list), 1)
+            properties = self.client._enqueue.call_args_list[0][0][0]["properties"]
             self._assert_system_prompt_captured(properties["$ai_input"])
 
     def test_openai_separate_system_parameter(self):
         """Test OpenAI with system prompt as separate parameter."""
         try:
-            from posthog.ai.openai import OpenAI
             from openai.types.chat import ChatCompletion, ChatCompletionMessage
             from openai.types.chat.chat_completion import Choice
             from openai.types.completion_usage import CompletionUsage
+
+            from posthog.ai.openai import OpenAI
         except ImportError:
             self.skipTest("OpenAI package not available")
 
@@ -142,18 +148,21 @@ class TestSystemPromptCapture(unittest.TestCase):
                 posthog_distinct_id="test-user",
             )
 
-            self.assertEqual(len(self.client.capture.call_args_list), 1)
-            properties = self.client.capture.call_args_list[0][1]["properties"]
+            self.assertEqual(len(self.client._enqueue.call_args_list), 1)
+            properties = self.client._enqueue.call_args_list[0][0][0]["properties"]
             self._assert_system_prompt_captured(properties["$ai_input"])
 
     def test_openai_streaming_system_parameter(self):
         """Test OpenAI streaming with system parameter."""
         try:
-            from posthog.ai.openai import OpenAI
-            from openai.types.chat.chat_completion_chunk import ChatCompletionChunk
+            from openai.types.chat.chat_completion_chunk import (
+                ChatCompletionChunk,
+                ChoiceDelta,
+            )
             from openai.types.chat.chat_completion_chunk import Choice as ChoiceChunk
-            from openai.types.chat.chat_completion_chunk import ChoiceDelta
             from openai.types.completion_usage import CompletionUsage
+
+            from posthog.ai.openai import OpenAI
         except ImportError:
             self.skipTest("OpenAI package not available")
 
@@ -206,8 +215,8 @@ class TestSystemPromptCapture(unittest.TestCase):
 
             list(response_generator)  # Consume generator
 
-            self.assertEqual(len(self.client.capture.call_args_list), 1)
-            properties = self.client.capture.call_args_list[0][1]["properties"]
+            self.assertEqual(len(self.client._enqueue.call_args_list), 1)
+            properties = self.client._enqueue.call_args_list[0][0][0]["properties"]
             self._assert_system_prompt_captured(properties["$ai_input"])
 
     # Anthropic Tests
@@ -239,8 +248,8 @@ class TestSystemPromptCapture(unittest.TestCase):
                 posthog_distinct_id="test-user",
             )
 
-            self.assertEqual(len(self.client.capture.call_args_list), 1)
-            properties = self.client.capture.call_args_list[0][1]["properties"]
+            self.assertEqual(len(self.client._enqueue.call_args_list), 1)
+            properties = self.client._enqueue.call_args_list[0][0][0]["properties"]
             self._assert_system_prompt_captured(properties["$ai_input"])
 
     def test_anthropic_separate_system_parameter(self):
@@ -269,8 +278,8 @@ class TestSystemPromptCapture(unittest.TestCase):
                 posthog_distinct_id="test-user",
             )
 
-            self.assertEqual(len(self.client.capture.call_args_list), 1)
-            properties = self.client.capture.call_args_list[0][1]["properties"]
+            self.assertEqual(len(self.client._enqueue.call_args_list), 1)
+            properties = self.client._enqueue.call_args_list[0][0][0]["properties"]
             self._assert_system_prompt_captured(properties["$ai_input"])
 
     # Gemini Tests
@@ -310,8 +319,8 @@ class TestSystemPromptCapture(unittest.TestCase):
                 posthog_distinct_id="test-user",
             )
 
-            self.assertEqual(len(self.client.capture.call_args_list), 1)
-            properties = self.client.capture.call_args_list[0][1]["properties"]
+            self.assertEqual(len(self.client._enqueue.call_args_list), 1)
+            properties = self.client._enqueue.call_args_list[0][0][0]["properties"]
             self._assert_system_prompt_captured(properties["$ai_input"])
 
     def test_gemini_system_instruction_parameter(self):
@@ -349,6 +358,6 @@ class TestSystemPromptCapture(unittest.TestCase):
                 posthog_distinct_id="test-user",
             )
 
-            self.assertEqual(len(self.client.capture.call_args_list), 1)
-            properties = self.client.capture.call_args_list[0][1]["properties"]
+            self.assertEqual(len(self.client._enqueue.call_args_list), 1)
+            properties = self.client._enqueue.call_args_list[0][0][0]["properties"]
             self._assert_system_prompt_captured(properties["$ai_input"])

posthog/test/test_client.py

@@ -9,7 +9,7 @@ from parameterized import parameterized
 
 from posthog.client import Client
 from posthog.contexts import get_context_session_id, new_context, set_context_session
-from posthog.request import APIError
+from posthog.request import APIError, GetResponse
 from posthog.test.test_utils import FAKE_TEST_API_KEY
 from posthog.types import FeatureFlag, LegacyFlagMetadata
 from posthog.version import VERSION
@@ -198,12 +198,6 @@ class TestClient(unittest.TestCase):
             print(capture_call)
             self.assertEqual(capture_call[1]["distinct_id"], "distinct_id")
             self.assertEqual(capture_call[0][0], "$exception")
-            self.assertEqual(
-                capture_call[1]["properties"]["$exception_type"], "Exception"
-            )
-            self.assertEqual(
-                capture_call[1]["properties"]["$exception_message"], "test exception"
-            )
             self.assertEqual(
                 capture_call[1]["properties"]["$exception_list"][0]["mechanism"][
                     "type"
@@ -415,7 +409,9 @@ class TestClient(unittest.TestCase):
             )
             client.feature_flags = [multivariate_flag, basic_flag, false_flag]
 
-            msg_uuid = client.capture("python test event", distinct_id="distinct_id")
+            msg_uuid = client.capture(
+                "python test event", distinct_id="distinct_id", send_feature_flags=True
+            )
             self.assertIsNotNone(msg_uuid)
             self.assertFalse(self.failed)
 
@@ -483,6 +479,20 @@ class TestClient(unittest.TestCase):
             self.assertEqual(client.cohorts, {})
             self.assertIn("PostHog feature flags quota limited", logs.output[0])
 
+    @mock.patch("posthog.client.get")
+    def test_load_feature_flags_unauthorized(self, patch_get):
+        patch_get.side_effect = APIError(401, "Unauthorized")
+
+        client = Client(FAKE_TEST_API_KEY, personal_api_key="test")
+        with self.assertLogs("posthog", level="ERROR") as logs:
+            client._load_feature_flags()
+
+            self.assertEqual(client.feature_flags, [])
+            self.assertEqual(client.feature_flags_by_key, {})
+            self.assertEqual(client.group_type_mapping, {})
+            self.assertEqual(client.cohorts, {})
+            self.assertIn("please set a valid personal_api_key", logs.output[0])
+
     @mock.patch("posthog.client.flags")
     def test_dont_override_capture_with_local_flags(self, patch_flags):
         patch_flags.return_value = {"featureFlags": {"beta-feature": "random-variant"}}
@@ -571,6 +581,7 @@ class TestClient(unittest.TestCase):
                 "python test event",
                 distinct_id="distinct_id",
                 properties={"$feature/beta-feature-local": "my-custom-variant"},
+                send_feature_flags=True,
             )
             self.assertIsNotNone(msg_uuid)
             self.assertFalse(self.failed)
@@ -651,6 +662,7 @@ class TestClient(unittest.TestCase):
                 person_properties={},
                 group_properties={},
                 geoip_disable=True,
+                device_id=None,
             )
 
     @mock.patch("posthog.client.flags")
@@ -715,6 +727,7 @@ class TestClient(unittest.TestCase):
                 person_properties={},
                 group_properties={},
                 geoip_disable=False,
+                device_id=None,
             )
 
     @mock.patch("posthog.client.flags")
@@ -752,6 +765,178 @@ class TestClient(unittest.TestCase):
 
             self.assertEqual(patch_flags.call_count, 0)
 
+    @mock.patch("posthog.client.flags")
+    def test_capture_with_send_feature_flags_false_and_local_evaluation_doesnt_send_flags(
+        self, patch_flags
+    ):
+        """Test that send_feature_flags=False with local evaluation enabled does NOT send flags"""
+        patch_flags.return_value = {"featureFlags": {"beta-feature": "remote-variant"}}
+
+        multivariate_flag = {
+            "id": 1,
+            "name": "Beta Feature",
+            "key": "beta-feature-local",
+            "active": True,
+            "rollout_percentage": 100,
+            "filters": {
+                "groups": [
+                    {
+                        "rollout_percentage": 100,
+                    },
+                ],
+                "multivariate": {
+                    "variants": [
+                        {
+                            "key": "first-variant",
+                            "name": "First Variant",
+                            "rollout_percentage": 50,
+                        },
+                        {
+                            "key": "second-variant",
+                            "name": "Second Variant",
+                            "rollout_percentage": 50,
+                        },
+                    ]
+                },
+            },
+        }
+        simple_flag = {
+            "id": 2,
+            "name": "Simple Flag",
+            "key": "simple-flag",
+            "active": True,
+            "filters": {
+                "groups": [
+                    {
+                        "rollout_percentage": 100,
+                    }
+                ],
+            },
+        }
+
+        with mock.patch("posthog.client.batch_post") as mock_post:
+            client = Client(
+                FAKE_TEST_API_KEY,
+                on_error=self.set_fail,
+                personal_api_key=FAKE_TEST_API_KEY,
+                sync_mode=True,
+            )
+            client.feature_flags = [multivariate_flag, simple_flag]
+
+            msg_uuid = client.capture(
+                "python test event",
+                distinct_id="distinct_id",
+                send_feature_flags=False,
+            )
+            self.assertIsNotNone(msg_uuid)
+            self.assertFalse(self.failed)
+
+            # Get the enqueued message from the mock
+            mock_post.assert_called_once()
+            batch_data = mock_post.call_args[1]["batch"]
+            msg = batch_data[0]
+
+            self.assertEqual(msg["event"], "python test event")
+            self.assertEqual(msg["distinct_id"], "distinct_id")
+
+            # CRITICAL: Verify local flags are NOT included in the event
+            self.assertNotIn("$feature/beta-feature-local", msg["properties"])
+            self.assertNotIn("$feature/simple-flag", msg["properties"])
+            self.assertNotIn("$active_feature_flags", msg["properties"])
+
+            # CRITICAL: Verify the /flags API was NOT called
+            self.assertEqual(patch_flags.call_count, 0)
+
+    @mock.patch("posthog.client.flags")
+    def test_capture_with_send_feature_flags_true_and_local_evaluation_uses_local_flags(
+        self, patch_flags
+    ):
+        """Test that send_feature_flags=True with local evaluation enabled uses local flags without API call"""
+        patch_flags.return_value = {"featureFlags": {"remote-flag": "remote-variant"}}
+
+        multivariate_flag = {
+            "id": 1,
+            "name": "Beta Feature",
+            "key": "beta-feature-local",
+            "active": True,
+            "rollout_percentage": 100,
+            "filters": {
+                "groups": [
+                    {
+                        "rollout_percentage": 100,
+                    },
+                ],
+                "multivariate": {
+                    "variants": [
+                        {
+                            "key": "first-variant",
+                            "name": "First Variant",
+                            "rollout_percentage": 50,
+                        },
+                        {
+                            "key": "second-variant",
+                            "name": "Second Variant",
+                            "rollout_percentage": 50,
+                        },
+                    ]
+                },
+            },
+        }
+        simple_flag = {
+            "id": 2,
+            "name": "Simple Flag",
+            "key": "simple-flag",
+            "active": True,
+            "filters": {
+                "groups": [
+                    {
+                        "rollout_percentage": 100,
+                    }
+                ],
+            },
+        }
+
+        with mock.patch("posthog.client.batch_post") as mock_post:
+            client = Client(
+                FAKE_TEST_API_KEY,
+                on_error=self.set_fail,
+                personal_api_key=FAKE_TEST_API_KEY,
+                sync_mode=True,
+            )
+            client.feature_flags = [multivariate_flag, simple_flag]
+
+            msg_uuid = client.capture(
+                "python test event",
+                distinct_id="distinct_id",
+                send_feature_flags=True,
+            )
+            self.assertIsNotNone(msg_uuid)
+            self.assertFalse(self.failed)
+
+            # Get the enqueued message from the mock
+            mock_post.assert_called_once()
+            batch_data = mock_post.call_args[1]["batch"]
+            msg = batch_data[0]
+
+            self.assertEqual(msg["event"], "python test event")
+            self.assertEqual(msg["distinct_id"], "distinct_id")
+
+            # Verify local flags are included in the event
+            self.assertIn("$feature/beta-feature-local", msg["properties"])
+            self.assertIn("$feature/simple-flag", msg["properties"])
+            self.assertEqual(msg["properties"]["$feature/simple-flag"], True)
+
+            # Verify active feature flags are set correctly
+            active_flags = msg["properties"]["$active_feature_flags"]
+            self.assertIn("beta-feature-local", active_flags)
+            self.assertIn("simple-flag", active_flags)
+
+            # The remote flag should NOT be included since we used local evaluation
+            self.assertNotIn("$feature/remote-flag", msg["properties"])
+
+            # CRITICAL: Verify the /flags API was NOT called
+            self.assertEqual(patch_flags.call_count, 0)
+
     @mock.patch("posthog.client.flags")
     def test_capture_with_send_feature_flags_options_only_evaluate_locally_true(
         self, patch_flags
@@ -1742,6 +1927,7 @@ class TestClient(unittest.TestCase):
             person_properties={"distinct_id": "some_id"},
             group_properties={},
             geoip_disable=True,
+            device_id=None,
             flag_keys_to_evaluate=["random_key"],
         )
         patch_flags.reset_mock()
@@ -1757,6 +1943,7 @@ class TestClient(unittest.TestCase):
             person_properties={"distinct_id": "feature_enabled_distinct_id"},
             group_properties={},
             geoip_disable=True,
+            device_id=None,
             flag_keys_to_evaluate=["random_key"],
         )
         patch_flags.reset_mock()
@@ -1770,6 +1957,7 @@ class TestClient(unittest.TestCase):
             person_properties={"distinct_id": "all_flags_payloads_id"},
             group_properties={},
             geoip_disable=False,
+            device_id=None,
         )
 
     @mock.patch("posthog.client.Poller")
@@ -1818,6 +2006,7 @@ class TestClient(unittest.TestCase):
                 "instance": {"$group_key": "app.posthog.com"},
             },
             geoip_disable=False,
+            device_id=None,
             flag_keys_to_evaluate=["random_key"],
         )
 
@@ -1845,6 +2034,7 @@ class TestClient(unittest.TestCase):
                 "instance": {"$group_key": "app.posthog.com"},
             },
             geoip_disable=False,
+            device_id=None,
             flag_keys_to_evaluate=["random_key"],
         )
 
@@ -1862,8 +2052,116 @@ class TestClient(unittest.TestCase):
             person_properties={"distinct_id": "some_id"},
             group_properties={},
             geoip_disable=False,
+            device_id=None,
         )
 
+    @parameterized.expand(
+        [
+            # method, method_args, expected_person_props, expected_flag_keys
+            (
+                "get_feature_flag",
+                ["random_key", "some_id"],
+                {"distinct_id": "some_id"},
+                ["random_key"],
+            ),
+            (
+                "feature_enabled",
+                ["random_key", "some_id"],
+                {"distinct_id": "some_id"},
+                ["random_key"],
+            ),
+            (
+                "get_all_flags_and_payloads",
+                ["some_id"],
+                {"distinct_id": "some_id"},
+                None,
+            ),
+            ("get_all_flags", ["some_id"], {"distinct_id": "some_id"}, None),
+            ("get_flags_decision", ["some_id"], {}, None),
+        ]
+    )
+    @mock.patch("posthog.client.flags")
+    def test_device_id_is_passed_to_flags_request(
+        self,
+        method,
+        method_args,
+        expected_person_props,
+        expected_flag_keys,
+        patch_flags,
+    ):
+        """Test that device_id is properly passed to the flags request when provided."""
+        patch_flags.return_value = {"featureFlags": {"beta-feature": "random-variant"}}
+        client = Client(FAKE_TEST_API_KEY, on_error=self.set_fail)
+
+        getattr(client, method)(*method_args, device_id="test-device-123")
+
+        expected_call = {
+            "distinct_id": "some_id",
+            "groups": {},
+            "person_properties": expected_person_props,
+            "group_properties": {},
+            "geoip_disable": True,
+            "device_id": "test-device-123",
+        }
+        if expected_flag_keys:
+            expected_call["flag_keys_to_evaluate"] = expected_flag_keys
+
+        patch_flags.assert_called_with(
+            "random_key", "https://us.i.posthog.com", timeout=3, **expected_call
+        )
+
+    @mock.patch("posthog.client.flags")
+    def test_device_id_from_context_is_used_in_flags_request(self, patch_flags):
+        """Test that device_id from context is used in flags request when not explicitly provided."""
+        from posthog.contexts import new_context, set_context_device_id
+
+        patch_flags.return_value = {
+            "featureFlags": {
+                "beta-feature": "random-variant",
+            }
+        }
+        client = Client(
+            FAKE_TEST_API_KEY,
+            on_error=self.set_fail,
+        )
+
+        # Test that device_id from context is used
+        with new_context():
+            set_context_device_id("context-device-id")
+            client.get_feature_flag("random_key", "some_id")
+            patch_flags.assert_called_with(
+                "random_key",
+                "https://us.i.posthog.com",
+                timeout=3,
+                distinct_id="some_id",
+                groups={},
+                person_properties={"distinct_id": "some_id"},
+                group_properties={},
+                geoip_disable=True,
+                device_id="context-device-id",
+                flag_keys_to_evaluate=["random_key"],
+            )
+
+        # Test that explicit device_id overrides context
+        patch_flags.reset_mock()
+        with new_context():
+            set_context_device_id("context-device-id")
+            client.get_feature_flag(
+                "random_key", "some_id", device_id="explicit-device-id"
+            )
+            patch_flags.assert_called_with(
+                "random_key",
+                "https://us.i.posthog.com",
+                timeout=3,
+                distinct_id="some_id",
+                groups={},
+                person_properties={"distinct_id": "some_id"},
+                group_properties={},
+                geoip_disable=True,
+                device_id="explicit-device-id",
+                flag_keys_to_evaluate=["random_key"],
+            )
+
     @parameterized.expand(
         [
             # name, sys_platform, version_info, expected_runtime, expected_version, expected_os, expected_os_version, platform_method, platform_return, distro_info
@@ -2095,13 +2393,21 @@ class TestClient(unittest.TestCase):
         self, patch_get, patch_poller
     ):
         """Test that when enable_local_evaluation=False, the poller is not started"""
-        patch_get.return_value = {
-            "flags": [
-                {"id": 1, "name": "Beta Feature", "key": "beta-feature", "active": True}
-            ],
-            "group_type_mapping": {},
-            "cohorts": {},
-        }
+        patch_get.return_value = GetResponse(
+            data={
+                "flags": [
+                    {
+                        "id": 1,
+                        "name": "Beta Feature",
+                        "key": "beta-feature",
+                        "active": True,
+                    }
+                ],
+                "group_type_mapping": {},
+                "cohorts": {},
+            },
+            etag='"test-etag"',
+        )
 
         client = Client(
             FAKE_TEST_API_KEY,
@@ -2123,13 +2429,21 @@ class TestClient(unittest.TestCase):
     @mock.patch("posthog.client.get")
     def test_enable_local_evaluation_true_starts_poller(self, patch_get, patch_poller):
         """Test that when enable_local_evaluation=True (default), the poller is started"""
-        patch_get.return_value = {
-            "flags": [
-                {"id": 1, "name": "Beta Feature", "key": "beta-feature", "active": True}
-            ],
-            "group_type_mapping": {},
-            "cohorts": {},
-        }
+        patch_get.return_value = GetResponse(
+            data={
+                "flags": [
+                    {
+                        "id": 1,
+                        "name": "Beta Feature",
+                        "key": "beta-feature",
+                        "active": True,
+                    }
+                ],
+                "group_type_mapping": {},
+                "cohorts": {},
+            },
+            etag='"test-etag"',
+        )
 
         client = Client(
             FAKE_TEST_API_KEY,

posthog/test/test_consumer.py

@@ -1,8 +1,10 @@
 import json
 import time
 import unittest
+from typing import Any
 
 import mock
+from parameterized import parameterized
 
 try:
     from queue import Queue
@@ -14,15 +16,19 @@ from posthog.request import APIError
 from posthog.test.test_utils import TEST_API_KEY
 
 
+def _track_event(event_name: str = "python event") -> dict[str, str]:
+    return {"type": "track", "event": event_name, "distinct_id": "distinct_id"}
+
+
 class TestConsumer(unittest.TestCase):
-    def test_next(self):
+    def test_next(self) -> None:
         q = Queue()
         consumer = Consumer(q, "")
         q.put(1)
         next = consumer.next()
         self.assertEqual(next, [1])
 
-    def test_next_limit(self):
+    def test_next_limit(self) -> None:
         q = Queue()
         flush_at = 50
         consumer = Consumer(q, "", flush_at)
@@ -31,7 +37,7 @@ class TestConsumer(unittest.TestCase):
         next = consumer.next()
         self.assertEqual(next, list(range(flush_at)))
 
-    def test_dropping_oversize_msg(self):
+    def test_dropping_oversize_msg(self) -> None:
         q = Queue()
         consumer = Consumer(q, "")
         oversize_msg = {"m": "x" * MAX_MSG_SIZE}
@@ -40,15 +46,14 @@ class TestConsumer(unittest.TestCase):
         self.assertEqual(next, [])
         self.assertTrue(q.empty())
 
-    def test_upload(self):
+    def test_upload(self) -> None:
         q = Queue()
         consumer = Consumer(q, TEST_API_KEY)
-        track = {"type": "track", "event": "python event", "distinct_id": "distinct_id"}
-        q.put(track)
+        q.put(_track_event())
         success = consumer.upload()
         self.assertTrue(success)
 
-    def test_flush_interval(self):
+    def test_flush_interval(self) -> None:
         # Put _n_ items in the queue, pausing a little bit more than
         # _flush_interval_ after each one.
         # The consumer should upload _n_ times.
@@ -57,17 +62,12 @@ class TestConsumer(unittest.TestCase):
         consumer = Consumer(q, TEST_API_KEY, flush_at=10, flush_interval=flush_interval)
         with mock.patch("posthog.consumer.batch_post") as mock_post:
             consumer.start()
-            for i in range(0, 3):
-                track = {
-                    "type": "track",
-                    "event": "python event %d" % i,
-                    "distinct_id": "distinct_id",
-                }
-                q.put(track)
+            for i in range(3):
+                q.put(_track_event("python event %d" % i))
                 time.sleep(flush_interval * 1.1)
             self.assertEqual(mock_post.call_count, 3)
 
-    def test_multiple_uploads_per_interval(self):
+    def test_multiple_uploads_per_interval(self) -> None:
         # Put _flush_at*2_ items in the queue at once, then pause for
         # _flush_interval_. The consumer should upload 2 times.
         q = Queue()
@@ -78,88 +78,60 @@ class TestConsumer(unittest.TestCase):
         )
         with mock.patch("posthog.consumer.batch_post") as mock_post:
             consumer.start()
-            for i in range(0, flush_at * 2):
-                track = {
-                    "type": "track",
-                    "event": "python event %d" % i,
-                    "distinct_id": "distinct_id",
-                }
-                q.put(track)
+            for i in range(flush_at * 2):
+                q.put(_track_event("python event %d" % i))
             time.sleep(flush_interval * 1.1)
             self.assertEqual(mock_post.call_count, 2)
 
-    def test_request(self):
+    def test_request(self) -> None:
         consumer = Consumer(None, TEST_API_KEY)
-        track = {"type": "track", "event": "python event", "distinct_id": "distinct_id"}
-        consumer.request([track])
+        consumer.request([_track_event()])
 
-    def _test_request_retry(self, consumer, expected_exception, exception_count):
-        def mock_post(*args, **kwargs):
-            mock_post.call_count += 1
-            if mock_post.call_count <= exception_count:
-                raise expected_exception
+    def _run_retry_test(
+        self, exception: Exception, exception_count: int, retries: int = 10
+    ) -> None:
+        call_count = [0]
 
-        mock_post.call_count = 0
+        def mock_post(*args: Any, **kwargs: Any) -> None:
+            call_count[0] += 1
+            if call_count[0] <= exception_count:
+                raise exception
 
+        consumer = Consumer(None, TEST_API_KEY, retries=retries)
         with mock.patch(
             "posthog.consumer.batch_post", mock.Mock(side_effect=mock_post)
         ):
-            track = {
-                "type": "track",
-                "event": "python event",
-                "distinct_id": "distinct_id",
-            }
-            # request() should succeed if the number of exceptions raised is
-            # less than the retries paramater.
-            if exception_count <= consumer.retries:
-                consumer.request([track])
+            if exception_count <= retries:
+                consumer.request([_track_event()])
             else:
-                # if exceptions are raised more times than the retries
-                # parameter, we expect the exception to be returned to
-                # the caller.
-                try:
-                    consumer.request([track])
-                except type(expected_exception) as exc:
-                    self.assertEqual(exc, expected_exception)
-                else:
-                    self.fail(
-                        "request() should raise an exception if still failing after %d retries"
-                        % consumer.retries
-                    )
+                with self.assertRaises(type(exception)):
+                    consumer.request([_track_event()])
 
-    def test_request_retry(self):
-        # we should retry on general errors
-        consumer = Consumer(None, TEST_API_KEY)
-        self._test_request_retry(consumer, Exception("generic exception"), 2)
+    @parameterized.expand(
+        [
+            ("general_errors", Exception("generic exception"), 2),
+            ("server_errors", APIError(500, "Internal Server Error"), 2),
+            ("rate_limit_errors", APIError(429, "Too Many Requests"), 2),
+        ]
+    )
+    def test_request_retries_on_retriable_errors(
+        self, _name: str, exception: Exception, exception_count: int
+    ) -> None:
+        self._run_retry_test(exception, exception_count)
 
-        # we should retry on server errors
-        consumer = Consumer(None, TEST_API_KEY)
-        self._test_request_retry(consumer, APIError(500, "Internal Server Error"), 2)
+    def test_request_does_not_retry_client_errors(self) -> None:
+        with self.assertRaises(APIError):
+            self._run_retry_test(APIError(400, "Client Errors"), 1)
 
-        # we should retry on HTTP 429 errors
-        consumer = Consumer(None, TEST_API_KEY)
-        self._test_request_retry(consumer, APIError(429, "Too Many Requests"), 2)
+    def test_request_fails_when_exceptions_exceed_retries(self) -> None:
+        self._run_retry_test(APIError(500, "Internal Server Error"), 4, retries=3)
 
-        # we should NOT retry on other client errors
-        consumer = Consumer(None, TEST_API_KEY)
-        api_error = APIError(400, "Client Errors")
-        try:
-            self._test_request_retry(consumer, api_error, 1)
-        except APIError:
-            pass
-        else:
-            self.fail("request() should not retry on client errors")
-
-        # test for number of exceptions raise > retries value
-        consumer = Consumer(None, TEST_API_KEY, retries=3)
-        self._test_request_retry(consumer, APIError(500, "Internal Server Error"), 3)
-
-    def test_pause(self):
+    def test_pause(self) -> None:
         consumer = Consumer(None, TEST_API_KEY)
         consumer.pause()
         self.assertFalse(consumer.running)
 
-    def test_max_batch_size(self):
+    def test_max_batch_size(self) -> None:
         q = Queue()
         consumer = Consumer(q, TEST_API_KEY, flush_at=100000, flush_interval=3)
         properties = {}
@@ -175,7 +147,7 @@ class TestConsumer(unittest.TestCase):
         # Let's capture 8MB of data to trigger two batches
         n_msgs = int(8_000_000 / msg_size)
 
-        def mock_post_fn(_, data, **kwargs):
+        def mock_post_fn(_: str, data: str, **kwargs: Any) -> mock.Mock:
             res = mock.Mock()
             res.status_code = 200
             request_size = len(data.encode())
@@ -194,3 +166,91 @@ class TestConsumer(unittest.TestCase):
                 q.put(track)
             q.join()
             self.assertEqual(mock_post.call_count, 2)
+
+    def test_request_sleeps_with_retry_after(self) -> None:
+        error = APIError(429, "Too Many Requests", retry_after=5.0)
+        call_count = [0]
+
+        def mock_post(*args: Any, **kwargs: Any) -> None:
+            call_count[0] += 1
+            if call_count[0] <= 1:
+                raise error
+
+        consumer = Consumer(None, TEST_API_KEY, retries=3)
+        with (
+            mock.patch("posthog.consumer.batch_post", side_effect=mock_post),
+            mock.patch("posthog.consumer.time.sleep") as mock_sleep,
+        ):
+            consumer.request([_track_event()])
+            mock_sleep.assert_called_once_with(5.0)
+
+    def test_request_uses_exponential_backoff_without_retry_after(self) -> None:
+        error = APIError(503, "Service Unavailable")
+        call_count = [0]
+
+        def mock_post(*args: Any, **kwargs: Any) -> None:
+            call_count[0] += 1
+            if call_count[0] <= 3:
+                raise error
+
+        consumer = Consumer(None, TEST_API_KEY, retries=3)
+        with (
+            mock.patch("posthog.consumer.batch_post", side_effect=mock_post),
+            mock.patch("posthog.consumer.time.sleep") as mock_sleep,
+        ):
+            consumer.request([_track_event()])
+            self.assertEqual(
+                mock_sleep.call_args_list,
+                [
+                    mock.call(1),  # 2^0
+                    mock.call(2),  # 2^1
+                    mock.call(4),  # 2^2
+                ],
+            )
+
+    def test_request_retries_on_408(self) -> None:
+        call_count = [0]
+
+        def mock_post(*args: Any, **kwargs: Any) -> None:
+            call_count[0] += 1
+            if call_count[0] <= 1:
+                raise APIError(408, "Request Timeout")
+
+        consumer = Consumer(None, TEST_API_KEY, retries=3)
+        with (
+            mock.patch("posthog.consumer.batch_post", side_effect=mock_post),
+            mock.patch("posthog.consumer.time.sleep"),
+        ):
+            consumer.request([_track_event()])
+            self.assertEqual(call_count[0], 2)
+
+    @parameterized.expand(
+        [
+            ("on_error_succeeds", False),
+            ("on_error_raises", True),
+        ]
+    )
+    def test_upload_exception_calls_on_error_and_does_not_raise(
+        self, _name: str, on_error_raises: bool
+    ) -> None:
+        on_error_called: list[tuple[Exception, list[dict[str, str]]]] = []
+
+        def on_error(e: Exception, batch: list[dict[str, str]]) -> None:
+            on_error_called.append((e, batch))
+            if on_error_raises:
+                raise Exception("on_error failed")
+
+        q = Queue()
+        consumer = Consumer(q, TEST_API_KEY, on_error=on_error)
+        track = _track_event()
+        q.put(track)
+
+        with mock.patch.object(
+            consumer, "request", side_effect=Exception("request failed")
+        ):
+            result = consumer.upload()
+
+        self.assertFalse(result)
+        self.assertEqual(len(on_error_called), 1)
+        self.assertEqual(str(on_error_called[0][0]), "request failed")
+        self.assertEqual(on_error_called[0][1], [track])

posthog/test/test_contexts.py

@@ -191,6 +191,32 @@ class TestContexts(unittest.TestCase):
             assert get_context_distinct_id() == "user123"
             assert get_context_session_id() == "session456"
 
+    def test_child_tags_override_parent_tags_in_non_fresh_context(self):
+        with new_context(fresh=True):
+            tag("shared_key", "parent_value")
+            tag("parent_only", "parent")
+
+            with new_context(fresh=False):
+                # Child should inherit parent tags
+                assert get_tags()["parent_only"] == "parent"
+
+                # Child sets same key - should override parent
+                tag("shared_key", "child_value")
+                tag("child_only", "child")
+
+                tags = get_tags()
+                # Child value should win for shared key
+                assert tags["shared_key"] == "child_value"
+                # Both parent and child tags should be present
+                assert tags["parent_only"] == "parent"
+                assert tags["child_only"] == "child"
+
+            # Parent context should be unchanged
+            parent_tags = get_tags()
+            assert parent_tags["shared_key"] == "parent_value"
+            assert parent_tags["parent_only"] == "parent"
+            assert "child_only" not in parent_tags
+
     def test_scoped_decorator_with_context_ids(self):
         @scoped()
         def function_with_context():

posthog/test/test_exception_capture.py

@@ -59,8 +59,29 @@ def test_code_variables_capture(tmpdir):
         my_number = 42
         my_bool = True
         my_dict = {"name": "test", "value": 123}
+        my_sensitive_dict = {
+            "safe_key": "safe_value",
+            "password": "secret123",  # key matches pattern -> should be masked
+            "other_key": "contains_password_here",  # value matches pattern -> should be masked
+        }
+        my_nested_dict = {
+            "level1": {
+                "level2": {
+                    "api_key": "nested_secret",  # deeply nested key matches
+                    "data": "contains_token_here",  # deeply nested value matches
+                    "safe": "visible",
+                }
+            }
+        }
+        my_list = ["safe_item", "has_password_inside", "another_safe"]
+        my_tuple = ("tuple_safe", "secret_in_value", "tuple_also_safe")
+        my_list_of_dicts = [
+            {"id": 1, "password": "list_dict_secret"},
+            {"id": 2, "value": "safe_value"},
+        ]
         my_obj = UnserializableObject()
-        my_password = "secret123"  # Should be masked by default
+        my_password = "secret123"  # Should be masked by default (name matches)
+        my_innocent_var = "contains_password_here"  # Should be masked by default (value matches)
         __should_be_ignored = "hidden"  # Should be ignored by default
         
         1/0  # Trigger exception
@@ -96,8 +117,31 @@ def test_code_variables_capture(tmpdir):
     assert b"'my_number': 42" in output
     assert b"'my_bool': 'True'" in output
     assert b'"my_dict": "{\\"name\\": \\"test\\", \\"value\\": 123}"' in output
-    assert b'"my_obj": "<UnserializableObject>"' in output
+    assert (
+        b'{\\"safe_key\\": \\"safe_value\\", \\"password\\": \\"$$_posthog_redacted_based_on_masking_rules_$$\\", \\"other_key\\": \\"$$_posthog_redacted_based_on_masking_rules_$$\\"}'
+        in output
+    )
+    assert (
+        b'{\\"level1\\": {\\"level2\\": {\\"api_key\\": \\"$$_posthog_redacted_based_on_masking_rules_$$\\", \\"data\\": \\"$$_posthog_redacted_based_on_masking_rules_$$\\", \\"safe\\": \\"visible\\"}}}'
+        in output
+    )
+    assert (
+        b'[\\"safe_item\\", \\"$$_posthog_redacted_based_on_masking_rules_$$\\", \\"another_safe\\"]'
+        in output
+    )
+    assert (
+        b'[\\"tuple_safe\\", \\"$$_posthog_redacted_based_on_masking_rules_$$\\", \\"tuple_also_safe\\"]'
+        in output
+    )
+    assert (
+        b'[{\\"id\\": 1, \\"password\\": \\"$$_posthog_redacted_based_on_masking_rules_$$\\"}, {\\"id\\": 2, \\"value\\": \\"safe_value\\"}]'
+        in output
+    )
+    assert b"<__main__.UnserializableObject object at" in output
     assert b"'my_password': '$$_posthog_redacted_based_on_masking_rules_$$'" in output
+    assert (
+        b"'my_innocent_var': '$$_posthog_redacted_based_on_masking_rules_$$'" in output
+    )
     assert b"'__should_be_ignored':" not in output
 
     # Variables from intermediate_function frame
@@ -332,3 +376,314 @@ def test_code_variables_enabled_then_disabled_in_context(tmpdir):
     assert '"code_variables":' not in output
     assert "'my_var'" not in output
     assert "'important_value'" not in output
+
+
+def test_code_variables_repr_fallback(tmpdir):
+    app = tmpdir.join("app.py")
+    app.write(
+        dedent(
+            """
+    import os
+    import re
+    from datetime import datetime, timedelta
+    from decimal import Decimal
+    from fractions import Fraction
+    from posthog import Posthog
+    
+    class CustomReprClass:
+        def __repr__(self):
+            return '<CustomReprClass: custom representation>'
+    
+    posthog = Posthog(
+        'phc_x', 
+        host='https://eu.i.posthog.com', 
+        debug=True, 
+        enable_exception_autocapture=True,
+        capture_exception_code_variables=True,
+        project_root=os.path.dirname(os.path.abspath(__file__))
+    )
+    
+    def trigger_error():
+        my_regex = re.compile(r'\\d+')
+        my_datetime = datetime(2024, 1, 15, 10, 30, 45)
+        my_timedelta = timedelta(days=5, hours=3)
+        my_decimal = Decimal('123.456')
+        my_fraction = Fraction(3, 4)
+        my_set = {1, 2, 3}
+        my_frozenset = frozenset([4, 5, 6])
+        my_bytes = b'hello bytes'
+        my_bytearray = bytearray(b'mutable bytes')
+        my_memoryview = memoryview(b'memory view')
+        my_complex = complex(3, 4)
+        my_range = range(10)
+        my_custom = CustomReprClass()
+        my_lambda = lambda x: x * 2
+        my_function = trigger_error
+        
+        1/0
+    
+    trigger_error()
+    """
+        )
+    )
+
+    with pytest.raises(subprocess.CalledProcessError) as excinfo:
+        subprocess.check_output([sys.executable, str(app)], stderr=subprocess.STDOUT)
+
+    output = excinfo.value.output.decode("utf-8")
+
+    assert "ZeroDivisionError" in output
+    assert "code_variables" in output
+
+    assert "re.compile(" in output and "\\\\d+" in output
+    assert "datetime.datetime(2024, 1, 15, 10, 30, 45)" in output
+    assert "datetime.timedelta(days=5, seconds=10800)" in output
+    assert "Decimal('123.456')" in output
+    assert "Fraction(3, 4)" in output
+    assert "{1, 2, 3}" in output
+    assert "frozenset({4, 5, 6})" in output
+    assert "b'hello bytes'" in output
+    assert "bytearray(b'mutable bytes')" in output
+    assert "<memory at" in output
+    assert "(3+4j)" in output
+    assert "range(0, 10)" in output
+    assert "<CustomReprClass: custom representation>" in output
+    assert "<lambda>" in output
+    assert "<function trigger_error at" in output
+
+
+def test_code_variables_too_long_string_value_replaced(tmpdir):
+    app = tmpdir.join("app.py")
+    app.write(
+        dedent(
+            """
+    import os
+    from posthog import Posthog
+
+    posthog = Posthog(
+        'phc_x',
+        host='https://eu.i.posthog.com',
+        debug=True,
+        enable_exception_autocapture=True,
+        capture_exception_code_variables=True,
+        project_root=os.path.dirname(os.path.abspath(__file__))
+    )
+
+    def trigger_error():
+        short_value = "I am short"
+        long_value = "x" * 20000
+        long_blob = "password_" + "a" * 20000
+
+        1/0
+
+    trigger_error()
+    """
+        )
+    )
+
+    with pytest.raises(subprocess.CalledProcessError) as excinfo:
+        subprocess.check_output([sys.executable, str(app)], stderr=subprocess.STDOUT)
+
+    output = excinfo.value.output.decode("utf-8")
+
+    assert "ZeroDivisionError" in output
+    assert "code_variables" in output
+
+    assert "'short_value': 'I am short'" in output
+
+    assert "$$_posthog_value_too_long_$$" in output
+
+    assert "'long_blob': '$$_posthog_value_too_long_$$'" in output
+
+
+def test_code_variables_too_long_string_in_nested_dict(tmpdir):
+    app = tmpdir.join("app.py")
+    app.write(
+        dedent(
+            """
+    import os
+    from posthog import Posthog
+
+    posthog = Posthog(
+        'phc_x',
+        host='https://eu.i.posthog.com',
+        debug=True,
+        enable_exception_autocapture=True,
+        capture_exception_code_variables=True,
+        project_root=os.path.dirname(os.path.abspath(__file__))
+    )
+
+    def trigger_error():
+        my_data = {
+            "short_key": "short_val",
+            "long_key": "y" * 20000,
+            "nested": {
+                "deep_long": "z" * 20000,
+                "deep_short": "ok",
+            },
+        }
+
+        1/0
+
+    trigger_error()
+    """
+        )
+    )
+
+    with pytest.raises(subprocess.CalledProcessError) as excinfo:
+        subprocess.check_output([sys.executable, str(app)], stderr=subprocess.STDOUT)
+
+    output = excinfo.value.output.decode("utf-8")
+
+    assert "ZeroDivisionError" in output
+    assert "code_variables" in output
+
+    assert "short_val" in output
+    assert "ok" in output
+
+    assert "$$_posthog_value_too_long_$$" in output
+    assert "y" * 1000 not in output
+    assert "z" * 1000 not in output
+
+
+def test_mask_sensitive_data_too_long_dict_key():
+    from posthog.exception_utils import (
+        CODE_VARIABLES_TOO_LONG_VALUE,
+        _compile_patterns,
+        _mask_sensitive_data,
+    )
+
+    compiled_mask = _compile_patterns([r"(?i)password"])
+
+    result = _mask_sensitive_data(
+        {
+            "short": "visible",
+            "k" * 20000: "hidden_val",
+            "password": "secret",
+        },
+        compiled_mask,
+    )
+
+    assert result["short"] == "visible"
+    # This then gets shortened by the JSON truncation at 1024 chars anyways so no worries
+    assert result["k" * 20000] == CODE_VARIABLES_TOO_LONG_VALUE
+    assert result["password"] == "$$_posthog_redacted_based_on_masking_rules_$$"
+
+
+def test_mask_sensitive_data_circular_ref():
+    from posthog.exception_utils import _compile_patterns, _mask_sensitive_data
+
+    compiled_mask = _compile_patterns([r"(?i)password"])
+
+    # Circular dict
+    circular_dict = {"key": "value"}
+    circular_dict["self"] = circular_dict
+
+    result = _mask_sensitive_data(circular_dict, compiled_mask)
+    assert result["key"] == "value"
+    assert result["self"] == "<circular ref>"
+
+    # Circular list
+    circular_list = ["item"]
+    circular_list.append(circular_list)
+
+    result = _mask_sensitive_data(circular_list, compiled_mask)
+    assert result[0] == "item"
+    assert result[1] == "<circular ref>"
+
+
+def test_compile_patterns_fast_path_and_regex_fallback():
+    from posthog.exception_utils import _compile_patterns, _pattern_matches
+
+    # Simple case-insensitive patterns should become substrings
+    simple_only = _compile_patterns([r"(?i)password", r"(?i)token", r"(?i)jwt"])
+    substrings, regexes = simple_only
+    assert substrings == ["password", "token", "jwt"]
+    assert regexes == []
+
+    assert _pattern_matches("my_password_var", simple_only) is True
+    assert _pattern_matches("MY_TOKEN", simple_only) is True
+    assert _pattern_matches("safe_variable", simple_only) is False
+
+    # Complex regex patterns should stay as compiled regexes
+    complex_only = _compile_patterns([r"^__.*", r"\d{3,}", r"^sk_live_"])
+    substrings, regexes = complex_only
+    assert substrings == []
+    assert len(regexes) == 3
+
+    assert _pattern_matches("__dunder", complex_only) is True
+    assert _pattern_matches("has_999_numbers", complex_only) is True
+    assert _pattern_matches("sk_live_abc123", complex_only) is True
+    assert _pattern_matches("normal_var", complex_only) is False
+
+    # Mixed: simple substrings + complex regexes together
+    mixed = _compile_patterns(
+        [
+            r"(?i)secret",  # simple
+            r"(?i)api_key",  # simple
+            r"^__.*",  # regex
+            r"\btoken_\w+",  # regex
+        ]
+    )
+    substrings, regexes = mixed
+    assert substrings == ["secret", "api_key"]
+    assert len(regexes) == 2
+
+    # Substring matches
+    assert _pattern_matches("my_secret", mixed) is True
+    assert _pattern_matches("API_KEY_VALUE", mixed) is True
+
+    # Regex matches
+    assert _pattern_matches("__private", mixed) is True
+    assert _pattern_matches("token_abc", mixed) is True
+
+    # No match
+    assert _pattern_matches("safe_var", mixed) is False
+
+
+def test_mask_sensitive_data_large_dict_replaced():
+    from posthog.exception_utils import (
+        CODE_VARIABLES_TOO_LONG_VALUE,
+        _compile_patterns,
+        _mask_sensitive_data,
+    )
+
+    compiled_mask = _compile_patterns([r"(?i)password"])
+
+    large_dict = {f"key_{i}": f"value_{i}" for i in range(300)}
+
+    result = _mask_sensitive_data(large_dict, compiled_mask)
+
+    assert result == CODE_VARIABLES_TOO_LONG_VALUE
+
+
+def test_mask_sensitive_data_large_list_replaced():
+    from posthog.exception_utils import (
+        CODE_VARIABLES_TOO_LONG_VALUE,
+        _compile_patterns,
+        _mask_sensitive_data,
+    )
+
+    compiled_mask = _compile_patterns([r"(?i)password"])
+
+    large_list = [f"item_{i}" for i in range(300)]
+
+    result = _mask_sensitive_data(large_list, compiled_mask)
+
+    assert result == CODE_VARIABLES_TOO_LONG_VALUE
+
+
+def test_mask_sensitive_data_large_tuple_replaced():
+    from posthog.exception_utils import (
+        CODE_VARIABLES_TOO_LONG_VALUE,
+        _compile_patterns,
+        _mask_sensitive_data,
+    )
+
+    compiled_mask = _compile_patterns([r"(?i)password"])
+
+    large_tuple = tuple(f"item_{i}" for i in range(300))
+
+    result = _mask_sensitive_data(large_tuple, compiled_mask)
+
+    assert result == CODE_VARIABLES_TOO_LONG_VALUE

posthog/test/test_feature_flag_result.py

@@ -4,7 +4,13 @@ import mock
 
 from posthog.client import Client
 from posthog.test.test_utils import FAKE_TEST_API_KEY
-from posthog.types import FeatureFlag, FeatureFlagResult, FlagMetadata, FlagReason
+from posthog.types import (
+    FeatureFlag,
+    FeatureFlagError,
+    FeatureFlagResult,
+    FlagMetadata,
+    FlagReason,
+)
 
 
 class TestFeatureFlagResult(unittest.TestCase):
@@ -189,7 +195,6 @@ class TestGetFeatureFlagResult(unittest.TestCase):
 
     def set_fail(self, e, batch):
         """Mark the failure handler"""
-        print("FAIL", e, batch)  # noqa: T201
         self.failed = True
 
     def setUp(self):
@@ -241,6 +246,9 @@ class TestGetFeatureFlagResult(unittest.TestCase):
             groups={},
             disable_geoip=None,
         )
+        # Verify error property is NOT present on successful evaluation
+        captured_properties = patch_capture.call_args[1]["properties"]
+        self.assertNotIn("$feature_flag_error", captured_properties)
 
     @mock.patch.object(Client, "capture")
     def test_get_feature_flag_result_variant_local_evaluation(self, patch_capture):
@@ -295,6 +303,9 @@ class TestGetFeatureFlagResult(unittest.TestCase):
             groups={},
             disable_geoip=None,
         )
+        # Verify error property is NOT present on successful evaluation
+        captured_properties = patch_capture.call_args[1]["properties"]
+        self.assertNotIn("$feature_flag_error", captured_properties)
 
         another_flag_result = self.client.get_feature_flag_result(
             "person-flag", "another-distinct-id", person_properties={"region": "USA"}
@@ -360,6 +371,9 @@ class TestGetFeatureFlagResult(unittest.TestCase):
             groups={},
             disable_geoip=None,
         )
+        # Verify error property is NOT present on successful evaluation
+        captured_properties = patch_capture.call_args[1]["properties"]
+        self.assertNotIn("$feature_flag_error", captured_properties)
 
     @mock.patch("posthog.client.flags")
     @mock.patch.object(Client, "capture")
@@ -403,6 +417,9 @@ class TestGetFeatureFlagResult(unittest.TestCase):
             groups={},
             disable_geoip=None,
         )
+        # Verify error property is NOT present on successful evaluation
+        captured_properties = patch_capture.call_args[1]["properties"]
+        self.assertNotIn("$feature_flag_error", captured_properties)
 
     @mock.patch("posthog.client.flags")
     @mock.patch.object(Client, "capture")
@@ -438,6 +455,428 @@ class TestGetFeatureFlagResult(unittest.TestCase):
                 "$feature_flag_response": None,
                 "locally_evaluated": False,
                 "$feature/no-person-flag": None,
+                "$feature_flag_error": FeatureFlagError.FLAG_MISSING,
+            },
+            groups={},
+            disable_geoip=None,
+        )
+
+    @mock.patch("posthog.client.flags")
+    @mock.patch.object(Client, "capture")
+    def test_get_feature_flag_result_with_errors_while_computing_flags(
+        self, patch_capture, patch_flags
+    ):
+        """Test that errors_while_computing_flags is included in the $feature_flag_called event.
+
+        When the server returns errorsWhileComputingFlags=true, it indicates that there
+        was an error computing one or more flags. We include this in the event so users
+        can identify and debug flag evaluation issues.
+        """
+        patch_flags.return_value = {
+            "flags": {
+                "my-flag": {
+                    "key": "my-flag",
+                    "enabled": True,
+                    "variant": None,
+                    "reason": {"description": "Matched condition set 1"},
+                    "metadata": {"id": 1, "version": 1, "payload": None},
+                },
+            },
+            "requestId": "test-request-id-789",
+            "errorsWhileComputingFlags": True,
+        }
+
+        flag_result = self.client.get_feature_flag_result("my-flag", "some-distinct-id")
+
+        self.assertEqual(flag_result.enabled, True)
+        patch_capture.assert_called_with(
+            "$feature_flag_called",
+            distinct_id="some-distinct-id",
+            properties={
+                "$feature_flag": "my-flag",
+                "$feature_flag_response": True,
+                "locally_evaluated": False,
+                "$feature/my-flag": True,
+                "$feature_flag_request_id": "test-request-id-789",
+                "$feature_flag_reason": "Matched condition set 1",
+                "$feature_flag_id": 1,
+                "$feature_flag_version": 1,
+                "$feature_flag_error": FeatureFlagError.ERRORS_WHILE_COMPUTING,
+            },
+            groups={},
+            disable_geoip=None,
+        )
+
+    @mock.patch("posthog.client.flags")
+    @mock.patch.object(Client, "capture")
+    def test_get_feature_flag_result_flag_not_in_response(
+        self, patch_capture, patch_flags
+    ):
+        """Test that when a flag is not in the API response, we capture flag_missing error.
+
+        This happens when a flag doesn't exist or the user doesn't match any conditions.
+        """
+        patch_flags.return_value = {
+            "flags": {
+                "other-flag": {
+                    "key": "other-flag",
+                    "enabled": True,
+                    "variant": None,
+                    "reason": {"description": "Matched condition set 1"},
+                    "metadata": {"id": 1, "version": 1, "payload": None},
+                },
+            },
+            "requestId": "test-request-id-456",
+        }
+
+        flag_result = self.client.get_feature_flag_result(
+            "missing-flag", "some-distinct-id"
+        )
+
+        self.assertIsNone(flag_result)
+        patch_capture.assert_called_with(
+            "$feature_flag_called",
+            distinct_id="some-distinct-id",
+            properties={
+                "$feature_flag": "missing-flag",
+                "$feature_flag_response": None,
+                "locally_evaluated": False,
+                "$feature/missing-flag": None,
+                "$feature_flag_request_id": "test-request-id-456",
+                "$feature_flag_error": FeatureFlagError.FLAG_MISSING,
+            },
+            groups={},
+            disable_geoip=None,
+        )
+
+    @mock.patch("posthog.client.flags")
+    @mock.patch.object(Client, "capture")
+    def test_get_feature_flag_result_errors_computing_and_flag_missing(
+        self, patch_capture, patch_flags
+    ):
+        """Test that both errors are reported when errorsWhileComputingFlags=true AND flag is missing.
+
+        This can happen when the server encounters errors computing flags AND the requested
+        flag is not in the response. Both conditions should be reported for debugging.
+        """
+        patch_flags.return_value = {
+            "flags": {},  # Flag is missing
+            "requestId": "test-request-id-999",
+            "errorsWhileComputingFlags": True,  # But errors also occurred
+        }
+
+        flag_result = self.client.get_feature_flag_result(
+            "missing-flag", "some-distinct-id"
+        )
+
+        self.assertIsNone(flag_result)
+        patch_capture.assert_called_with(
+            "$feature_flag_called",
+            distinct_id="some-distinct-id",
+            properties={
+                "$feature_flag": "missing-flag",
+                "$feature_flag_response": None,
+                "locally_evaluated": False,
+                "$feature/missing-flag": None,
+                "$feature_flag_request_id": "test-request-id-999",
+                "$feature_flag_error": f"{FeatureFlagError.ERRORS_WHILE_COMPUTING},{FeatureFlagError.FLAG_MISSING}",
+            },
+            groups={},
+            disable_geoip=None,
+        )
+
+    @mock.patch("posthog.client.flags")
+    @mock.patch.object(Client, "capture")
+    def test_get_feature_flag_result_unknown_error(self, patch_capture, patch_flags):
+        """Test that unexpected exceptions are captured as unknown_error."""
+        patch_flags.side_effect = Exception("Unexpected error")
+
+        flag_result = self.client.get_feature_flag_result("my-flag", "some-distinct-id")
+
+        self.assertIsNone(flag_result)
+        patch_capture.assert_called_with(
+            "$feature_flag_called",
+            distinct_id="some-distinct-id",
+            properties={
+                "$feature_flag": "my-flag",
+                "$feature_flag_response": None,
+                "locally_evaluated": False,
+                "$feature/my-flag": None,
+                "$feature_flag_error": FeatureFlagError.UNKNOWN_ERROR,
+            },
+            groups={},
+            disable_geoip=None,
+        )
+
+    @mock.patch("posthog.client.flags")
+    @mock.patch.object(Client, "capture")
+    def test_get_feature_flag_result_timeout_error(self, patch_capture, patch_flags):
+        """Test that timeout errors are captured specifically."""
+        from posthog.request import RequestsTimeout
+
+        patch_flags.side_effect = RequestsTimeout("Request timed out")
+
+        flag_result = self.client.get_feature_flag_result("my-flag", "some-distinct-id")
+
+        self.assertIsNone(flag_result)
+        patch_capture.assert_called_with(
+            "$feature_flag_called",
+            distinct_id="some-distinct-id",
+            properties={
+                "$feature_flag": "my-flag",
+                "$feature_flag_response": None,
+                "locally_evaluated": False,
+                "$feature/my-flag": None,
+                "$feature_flag_error": FeatureFlagError.TIMEOUT,
+            },
+            groups={},
+            disable_geoip=None,
+        )
+
+    @mock.patch("posthog.client.flags")
+    @mock.patch.object(Client, "capture")
+    def test_get_feature_flag_result_connection_error(self, patch_capture, patch_flags):
+        """Test that connection errors are captured specifically."""
+        from posthog.request import RequestsConnectionError
+
+        patch_flags.side_effect = RequestsConnectionError("Connection refused")
+
+        flag_result = self.client.get_feature_flag_result("my-flag", "some-distinct-id")
+
+        self.assertIsNone(flag_result)
+        patch_capture.assert_called_with(
+            "$feature_flag_called",
+            distinct_id="some-distinct-id",
+            properties={
+                "$feature_flag": "my-flag",
+                "$feature_flag_response": None,
+                "locally_evaluated": False,
+                "$feature/my-flag": None,
+                "$feature_flag_error": FeatureFlagError.CONNECTION_ERROR,
+            },
+            groups={},
+            disable_geoip=None,
+        )
+
+    @mock.patch("posthog.client.flags")
+    @mock.patch.object(Client, "capture")
+    def test_get_feature_flag_result_api_error(self, patch_capture, patch_flags):
+        """Test that API errors include the status code."""
+        from posthog.request import APIError
+
+        patch_flags.side_effect = APIError(500, "Internal server error")
+
+        flag_result = self.client.get_feature_flag_result("my-flag", "some-distinct-id")
+
+        self.assertIsNone(flag_result)
+        patch_capture.assert_called_with(
+            "$feature_flag_called",
+            distinct_id="some-distinct-id",
+            properties={
+                "$feature_flag": "my-flag",
+                "$feature_flag_response": None,
+                "locally_evaluated": False,
+                "$feature/my-flag": None,
+                "$feature_flag_error": FeatureFlagError.api_error(500),
+            },
+            groups={},
+            disable_geoip=None,
+        )
+
+    @mock.patch("posthog.client.flags")
+    @mock.patch.object(Client, "capture")
+    def test_get_feature_flag_result_quota_limited(self, patch_capture, patch_flags):
+        """Test that quota limit errors are captured specifically."""
+        from posthog.request import QuotaLimitError
+
+        patch_flags.side_effect = QuotaLimitError(429, "Rate limit exceeded")
+
+        flag_result = self.client.get_feature_flag_result("my-flag", "some-distinct-id")
+
+        self.assertIsNone(flag_result)
+        patch_capture.assert_called_with(
+            "$feature_flag_called",
+            distinct_id="some-distinct-id",
+            properties={
+                "$feature_flag": "my-flag",
+                "$feature_flag_response": None,
+                "locally_evaluated": False,
+                "$feature/my-flag": None,
+                "$feature_flag_error": FeatureFlagError.QUOTA_LIMITED,
+            },
+            groups={},
+            disable_geoip=None,
+        )
+
+
+class TestFeatureFlagErrorWithStaleCacheFallback(unittest.TestCase):
+    """Tests for stale cache fallback behavior when flag evaluation fails.
+
+    When the PostHog API is unavailable (timeout, connection error, etc.), the SDK
+    falls back to stale cached flag values if available. These tests verify that:
+    1. The stale cached value is returned when an error occurs
+    2. The $feature_flag_error property is still set (for debugging)
+    3. The response reflects the cached value, not None
+    """
+
+    def set_fail(self, e, batch):
+        """Mark the failure handler"""
+        self.failed = True
+
+    def setUp(self):
+        self.failed = False
+        # Create client with memory-based flag cache enabled
+        self.client = Client(
+            FAKE_TEST_API_KEY,
+            on_error=self.set_fail,
+            flag_fallback_cache_url="memory://local/?ttl=300&size=10000",
+        )
+
+    def _populate_stale_cache(self, distinct_id, flag_key, flag_result):
+        """Pre-populate the flag cache with a value that will be used for stale fallback."""
+        self.client.flag_cache.set_cached_flag(
+            distinct_id,
+            flag_key,
+            flag_result,
+            flag_definition_version=self.client.flag_definition_version,
+        )
+
+    @mock.patch("posthog.client.flags")
+    @mock.patch.object(Client, "capture")
+    def test_timeout_error_returns_stale_cached_value(self, patch_capture, patch_flags):
+        """Test that timeout errors return stale cached value when available."""
+        from posthog.request import RequestsTimeout
+
+        # Pre-populate cache with a flag result
+        cached_result = FeatureFlagResult.from_value_and_payload(
+            "my-flag", "cached-variant", '{"from": "cache"}'
+        )
+        self._populate_stale_cache("some-distinct-id", "my-flag", cached_result)
+
+        # Simulate timeout error
+        patch_flags.side_effect = RequestsTimeout("Request timed out")
+
+        flag_result = self.client.get_feature_flag_result("my-flag", "some-distinct-id")
+
+        # Should return the stale cached value
+        self.assertIsNotNone(flag_result)
+        self.assertEqual(flag_result.variant, "cached-variant")
+        self.assertEqual(flag_result.payload, {"from": "cache"})
+
+        # Error should still be tracked for debugging
+        patch_capture.assert_called_with(
+            "$feature_flag_called",
+            distinct_id="some-distinct-id",
+            properties={
+                "$feature_flag": "my-flag",
+                "$feature_flag_response": "cached-variant",
+                "locally_evaluated": False,
+                "$feature/my-flag": "cached-variant",
+                "$feature_flag_payload": {"from": "cache"},
+                "$feature_flag_error": FeatureFlagError.TIMEOUT,
+            },
+            groups={},
+            disable_geoip=None,
+        )
+
+    @mock.patch("posthog.client.flags")
+    @mock.patch.object(Client, "capture")
+    def test_connection_error_returns_stale_cached_value(
+        self, patch_capture, patch_flags
+    ):
+        """Test that connection errors return stale cached value when available."""
+        from posthog.request import RequestsConnectionError
+
+        # Pre-populate cache with a boolean flag result
+        cached_result = FeatureFlagResult.from_value_and_payload("my-flag", True, None)
+        self._populate_stale_cache("some-distinct-id", "my-flag", cached_result)
+
+        # Simulate connection error
+        patch_flags.side_effect = RequestsConnectionError("Connection refused")
+
+        flag_result = self.client.get_feature_flag_result("my-flag", "some-distinct-id")
+
+        # Should return the stale cached value
+        self.assertIsNotNone(flag_result)
+        self.assertEqual(flag_result.enabled, True)
+        self.assertIsNone(flag_result.variant)
+
+        # Error should still be tracked
+        patch_capture.assert_called_with(
+            "$feature_flag_called",
+            distinct_id="some-distinct-id",
+            properties={
+                "$feature_flag": "my-flag",
+                "$feature_flag_response": True,
+                "locally_evaluated": False,
+                "$feature/my-flag": True,
+                "$feature_flag_error": FeatureFlagError.CONNECTION_ERROR,
+            },
+            groups={},
+            disable_geoip=None,
+        )
+
+    @mock.patch("posthog.client.flags")
+    @mock.patch.object(Client, "capture")
+    def test_api_error_returns_stale_cached_value(self, patch_capture, patch_flags):
+        """Test that API errors return stale cached value when available."""
+        from posthog.request import APIError
+
+        # Pre-populate cache
+        cached_result = FeatureFlagResult.from_value_and_payload(
+            "my-flag", "control", None
+        )
+        self._populate_stale_cache("some-distinct-id", "my-flag", cached_result)
+
+        # Simulate API error
+        patch_flags.side_effect = APIError(503, "Service unavailable")
+
+        flag_result = self.client.get_feature_flag_result("my-flag", "some-distinct-id")
+
+        # Should return the stale cached value
+        self.assertIsNotNone(flag_result)
+        self.assertEqual(flag_result.variant, "control")
+
+        # Error should still be tracked with status code
+        patch_capture.assert_called_with(
+            "$feature_flag_called",
+            distinct_id="some-distinct-id",
+            properties={
+                "$feature_flag": "my-flag",
+                "$feature_flag_response": "control",
+                "locally_evaluated": False,
+                "$feature/my-flag": "control",
+                "$feature_flag_error": FeatureFlagError.api_error(503),
+            },
+            groups={},
+            disable_geoip=None,
+        )
+
+    @mock.patch("posthog.client.flags")
+    @mock.patch.object(Client, "capture")
+    def test_error_without_cache_returns_none(self, patch_capture, patch_flags):
+        """Test that errors return None when no stale cache is available."""
+        from posthog.request import RequestsTimeout
+
+        # Do NOT populate cache - no fallback available
+
+        patch_flags.side_effect = RequestsTimeout("Request timed out")
+
+        flag_result = self.client.get_feature_flag_result("my-flag", "some-distinct-id")
+
+        # Should return None since no cache available
+        self.assertIsNone(flag_result)
+
+        # Error should still be tracked
+        patch_capture.assert_called_with(
+            "$feature_flag_called",
+            distinct_id="some-distinct-id",
+            properties={
+                "$feature_flag": "my-flag",
+                "$feature_flag_response": None,
+                "locally_evaluated": False,
+                "$feature/my-flag": None,
+                "$feature_flag_error": FeatureFlagError.TIMEOUT,
             },
             groups={},
             disable_geoip=None,

posthog/test/test_feature_flags.py

@@ -11,7 +11,7 @@ from posthog.feature_flags import (
     match_property,
     relative_date_parse_for_feature_flag_matching,
 )
-from posthog.request import APIError
+from posthog.request import APIError, GetResponse
 from posthog.test.test_utils import FAKE_TEST_API_KEY
 
 
@@ -233,6 +233,27 @@ class TestLocalEvaluation(unittest.TestCase):
 
         self.assertEqual(patch_flags.call_count, 1)
 
+    def test_group_flag_is_inconclusive_when_group_properties_missing(self):
+        feature_flag = {
+            "id": 1,
+            "name": "Group Flag Without Property Filters",
+            "key": "group-flag-no-props",
+            "active": True,
+            "filters": {
+                "aggregation_group_type_index": 0,
+                "groups": [{"properties": [], "rollout_percentage": 100}],
+            },
+        }
+        self.client.group_type_mapping = {"0": "company"}
+
+        with self.assertRaises(InconclusiveMatchError):
+            self.client._compute_flag_locally(
+                feature_flag,
+                "some-distinct-id",
+                groups={"company": "acme"},
+                group_properties={},
+            )
+
     @mock.patch("posthog.client.flags")
     @mock.patch("posthog.client.get")
     def test_flag_with_complex_definition(self, patch_get, patch_flags):
@@ -2348,23 +2369,27 @@ class TestLocalEvaluation(unittest.TestCase):
     @mock.patch("posthog.client.Poller")
     @mock.patch("posthog.client.get")
     def test_load_feature_flags(self, patch_get, patch_poll):
-        patch_get.return_value = {
-            "flags": [
-                {
-                    "id": 1,
-                    "name": "Beta Feature",
-                    "key": "beta-feature",
-                    "active": True,
-                },
-                {
-                    "id": 2,
-                    "name": "Alpha Feature",
-                    "key": "alpha-feature",
-                    "active": False,
-                },
-            ],
-            "group_type_mapping": {"0": "company"},
-        }
+        patch_get.return_value = GetResponse(
+            data={
+                "flags": [
+                    {
+                        "id": 1,
+                        "name": "Beta Feature",
+                        "key": "beta-feature",
+                        "active": True,
+                    },
+                    {
+                        "id": 2,
+                        "name": "Alpha Feature",
+                        "key": "alpha-feature",
+                        "active": False,
+                    },
+                ],
+                "group_type_mapping": {"0": "company"},
+                "cohorts": {},
+            },
+            etag='"abc123"',
+        )
         client = Client(FAKE_TEST_API_KEY, personal_api_key="test")
         with freeze_time("2020-01-01T12:01:00.0000Z"):
             client.load_feature_flags()
@@ -2375,8 +2400,144 @@ class TestLocalEvaluation(unittest.TestCase):
             client._last_feature_flag_poll.isoformat(), "2020-01-01T12:01:00+00:00"
         )
         self.assertEqual(patch_poll.call_count, 1)
+        # Verify ETag is stored
+        self.assertEqual(client._flags_etag, '"abc123"')
 
-    def test_load_feature_flags_wrong_key(self):
+    @mock.patch("posthog.client.Poller")
+    @mock.patch("posthog.client.get")
+    def test_load_feature_flags_sends_etag_on_subsequent_requests(
+        self, patch_get, patch_poll
+    ):
+        """Test that the ETag is sent in If-None-Match header on subsequent requests"""
+        patch_get.return_value = GetResponse(
+            data={
+                "flags": [{"id": 1, "key": "beta-feature", "active": True}],
+                "group_type_mapping": {},
+                "cohorts": {},
+            },
+            etag='"initial-etag"',
+        )
+        client = Client(FAKE_TEST_API_KEY, personal_api_key="test")
+        client.load_feature_flags()
+
+        # First call should have no etag
+        first_call_kwargs = patch_get.call_args_list[0][1]
+        self.assertIsNone(first_call_kwargs.get("etag"))
+
+        # Simulate second call
+        client._load_feature_flags()
+
+        # Second call should have the etag
+        second_call_kwargs = patch_get.call_args_list[1][1]
+        self.assertEqual(second_call_kwargs.get("etag"), '"initial-etag"')
+
+    @mock.patch("posthog.client.Poller")
+    @mock.patch("posthog.client.get")
+    def test_load_feature_flags_304_not_modified(self, patch_get, patch_poll):
+        """Test that 304 Not Modified responses skip flag processing"""
+        # First response with flags
+        initial_response = GetResponse(
+            data={
+                "flags": [{"id": 1, "key": "beta-feature", "active": True}],
+                "group_type_mapping": {"0": "company"},
+                "cohorts": {},
+            },
+            etag='"test-etag"',
+        )
+        # Second response is 304 Not Modified
+        not_modified_response = GetResponse(
+            data=None,
+            etag='"test-etag"',
+            not_modified=True,
+        )
+        patch_get.side_effect = [initial_response, not_modified_response]
+
+        client = Client(FAKE_TEST_API_KEY, personal_api_key="test")
+        client.load_feature_flags()
+
+        # Verify initial flags are loaded
+        self.assertEqual(len(client.feature_flags), 1)
+        self.assertEqual(client.feature_flags[0]["key"], "beta-feature")
+        self.assertEqual(client.group_type_mapping, {"0": "company"})
+
+        # Second call with 304
+        client._load_feature_flags()
+
+        # Flags should still be the same (not cleared)
+        self.assertEqual(len(client.feature_flags), 1)
+        self.assertEqual(client.feature_flags[0]["key"], "beta-feature")
+        self.assertEqual(client.group_type_mapping, {"0": "company"})
+
+    @mock.patch("posthog.client.Poller")
+    @mock.patch("posthog.client.get")
+    def test_load_feature_flags_etag_updated_on_new_response(
+        self, patch_get, patch_poll
+    ):
+        """Test that ETag is updated when flags change"""
+        patch_get.side_effect = [
+            GetResponse(
+                data={
+                    "flags": [{"id": 1, "key": "flag-v1", "active": True}],
+                    "group_type_mapping": {},
+                    "cohorts": {},
+                },
+                etag='"etag-v1"',
+            ),
+            GetResponse(
+                data={
+                    "flags": [{"id": 1, "key": "flag-v2", "active": True}],
+                    "group_type_mapping": {},
+                    "cohorts": {},
+                },
+                etag='"etag-v2"',
+            ),
+        ]
+
+        client = Client(FAKE_TEST_API_KEY, personal_api_key="test")
+        client.load_feature_flags()
+        self.assertEqual(client._flags_etag, '"etag-v1"')
+
+        client._load_feature_flags()
+        self.assertEqual(client._flags_etag, '"etag-v2"')
+        self.assertEqual(client.feature_flags[0]["key"], "flag-v2")
+
+    @mock.patch("posthog.client.Poller")
+    @mock.patch("posthog.client.get")
+    def test_load_feature_flags_clears_etag_when_server_stops_sending(
+        self, patch_get, patch_poll
+    ):
+        """Test that ETag is cleared when server stops sending it"""
+        patch_get.side_effect = [
+            GetResponse(
+                data={
+                    "flags": [{"id": 1, "key": "flag-v1", "active": True}],
+                    "group_type_mapping": {},
+                    "cohorts": {},
+                },
+                etag='"etag-v1"',
+            ),
+            GetResponse(
+                data={
+                    "flags": [{"id": 1, "key": "flag-v2", "active": True}],
+                    "group_type_mapping": {},
+                    "cohorts": {},
+                },
+                etag=None,  # Server stopped sending ETag
+            ),
+        ]
+
+        client = Client(FAKE_TEST_API_KEY, personal_api_key="test")
+        client.load_feature_flags()
+        self.assertEqual(client._flags_etag, '"etag-v1"')
+
+        client._load_feature_flags()
+        self.assertIsNone(client._flags_etag)
+        self.assertEqual(client.feature_flags[0]["key"], "flag-v2")
+
+    @mock.patch("posthog.client.Poller")
+    @mock.patch("posthog.client.get")
+    def test_load_feature_flags_wrong_key(self, patch_get, _patch_poll):
+        patch_get.side_effect = APIError(401, "Unauthorized")
         client = Client(FAKE_TEST_API_KEY, personal_api_key=FAKE_TEST_API_KEY)
 
         with self.assertLogs("posthog", level="ERROR") as logs:
@@ -2925,6 +3086,7 @@ class TestLocalEvaluation(unittest.TestCase):
                 "some-distinct-id",
                 match_value=True,
                 person_properties={"region": "USA"},
+                send_feature_flag_events=True,
             ),
             300,
         )
@@ -3082,6 +3244,541 @@ class TestLocalEvaluation(unittest.TestCase):
         # Verify API was called (fallback occurred)
         self.assertEqual(patch_flags.call_count, 1)
 
+    @mock.patch("posthog.client.flags")
+    def test_device_id_bucketing_uses_device_id_for_hash(self, patch_flags):
+        """
+        When a flag has bucketing_identifier: "device_id", the device_id should be
+        used for hashing instead of distinct_id.
+        """
+        client = Client(FAKE_TEST_API_KEY, personal_api_key=FAKE_TEST_API_KEY)
+
+        # This flag uses device_id for bucketing
+        client.feature_flags = [
+            {
+                "id": 1,
+                "key": "device-bucketed-flag",
+                "active": True,
+                "filters": {
+                    "bucketing_identifier": "device_id",
+                    "groups": [
+                        {
+                            "properties": [],
+                            "rollout_percentage": 100,
+                        }
+                    ],
+                },
+            }
+        ]
+
+        # Same distinct_id with different device_ids should produce different results
+        # (based on rollout percentage, we check consistency)
+        result1 = client.get_feature_flag(
+            "device-bucketed-flag", "user-123", device_id="device-A"
+        )
+        result2 = client.get_feature_flag(
+            "device-bucketed-flag", "user-123", device_id="device-A"
+        )
+
+        # Same device_id should give consistent results
+        self.assertEqual(result1, result2)
+
+        # No API fallback should occur
+        self.assertEqual(patch_flags.call_count, 0)
+
+    def test_match_feature_flag_properties_without_bucketing_value_is_deprecated(
+        self,
+    ):
+        """
+        match_feature_flag_properties should preserve backward compatibility when
+        bucketing_value is omitted, while warning about deprecation.
+        """
+        from posthog.feature_flags import match_feature_flag_properties
+
+        flag = {
+            "id": 1,
+            "key": "device-bucketed-flag",
+            "active": True,
+            "filters": {
+                "bucketing_identifier": "device_id",
+                "groups": [
+                    {
+                        "properties": [],
+                        "rollout_percentage": 100,
+                    }
+                ],
+            },
+        }
+
+        with self.assertWarnsRegex(
+            DeprecationWarning, "without bucketing_value is deprecated"
+        ):
+            result = match_feature_flag_properties(
+                flag,
+                "user-123",
+                {},
+                device_id="device-123",
+            )
+
+        self.assertTrue(result)
+
+    @mock.patch("posthog.client.flags")
+    def test_device_id_bucketing_same_device_different_users_same_result(
+        self, patch_flags
+    ):
+        """
+        When a flag uses device_id bucketing, different distinct_ids with the same
+        device_id should get the same result.
+        """
+        client = Client(FAKE_TEST_API_KEY, personal_api_key=FAKE_TEST_API_KEY)
+
+        client.feature_flags = [
+            {
+                "id": 1,
+                "key": "device-bucketed-flag",
+                "active": True,
+                "filters": {
+                    "bucketing_identifier": "device_id",
+                    "groups": [
+                        {
+                            "properties": [],
+                            "rollout_percentage": 50,
+                        }
+                    ],
+                },
+            }
+        ]
+
+        # Different distinct_ids with the same device_id should get the same result
+        result1 = client.get_feature_flag(
+            "device-bucketed-flag", "user-A", device_id="shared-device"
+        )
+        result2 = client.get_feature_flag(
+            "device-bucketed-flag", "user-B", device_id="shared-device"
+        )
+        result3 = client.get_feature_flag(
+            "device-bucketed-flag", "user-C", device_id="shared-device"
+        )
+
+        # All should be the same since device_id is the same
+        self.assertEqual(result1, result2)
+        self.assertEqual(result2, result3)
+
+        self.assertEqual(patch_flags.call_count, 0)
+
+    @mock.patch("posthog.client.flags")
+    def test_device_id_bucketing_fallback_when_device_id_missing(self, patch_flags):
+        """
+        When a flag requires device_id for bucketing but none is provided,
+        it should fallback to server evaluation.
+        """
+        patch_flags.return_value = {"featureFlags": {"device-bucketed-flag": True}}
+        client = Client(FAKE_TEST_API_KEY, personal_api_key=FAKE_TEST_API_KEY)
+
+        client.feature_flags = [
+            {
+                "id": 1,
+                "key": "device-bucketed-flag",
+                "active": True,
+                "filters": {
+                    "bucketing_identifier": "device_id",
+                    "groups": [
+                        {
+                            "properties": [],
+                            "rollout_percentage": 100,
+                        }
+                    ],
+                },
+            }
+        ]
+
+        # No device_id provided - should fallback to API
+        result = client.get_feature_flag("device-bucketed-flag", "user-123")
+
+        self.assertTrue(result)
+        # API should have been called
+        self.assertEqual(patch_flags.call_count, 1)
+
+    @mock.patch("posthog.client.flags")
+    def test_device_id_bucketing_returns_none_when_only_evaluate_locally_and_no_device_id(
+        self, patch_flags
+    ):
+        """
+        When only_evaluate_locally=True and device_id is required but missing,
+        should return None instead of falling back to API.
+        """
+        client = Client(FAKE_TEST_API_KEY, personal_api_key=FAKE_TEST_API_KEY)
+
+        client.feature_flags = [
+            {
+                "id": 1,
+                "key": "device-bucketed-flag",
+                "active": True,
+                "filters": {
+                    "bucketing_identifier": "device_id",
+                    "groups": [
+                        {
+                            "properties": [],
+                            "rollout_percentage": 100,
+                        }
+                    ],
+                },
+            }
+        ]
+
+        # No device_id + only_evaluate_locally should return None
+        result = client.get_feature_flag(
+            "device-bucketed-flag", "user-123", only_evaluate_locally=True
+        )
+
+        self.assertIsNone(result)
+        # API should NOT have been called
+        self.assertEqual(patch_flags.call_count, 0)
+
+    @mock.patch("posthog.client.flags")
+    def test_default_bucketing_identifier_uses_distinct_id(self, patch_flags):
+        """
+        When bucketing_identifier is not set or is 'distinct_id', should use
+        distinct_id for hashing (default behavior).
+        """
+        client = Client(FAKE_TEST_API_KEY, personal_api_key=FAKE_TEST_API_KEY)
+
+        # Flag without bucketing_identifier (defaults to distinct_id)
+        client.feature_flags = [
+            {
+                "id": 1,
+                "key": "normal-flag",
+                "active": True,
+                "filters": {
+                    "groups": [
+                        {
+                            "properties": [],
+                            "rollout_percentage": 50,
+                        }
+                    ],
+                },
+            }
+        ]
+
+        # Different distinct_ids should potentially produce different results
+        # but same distinct_id should produce same result
+        result1 = client.get_feature_flag("normal-flag", "user-A")
+        result2 = client.get_feature_flag("normal-flag", "user-A")
+
+        self.assertEqual(result1, result2)
+        self.assertEqual(patch_flags.call_count, 0)
+
+    @mock.patch("posthog.client.flags")
+    def test_device_id_bucketing_with_multivariate_flag(self, patch_flags):
+        """
+        Multivariate flag variant selection should use device_id when
+        bucketing_identifier is set to device_id.
+        """
+        client = Client(FAKE_TEST_API_KEY, personal_api_key=FAKE_TEST_API_KEY)
+
+        client.feature_flags = [
+            {
+                "id": 1,
+                "key": "multivariate-device-flag",
+                "active": True,
+                "filters": {
+                    "bucketing_identifier": "device_id",
+                    "groups": [
+                        {
+                            "properties": [],
+                            "rollout_percentage": 100,
+                        }
+                    ],
+                    "multivariate": {
+                        "variants": [
+                            {"key": "control", "rollout_percentage": 50},
+                            {"key": "test", "rollout_percentage": 50},
+                        ]
+                    },
+                },
+            }
+        ]
+
+        # Same device_id should give same variant
+        result1 = client.get_feature_flag(
+            "multivariate-device-flag", "user-A", device_id="device-1"
+        )
+        result2 = client.get_feature_flag(
+            "multivariate-device-flag", "user-B", device_id="device-1"
+        )
+
+        # Both should get the same variant because device_id is the same
+        self.assertEqual(result1, result2)
+        self.assertIn(result1, ["control", "test"])
+
+        self.assertEqual(patch_flags.call_count, 0)
+
+    @mock.patch("posthog.client.flags")
+    def test_device_id_bucketing_from_context(self, patch_flags):
+        """
+        When device_id is not passed as a parameter but is set in the context,
+        it should be resolved from context.
+        """
+        from posthog.contexts import new_context, set_context_device_id
+
+        client = Client(FAKE_TEST_API_KEY, personal_api_key=FAKE_TEST_API_KEY)
+
+        client.feature_flags = [
+            {
+                "id": 1,
+                "key": "device-bucketed-flag",
+                "active": True,
+                "filters": {
+                    "bucketing_identifier": "device_id",
+                    "groups": [
+                        {
+                            "properties": [],
+                            "rollout_percentage": 100,
+                        }
+                    ],
+                },
+            }
+        ]
+
+        # Set device_id in context
+        with new_context():
+            set_context_device_id("context-device-id")
+            result = client.get_feature_flag("device-bucketed-flag", "user-123")
+
+        # Should evaluate locally using the context device_id
+        self.assertTrue(result)
+        self.assertEqual(patch_flags.call_count, 0)
+
+    @mock.patch("posthog.client.flags")
+    def test_group_flags_ignore_bucketing_identifier(self, patch_flags):
+        """
+        Group flags should continue to use the group identifier for hashing,
+        regardless of the bucketing_identifier setting.
+        """
+        client = Client(FAKE_TEST_API_KEY, personal_api_key=FAKE_TEST_API_KEY)
+
+        client.feature_flags = [
+            {
+                "id": 1,
+                "key": "group-flag",
+                "active": True,
+                "filters": {
+                    "aggregation_group_type_index": 0,
+                    "bucketing_identifier": "device_id",  # Should be ignored for group flags
+                    "groups": [
+                        {
+                            "properties": [],
+                            "rollout_percentage": 100,
+                        }
+                    ],
+                },
+            }
+        ]
+        client.group_type_mapping = {"0": "company"}
+
+        # Even with bucketing_identifier set to device_id, group flag should use group identifier
+        result = client.get_feature_flag(
+            "group-flag",
+            "user-123",
+            groups={"company": "acme-inc"},
+            device_id="some-device",
+        )
+
+        self.assertTrue(result)
+        self.assertEqual(patch_flags.call_count, 0)
+
+    @mock.patch("posthog.client.flags")
+    def test_group_flag_dependency_receives_device_id(self, patch_flags):
+        """
+        Group flag dependency evaluation should receive device_id so dependent
+        device_id-bucketed flags can be evaluated locally.
+        """
+        patch_flags.return_value = {"featureFlags": {"group-parent-flag": "from-api"}}
+        client = Client(FAKE_TEST_API_KEY, personal_api_key=FAKE_TEST_API_KEY)
+
+        client.feature_flags = [
+            {
+                "id": 1,
+                "key": "device-dependent-flag",
+                "active": True,
+                "filters": {
+                    "bucketing_identifier": "device_id",
+                    "groups": [
+                        {
+                            "properties": [],
+                            "rollout_percentage": 100,
+                        }
+                    ],
+                },
+            },
+            {
+                "id": 2,
+                "key": "group-parent-flag",
+                "active": True,
+                "filters": {
+                    "aggregation_group_type_index": 0,
+                    "groups": [
+                        {
+                            "properties": [
+                                {
+                                    "key": "device-dependent-flag",
+                                    "operator": "flag_evaluates_to",
+                                    "value": True,
+                                    "type": "flag",
+                                    "dependency_chain": ["device-dependent-flag"],
+                                }
+                            ],
+                            "rollout_percentage": 100,
+                        }
+                    ],
+                },
+            },
+        ]
+        client.group_type_mapping = {"0": "company"}
+
+        result = client.get_feature_flag(
+            "group-parent-flag",
+            "user-123",
+            groups={"company": "acme-inc"},
+            device_id="device-123",
+        )
+
+        self.assertTrue(result)
+        self.assertEqual(patch_flags.call_count, 0)
+
+    @mock.patch("posthog.client.flags")
+    def test_group_flag_dependency_ignores_device_id_bucketing_identifier(
+        self, patch_flags
+    ):
+        """
+        Group flag dependencies should keep bucketing by group key, even when
+        the dependent group flag has bucketing_identifier set to device_id.
+        """
+        patch_flags.return_value = {"featureFlags": {"parent-group-flag": "from-api"}}
+        client = Client(FAKE_TEST_API_KEY, personal_api_key=FAKE_TEST_API_KEY)
+
+        client.feature_flags = [
+            {
+                "id": 1,
+                "key": "child-group-flag",
+                "active": True,
+                "filters": {
+                    "aggregation_group_type_index": 0,
+                    "bucketing_identifier": "device_id",
+                    "groups": [
+                        {
+                            "properties": [],
+                            "rollout_percentage": 100,
+                        }
+                    ],
+                },
+            },
+            {
+                "id": 2,
+                "key": "parent-group-flag",
+                "active": True,
+                "filters": {
+                    "aggregation_group_type_index": 0,
+                    "groups": [
+                        {
+                            "properties": [
+                                {
+                                    "key": "child-group-flag",
+                                    "operator": "flag_evaluates_to",
+                                    "value": True,
+                                    "type": "flag",
+                                    "dependency_chain": ["child-group-flag"],
+                                }
+                            ],
+                            "rollout_percentage": 100,
+                        }
+                    ],
+                },
+            },
+        ]
+        client.group_type_mapping = {"0": "company"}
+
+        result = client.get_feature_flag(
+            "parent-group-flag",
+            "user-123",
+            groups={"company": "acme-inc"},
+        )
+
+        self.assertTrue(result)
+        self.assertEqual(patch_flags.call_count, 0)
+
+    @mock.patch("posthog.client.flags")
+    def test_get_all_flags_with_device_id_bucketing(self, patch_flags):
+        """
+        get_all_flags_and_payloads should properly handle flags with device_id bucketing.
+        """
+        client = Client(FAKE_TEST_API_KEY, personal_api_key=FAKE_TEST_API_KEY)
+
+        client.feature_flags = [
+            {
+                "id": 1,
+                "key": "normal-flag",
+                "active": True,
+                "filters": {
+                    "groups": [{"properties": [], "rollout_percentage": 100}],
+                },
+            },
+            {
+                "id": 2,
+                "key": "device-flag",
+                "active": True,
+                "filters": {
+                    "bucketing_identifier": "device_id",
+                    "groups": [{"properties": [], "rollout_percentage": 100}],
+                },
+            },
+        ]
+
+        # With device_id provided, both flags should be evaluated locally
+        result = client.get_all_flags("user-123", device_id="my-device")
+
+        self.assertEqual(result["normal-flag"], True)
+        self.assertEqual(result["device-flag"], True)
+        self.assertEqual(patch_flags.call_count, 0)
+
+    @mock.patch("posthog.client.flags")
+    def test_get_all_flags_fallback_when_device_id_missing_for_some_flags(
+        self, patch_flags
+    ):
+        """
+        When some flags require device_id but it's not provided, those flags
+        should trigger fallback while others can be evaluated locally.
+        """
+        patch_flags.return_value = {
+            "featureFlags": {"normal-flag": True, "device-flag": "from-api"}
+        }
+        client = Client(FAKE_TEST_API_KEY, personal_api_key=FAKE_TEST_API_KEY)
+
+        client.feature_flags = [
+            {
+                "id": 1,
+                "key": "normal-flag",
+                "active": True,
+                "filters": {
+                    "groups": [{"properties": [], "rollout_percentage": 100}],
+                },
+            },
+            {
+                "id": 2,
+                "key": "device-flag",
+                "active": True,
+                "filters": {
+                    "bucketing_identifier": "device_id",
+                    "groups": [{"properties": [], "rollout_percentage": 100}],
+                },
+            },
+        ]
+
+        # Without device_id, device-flag can't be evaluated locally
+        client.get_all_flags("user-123")
+
+        # Should fallback to API for all flags when any can't be evaluated locally
+        self.assertEqual(patch_flags.call_count, 1)
+
 
 class TestMatchProperties(unittest.TestCase):
     def property(self, key, value, operator=None):
@@ -3859,6 +4556,7 @@ class TestCaptureCalls(unittest.TestCase):
                 },
             },
             "requestId": "18043bf7-9cf6-44cd-b959-9662ee20d371",
+            "evaluatedAt": 1234567890,
         }
         client = Client(FAKE_TEST_API_KEY)
 
@@ -3878,6 +4576,7 @@ class TestCaptureCalls(unittest.TestCase):
                 "$feature_flag_id": 23,
                 "$feature_flag_version": 42,
                 "$feature_flag_request_id": "18043bf7-9cf6-44cd-b959-9662ee20d371",
+                "$feature_flag_evaluated_at": 1234567890,
             },
             groups={},
             disable_geoip=None,
@@ -3912,7 +4611,9 @@ class TestCaptureCalls(unittest.TestCase):
 
         self.assertEqual(
             client.get_feature_flag_payload(
-                "decide-flag-with-payload", "some-distinct-id"
+                "decide-flag-with-payload",
+                "some-distinct-id",
+                send_feature_flag_events=True,
             ),
             {"foo": "bar"},
         )
@@ -3988,9 +4689,10 @@ class TestCaptureCalls(unittest.TestCase):
 
     @mock.patch.object(Client, "capture")
     @mock.patch("posthog.client.flags")
-    def test_capture_is_called_in_get_feature_flag_payload(
+    def test_get_feature_flag_payload_does_not_send_feature_flag_called_events(
         self, patch_flags, patch_capture
     ):
+        """Test that get_feature_flag_payload does NOT send $feature_flag_called events"""
         patch_flags.return_value = {
             "featureFlags": {"person-flag": True},
             "featureFlagPayloads": {"person-flag": 300},
@@ -4012,68 +4714,18 @@ class TestCaptureCalls(unittest.TestCase):
                             "rollout_percentage": 100,
                         }
                     ],
+                    "payloads": {"true": '"payload"'},
                 },
             }
         ]
 
-        # Call get_feature_flag_payload with match_value=None to trigger get_feature_flag
-        client.get_feature_flag_payload(
+        payload = client.get_feature_flag_payload(
             key="person-flag",
             distinct_id="some-distinct-id",
             person_properties={"region": "USA", "name": "Aloha"},
         )
-
-        # Assert that capture was called once, with the correct parameters
-        self.assertEqual(patch_capture.call_count, 1)
-        patch_capture.assert_called_with(
-            "$feature_flag_called",
-            distinct_id="some-distinct-id",
-            properties={
-                "$feature_flag": "person-flag",
-                "$feature_flag_response": True,
-                "locally_evaluated": True,
-                "$feature/person-flag": True,
-            },
-            groups={},
-            disable_geoip=None,
-        )
-
-        # Reset mocks for further tests
-        patch_capture.reset_mock()
-        patch_flags.reset_mock()
-
-        # Call get_feature_flag_payload again for the same user; capture should not be called again because we've already reported an event for this distinct_id + flag
-        client.get_feature_flag_payload(
-            key="person-flag",
-            distinct_id="some-distinct-id",
-            person_properties={"region": "USA", "name": "Aloha"},
-        )
-
+        self.assertIsNotNone(payload)
         self.assertEqual(patch_capture.call_count, 0)
-        patch_capture.reset_mock()
-
-        # Call get_feature_flag_payload for a different user; capture should be called
-        client.get_feature_flag_payload(
-            key="person-flag",
-            distinct_id="some-distinct-id2",
-            person_properties={"region": "USA", "name": "Aloha"},
-        )
-
-        self.assertEqual(patch_capture.call_count, 1)
-        patch_capture.assert_called_with(
-            "$feature_flag_called",
-            distinct_id="some-distinct-id2",
-            properties={
-                "$feature_flag": "person-flag",
-                "$feature_flag_response": True,
-                "locally_evaluated": True,
-                "$feature/person-flag": True,
-            },
-            groups={},
-            disable_geoip=None,
-        )
-
-        patch_capture.reset_mock()
 
     @mock.patch("posthog.client.flags")
     def test_fallback_to_api_in_get_feature_flag_payload_when_flag_has_static_cohort(
@@ -4173,13 +4825,15 @@ class TestCaptureCalls(unittest.TestCase):
             disable_geoip=False,
         )
 
-    @mock.patch("posthog.client.MAX_DICT_SIZE", 100)
     @mock.patch.object(Client, "capture")
     @mock.patch("posthog.client.flags")
     def test_capture_multiple_users_doesnt_out_of_memory(
         self, patch_flags, patch_capture
     ):
         client = Client(FAKE_TEST_API_KEY, personal_api_key=FAKE_TEST_API_KEY)
+        # Set on the instance to avoid relying on module-constant patching behavior
+        # across Python/runtime implementations.
+        client.distinct_ids_feature_flags_reported.max_size = 100
         client.feature_flags = [
             {
                 "id": 1,

posthog/test/test_flag_definition_cache.py

@@ -0,0 +1,612 @@
+"""
+Tests for FlagDefinitionCacheProvider functionality.
+
+These tests follow the patterns from the TypeScript implementation in posthog-js/packages/node.
+"""
+
+import threading
+import unittest
+from typing import Optional
+from unittest import mock
+
+from posthog.client import Client
+from posthog.flag_definition_cache import (
+    FlagDefinitionCacheData,
+    FlagDefinitionCacheProvider,
+)
+from posthog.request import GetResponse
+from posthog.test.test_utils import FAKE_TEST_API_KEY
+
+
+class MockCacheProvider:
+    """A mock implementation of FlagDefinitionCacheProvider for testing."""
+
+    def __init__(self):
+        self.stored_data: Optional[FlagDefinitionCacheData] = None
+        self.should_fetch_return_value = True
+        self.get_call_count = 0
+        self.should_fetch_call_count = 0
+        self.on_received_call_count = 0
+        self.shutdown_call_count = 0
+        self.should_fetch_error: Optional[Exception] = None
+        self.get_error: Optional[Exception] = None
+        self.on_received_error: Optional[Exception] = None
+        self.shutdown_error: Optional[Exception] = None
+
+    def get_flag_definitions(self) -> Optional[FlagDefinitionCacheData]:
+        self.get_call_count += 1
+        if self.get_error:
+            raise self.get_error
+        return self.stored_data
+
+    def should_fetch_flag_definitions(self) -> bool:
+        self.should_fetch_call_count += 1
+        if self.should_fetch_error:
+            raise self.should_fetch_error
+        return self.should_fetch_return_value
+
+    def on_flag_definitions_received(self, data: FlagDefinitionCacheData) -> None:
+        self.on_received_call_count += 1
+        if self.on_received_error:
+            raise self.on_received_error
+        self.stored_data = data
+
+    def shutdown(self) -> None:
+        self.shutdown_call_count += 1
+        if self.shutdown_error:
+            raise self.shutdown_error
+
+
+class TestFlagDefinitionCacheProvider(unittest.TestCase):
+    """Tests for the FlagDefinitionCacheProvider protocol."""
+
+    @classmethod
+    def setUpClass(cls):
+        # Prevent real HTTP requests
+        cls.client_post_patcher = mock.patch("posthog.client.batch_post")
+        cls.consumer_post_patcher = mock.patch("posthog.consumer.batch_post")
+        cls.client_post_patcher.start()
+        cls.consumer_post_patcher.start()
+
+    @classmethod
+    def tearDownClass(cls):
+        cls.client_post_patcher.stop()
+        cls.consumer_post_patcher.stop()
+
+    def setUp(self):
+        self.cache_provider = MockCacheProvider()
+        self.sample_flags_data: FlagDefinitionCacheData = {
+            "flags": [
+                {"key": "test-flag", "active": True, "filters": {}},
+                {"key": "another-flag", "active": False, "filters": {}},
+            ],
+            "group_type_mapping": {"0": "company", "1": "project"},
+            "cohorts": {"1": {"properties": []}},
+        }
+
+    def tearDown(self):
+        # Ensure client cleanup
+        pass
+
+    def _create_client_with_cache(self) -> Client:
+        """Create a client with the mock cache provider."""
+        return Client(
+            FAKE_TEST_API_KEY,
+            personal_api_key="test-personal-key",
+            flag_definition_cache_provider=self.cache_provider,
+            sync_mode=True,
+            enable_local_evaluation=False,  # Disable poller for tests
+        )
+
+
+class TestCacheInitialization(TestFlagDefinitionCacheProvider):
+    """Tests for cache initialization behavior."""
+
+    @mock.patch("posthog.client.get")
+    def test_uses_cached_data_when_should_fetch_returns_false(self, mock_get):
+        """When should_fetch returns False and cache has data, use cached data."""
+        self.cache_provider.should_fetch_return_value = False
+        self.cache_provider.stored_data = self.sample_flags_data
+
+        client = self._create_client_with_cache()
+        client._load_feature_flags()
+
+        # Should not call API
+        mock_get.assert_not_called()
+
+        # Should have called cache methods
+        self.assertEqual(self.cache_provider.should_fetch_call_count, 1)
+        self.assertEqual(self.cache_provider.get_call_count, 1)
+
+        # Flags should be loaded from cache
+        self.assertEqual(len(client.feature_flags), 2)
+        self.assertEqual(client.feature_flags[0]["key"], "test-flag")
+
+        client.join()
+
+    @mock.patch("posthog.client.get")
+    def test_fetches_from_api_when_should_fetch_returns_true(self, mock_get):
+        """When should_fetch returns True, fetch from API."""
+        self.cache_provider.should_fetch_return_value = True
+
+        mock_get.return_value = GetResponse(
+            data=self.sample_flags_data, etag="test-etag", not_modified=False
+        )
+
+        client = self._create_client_with_cache()
+        client._load_feature_flags()
+
+        # Should call API
+        mock_get.assert_called_once()
+
+        # Should have called should_fetch but not get
+        self.assertEqual(self.cache_provider.should_fetch_call_count, 1)
+        self.assertEqual(self.cache_provider.get_call_count, 0)
+
+        # Should have called on_received to store in cache
+        self.assertEqual(self.cache_provider.on_received_call_count, 1)
+
+        client.join()
+
+    @mock.patch("posthog.client.get")
+    def test_emergency_fallback_when_cache_empty_and_no_flags(self, mock_get):
+        """When should_fetch=False but cache is empty and no flags loaded, fetch anyway."""
+        self.cache_provider.should_fetch_return_value = False
+        self.cache_provider.stored_data = None  # Empty cache
+
+        mock_get.return_value = GetResponse(
+            data=self.sample_flags_data, etag="test-etag", not_modified=False
+        )
+
+        client = self._create_client_with_cache()
+        client._load_feature_flags()
+
+        # Should call API due to emergency fallback
+        mock_get.assert_called_once()
+
+        # Should have called on_received
+        self.assertEqual(self.cache_provider.on_received_call_count, 1)
+
+        client.join()
+
+    @mock.patch("posthog.client.get")
+    def test_preserves_existing_flags_when_cache_returns_none(self, mock_get):
+        """When cache returns None but client has flags, preserve existing flags."""
+        self.cache_provider.should_fetch_return_value = False
+        self.cache_provider.stored_data = None  # Empty cache
+
+        client = self._create_client_with_cache()
+
+        # Pre-load flags (simulating a previous successful fetch)
+        client.feature_flags = self.sample_flags_data["flags"]
+        client.group_type_mapping = self.sample_flags_data["group_type_mapping"]
+        client.cohorts = self.sample_flags_data["cohorts"]
+
+        client._load_feature_flags()
+
+        # Should NOT call API since we already have flags
+        mock_get.assert_not_called()
+
+        # Existing flags should be preserved
+        self.assertEqual(len(client.feature_flags), 2)
+        self.assertEqual(client.feature_flags[0]["key"], "test-flag")
+
+        client.join()
+
+
+class TestFetchCoordination(TestFlagDefinitionCacheProvider):
+    """Tests for fetch coordination between workers."""
+
+    @mock.patch("posthog.client.get")
+    def test_calls_should_fetch_before_each_poll(self, mock_get):
+        """should_fetch_flag_definitions is called before each poll cycle."""
+        self.cache_provider.should_fetch_return_value = True
+
+        mock_get.return_value = GetResponse(
+            data=self.sample_flags_data, etag="test-etag", not_modified=False
+        )
+
+        client = self._create_client_with_cache()
+
+        # First poll
+        client._load_feature_flags()
+        self.assertEqual(self.cache_provider.should_fetch_call_count, 1)
+
+        # Second poll
+        client._load_feature_flags()
+        self.assertEqual(self.cache_provider.should_fetch_call_count, 2)
+
+        client.join()
+
+    @mock.patch("posthog.client.get")
+    def test_does_not_call_on_received_when_fetch_skipped(self, mock_get):
+        """on_flag_definitions_received is NOT called when fetch is skipped."""
+        self.cache_provider.should_fetch_return_value = False
+        self.cache_provider.stored_data = self.sample_flags_data
+
+        client = self._create_client_with_cache()
+        client._load_feature_flags()
+
+        # Should not call on_received since we didn't fetch
+        self.assertEqual(self.cache_provider.on_received_call_count, 0)
+
+        client.join()
+
+    @mock.patch("posthog.client.get")
+    def test_stores_data_in_cache_after_api_fetch(self, mock_get):
+        """on_flag_definitions_received receives the fetched data."""
+        self.cache_provider.should_fetch_return_value = True
+
+        mock_get.return_value = GetResponse(
+            data=self.sample_flags_data, etag="test-etag", not_modified=False
+        )
+
+        client = self._create_client_with_cache()
+        client._load_feature_flags()
+
+        # Should have stored data in cache
+        self.assertEqual(self.cache_provider.on_received_call_count, 1)
+        self.assertIsNotNone(self.cache_provider.stored_data)
+        self.assertEqual(len(self.cache_provider.stored_data["flags"]), 2)
+
+        client.join()
+
+    @mock.patch("posthog.client.get")
+    def test_304_not_modified_does_not_update_cache(self, mock_get):
+        """When API returns 304 Not Modified, cache should not be updated."""
+        self.cache_provider.should_fetch_return_value = True
+
+        # First fetch to populate flags and ETag
+        mock_get.return_value = GetResponse(
+            data=self.sample_flags_data, etag="test-etag", not_modified=False
+        )
+
+        client = self._create_client_with_cache()
+        client._load_feature_flags()
+
+        # Verify initial fetch worked
+        self.assertEqual(self.cache_provider.on_received_call_count, 1)
+        self.assertEqual(len(client.feature_flags), 2)
+
+        # Second fetch returns 304 Not Modified
+        mock_get.return_value = GetResponse(
+            data=None, etag="test-etag", not_modified=True
+        )
+
+        client._load_feature_flags()
+
+        # API was called twice
+        self.assertEqual(mock_get.call_count, 2)
+
+        # should_fetch was called twice
+        self.assertEqual(self.cache_provider.should_fetch_call_count, 2)
+
+        # on_received should NOT be called again (304 = no new data)
+        self.assertEqual(self.cache_provider.on_received_call_count, 1)
+
+        # Flags should still be present
+        self.assertEqual(len(client.feature_flags), 2)
+
+        client.join()
+
+
+class TestErrorHandling(TestFlagDefinitionCacheProvider):
+    """Tests for error handling in cache provider operations."""
+
+    @mock.patch("posthog.client.get")
+    def test_should_fetch_error_defaults_to_fetching(self, mock_get):
+        """When should_fetch throws an error, default to fetching from API."""
+        self.cache_provider.should_fetch_error = Exception("Lock acquisition failed")
+
+        mock_get.return_value = GetResponse(
+            data=self.sample_flags_data, etag="test-etag", not_modified=False
+        )
+
+        client = self._create_client_with_cache()
+        client._load_feature_flags()
+
+        # Should still fetch from API
+        mock_get.assert_called_once()
+
+        # Flags should be loaded
+        self.assertEqual(len(client.feature_flags), 2)
+
+        client.join()
+
+    @mock.patch("posthog.client.get")
+    def test_get_error_falls_back_to_api_fetch(self, mock_get):
+        """When get_flag_definitions throws an error, fetch from API."""
+        self.cache_provider.should_fetch_return_value = False
+        self.cache_provider.get_error = Exception("Cache read failed")
+
+        mock_get.return_value = GetResponse(
+            data=self.sample_flags_data, etag="test-etag", not_modified=False
+        )
+
+        client = self._create_client_with_cache()
+        client._load_feature_flags()
+
+        # Should fall back to API
+        mock_get.assert_called_once()
+
+        client.join()
+
+    @mock.patch("posthog.client.get")
+    def test_on_received_error_keeps_flags_in_memory(self, mock_get):
+        """When on_flag_definitions_received throws, flags are still in memory."""
+        self.cache_provider.should_fetch_return_value = True
+        self.cache_provider.on_received_error = Exception("Cache write failed")
+
+        mock_get.return_value = GetResponse(
+            data=self.sample_flags_data, etag="test-etag", not_modified=False
+        )
+
+        client = self._create_client_with_cache()
+        client._load_feature_flags()
+
+        # Flags should still be loaded in memory despite cache error
+        self.assertEqual(len(client.feature_flags), 2)
+        self.assertEqual(client.feature_flags[0]["key"], "test-flag")
+
+        client.join()
+
+    @mock.patch("posthog.client.get")
+    def test_shutdown_error_is_logged_but_continues(self, mock_get):
+        """When shutdown throws an error, it's logged but shutdown continues."""
+        self.cache_provider.shutdown_error = Exception("Lock release failed")
+
+        mock_get.return_value = GetResponse(
+            data=self.sample_flags_data, etag="test-etag", not_modified=False
+        )
+
+        client = self._create_client_with_cache()
+        client._load_feature_flags()
+
+        # Should not raise when joining
+        client.join()
+
+        # Shutdown was called
+        self.assertEqual(self.cache_provider.shutdown_call_count, 1)
+
+
+class TestShutdownLifecycle(TestFlagDefinitionCacheProvider):
+    """Tests for shutdown lifecycle."""
+
+    @mock.patch("posthog.client.get")
+    def test_shutdown_calls_cache_provider_shutdown(self, mock_get):
+        """Client shutdown calls cache provider shutdown."""
+        mock_get.return_value = GetResponse(
+            data=self.sample_flags_data, etag="test-etag", not_modified=False
+        )
+
+        client = self._create_client_with_cache()
+        client._load_feature_flags()
+
+        # Shutdown
+        client.join()
+
+        self.assertEqual(self.cache_provider.shutdown_call_count, 1)
+
+    @mock.patch("posthog.client.get")
+    def test_shutdown_called_even_without_fetching(self, mock_get):
+        """Shutdown is called even when cache was used instead of fetching."""
+        self.cache_provider.should_fetch_return_value = False
+        self.cache_provider.stored_data = self.sample_flags_data
+
+        client = self._create_client_with_cache()
+        client._load_feature_flags()
+        client.join()
+
+        # Shutdown should still be called
+        self.assertEqual(self.cache_provider.shutdown_call_count, 1)
+
+    @mock.patch("posthog.client.get")
+    def test_multiple_join_calls_only_shutdown_once(self, mock_get):
+        """Calling join() multiple times should only call cache provider shutdown once."""
+        mock_get.return_value = GetResponse(
+            data=self.sample_flags_data, etag="test-etag", not_modified=False
+        )
+
+        client = self._create_client_with_cache()
+        client._load_feature_flags()
+
+        # Call join multiple times
+        client.join()
+        client.join()
+        client.join()
+
+        # Shutdown should be called each time (current behavior - no guard)
+        # This test documents the current behavior
+        self.assertGreaterEqual(self.cache_provider.shutdown_call_count, 1)
+
+
+class TestBackwardCompatibility(TestFlagDefinitionCacheProvider):
+    """Tests for backward compatibility without cache provider."""
+
+    @mock.patch("posthog.client.get")
+    def test_works_without_cache_provider(self, mock_get):
+        """Client works normally without a cache provider configured."""
+        mock_get.return_value = GetResponse(
+            data=self.sample_flags_data, etag="test-etag", not_modified=False
+        )
+
+        # Create client without cache provider
+        client = Client(
+            FAKE_TEST_API_KEY,
+            personal_api_key="test-personal-key",
+            sync_mode=True,
+            enable_local_evaluation=False,
+        )
+        client._load_feature_flags()
+
+        # Should fetch from API
+        mock_get.assert_called_once()
+
+        # Flags should be loaded
+        self.assertEqual(len(client.feature_flags), 2)
+
+        client.join()
+
+
+class TestDataIntegrity(TestFlagDefinitionCacheProvider):
+    """Tests for data integrity between cache and client state."""
+
+    @mock.patch("posthog.client.get")
+    def test_cached_flags_available_for_evaluation(self, mock_get):
+        """Flags loaded from cache are available for local evaluation."""
+        self.cache_provider.should_fetch_return_value = False
+        self.cache_provider.stored_data = {
+            "flags": [
+                {
+                    "key": "test-flag",
+                    "active": True,
+                    "filters": {
+                        "groups": [
+                            {
+                                "properties": [],
+                                "rollout_percentage": 100,
+                            }
+                        ]
+                    },
+                }
+            ],
+            "group_type_mapping": {},
+            "cohorts": {},
+        }
+
+        client = self._create_client_with_cache()
+        client._load_feature_flags()
+
+        # Flag should be accessible
+        self.assertEqual(len(client.feature_flags), 1)
+        self.assertEqual(client.feature_flags_by_key["test-flag"]["key"], "test-flag")
+
+        client.join()
+
+    @mock.patch("posthog.client.get")
+    def test_group_type_mapping_loaded_from_cache(self, mock_get):
+        """Group type mapping is correctly loaded from cache."""
+        self.cache_provider.should_fetch_return_value = False
+        self.cache_provider.stored_data = self.sample_flags_data
+
+        client = self._create_client_with_cache()
+        client._load_feature_flags()
+
+        self.assertEqual(client.group_type_mapping["0"], "company")
+        self.assertEqual(client.group_type_mapping["1"], "project")
+
+        client.join()
+
+    @mock.patch("posthog.client.get")
+    def test_cohorts_loaded_from_cache(self, mock_get):
+        """Cohorts are correctly loaded from cache."""
+        self.cache_provider.should_fetch_return_value = False
+        self.cache_provider.stored_data = self.sample_flags_data
+
+        client = self._create_client_with_cache()
+        client._load_feature_flags()
+
+        self.assertIn("1", client.cohorts)
+
+        client.join()
+
+    @mock.patch("posthog.client.get")
+    def test_cache_updated_when_api_returns_new_data(self, mock_get):
+        """State transition: cache has old data -> API returns new -> cache updated."""
+        # Start with old cached data
+        old_flags_data: FlagDefinitionCacheData = {
+            "flags": [{"key": "old-flag", "active": True, "filters": {}}],
+            "group_type_mapping": {},
+            "cohorts": {},
+        }
+        self.cache_provider.stored_data = old_flags_data
+        self.cache_provider.should_fetch_return_value = False
+
+        client = self._create_client_with_cache()
+
+        # First load from cache
+        client._load_feature_flags()
+        self.assertEqual(client.feature_flags[0]["key"], "old-flag")
+        self.assertEqual(self.cache_provider.on_received_call_count, 0)
+
+        # Now trigger API fetch with new data
+        self.cache_provider.should_fetch_return_value = True
+        new_flags_data: FlagDefinitionCacheData = {
+            "flags": [{"key": "new-flag", "active": True, "filters": {}}],
+            "group_type_mapping": {"0": "company"},
+            "cohorts": {"1": {"properties": []}},
+        }
+        mock_get.return_value = GetResponse(
+            data=new_flags_data, etag="new-etag", not_modified=False
+        )
+
+        client._load_feature_flags()
+
+        # Verify new flags loaded
+        self.assertEqual(client.feature_flags[0]["key"], "new-flag")
+        self.assertEqual(client.group_type_mapping["0"], "company")
+
+        # Verify cache was updated
+        self.assertEqual(self.cache_provider.on_received_call_count, 1)
+        self.assertEqual(self.cache_provider.stored_data["flags"][0]["key"], "new-flag")
+
+        client.join()
+
+
+class TestConcurrency(TestFlagDefinitionCacheProvider):
+    """Tests for thread safety and concurrent access."""
+
+    @mock.patch("posthog.client.get")
+    def test_concurrent_load_feature_flags_is_thread_safe(self, mock_get):
+        """Multiple threads calling _load_feature_flags should not cause errors."""
+        mock_get.return_value = GetResponse(
+            data=self.sample_flags_data, etag="test-etag", not_modified=False
+        )
+
+        client = self._create_client_with_cache()
+        errors = []
+
+        def load_flags():
+            try:
+                client._load_feature_flags()
+            except Exception as e:
+                errors.append(e)
+
+        # Launch 5 threads concurrently
+        threads = [threading.Thread(target=load_flags) for _ in range(5)]
+        for t in threads:
+            t.start()
+        for t in threads:
+            t.join()
+
+        # Should complete without errors
+        self.assertEqual(len(errors), 0, f"Unexpected errors: {errors}")
+
+        # Flags should be loaded
+        self.assertIsNotNone(client.feature_flags)
+        self.assertEqual(len(client.feature_flags), 2)
+
+        client.join()
+
+
+class TestProtocolCompliance(unittest.TestCase):
+    """Tests for Protocol compliance."""
+
+    def test_mock_provider_is_protocol_instance(self):
+        """MockCacheProvider satisfies FlagDefinitionCacheProvider protocol."""
+        provider = MockCacheProvider()
+        self.assertIsInstance(provider, FlagDefinitionCacheProvider)
+
+    def test_incomplete_provider_is_not_protocol_instance(self):
+        """Class missing methods is not a FlagDefinitionCacheProvider."""
+
+        class IncompleteProvider:
+            def get_flag_definitions(self):
+                return None
+
+        provider = IncompleteProvider()
+        self.assertNotIsInstance(provider, FlagDefinitionCacheProvider)
+
+
+if __name__ == "__main__":
+    unittest.main()

posthog/test/test_request.py

@@ -6,16 +6,60 @@ import mock
 import pytest
 import requests
 
+import posthog.request as request_module
 from posthog.request import (
+    APIError,
     DatetimeSerializer,
+    GetResponse,
+    KEEP_ALIVE_SOCKET_OPTIONS,
     QuotaLimitError,
+    _mask_tokens_in_url,
     batch_post,
     decide,
     determine_server_host,
+    disable_connection_reuse,
+    enable_keep_alive,
+    flags,
+    get,
+    set_socket_options,
 )
 from posthog.test.test_utils import TEST_API_KEY
 
 
+@pytest.mark.parametrize(
+    "url, expected",
+    [
+        # Token with params after - masks keeping first 10 chars
+        (
+            "https://example.com/api/flags?token=phc_abc123xyz789&send_cohorts",
+            "https://example.com/api/flags?token=phc_abc123...&send_cohorts",
+        ),
+        # Token at end of URL
+        (
+            "https://example.com/api/flags?token=phc_abc123xyz789",
+            "https://example.com/api/flags?token=phc_abc123...",
+        ),
+        # No token - unchanged
+        (
+            "https://example.com/api/flags?other=value",
+            "https://example.com/api/flags?other=value",
+        ),
+        # Short token (<10 chars) - unchanged
+        (
+            "https://example.com/api/flags?token=short",
+            "https://example.com/api/flags?token=short",
+        ),
+        # Exactly 10 char token - gets ellipsis
+        (
+            "https://example.com/api/flags?token=1234567890",
+            "https://example.com/api/flags?token=1234567890...",
+        ),
+    ],
+)
+def test_mask_tokens_in_url(url, expected):
+    assert _mask_tokens_in_url(url) == expected
+
+
 class TestRequests(unittest.TestCase):
     def test_valid_request(self):
         res = batch_post(
@@ -107,6 +151,184 @@ class TestRequests(unittest.TestCase):
             self.assertEqual(response["featureFlags"], {"flag1": True})
 
 
+class TestGet(unittest.TestCase):
+    """Unit tests for the get() function HTTP-level behavior."""
+
+    @mock.patch("posthog.request._session.get")
+    def test_get_returns_data_and_etag(self, mock_get):
+        """Test that get() returns GetResponse with data and etag from headers."""
+        mock_response = requests.Response()
+        mock_response.status_code = 200
+        mock_response.headers["ETag"] = '"abc123"'
+        mock_response._content = json.dumps({"flags": [{"key": "test-flag"}]}).encode(
+            "utf-8"
+        )
+        mock_get.return_value = mock_response
+
+        response = get("api_key", "/test-url", host="https://example.com")
+
+        self.assertIsInstance(response, GetResponse)
+        self.assertEqual(response.data, {"flags": [{"key": "test-flag"}]})
+        self.assertEqual(response.etag, '"abc123"')
+        self.assertFalse(response.not_modified)
+
+    @mock.patch("posthog.request._session.get")
+    def test_get_sends_if_none_match_header_when_etag_provided(self, mock_get):
+        """Test that If-None-Match header is sent when etag parameter is provided."""
+        mock_response = requests.Response()
+        mock_response.status_code = 200
+        mock_response.headers["ETag"] = '"new-etag"'
+        mock_response._content = json.dumps({"flags": []}).encode("utf-8")
+        mock_get.return_value = mock_response
+
+        get("api_key", "/test-url", host="https://example.com", etag='"previous-etag"')
+
+        call_kwargs = mock_get.call_args[1]
+        self.assertEqual(call_kwargs["headers"]["If-None-Match"], '"previous-etag"')
+
+    @mock.patch("posthog.request._session.get")
+    def test_get_does_not_send_if_none_match_when_no_etag(self, mock_get):
+        """Test that If-None-Match header is not sent when no etag provided."""
+        mock_response = requests.Response()
+        mock_response.status_code = 200
+        mock_response._content = json.dumps({"flags": []}).encode("utf-8")
+        mock_get.return_value = mock_response
+
+        get("api_key", "/test-url", host="https://example.com")
+
+        call_kwargs = mock_get.call_args[1]
+        self.assertNotIn("If-None-Match", call_kwargs["headers"])
+
+    @mock.patch("posthog.request._session.get")
+    def test_get_handles_304_not_modified(self, mock_get):
+        """Test that 304 Not Modified response returns not_modified=True with no data."""
+        mock_response = requests.Response()
+        mock_response.status_code = 304
+        mock_response.headers["ETag"] = '"unchanged-etag"'
+        mock_get.return_value = mock_response
+
+        response = get(
+            "api_key", "/test-url", host="https://example.com", etag='"unchanged-etag"'
+        )
+
+        self.assertIsInstance(response, GetResponse)
+        self.assertIsNone(response.data)
+        self.assertEqual(response.etag, '"unchanged-etag"')
+        self.assertTrue(response.not_modified)
+
+    @mock.patch("posthog.request._session.get")
+    def test_get_304_without_etag_header_uses_request_etag(self, mock_get):
+        """Test that 304 response without ETag header falls back to request etag."""
+        mock_response = requests.Response()
+        mock_response.status_code = 304
+        # Server doesn't return ETag header on 304
+        mock_get.return_value = mock_response
+
+        response = get(
+            "api_key", "/test-url", host="https://example.com", etag='"original-etag"'
+        )
+
+        self.assertTrue(response.not_modified)
+        self.assertEqual(response.etag, '"original-etag"')
+
+    @mock.patch("posthog.request._session.get")
+    def test_get_200_without_etag_header(self, mock_get):
+        """Test that 200 response without ETag header returns None for etag."""
+        mock_response = requests.Response()
+        mock_response.status_code = 200
+        mock_response._content = json.dumps({"flags": []}).encode("utf-8")
+        # No ETag header
+        mock_get.return_value = mock_response
+
+        response = get("api_key", "/test-url", host="https://example.com")
+
+        self.assertFalse(response.not_modified)
+        self.assertIsNone(response.etag)
+        self.assertEqual(response.data, {"flags": []})
+
+    @mock.patch("posthog.request._session.get")
+    def test_get_error_response_raises_api_error(self, mock_get):
+        """Test that error responses raise APIError."""
+        mock_response = requests.Response()
+        mock_response.status_code = 401
+        mock_response._content = json.dumps({"detail": "Unauthorized"}).encode("utf-8")
+        mock_get.return_value = mock_response
+
+        with self.assertRaises(APIError) as ctx:
+            get("bad_key", "/test-url", host="https://example.com")
+
+        self.assertEqual(ctx.exception.status, 401)
+        self.assertEqual(ctx.exception.message, "Unauthorized")
+
+    @mock.patch("posthog.request._session.get")
+    def test_get_sends_authorization_header(self, mock_get):
+        """Test that Authorization header is sent with Bearer token."""
+        mock_response = requests.Response()
+        mock_response.status_code = 200
+        mock_response._content = json.dumps({}).encode("utf-8")
+        mock_get.return_value = mock_response
+
+        get("my-api-key", "/test-url", host="https://example.com")
+
+        call_kwargs = mock_get.call_args[1]
+        self.assertEqual(call_kwargs["headers"]["Authorization"], "Bearer my-api-key")
+
+    @mock.patch("posthog.request._session.get")
+    def test_get_sends_user_agent_header(self, mock_get):
+        """Test that User-Agent header is sent."""
+        mock_response = requests.Response()
+        mock_response.status_code = 200
+        mock_response._content = json.dumps({}).encode("utf-8")
+        mock_get.return_value = mock_response
+
+        get("api_key", "/test-url", host="https://example.com")
+
+        call_kwargs = mock_get.call_args[1]
+        self.assertIn("User-Agent", call_kwargs["headers"])
+        self.assertTrue(
+            call_kwargs["headers"]["User-Agent"].startswith("posthog-python/")
+        )
+
+    @mock.patch("posthog.request._session.get")
+    def test_get_passes_timeout(self, mock_get):
+        """Test that timeout parameter is passed to the request."""
+        mock_response = requests.Response()
+        mock_response.status_code = 200
+        mock_response._content = json.dumps({}).encode("utf-8")
+        mock_get.return_value = mock_response
+
+        get("api_key", "/test-url", host="https://example.com", timeout=30)
+
+        call_kwargs = mock_get.call_args[1]
+        self.assertEqual(call_kwargs["timeout"], 30)
+
+    @mock.patch("posthog.request._session.get")
+    def test_get_constructs_full_url(self, mock_get):
+        """Test that host and url are combined correctly."""
+        mock_response = requests.Response()
+        mock_response.status_code = 200
+        mock_response._content = json.dumps({}).encode("utf-8")
+        mock_get.return_value = mock_response
+
+        get("api_key", "/api/flags", host="https://example.com")
+
+        call_args = mock_get.call_args[0]
+        self.assertEqual(call_args[0], "https://example.com/api/flags")
+
+    @mock.patch("posthog.request._session.get")
+    def test_get_removes_trailing_slash_from_host(self, mock_get):
+        """Test that trailing slash is removed from host."""
+        mock_response = requests.Response()
+        mock_response.status_code = 200
+        mock_response._content = json.dumps({}).encode("utf-8")
+        mock_get.return_value = mock_response
+
+        get("api_key", "/api/flags", host="https://example.com/")
+
+        call_args = mock_get.call_args[0]
+        self.assertEqual(call_args[0], "https://example.com/api/flags")
+
+
 @pytest.mark.parametrize(
     "host, expected",
     [
@@ -128,3 +350,317 @@ class TestRequests(unittest.TestCase):
 )
 def test_routing_to_custom_host(host, expected):
     assert determine_server_host(host) == expected
+
+
+def test_enable_keep_alive_sets_socket_options():
+    try:
+        enable_keep_alive()
+        from posthog.request import _session
+
+        adapter = _session.get_adapter("https://example.com")
+        assert adapter.socket_options == KEEP_ALIVE_SOCKET_OPTIONS
+    finally:
+        set_socket_options(None)
+
+
+def test_set_socket_options_clears_with_none():
+    try:
+        enable_keep_alive()
+        set_socket_options(None)
+        from posthog.request import _session
+
+        adapter = _session.get_adapter("https://example.com")
+        assert adapter.socket_options is None
+    finally:
+        set_socket_options(None)
+
+
+def test_disable_connection_reuse_creates_fresh_sessions():
+    try:
+        disable_connection_reuse()
+        session1 = request_module._get_session()
+        session2 = request_module._get_session()
+        assert session1 is not session2
+    finally:
+        request_module._pooling_enabled = True
+
+
+def test_set_socket_options_is_idempotent():
+    try:
+        enable_keep_alive()
+        session1 = request_module._session
+        enable_keep_alive()
+        session2 = request_module._session
+        assert session1 is session2
+    finally:
+        set_socket_options(None)
+
+
+class TestFlagsSession(unittest.TestCase):
+    """Tests for flags session configuration."""
+
+    def test_retry_status_forcelist_excludes_rate_limits(self):
+        """Verify 429 (rate limit) is NOT retried - need to wait, not hammer."""
+        from posthog.request import RETRY_STATUS_FORCELIST
+
+        self.assertNotIn(429, RETRY_STATUS_FORCELIST)
+
+    def test_retry_status_forcelist_excludes_quota_errors(self):
+        """Verify 402 (payment required/quota) is NOT retried - won't resolve."""
+        from posthog.request import RETRY_STATUS_FORCELIST
+
+        self.assertNotIn(402, RETRY_STATUS_FORCELIST)
+
+    @mock.patch("posthog.request._get_flags_session")
+    def test_flags_uses_flags_session(self, mock_get_flags_session):
+        """flags() uses the dedicated flags session, not the general session."""
+        mock_response = requests.Response()
+        mock_response.status_code = 200
+        mock_response._content = json.dumps(
+            {
+                "featureFlags": {"test-flag": True},
+                "featureFlagPayloads": {},
+                "errorsWhileComputingFlags": False,
+            }
+        ).encode("utf-8")
+
+        mock_session = mock.MagicMock()
+        mock_session.post.return_value = mock_response
+        mock_get_flags_session.return_value = mock_session
+
+        result = flags("test-key", "https://test.posthog.com", distinct_id="user123")
+
+        self.assertEqual(result["featureFlags"]["test-flag"], True)
+        mock_get_flags_session.assert_called_once()
+        mock_session.post.assert_called_once()
+
+    @mock.patch("posthog.request._get_flags_session")
+    def test_flags_no_retry_on_quota_limit(self, mock_get_flags_session):
+        """flags() raises QuotaLimitError without retrying (at application level)."""
+        mock_response = requests.Response()
+        mock_response.status_code = 200
+        mock_response._content = json.dumps(
+            {
+                "quotaLimited": ["feature_flags"],
+                "featureFlags": {},
+                "featureFlagPayloads": {},
+                "errorsWhileComputingFlags": False,
+            }
+        ).encode("utf-8")
+
+        mock_session = mock.MagicMock()
+        mock_session.post.return_value = mock_response
+        mock_get_flags_session.return_value = mock_session
+
+        with self.assertRaises(QuotaLimitError):
+            flags("test-key", "https://test.posthog.com", distinct_id="user123")
+
+        # QuotaLimitError is raised after response is received, not retried
+        self.assertEqual(mock_session.post.call_count, 1)
+
+
+class TestFlagsSessionNetworkRetries(unittest.TestCase):
+    """Tests for network failure retries in the flags session."""
+
+    def test_flags_session_retry_config_includes_connection_errors(self):
+        """
+        Verify that the flags session is configured to retry on connection errors.
+
+        The urllib3 Retry adapter with connect=2 and read=2 automatically
+        retries on network-level failures (DNS failures, connection refused,
+        connection reset, etc.) up to 2 times each.
+        """
+        from posthog.request import _build_flags_session
+
+        session = _build_flags_session()
+
+        # Get the adapter for https://
+        adapter = session.get_adapter("https://test.posthog.com")
+
+        # Verify retry configuration
+        retry = adapter.max_retries
+        self.assertEqual(retry.total, 2, "Should have 2 total retries")
+        self.assertEqual(retry.connect, 2, "Should retry connection errors twice")
+        self.assertEqual(retry.read, 2, "Should retry read errors twice")
+        self.assertIn("POST", retry.allowed_methods, "Should allow POST retries")
+
+    def test_flags_session_retries_on_server_errors(self):
+        """
+        Verify that transient server errors (5xx) trigger retries.
+
+        This tests the status_forcelist configuration which specifies
+        which HTTP status codes should trigger a retry.
+        """
+        from posthog.request import _build_flags_session, RETRY_STATUS_FORCELIST
+
+        session = _build_flags_session()
+        adapter = session.get_adapter("https://test.posthog.com")
+        retry = adapter.max_retries
+
+        # Verify the status codes that trigger retries
+        self.assertEqual(
+            set(retry.status_forcelist),
+            set(RETRY_STATUS_FORCELIST),
+            "Should retry on transient server errors",
+        )
+
+        # Verify specific codes are included
+        self.assertIn(500, retry.status_forcelist)
+        self.assertIn(502, retry.status_forcelist)
+        self.assertIn(503, retry.status_forcelist)
+        self.assertIn(504, retry.status_forcelist)
+
+        # Verify rate limits and quota errors are NOT retried
+        self.assertNotIn(429, retry.status_forcelist)
+        self.assertNotIn(402, retry.status_forcelist)
+
+    def test_flags_session_has_backoff(self):
+        """
+        Verify that retries use exponential backoff to avoid thundering herd.
+        """
+        from posthog.request import _build_flags_session
+
+        session = _build_flags_session()
+        adapter = session.get_adapter("https://test.posthog.com")
+        retry = adapter.max_retries
+
+        self.assertEqual(
+            retry.backoff_factor,
+            0.5,
+            "Should use 0.5s backoff factor (0.5s, 1s delays)",
+        )
+
+
+class TestFlagsSessionRetryIntegration(unittest.TestCase):
+    """Integration tests that verify actual retry behavior with a local server."""
+
+    def test_retries_on_503_then_succeeds(self):
+        """
+        Verify that 503 errors trigger retries and eventually succeed.
+
+        Uses a local HTTP server that fails twice with 503, then succeeds.
+        This tests the full retry flow including backoff timing.
+        """
+        import threading
+        from http.server import HTTPServer, BaseHTTPRequestHandler
+        from socketserver import ThreadingMixIn
+        from urllib3.util.retry import Retry
+        from posthog.request import HTTPAdapterWithSocketOptions, RETRY_STATUS_FORCELIST
+
+        request_count = 0
+
+        class RetryTestHandler(BaseHTTPRequestHandler):
+            protocol_version = "HTTP/1.1"
+
+            def do_POST(self):
+                nonlocal request_count
+                request_count += 1
+
+                # Read and discard request body to prevent connection issues
+                content_length = int(self.headers.get("Content-Length", 0))
+                if content_length > 0:
+                    self.rfile.read(content_length)
+
+                if request_count <= 2:
+                    self.send_response(503)
+                    self.send_header("Content-Type", "application/json")
+                    body = b'{"error": "Service unavailable"}'
+                    self.send_header("Content-Length", str(len(body)))
+                    self.end_headers()
+                    self.wfile.write(body)
+                else:
+                    self.send_response(200)
+                    self.send_header("Content-Type", "application/json")
+                    body = (
+                        b'{"featureFlags": {"test": true}, "featureFlagPayloads": {}}'
+                    )
+                    self.send_header("Content-Length", str(len(body)))
+                    self.end_headers()
+                    self.wfile.write(body)
+
+            def log_message(self, format, *args):
+                pass  # Suppress logging
+
+        # Use ThreadingMixIn for cleaner shutdown
+        class ThreadedHTTPServer(ThreadingMixIn, HTTPServer):
+            daemon_threads = True
+
+        # Start server on a random available port
+        server = ThreadedHTTPServer(("127.0.0.1", 0), RetryTestHandler)
+        port = server.server_address[1]
+        server_thread = threading.Thread(target=server.serve_forever)
+        server_thread.daemon = True
+        server_thread.start()
+
+        try:
+            # Build session with same retry config as _build_flags_session
+            # but mounted on http:// for local testing
+            adapter = HTTPAdapterWithSocketOptions(
+                max_retries=Retry(
+                    total=2,
+                    connect=2,
+                    read=2,
+                    backoff_factor=0.01,  # Fast backoff for testing
+                    status_forcelist=RETRY_STATUS_FORCELIST,
+                    allowed_methods=["POST"],
+                ),
+            )
+            session = requests.Session()
+            session.mount("http://", adapter)
+
+            response = session.post(
+                f"http://127.0.0.1:{port}/flags/?v=2",
+                json={"distinct_id": "user123"},
+                timeout=5,
+            )
+
+            # Should succeed on 3rd attempt
+            self.assertEqual(response.status_code, 200)
+            self.assertEqual(request_count, 3)  # 1 initial + 2 retries
+        finally:
+            server.shutdown()
+            server.server_close()
+
+    def test_connection_errors_are_retried(self):
+        """
+        Verify that connection errors (no server) trigger retries.
+
+        Binds a socket to get a guaranteed available port, then closes it
+        so connection attempts fail with ConnectionError.
+        """
+        import socket
+        import time
+        from urllib3.util.retry import Retry
+        from posthog.request import HTTPAdapterWithSocketOptions, RETRY_STATUS_FORCELIST
+
+        # Get an available port by binding then closing a socket
+        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        sock.bind(("127.0.0.1", 0))
+        port = sock.getsockname()[1]
+        sock.close()  # Port is now available but nothing is listening
+
+        adapter = HTTPAdapterWithSocketOptions(
+            max_retries=Retry(
+                total=2,
+                connect=2,
+                read=2,
+                backoff_factor=0.05,  # Very fast for testing
+                status_forcelist=RETRY_STATUS_FORCELIST,
+                allowed_methods=["POST"],
+            ),
+        )
+        session = requests.Session()
+        session.mount("http://", adapter)
+
+        start = time.time()
+        with self.assertRaises(requests.exceptions.ConnectionError):
+            session.post(
+                f"http://127.0.0.1:{port}/flags/?v=2",
+                json={"distinct_id": "user123"},
+                timeout=1,
+            )
+        elapsed = time.time() - start
+
+        # With 3 attempts and backoff, should take more than instant
+        # but less than timeout (confirms retries happened)
+        self.assertGreater(elapsed, 0.05, "Should have some delay from retries")

posthog/test/test_utils.py

@@ -1,3 +1,4 @@
+import sys
 import time
 import unittest
 from dataclasses import dataclass
@@ -122,7 +123,9 @@ class TestUtils(unittest.TestCase):
             "bar": 2,
             "baz": None,
         }
-        assert utils.clean(ModelV1(foo=1, bar="2")) == {"foo": 1, "bar": "2"}
+        # Pydantic V1 is not compatible with Python 3.14+
+        if sys.version_info < (3, 14):
+            assert utils.clean(ModelV1(foo=1, bar="2")) == {"foo": 1, "bar": "2"}
         assert utils.clean(NestedModel(foo=ModelV2(foo="1", bar=2, baz="3"))) == {
             "foo": {"foo": "1", "bar": 2, "baz": "3"}
         }

posthog/types.py

@@ -123,6 +123,7 @@ class FlagsResponse(TypedDict, total=False):
     errorsWhileComputingFlags: bool
     requestId: str
     quotaLimit: Optional[List[str]]
+    evaluatedAt: Optional[int]
 
 
 class FlagsAndPayloads(TypedDict, total=True):
@@ -306,3 +307,42 @@ def to_payloads(response: FlagsResponse) -> Optional[dict[str, str]]:
         and value.enabled
         and value.metadata.payload is not None
     }
+
+
+class FeatureFlagError:
+    """Error type constants for the $feature_flag_error property.
+
+    These values are sent in analytics events to track flag evaluation failures.
+    They should not be changed without considering impact on existing dashboards
+    and queries that filter on these values.
+
+    Error values:
+        ERRORS_WHILE_COMPUTING: Server returned errorsWhileComputingFlags=true
+        FLAG_MISSING: Requested flag not in API response
+        QUOTA_LIMITED: Rate/quota limit exceeded
+        TIMEOUT: Request timed out
+        CONNECTION_ERROR: Network connectivity issue
+        UNKNOWN_ERROR: Unexpected exceptions
+
+    For API errors with status codes, use the api_error() method which returns
+    a string like "api_error_500".
+    """
+
+    ERRORS_WHILE_COMPUTING = "errors_while_computing_flags"
+    FLAG_MISSING = "flag_missing"
+    QUOTA_LIMITED = "quota_limited"
+    TIMEOUT = "timeout"
+    CONNECTION_ERROR = "connection_error"
+    UNKNOWN_ERROR = "unknown_error"
+
+    @staticmethod
+    def api_error(status: Union[int, str]) -> str:
+        """Generate API error string with status code.
+
+        Args:
+            status: HTTP status code from the API error
+
+        Returns:
+            Error string like "api_error_500"
+        """
+        return f"api_error_{status}"

posthog/version.py

@@ -1,4 +1 @@
-VERSION = "7.0.0"
-
-if __name__ == "__main__":
-    print(VERSION, end="")  # noqa: T201
+VERSION = "7.9.2"

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "posthog"
-dynamic = ["version"]
+version = "7.9.2"
 description = "Integrate PostHog into any python application."
 authors = [{ name = "PostHog", email = "hey@posthog.com" }]
 maintainers = [{ name = "PostHog", email = "hey@posthog.com" }]
@@ -21,6 +21,7 @@ classifiers = [
     "Programming Language :: Python :: 3.11",
     "Programming Language :: Python :: 3.12",
     "Programming Language :: Python :: 3.13",
+    "Programming Language :: Python :: 3.14",
 ]
 dependencies = [
     "requests>=2.7,<3.0",
@@ -83,15 +84,15 @@ packages = [
     "posthog.ai",
     "posthog.ai.langchain",
     "posthog.ai.openai",
+    "posthog.ai.openai_agents",
     "posthog.ai.anthropic",
     "posthog.ai.gemini",
     "posthog.test",
+    "posthog.test.ai",
+    "posthog.test.ai.openai_agents",
     "posthog.integrations",
 ]
 
-[tool.setuptools.dynamic]
-version = { attr = "posthog.version.VERSION" }
-
 [tool.pytest.ini_options]
 asyncio_mode = "auto"
 asyncio_default_fixture_loop_scope = "function"

references/posthog-python-references-latest.json

@@ -2,7 +2,7 @@
   "id": "posthog-python",
   "hogRef": "0.3",
   "info": {
-    "version": "6.9.3",
+    "version": "7.0.1",
     "id": "posthog-python",
     "title": "PostHog Python SDK",
     "description": "Integrate PostHog into any python application.",

sdk_compliance_adapter/Dockerfile

@@ -0,0 +1,22 @@
+FROM python:3.12-slim
+
+WORKDIR /app
+
+# Copy the SDK source code
+COPY posthog/ /app/sdk/posthog/
+COPY setup.py pyproject.toml README.md LICENSE /app/sdk/
+
+# Install the SDK from source
+RUN cd /app/sdk && pip install --no-cache-dir -e .
+
+# Install adapter dependencies
+RUN pip install --no-cache-dir flask python-dateutil
+
+# Copy adapter code
+COPY sdk_compliance_adapter/adapter.py /app/adapter.py
+
+# Expose port 8080
+EXPOSE 8080
+
+# Run the adapter
+CMD ["python", "/app/adapter.py"]

sdk_compliance_adapter/README.md

@@ -0,0 +1,76 @@
+# PostHog Python SDK Test Adapter
+
+This adapter wraps the posthog-python SDK for compliance testing with the [PostHog SDK Test Harness](https://github.com/PostHog/posthog-sdk-test-harness).
+
+## What is This?
+
+This is a simple Flask app that:
+1. Wraps the posthog-python SDK
+2. Exposes a REST API for the test harness to control
+3. Tracks internal SDK state for test assertions
+
+## Running Tests
+
+Tests run automatically in CI via GitHub Actions. See the test harness repo for details.
+
+### Locally with Docker Compose
+
+```bash
+# From the posthog-python/sdk_compliance_adapter directory
+docker-compose up --build --abort-on-container-exit
+```
+
+This will:
+1. Build the Python SDK adapter
+2. Pull the test harness image
+3. Run all compliance tests
+4. Show results
+
+### Manually with Docker
+
+```bash
+# Create network
+docker network create test-network
+
+# Build and run adapter
+docker build -f sdk_compliance_adapter/Dockerfile -t posthog-python-adapter .
+docker run -d --name sdk-adapter --network test-network -p 8080:8080 posthog-python-adapter
+
+# Run test harness
+docker run --rm \
+  --name test-harness \
+  --network test-network \
+  ghcr.io/posthog/sdk-test-harness:latest \
+  run --adapter-url http://sdk-adapter:8080 --mock-url http://test-harness:8081
+
+# Cleanup
+docker stop sdk-adapter && docker rm sdk-adapter
+docker network rm test-network
+```
+
+## Adapter Implementation
+
+See [adapter.py](adapter.py) for the implementation.
+
+The adapter implements the standard SDK adapter interface defined in the [test harness CONTRACT](https://github.com/PostHog/posthog-sdk-test-harness/blob/main/CONTRACT.yaml):
+
+- `GET /health` - Return SDK information
+- `POST /init` - Initialize SDK with config
+- `POST /capture` - Capture an event
+- `POST /flush` - Flush pending events
+- `GET /state` - Return internal state
+- `POST /reset` - Reset SDK state
+
+### Key Implementation Details
+
+**Request Tracking**: The adapter monkey-patches `batch_post` to track all HTTP requests made by the SDK, including retries.
+
+**State Management**: Thread-safe state tracking for events captured vs sent, retry attempts, and errors.
+
+**UUID Tracking**: Extracts and tracks UUIDs from batches to verify deduplication.
+
+## Documentation
+
+For complete documentation on the test harness and how to implement adapters, see:
+- [PostHog SDK Test Harness](https://github.com/PostHog/posthog-sdk-test-harness)
+- [Adapter Implementation Guide](https://github.com/PostHog/posthog-sdk-test-harness/blob/main/ADAPTER_GUIDE.md)

sdk_compliance_adapter/adapter.py

@@ -0,0 +1,383 @@
+"""
+PostHog Python SDK Test Adapter
+
+This adapter implements the SDK Test Adapter Interface defined in the PostHog Capture API Contract.
+It wraps the posthog-python SDK and exposes a REST API for the test harness to exercise.
+"""
+
+import logging
+import os
+import threading
+import time
+from typing import Any, Dict, List, Optional
+
+from flask import Flask, jsonify, request
+
+from posthog import Client
+from posthog.request import batch_post as original_batch_post
+from posthog.version import VERSION
+
+# Configure logging
+logging.basicConfig(
+    level=logging.DEBUG, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
+)
+logger = logging.getLogger(__name__)
+
+app = Flask(__name__)
+
+
+class RequestInfo:
+    """Information about an HTTP request made by the SDK"""
+
+    def __init__(
+        self,
+        timestamp_ms: int,
+        status_code: int,
+        retry_attempt: int,
+        event_count: int,
+        uuid_list: List[str],
+    ):
+        self.timestamp_ms = timestamp_ms
+        self.status_code = status_code
+        self.retry_attempt = retry_attempt
+        self.event_count = event_count
+        self.uuid_list = uuid_list
+
+    def to_dict(self) -> Dict[str, Any]:
+        return {
+            "timestamp_ms": self.timestamp_ms,
+            "status_code": self.status_code,
+            "retry_attempt": self.retry_attempt,
+            "event_count": self.event_count,
+            "uuid_list": self.uuid_list,
+        }
+
+
+class SDKState:
+    """Tracks SDK internal state for test assertions"""
+
+    def __init__(self):
+        self.lock = threading.Lock()
+        self.pending_events = 0
+        self.total_events_captured = 0
+        self.total_events_sent = 0
+        self.total_retries = 0
+        self.last_error: Optional[str] = None
+        self.requests_made: List[RequestInfo] = []
+        self.client: Optional[Client] = None
+        self.retry_attempts: Dict[str, int] = {}  # Track retry attempts by batch ID
+
+    def reset(self):
+        """Reset all state"""
+        with self.lock:
+            self.pending_events = 0
+            self.total_events_captured = 0
+            self.total_events_sent = 0
+            self.total_retries = 0
+            self.last_error = None
+            self.requests_made = []
+            self.retry_attempts = {}
+            if self.client:
+                # Flush and shutdown existing client
+                try:
+                    self.client.shutdown()
+                except Exception as e:
+                    logger.warning(f"Error shutting down client: {e}")
+                self.client = None
+
+    def increment_captured(self):
+        """Increment total events captured"""
+        with self.lock:
+            self.total_events_captured += 1
+            self.pending_events += 1
+
+    def record_request(self, status_code: int, batch: List[Dict], batch_id: str):
+        """Record an HTTP request made by the SDK"""
+        with self.lock:
+            # Determine retry attempt for this batch
+            retry_attempt = self.retry_attempts.get(batch_id, 0)
+
+            # Extract UUIDs from batch
+            uuid_list = [event.get("uuid", "") for event in batch]
+
+            request_info = RequestInfo(
+                timestamp_ms=int(time.time() * 1000),
+                status_code=status_code,
+                retry_attempt=retry_attempt,
+                event_count=len(batch),
+                uuid_list=uuid_list,
+            )
+            self.requests_made.append(request_info)
+
+            # Update counters
+            if status_code == 200:
+                # Success - clear pending events
+                self.total_events_sent += len(batch)
+                self.pending_events = max(0, self.pending_events - len(batch))
+                # Remove batch from retry tracking
+                self.retry_attempts.pop(batch_id, None)
+            else:
+                # Failure - increment retry count
+                self.retry_attempts[batch_id] = retry_attempt + 1
+                if retry_attempt > 0:
+                    self.total_retries += 1
+
+    def record_error(self, error: str):
+        """Record an error"""
+        with self.lock:
+            self.last_error = error
+
+    def get_state(self) -> Dict[str, Any]:
+        """Get current state as dict"""
+        with self.lock:
+            return {
+                "pending_events": self.pending_events,
+                "total_events_captured": self.total_events_captured,
+                "total_events_sent": self.total_events_sent,
+                "total_retries": self.total_retries,
+                "last_error": self.last_error,
+                "requests_made": [r.to_dict() for r in self.requests_made],
+            }
+
+
+# Global state
+state = SDKState()
+
+
+def create_batch_id(batch: List[Dict]) -> str:
+    """Create a unique ID for a batch based on UUIDs"""
+    uuids = sorted([event.get("uuid", "") for event in batch])
+    return "-".join(uuids[:3])  # Use first 3 UUIDs as batch ID
+
+
+def patched_batch_post(
+    api_key: str,
+    host: Optional[str] = None,
+    gzip: bool = False,
+    timeout: int = 15,
+    **kwargs,
+):
+    """Patched version of batch_post that tracks requests"""
+    batch = kwargs.get("batch", [])
+    batch_id = create_batch_id(batch)
+
+    try:
+        # Call original batch_post
+        response = original_batch_post(api_key, host, gzip, timeout, **kwargs)
+        # Record successful request
+        state.record_request(200, batch, batch_id)
+        return response
+    except Exception as e:
+        # Record failed request
+        status_code = (
+            getattr(e, "status_code", 500) if hasattr(e, "status_code") else 500
+        )
+        state.record_request(status_code, batch, batch_id)
+        state.record_error(str(e))
+        raise
+
+
+# Monkey-patch the batch_post function
+import posthog.request  # noqa: E402
+
+posthog.request.batch_post = patched_batch_post
+
+# Also patch in consumer module
+import posthog.consumer  # noqa: E402
+
+posthog.consumer.batch_post = patched_batch_post
+
+
+@app.route("/health", methods=["GET"])
+def health():
+    """Health check endpoint"""
+    return jsonify(
+        {
+            "sdk_name": "posthog-python",
+            "sdk_version": VERSION,
+            "adapter_version": "1.0.0",
+        }
+    )
+
+
+@app.route("/init", methods=["POST"])
+def init():
+    """Initialize the SDK client"""
+    try:
+        data = request.json or {}
+
+        # Reset state
+        state.reset()
+
+        # Extract config
+        api_key = data.get("api_key")
+        host = data.get("host")
+        flush_at = data.get("flush_at", 100)
+        flush_interval_ms = data.get("flush_interval_ms", 500)
+        max_retries = data.get("max_retries", 3)
+        enable_compression = data.get("enable_compression", False)
+
+        if not api_key:
+            return jsonify({"error": "api_key is required"}), 400
+        if not host:
+            return jsonify({"error": "host is required"}), 400
+
+        # Convert flush_interval from ms to seconds
+        flush_interval = flush_interval_ms / 1000.0
+
+        # Create client
+        client = Client(
+            project_api_key=api_key,
+            host=host,
+            flush_at=flush_at,
+            flush_interval=flush_interval,
+            gzip=enable_compression,
+            max_retries=max_retries,
+            debug=True,
+        )
+
+        state.client = client
+
+        logger.info(
+            f"Initialized SDK with api_key={api_key[:10]}..., host={host}, "
+            f"flush_at={flush_at}, flush_interval={flush_interval}, "
+            f"max_retries={max_retries}, gzip={enable_compression}"
+        )
+
+        return jsonify({"success": True})
+    except Exception as e:
+        logger.exception("Error initializing SDK")
+        return jsonify({"error": str(e)}), 500
+
+
+@app.route("/capture", methods=["POST"])
+def capture():
+    """Capture a single event"""
+    try:
+        if not state.client:
+            return jsonify({"error": "SDK not initialized"}), 400
+
+        data = request.json or {}
+
+        distinct_id = data.get("distinct_id")
+        event = data.get("event")
+        properties = data.get("properties")
+        timestamp = data.get("timestamp")
+
+        if not distinct_id:
+            return jsonify({"error": "distinct_id is required"}), 400
+        if not event:
+            return jsonify({"error": "event is required"}), 400
+
+        # Capture event
+        kwargs = {"distinct_id": distinct_id, "properties": properties}
+        if timestamp:
+            # Parse ISO8601 timestamp
+            from dateutil.parser import parse  # type: ignore[import-untyped]
+
+            kwargs["timestamp"] = parse(timestamp)
+
+        uuid = state.client.capture(event, **kwargs)
+
+        # Track that we captured an event
+        state.increment_captured()
+
+        logger.info(f"Captured event: {event} for {distinct_id}, uuid={uuid}")
+
+        return jsonify({"success": True, "uuid": uuid})
+    except Exception as e:
+        logger.exception("Error capturing event")
+        state.record_error(str(e))
+        return jsonify({"error": str(e)}), 500
+
+
+@app.route("/identify", methods=["POST"])
+def identify():
+    """Identify a user"""
+    try:
+        if not state.client:
+            return jsonify({"error": "SDK not initialized"}), 400
+
+        data = request.json or {}
+
+        distinct_id = data.get("distinct_id")
+        properties = data.get("properties")
+        properties_set_once = data.get("properties_set_once")
+
+        if not distinct_id:
+            return jsonify({"error": "distinct_id is required"}), 400
+
+        # Use the identify pattern - set + set_once
+        if properties:
+            state.client.set(distinct_id=distinct_id, properties=properties)
+            state.increment_captured()
+
+        if properties_set_once:
+            state.client.set_once(
+                distinct_id=distinct_id, properties=properties_set_once
+            )
+            state.increment_captured()
+
+        logger.info(f"Identified user: {distinct_id}")
+
+        return jsonify({"success": True})
+    except Exception as e:
+        logger.exception("Error identifying user")
+        state.record_error(str(e))
+        return jsonify({"error": str(e)}), 500
+
+
+@app.route("/flush", methods=["POST"])
+def flush():
+    """Force flush all pending events"""
+    try:
+        if not state.client:
+            return jsonify({"error": "SDK not initialized"}), 400
+
+        # Flush and wait
+        state.client.flush()
+
+        # Wait a bit for flush to complete
+        # The flush() method triggers queue.join() which blocks until all items are processed
+        time.sleep(0.5)
+
+        logger.info("Flushed pending events")
+
+        return jsonify({"success": True, "events_flushed": state.total_events_sent})
+    except Exception as e:
+        logger.exception("Error flushing events")
+        state.record_error(str(e))
+        return jsonify({"error": str(e), "errors": [str(e)]}, 500)
+
+
+@app.route("/state", methods=["GET"])
+def get_state():
+    """Get internal SDK state"""
+    try:
+        return jsonify(state.get_state())
+    except Exception as e:
+        logger.exception("Error getting state")
+        return jsonify({"error": str(e)}), 500
+
+
+@app.route("/reset", methods=["POST"])
+def reset():
+    """Reset SDK state"""
+    try:
+        state.reset()
+        logger.info("Reset SDK state")
+        return jsonify({"success": True})
+    except Exception as e:
+        logger.exception("Error resetting state")
+        return jsonify({"error": str(e)}), 500
+
+
+def main():
+    """Main entry point"""
+    port = int(os.environ.get("PORT", 8080))
+    logger.info(f"Starting SDK Test Adapter on port {port}")
+    app.run(host="0.0.0.0", port=port, debug=False)
+
+
+if __name__ == "__main__":
+    main()

sdk_compliance_adapter/docker-compose.yml

@@ -0,0 +1,25 @@
+version: "3.8"
+
+services:
+  # PostHog Python SDK adapter
+  sdk-adapter:
+    build:
+      context: ..
+      dockerfile: sdk_compliance_adapter/Dockerfile
+    ports:
+      - "8080:8080"
+    networks:
+      - test-network
+
+  # Test harness
+  test-harness:
+    image: ghcr.io/posthog/sdk-test-harness:latest
+    command: ["run", "--adapter-url", "http://sdk-adapter:8080", "--mock-url", "http://test-harness:8081"]
+    networks:
+      - test-network
+    depends_on:
+      - sdk-adapter
+
+networks:
+  test-network:
+    driver: bridge

sdk_compliance_adapter/requirements.txt

@@ -0,0 +1,3 @@
+# SDK Test Adapter dependencies
+flask>=3.0.0
+python-dateutil>=2.8.0

setup_analytics.py

@@ -23,6 +23,7 @@ with open("pyproject.toml", "rb") as f:
 
 # Override specific values
 config["project"]["name"] = "posthoganalytics"
+config["project"]["readme"] = "README_ANALYTICS.md"
 config["tool"]["setuptools"]["dynamic"]["version"] = {
     "attr": "posthoganalytics.version.VERSION"
 }