feat: Add environment ID support for hooks.

contract-tests/hook.py

@@ -33,6 +33,7 @@ def __post(self, stage: str, series_context: EvaluationSeriesContext, data: dict
                 'context': series_context.context.to_dict(),
                 'defaultValue': series_context.default_value,
                 'method': series_context.method,
+                'environmentId': series_context.environment_id,
             },
             'evaluationSeriesData': data,
             'stage': stage,

contract-tests/service.py

@@ -77,6 +77,7 @@ def status():
             'instance-id',
             'anonymous-redaction',
             'evaluation-hooks',
+            'hook-environment-id',
             'omit-anonymous-contexts',
             'client-prereq-events',
             'persistent-data-store-redis',

ldclient/client.py

@@ -709,7 +709,7 @@ def __evaluate_with_hooks(self, key: str, context: Context, default_value: Any,
 
             hooks = self.__hooks.copy()
 
-        series_context = EvaluationSeriesContext(key=key, context=context, default_value=default_value, method=method)
+        series_context = EvaluationSeriesContext(key=key, context=context, default_value=default_value, method=method, environment_id=self._data_system.environment_id)
         hook_data = self.__execute_before_evaluation(hooks, series_context)
         evaluation_result = block()
         self.__execute_after_evaluation(hooks, series_context, hook_data, evaluation_result.evaluation_detail)

ldclient/hook.py

@@ -1,6 +1,6 @@
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
-from typing import Any
+from typing import Any, Optional
 
 from ldclient.context import Context
 from ldclient.evaluation import EvaluationDetail
@@ -17,6 +17,7 @@ class EvaluationSeriesContext:
     context: Context  #: The context used during evaluation.
     default_value: Any  #: The default value provided to the evaluation method
     method: str  #: The string version of the method which triggered the evaluation series.
+    environment_id: Optional[str] = None  #: The environment ID the SDK is connected to, if available.
 
 
 @dataclass

ldclient/impl/datasource/datasource_common.py

@@ -5,8 +5,9 @@
 # currently excluded from documentation - see docs/README.md
 
 from collections import namedtuple
-from typing import Optional
+from typing import Mapping, Optional, Protocol, runtime_checkable
 
+from ldclient.impl.util import _LD_ENVID_HEADER
 from ldclient.interfaces import DataSourceUpdateSink, FeatureStore
 from ldclient.versioned_data_kind import FEATURES, SEGMENTS
 
@@ -34,6 +35,32 @@ def sink_or_store(sink: Optional[DataSourceUpdateSink], store: FeatureStore):
     return sink
 
 
+@runtime_checkable
+class EnvironmentIdSink(Protocol):
+    """
+    Implemented by data source update sinks which can record the environment ID
+    reported by LaunchDarkly. This is separate from
+    :class:`ldclient.interfaces.DataSourceUpdateSink` so that externally
+    implemented sinks remain compatible.
+    """
+
+    def set_environment_id(self, environment_id: str) -> None:
+        ...
+
+
+def record_environment_id(sink, headers: Optional[Mapping[str, str]]):
+    """
+    Records the environment ID from a set of LaunchDarkly response headers, if
+    both the headers and the sink provide one.
+    """
+    if headers is None or not isinstance(sink, EnvironmentIdSink):
+        return
+
+    environment_id = headers.get(_LD_ENVID_HEADER)
+    if isinstance(environment_id, str) and environment_id != '':
+        sink.set_environment_id(environment_id)
+
+
 def parse_path(path: str):
     for kind in [FEATURES, SEGMENTS]:
         if path.startswith(kind.stream_api_path):

ldclient/impl/datasource/feature_requester.py

@@ -4,6 +4,7 @@
 
 import json
 from collections import namedtuple
+from typing import Mapping, Optional, Tuple
 from urllib import parse
 
 import urllib3
@@ -27,6 +28,10 @@ def __init__(self, config):
             self._poll_uri += '?%s' % parse.urlencode({'filter': config.payload_filter_key})
 
     def get_all_data(self):
+        (data, _) = self.get_all_data_with_headers()
+        return data
+
+    def get_all_data_with_headers(self) -> Tuple[dict, Optional[Mapping[str, str]]]:
         uri = self._poll_uri
         hdrs = _headers(self._config)
         cache_entry = self._cache.get(uri)
@@ -47,4 +52,4 @@ def get_all_data(self):
                 self._cache[uri] = CacheEntry(data=data, etag=etag)
         log.debug("%s response status:[%d] From cache? [%s] ETag:[%s]", uri, r.status, from_cache, etag)
 
-        return {FEATURES: data['flags'], SEGMENTS: data['segments']}
+        return ({FEATURES: data['flags'], SEGMENTS: data['segments']}, r.headers)

ldclient/impl/datasource/polling.py

@@ -6,10 +6,13 @@
 
 import time
 from threading import Event
-from typing import Optional
+from typing import Any, Mapping, Optional, Protocol, Tuple, runtime_checkable
 
 from ldclient.config import Config
-from ldclient.impl.datasource.datasource_common import sink_or_store
+from ldclient.impl.datasource.datasource_common import (
+    record_environment_id,
+    sink_or_store
+)
 from ldclient.impl.repeating_task import RepeatingTask
 from ldclient.impl.util import (
     UnsuccessfulResponseException,
@@ -28,6 +31,12 @@
 )
 
 
+@runtime_checkable
+class _FeatureRequesterWithHeaders(Protocol):
+    def get_all_data_with_headers(self) -> Tuple[Any, Optional[Mapping[str, str]]]:
+        ...
+
+
 class PollingUpdateProcessor(UpdateProcessor):
     def __init__(self, config: Config, requester: FeatureRequester, store: FeatureStore, ready: Event):
         self._config = config
@@ -58,7 +67,8 @@ def __stop_with_error_info(self, error: Optional[DataSourceErrorInfo]):
 
     def _poll(self):
         try:
-            all_data = self._requester.get_all_data()
+            (all_data, headers) = self._get_all_data_with_headers()
+            record_environment_id(self._data_source_update_sink, headers)
             sink_or_store(self._data_source_update_sink, self._store).init(all_data)
             if not self._ready.is_set() and self._store.initialized:
                 log.info("PollingUpdateProcessor initialized ok")
@@ -84,3 +94,13 @@ def _poll(self):
 
             if self._data_source_update_sink is not None:
                 self._data_source_update_sink.update_status(DataSourceState.INTERRUPTED, DataSourceErrorInfo(DataSourceErrorKind.UNKNOWN, 0, time.time(), str(e)))
+
+    def _get_all_data_with_headers(self) -> Tuple[Any, Optional[Mapping[str, str]]]:
+        """
+        Externally provided feature requesters are not required to surface
+        response headers, so fall back to the data-only method.
+        """
+        if isinstance(self._requester, _FeatureRequesterWithHeaders):
+            return self._requester.get_all_data_with_headers()
+
+        return (self._requester.get_all_data(), None)

ldclient/impl/datasource/status.py

@@ -26,12 +26,22 @@ def __init__(self, store: FeatureStore, status_listeners: Listeners, flag_change
 
         self.__lock = ReadWriteLock()
         self.__status = DataSourceStatus(DataSourceState.INITIALIZING, time.time(), None)
+        self.__environment_id: Optional[str] = None
 
     @property
     def status(self) -> DataSourceStatus:
         with self.__lock.read():
             return self.__status
 
+    @property
+    def environment_id(self) -> Optional[str]:
+        with self.__lock.read():
+            return self.__environment_id
+
+    def set_environment_id(self, environment_id: str) -> None:
+        with self.__lock.write():
+            self.__environment_id = environment_id
+
     def init(self, all_data: Mapping[VersionedDataKind, Mapping[str, dict]]):
         old_data = None
 

ldclient/impl/datasource/streaming.py

@@ -5,7 +5,7 @@
 from urllib import parse
 
 from ld_eventsource import SSEClient
-from ld_eventsource.actions import Event, Fault
+from ld_eventsource.actions import Event, Fault, Start
 from ld_eventsource.config import (
     ConnectStrategy,
     ErrorStrategy,
@@ -16,6 +16,7 @@
 from ldclient.impl.datasource.datasource_common import (
     STREAM_ALL_PATH,
     parse_path,
+    record_environment_id,
     sink_or_store
 )
 from ldclient.impl.http import HTTPFactory, _http_factory
@@ -62,7 +63,9 @@ def run(self):
         self._sse = self._create_sse_client()
         self._connection_attempt_start_time = time.time()
         for action in self._sse.all:
-            if isinstance(action, Event):
+            if isinstance(action, Start):
+                record_environment_id(self._data_source_update_sink, action.headers)
+            elif isinstance(action, Event):
                 message_ok = False
                 try:
                     message_ok = self._process_message(sink_or_store(self._data_source_update_sink, self._store), action)

ldclient/impl/datasystem/__init__.py

@@ -6,7 +6,7 @@
 from abc import abstractmethod
 from enum import Enum
 from threading import Event
-from typing import TYPE_CHECKING, Protocol, runtime_checkable
+from typing import TYPE_CHECKING, Optional, Protocol, runtime_checkable
 
 if TYPE_CHECKING:
     from ldclient.impl.aio.concurrency import AsyncEvent
@@ -146,6 +146,18 @@ def store(self) -> ReadOnlyStore:
         """
         raise NotImplementedError
 
+    @property
+    @abstractmethod
+    def environment_id(self) -> Optional[str]:
+        """
+        Returns the environment ID reported by LaunchDarkly, if it is known.
+
+        This is only available once a connection to LaunchDarkly has provided
+        it, and it will be None when the SDK is offline, using an unsupported
+        data source, or connected to a service which does not report it.
+        """
+        raise NotImplementedError
+
 
 class AsyncDataSystem(Protocol):
     """

ldclient/impl/datasystem/fdv1.py

@@ -98,6 +98,10 @@ def stop(self):
     def store(self) -> ReadOnlyStore:
         return self._store_wrapper
 
+    @property
+    def environment_id(self) -> Optional[str]:
+        return self._data_source_update_sink.environment_id
+
     def set_diagnostic_accumulator(self, diagnostic_accumulator: DiagnosticAccumulator):
         """
         Sets the diagnostic accumulator for streaming initialization metrics.

ldclient/impl/datasystem/fdv2.py

@@ -95,6 +95,7 @@ def __init__(
         self._lock = ReadWriteLock()
         self._active_synchronizer: Optional[Synchronizer] = None
         self._threads: List[Thread] = []
+        self._environment_id: Optional[str] = None
 
         # Track configuration
         self._configured_with_data_sources = (
@@ -240,6 +241,8 @@ def _run_initializers(self, set_on_ready: Event) -> bool:
                 basis = basis_result.value
                 log.info("Initialized via %s", initializer.name)
 
+                self._record_environment_id(basis.environment_id)
+
                 # Apply the basis to the store
                 self._store.apply(basis.change_set, basis.persist)
 
@@ -412,6 +415,9 @@ def reader(self: 'FDv2'):
                 if self._stop_event.is_set():
                     return ConditionDirective.FALLBACK
 
+                if update.state == DataSourceState.VALID:
+                    self._record_environment_id(update.environment_id)
+
                 # Handle the update
                 if update.change_set is not None:
                     self._store.apply(update.change_set, True)
@@ -496,6 +502,19 @@ def _persistent_store_outage_recovery(self, data_store_status: DataStoreStatus):
         if err is not None:
             log.error("Failed to reinitialize data store", exc_info=err)
 
+    def _record_environment_id(self, environment_id: Optional[str]):
+        if not isinstance(environment_id, str) or environment_id == '':
+            return
+
+        with self._lock.write():
+            self._environment_id = environment_id
+
+    @property
+    def environment_id(self) -> Optional[str]:
+        """Get the environment ID reported by LaunchDarkly, if known."""
+        with self._lock.read():
+            return self._environment_id
+
     @property
     def store(self) -> ReadOnlyStore:
         """Get the underlying store for flag evaluation."""

ldclient/interfaces.py

@@ -570,6 +570,12 @@ def get_all(self):
         """
         pass
 
+    def get_all_data(self) -> Mapping[VersionedDataKind, Mapping[str, dict]]:
+        """
+        Fetches all feature flag and segment data.
+        """
+        raise NotImplementedError
+
 
 class AsyncFeatureRequester(ABC):
     """

ldclient/testing/impl/datasource/test_feature_requester.py

@@ -25,6 +25,19 @@ def test_get_all_data_returns_data():
         assert result == expected_data
 
 
+def test_get_all_data_with_headers_returns_response_headers():
+    with start_server() as server:
+        config = Config(sdk_key='sdk-key', base_uri=server.uri)
+        fr = FeatureRequesterImpl(config)
+
+        resp_data = {'flags': {}, 'segments': {}}
+        server.for_path('/sdk/latest-all', JsonResponse(resp_data, {'X-LD-EnvID': 'env-abc-123'}))
+
+        (data, headers) = fr.get_all_data_with_headers()
+        assert data == {FEATURES: {}, SEGMENTS: {}}
+        assert headers.get('X-LD-EnvID') == 'env-abc-123'
+
+
 def test_get_all_data_sends_headers():
     with start_server() as server:
         config = Config(sdk_key='sdk-key', base_uri=server.uri)

ldclient/testing/impl/datasource/test_polling_processor.py

@@ -145,3 +145,38 @@ def verify_recoverable_http_error(http_status_code, ignore_mock):
         assert status.state == DataSourceState.INITIALIZING
         assert status.error.kind == DataSourceErrorKind.ERROR_RESPONSE
         assert status.error.status_code == http_status_code
+
+
+class MockFeatureRequesterWithHeaders(MockFeatureRequester):
+    def __init__(self, headers):
+        super().__init__()
+        self.headers = headers
+
+    def get_all_data_with_headers(self):
+        return (self.get_all_data(), self.headers)
+
+
+def test_records_environment_id_from_polling_headers():
+    global mock_requester
+    mock_requester = MockFeatureRequesterWithHeaders({'X-LD-EnvID': 'env-abc-123'})
+    mock_requester.all_data = {FEATURES: {}, SEGMENTS: {}}
+
+    config = Config("SDK_KEY")
+    sink = DataSourceUpdateSinkImpl(store, Listeners(), Listeners())
+    config._data_source_update_sink = sink
+    setup_processor(config)
+    assert ready.wait(2)
+
+    assert sink.environment_id == 'env-abc-123'
+
+
+def test_environment_id_is_none_when_requester_provides_no_headers():
+    mock_requester.all_data = {FEATURES: {}, SEGMENTS: {}}
+
+    config = Config("SDK_KEY")
+    sink = DataSourceUpdateSinkImpl(store, Listeners(), Listeners())
+    config._data_source_update_sink = sink
+    setup_processor(config)
+    assert ready.wait(2)
+
+    assert sink.environment_id is None