gatewaysdk 0.3.2__py3-none-any.whl

gatewaysdk/experiments/client.py

@@ -0,0 +1,877 @@
+# Copyright (c) Gateway. All rights reserved.
+
+"""Main client API for A/B experiments.
+
+This module provides the ExperimentClient class that users interact with
+to register experiments and get variant assignments. It also manages the
+context variable that the span processor reads from.
+
+The client supports two modes:
+1. **Local-only**: Experiments registered in code stay in-memory
+2. **Backend-synced**: Experiments are fetched from and synced to the Gateway backend
+
+Usage::
+
+    import gatewaysdk as msk
+
+    # Initialize with backend sync (recommended)
+    msk.init_experiments(
+        api_endpoint="https://gateway.example.com/api/public/experiments/exposures/batch",
+        public_key="pk-lf-...",
+        secret_key="sk-lf-...",
+    )
+
+    # Option A: Register experiment in code (syncs to backend in background)
+    msk.experiments.register(
+        key="prompt-test-v1",
+        variants=[
+            {"key": "control", "weight": 50, "config": {"prompt": "..."}},
+            {"key": "treatment", "weight": 50, "config": {"prompt": "..."}},
+        ],
+        traffic_percent=100,
+    )
+
+    # Option B: Fetch experiment from backend (defined in UI)
+    msk.experiments.fetch("prompt-test-v1")
+
+    # Get variant for a user (~0ms, purely in-memory)
+    variant = msk.experiments.get_variant("prompt-test-v1", user_id="user-123")
+    if variant:
+        print(f"Assigned to {variant.key}")
+        system_prompt = variant.config["prompt"]
+"""
+
+from __future__ import annotations
+
+import logging
+import threading
+from contextvars import ContextVar
+from datetime import datetime, timezone
+from typing import TYPE_CHECKING, Any, Dict, List, Optional
+
+import requests
+
+from .assignment import assign_variant
+from .types import ExperimentConfig, Exposure, Variant
+
+if TYPE_CHECKING:
+    from .exposure import ExposureLogger
+
+logger = logging.getLogger(__name__)
+
+
+# =============================================================================
+# Exceptions
+# =============================================================================
+
+
+class BalancedAssignmentError(Exception):
+    """Raised when server-side balanced assignment fails due to network/server issues.
+
+    This exception is raised by get_variant_balanced() when the server cannot be reached
+    or returns an error. The caller (get_variant) catches this and falls back to
+    local hash-based assignment for graceful degradation.
+
+    This is NOT raised when the user is legitimately not in traffic (server returns
+    variantKey: null) — in that case, None is returned normally.
+    """
+
+    pass
+
+
+# =============================================================================
+# Context Variable for Experiment Assignment
+# =============================================================================
+
+# Stores the current experiment context for span injection.
+# The TracingContextSpanProcessor reads this and injects attributes into spans.
+_current_experiment: ContextVar[Optional[Dict[str, Any]]] = ContextVar("msk_current_experiment", default=None)
+
+
+def get_current_experiment() -> Optional[Dict[str, Any]]:
+    """Get the current experiment assignment context.
+
+    Returns a dict with keys:
+    - experiment_id: The experiment key
+    - variant_key: The assigned variant key
+    - user_id: The user identifier
+
+    Returns None if no experiment is active.
+
+    This is called by TracingContextSpanProcessor to inject experiment
+    attributes into spans.
+    """
+    return _current_experiment.get()
+
+
+def _set_current_experiment(ctx: Optional[Dict[str, Any]]) -> None:
+    """Set the current experiment context.
+
+    Internal function - called by ExperimentClient.get_variant().
+    """
+    _current_experiment.set(ctx)
+
+
+def clear_current_experiment() -> None:
+    """Clear the current experiment context.
+
+    Call this when you want to stop attributing traces to an experiment,
+    for example when switching users or ending an experiment session.
+    """
+    _current_experiment.set(None)
+
+
+# =============================================================================
+# Experiment Client
+# =============================================================================
+
+
+class ExperimentClient:
+    """Client for managing A/B experiments.
+
+    This class provides the main API for:
+    - Registering experiments with variants (locally and to backend)
+    - Fetching experiments from the backend
+    - Getting variant assignments for users (~0ms, in-memory)
+    - Managing exposure logging (async, background thread)
+    - Background config polling (keeps in-memory cache fresh)
+
+    The client is typically accessed via the gatewaysdk module:
+        msk.experiments.register(...)
+        msk.experiments.fetch(...)
+        msk.experiments.get_variant(...)
+
+    Attributes:
+        exposure_logger: The background logger for exposure events
+        api_base_url: Base URL for the experiments API
+        public_key: Gateway public key for authentication
+        secret_key: Gateway secret key for authentication
+        poll_interval: Seconds between background config refreshes (default 30)
+    """
+
+    def __init__(
+        self,
+        exposure_logger: Optional["ExposureLogger"] = None,
+        api_base_url: Optional[str] = None,
+        public_key: Optional[str] = None,
+        secret_key: Optional[str] = None,
+        poll_interval: float = 30.0,
+    ):
+        """Initialize the experiment client.
+
+        Args:
+            exposure_logger: Optional ExposureLogger for sending exposure events.
+                            If not provided, exposures will be logged but not sent.
+            api_base_url: Base URL for the experiments API (e.g., "https://gateway.example.com").
+                         If not provided, backend sync is disabled.
+            public_key: Gateway public key for API authentication.
+            secret_key: Gateway secret key for API authentication.
+            poll_interval: Seconds between background polls for config updates.
+                          Default is 30 seconds. Set to 0 to disable polling.
+        """
+        # PREVIOUSLY: plain dict, written from main thread only.
+        # NOW: protected by a lock because the background refresh thread also
+        # writes to it while get_variant() reads from the main thread.
+        self._experiments: Dict[str, ExperimentConfig] = {}
+        self._experiments_lock = threading.Lock()
+
+        self._exposure_logger = exposure_logger
+        self._api_base_url = api_base_url
+        self._public_key = public_key
+        self._secret_key = secret_key
+        self._poll_interval = poll_interval
+
+        # Background polling state
+        self._poll_thread: Optional[threading.Thread] = None
+        self._poll_shutdown = threading.Event()
+        self._poll_started = False
+
+    def set_exposure_logger(self, logger: "ExposureLogger") -> None:
+        """Set the exposure logger after initialization.
+
+        This is useful when the client is created before the API configuration
+        is known (e.g., during module initialization).
+        """
+        self._exposure_logger = logger
+
+    def configure_backend(
+        self,
+        api_base_url: str,
+        public_key: str,
+        secret_key: str,
+    ) -> None:
+        """Configure backend connectivity for fetching/syncing experiments.
+
+        Args:
+            api_base_url: Base URL for the experiments API.
+            public_key: Gateway public key.
+            secret_key: Gateway secret key.
+        """
+        self._api_base_url = api_base_url
+        self._public_key = public_key
+        self._secret_key = secret_key
+        logger.debug(f"Configured experiments backend: {api_base_url}")
+
+    def _get_auth_header(self) -> Dict[str, str]:
+        """Get authorization header for API requests."""
+        if not self._public_key or not self._secret_key:
+            return {}
+        import base64
+
+        credentials = base64.b64encode(f"{self._public_key}:{self._secret_key}".encode()).decode()
+        return {"Authorization": f"Basic {credentials}"}
+
+    @staticmethod
+    def _parse_experiment_response(data: Dict[str, Any]) -> ExperimentConfig:
+        """Parse an experiment JSON response into an ExperimentConfig."""
+        return ExperimentConfig(
+            key=data["key"],
+            variants=[
+                Variant(
+                    key=v["key"],
+                    weight=v["weight"],
+                    config=v.get("config", {}),
+                )
+                for v in data["variants"]
+            ],
+            traffic_percent=data["trafficPercent"],
+        )
+
+    # -------------------------------------------------------------------------
+    # Background Config Polling
+    # -------------------------------------------------------------------------
+
+    def start_background_refresh(self) -> None:
+        """Start a background daemon thread that polls the backend for config updates.
+
+        WHY: Previously, config was only loaded when the caller explicitly called
+        fetch() or fetch_all() — a blocking HTTP call. Now, a daemon thread runs
+        silently in the background and atomically swaps in fresh configs every
+        poll_interval seconds. get_variant() always reads from warm memory (~0ms).
+
+        Safe to call multiple times — only starts once. No-op if poll_interval=0
+        or backend is not configured.
+        """
+        if self._poll_started or not self._api_base_url or self._poll_interval <= 0:
+            return
+
+        self._poll_shutdown.clear()
+        self._poll_thread = threading.Thread(
+            target=self._poll_loop,
+            daemon=True,
+            name="msk-experiment-refresh",
+        )
+        self._poll_thread.start()
+        self._poll_started = True
+        logger.debug(f"Background experiment refresh started (interval={self._poll_interval}s)")
+
+    def stop_background_refresh(self) -> None:
+        """Stop the background refresh thread. Blocks until it exits (up to 5s)."""
+        if not self._poll_started:
+            return
+        self._poll_shutdown.set()
+        if self._poll_thread:
+            self._poll_thread.join(timeout=5.0)
+        self._poll_started = False
+        logger.debug("Background experiment refresh stopped")
+
+    def _poll_loop(self) -> None:
+        """Background loop: refresh all known experiment keys every poll_interval."""
+        # Run immediately on first iteration, then wait between subsequent ones.
+        while not self._poll_shutdown.is_set():
+            self._refresh_all_from_backend()
+            # wait() returns True if shutdown was requested, False on timeout.
+            # Using wait() instead of sleep() so shutdown is instant.
+            self._poll_shutdown.wait(timeout=self._poll_interval)
+
+    def _refresh_all_from_backend(self) -> None:
+        """Silently re-fetch all registered experiment keys from the backend.
+
+        Called by the background poll loop. Updates _experiments atomically
+        under a lock so get_variant() never sees a partial state.
+        """
+        if not self._api_base_url:
+            return
+
+        with self._experiments_lock:
+            keys = list(self._experiments.keys())
+
+        if not keys:
+            return
+
+        headers = self._get_auth_header()
+        updated = 0
+
+        for key in keys:
+            url = f"{self._api_base_url}/api/public/experiments/{key}"
+            try:
+                response = requests.get(url, headers=headers, timeout=5)
+                if response.status_code == 200:
+                    config = self._parse_experiment_response(response.json())
+                    with self._experiments_lock:
+                        self._experiments[key] = config
+                    updated += 1
+                elif response.status_code == 404:
+                    logger.debug(f"Experiment '{key}' no longer found in backend")
+                else:
+                    logger.debug(f"Background refresh for '{key}' got {response.status_code}")
+            except Exception as e:
+                # Never crash the background thread — just log and move on.
+                logger.debug(f"Background refresh failed for '{key}': {e}")
+
+        if updated:
+            logger.debug(f"Background refresh updated {updated}/{len(keys)} experiments")
+
+    # -------------------------------------------------------------------------
+    # Register / Sync
+    # -------------------------------------------------------------------------
+
+    def register(
+        self,
+        key: str,
+        variants: List[Dict[str, Any]],
+        traffic_percent: int = 100,
+        sync_to_backend: bool = True,
+    ) -> None:
+        """Register an experiment configuration.
+
+        Experiments are defined in code and can optionally be synced to the
+        backend database. This makes them visible in the UI and allows for
+        evaluator configuration.
+
+        PREVIOUSLY: if sync_to_backend=True, this blocked the caller for the
+        full HTTP round-trip (~100–500ms) before returning.
+        NOW: the backend sync is fire-and-forget (daemon thread). register()
+        returns immediately after writing to in-memory state.
+
+        Args:
+            key: Unique experiment identifier (e.g., "prompt-test-v1").
+                 Should be descriptive and versioned.
+            variants: List of variant dicts, each with:
+                - key: Variant identifier (e.g., "control", "treatment")
+                - weight: Relative weight for traffic (e.g., 50 = 50%)
+                - config: Arbitrary config dict passed to your code
+            traffic_percent: Percentage of users to include (0-100).
+                            Use < 100 to gradually roll out experiments.
+            sync_to_backend: If True and backend is configured, syncs the
+                            experiment to the Gateway database. Default True.
+
+        Example::
+
+            msk.experiments.register(
+                key="prompt-test-v1",
+                variants=[
+                    {"key": "control", "weight": 50, "config": {"prompt": "You are a helpful assistant."}},
+                    {"key": "treatment", "weight": 50, "config": {"prompt": "You are a concise assistant."}},
+                ],
+                traffic_percent=100,
+            )
+        """
+        if traffic_percent < 0 or traffic_percent > 100:
+            raise ValueError(f"traffic_percent must be 0-100, got {traffic_percent}")
+
+        if len(variants) < 2:
+            raise ValueError("Experiments require at least 2 variants")
+
+        config = ExperimentConfig(
+            key=key,
+            variants=[
+                Variant(
+                    key=v["key"],
+                    weight=v.get("weight", 1),
+                    config=v.get("config", {}),
+                )
+                for v in variants
+            ],
+            traffic_percent=traffic_percent,
+        )
+
+        with self._experiments_lock:
+            self._experiments[key] = config
+
+        logger.debug(f"Registered experiment '{key}' with {len(variants)} variants, " f"{traffic_percent}% traffic")
+
+        # PREVIOUSLY: blocking requests.post() here — caller waited for it.
+        # NOW: spawns a daemon thread and returns immediately.
+        if sync_to_backend and self._api_base_url:
+            self._sync_experiment_to_backend_async(key, variants, traffic_percent)
+
+    def _sync_experiment_to_backend_async(
+        self,
+        key: str,
+        variants: List[Dict[str, Any]],
+        traffic_percent: int,
+    ) -> None:
+        """Fire-and-forget backend sync — runs in a daemon thread.
+
+        PREVIOUSLY: this was _sync_experiment_to_backend() and was called
+        directly (blocking). Now it spawns a thread and returns in ~0ms.
+        """
+
+        def _do_sync() -> None:
+            url = f"{self._api_base_url}/api/public/experiments"
+            headers = {
+                "Content-Type": "application/json",
+                **self._get_auth_header(),
+            }
+            payload = {
+                "key": key,
+                "trafficPercent": traffic_percent,
+                "variants": [
+                    {
+                        "key": v["key"],
+                        "weight": v.get("weight", 1),
+                        "config": v.get("config", {}),
+                    }
+                    for v in variants
+                ],
+            }
+            try:
+                response = requests.post(url, json=payload, headers=headers, timeout=5)
+                if response.status_code in (200, 201):
+                    data = response.json()
+                    action = "created" if data.get("created") else "updated"
+                    logger.info(f"Experiment '{key}' {action} in backend (id={data.get('id')})")
+                else:
+                    logger.warning(
+                        f"Failed to sync experiment '{key}' to backend: " f"{response.status_code} {response.text}"
+                    )
+            except Exception as e:
+                logger.warning(f"Failed to sync experiment '{key}' to backend: {e}")
+
+        thread = threading.Thread(
+            target=_do_sync,
+            daemon=True,
+            name=f"msk-sync-{key}",
+        )
+        thread.start()
+
+    # -------------------------------------------------------------------------
+    # Fetch (explicit blocking fetch — for startup patterns)
+    # -------------------------------------------------------------------------
+
+    def fetch(
+        self,
+        experiment_key: str,
+        auto_register: bool = True,
+    ) -> Optional[ExperimentConfig]:
+        """Fetch an experiment configuration from the backend.
+
+        This is intentionally synchronous — use it at startup when you need
+        the config to be available before proceeding. For runtime updates,
+        rely on start_background_refresh() instead.
+
+        Args:
+            experiment_key: The experiment identifier to fetch.
+            auto_register: If True, automatically register the fetched
+                          experiment locally so get_variant() works.
+
+        Returns:
+            The ExperimentConfig if found, None otherwise.
+
+        Example::
+
+            # Fetch experiment defined in the UI
+            config = msk.experiments.fetch("prompt-test-v1")
+            if config:
+                variant = msk.experiments.get_variant("prompt-test-v1", user_id="user-123")
+        """
+        if not self._api_base_url:
+            logger.warning("Backend not configured, cannot fetch experiment")
+            return None
+
+        url = f"{self._api_base_url}/api/public/experiments/{experiment_key}"
+        headers = self._get_auth_header()
+
+        try:
+            response = requests.get(url, headers=headers, timeout=5)
+            if response.status_code == 200:
+                config = self._parse_experiment_response(response.json())
+
+                if auto_register:
+                    with self._experiments_lock:
+                        self._experiments[experiment_key] = config
+                    logger.info(
+                        f"Fetched and registered experiment '{experiment_key}' "
+                        f"from backend ({len(config.variants)} variants)"
+                    )
+                else:
+                    logger.info(f"Fetched experiment '{experiment_key}' from backend")
+
+                return config
+            elif response.status_code == 404:
+                logger.warning(f"Experiment '{experiment_key}' not found in backend")
+                return None
+            else:
+                logger.warning(
+                    f"Failed to fetch experiment '{experiment_key}': " f"{response.status_code} {response.text}"
+                )
+                return None
+        except Exception as e:
+            logger.warning(f"Failed to fetch experiment '{experiment_key}': {e}")
+            return None
+
+    def fetch_all(self, auto_register: bool = True) -> List[ExperimentConfig]:
+        """Fetch all active experiments from the backend.
+
+        The listing endpoint returns summaries only (no variant details), so
+        each experiment must be fetched individually via ``fetch()``.
+
+        Useful at startup to pre-populate the in-memory cache. After calling
+        this, ``start_background_refresh()`` will keep configs fresh.
+
+        Args:
+            auto_register: If True, automatically register fetched experiments.
+
+        Returns:
+            List of ExperimentConfig objects that were successfully fetched.
+        """
+        if not self._api_base_url:
+            logger.warning("Backend not configured, cannot fetch experiments")
+            return []
+
+        url = f"{self._api_base_url}/api/public/experiments"
+        headers = self._get_auth_header()
+
+        try:
+            response = requests.get(url, headers=headers, timeout=10)
+            if response.status_code != 200:
+                logger.warning(f"Failed to fetch experiments: {response.status_code} {response.text}")
+                return []
+
+            experiment_keys = [exp["key"] for exp in response.json().get("experiments", [])]
+        except Exception as e:
+            logger.warning(f"Failed to fetch experiments: {e}")
+            return []
+
+        configs: List[ExperimentConfig] = []
+        for key in experiment_keys:
+            config = self.fetch(key, auto_register=auto_register)
+            if config:
+                configs.append(config)
+
+        logger.info(f"Fetched {len(configs)}/{len(experiment_keys)} experiments from backend")
+        return configs
+
+    # -------------------------------------------------------------------------
+    # Get Variant (~0ms hot path)
+    # -------------------------------------------------------------------------
+
+    def get_variant(
+        self,
+        experiment_key: str,
+        user_id: str,
+    ) -> Optional[Variant]:
+        """Get the variant for a user in an experiment.
+
+        This is the main entry point for A/B testing. It:
+        1. Deterministically assigns the user to a variant (or None if not in traffic)
+        2. Sets the context variable for span injection
+        3. Logs an exposure event (sent to backend asynchronously)
+
+        When the backend is configured, uses server-side balanced assignment
+        (Redis counter) for exact proportional allocation at any sample size.
+        Falls back to local hash-based assignment when the backend is unavailable.
+
+        Args:
+            experiment_key: The experiment identifier (must be registered first)
+            user_id: The user/subject identifier. This should be a stable ID
+                    that uniquely identifies the user across sessions.
+
+        Returns:
+            The assigned Variant with its config, or None if:
+            - The experiment is not registered
+            - The user is not in the traffic allocation
+
+        Example::
+
+            variant = msk.experiments.get_variant("prompt-test-v1", user_id="user-alice-789")
+            if variant:
+                print(f"User assigned to {variant.key}")
+                system_prompt = variant.config["prompt"]
+                # Run your agent with this prompt
+            else:
+                # User not in experiment, use default behavior
+                system_prompt = DEFAULT_PROMPT
+        """
+        if self._api_base_url:
+            try:
+                result = self.get_variant_balanced(experiment_key, user_id)
+                # result is None only when user is legitimately not in traffic
+                # (server returned variantKey: null) — do NOT fall back
+                return result
+            except BalancedAssignmentError:
+                # Network error, timeout, server error — fall back to hash-based
+                logger.warning(
+                    f"Balanced assignment failed for '{experiment_key}', " f"falling back to hash-based assignment"
+                )
+
+        return self._get_variant_hash(experiment_key, user_id)
+
+    def _get_variant_hash(
+        self,
+        experiment_key: str,
+        user_id: str,
+    ) -> Optional[Variant]:
+        """Hash-based local assignment (fallback when backend is unavailable)."""
+        with self._experiments_lock:
+            config = self._experiments.get(experiment_key)
+
+        if config is None:
+            logger.warning(
+                f"Experiment '{experiment_key}' not registered. "
+                f"Call msk.experiments.register() or msk.experiments.fetch() first."
+            )
+            return None
+
+        variant = assign_variant(config, user_id)
+        if variant is None:
+            logger.debug(f"User '{user_id}' not in traffic for '{experiment_key}'")
+            return None
+
+        _set_current_experiment(
+            {
+                "experiment_id": experiment_key,
+                "variant_key": variant.key,
+                "user_id": user_id,
+            }
+        )
+
+        if self._exposure_logger:
+            self._exposure_logger.log(
+                Exposure(
+                    user_id=user_id,
+                    experiment_id=experiment_key,
+                    variant_key=variant.key,
+                    timestamp=datetime.now(timezone.utc).isoformat(),
+                )
+            )
+        else:
+            logger.debug(
+                f"No exposure logger configured, exposure not sent "
+                f"(experiment={experiment_key}, variant={variant.key}, user={user_id})"
+            )
+
+        logger.debug(f"Assigned user '{user_id}' to variant '{variant.key}' " f"in experiment '{experiment_key}'")
+
+        return variant
+
+    def assign_with_override(
+        self,
+        experiment_key: str,
+        user_id: str,
+        variant_key: str,
+    ) -> Optional[Variant]:
+        """Assign a user to a specific variant, bypassing hash-based selection.
+
+        Use this when you have pre-computed balanced assignments (e.g. via
+        ``compute_balanced_assignments``) and want to force a specific variant
+        while still setting the tracing context and logging the exposure.
+
+        Args:
+            experiment_key: The experiment identifier (must be registered first).
+            user_id: The user/subject identifier.
+            variant_key: The variant key to force (e.g. "control", "treatment").
+
+        Returns:
+            The matching Variant with its config, or None if the experiment is
+            not registered or variant_key is not found.
+
+        Example::
+
+            assignments = compute_balanced_assignments(config, all_users)
+            for user_id, variant in assignments.items():
+                if variant:
+                    msk.experiments.assign_with_override(
+                        "my-exp", user_id, variant.key
+                    )
+        """
+        with self._experiments_lock:
+            config = self._experiments.get(experiment_key)
+
+        if config is None:
+            logger.warning(
+                f"Experiment '{experiment_key}' not registered. "
+                f"Call msk.experiments.register() or msk.experiments.fetch() first."
+            )
+            return None
+
+        variant = next((v for v in config.variants if v.key == variant_key), None)
+        if variant is None:
+            logger.warning(f"Variant '{variant_key}' not found in experiment '{experiment_key}'")
+            return None
+
+        _set_current_experiment(
+            {
+                "experiment_id": experiment_key,
+                "variant_key": variant.key,
+                "user_id": user_id,
+            }
+        )
+
+        if self._exposure_logger:
+            self._exposure_logger.log(
+                Exposure(
+                    user_id=user_id,
+                    experiment_id=experiment_key,
+                    variant_key=variant.key,
+                    timestamp=datetime.now(timezone.utc).isoformat(),
+                )
+            )
+
+        logger.debug(f"Force-assigned user '{user_id}' to variant '{variant.key}' " f"in experiment '{experiment_key}'")
+
+        return variant
+
+    def get_variant_balanced(
+        self,
+        experiment_key: str,
+        user_id: str,
+    ) -> Optional[Variant]:
+        """Get variant using server-side balanced assignment (guaranteed 50/50).
+
+        Unlike ``get_variant`` which uses local hash-based assignment (can produce
+        unequal splits at small sample sizes), this method calls the server which
+        uses a Redis atomic counter to guarantee exact proportional allocation.
+
+        Trade-off: requires a network call (~10-50ms) vs hash-based (~0ms local).
+        Use this when you have small sample sizes and need exact balance.
+
+        The assignment is:
+        1. Deterministic per user — same user always gets same variant
+        2. Perfectly balanced — guaranteed 50/50 (or proportional to weights)
+        3. Persistent — stored in Redis, survives restarts
+
+        Args:
+            experiment_key: The experiment identifier (must exist in backend).
+            user_id: The user/subject identifier.
+
+        Returns:
+            The assigned Variant with its config, or None if:
+            - User is legitimately not in traffic (server returned variantKey: null)
+
+        Raises:
+            BalancedAssignmentError: If network error, timeout, or server error occurs.
+                Caller can catch this to fall back to hash-based assignment.
+
+        Example::
+
+            # Use balanced assignment for small-sample experiments
+            variant = msk.experiments.get_variant_balanced(
+                "prompt-test-v1", user_id="user-alice-789"
+            )
+            if variant:
+                print(f"User assigned to {variant.key} (server-balanced)")
+        """
+        if not self._api_base_url:
+            logger.warning(
+                "Backend not configured — cannot use balanced assignment. "
+                "Falling back to local hash-based assignment."
+            )
+            return self.get_variant(experiment_key, user_id)
+
+        url = f"{self._api_base_url}/api/public/experiments/assign"
+        headers = self._get_auth_header()
+        payload = {
+            "experimentKey": experiment_key,
+            "userId": user_id,
+        }
+
+        try:
+            response = requests.post(url, json=payload, headers=headers, timeout=10)
+
+            if response.status_code != 200:
+                error_msg = response.text[:200] if response.text else "Unknown error"
+                logger.error(f"Server-balanced assignment failed: {response.status_code} — {error_msg}")
+                raise BalancedAssignmentError(f"Server returned {response.status_code}: {error_msg}")
+
+            data = response.json()
+
+            if data.get("variantKey") is None:
+                logger.debug(f"User '{user_id}' not in traffic for '{experiment_key}' (server response)")
+                return None
+
+            variant = Variant(
+                key=data["variantKey"],
+                weight=0,  # Weight not returned by assignment endpoint
+                config=data.get("config", {}),
+            )
+
+            # Set context for span processor
+            _set_current_experiment(
+                {
+                    "experiment_id": experiment_key,
+                    "variant_key": variant.key,
+                    "user_id": user_id,
+                }
+            )
+
+            # Log exposure
+            if self._exposure_logger:
+                self._exposure_logger.log(
+                    Exposure(
+                        user_id=user_id,
+                        experiment_id=experiment_key,
+                        variant_key=variant.key,
+                        timestamp=datetime.now(timezone.utc).isoformat(),
+                    )
+                )
+
+            is_new = data.get("isNewAssignment", True)
+            logger.debug(
+                f"Server-balanced assignment: user '{user_id}' → '{variant.key}' "
+                f"in experiment '{experiment_key}' (new={is_new})"
+            )
+
+            return variant
+
+        except requests.exceptions.Timeout:
+            logger.error(f"Server-balanced assignment timed out for '{experiment_key}'")
+            raise BalancedAssignmentError(f"Timeout for '{experiment_key}'")
+        except requests.exceptions.RequestException as e:
+            logger.error(f"Server-balanced assignment failed for '{experiment_key}': {e}")
+            raise BalancedAssignmentError(f"Network error for '{experiment_key}': {e}")
+        except BalancedAssignmentError:
+            raise  # Re-raise our custom exception
+        except Exception as e:
+            logger.error(f"Unexpected error in server-balanced assignment: {e}")
+            raise BalancedAssignmentError(f"Unexpected error for '{experiment_key}': {e}")
+
+    # -------------------------------------------------------------------------
+    # Introspection
+    # -------------------------------------------------------------------------
+
+    def is_registered(self, experiment_key: str) -> bool:
+        """Check if an experiment is registered."""
+        with self._experiments_lock:
+            return experiment_key in self._experiments
+
+    def list_experiments(self) -> List[str]:
+        """List all registered experiment keys."""
+        with self._experiments_lock:
+            return list(self._experiments.keys())
+
+
+# =============================================================================
+# Module-level instance (populated by gatewaysdk.init())
+# =============================================================================
+
+# This will be set by the main gatewaysdk module during init()
+_experiments_client: Optional[ExperimentClient] = None
+
+
+def get_experiments_client() -> Optional[ExperimentClient]:
+    """Get the global experiments client instance.
+
+    Returns None if gatewaysdk has not been initialized.
+    """
+    return _experiments_client
+
+
+def _set_experiments_client(client: Optional[ExperimentClient]) -> None:
+    """Set the global experiments client instance.
+
+    Internal function - called by gatewaysdk.init().
+    """
+    global _experiments_client
+    _experiments_client = client