diff --git a/packages/client/agents.md b/packages/client/agents.md index b8673873..32773398 100644 --- a/packages/client/agents.md +++ b/packages/client/agents.md @@ -807,7 +807,7 @@ When `enabled` is `False`, `config` is always `None`. When `enabled` is `True` b - **Explicit initialization — SDK path.** `await init_client(options?)` dynamically imports `launchdarkly-server-sdk` at runtime (optional peer dep). If the package is not installed it raises with a clear message. - **Explicit initialization — BYOC path.** `await init_client(client)` accepts any pre-initialized object that satisfies `LDClientInterface` — this is the path for custom or edge environments whose SDK has different init semantics. - `get_client()` raises `RuntimeError` if `init_client()` has not resolved. -- `await shutdown()` must be called before process exit. It flushes OTel spans, flushes LD events, and closes the LD client. +- `await shutdown()` must be called before process exit. It flushes OTel spans, flushes LD events, and closes the LD client. It also releases the process-global OTel tracer provider, so a later `init_client()` can install its own — `trace.set_tracer_provider` is once-guarded, and leaving it set would route every later span to the provider just torn down. Only released when this SDK's own `set_tracer_provider` actually took effect — the set is once-guarded, so when another library registered first ours is refused (a warning is emitted) and that provider is left alone rather than torn down. --- diff --git a/packages/client/src/launchdarkly_ai_server/lifecycle.py b/packages/client/src/launchdarkly_ai_server/lifecycle.py index 736dd32b..b2c44111 100644 --- a/packages/client/src/launchdarkly_ai_server/lifecycle.py +++ b/packages/client/src/launchdarkly_ai_server/lifecycle.py @@ -24,6 +24,11 @@ def _env(name: str) -> str | None: _client: Any = None _tracer_provider: Any = None +# True only when *this* SDK's call to trace.set_tracer_provider actually took +# effect. "We built a provider" is not the same as "we own the global": the set +# is once-guarded, so when another library registered first ours is refused and +# the global stays theirs. Only the owner may release it on shutdown. +_owns_otel_globals: bool = False def get_client() -> Any: @@ -49,7 +54,7 @@ def _setup_telemetry(sdk_key: str, options: InitClientOptions | None = None) -> - Registers W3C trace context and baggage propagators. - Configures GZIP compression on the OTLP exporter. """ - global _tracer_provider + global _tracer_provider, _owns_otel_globals opts = options or {} @@ -113,6 +118,18 @@ def _setup_telemetry(sdk_key: str, options: InitClientOptions | None = None) -> pass trace.set_tracer_provider(provider) + # The set is refused, with a warning from OTel, when another library got + # there first. Record whether it actually took: shutdown must not release + # a global it never owned, and the caller's telemetry options are moot if + # someone else's provider is the one handing out tracers. + _owns_otel_globals = trace.get_tracer_provider() is provider + if not _owns_otel_globals: + logger.warning( + "An OpenTelemetry tracer provider was already registered by " + "something else in this process, so LaunchDarkly's telemetry " + "configuration is not in effect; spans will go wherever that " + "provider sends them." + ) _tracer_provider = provider return provider @@ -167,8 +184,13 @@ async def _resolve_client(opts: InitClientOptions, client: Any) -> Any: # BYOC path — pre-initialized client if client is not None: - _client = client + # Adopted only after telemetry setup succeeds. ``_client`` is the + # idempotency guard above, so assigning it first meant a setup that + # raised (a malformed OTEL_EXPORTER_OTLP_TIMEOUT does) left it set: the + # next call returned it as a silent success with no telemetry, hiding + # the config error. The caller owns this client, so it is not closed. _setup_telemetry(opts.get("sdkKey", "byoc"), opts) + _client = client flush_ai_sdk_info(_client) return _client @@ -214,12 +236,56 @@ async def _resolve_client(opts: InitClientOptions, client: Any) -> Any: # start_wait caps the blocking init time; matches the TS SDK's 10 s timeout. ld_client = client_cls(ld_config, start_wait=10) + # As on the BYOC path: only a fully set-up client becomes the singleton. We + # built this one, so close it on failure — it holds a streaming connection + # that would otherwise outlive the attempt. + try: + _setup_telemetry(sdk_key, opts) + except Exception: + try: + close_result = ld_client.close() + if inspect.isawaitable(close_result): + await close_result + except Exception: + pass + raise _client = ld_client - _setup_telemetry(sdk_key, opts) flush_ai_sdk_info(_client) return _client +def _release_otel_globals() -> None: + """ + Releases the process-global tracer provider that ``_setup_telemetry`` + installed, so a later ``init_client`` can install its own. + + ``trace.set_tracer_provider`` is guarded by a ``Once``: a second call logs + "Overriding of current TracerProvider is not allowed" and keeps the provider + already in place. Without this, an init/shutdown/init cycle would leave every + span routed to the provider that was already shut down, and export nothing. + + opentelemetry-python exposes no public API to unset it, so this reaches for + the module globals — both the slot and the ``Once`` that guards it, since + clearing the slot alone leaves the guard tripped and the next set a no-op. + + Callers must gate this on ``_owns_otel_globals``. Having built a provider is + not enough: when another library registered first, our set was refused and + the global is still theirs, so releasing it here would tear down the host + application's tracing and leave the global a no-op proxy. + + The global text map propagator needs no equivalent: ``set_global_textmap`` + is a plain assignment with no ``Once``, so the next setup overwrites it. + """ + try: + from opentelemetry import trace + from opentelemetry.util._once import Once + + trace._TRACER_PROVIDER = None + trace._TRACER_PROVIDER_SET_ONCE = Once() + except Exception: # pragma: no cover - defensive, OTel absent or restructured + logger.debug("Could not release the global OTel tracer provider", exc_info=True) + + async def shutdown() -> None: """ Shuts down the singleton client. Idempotent — safe to call multiple times @@ -227,24 +293,35 @@ async def shutdown() -> None: Also clears the configured skill store; pass ``skillStore`` again to the next ``init_client`` to keep using the skill accessors. + + When telemetry was running, this also releases the process-global tracer + provider so a later ``init_client`` can install its own — see + ``_release_otel_globals``. """ - global _client, _tracer_provider + global _client, _tracer_provider, _owns_otel_globals local_client = _client local_provider = _tracer_provider + owned_globals = _owns_otel_globals skills._clear_state() # Null the singleton before any awaits so a second call is a no-op _client = None _tracer_provider = None + _owns_otel_globals = False reset_ai_sdk_info() if local_provider is not None: + # Shut the provider down either way — we built it, and it owns an + # exporter and a batch timer — but only release the global registration + # when it was ours to take. try: local_provider.shutdown() except Exception: pass + if owned_globals: + _release_otel_globals() if local_client is not None: try: @@ -269,10 +346,16 @@ def _set_client_for_testing(c: Any) -> None: def _reset_for_testing() -> None: """Test helper — clear all singleton state.""" - global _client, _tracer_provider + global _client, _tracer_provider, _owns_otel_globals + owned_globals = _owns_otel_globals _client = None _tracer_provider = None + _owns_otel_globals = False skills._clear_state() + # Mirrors shutdown(): without this a suite that inits more than once leaves + # every later span on the first test's provider. + if owned_globals: + _release_otel_globals() async def inspect_config( diff --git a/packages/client/tests/test_lifecycle.py b/packages/client/tests/test_lifecycle.py index 60bf35f1..589bb8cc 100644 --- a/packages/client/tests/test_lifecycle.py +++ b/packages/client/tests/test_lifecycle.py @@ -28,6 +28,34 @@ def reset_singleton() -> None: reset_ai_sdk_info(clear_known=True) +@pytest.fixture +def restore_otel_globals() -> Any: + """Snapshot and restore the real OTel trace globals around one test. + + The autouse reset only releases them when lifecycle actually installed a + provider, so tests that set them directly have to put them back themselves + or they leak into every later test in the session. + """ + from opentelemetry import trace as otel_trace + from opentelemetry.util._once import Once + + def _install(provider: Any) -> None: + # The Once is only meaningful alongside the slot it guards: tripped iff a + # provider is installed. Restoring the *same* Once object would hand back + # one this test already tripped, so build a fresh one each time. + otel_trace._TRACER_PROVIDER = provider + once = Once() + if provider is not None: + once.do_once(lambda: None) + otel_trace._TRACER_PROVIDER_SET_ONCE = once + + saved_provider = otel_trace._TRACER_PROVIDER + # Start from a clean slate — an earlier test may have left the guard tripped. + _install(None) + yield otel_trace + _install(saved_provider) + + def _make_stub_client() -> MagicMock: stub = MagicMock() stub.variation = AsyncMock(return_value=None) @@ -73,6 +101,41 @@ async def test_get_client_returns_passed_client(self) -> None: await init_client(client=stub) assert get_client() is stub + async def test_repeat_call_does_not_rerun_telemetry_setup(self) -> None: + # The idempotency check sits ahead of this path on purpose. OTel's global + # tracer provider is once-guarded, so a second _setup_telemetry would + # build a provider that receives no spans while taking over the handle + # shutdown() flushes — silently dropping the first provider's buffer. + stub = _make_stub_client() + with patch.object( + lifecycle_module, "_setup_telemetry", return_value=None + ) as setup: + await init_client({"serviceName": "first"}, stub) + await init_client({}, stub) + await init_client(None, _make_stub_client()) + assert setup.call_count == 1 + assert setup.call_args.args[1] == {"serviceName": "first"} + + async def test_repeat_call_does_not_swap_the_client(self) -> None: + first = _make_stub_client() + second = _make_stub_client() + with patch.object(lifecycle_module, "_setup_telemetry", return_value=None): + assert await init_client(client=first) is first + assert await init_client(client=second) is first + + async def test_a_failed_telemetry_setup_does_not_adopt_the_client(self) -> None: + stub = _make_stub_client() + with patch.object( + lifecycle_module, "_setup_telemetry", side_effect=ValueError("bad config") + ): + with pytest.raises(ValueError): + await init_client(client=stub) + assert lifecycle_module._client is None + # The caller owns a BYOC client, so a failed init must not close it. + stub.close.assert_not_called() + with patch.object(lifecycle_module, "_setup_telemetry", return_value=None): + assert await init_client(client=stub) is stub + async def test_flushes_registered_ai_package_information(self) -> None: stub = _make_stub_client() register_ai_sdk_package("launchdarkly-ai-server", "0.1.3") @@ -172,6 +235,32 @@ async def test_is_idempotent(self) -> None: # LDClient is only instantiated on first init; singleton is reused assert mock_ld.LDClient.call_count == 1 + async def test_a_failed_telemetry_setup_leaves_no_client_behind( + self, restore_otel_globals: Any + ) -> None: + # A malformed OTEL_EXPORTER_OTLP_TIMEOUT makes the real _setup_telemetry + # raise. _client used to be assigned first, so the next call returned the + # half-built client as a silent success with no telemetry, and its + # connection was never closed. + failed = _make_stub_client() + healthy = _make_stub_client() + mock_ld = MagicMock() + mock_ld.Config = MagicMock(return_value=MagicMock()) + mock_ld.LDClient = MagicMock(side_effect=[failed, healthy]) + with patch("importlib.import_module", return_value=mock_ld): + bad_env = {"LD_SDK_KEY": "k", "OTEL_EXPORTER_OTLP_TIMEOUT": "soon"} + with patch.dict(os.environ, bad_env): + with pytest.raises(ValueError): + await init_client() + assert lifecycle_module._client is None + failed.close.assert_awaited_once() + + with patch.dict(os.environ, {"LD_SDK_KEY": "k"}): + with patch.object( + lifecycle_module, "_setup_telemetry", return_value=None + ): + assert await init_client() is healthy + async def test_returns_initialized_client(self) -> None: stub = _make_stub_client() mock_ld = MagicMock() @@ -258,6 +347,97 @@ async def test_allows_reinitialization(self) -> None: await init_client(client=stub2) assert get_client() is stub2 + async def test_releases_the_global_tracer_provider( + self, restore_otel_globals: Any + ) -> None: + # Without this, the next set_tracer_provider is refused and every later + # span routes to the provider shutdown() just tore down. + otel_trace = restore_otel_globals + provider = MagicMock() + with patch.object(lifecycle_module, "_setup_telemetry", return_value=provider): + await init_client(client=_make_stub_client()) + # _setup_telemetry is patched, so stand in for what it would have left, + # ownership flag included — this is the case where our set did take. + lifecycle_module._tracer_provider = provider + lifecycle_module._owns_otel_globals = True + otel_trace._TRACER_PROVIDER = provider + otel_trace._TRACER_PROVIDER_SET_ONCE.do_once(lambda: None) + + await shutdown() + + provider.shutdown.assert_called_once() + assert otel_trace._TRACER_PROVIDER is None + assert otel_trace._TRACER_PROVIDER_SET_ONCE._done is False + + async def test_a_full_init_shutdown_init_cycle_lands_on_a_live_provider( + self, restore_otel_globals: Any + ) -> None: + # End to end over the real _setup_telemetry: the second cycle's provider + # must be the one the global trace API hands out, not the dead first one. + # The OTLP exporter is stubbed so this never reaches the network. + otel_trace = restore_otel_globals + from opentelemetry.exporter.otlp.proto.http import trace_exporter + from opentelemetry.sdk.trace.export import SpanExporter, SpanExportResult + + class _NoopExporter(SpanExporter): + def __init__(self, *args: Any, **kwargs: Any) -> None: + pass + + def export(self, spans: Any) -> Any: + return SpanExportResult.SUCCESS + + def shutdown(self) -> None: + pass + + seen = [] + with patch.object(trace_exporter, "OTLPSpanExporter", _NoopExporter): + for name in ("cycle1", "cycle2"): + await init_client( + {"sdkKey": "k", "serviceName": name}, _make_stub_client() + ) + resource = otel_trace.get_tracer_provider().resource + seen.append(resource.attributes.get("service.name")) + await shutdown() + assert seen == ["cycle1", "cycle2"] + + async def test_does_not_release_a_tracer_provider_another_library_registered( + self, restore_otel_globals: Any + ) -> None: + # Building a provider is not owning the global. When something else in + # the process registered first, our set_tracer_provider is refused and + # the global stays theirs — releasing it on shutdown would tear down the + # host application's tracing and leave a no-op proxy behind. + otel_trace = restore_otel_globals + from opentelemetry.sdk.resources import Resource + from opentelemetry.sdk.trace import TracerProvider + + foreign = TracerProvider(resource=Resource.create({"service.name": "foreign"})) + otel_trace.set_tracer_provider(foreign) + + await init_client({"sdkKey": "k", "serviceName": "ld"}, _make_stub_client()) + assert otel_trace.get_tracer_provider() is foreign + assert lifecycle_module._owns_otel_globals is False + + await shutdown() + + assert otel_trace.get_tracer_provider() is foreign + assert otel_trace._TRACER_PROVIDER is foreign + + async def test_does_not_release_otel_globals_when_telemetry_never_started( + self, restore_otel_globals: Any + ) -> None: + # setup returned None (OTel absent), so those globals are not ours to + # clear — another library in the process may own them. + otel_trace = restore_otel_globals + sentinel = MagicMock() + otel_trace._TRACER_PROVIDER = sentinel + otel_trace._TRACER_PROVIDER_SET_ONCE.do_once(lambda: None) + with patch.object(lifecycle_module, "_setup_telemetry", return_value=None): + await init_client(client=_make_stub_client()) + assert lifecycle_module._owns_otel_globals is False + await shutdown() + assert otel_trace._TRACER_PROVIDER is sentinel + async def test_reemits_registered_packages_after_shutdown(self) -> None: stub1 = _make_stub_client() stub2 = _make_stub_client()