fix: double counting anthropic langchain (#399)

* add in_app configuration for python SDK

* bump version

* add in_app_modules to init script as well

.github/workflows/ci.yml

@@ -36,6 +36,10 @@ jobs:
               run: |
                   ruff format --check .
 
+            - name: Lint with ruff
+              run: |
+                  ruff check .
+
             - name: Check types with mypy
               run: |
                   mypy --no-site-packages --config-file mypy.ini  . | mypy-baseline filter
@@ -45,7 +49,7 @@ jobs:
         runs-on: ubuntu-latest
         strategy:
             matrix:
-                python-version: ['3.9', '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

.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

@@ -12,15 +12,14 @@ jobs:
   release:
     name: Publish release
     runs-on: ubuntu-latest
-    env:
-      TWINE_USERNAME: __token__
-      TWINE_PASSWORD: ${{ secrets.PYPI_API_TOKEN }}
+    permissions:
+      contents: write
+      id-token: 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
@@ -40,12 +39,12 @@ jobs:
         run: uv sync --extra dev
 
       - name: Push releases to PyPI
+        env:
+          TWINE_USERNAME: __token__
         run: uv run make release && uv run make release_analytics
 
       - name: Create GitHub release
         uses: actions/create-release@0cb9c9b65d5d1901c1f53e5e66eaf4afd303e70e # v1
-        env:
-          GITHUB_TOKEN: ${{ secrets.POSTHOG_BOT_PAT }}
         with:
           tag_name: v${{ env.REPO_VERSION }}
           release_name: ${{ env.REPO_VERSION }}

CHANGELOG.md

@@ -1,3 +1,80 @@
+# 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
+
+- chore(llma): update LLM provider SDKs to latest major versions
+  - openai: 1.102.0 → 2.7.1
+  - anthropic: 0.64.0 → 0.72.0
+  - google-genai: 1.32.0 → 1.49.0
+  - langchain-core: 0.3.75 → 1.0.3
+  - langchain-openai: 0.3.32 → 1.0.2
+  - langchain-anthropic: 0.3.19 → 1.0.1
+  - langchain-community: 0.3.29 → 0.4.1
+  - langgraph: 0.6.6 → 1.0.2
+
+# 6.9.3 - 2025-11-10
+
+- feat(ph-ai): PostHog properties dict in GenerationMetadata
+
+# 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
+
+- fix(error-tracking): pass code variables config from init to client
+
 # 6.9.0 - 2025-11-06
 
 - feat(error-tracking): add local variables capture

README.md

@@ -30,8 +30,8 @@ We recommend using [uv](https://docs.astral.sh/uv/). It's super fast.
 ## PostHog recommends `uv` so...
 
 ```bash
-uv python install 3.9.19
-uv python pin 3.9.19
+uv python install 3.12
+uv python pin 3.12
 uv venv
 source env/bin/activate
 uv sync --extra dev --extra test

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,21 +1,61 @@
 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.feature_flags import InconclusiveMatchError, RequiresServerEvaluation
-from posthog.types import FeatureFlag, FlagsAndPayloads, FeatureFlagResult
+from posthog.contexts import (
+    set_code_variables_mask_patterns_context as inner_set_code_variables_mask_patterns_context,
+)
+from posthog.contexts import (
+    set_context_session as inner_set_context_session,
+)
+from posthog.contexts import (
+    tag as inner_tag,
+)
+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
 
 __version__ = VERSION
@@ -177,6 +217,11 @@ enable_local_evaluation = True  # type: bool
 
 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
 # it impossible to write `posthog.capture(distinct-id, event-name)` - basically, to enforce
@@ -771,6 +816,10 @@ def setup() -> Client:
             enable_exception_autocapture=enable_exception_autocapture,
             log_captured_exceptions=log_captured_exceptions,
             enable_local_evaluation=enable_local_evaluation,
+            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/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

@@ -29,35 +29,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 +114,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 +152,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 +182,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 +236,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(
                         {

posthog/ai/langchain/callbacks.py

@@ -1,8 +1,8 @@
 try:
-    import langchain  # noqa: F401
+    import langchain_core  # noqa: F401
 except ImportError:
     raise ModuleNotFoundError(
-        "Please install LangChain to use this feature: 'pip install langchain'"
+        "Please install LangChain to use this feature: 'pip install langchain-core'"
     )
 
 import json
@@ -79,6 +79,8 @@ class GenerationMetadata(SpanMetadata):
     """Base URL of the provider's API used in the run."""
     tools: Optional[List[Dict[str, Any]]] = None
     """Tools provided to the model."""
+    posthog_properties: Optional[Dict[str, Any]] = None
+    """PostHog properties of the run."""
 
 
 RunMetadata = Union[SpanMetadata, GenerationMetadata]
@@ -420,6 +422,8 @@ class CallbackHandler(BaseCallbackHandler):
                 generation.model = model
             if provider := metadata.get("ls_provider"):
                 generation.provider = provider
+
+            generation.posthog_properties = metadata.get("posthog_properties")
         try:
             base_url = serialized["kwargs"]["openai_api_base"]
             if base_url is not None:
@@ -566,6 +570,9 @@ class CallbackHandler(BaseCallbackHandler):
             "$ai_framework": "langchain",
         }
 
+        if isinstance(run.posthog_properties, dict):
+            event_properties.update(run.posthog_properties)
+
         if run.tools:
             event_properties["$ai_tools"] = run.tools
 
@@ -575,7 +582,7 @@ class CallbackHandler(BaseCallbackHandler):
             event_properties["$ai_is_error"] = True
         else:
             # Add usage
-            usage = _parse_usage(output)
+            usage = _parse_usage(output, run.provider, run.model)
             event_properties["$ai_input_tokens"] = usage.input_tokens
             event_properties["$ai_output_tokens"] = usage.output_tokens
             event_properties["$ai_cache_creation_input_tokens"] = (
@@ -696,6 +703,8 @@ class ModelUsage:
 
 def _parse_usage_model(
     usage: Union[BaseModel, dict],
+    provider: Optional[str] = None,
+    model: Optional[str] = None,
 ) -> ModelUsage:
     if isinstance(usage, BaseModel):
         usage = usage.__dict__
@@ -764,16 +773,32 @@ def _parse_usage_model(
             for mapped_key, dataclass_key in field_mapping.items()
         },
     )
-    # In LangChain, input_tokens is the sum of input and cache read tokens.
-    # Our cost calculation expects them to be separate, for Anthropic.
-    if 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
+    # 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.
+    # 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":
+        is_anthropic = True
+    elif model and "anthropic" in model.lower():
+        is_anthropic = True
+
+    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
 
 
-def _parse_usage(response: LLMResult) -> ModelUsage:
+def _parse_usage(
+    response: LLMResult, provider: Optional[str] = None, model: Optional[str] = None
+) -> ModelUsage:
     # langchain-anthropic uses the usage field
     llm_usage_keys = ["token_usage", "usage"]
     llm_usage: ModelUsage = ModelUsage(
@@ -787,13 +812,15 @@ def _parse_usage(response: LLMResult) -> ModelUsage:
     if response.llm_output is not None:
         for key in llm_usage_keys:
             if response.llm_output.get(key):
-                llm_usage = _parse_usage_model(response.llm_output[key])
+                llm_usage = _parse_usage_model(
+                    response.llm_output[key], provider, model
+                )
                 break
 
     if hasattr(response, "generations"):
         for generation in response.generations:
             if "usage" in generation:
-                llm_usage = _parse_usage_model(generation["usage"])
+                llm_usage = _parse_usage_model(generation["usage"], provider, model)
                 break
 
             for generation_chunk in generation:
@@ -801,7 +828,9 @@ def _parse_usage(response: LLMResult) -> ModelUsage:
                     "usage_metadata" in generation_chunk.generation_info
                 ):
                     llm_usage = _parse_usage_model(
-                        generation_chunk.generation_info["usage_metadata"]
+                        generation_chunk.generation_info["usage_metadata"],
+                        provider,
+                        model,
                     )
                     break
 
@@ -828,7 +857,7 @@ def _parse_usage(response: LLMResult) -> ModelUsage:
                     bedrock_anthropic_usage or bedrock_titan_usage or ollama_usage
                 )
                 if chunk_usage:
-                    llm_usage = _parse_usage_model(chunk_usage)
+                    llm_usage = _parse_usage_model(chunk_usage, provider, model)
                     break
 
     return llm_usage

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

@@ -67,6 +67,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(
                 {

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
 
@@ -83,6 +96,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 +118,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 +130,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 +142,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 +153,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 +205,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/utils.py

@@ -285,7 +285,7 @@ def call_llm_and_track_usage(
 
         event_properties = {
             "$ai_provider": provider,
-            "$ai_model": kwargs.get("model"),
+            "$ai_model": kwargs.get("model") or getattr(response, "model", None),
             "$ai_model_parameters": get_model_params(kwargs),
             "$ai_input": with_privacy_mode(
                 ph_client, posthog_privacy_mode, sanitized_messages
@@ -396,7 +396,7 @@ async def call_llm_and_track_usage_async(
 
         event_properties = {
             "$ai_provider": provider,
-            "$ai_model": kwargs.get("model"),
+            "$ai_model": kwargs.get("model") or getattr(response, "model", None),
             "$ai_model_parameters": get_model_params(kwargs),
             "$ai_input": with_privacy_mode(
                 ph_client, posthog_privacy_mode, sanitized_messages

posthog/client.py

@@ -2,53 +2,62 @@ 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_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,
 )
+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 +193,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 +233,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 +243,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 +266,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 +309,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,
@@ -621,7 +636,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 +668,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 +682,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
 
@@ -725,21 +752,7 @@ class Client(object):
         Examples:
             ```python
             # Set with distinct id
-            posthog.capture(
-                'event_name',
-                distinct_id='user-distinct-id',
-                properties={
-                    '$set': {'name': 'Max Hedgehog'},
-                    '$set_once': {'initial_url': '/blog'}
-                }
-            )
-            ```
-            ```python
-            # Set using context
-            from posthog import new_context, identify_context
-            with new_context():
-                identify_context('user-distinct-id')
-                posthog.capture('event_name')
+            posthog.set(distinct_id='user123', properties={'name': 'Max Hedgehog'})
             ```
 
         Category:
@@ -987,15 +1000,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,
             }
@@ -1160,17 +1170,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.
@@ -1186,7 +1204,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 {}
@@ -1196,19 +1278,41 @@ 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:
@@ -1308,7 +1412,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(
@@ -1400,6 +1505,19 @@ class Client(object):
             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,
@@ -1432,6 +1550,8 @@ class Client(object):
         flag_result = None
         flag_details = None
         request_id = None
+        evaluated_at = None
+        feature_flag_error: Optional[str] = None
 
         flag_value = self._locally_evaluate_flag(
             key, distinct_id, groups, person_properties, group_properties
@@ -1456,14 +1576,24 @@ 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,
+                    )
                 )
+                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
                 )
@@ -1477,19 +1607,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(
@@ -1501,7 +1638,9 @@ class Client(object):
                 groups,
                 disable_geoip,
                 request_id,
+                evaluated_at,
                 flag_details,
+                feature_flag_error,
             )
 
         return flag_result
@@ -1653,7 +1792,7 @@ 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,
     ):
         """
@@ -1667,7 +1806,7 @@ 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.
 
         Examples:
@@ -1683,6 +1822,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,
@@ -1704,9 +1851,10 @@ class Client(object):
         person_properties: dict[str, str],
         group_properties: dict[str, str],
         disable_geoip: Optional[bool],
-    ) -> tuple[Optional[FeatureFlag], Optional[str]]:
+    ) -> 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,
@@ -1717,9 +1865,11 @@ class Client(object):
             flag_keys_to_evaluate=[key],
         )
         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,
@@ -1731,7 +1881,9 @@ 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)}"
@@ -1754,6 +1906,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
@@ -1764,6 +1918,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",
@@ -2029,9 +2185,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/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
@@ -54,6 +54,10 @@ DEFAULT_CODE_VARIABLES_MASK_PATTERNS = [
     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"^__.*"]
@@ -929,7 +933,7 @@ def _compile_patterns(patterns):
     for pattern in patterns:
         try:
             compiled.append(re.compile(pattern))
-        except:
+        except Exception:
             pass
     return compiled
 
@@ -941,7 +945,31 @@ def _pattern_matches(name, patterns):
     return False
 
 
-def _serialize_variable_value(value, limiter, max_length=1024):
+def _mask_sensitive_data(value, compiled_mask):
+    if not compiled_mask:
+        return value
+
+    if isinstance(value, dict):
+        result = {}
+        for k, v in value.items():
+            key_str = str(k) if not isinstance(k, str) else k
+            if _pattern_matches(key_str, compiled_mask):
+                result[k] = CODE_VARIABLES_REDACTED_VALUE
+            else:
+                result[k] = _mask_sensitive_data(v, compiled_mask)
+        return result
+    elif isinstance(value, (list, tuple)):
+        masked_items = [_mask_sensitive_data(item, compiled_mask) for item in value]
+        return type(value)(masked_items)
+    elif isinstance(value, str):
+        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 +982,13 @@ def _serialize_variable_value(value, limiter, max_length=1024):
             limiter.add(result_size)
             return value
         elif isinstance(value, str):
-            result = value
+            if 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 +1001,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 +1075,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 +1087,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/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
+import re
+import socket
+from dataclasses import dataclass
 from datetime import date, datetime
 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")
@@ -124,8 +262,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 +285,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,15 +309,42 @@ 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):
@@ -187,6 +361,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/gemini/test_gemini.py

@@ -407,7 +407,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 +420,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()

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/__init__.py

@@ -1,5 +1,5 @@
 import pytest
 
-pytest.importorskip("langchain")
+pytest.importorskip("langchain_core")
 pytest.importorskip("langchain_community")
 pytest.importorskip("langgraph")

posthog/test/ai/langchain/test_callbacks.py

@@ -113,6 +113,7 @@ def test_metadata_capture(mock_client):
         base_url="https://us.posthog.com",
         name="test",
         end_time=None,
+        posthog_properties=None,
     )
     assert callbacks._runs[run_id] == expected
     with patch("time.time", return_value=1234567891):
@@ -1124,9 +1125,9 @@ def test_anthropic_chain(mock_client):
     )
     chain = prompt | ChatAnthropic(
         api_key=ANTHROPIC_API_KEY,
-        model="claude-3-opus-20240229",
+        model="claude-sonnet-4-5-20250929",
         temperature=0,
-        max_tokens=1,
+        max_tokens=1024,
     )
     callbacks = CallbackHandler(
         mock_client,
@@ -1149,12 +1150,12 @@ def test_anthropic_chain(mock_client):
     assert gen_args["event"] == "$ai_generation"
     assert gen_props["$ai_trace_id"] == "test-trace-id"
     assert gen_props["$ai_provider"] == "anthropic"
-    assert gen_props["$ai_model"] == "claude-3-opus-20240229"
+    assert gen_props["$ai_model"] == "claude-sonnet-4-5-20250929"
     assert gen_props["foo"] == "bar"
 
     assert gen_props["$ai_model_parameters"] == {
         "temperature": 0.0,
-        "max_tokens": 1,
+        "max_tokens": 1024,
         "streaming": False,
     }
     assert gen_props["$ai_input"] == [
@@ -1170,7 +1171,7 @@ def test_anthropic_chain(mock_client):
         <= approximate_latency
     )
     assert gen_props["$ai_input_tokens"] == 17
-    assert gen_props["$ai_output_tokens"] == 1
+    assert gen_props["$ai_output_tokens"] == 4
 
     assert trace_args["event"] == "$ai_trace"
     assert trace_props["$ai_input_state"] == {}
@@ -1187,9 +1188,9 @@ async def test_async_anthropic_streaming(mock_client):
     )
     chain = prompt | ChatAnthropic(
         api_key=ANTHROPIC_API_KEY,
-        model="claude-3-opus-20240229",
+        model="claude-sonnet-4-5-20250929",
         temperature=0,
-        max_tokens=1,
+        max_tokens=1024,
         streaming=True,
         stream_usage=True,
     )
@@ -1269,6 +1270,7 @@ def test_metadata_tools(mock_client):
         name="test",
         tools=tools,
         end_time=None,
+        posthog_properties=None,
     )
     assert callbacks._runs[run_id] == expected
     with patch("time.time", return_value=1234567891):
@@ -1584,13 +1586,147 @@ def test_anthropic_cache_write_and_read_tokens(mock_client):
     generation_props = generation_args["properties"]
 
     assert generation_args["event"] == "$ai_generation"
-    assert generation_props["$ai_input_tokens"] == 400
+    assert (
+        generation_props["$ai_input_tokens"] == 1200
+    )  # No provider metadata, no subtraction
     assert generation_props["$ai_output_tokens"] == 30
     assert generation_props["$ai_cache_creation_input_tokens"] == 0
     assert generation_props["$ai_cache_read_input_tokens"] == 800
     assert generation_props["$ai_reasoning_tokens"] == 0
 
 
+def test_anthropic_provider_subtracts_cache_tokens(mock_client):
+    """Test that Anthropic provider correctly subtracts cache 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 tokens: 1200 input (includes 800 cached)
+    response = LLMResult(
+        generations=[
+            [
+                ChatGeneration(
+                    message=AIMessage(content="Response"),
+                    generation_info={
+                        "usage_metadata": {
+                            "input_tokens": 1200,
+                            "output_tokens": 50,
+                            "cache_read_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"] == 400  # 1200 - 800
+    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(
@@ -1626,7 +1762,7 @@ def test_openai_cache_read_tokens(mock_client):
     generation_props = generation_args["properties"]
 
     assert generation_args["event"] == "$ai_generation"
-    assert generation_props["$ai_input_tokens"] == 50
+    assert generation_props["$ai_input_tokens"] == 150  # No subtraction for OpenAI
     assert generation_props["$ai_output_tokens"] == 40
     assert generation_props["$ai_cache_read_input_tokens"] == 100
     assert generation_props["$ai_cache_creation_input_tokens"] == 0
@@ -1708,7 +1844,7 @@ def test_combined_reasoning_and_cache_tokens(mock_client):
     generation_props = generation_args["properties"]
 
     assert generation_args["event"] == "$ai_generation"
-    assert generation_props["$ai_input_tokens"] == 200
+    assert generation_props["$ai_input_tokens"] == 500  # No subtraction for OpenAI
     assert generation_props["$ai_output_tokens"] == 100
     assert generation_props["$ai_cache_read_input_tokens"] == 300
     assert generation_props["$ai_cache_creation_input_tokens"] == 0
@@ -1716,7 +1852,7 @@ def test_combined_reasoning_and_cache_tokens(mock_client):
 
 
 @pytest.mark.skipif(not OPENAI_API_KEY, reason="OPENAI_API_KEY is not set")
-def test_openai_reasoning_tokens(mock_client):
+def test_openai_reasoning_tokens_o4_mini(mock_client):
     model = ChatOpenAI(
         api_key=OPENAI_API_KEY, model="o4-mini", max_completion_tokens=10
     )
@@ -1917,8 +2053,8 @@ def test_cache_read_tokens_subtraction_from_input_tokens(mock_client):
     generation_props = generation_args["properties"]
 
     assert generation_args["event"] == "$ai_generation"
-    # Input tokens should be reduced: 150 - 100 = 50
-    assert generation_props["$ai_input_tokens"] == 50
+    # Input tokens not reduced without provider metadata
+    assert generation_props["$ai_input_tokens"] == 150
     assert generation_props["$ai_output_tokens"] == 40
     assert generation_props["$ai_cache_read_input_tokens"] == 100
 
@@ -1959,8 +2095,8 @@ def test_cache_read_tokens_subtraction_prevents_negative(mock_client):
     generation_props = generation_args["properties"]
 
     assert generation_args["event"] == "$ai_generation"
-    # Input tokens should be 0, not negative: max(80 - 100, 0) = 0
-    assert generation_props["$ai_input_tokens"] == 0
+    # Input tokens not reduced without provider metadata
+    assert generation_props["$ai_input_tokens"] == 80
     assert generation_props["$ai_output_tokens"] == 20
     assert generation_props["$ai_cache_read_input_tokens"] == 100
 
@@ -2045,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")])
 
@@ -2126,3 +2264,180 @@ def test_agent_action_and_finish_imports():
     assert mock_client.capture.call_count == 1
     call_args = mock_client.capture.call_args[1]
     assert call_args["event"] == "$ai_span"
+
+
+def test_posthog_properties_field_in_generation_metadata(mock_client):
+    """Test that posthog_properties is properly stored in GenerationMetadata."""
+    callbacks = CallbackHandler(mock_client)
+    run_id = uuid.uuid4()
+
+    # Test with billable=True
+    with patch("time.time", return_value=1234567890):
+        callbacks._set_llm_metadata(
+            {"kwargs": {"openai_api_base": "https://api.openai.com"}},
+            run_id,
+            messages=[{"role": "user", "content": "Test message"}],
+            invocation_params={"temperature": 0.5},
+            metadata={
+                "ls_model_name": "gpt-4o",
+                "ls_provider": "openai",
+                "posthog_properties": {"$ai_billable": True},
+            },
+            name="test",
+        )
+
+    expected = GenerationMetadata(
+        model="gpt-4o",
+        input=[{"role": "user", "content": "Test message"}],
+        start_time=1234567890,
+        model_params={"temperature": 0.5},
+        provider="openai",
+        base_url="https://api.openai.com",
+        name="test",
+        posthog_properties={"$ai_billable": True},
+        end_time=None,
+    )
+    assert callbacks._runs[run_id] == expected
+    assert callbacks._runs[run_id].posthog_properties == {"$ai_billable": True}
+
+    callbacks._pop_run_metadata(run_id)
+
+    # Test with billable=False (explicit)
+    run_id2 = uuid.uuid4()
+    with patch("time.time", return_value=1234567890):
+        callbacks._set_llm_metadata(
+            {"kwargs": {"openai_api_base": "https://api.openai.com"}},
+            run_id2,
+            messages=[{"role": "user", "content": "Test message"}],
+            invocation_params={"temperature": 0.5},
+            metadata={
+                "ls_model_name": "gpt-4o",
+                "ls_provider": "openai",
+                "posthog_properties": {"$ai_billable": False},
+            },
+            name="test",
+        )
+
+    assert callbacks._runs[run_id2].posthog_properties == {"$ai_billable": False}
+    callbacks._pop_run_metadata(run_id2)
+
+    # Test when posthog_properties not provided
+    run_id3 = uuid.uuid4()
+    with patch("time.time", return_value=1234567890):
+        callbacks._set_llm_metadata(
+            {"kwargs": {"openai_api_base": "https://api.openai.com"}},
+            run_id3,
+            messages=[{"role": "user", "content": "Test message"}],
+            invocation_params={"temperature": 0.5},
+            metadata={"ls_model_name": "gpt-4o", "ls_provider": "openai"},
+            name="test",
+        )
+
+    assert callbacks._runs[run_id3].posthog_properties is None
+
+
+def test_billable_property_in_generation_event(mock_client):
+    """Test that the billable property is captured in the $ai_generation event."""
+    callbacks = CallbackHandler(mock_client)
+
+    # We need to test the _set_llm_metadata directly since FakeMessagesListChatModel
+    # doesn't support metadata in the same way as real models
+    run_id = uuid.uuid4()
+    with patch("time.time", return_value=1234567890):
+        callbacks._set_llm_metadata(
+            {},
+            run_id,
+            messages=[{"role": "user", "content": "Test"}],
+            metadata={
+                "posthog_properties": {"$ai_billable": True},
+                "ls_model_name": "test-model",
+            },
+            invocation_params={},
+        )
+
+    mock_response = MagicMock()
+    mock_response.generations = [[MagicMock()]]
+
+    with patch("time.time", return_value=1234567891):
+        run = callbacks._pop_run_metadata(run_id)
+
+    callbacks._capture_generation(
+        trace_id=run_id,
+        run_id=run_id,
+        run=run,
+        output=mock_response,
+        parent_run_id=None,
+    )
+
+    assert mock_client.capture.call_count == 1
+    call_args = mock_client.capture.call_args[1]
+    props = call_args["properties"]
+
+    assert call_args["event"] == "$ai_generation"
+    assert props["$ai_billable"] is True
+
+
+def test_billable_defaults_to_false_in_event(mock_client):
+    """Test that $ai_billable is not present when not specified."""
+    prompt = ChatPromptTemplate.from_messages([("user", "Test query")])
+    model = FakeMessagesListChatModel(
+        responses=[AIMessage(content="Test response")],
+    )
+
+    callbacks = [CallbackHandler(mock_client)]
+    chain = prompt | model
+    chain.invoke({}, config={"callbacks": callbacks})
+
+    generation_call = None
+    for call in mock_client.capture.call_args_list:
+        if call[1]["event"] == "$ai_generation":
+            generation_call = call
+            break
+
+    assert generation_call is not None
+    props = generation_call[1]["properties"]
+    assert "$ai_billable" not in props
+
+
+def test_billable_with_real_chain(mock_client):
+    """Test billable tracking through a complete chain execution with mocked metadata."""
+    callbacks = CallbackHandler(mock_client)
+    run_id = uuid.uuid4()
+
+    with patch("time.time", return_value=1000.0):
+        callbacks._set_llm_metadata(
+            {},
+            run_id,
+            messages=[{"role": "user", "content": "What's the weather?"}],
+            metadata={
+                "ls_model_name": "fake-model",
+                "ls_provider": "fake",
+                "posthog_properties": {"$ai_billable": True},
+            },
+            invocation_params={"temperature": 0.7},
+        )
+
+    assert callbacks._runs[run_id].posthog_properties == {"$ai_billable": True}
+
+    mock_response = MagicMock()
+    mock_response.generations = [[MagicMock()]]
+
+    with patch("time.time", return_value=1001.0):
+        run = callbacks._pop_run_metadata(run_id)
+
+    callbacks._capture_generation(
+        trace_id=run_id,
+        run_id=run_id,
+        run=run,
+        output=mock_response,
+        parent_run_id=None,
+    )
+
+    assert mock_client.capture.call_count == 1
+    call_args = mock_client.capture.call_args[1]
+    props = call_args["properties"]
+
+    assert call_args["event"] == "$ai_generation"
+    assert props["$ai_billable"] is True
+    assert props["$ai_model"] == "fake-model"
+    assert props["$ai_provider"] == "fake"

posthog/test/ai/openai/test_openai.py

@@ -1676,3 +1676,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/test_sanitization.py

@@ -1,3 +1,4 @@
+import os
 import unittest
 
 from posthog.ai.sanitization import (
@@ -331,5 +332,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/integrations/test_middleware.py

@@ -315,7 +315,9 @@ class TestPosthogContextMiddlewareSync(unittest.TestCase):
         get_response = Mock(return_value=mock_response)
 
         # Create middleware with request filter that filters all requests
-        request_filter = lambda req: False
+        def request_filter(req):
+            return False
+
         middleware = PosthogContextMiddleware.__new__(PosthogContextMiddleware)
         middleware.get_response = get_response
         middleware._is_coroutine = False

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)
 
@@ -571,6 +567,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)
@@ -752,6 +749,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
@@ -2095,13 +2264,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 +2300,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_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,77 @@ 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

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
 
 
@@ -2348,23 +2348,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,6 +2379,139 @@ 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"')
+
+    @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")
 
     def test_load_feature_flags_wrong_key(self):
         client = Client(FAKE_TEST_API_KEY, personal_api_key=FAKE_TEST_API_KEY)
@@ -2925,6 +3062,7 @@ class TestLocalEvaluation(unittest.TestCase):
                 "some-distinct-id",
                 match_value=True,
                 person_properties={"region": "USA"},
+                send_feature_flag_events=True,
             ),
             300,
         )
@@ -3859,6 +3997,7 @@ class TestCaptureCalls(unittest.TestCase):
                 },
             },
             "requestId": "18043bf7-9cf6-44cd-b959-9662ee20d371",
+            "evaluatedAt": 1234567890,
         }
         client = Client(FAKE_TEST_API_KEY)
 
@@ -3878,6 +4017,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 +4052,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 +4130,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 +4155,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(

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,4 @@
-VERSION = "6.9.0"
+VERSION = "7.4.3"
 
 if __name__ == "__main__":
     print(VERSION, end="")  # noqa: T201

pyproject.toml

@@ -10,18 +10,18 @@ authors = [{ name = "PostHog", email = "hey@posthog.com" }]
 maintainers = [{ name = "PostHog", email = "hey@posthog.com" }]
 license = { text = "MIT" }
 readme = "README.md"
-requires-python = ">=3.9"
+requires-python = ">=3.10"
 classifiers = [
     "Development Status :: 5 - Production/Stable",
     "Intended Audience :: Developers",
     "Operating System :: OS Independent",
     "License :: OSI Approved :: MIT License",
     "Programming Language :: Python",
-    "Programming Language :: Python :: 3.9",
     "Programming Language :: Python :: 3.10",
     "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",
@@ -66,13 +66,13 @@ test = [
     "pytest-timeout",
     "pytest-asyncio",
     "django",
-    "openai",
-    "anthropic",
-    "langgraph>=0.4.8",
-    "langchain-core>=0.3.65",
-    "langchain-community>=0.3.25",
-    "langchain-openai>=0.3.22",
-    "langchain-anthropic>=0.3.15",
+    "openai>=2.0",
+    "anthropic>=0.72",
+    "langgraph>=1.0",
+    "langchain-core>=1.0",
+    "langchain-community>=0.4",
+    "langchain-openai>=1.0",
+    "langchain-anthropic>=1.0",
     "google-genai",
     "pydantic",
     "parameterized>=0.8.1",

references/posthog-python-references-latest.json

@@ -2,7 +2,7 @@
   "id": "posthog-python",
   "hogRef": "0.3",
   "info": {
-    "version": "6.8.0",
+    "version": "7.0.1",
     "id": "posthog-python",
     "title": "PostHog Python SDK",
     "description": "Integrate PostHog into any python application.",
@@ -503,6 +503,24 @@
               "description": "",
               "isOptional": false,
               "type": "bool"
+            },
+            {
+              "name": "capture_exception_code_variables",
+              "description": "",
+              "isOptional": false,
+              "type": "bool"
+            },
+            {
+              "name": "code_variables_mask_patterns",
+              "description": "",
+              "isOptional": false,
+              "type": "any"
+            },
+            {
+              "name": "code_variables_ignore_patterns",
+              "description": "",
+              "isOptional": false,
+              "type": "any"
             }
           ],
           "showDocs": true,
@@ -1478,12 +1496,7 @@
             {
               "id": "example_1",
               "name": "Set with distinct id",
-              "code": "# Set with distinct id\nposthog.capture(\n    'event_name',\n    distinct_id='user-distinct-id',\n    properties={\n        '$set': {'name': 'Max Hedgehog'},\n        '$set_once': {'initial_url': '/blog'}\n    }\n)"
-            },
-            {
-              "id": "example_2",
-              "name": "Set using context",
-              "code": "# Set using context\nfrom posthog import new_context, identify_context\nwith new_context():\n    identify_context('user-distinct-id')\n    posthog.capture('event_name')"
+              "code": "# Set with distinct id\nposthog.set(distinct_id='user123', properties={'name': 'Max Hedgehog'})"
             }
           ]
         },
@@ -2138,6 +2151,12 @@
               "description": "Whether to capture exceptions raised within the context (default: True)",
               "isOptional": false,
               "type": "bool"
+            },
+            {
+              "name": "client",
+              "description": "Optional Posthog client instance to use for this context (default: None)",
+              "isOptional": false,
+              "type": "any"
             }
           ],
           "showDocs": true,
@@ -2216,6 +2235,69 @@
             }
           ]
         },
+        {
+          "id": "set_capture_exception_code_variables_context",
+          "title": "set_capture_exception_code_variables_context",
+          "description": "Set whether code variables are captured for the current context.",
+          "details": "",
+          "category": null,
+          "params": [
+            {
+              "name": "enabled",
+              "description": "",
+              "isOptional": true,
+              "type": "bool"
+            }
+          ],
+          "showDocs": true,
+          "releaseTag": "public",
+          "returnType": {
+            "id": "return_type",
+            "name": "None"
+          }
+        },
+        {
+          "id": "set_code_variables_ignore_patterns_context",
+          "title": "set_code_variables_ignore_patterns_context",
+          "description": "Variable names matching these patterns will be ignored completely when capturing code variables.",
+          "details": "",
+          "category": null,
+          "params": [
+            {
+              "name": "ignore_patterns",
+              "description": "",
+              "isOptional": true,
+              "type": "list"
+            }
+          ],
+          "showDocs": true,
+          "releaseTag": "public",
+          "returnType": {
+            "id": "return_type",
+            "name": "None"
+          }
+        },
+        {
+          "id": "set_code_variables_mask_patterns_context",
+          "title": "set_code_variables_mask_patterns_context",
+          "description": "Variable names matching these patterns will be masked with *** when capturing code variables.",
+          "details": "",
+          "category": null,
+          "params": [
+            {
+              "name": "mask_patterns",
+              "description": "",
+              "isOptional": true,
+              "type": "list"
+            }
+          ],
+          "showDocs": true,
+          "releaseTag": "public",
+          "returnType": {
+            "id": "return_type",
+            "name": "None"
+          }
+        },
         {
           "id": "set_context_session",
           "title": "set_context_session",

setup.py

@@ -14,7 +14,7 @@ long_description = """
 PostHog is developer-friendly, self-hosted product analytics.
 posthog-python is the python package.
 
-This package requires Python 3.9 or higher.
+This package requires Python 3.10 or higher.
 """
 
 # Minimal setup.py for backward compatibility

setup_analytics.py

@@ -47,7 +47,7 @@ long_description = """
 PostHog is developer-friendly, self-hosted product analytics.
 posthog-python is the python package.
 
-This package requires Python 3.9 or higher.
+This package requires Python 3.10 or higher.
 """
 
 # Minimal setup.py for backward compatibility