Contextual bandit support

CHANGELOG.md

@@ -1,5 +1,22 @@
 # Changelog
 
+## Unreleased
+
+### Features
+
+* Contextual bandit support in both clients, at behavioral parity with the JavaScript SDK:
+  * Consumes the `contextualBandits` payload section (and `encryptedContextualBandits`), evaluates `contextualBanditRef`/`contextualVariations` rules with per-leaf weight substitution, and reports `leafId`, `variationWeights`, and `banditVersion` on experiment results for exposure logging.
+  * New `set_payload()` on `GrowthBook` and `GrowthBookClient` for seeding full SDK payloads; only the sections present are overwritten, and encrypted sections are decrypted with the configured `decryption_key` (JS `setPayload` semantics).
+  * Payload refreshes missing a section (`savedGroups`, `contextualBandits`) preserve the previous value at every layer instead of wiping it; the synchronous client serializes payload writers and publishes one coherent evaluation snapshot per update (evaluations stay lock-free), matching the async client.
+  * Malformed bandit definitions or leaves degrade to the rule's aggregate weights instead of raising during evaluation; reported `variationWeights` always match the weights bucketing actually used, and a single, total validity rule governs every weight vector: `getBucketRanges` normalizes vectors with negative, non-finite, boolean, non-numeric, or float-overflowing entries (not just wrong length/sum) to equal weights — never raising, even on arbitrary-precision integers, so bucket ranges can never be inverted, and bandit leaf/aggregate/override propensities always describe the vector actually used. Bandit identifiers get the same treatment: a leaf with a non-integer `leafId` is malformed (aggregate-weights fallback), and a non-integer `banditVersion` is omitted, so exposure metadata never carries invalid attribution ids into bandit training. A rule that pairs `contextualBanditRef` with explicit `ranges` buckets on the ranges (unchanged, matching the JS SDK) but drops the bandit metadata, since leaf propensities cannot describe a ranges-governed assignment.
+  * Contextual bandit exposure metadata survives remote evaluation: `rule.tracks` results from the proxy keep `leafId`, `variationWeights`, and `banditVersion` when replayed through the tracking callback, held to the same validity rules as locally evaluated exposures (invalid identifiers or weight vectors are dropped, not forwarded).
+  * Async evaluations that perform I/O (remote eval or a sticky bucket service) freeze every mutable evaluation input — attributes, groups, overrides, forced variations and features, nested containers included — before their first await, so mutating a `UserContext` mid-flight cannot change leaf routing, remote payloads, or forced assignments. Plain CDN evaluations never yield and skip the copy entirely; deferred callbacks get fire-time snapshots in both clients. `preload_remote_eval` takes the same call-time snapshot, so mutating the context after preloading can no longer cache one attribute state's response under another's key via the stale-while-revalidate background refresh.
+
+### Bug Fixes
+
+* `savedGroups` from a feature refresh were applied to the evaluation context one refresh late in the synchronous client.
+* The built-in tracking plugin now sends the exposure-time user context attributes with experiment events (previously async client events had none).
+
 ## [3.0.0](https://github.com/growthbook/growthbook-python/compare/v2.4.0...v3.0.0) (2026-08-26)
 
 

README.md

@@ -611,6 +611,34 @@ Behaviors to be aware of:
 
 With `GrowthBookClient`, the `on_experiment_viewed` and `on_feature_usage` options — and callbacks registered via `client.subscribe()` — may be either regular functions or coroutines. Coroutine callbacks are scheduled on the event loop without blocking evaluation, and a tracking callback that raises is retried on the next evaluation of the same experiment/user pair.
 
+## Contextual Bandits
+
+Contextual bandit experiments (GrowthBook Enterprise) learn which variation works best for each kind of user. All learning happens server-side: GrowthBook trains a model that partitions users into leaves and computes per-leaf variation weights. The SDK routes each user to the first leaf whose condition matches their attributes and buckets them with that leaf's weights — there is no model or sampling client-side.
+
+Support is automatic once features are loaded from the GrowthBook API. Payloads for contextual bandit experiments contain:
+
+- A top-level `contextualBandits` map (or `encryptedContextualBandits`, decrypted with your `decryption_key` like the other encrypted sections) holding each bandit's `banditVersion` and its per-leaf `condition`/`weights`.
+- Feature rules carrying `contextualBanditRef` (which bandit to use) and `contextualVariations` (the variation values).
+
+Payloads can also be seeded or updated manually. `set_payload` (both clients) overwrites only the sections present, so a payload without `contextualBandits` preserves the current map — the same semantics background refreshes use:
+
+```python
+gb.set_payload({
+    "features": {...},
+    "contextualBandits": {...},  # omitted sections are preserved
+})
+```
+
+When a user is exposed through a contextual bandit rule, the experiment `Result` (including the result passed to `on_experiment_viewed`) carries three extra fields:
+
+- **leafId** — the matched leaf (`-1` when no leaf matched and the rule's aggregate weights were used)
+- **variationWeights** — the weight vector actually used for bucketing (the assignment propensities)
+- **banditVersion** — the version of the bandit model that produced the weights
+
+Log these to your data warehouse in your tracking callback — along with the `user_context` attributes used at exposure time — so GrowthBook can train the bandit and attribute exposures to the right model version. Because weights change as the bandit learns, users may be re-bucketed during the experiment; sticky bucketing is disabled server-side for these rules and analysis attributes each user to their first exposure.
+
+Older SDK versions ignore contextual bandit rules entirely and serve the feature's default value.
+
 ## Inline Experiments
 
 Instead of declaring all features up-front and referencing them by ids in your code, you can also just run an experiment directly. This is done with the `run` method:

growthbook/__init__.py

@@ -28,7 +28,12 @@
     feature_repo,
 )
 
-from .common_types import AbstractAsyncStickyBucketService
+from .common_types import (
+    AbstractAsyncStickyBucketService,
+    ContextualBanditAssignment,
+    ContextualBanditContext,
+    ContextualBanditDefinition,
+)
 
 from .growthbook_client import (
     GrowthBookClient,
@@ -63,6 +68,9 @@
     "Feature",
     "FeatureResult",
     "FeatureRule",
+    "ContextualBanditAssignment",
+    "ContextualBanditContext",
+    "ContextualBanditDefinition",
     # Typing helpers
     "JSONValue",
     "TrackingCallback",

growthbook/codegen.py

@@ -64,7 +64,16 @@ def python_type_for(default_value: Any) -> str:
 # endpoint payload from a bare {feature_key: definition} map that happens to
 # contain a feature literally named "features".
 _ENDPOINT_KEYS = frozenset(
-    {"features", "savedGroups", "encryptedFeatures", "encryptedSavedGroups", "dateUpdated", "status"}
+    {
+        "features",
+        "savedGroups",
+        "contextualBandits",
+        "encryptedFeatures",
+        "encryptedSavedGroups",
+        "encryptedContextualBandits",
+        "dateUpdated",
+        "status",
+    }
 )
 
 
@@ -93,9 +102,12 @@ def infer_feature_type(definition: Any) -> str:
     for rule in definition.get("rules") or []:
         if isinstance(rule, dict):
             candidates.append(rule.get("force"))
-            variations = rule.get("variations")
-            if isinstance(variations, list):
-                candidates.extend(variations)
+            # Contextual bandit rules carry their values under
+            # contextualVariations instead of variations.
+            for key in ("variations", "contextualVariations"):
+                variations = rule.get(key)
+                if isinstance(variations, list):
+                    candidates.extend(variations)
     types = {python_type_for(v) for v in candidates if v is not None}
     if len(types) == 1:
         return types.pop()

growthbook/common_types.py

@@ -1,6 +1,6 @@
 #!/usr/bin/env python
 
-from dataclasses import dataclass, field
+from dataclasses import dataclass, field, replace
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -56,6 +56,33 @@ class Filter(TypedDict, total=False):
     hashVersion: int
     attribute: str
 
+
+# Contextual bandit payload types. The server ships a top-level
+# "contextualBandits" map keyed by bandit id; each entry holds per-leaf
+# weight vectors computed server-side (Thompson sampling within decision-tree
+# leaves). The SDK only picks the first leaf whose condition matches and
+# buckets with that leaf's weights — no model evaluation happens client-side.
+
+class ContextualBanditContext(TypedDict, total=False):
+    leafId: Required[int]
+    condition: Dict[str, Any]
+    weights: Required[List[float]]
+
+
+class ContextualBanditDefinition(TypedDict, total=False):
+    banditVersion: int
+    contexts: List[ContextualBanditContext]
+
+
+# Assignment metadata attached to an Experiment/Result for a contextual
+# bandit rule. banditVersion is omitted (never None) when the definition
+# doesn't carry one — serialization must match the JS SDK byte-for-byte.
+class ContextualBanditAssignment(TypedDict, total=False):
+    leafId: Required[int]
+    variationWeights: Required[List[float]]
+    banditVersion: int
+
+
 class Experiment(Generic[T]):
     def __init__(
         self,
@@ -85,6 +112,7 @@ def __init__(
         minBucketVersion: Optional[int] = None,
         parentConditions: Optional[List[Dict[str, Any]]] = None,
         customFields: Optional[Dict[str, Any]] = None,
+        contextualBandit: Optional[ContextualBanditAssignment] = None,
         # NoReturn makes literal unknown kwargs a checker error (like TS excess
         # property checks) while **dict payload splats (typed Any) still pass;
         # at runtime unknown payload keys are swallowed as before.
@@ -113,6 +141,7 @@ def __init__(
         # Custom Fields defined for the experiment in the GrowthBook UI.
         # Arrives from the API as a flat dict (e.g. {"cfl_abc123": "value"}).
         self.customFields = customFields or {}
+        self.contextualBandit = contextualBandit
 
         self.fallbackAttribute = None
         if not self.disableStickyBucketing:
@@ -156,6 +185,8 @@ def to_dict(self) -> Dict[str, Any]:
             obj["parentConditions"] = self.parentConditions
         if self.customFields:
             obj["customFields"] = self.customFields
+        if self.contextualBandit is not None:
+            obj["contextualBandit"] = self.contextualBandit
 
         return obj
 
@@ -194,6 +225,9 @@ def __init__(
         meta: Optional[VariationMeta] = None,
         bucket: Optional[float] = None,
         stickyBucketUsed: bool = False,
+        leafId: Optional[int] = None,
+        variationWeights: Optional[List[float]] = None,
+        banditVersion: Optional[int] = None,
     ) -> None:
         self.variationId = variationId
         self.inExperiment = inExperiment
@@ -204,6 +238,11 @@ def __init__(
         self.featureId = featureId or None
         self.bucket = bucket
         self.stickyBucketUsed = stickyBucketUsed
+        # Contextual bandit exposure metadata; set only for real hashed
+        # exposures of a contextual bandit rule.
+        self.leafId = leafId
+        self.variationWeights = variationWeights
+        self.banditVersion = banditVersion
 
         self.key = str(variationId)
         self.name = ""
@@ -236,6 +275,14 @@ def to_dict(self) -> Dict[str, Any]:
             obj["name"] = self.name
         if self.passthrough:
             obj["passthrough"] = True
+        # The fallback leafId is -1, so these must be gated on None, not
+        # truthiness.
+        if self.leafId is not None:
+            obj["leafId"] = self.leafId
+        if self.variationWeights is not None:
+            obj["variationWeights"] = self.variationWeights
+        if self.banditVersion is not None:
+            obj["banditVersion"] = self.banditVersion
 
         return obj
 
@@ -312,6 +359,8 @@ def __init__(
         minBucketVersion: Optional[int] = None,
         parentConditions: Optional[List[Dict[str, Any]]] = None,
         tracks: Optional[List[Dict[str, Any]]] = None,
+        contextualBanditRef: Optional[str] = None,
+        contextualVariations: Optional[List[Any]] = None,
         # See Experiment.__init__: checker-strict, runtime-permissive.
         **_ignored: NoReturn,
     ) -> None:
@@ -344,6 +393,11 @@ def __init__(
         # Remote-eval rules carry pre-evaluated experiment tracking events on
         # the force branch; see _fireRuleTracks in core.py.
         self.tracks = tracks
+        # Contextual bandit rules carry their variations under
+        # contextualVariations (capability-gated key) so bandit-unaware SDKs
+        # skip the rule instead of bucketing on stale weights.
+        self.contextualBanditRef = contextualBanditRef
+        self.contextualVariations = contextualVariations
 
     def to_dict(self) -> Dict[str, Any]:
         data: Dict[str, Any] = {}
@@ -393,6 +447,10 @@ def to_dict(self) -> Dict[str, Any]:
             data["parentConditions"] = self.parentConditions
         if self.tracks:
             data["tracks"] = self.tracks
+        if self.contextualBanditRef:
+            data["contextualBanditRef"] = self.contextualBanditRef
+        if self.contextualVariations is not None:
+            data["contextualVariations"] = self.contextualVariations
 
         return data
 
@@ -529,6 +587,51 @@ def __call__(
 ]
 
 
+def snapshot_attributes(attributes: Dict[str, Any]) -> Dict[str, Any]:
+    """Recursive copy of a JSON-compatible attributes dict (dicts and lists
+    are copied, scalars shared). Freezes the values used for bucketing and
+    contextual bandit leaf routing so later caller mutations — including
+    nested ones — can't change what an in-flight evaluation or a deferred
+    callback observes."""
+    def copy_value(value: Any) -> Any:
+        if isinstance(value, dict):
+            return {k: copy_value(v) for k, v in value.items()}
+        if isinstance(value, list):
+            return [copy_value(v) for v in value]
+        return value
+
+    return {k: copy_value(v) for k, v in attributes.items()}
+
+
+def snapshot_user_context(user: "UserContext") -> "UserContext":
+    """Call-time snapshot of every mutable evaluation input on a UserContext.
+
+    The single freeze operation used wherever a context's data outlives the
+    caller's synchronous control flow: evaluations that await (remote-eval
+    fetch, sticky-bucket I/O) and remote-eval cache preloads, whose POST body
+    can be serialized by a background SWR refresh long after the caller has
+    moved on. Without it, a later caller mutation could change the request
+    body while the cache key still describes the original attribute state."""
+    return replace(
+        user,
+        attributes=snapshot_attributes(user.attributes),
+        groups=snapshot_attributes(user.groups),
+        forced_variations=snapshot_attributes(user.forced_variations),
+        forced_features=snapshot_attributes(user.forced_features),
+        overrides=snapshot_attributes(user.overrides),
+    )
+
+
+def tracking_user_context(user: "UserContext") -> "UserContext":
+    """Exposure-time snapshot of a user context for tracking and
+    feature-usage callbacks (JS SDK: getTrackingUserContext).
+
+    Attributes are copied recursively (the JS SDK only shallow-spreads;
+    Python's threaded callers make nested mutation of a deferred callback's
+    payload a realistic hazard, so this diverges deliberately)."""
+    return replace(user, attributes=snapshot_attributes(user.attributes))
+
+
 @dataclass
 class Options:
     url: Optional[str] = None
@@ -573,6 +676,10 @@ class GlobalContext:
     options: Options
     features: Dict[str, "Feature"] = field(default_factory=dict)
     saved_groups: Dict[str, Any] = field(default_factory=dict)
+    # Payload "contextualBandits" map ({bandit id -> definition}), kept as
+    # dicts (not materialized into classes) like saved_groups; the TypedDict
+    # value type documents the wire shape for readers and checkers.
+    contextual_bandits: Dict[str, ContextualBanditDefinition] = field(default_factory=dict)
 
 @dataclass
 class EvaluationContext: