diff --git a/packages/opentelemetry-instrumentation-anthropic/opentelemetry/instrumentation/anthropic/__init__.py b/packages/opentelemetry-instrumentation-anthropic/opentelemetry/instrumentation/anthropic/__init__.py index f3096e47ba..e46713e513 100644 --- a/packages/opentelemetry-instrumentation-anthropic/opentelemetry/instrumentation/anthropic/__init__.py +++ b/packages/opentelemetry-instrumentation-anthropic/opentelemetry/instrumentation/anthropic/__init__.py @@ -13,6 +13,10 @@ emit_input_events, emit_response_events, ) +from opentelemetry.instrumentation.anthropic.retry_handler import ( + instrument_retry_emitter, + uninstrument_retry_emitter, +) from opentelemetry.instrumentation.anthropic.safety import ( _apply_completion_safety, _apply_prompt_safety, @@ -50,6 +54,7 @@ Meters, SpanAttributes, ) +from opentelemetry import trace from opentelemetry.trace import Span, SpanKind, Tracer, get_tracer from opentelemetry.trace.status import Status, StatusCode from typing_extensions import Coroutine @@ -548,8 +553,23 @@ def _wrap( kwargs = _apply_prompt_safety(span, kwargs, name) _handle_input(span, event_logger, kwargs) start_time = time.time() + # ST-10.4 (review-driven 2026-05-17): make the anthropic.chat span + # the AMBIENT OTel context for the duration of the SDK call so the + # FortifyRoot retry handler (wrapping + # ``anthropic._base_client.SyncHttpxClientWrapper.send``) can + # resolve this span as the retry_attempt parent via + # ``trace.get_current_span()``. Without this, the wrapped() + # call sees whatever ambient was active BEFORE this wrapper + # opened anthropic.chat — typically the invalid default span — + # so the retry hook's no-parent guard fires and no + # retry_attempt span is emitted. Mirrors the same pattern openai + # chat_wrapper has used since vendoring; bedrock non-streaming + # already uses ``start_as_current_span``. ``end_on_exit=False`` + # because we close the span manually below after streaming / + # response handling. try: - response = wrapped(*args, **kwargs) + with trace.use_span(span, end_on_exit=False): + response = wrapped(*args, **kwargs) except Exception as e: # pylint: disable=broad-except end_time = time.time() attributes = error_metrics_attributes(e) @@ -672,8 +692,11 @@ async def _awrap( kwargs = await asyncio.to_thread(_apply_prompt_safety, span, kwargs, name) # FR: async safety await _ahandle_input(span, event_logger, kwargs) start_time = time.time() + # ST-10.4 (review-driven 2026-05-17): see sync _wrap above for + # rationale on use_span(end_on_exit=False) around the wrapped call. try: - response = await wrapped(*args, **kwargs) + with trace.use_span(span, end_on_exit=False): + response = await wrapped(*args, **kwargs) except Exception as e: # pylint: disable=broad-except end_time = time.time() attributes = error_metrics_attributes(e) @@ -867,7 +890,16 @@ def _instrument(self, **kwargs): except Exception: pass # that's ok, we don't want to fail if some methods do not exist + # ST-10.4: per-attempt retry_attempt emission via private + # ``anthropic._base_client`` httpx wrapper classes. Guarded + # against missing private symbols (logs warning + skips emission). + # Pass the same tracer_provider the rest of the instrumentor + # uses so retry_attempt spans land in the same exporter as the + # anthropic logical span (review-driven 2026-05-16 fix). + instrument_retry_emitter(tracer_provider=tracer_provider) + def _uninstrument(self, **kwargs): + uninstrument_retry_emitter() # ST-10.4 symmetry for wrapped_method in WRAPPED_METHODS: wrap_package = wrapped_method.get("package") wrap_object = wrapped_method.get("object") diff --git a/packages/opentelemetry-instrumentation-anthropic/opentelemetry/instrumentation/anthropic/retry_handler.py b/packages/opentelemetry-instrumentation-anthropic/opentelemetry/instrumentation/anthropic/retry_handler.py new file mode 100644 index 0000000000..a94aff6d03 --- /dev/null +++ b/packages/opentelemetry-instrumentation-anthropic/opentelemetry/instrumentation/anthropic/retry_handler.py @@ -0,0 +1,558 @@ +"""ST-10.4 retry-aware emission for the Anthropic direct SDK. + +Per RETRY_LOOP.md §4.4 Anthropic row + §4.7 suppression discipline + +ST-10.0 hook-table addendum (in phase_st10_retryloop.txt): + + - Hook ``anthropic._base_client.SyncHttpxClientWrapper.send`` and + ``AsyncHttpxClientWrapper.send`` — fires once per HTTP attempt + inside the SDK's internal retry loop. Avoids global httpx + monkey-patching (which would affect every httpx caller in the + process — including non-LLM ones). + - Private-symbol guard: ``anthropic._base_client`` is private and + may shift across SDK versions. We wrap defensively; if the + private symbol is missing or incompatibly typed, log a warning + and skip retry-attempt emission for that variant. Normal anthropic + instrumentation continues unaffected. + - Endpoint allow-list: only LLM endpoints emit retry_attempt spans + (``/v1/messages`` for the modern Messages API, + ``/v1/complete`` for legacy completions). Anything else (token + refresh, model listing, etc.) does NOT emit retry_attempt — per + §4.4.1 allow-listing requirement. + - Suppression discipline (§4.7): check BOTH the OTel context + ``SUPPRESS_LANGUAGE_MODEL_INSTRUMENTATION_KEY`` AND the shared + ``is_framework_owned()`` registry. If either says "suppress", + skip emission. Prevents framework retries (LiteLLM / LangChain / + LlamaIndex) from DOUBLE-emitting when both framework wrappers + and direct-SDK wrappers are active. + - Parent resolution: use the current OTel ambient span. If invalid + or absent, skip gracefully (no orphan retry_attempt span). The + §4.5 backend dedup degrades gracefully. + +Tests: see ``tests/test_retry_attempt_emission.py``. +""" + +from __future__ import annotations + +import logging +import threading +from typing import Any, Optional + +from opentelemetry import context as context_api +from opentelemetry import trace +from opentelemetry.instrumentation.anthropic.version import __version__ +from opentelemetry.instrumentation.fortifyroot import ( + is_framework_owned, +) +from opentelemetry.instrumentation.utils import unwrap +from opentelemetry.semconv_ai import SUPPRESS_LANGUAGE_MODEL_INSTRUMENTATION_KEY +from opentelemetry.trace import SpanKind, Status, StatusCode +from wrapt import wrap_function_wrapper + +logger = logging.getLogger(__name__) + + +# ST-10 §4.4 / §4.5 constants. +_FR_RETRY_ATTEMPT_SPAN_NAME = "fortifyroot.anthropic.retry_attempt" +_FR_SPAN_ROLE_KEY = "fortifyroot.span.role" +_FR_SPAN_ROLE_RETRY_ATTEMPT = "retry_attempt" +_FR_HAS_RETRY_ATTEMPT_CHILD_KEY = "fortifyroot.span.has_retry_attempt_child" + +_ANTHROPIC_BASE_CLIENT_MODULE = "anthropic._base_client" +_SYNC_WRAPPER_CLASS = "SyncHttpxClientWrapper" +_ASYNC_WRAPPER_CLASS = "AsyncHttpxClientWrapper" +_WRAPPED_METHOD = "send" + +# §4.4.1 endpoint allow-list — match against the URL path suffix. +# The Anthropic SDK targets ``api.anthropic.com`` for the public +# endpoints and ``bedrock-runtime.*.amazonaws.com`` when routed through +# the AWS SDK; in both cases the path identifies the operation. +_LLM_PATH_SUFFIXES = ( + "/v1/messages", # modern Messages API + "/v1/complete", # legacy completions + "/messages", # generic suffix (covers /v1/messages and Azure-style paths) + "/complete", +) + + +_state_lock = threading.Lock() +_installed = False + +# Tracer provider the wrap was installed with. ``None`` → fall back to +# the global TracerProvider. Set via ``instrument_retry_emitter(tracer_provider=...)`` +# so retry spans go to the same provider as the anthropic logical span +# (e.g. ``anthropic.chat``) emitted by ``_wrap``. Without this, a +# consumer passing an explicit provider to +# ``AnthropicInstrumentor().instrument(tracer_provider=...)`` gets the +# parent span on their provider but the retry_attempt span sent into +# a no-op tracer. +_tracer_provider = None + + +def _request_path(request: Any) -> str: + try: + url = getattr(request, "url", None) + if url is None: + return "" + path = getattr(url, "path", None) + if isinstance(path, str): + return path + return str(url) + except Exception: + return "" + + +def _is_llm_endpoint(path: str) -> bool: + if not path: + return False + for suffix in _LLM_PATH_SUFFIXES: + if path.endswith(suffix): + return True + return False + + +def _operation_for_path(path: str) -> str: + if path.endswith("/complete"): + return "text_completion" + return "chat" # /v1/messages and variants + + +def _resolve_model_from_request(request: Any) -> Optional[str]: + """Best-effort: pull ``model`` from the JSON request body.""" + try: + content = getattr(request, "content", None) + if content is None or not isinstance(content, (bytes, bytearray)): + return None + import json + body = json.loads(content.decode("utf-8")) + if isinstance(body, dict): + model = body.get("model") + if isinstance(model, str) and model: + return model + except Exception: + return None + return None + + +def _server_attrs_from_request(request: Any) -> dict[str, Any]: + out: dict[str, Any] = {} + try: + url = getattr(request, "url", None) + if url is None: + return out + host = getattr(url, "host", None) + if isinstance(host, str) and host: + out["server.address"] = host + port = getattr(url, "port", None) + if port is None: + scheme = getattr(url, "scheme", None) + port = 443 if scheme == "https" else 80 if scheme == "http" else None + if isinstance(port, int): + out["server.port"] = port + except Exception: + pass + return out + + +def _is_suppressed() -> bool: + try: + if context_api.get_value(SUPPRESS_LANGUAGE_MODEL_INSTRUMENTATION_KEY): + return True + except Exception: + pass + try: + if is_framework_owned(): + return True + except Exception: + pass + return False + + +def _resolve_parent_span() -> Optional["trace.Span"]: + current = trace.get_current_span() + if current is None: + return None + ctx = current.get_span_context() + if ctx is None or not ctx.is_valid: + return None + return current + + +def _set_parent_marker(parent_span: "trace.Span") -> None: + try: + parent_span.set_attribute(_FR_HAS_RETRY_ATTEMPT_CHILD_KEY, True) + except Exception: + logger.debug("failed to set has_retry_attempt_child on parent", exc_info=True) + + +def _start_attempt_span(request: Any, parent_span: "trace.Span") -> "trace.Span": + path = _request_path(request) + attrs: dict[str, Any] = { + _FR_SPAN_ROLE_KEY: _FR_SPAN_ROLE_RETRY_ATTEMPT, + # Match the existing Anthropic instrumentor's gen_ai.system + # value ("Anthropic" title-case — see + # ``anthropic/__init__.py`` ``_wrap``). The earlier draft of + # this handler used lower-case "anthropic" for cross-handler + # consistency, but that mismatched the upstream Anthropic + # instrumentor and the fr-system-tests ProviderModel + # ``gen_ai_system="Anthropic"`` assertion — ST-10.4 + # review-driven fix 2026-05-17. Backend provider grouping is + # case-sensitive in practice; canonicalisation happens via + # ``event_provider`` (which is always lower-case in tests). + "gen_ai.system": "Anthropic", + "gen_ai.operation.name": _operation_for_path(path), + } + model = _resolve_model_from_request(request) + if model: + attrs["gen_ai.request.model"] = model + attrs.update(_server_attrs_from_request(request)) + + tracer = trace.get_tracer(__name__, __version__, _tracer_provider) + parent_ctx = trace.set_span_in_context(parent_span) + span = tracer.start_span( + _FR_RETRY_ATTEMPT_SPAN_NAME, + kind=SpanKind.CLIENT, + attributes=attrs, + context=parent_ctx, + ) + return span + + +def _finalize_success(span: "trace.Span", response: Any, *, is_streaming: bool = False) -> None: + """Apply response-side attributes + status. + + For non-streaming responses (regardless of 2xx vs non-2xx), parse + the JSON body and extract usage tokens + response id + response + model. Per RETRY_LOOP.md §4.4 token-usage rule (around line 164): + wrappers MUST extract usage from the response body whenever it's + present, regardless of whether the attempt succeeded — some + failures DO consume tokens and the provider returns usage in the + error body. §4.5 backend dedup makes the qualifying retry_attempt + canonical (even single-attempt), so usage must live on this span. + + Streaming responses (``stream=True`` passed to ``send()``) skip + body reading to avoid consuming the SSE stream before the SDK can + iterate it. Streaming usage capture is a deferred follow-up. + """ + try: + status_code = getattr(response, "status_code", None) + if isinstance(status_code, int): + span.set_attribute("http.status_code", status_code) + + # Header-based response id (works for streaming too, both success and error). + try: + headers = getattr(response, "headers", None) or {} + rid = None + if hasattr(headers, "get"): + rid = headers.get("request-id") or headers.get("x-request-id") + if isinstance(rid, str) and rid: + span.set_attribute("gen_ai.response.id", rid) + except Exception: + pass + + # Non-streaming body parse — applies to BOTH 2xx and non-2xx per the + # §4.4 token-usage rule. ``_extract_usage_from_body`` is fully + # defensive (missing fields / parse failure / no ``usage`` block all + # degrade silently). + if not is_streaming: + _extract_usage_from_body(span, response) + + if isinstance(status_code, int) and 200 <= status_code < 300: + span.set_status(Status(StatusCode.OK)) + else: + err_type = "anthropic.APIStatusError" + if isinstance(status_code, int): + if status_code == 429: + err_type = "anthropic.RateLimitError" + elif 500 <= status_code < 600: + err_type = "anthropic.InternalServerError" + elif status_code == 401: + err_type = "anthropic.AuthenticationError" + elif status_code == 403: + err_type = "anthropic.PermissionDeniedError" + span.set_attribute("error.type", err_type) + span.set_status(Status(StatusCode.ERROR, f"http {status_code}")) + except Exception: + logger.debug("failed to set response attrs on anthropic retry_attempt", exc_info=True) + + +def _extract_usage_from_body(span: "trace.Span", response: Any) -> None: + """Parse a non-streaming Anthropic response body and copy usage / + response id / response model attrs to the retry_attempt span. + + Anthropic Messages response: ``{"id": "...", "model": "...", + "usage": {"input_tokens": N, "output_tokens": M, ...}}``. Legacy + /v1/complete may not carry usage; tolerate absence. + """ + try: + body = response.json() + except Exception: + return + if not isinstance(body, dict): + return + try: + rid = body.get("id") + if isinstance(rid, str) and rid: + span.set_attribute("gen_ai.response.id", rid) + rmodel = body.get("model") + if isinstance(rmodel, str) and rmodel: + span.set_attribute("gen_ai.response.model", rmodel) + usage = body.get("usage") + if isinstance(usage, dict): + pt = usage.get("input_tokens") + ct = usage.get("output_tokens") + if isinstance(pt, int): + span.set_attribute("gen_ai.usage.input_tokens", pt) + if isinstance(ct, int): + span.set_attribute("gen_ai.usage.output_tokens", ct) + except Exception: + logger.debug("failed to extract usage from anthropic response body", exc_info=True) + + +def _finalize_error(span: "trace.Span", error: BaseException) -> None: + try: + error_type = type(error).__name__ + mod = type(error).__module__ + if isinstance(mod, str) and mod and mod != "builtins": + error_type = f"{mod.split('.')[0]}.{error_type}" + span.set_attribute("error.type", error_type) + status_code = getattr(error, "status_code", None) or getattr( + getattr(error, "response", None), "status_code", None + ) + if isinstance(status_code, int): + span.set_attribute("http.status_code", status_code) + try: + span.record_exception(error) + except Exception: + pass + span.set_status(Status(StatusCode.ERROR, str(error))) + except Exception: + logger.debug("failed to set error attrs on anthropic retry_attempt", exc_info=True) + + +def _should_emit_for(request: Any) -> bool: + if _is_suppressed(): + return False + path = _request_path(request) + if not _is_llm_endpoint(path): + return False + return True + + +def _sync_send_wrapper(wrapped, instance, args, kwargs): + """Wraps ``anthropic._base_client.SyncHttpxClientWrapper.send``. + + Direct-SDK wrappers do NOT register tokens in the §4.7.1 framework + registry (the registry's contract reserves registration for FRAMEWORK + wrappers — LiteLLM / LangChain / LlamaIndex). They only consult via + ``is_framework_owned()``. Self-registering would falsely suppress + concurrent direct-SDK calls on the same OS thread, which manifests + most visibly under asyncio (multiple tasks sharing one thread). + + Streaming skip (ST-10.4 review-driven 2026-05-17): when + ``stream=True`` is passed to send, this wrapper SKIPS retry_attempt + emission. Streaming retry_attempts cannot carry usage (SSE stream + can't be peeked) but §4.5 dedup would still promote them to the + canonical LLMUsageEvent, producing zero-token events. Leaving the + parent ``anthropic.chat`` span as the canonical (it gets full + usage from the Anthropic streaming wrapper). Streaming retry-loop + detection is the deferred follow-up + ``ST-10.4-FOLLOWUP-streaming-usage``. + """ + request = args[0] if args else kwargs.get("request") + if request is None or not _should_emit_for(request): + return wrapped(*args, **kwargs) + + if bool(kwargs.get("stream", False)): + return wrapped(*args, **kwargs) + + parent = _resolve_parent_span() + if parent is None: + return wrapped(*args, **kwargs) + + is_streaming = False + span = _start_attempt_span(request, parent) + _set_parent_marker(parent) + + try: + response = wrapped(*args, **kwargs) + except BaseException as exc: + try: + _finalize_error(span, exc) + finally: + try: + span.end() + except Exception: + pass + raise + try: + _finalize_success(span, response, is_streaming=is_streaming) + finally: + try: + span.end() + except Exception: + pass + return response + + +async def _async_send_wrapper(wrapped, instance, args, kwargs): + """Wraps ``anthropic._base_client.AsyncHttpxClientWrapper.send``. + + Same no-self-registration contract and streaming-skip as + ``_sync_send_wrapper``. + """ + request = args[0] if args else kwargs.get("request") + if request is None or not _should_emit_for(request): + return await wrapped(*args, **kwargs) + + if bool(kwargs.get("stream", False)): + return await wrapped(*args, **kwargs) + + parent = _resolve_parent_span() + if parent is None: + return await wrapped(*args, **kwargs) + + is_streaming = False + span = _start_attempt_span(request, parent) + _set_parent_marker(parent) + + try: + response = await wrapped(*args, **kwargs) + except BaseException as exc: + try: + _finalize_error(span, exc) + finally: + try: + span.end() + except Exception: + pass + raise + try: + _finalize_success(span, response, is_streaming=is_streaming) + finally: + try: + span.end() + except Exception: + pass + return response + + +def _has_wrappable_symbol(module_name: str, class_name: str, method_name: str) -> bool: + try: + import importlib + module = importlib.import_module(module_name) + except Exception: + return False + cls = getattr(module, class_name, None) + if cls is None or not isinstance(cls, type): + return False + method = getattr(cls, method_name, None) + if method is None: + return False + if not callable(method): + return False + return True + + +def instrument_retry_emitter(tracer_provider=None) -> None: + """Install the Anthropic retry_attempt emitter. Idempotent. + + Wraps: + - anthropic._base_client.SyncHttpxClientWrapper.send + - anthropic._base_client.AsyncHttpxClientWrapper.send + + ``tracer_provider`` is the TracerProvider retry spans should be + emitted through. Caller should pass the same provider the parent + ``AnthropicInstrumentor`` was configured with so retry_attempt + spans land in the same exporter chain as the anthropic logical + span. If ``None``, falls back to the global TracerProvider. + + If a private symbol is missing or incompatible, logs a warning and + skips wrapping that variant. Normal anthropic instrumentation is NOT + affected. + """ + global _installed, _tracer_provider + with _state_lock: + if _installed: + return + _tracer_provider = tracer_provider + sync_ok = _has_wrappable_symbol( + _ANTHROPIC_BASE_CLIENT_MODULE, _SYNC_WRAPPER_CLASS, _WRAPPED_METHOD, + ) + async_ok = _has_wrappable_symbol( + _ANTHROPIC_BASE_CLIENT_MODULE, _ASYNC_WRAPPER_CLASS, _WRAPPED_METHOD, + ) + if not sync_ok: + logger.warning( + "ST-10.4: anthropic._base_client.%s.%s missing/incompatible; " + "skipping sync retry_attempt emission. Normal anthropic " + "instrumentation is unaffected.", + _SYNC_WRAPPER_CLASS, _WRAPPED_METHOD, + ) + else: + try: + wrap_function_wrapper( + _ANTHROPIC_BASE_CLIENT_MODULE, + f"{_SYNC_WRAPPER_CLASS}.{_WRAPPED_METHOD}", + _sync_send_wrapper, + ) + except Exception as e: + logger.warning( + "ST-10.4: failed to wrap anthropic sync httpx send (%s); " + "retry_attempt emission disabled for sync path", + e, + ) + if not async_ok: + logger.warning( + "ST-10.4: anthropic._base_client.%s.%s missing/incompatible; " + "skipping async retry_attempt emission. Normal anthropic " + "instrumentation is unaffected.", + _ASYNC_WRAPPER_CLASS, _WRAPPED_METHOD, + ) + else: + try: + wrap_function_wrapper( + _ANTHROPIC_BASE_CLIENT_MODULE, + f"{_ASYNC_WRAPPER_CLASS}.{_WRAPPED_METHOD}", + _async_send_wrapper, + ) + except Exception as e: + logger.warning( + "ST-10.4: failed to wrap anthropic async httpx send (%s); " + "retry_attempt emission disabled for async path", + e, + ) + _installed = True + + +def uninstrument_retry_emitter() -> None: + """Remove the Anthropic retry_attempt wraps. Idempotent.""" + global _installed, _tracer_provider + with _state_lock: + if not _installed: + return + _tracer_provider = None + for cls in (_SYNC_WRAPPER_CLASS, _ASYNC_WRAPPER_CLASS): + try: + unwrap(f"{_ANTHROPIC_BASE_CLIENT_MODULE}.{cls}", _WRAPPED_METHOD) + except Exception: + logger.debug( + "ST-10.4: anthropic unwrap of %s.%s failed (likely " + "wrap was never installed for this variant)", + cls, _WRAPPED_METHOD, + exc_info=True, + ) + _installed = False + + +def _is_installed_for_test() -> bool: + with _state_lock: + return _installed + + +__all__ = [ + "instrument_retry_emitter", + "uninstrument_retry_emitter", + "_FR_RETRY_ATTEMPT_SPAN_NAME", + "_FR_HAS_RETRY_ATTEMPT_CHILD_KEY", + "_LLM_PATH_SUFFIXES", +] diff --git a/packages/opentelemetry-instrumentation-anthropic/tests/conftest.py b/packages/opentelemetry-instrumentation-anthropic/tests/conftest.py index 0e1e7926a2..e9a13450e2 100644 --- a/packages/opentelemetry-instrumentation-anthropic/tests/conftest.py +++ b/packages/opentelemetry-instrumentation-anthropic/tests/conftest.py @@ -24,9 +24,39 @@ pytest_plugins = [] +class _NoFortifyRootSpanExporter(InMemorySpanExporter): + """ST-10.4 (2026-05-17): filter out any span whose name starts + with ``fortifyroot.`` from the upstream-test span exporter. + + ST-10.4 added per-attempt ``fortifyroot.anthropic.retry_attempt`` + sibling spans under every anthropic logical call (once the + Anthropic ``_wrap`` started using ``trace.use_span`` so the retry + handler can find the parent span). Upstream / legacy Anthropic + tests assert exact span-name lists (e.g. + ``all(span.name == "anthropic.chat" for span in spans)``); without + a filter, every such assertion would now fail because the + retry_attempt sibling is also exported. + + Mirrors the OpenAI test conftest pattern (added 2026-05-16) and + the LangChain CI-hardening pattern (2026-05-15). ST-10.4 unit + tests in ``tests/test_retry_attempt_emission.py`` use their own + ``fresh_tracer`` fixture (not this one), so they continue to see + retry_attempt spans and aren't affected by the filter. + """ + + def get_finished_spans(self): # type: ignore[override] + # Filter by role rather than name prefix so legitimate + # fortifyroot.*.safety / .llm_wrapper / etc. spans remain + # visible to tests that inspect them. + return tuple( + s for s in super().get_finished_spans() + if (s.attributes or {}).get("fortifyroot.span.role") != "retry_attempt" + ) + + @pytest.fixture(scope="function", name="span_exporter") def fixture_span_exporter(): - exporter = InMemorySpanExporter() + exporter = _NoFortifyRootSpanExporter() yield exporter diff --git a/packages/opentelemetry-instrumentation-anthropic/tests/test_retry_attempt_emission.py b/packages/opentelemetry-instrumentation-anthropic/tests/test_retry_attempt_emission.py new file mode 100644 index 0000000000..a6bb6f218c --- /dev/null +++ b/packages/opentelemetry-instrumentation-anthropic/tests/test_retry_attempt_emission.py @@ -0,0 +1,651 @@ +"""Tests for ST-10.4 Anthropic direct-SDK retry-attempt emission. + +Covers (per RETRY_LOOP.md §4.4 Anthropic row + §4.7 suppression): + - Instrumentor symmetry: install/uninstall flips state cleanly and is + idempotent. + - Single-attempt happy path: ONE retry_attempt span under the active + parent; parent gets has_retry_attempt_child=true; span carries + role / gen_ai.system=anthropic / gen_ai.request.model / http.status_code. + - Multi-attempt retry path: N siblings under one parent. + - Error-attempt: status=ERROR, http.status_code, error.type. + - No-parent guard. + - §4.7 suppression — context-API and framework-registry both skip. + - §4.4.1 endpoint allow-list: non-LLM SDK traffic does NOT emit. + - Private-symbol missing guard: warning + skip, no crash. + +UNIT tests — drive ``_sync_send_wrapper`` / ``_async_send_wrapper`` +directly with fake httpx Request / Response objects. +""" + +from __future__ import annotations + +import asyncio +from types import SimpleNamespace +from typing import Optional +from unittest.mock import patch + +import pytest +from opentelemetry import context as context_api +from opentelemetry import trace +from opentelemetry.instrumentation.anthropic.retry_handler import ( + _FR_HAS_RETRY_ATTEMPT_CHILD_KEY, + _FR_RETRY_ATTEMPT_SPAN_NAME, + _async_send_wrapper, + _has_wrappable_symbol, + _is_installed_for_test, + _sync_send_wrapper, + instrument_retry_emitter, + uninstrument_retry_emitter, +) +from opentelemetry.instrumentation.fortifyroot import ( + is_framework_owned, + register_framework_attempt, + retry_registry, + unregister_framework_attempt, +) +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import ( + InMemorySpanExporter, +) +from opentelemetry.semconv_ai import SUPPRESS_LANGUAGE_MODEL_INSTRUMENTATION_KEY + + +@pytest.fixture +def fresh_tracer(): + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + try: + trace.set_tracer_provider(provider) + except Exception: + pass + current = trace.get_tracer_provider() + if current is not provider: + try: + current.add_span_processor(SimpleSpanProcessor(exporter)) + except Exception: + pass + yield trace.get_tracer("test"), exporter, current + + +@pytest.fixture(autouse=True) +def reset_state(): + retry_registry._reset_for_test() + try: + uninstrument_retry_emitter() + except Exception: + pass + yield + retry_registry._reset_for_test() + try: + uninstrument_retry_emitter() + except Exception: + pass + + +def _make_request(path: str = "/v1/messages", host: str = "api.anthropic.com", + port: Optional[int] = None, scheme: str = "https", + model: Optional[str] = "claude-haiku-4-5") -> SimpleNamespace: + import json + body = {"model": model, "messages": [{"role": "user", "content": "hi"}], "max_tokens": 10} if model else {} + content = json.dumps(body).encode("utf-8") + url = SimpleNamespace(path=path, host=host, port=port, scheme=scheme) + return SimpleNamespace(url=url, content=content) + + +def _make_response(status_code: int = 200, request_id: str = "req-abc", + body: Optional[dict] = None) -> SimpleNamespace: + """Fake ``httpx.Response``. Non-streaming retry_attempt path now + parses the body for usage / response.id / response.model — see the + 2026-05-13 review-driven C2 fix. + """ + if body is None: + body = { + "id": f"msg_{request_id}", + "model": "claude-haiku-4-5-20251001", + "usage": {"input_tokens": 9, "output_tokens": 4}, + } + headers = {"request-id": request_id, "x-request-id": request_id} + return SimpleNamespace( + status_code=status_code, + headers=headers, + json=lambda: body, + ) + + +def _retry_spans(exporter: InMemorySpanExporter): + return [s for s in exporter.get_finished_spans() + if s.name == _FR_RETRY_ATTEMPT_SPAN_NAME] + + +# --------------------------------------------------------------------------- +# Instrumentor symmetry. +# --------------------------------------------------------------------------- + +def test_install_then_uninstall_is_idempotent_and_symmetric(): + assert not _is_installed_for_test() + + instrument_retry_emitter() + assert _is_installed_for_test() + + instrument_retry_emitter() # idempotent + assert _is_installed_for_test() + + uninstrument_retry_emitter() + assert not _is_installed_for_test() + + uninstrument_retry_emitter() # idempotent + assert not _is_installed_for_test() + + +def test_install_actually_wraps_the_anthropic_private_send(): + """REGRESSION GUARD: install must wrap + ``anthropic._base_client.SyncHttpxClientWrapper.send``. Without this + check the install path could silently no-op.""" + pytest.importorskip("anthropic") + from anthropic import _base_client as base + + if not hasattr(base, "SyncHttpxClientWrapper") or not hasattr( + base.SyncHttpxClientWrapper, "send" + ): + pytest.skip("anthropic._base_client.SyncHttpxClientWrapper.send not present") + + instrument_retry_emitter() + try: + send = base.SyncHttpxClientWrapper.send + assert hasattr(send, "__wrapped__"), ( + "after instrument, SyncHttpxClientWrapper.send must be a wrapt wrapper" + ) + finally: + uninstrument_retry_emitter() + + send_after = base.SyncHttpxClientWrapper.send + assert not hasattr(send_after, "__wrapped__") + + +# --------------------------------------------------------------------------- +# Private-symbol guard. +# --------------------------------------------------------------------------- + +def test_missing_private_symbol_logs_warning_and_does_not_crash(caplog): + import opentelemetry.instrumentation.anthropic.retry_handler as rh + + with patch.object(rh, "_has_wrappable_symbol", return_value=False): + with caplog.at_level("WARNING"): + instrument_retry_emitter() # must not raise + assert _is_installed_for_test() + msgs = " ".join(r.getMessage() for r in caplog.records) + assert "missing/incompatible" in msgs + uninstrument_retry_emitter() + + +def test_has_wrappable_symbol_returns_true_for_real_anthropic(): + pytest.importorskip("anthropic") + assert _has_wrappable_symbol("anthropic._base_client", "SyncHttpxClientWrapper", "send") + assert _has_wrappable_symbol("anthropic._base_client", "AsyncHttpxClientWrapper", "send") + + +def test_has_wrappable_symbol_returns_false_for_unknown(): + assert not _has_wrappable_symbol("anthropic._base_client", "NoSuchClass", "send") + assert not _has_wrappable_symbol("anthropic._base_client", "SyncHttpxClientWrapper", "no_such_method") + + +# --------------------------------------------------------------------------- +# Single-attempt happy path. +# --------------------------------------------------------------------------- + +def test_single_attempt_emits_one_span_with_marker(fresh_tracer): + tracer, exporter, _ = fresh_tracer + request = _make_request(model="claude-haiku-4-5", path="/v1/messages") + response = _make_response(status_code=200, request_id="req-1") + + parent = tracer.start_span("anthropic.chat") + with trace.use_span(parent, end_on_exit=False): + result = _sync_send_wrapper(lambda *a, **kw: response, None, (request,), {}) + parent.end() + + assert result is response + + spans = exporter.get_finished_spans() + parent_exported = next(s for s in spans if s.name == "anthropic.chat") + retry_spans = _retry_spans(exporter) + assert len(retry_spans) == 1 + rs = retry_spans[0] + assert rs.parent.span_id == parent_exported.context.span_id + assert rs.attributes.get("fortifyroot.span.role") == "retry_attempt" + assert rs.attributes.get("gen_ai.system") == "Anthropic" + assert rs.attributes.get("gen_ai.request.model") == "claude-haiku-4-5" + assert rs.attributes.get("gen_ai.operation.name") == "chat" + assert rs.attributes.get("http.status_code") == 200 + # Response-body-derived attrs (replaces header-only id with the body's msg_*). + assert rs.attributes.get("gen_ai.response.id") == "msg_req-1" + assert rs.attributes.get("gen_ai.response.model") == "claude-haiku-4-5-20251001" + assert rs.attributes.get("gen_ai.usage.input_tokens") == 9 + assert rs.attributes.get("gen_ai.usage.output_tokens") == 4 + assert rs.attributes.get("server.address") == "api.anthropic.com" + assert rs.attributes.get("server.port") == 443 + assert parent_exported.attributes.get(_FR_HAS_RETRY_ATTEMPT_CHILD_KEY) is True + + +# --------------------------------------------------------------------------- +# Multi-attempt retry path. +# --------------------------------------------------------------------------- + +def test_three_attempts_share_parent(fresh_tracer): + tracer, exporter, _ = fresh_tracer + request = _make_request(model="claude-haiku-4-5") + resp_429 = _make_response(status_code=429) + resp_200 = _make_response(status_code=200, request_id="r-ok") + + parent = tracer.start_span("anthropic.chat") + with trace.use_span(parent, end_on_exit=False): + _sync_send_wrapper(lambda *a, **kw: resp_429, None, (request,), {}) + _sync_send_wrapper(lambda *a, **kw: resp_429, None, (request,), {}) + _sync_send_wrapper(lambda *a, **kw: resp_200, None, (request,), {}) + parent.end() + + spans = exporter.get_finished_spans() + parent_exported = next(s for s in spans if s.name == "anthropic.chat") + retry_spans = _retry_spans(exporter) + assert len(retry_spans) == 3 + + parent_ids = {s.parent.span_id for s in retry_spans} + assert parent_ids == {parent_exported.context.span_id} + + from opentelemetry.trace import StatusCode + error_count = sum(1 for s in retry_spans if s.status.status_code == StatusCode.ERROR) + ok_count = sum(1 for s in retry_spans if s.status.status_code == StatusCode.OK) + assert error_count == 2 and ok_count == 1 + + for s in retry_spans: + if s.status.status_code == StatusCode.ERROR: + assert s.attributes.get("http.status_code") == 429 + assert s.attributes.get("error.type") == "anthropic.RateLimitError" + + +# --------------------------------------------------------------------------- +# Error / exception path. +# --------------------------------------------------------------------------- + +def test_exception_path_records_error_and_reraises(fresh_tracer): + tracer, exporter, _ = fresh_tracer + request = _make_request(model="claude-haiku-4-5") + + class ConnectError(Exception): + pass + + err = ConnectError("connect failed") + + parent = tracer.start_span("anthropic.chat") + with trace.use_span(parent, end_on_exit=False): + with pytest.raises(ConnectError): + _sync_send_wrapper(lambda *a, **kw: (_ for _ in ()).throw(err), None, (request,), {}) + parent.end() + + from opentelemetry.trace import StatusCode + retry_spans = _retry_spans(exporter) + assert len(retry_spans) == 1 + rs = retry_spans[0] + assert rs.status.status_code == StatusCode.ERROR + assert "ConnectError" in (rs.attributes.get("error.type") or "") + + +# --------------------------------------------------------------------------- +# Marker timing (§4.5). +# --------------------------------------------------------------------------- + +def test_marker_set_AFTER_first_attempt_not_at_parent_creation(fresh_tracer): + tracer, exporter, _ = fresh_tracer + request = _make_request() + parent = tracer.start_span("anthropic.chat") + assert _FR_HAS_RETRY_ATTEMPT_CHILD_KEY not in dict(parent.attributes or {}) + with trace.use_span(parent, end_on_exit=False): + _sync_send_wrapper(lambda *a, **kw: _make_response(200), None, (request,), {}) + assert dict(parent.attributes or {}).get(_FR_HAS_RETRY_ATTEMPT_CHILD_KEY) is True + parent.end() + + +def test_marker_NOT_set_when_no_attempts_fire(fresh_tracer): + tracer, exporter, _ = fresh_tracer + parent = tracer.start_span("anthropic.chat") + parent.end() + parent_exported = next( + s for s in exporter.get_finished_spans() if s.name == "anthropic.chat" + ) + assert _FR_HAS_RETRY_ATTEMPT_CHILD_KEY not in (parent_exported.attributes or {}) + + +# --------------------------------------------------------------------------- +# No-parent guard. +# --------------------------------------------------------------------------- + +def test_no_parent_does_not_emit(fresh_tracer): + tracer, exporter, _ = fresh_tracer + request = _make_request() + result = _sync_send_wrapper(lambda *a, **kw: _make_response(200), None, (request,), {}) + assert result.status_code == 200 + assert len(_retry_spans(exporter)) == 0 + + +# --------------------------------------------------------------------------- +# §4.7 suppression. +# --------------------------------------------------------------------------- + +def test_context_suppression_skips_emission(fresh_tracer): + tracer, exporter, _ = fresh_tracer + request = _make_request() + parent = tracer.start_span("anthropic.chat") + token = context_api.attach( + context_api.set_value(SUPPRESS_LANGUAGE_MODEL_INSTRUMENTATION_KEY, True) + ) + try: + with trace.use_span(parent, end_on_exit=False): + result = _sync_send_wrapper(lambda *a, **kw: _make_response(200), None, (request,), {}) + finally: + context_api.detach(token) + parent.end() + assert result.status_code == 200 + assert len(_retry_spans(exporter)) == 0 + + +def test_framework_registry_suppression_skips_emission(fresh_tracer): + tracer, exporter, _ = fresh_tracer + request = _make_request() + parent = tracer.start_span("anthropic.chat") + tok = register_framework_attempt() + assert is_framework_owned() + try: + with trace.use_span(parent, end_on_exit=False): + result = _sync_send_wrapper(lambda *a, **kw: _make_response(200), None, (request,), {}) + finally: + unregister_framework_attempt(tok) + parent.end() + assert result.status_code == 200 + assert len(_retry_spans(exporter)) == 0 + + +def test_direct_sdk_wrapper_does_not_self_register_in_framework_registry(fresh_tracer): + """REGRESSION GUARD (review-driven fix 2026-05-13): direct-SDK + wrappers MUST NOT register tokens in the §4.7.1 framework registry. + Self-registration causes false suppression of concurrent direct-SDK + calls sharing the same thread (asyncio repro in + ``test_two_concurrent_async_sends_each_emit_a_retry_attempt``).""" + tracer, _, _ = fresh_tracer + request = _make_request() + + seen_during: list[bool] = [] + + def wrapped(*a, **kw): + seen_during.append(is_framework_owned()) + return _make_response(200) + + parent = tracer.start_span("anthropic.chat") + with trace.use_span(parent, end_on_exit=False): + assert not is_framework_owned() + _sync_send_wrapper(wrapped, None, (request,), {}) + assert not is_framework_owned() + parent.end() + assert seen_during == [False], ( + "direct-SDK wrapper must NOT self-register; got is_framework_owned() == True " + "during the wrapped send, which would falsely suppress concurrent calls" + ) + + +def test_two_concurrent_async_sends_each_emit_a_retry_attempt(fresh_tracer): + """REGRESSION GUARD (review-driven fix 2026-05-13): two asyncio + tasks sharing one OS thread MUST each emit their own retry_attempt.""" + tracer, exporter, _ = fresh_tracer + + request = _make_request(path="/v1/messages", model="claude-haiku-4-5") + parent = tracer.start_span("anthropic.chat") + + async def one_send(): + async def wrapped(*a, **kw): + await asyncio.sleep(0) + return _make_response(200) + with trace.use_span(parent, end_on_exit=False): + return await _async_send_wrapper(wrapped, None, (request,), {}) + + async def run(): + return await asyncio.gather(one_send(), one_send()) + + loop = asyncio.new_event_loop() + try: + results = loop.run_until_complete(run()) + finally: + loop.close() + parent.end() + + assert all(r.status_code == 200 for r in results) + retry_spans = _retry_spans(exporter) + assert len(retry_spans) == 2, ( + f"expected 2 retry_attempt spans, got {len(retry_spans)} — async self-suppression regression" + ) + + +# --------------------------------------------------------------------------- +# §4.4.1 endpoint allow-list. +# --------------------------------------------------------------------------- + +@pytest.mark.parametrize("path", [ + "/v1/models", + "/v1/organizations/usage", + "/", + "", +]) +def test_non_llm_endpoints_do_not_emit(fresh_tracer, path): + tracer, exporter, _ = fresh_tracer + request = _make_request(path=path, model=None) + parent = tracer.start_span("anthropic.chat") + with trace.use_span(parent, end_on_exit=False): + _sync_send_wrapper(lambda *a, **kw: _make_response(200), None, (request,), {}) + parent.end() + assert len(_retry_spans(exporter)) == 0, f"non-LLM endpoint {path!r} must not emit" + + +@pytest.mark.parametrize("path,expected_op", [ + ("/v1/messages", "chat"), + ("/v1/complete", "text_completion"), +]) +def test_llm_endpoints_emit_with_correct_operation(fresh_tracer, path, expected_op): + tracer, exporter, _ = fresh_tracer + request = _make_request(path=path, model="claude-haiku-4-5") + parent = tracer.start_span("anthropic.chat") + with trace.use_span(parent, end_on_exit=False): + _sync_send_wrapper(lambda *a, **kw: _make_response(200), None, (request,), {}) + parent.end() + retry_spans = _retry_spans(exporter) + assert len(retry_spans) == 1 + assert retry_spans[0].attributes.get("gen_ai.operation.name") == expected_op + + +# --------------------------------------------------------------------------- +# §4.5-driven usage-extraction policy (review-driven fix 2026-05-13). +# --------------------------------------------------------------------------- + +def test_non_streaming_response_body_populates_usage_id_and_model(fresh_tracer): + """Backend dedup makes the retry_attempt span the canonical + LLMUsageEvent (see ``proc_llm_extractor.go``). Anthropic non-streaming + responses MUST carry usage on the retry_attempt span.""" + tracer, exporter, _ = fresh_tracer + request = _make_request(model="claude-haiku-4-5") + response = _make_response( + status_code=200, + request_id="r-1", + body={ + "id": "msg_01ABC", + "model": "claude-haiku-4-5-20251001", + "usage": {"input_tokens": 41, "output_tokens": 17}, + }, + ) + parent = tracer.start_span("anthropic.chat") + with trace.use_span(parent, end_on_exit=False): + _sync_send_wrapper(lambda *a, **kw: response, None, (request,), {}) + parent.end() + + retry_spans = _retry_spans(exporter) + assert len(retry_spans) == 1 + rs = retry_spans[0] + assert rs.attributes.get("gen_ai.usage.input_tokens") == 41 + assert rs.attributes.get("gen_ai.usage.output_tokens") == 17 + assert rs.attributes.get("gen_ai.response.id") == "msg_01ABC" + assert rs.attributes.get("gen_ai.response.model") == "claude-haiku-4-5-20251001" + + +def test_streaming_request_skips_emission_entirely(fresh_tracer): + """ST-10.4 (review-driven 2026-05-17): when ``stream=True`` is + passed to send, the wrap SKIPS retry_attempt emission entirely. + See the openai analog for full rationale. The parent + ``anthropic.chat`` span stays the canonical LLMUsageEvent. + """ + tracer, exporter, _ = fresh_tracer + request = _make_request(model="claude-haiku-4-5") + + json_called = {"n": 0} + + def explode(): + json_called["n"] += 1 + raise RuntimeError("body must not be read on streaming response") + + headers = {"request-id": "req-stream-1", "x-request-id": "req-stream-1"} + response = SimpleNamespace(status_code=200, headers=headers, json=explode) + + parent = tracer.start_span("anthropic.chat") + with trace.use_span(parent, end_on_exit=False): + result = _sync_send_wrapper( + lambda *a, **kw: response, + None, + (request,), + {"stream": True}, + ) + parent.end() + + assert result is response + assert json_called["n"] == 0 + retry_spans = _retry_spans(exporter) + assert len(retry_spans) == 0, ( + f"streaming attempts MUST skip retry_attempt emission entirely; " + f"got {len(retry_spans)} span(s)." + ) + parent_exported = next( + s for s in exporter.get_finished_spans() if s.name == "anthropic.chat" + ) + assert _FR_HAS_RETRY_ATTEMPT_CHILD_KEY not in (parent_exported.attributes or {}) + + +def test_non_2xx_response_with_usage_in_body_populates_usage_tokens(fresh_tracer): + """REGRESSION GUARD (review-driven follow-up 2026-05-13): per + RETRY_LOOP.md §4.4 token-usage rule, wrappers MUST extract usage + from the response body whenever present, regardless of success. + Some failures consume tokens and the provider returns usage in + the error body.""" + tracer, exporter, _ = fresh_tracer + request = _make_request(model="claude-haiku-4-5") + # Anthropic-shaped 400 with usage (hypothetical context-length scenario). + response = _make_response( + status_code=400, + request_id="r-ctx-exceeded", + body={ + "type": "error", + "error": {"type": "invalid_request_error", "message": "context too long"}, + "id": "msg_ERR", + "model": "claude-haiku-4-5-20251001", + "usage": {"input_tokens": 200000, "output_tokens": 0}, + }, + ) + parent = tracer.start_span("anthropic.chat") + with trace.use_span(parent, end_on_exit=False): + _sync_send_wrapper(lambda *a, **kw: response, None, (request,), {}) + parent.end() + + from opentelemetry.trace import StatusCode + retry_spans = _retry_spans(exporter) + assert len(retry_spans) == 1 + rs = retry_spans[0] + assert rs.status.status_code == StatusCode.ERROR + assert rs.attributes.get("http.status_code") == 400 + assert rs.attributes.get("error.type") == "anthropic.APIStatusError" + assert rs.attributes.get("gen_ai.usage.input_tokens") == 200000 + assert rs.attributes.get("gen_ai.usage.output_tokens") == 0 + assert rs.attributes.get("gen_ai.response.id") == "msg_ERR" + + +def test_non_2xx_response_without_usage_in_body_omits_usage_tokens(fresh_tracer): + """Counterpart: error body without ``usage`` must leave the span's + usage attrs unset (per §4.4: omit when unknown).""" + tracer, exporter, _ = fresh_tracer + request = _make_request(model="claude-haiku-4-5") + response = _make_response( + status_code=500, + request_id="r-server-err", + body={"type": "error", "error": {"type": "api_error", "message": "boom"}}, + ) + parent = tracer.start_span("anthropic.chat") + with trace.use_span(parent, end_on_exit=False): + _sync_send_wrapper(lambda *a, **kw: response, None, (request,), {}) + parent.end() + + from opentelemetry.trace import StatusCode + retry_spans = _retry_spans(exporter) + assert len(retry_spans) == 1 + rs = retry_spans[0] + assert rs.status.status_code == StatusCode.ERROR + assert rs.attributes.get("http.status_code") == 500 + assert rs.attributes.get("error.type") == "anthropic.InternalServerError" + assert rs.attributes.get("gen_ai.usage.input_tokens") is None + assert rs.attributes.get("gen_ai.usage.output_tokens") is None + + +def test_non_streaming_response_with_malformed_body_does_not_crash(fresh_tracer): + """Body parse failure (non-JSON / malformed) must NOT raise.""" + tracer, exporter, _ = fresh_tracer + request = _make_request(model="claude-haiku-4-5") + + def explode(): + raise ValueError("malformed JSON") + + headers = {"request-id": "req-bad-body"} + response = SimpleNamespace(status_code=200, headers=headers, json=explode) + + parent = tracer.start_span("anthropic.chat") + with trace.use_span(parent, end_on_exit=False): + _sync_send_wrapper(lambda *a, **kw: response, None, (request,), {}) + parent.end() + + retry_spans = _retry_spans(exporter) + assert len(retry_spans) == 1 + from opentelemetry.trace import StatusCode + assert retry_spans[0].status.status_code == StatusCode.OK + assert retry_spans[0].attributes.get("http.status_code") == 200 + assert retry_spans[0].attributes.get("gen_ai.response.id") == "req-bad-body" + + +# --------------------------------------------------------------------------- +# Async path. +# --------------------------------------------------------------------------- + +def test_async_send_wrapper_emits(fresh_tracer): + tracer, exporter, _ = fresh_tracer + request = _make_request() + + async def wrapped(*a, **kw): + return _make_response(200) + + async def run(): + parent = tracer.start_span("anthropic.chat") + try: + with trace.use_span(parent, end_on_exit=False): + return await _async_send_wrapper(wrapped, None, (request,), {}) + finally: + parent.end() + + loop = asyncio.new_event_loop() + try: + result = loop.run_until_complete(run()) + finally: + loop.close() + assert result.status_code == 200 + retry_spans = _retry_spans(exporter) + assert len(retry_spans) == 1 + assert retry_spans[0].attributes.get("gen_ai.system") == "Anthropic" diff --git a/packages/opentelemetry-instrumentation-bedrock/opentelemetry/instrumentation/bedrock/__init__.py b/packages/opentelemetry-instrumentation-bedrock/opentelemetry/instrumentation/bedrock/__init__.py index acd1e762a2..3a60d38297 100644 --- a/packages/opentelemetry-instrumentation-bedrock/opentelemetry/instrumentation/bedrock/__init__.py +++ b/packages/opentelemetry-instrumentation-bedrock/opentelemetry/instrumentation/bedrock/__init__.py @@ -22,6 +22,9 @@ guardrail_handling, ) from opentelemetry.instrumentation.bedrock.prompt_caching import prompt_caching_handling +from opentelemetry.instrumentation.bedrock.retry_handler import ( + install_event_hooks_on_client, +) from opentelemetry.instrumentation.bedrock.reusable_streaming_body import ( ReusableStreamingBody, ) @@ -60,6 +63,7 @@ SUPPRESS_LANGUAGE_MODEL_INSTRUMENTATION_KEY, Meters, ) +from opentelemetry import trace from opentelemetry.trace import Span, SpanKind, get_tracer from wrapt import wrap_function_wrapper @@ -121,13 +125,21 @@ def is_metrics_enabled() -> bool: def _with_tracer_wrapper(func): - """Helper for providing tracer for wrapper functions.""" + """Helper for providing tracer for wrapper functions. + + ``tracer_provider`` (added 2026-05-16 for the ST-10.4 review-driven + fix) is plumbed alongside the tracer so the bedrock-runtime client + wrap can pass it to ``install_event_hooks_on_client`` — letting the + retry handler's event hooks emit spans through the same provider + the rest of the instrumentor uses. + """ def _with_tracer( tracer, metric_params, event_logger, to_wrap, + tracer_provider=None, ): def wrapper(wrapped, instance, args, kwargs): return func( @@ -135,6 +147,7 @@ def wrapper(wrapped, instance, args, kwargs): metric_params, event_logger, to_wrap, + tracer_provider, wrapped, instance, args, @@ -152,6 +165,7 @@ def _wrap( metric_params, event_logger, to_wrap, + tracer_provider, wrapped, instance, args, @@ -183,6 +197,13 @@ def _wrap( client.converse_stream = _instrumented_converse_stream( client.converse_stream, tracer, metric_params, event_logger ) + # ST-10.4: register per-attempt botocore event hooks on the + # bedrock-runtime client. Public botocore API; emits one + # retry_attempt sibling span per HTTP attempt under the outer + # bedrock.completion / bedrock.converse span. ``tracer_provider`` + # is passed so retry_attempt spans go to the same provider as + # the bedrock logical span — review-driven 2026-05-16 fix. + install_event_hooks_on_client(client, tracer_provider=tracer_provider) return client except Exception as e: end_time = time.time() @@ -229,8 +250,17 @@ def with_instrumentation(*args, **kwargs): span = tracer.start_span(_BEDROCK_INVOKE_SPAN_NAME, kind=SpanKind.CLIENT) + # ST-10.4: make the streaming span the AMBIENT OTel context for + # the duration of the underlying boto3 call so per-attempt + # botocore event hooks (before-send.bedrock-runtime.*) can + # resolve this span as the retry_attempt parent. The span is + # NOT ended here — the streaming wrapper installed in + # ``_handle_stream_call`` owns the eventual ``span.end()`` once + # the response stream is fully consumed (which is why + # ``start_as_current_span`` cannot be used directly). kwargs = _apply_invoke_prompt_safety(span, kwargs, _BEDROCK_INVOKE_SPAN_NAME) - response = fn(*args, **kwargs) + with trace.use_span(span, end_on_exit=False): + response = fn(*args, **kwargs) _handle_stream_call(span, kwargs, response, metric_params, event_logger) return response @@ -267,7 +297,10 @@ def with_instrumentation(*args, **kwargs): span = tracer.start_span(_BEDROCK_CONVERSE_SPAN_NAME, kind=SpanKind.CLIENT) kwargs = _apply_converse_prompt_safety(span, kwargs, _BEDROCK_CONVERSE_SPAN_NAME) - response = fn(*args, **kwargs) + # ST-10.4: see _instrumented_model_invoke_with_response_stream + # for the rationale on use_span(end_on_exit=False). + with trace.use_span(span, end_on_exit=False): + response = fn(*args, **kwargs) if span.is_recording(): _handle_converse_stream(span, kwargs, response, metric_params, event_logger) @@ -631,6 +664,7 @@ def _instrument(self, **kwargs): metric_params, event_logger, wrapped_method, + tracer_provider=tracer_provider, ), ) diff --git a/packages/opentelemetry-instrumentation-bedrock/opentelemetry/instrumentation/bedrock/retry_handler.py b/packages/opentelemetry-instrumentation-bedrock/opentelemetry/instrumentation/bedrock/retry_handler.py new file mode 100644 index 0000000000..7fee3d8df1 --- /dev/null +++ b/packages/opentelemetry-instrumentation-bedrock/opentelemetry/instrumentation/bedrock/retry_handler.py @@ -0,0 +1,542 @@ +"""ST-10.4 retry-aware emission for the Bedrock (botocore) direct SDK. + +Per RETRY_LOOP.md §4.4 Bedrock row + §4.7 suppression discipline + +ST-10.0 hook-table addendum (in phase_st10_retryloop.txt): + + - Use botocore's PUBLIC event-hook API on bedrock-runtime clients: + * ``before-send.bedrock-runtime.*`` fires once per HTTP attempt + (verified in botocore 1.35.x ``endpoint.py:_do_get_response``) + with the prepared ``request`` as a keyword arg. + * ``response-received.bedrock-runtime.*`` fires when the attempt + completes (success OR error) with ``http_response``, ``parsed``, + ``exception``, and ``context`` kwargs. + - We DO NOT monkey-patch botocore internals. Event hooks are + documented public API, low version-fragility. The earlier draft + referenced ``after-call.*`` — that event does NOT exist in current + botocore for per-attempt callbacks and has been removed. + - Per-attempt correlation: store the started span on + ``request.context`` under ``_CTX_SPAN_KEY`` (a per-request mutable + dict botocore propagates end-to-end). ``before-send`` writes; + ``response-received`` reads via the ``context`` kwarg, which + references the SAME dict. This avoids thread-local state and + survives async/aiobotocore variations cleanly. NB: we deliberately + do NOT register a framework-attempt token here — the §4.7.1 + registry's documented contract reserves registration for FRAMEWORK + wrappers (LiteLLM / LangChain / LlamaIndex). Direct-SDK wrappers + only CONSULT via ``is_framework_owned()``. (See the 2026-05-13 + review-driven C1 fix; an earlier draft of this module incorrectly + registered a token.) + - Endpoint allow-list: implicit via the event-name pattern + ``bedrock-runtime.*``. Only bedrock-runtime operations (Invoke / + Converse and their streaming variants) fire these events; non-LLM + AWS service traffic (S3, STS, etc.) is unaffected. + - Suppression discipline (§4.7): before emitting, check BOTH the + OTel context ``SUPPRESS_LANGUAGE_MODEL_INSTRUMENTATION_KEY`` AND + the shared ``is_framework_owned()`` registry (consult-only — see + above). + - Parent resolution: use the current OTel ambient span (the existing + Bedrock outer span ``bedrock.completion`` / ``bedrock.converse`` + when the user calls through the FR-wrapped Bedrock client). If no + valid ambient parent exists, skip emission. + +This module exports ``install_event_hooks_on_client(client)`` which the +Bedrock instrumentor calls inside its existing ``_wrap`` path when a +bedrock-runtime client is created. ``uninstall_event_hooks_on_client(client)`` +is provided for symmetry / tests, though in practice clients are +short-lived and don't need explicit unhook. + +Tests: see ``tests/test_retry_attempt_emission.py``. +""" + +from __future__ import annotations + +import logging +import threading +from typing import Any, Optional + +from opentelemetry import context as context_api +from opentelemetry import trace +from opentelemetry.instrumentation.bedrock.version import __version__ +from opentelemetry.instrumentation.fortifyroot import ( + is_framework_owned, +) +from opentelemetry.semconv_ai import SUPPRESS_LANGUAGE_MODEL_INSTRUMENTATION_KEY +from opentelemetry.trace import SpanKind, Status, StatusCode + +logger = logging.getLogger(__name__) + + +# ST-10 §4.4 / §4.5 constants. +_FR_RETRY_ATTEMPT_SPAN_NAME = "fortifyroot.bedrock.retry_attempt" +_FR_SPAN_ROLE_KEY = "fortifyroot.span.role" +_FR_SPAN_ROLE_RETRY_ATTEMPT = "retry_attempt" +_FR_HAS_RETRY_ATTEMPT_CHILD_KEY = "fortifyroot.span.has_retry_attempt_child" + +# Event-name patterns. ``*`` matches any operation under bedrock-runtime. +_BEFORE_SEND_PATTERN = "before-send.bedrock-runtime.*" +_RESPONSE_RECEIVED_PATTERN = "response-received.bedrock-runtime.*" + +# Stable correlation key used inside the per-request context dict. +_CTX_SPAN_KEY = "_fr_retry_attempt_span" +_CTX_TOKEN_KEY = "_fr_retry_attempt_framework_token" + +# Tracer provider the hooks should emit spans through. ``None`` → fall +# back to the global TracerProvider. Set via the optional +# ``tracer_provider`` arg to ``install_event_hooks_on_client`` so retry +# spans go to the same provider as the bedrock logical span (e.g. +# ``bedrock.completion``) emitted by the FR Bedrock instrumentor. +# Without this, a consumer passing an explicit provider to +# ``BedrockInstrumentor().instrument(tracer_provider=...)`` would get +# the parent span on their provider but the retry_attempt span sent +# into a no-op tracer. Module-level (not per-client) because all +# bedrock-runtime clients created under one instrumentor share the +# same provider config. (review-driven 2026-05-16 fix.) +_tracer_provider = None + + +def _is_suppressed() -> bool: + try: + if context_api.get_value(SUPPRESS_LANGUAGE_MODEL_INSTRUMENTATION_KEY): + return True + except Exception: + pass + try: + if is_framework_owned(): + return True + except Exception: + pass + return False + + +def _resolve_parent_span() -> Optional["trace.Span"]: + current = trace.get_current_span() + if current is None: + return None + ctx = current.get_span_context() + if ctx is None or not ctx.is_valid: + return None + return current + + +def _set_parent_marker(parent_span: "trace.Span") -> None: + try: + parent_span.set_attribute(_FR_HAS_RETRY_ATTEMPT_CHILD_KEY, True) + except Exception: + logger.debug("failed to set has_retry_attempt_child on parent", exc_info=True) + + +def _operation_from_event_name(event_name: str) -> str: + """Derive ``gen_ai.operation.name`` from a botocore event name like + ``before-send.bedrock-runtime.InvokeModel`` → ``chat`` / + ``text_completion`` mapping. + + Bedrock operations seen at the bedrock-runtime endpoint: + - InvokeModel / InvokeModelWithResponseStream → chat/completion + (depends on the model — we generalise to "chat") + - Converse / ConverseStream → chat + """ + if not isinstance(event_name, str): + return "chat" + name = event_name.split(".")[-1].lower() + if "invoke" in name or "converse" in name: + return "chat" + return "chat" + + +def _model_from_request(request: Any) -> Optional[str]: + """Best-effort: extract ``modelId`` from the request URL. + + Bedrock URL shape: ``/model/{model-id}/invoke`` or + ``/model/{model-id}/converse``. Some model IDs themselves contain + slashes after the vendor prefix, so we extract everything between + ``/model/`` and the trailing ``/invoke``/``/converse``. + """ + try: + url = getattr(request, "url", None) + if not isinstance(url, str): + return None + # Path-only form for endpoint requests: starts with /model/... + idx = url.find("/model/") + if idx < 0: + return None + rest = url[idx + len("/model/"):] + for suffix in ("/invoke-with-response-stream", "/invoke", "/converse-stream", "/converse"): + if rest.endswith(suffix): + return rest[: -len(suffix)] + # Fallback: trim any trailing query / slash. + slash = rest.find("/") + if slash > 0: + return rest[:slash] + return rest or None + except Exception: + return None + + +def _server_attrs_from_request(request: Any) -> dict[str, Any]: + out: dict[str, Any] = {} + try: + url = getattr(request, "url", None) + if not isinstance(url, str): + return out + # AWSPreparedRequest.url is a full URL string, not a parsed object. + from urllib.parse import urlparse + parsed = urlparse(url) + if parsed.hostname: + out["server.address"] = parsed.hostname + if parsed.port is not None: + out["server.port"] = parsed.port + elif parsed.scheme == "https": + out["server.port"] = 443 + elif parsed.scheme == "http": + out["server.port"] = 80 + except Exception: + pass + return out + + +def _start_attempt_span( + request: Any, + event_name: str, + parent_span: "trace.Span", +) -> "trace.Span": + attrs: dict[str, Any] = { + _FR_SPAN_ROLE_KEY: _FR_SPAN_ROLE_RETRY_ATTEMPT, + "gen_ai.system": "AWS", # matches existing Bedrock instrumentor + "gen_ai.operation.name": _operation_from_event_name(event_name), + } + model = _model_from_request(request) + if model: + attrs["gen_ai.request.model"] = model + attrs.update(_server_attrs_from_request(request)) + + tracer = trace.get_tracer(__name__, __version__, _tracer_provider) + parent_ctx = trace.set_span_in_context(parent_span) + span = tracer.start_span( + _FR_RETRY_ATTEMPT_SPAN_NAME, + kind=SpanKind.CLIENT, + attributes=attrs, + context=parent_ctx, + ) + return span + + +def _request_context(request: Any) -> Optional[dict]: + """Return the AWSPreparedRequest's per-request context dict, or None + if unavailable. botocore exposes this as ``request.context`` — + a mutable dict that propagates from request creation through to the + ``response-received`` event. + """ + ctx = getattr(request, "context", None) + if isinstance(ctx, dict): + return ctx + return None + + +def _finalize_success(span: "trace.Span", http_response: Any, parsed: Any) -> None: + try: + status_code = getattr(http_response, "status_code", None) + if isinstance(status_code, int): + span.set_attribute("http.status_code", status_code) + if isinstance(status_code, int) and 200 <= status_code < 300: + # Best-effort: response id from x-amzn-RequestId header. + try: + headers = getattr(http_response, "headers", None) or {} + if hasattr(headers, "get"): + rid = headers.get("x-amzn-RequestId") or headers.get("x-amzn-requestid") + if isinstance(rid, str) and rid: + span.set_attribute("gen_ai.response.id", rid) + except Exception: + pass + # Best-effort: token usage from parsed Converse response. + try: + if isinstance(parsed, dict): + usage = parsed.get("usage") or {} + if isinstance(usage, dict): + if isinstance(usage.get("inputTokens"), int): + span.set_attribute("gen_ai.usage.input_tokens", usage["inputTokens"]) + if isinstance(usage.get("outputTokens"), int): + span.set_attribute("gen_ai.usage.output_tokens", usage["outputTokens"]) + except Exception: + pass + span.set_status(Status(StatusCode.OK)) + else: + err_type = "botocore.HTTPStatusError" + if isinstance(status_code, int): + if status_code == 429: + err_type = "botocore.ThrottlingException" + elif 500 <= status_code < 600: + err_type = "botocore.InternalServerError" + elif status_code in (401, 403): + err_type = "botocore.AccessDeniedException" + span.set_attribute("error.type", err_type) + span.set_status(Status(StatusCode.ERROR, f"http {status_code}")) + except Exception: + logger.debug("failed to set success attrs on bedrock retry_attempt", exc_info=True) + + +def _finalize_error(span: "trace.Span", exception: BaseException) -> None: + try: + error_type = type(exception).__name__ + mod = type(exception).__module__ + if isinstance(mod, str) and mod and mod != "builtins": + error_type = f"{mod.split('.')[0]}.{error_type}" + span.set_attribute("error.type", error_type) + status_code = getattr(exception, "status_code", None) or getattr( + getattr(exception, "response", None), "status_code", None + ) + if isinstance(status_code, int): + span.set_attribute("http.status_code", status_code) + try: + span.record_exception(exception) + except Exception: + pass + span.set_status(Status(StatusCode.ERROR, str(exception))) + except Exception: + logger.debug("failed to set error attrs on bedrock retry_attempt", exc_info=True) + + +def _before_send_hook(event_name: Optional[str] = None, request: Any = None, **_kwargs) -> None: + """botocore ``before-send.bedrock-runtime.*`` event handler. + + Starts a retry_attempt span (when guards permit) and stows it on + ``request.context`` for retrieval by the paired ``response-received`` + hook. + + No self-registration in the §4.7.1 framework registry: the registry's + contract reserves registration for FRAMEWORK wrappers (LiteLLM / + LangChain / LlamaIndex). Direct-SDK wrappers only CONSULT via + ``is_framework_owned()``. Defensive cleanup of any prior context-stored + span on this request still runs, but it no longer touches the + framework registry. + + Streaming skip (ST-10.4 review-driven 2026-05-17): when the + botocore operation is a streaming one (event name ends with + ``Stream``: ``InvokeModelWithResponseStream`` / + ``ConverseStream``), this hook skips retry_attempt emission. + Rationale matches openai/anthropic: streaming retry_attempts + cannot carry token usage at attempt-end (usage arrives via the + stream-completion callback installed by the Bedrock streaming + wrapper, AFTER our hook has already finalised the span), but + §4.5 dedup would still promote them to canonical → zero-token + LLMUsageEvent. Leaving the parent ``bedrock.completion`` / + ``bedrock.converse`` span (which DOES get full usage from + ``stream_done``) as the canonical event. Streaming retry-loop + detection is the deferred follow-up + ``ST-10.4-FOLLOWUP-streaming-usage``. + """ + try: + if request is None: + return + if isinstance(event_name, str) and event_name.endswith("Stream"): + return + ctx = _request_context(request) + if ctx is None: + return + + # Defensive: if a prior attempt's span wasn't cleared (e.g. the + # paired ``response-received`` didn't fire — botocore normally + # guarantees pairing but we tolerate the weird sequence), end + # it as ERROR so the span doesn't leak. + prev_span = ctx.pop(_CTX_SPAN_KEY, None) + # Old payload may have a token field (pre-fix state) — drop it + # silently if present; we no longer register tokens. + ctx.pop(_CTX_TOKEN_KEY, None) + if prev_span is not None: + try: + prev_span.set_status( + Status(StatusCode.ERROR, "bedrock retry_attempt superseded (paired response-received missing)") + ) + prev_span.end() + except Exception: + pass + + if _is_suppressed(): + return + + parent = _resolve_parent_span() + if parent is None: + return + span = _start_attempt_span(request, event_name or "", parent) + _set_parent_marker(parent) + ctx[_CTX_SPAN_KEY] = span + except Exception: + logger.debug("ST-10.4: bedrock before-send hook failed", exc_info=True) + + +class _ResponseDictAdapter: + """Minimal adapter so ``_finalize_success`` can read ``status_code`` + + ``headers`` uniformly from either an httpx-style response object + (``.status_code``, ``.headers``) OR botocore's ``response_dict`` + (a plain dict with ``'status_code'`` + ``'headers'`` keys). + + botocore 1.42.x (and likely earlier point-releases) emits the + ``response-received`` event with ``response_dict=`` + ``parsed_response=`` + kwargs — NOT ``http_response=`` + ``parsed=``. The earlier hook + signature missed both rename pairs, leaving the retry_attempt span + with no status code, no OK/ERROR status, and ``IsError=False`` for + every attempt. Backend's RetryDetector then skipped the whole + sibling group on its ``failedCount == 0`` check. + """ + + __slots__ = ("status_code", "headers") + + def __init__(self, response_dict: dict) -> None: + sc = response_dict.get("status_code") + self.status_code = sc if isinstance(sc, int) else None + h = response_dict.get("headers") + self.headers = h if isinstance(h, dict) else {} + + +def _response_received_hook( + http_response: Any = None, + parsed: Any = None, + response_dict: Any = None, + parsed_response: Any = None, + context: Any = None, + exception: Any = None, + **_kwargs, +) -> None: + """botocore ``response-received.bedrock-runtime.*`` event handler. + + Finalises the retry_attempt span started by the paired + ``before-send`` hook. ``context`` is the per-request context dict + (same dict referenced by ``request.context`` at before-send time). + + Accepts BOTH botocore parameter-name conventions: + * ``http_response`` + ``parsed`` (older botocore signature the + fork-side unit tests drive directly). + * ``response_dict`` + ``parsed_response`` (botocore 1.42.x — the + names the real event emitter actually uses; verified via + ST-10.6 fr-system-tests probe). + + Per-attempt finalisation reads ``status_code`` and ``headers`` from + whichever shape is present; the rest of the body just falls through + to ``_finalize_success`` / ``_finalize_error``. + """ + try: + if not isinstance(context, dict): + return + span = context.pop(_CTX_SPAN_KEY, None) + # Old payload key (pre-fix state) — drop silently if present. + context.pop(_CTX_TOKEN_KEY, None) + if span is None: + return + + # Reconcile old/new botocore param names. If both are present + # (unlikely, but defensive), prefer the explicit older names so + # the existing unit tests' direct invocation path is not + # disturbed. + if http_response is None and isinstance(response_dict, dict): + http_response = _ResponseDictAdapter(response_dict) + if parsed is None and parsed_response is not None: + parsed = parsed_response + + try: + if exception is not None: + _finalize_error(span, exception) + else: + _finalize_success(span, http_response, parsed) + finally: + try: + span.end() + except Exception: + pass + except Exception: + logger.debug("ST-10.4: bedrock response-received hook failed", exc_info=True) + + +def install_event_hooks_on_client(client: Any, tracer_provider=None) -> None: + """Register retry-attempt event hooks on a freshly-created + bedrock-runtime client. Idempotent per-client — botocore's event + system de-dups identical (event-pattern, callable) registrations + via the unique-id mechanism. + + ``tracer_provider`` is stashed in module-level state and used by + the event hooks when creating retry_attempt spans. The same + provider is used for ALL clients in this process — all bedrock + clients created under one ``BedrockInstrumentor.instrument(...)`` + call share the same provider config. If ``None``, falls back to + the global TracerProvider. + + Called from ``BedrockInstrumentor._wrap`` after the bedrock-runtime + client is created and the FR per-method wraps are applied. + """ + global _tracer_provider + if tracer_provider is not None: + _tracer_provider = tracer_provider + try: + events = getattr(client, "meta", None) + events = getattr(events, "events", None) if events is not None else None + if events is None: + logger.debug( + "ST-10.4: bedrock client missing .meta.events; " + "skipping retry_attempt event-hook registration" + ) + return + events.register( + _BEFORE_SEND_PATTERN, + _before_send_hook, + unique_id="fortifyroot.bedrock.retry_attempt.before-send", + ) + events.register( + _RESPONSE_RECEIVED_PATTERN, + _response_received_hook, + unique_id="fortifyroot.bedrock.retry_attempt.response-received", + ) + except Exception: + logger.warning( + "ST-10.4: failed to register bedrock retry_attempt event hooks; " + "retry_attempt emission disabled for this client", + exc_info=True, + ) + + +def uninstall_event_hooks_on_client(client: Any) -> None: + """Remove the retry-attempt event hooks from a bedrock-runtime client. + Idempotent. Provided for tests and symmetry; production clients are + usually short-lived and don't need explicit unhook. + + Note: the module-level ``_tracer_provider`` is NOT cleared here — + it is shared across all bedrock clients created under one + instrumentor and remains relevant until the instrumentor itself is + uninstrumented. ``BedrockInstrumentor._uninstrument`` clears it + (see ``bedrock/__init__.py``). + """ + try: + events = getattr(client, "meta", None) + events = getattr(events, "events", None) if events is not None else None + if events is None: + return + try: + events.unregister( + _BEFORE_SEND_PATTERN, + unique_id="fortifyroot.bedrock.retry_attempt.before-send", + ) + except Exception: + pass + try: + events.unregister( + _RESPONSE_RECEIVED_PATTERN, + unique_id="fortifyroot.bedrock.retry_attempt.response-received", + ) + except Exception: + pass + except Exception: + logger.debug("ST-10.4: bedrock event-hook unregister failed", exc_info=True) + + +def _reset_tracer_provider_for_test() -> None: + """Test-only helper: clear the module-level tracer provider. Not part + of the public contract.""" + global _tracer_provider + _tracer_provider = None + + +__all__ = [ + "install_event_hooks_on_client", + "uninstall_event_hooks_on_client", + "_FR_RETRY_ATTEMPT_SPAN_NAME", + "_FR_HAS_RETRY_ATTEMPT_CHILD_KEY", + "_before_send_hook", + "_response_received_hook", +] diff --git a/packages/opentelemetry-instrumentation-bedrock/tests/test_retry_attempt_emission.py b/packages/opentelemetry-instrumentation-bedrock/tests/test_retry_attempt_emission.py new file mode 100644 index 0000000000..c2f30abf6f --- /dev/null +++ b/packages/opentelemetry-instrumentation-bedrock/tests/test_retry_attempt_emission.py @@ -0,0 +1,569 @@ +"""Tests for ST-10.4 Bedrock direct-SDK retry-attempt emission. + +Covers (per RETRY_LOOP.md §4.4 Bedrock row + §4.7 suppression): + - install_event_hooks_on_client registers both event-name patterns + on a bedrock-runtime client's botocore event system. + - Single-attempt happy path: before-send → response-received pair + emits ONE retry_attempt span; parent marker set; attrs include + model + http.status_code + gen_ai.usage.* from Converse response. + - Multi-attempt retry path: N siblings under one parent. + - Error attempt: response-received with ``exception`` sets ERROR + status + error.type. + - No-parent guard. + - §4.7 suppression — context-API and framework-registry both skip. + - Defensive: before-send firing twice without an intervening + response-received closes the orphaned prior span as ERROR. + +UNIT tests — drive ``_before_send_hook`` / ``_response_received_hook`` +directly with fake AWSPreparedRequest-shaped objects, so the test does +not depend on a real boto3 client / real AWS credentials / network. +""" + +from __future__ import annotations + +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest +from opentelemetry import context as context_api +from opentelemetry import trace +from opentelemetry.instrumentation.bedrock.retry_handler import ( + _CTX_SPAN_KEY, + _CTX_TOKEN_KEY, + _FR_HAS_RETRY_ATTEMPT_CHILD_KEY, + _FR_RETRY_ATTEMPT_SPAN_NAME, + _BEFORE_SEND_PATTERN, + _RESPONSE_RECEIVED_PATTERN, + _before_send_hook, + _response_received_hook, + install_event_hooks_on_client, + uninstall_event_hooks_on_client, +) +from opentelemetry.instrumentation.fortifyroot import ( + is_framework_owned, + register_framework_attempt, + retry_registry, + unregister_framework_attempt, +) +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import ( + InMemorySpanExporter, +) +from opentelemetry.semconv_ai import SUPPRESS_LANGUAGE_MODEL_INSTRUMENTATION_KEY + + +@pytest.fixture +def fresh_tracer(): + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + try: + trace.set_tracer_provider(provider) + except Exception: + pass + current = trace.get_tracer_provider() + if current is not provider: + try: + current.add_span_processor(SimpleSpanProcessor(exporter)) + except Exception: + pass + yield trace.get_tracer("test"), exporter, current + + +@pytest.fixture(autouse=True) +def reset_registry(): + retry_registry._reset_for_test() + yield + retry_registry._reset_for_test() + + +def _make_request(model: str = "anthropic.claude-haiku-4-5", + region: str = "us-east-1", + operation: str = "invoke") -> SimpleNamespace: + """Build a fake AWSPreparedRequest-shaped object. + + The retry hook reads: + - request.url (string with /model/{modelId}/{op}) + - request.context (mutable dict — propagated to response-received) + """ + url = f"https://bedrock-runtime.{region}.amazonaws.com/model/{model}/{operation}" + return SimpleNamespace(url=url, context={}) + + +def _make_http_response(status_code: int = 200, request_id: str = "amzn-req-1") -> SimpleNamespace: + headers = {"x-amzn-RequestId": request_id} + return SimpleNamespace(status_code=status_code, headers=headers) + + +def _retry_spans(exporter: InMemorySpanExporter): + return [s for s in exporter.get_finished_spans() + if s.name == _FR_RETRY_ATTEMPT_SPAN_NAME] + + +# --------------------------------------------------------------------------- +# Event-hook registration. +# --------------------------------------------------------------------------- + +def test_install_event_hooks_registers_both_patterns(): + """install_event_hooks_on_client(client) MUST register the + before-send AND response-received patterns on the client's + botocore event system.""" + client = MagicMock() + install_event_hooks_on_client(client) + + register_calls = client.meta.events.register.call_args_list + patterns_registered = [call.args[0] for call in register_calls] + assert _BEFORE_SEND_PATTERN in patterns_registered, ( + f"before-send pattern must be registered; got {patterns_registered}" + ) + assert _RESPONSE_RECEIVED_PATTERN in patterns_registered, ( + f"response-received pattern must be registered; got {patterns_registered}" + ) + + # Each registration uses a stable unique_id so botocore dedups. + for call in register_calls: + assert "unique_id" in call.kwargs, "unique_id must be supplied for dedup" + assert call.kwargs["unique_id"].startswith("fortifyroot.bedrock.retry_attempt.") + + +def test_uninstall_event_hooks_unregisters_both_patterns(): + client = MagicMock() + install_event_hooks_on_client(client) + uninstall_event_hooks_on_client(client) + + unregister_calls = client.meta.events.unregister.call_args_list + patterns_unregistered = [call.args[0] for call in unregister_calls] + assert _BEFORE_SEND_PATTERN in patterns_unregistered + assert _RESPONSE_RECEIVED_PATTERN in patterns_unregistered + + +def test_install_on_client_without_meta_events_is_a_noop_no_crash(): + """If the client doesn't have .meta.events, install logs debug + and returns gracefully — no exception.""" + client = SimpleNamespace() + install_event_hooks_on_client(client) # must not raise + uninstall_event_hooks_on_client(client) # must not raise + + +# --------------------------------------------------------------------------- +# Single-attempt happy path. +# --------------------------------------------------------------------------- + +def test_single_attempt_emits_one_span_with_marker(fresh_tracer): + tracer, exporter, _ = fresh_tracer + request = _make_request(model="anthropic.claude-haiku-4-5", operation="invoke") + + parent = tracer.start_span("bedrock.completion") + with trace.use_span(parent, end_on_exit=False): + # before-send fires first. + _before_send_hook( + event_name="before-send.bedrock-runtime.InvokeModel", + request=request, + ) + # Span stowed on request.context. + assert _CTX_SPAN_KEY in request.context + + # response-received fires with the same context dict. + _response_received_hook( + http_response=_make_http_response(200, "amzn-req-1"), + parsed={"usage": {"inputTokens": 12, "outputTokens": 7}}, + context=request.context, + exception=None, + ) + # Context cleaned up. + assert _CTX_SPAN_KEY not in request.context + assert _CTX_TOKEN_KEY not in request.context + + parent.end() + + spans = exporter.get_finished_spans() + parent_exported = next(s for s in spans if s.name == "bedrock.completion") + retry_spans = _retry_spans(exporter) + assert len(retry_spans) == 1 + rs = retry_spans[0] + assert rs.parent.span_id == parent_exported.context.span_id + assert rs.attributes.get("fortifyroot.span.role") == "retry_attempt" + assert rs.attributes.get("gen_ai.system") == "AWS" + assert rs.attributes.get("gen_ai.request.model") == "anthropic.claude-haiku-4-5" + assert rs.attributes.get("gen_ai.operation.name") == "chat" + assert rs.attributes.get("http.status_code") == 200 + assert rs.attributes.get("gen_ai.response.id") == "amzn-req-1" + assert rs.attributes.get("gen_ai.usage.input_tokens") == 12 + assert rs.attributes.get("gen_ai.usage.output_tokens") == 7 + assert rs.attributes.get("server.address") == "bedrock-runtime.us-east-1.amazonaws.com" + assert rs.attributes.get("server.port") == 443 + assert parent_exported.attributes.get(_FR_HAS_RETRY_ATTEMPT_CHILD_KEY) is True + + +# --------------------------------------------------------------------------- +# Multi-attempt retry path. +# --------------------------------------------------------------------------- + +def test_three_attempts_share_parent(fresh_tracer): + """429 → 429 → 200 chain produces 3 sibling retry_attempt spans + under one parent. Each attempt is a fresh AWSPreparedRequest with + its own context dict (matching botocore's per-attempt model).""" + tracer, exporter, _ = fresh_tracer + + parent = tracer.start_span("bedrock.converse") + with trace.use_span(parent, end_on_exit=False): + for status_code, _ in [(429, "a"), (429, "b"), (200, "c")]: + req = _make_request(operation="converse") + _before_send_hook( + event_name="before-send.bedrock-runtime.Converse", + request=req, + ) + _response_received_hook( + http_response=_make_http_response(status_code), + parsed={}, + context=req.context, + exception=None, + ) + parent.end() + + spans = exporter.get_finished_spans() + parent_exported = next(s for s in spans if s.name == "bedrock.converse") + retry_spans = _retry_spans(exporter) + assert len(retry_spans) == 3 + + parent_ids = {s.parent.span_id for s in retry_spans} + assert parent_ids == {parent_exported.context.span_id} + + from opentelemetry.trace import StatusCode + err = sum(1 for s in retry_spans if s.status.status_code == StatusCode.ERROR) + ok = sum(1 for s in retry_spans if s.status.status_code == StatusCode.OK) + assert err == 2 and ok == 1 + + for s in retry_spans: + if s.status.status_code == StatusCode.ERROR: + assert s.attributes.get("http.status_code") == 429 + assert s.attributes.get("error.type") == "botocore.ThrottlingException" + + +# --------------------------------------------------------------------------- +# Error / exception path. +# --------------------------------------------------------------------------- + +def test_exception_path_records_error(fresh_tracer): + """When response-received fires with ``exception`` set, the span + finalises as ERROR with error.type.""" + tracer, exporter, _ = fresh_tracer + + class ConnectError(Exception): + pass + + err = ConnectError("connect timeout") + + request = _make_request() + parent = tracer.start_span("bedrock.completion") + with trace.use_span(parent, end_on_exit=False): + _before_send_hook( + event_name="before-send.bedrock-runtime.InvokeModel", + request=request, + ) + _response_received_hook( + http_response=None, + parsed=None, + context=request.context, + exception=err, + ) + parent.end() + + from opentelemetry.trace import StatusCode + retry_spans = _retry_spans(exporter) + assert len(retry_spans) == 1 + rs = retry_spans[0] + assert rs.status.status_code == StatusCode.ERROR + assert "ConnectError" in (rs.attributes.get("error.type") or "") + + +# --------------------------------------------------------------------------- +# Marker timing (§4.5). +# --------------------------------------------------------------------------- + +def test_marker_NOT_set_when_no_attempts_fire(fresh_tracer): + tracer, exporter, _ = fresh_tracer + parent = tracer.start_span("bedrock.completion") + parent.end() + parent_exported = next( + s for s in exporter.get_finished_spans() if s.name == "bedrock.completion" + ) + assert _FR_HAS_RETRY_ATTEMPT_CHILD_KEY not in (parent_exported.attributes or {}) + + +# --------------------------------------------------------------------------- +# No-parent guard. +# --------------------------------------------------------------------------- + +def test_no_parent_does_not_emit(fresh_tracer): + tracer, exporter, _ = fresh_tracer + request = _make_request() + _before_send_hook( + event_name="before-send.bedrock-runtime.InvokeModel", + request=request, + ) + # No span stowed on context because no parent. + assert _CTX_SPAN_KEY not in request.context + _response_received_hook( + http_response=_make_http_response(200), + parsed={}, + context=request.context, + exception=None, + ) + assert len(_retry_spans(exporter)) == 0 + + +# --------------------------------------------------------------------------- +# §4.7 suppression. +# --------------------------------------------------------------------------- + +def test_context_suppression_skips_emission(fresh_tracer): + tracer, exporter, _ = fresh_tracer + request = _make_request() + parent = tracer.start_span("bedrock.completion") + token = context_api.attach( + context_api.set_value(SUPPRESS_LANGUAGE_MODEL_INSTRUMENTATION_KEY, True) + ) + try: + with trace.use_span(parent, end_on_exit=False): + _before_send_hook( + event_name="before-send.bedrock-runtime.InvokeModel", + request=request, + ) + _response_received_hook( + http_response=_make_http_response(200), + parsed={}, + context=request.context, + exception=None, + ) + finally: + context_api.detach(token) + parent.end() + assert len(_retry_spans(exporter)) == 0 + + +def test_framework_registry_suppression_skips_emission(fresh_tracer): + tracer, exporter, _ = fresh_tracer + request = _make_request() + parent = tracer.start_span("bedrock.completion") + tok = register_framework_attempt() + assert is_framework_owned() + try: + with trace.use_span(parent, end_on_exit=False): + _before_send_hook( + event_name="before-send.bedrock-runtime.InvokeModel", + request=request, + ) + _response_received_hook( + http_response=_make_http_response(200), + parsed={}, + context=request.context, + exception=None, + ) + finally: + unregister_framework_attempt(tok) + parent.end() + assert len(_retry_spans(exporter)) == 0 + + +def test_direct_sdk_wrapper_does_not_self_register_in_framework_registry(fresh_tracer): + """REGRESSION GUARD (review-driven fix 2026-05-13): the bedrock + retry hooks MUST NOT register tokens in the §4.7.1 framework + registry. Direct-SDK wrappers only CONSULT via + ``is_framework_owned()``.""" + tracer, _, _ = fresh_tracer + request = _make_request() + parent = tracer.start_span("bedrock.completion") + with trace.use_span(parent, end_on_exit=False): + assert not is_framework_owned() + _before_send_hook( + event_name="before-send.bedrock-runtime.InvokeModel", + request=request, + ) + assert not is_framework_owned(), ( + "before-send hook must NOT self-register a framework token" + ) + _response_received_hook( + http_response=_make_http_response(200), + parsed={}, + context=request.context, + exception=None, + ) + assert not is_framework_owned() + parent.end() + + +# --------------------------------------------------------------------------- +# Defensive: before-send fires twice (botocore pairs guaranteed but +# tolerate weird sequences). +# --------------------------------------------------------------------------- + +def test_double_before_send_closes_orphan_as_error(fresh_tracer): + tracer, exporter, _ = fresh_tracer + request = _make_request() + parent = tracer.start_span("bedrock.completion") + with trace.use_span(parent, end_on_exit=False): + _before_send_hook( + event_name="before-send.bedrock-runtime.InvokeModel", + request=request, + ) + first_span = request.context[_CTX_SPAN_KEY] + # A second before-send (without the paired response-received) + # supersedes the first — the orphan is closed as ERROR so it + # doesn't leak. + _before_send_hook( + event_name="before-send.bedrock-runtime.InvokeModel", + request=request, + ) + _response_received_hook( + http_response=_make_http_response(200), + parsed={}, + context=request.context, + exception=None, + ) + parent.end() + + retry_spans = _retry_spans(exporter) + assert len(retry_spans) == 2, f"expected 2 spans (orphan + completed); got {len(retry_spans)}" + # The orphan span (first one started) MUST be ERROR. + from opentelemetry.trace import StatusCode + errors = [s for s in retry_spans if s.status.status_code == StatusCode.ERROR] + oks = [s for s in retry_spans if s.status.status_code == StatusCode.OK] + assert len(errors) == 1 and len(oks) == 1 + # The orphan and first_span share the same span context. + assert first_span.get_span_context().span_id == errors[0].context.span_id + + +# --------------------------------------------------------------------------- +# Response-received with no prior before-send is a no-op. +# --------------------------------------------------------------------------- + +def test_response_received_with_no_prior_before_send_is_noop(fresh_tracer): + """Stray response-received (no paired before-send started a span) + must NOT crash. Defensive contract per the hook docstring.""" + _, exporter, _ = fresh_tracer + _response_received_hook( + http_response=_make_http_response(200), + parsed={}, + context={}, # empty context + exception=None, + ) + assert len(_retry_spans(exporter)) == 0 + + +# --------------------------------------------------------------------------- +# Defensive: request with no .context dict. +# --------------------------------------------------------------------------- + +def test_before_send_with_request_missing_context_is_noop(fresh_tracer): + """If a botocore request somehow arrives without a usable + ``.context`` dict (defensive — every AWSPreparedRequest in current + botocore has one), the hook must NOT crash and must NOT emit a + span (no place to stash state for the paired response-received).""" + tracer, exporter, _ = fresh_tracer + # SimpleNamespace with no `context` attribute. + request_no_context = SimpleNamespace( + url="https://bedrock-runtime.us-east-1.amazonaws.com/model/x/invoke" + ) + parent = tracer.start_span("bedrock.completion") + with trace.use_span(parent, end_on_exit=False): + # Must not raise. + _before_send_hook( + event_name="before-send.bedrock-runtime.InvokeModel", + request=request_no_context, + ) + parent.end() + assert len(_retry_spans(exporter)) == 0 + + +# --------------------------------------------------------------------------- +# M1: streaming-span ambient-context fix (review-driven 2026-05-13). +# --------------------------------------------------------------------------- + +def test_streaming_event_skips_emission_entirely(fresh_tracer): + """ST-10.4 (review-driven 2026-05-17): when the botocore operation + is a streaming one (event name ends with ``Stream``: + ``InvokeModelWithResponseStream`` / ``ConverseStream``), the + ``_before_send_hook`` SKIPS retry_attempt emission entirely. + + Rationale matches openai/anthropic streaming-skip: streaming + retry_attempts cannot carry usage at attempt-end (usage arrives + via the stream-completion callback in the Bedrock streaming + wrapper, AFTER our hook has finalised) but §4.5 dedup would + still promote them to canonical → zero-token LLMUsageEvents + breaking fr-system-tests. The parent + ``bedrock.completion`` / ``bedrock.converse`` span (which DOES + get full usage via ``stream_done``) stays canonical. Streaming + retry-loop detection is the deferred follow-up + ``ST-10.4-FOLLOWUP-streaming-usage``. + + This test supersedes the earlier + ``test_streaming_span_via_use_span_parents_retry_attempt`` from + the M1 fix round — the M1 ``trace.use_span`` wrapper in the + Bedrock instrumentor is still needed for non-streaming + operations and for future work that wants to plug streaming + retry-loop detection back in. + """ + tracer, exporter, _ = fresh_tracer + streaming_span = tracer.start_span("bedrock.completion") + request = _make_request(operation="invoke-with-response-stream") + + with trace.use_span(streaming_span, end_on_exit=False): + _before_send_hook( + event_name="before-send.bedrock-runtime.InvokeModelWithResponseStream", + request=request, + ) + # ``request.context`` should be UNTOUCHED — no span stashed. + assert _CTX_SPAN_KEY not in request.context + _response_received_hook( + http_response=_make_http_response(200), + parsed={}, + context=request.context, + exception=None, + ) + streaming_span.end() + + retry_spans = _retry_spans(exporter) + assert len(retry_spans) == 0, ( + f"streaming events MUST skip retry_attempt emission entirely; " + f"got {len(retry_spans)} span(s)." + ) + + # ConverseStream variant — same behaviour. + request2 = _make_request(operation="converse-stream") + with trace.use_span(streaming_span, end_on_exit=False): + _before_send_hook( + event_name="before-send.bedrock-runtime.ConverseStream", + request=request2, + ) + assert _CTX_SPAN_KEY not in request2.context + + +def test_streaming_span_without_use_span_skips_emission(fresh_tracer): + """Counterpoint to the M1 fix: WITHOUT ``trace.use_span``, the + streaming Bedrock span is NOT ambient → ``_resolve_parent_span`` + returns None (no valid ambient) → retry_attempt skipped. Documents + the pre-fix failure mode so a future regression on the instrumentor + side gets caught.""" + tracer, exporter, _ = fresh_tracer + streaming_span = tracer.start_span("bedrock.completion") # NOT made current + request = _make_request(operation="invoke-with-response-stream") + + # No use_span wrap — mimicking the pre-M1 instrumentor. + _before_send_hook( + event_name="before-send.bedrock-runtime.InvokeModelWithResponseStream", + request=request, + ) + _response_received_hook( + http_response=_make_http_response(200), + parsed={}, + context=request.context, + exception=None, + ) + streaming_span.end() + + assert len(_retry_spans(exporter)) == 0, ( + "without use_span(streaming_span), retry_attempt must be skipped — " + "no valid ambient parent" + ) diff --git a/packages/opentelemetry-instrumentation-fortifyroot/opentelemetry/instrumentation/fortifyroot/__init__.py b/packages/opentelemetry-instrumentation-fortifyroot/opentelemetry/instrumentation/fortifyroot/__init__.py index 09f79ea930..d43e368805 100644 --- a/packages/opentelemetry-instrumentation-fortifyroot/opentelemetry/instrumentation/fortifyroot/__init__.py +++ b/packages/opentelemetry-instrumentation-fortifyroot/opentelemetry/instrumentation/fortifyroot/__init__.py @@ -23,6 +23,12 @@ run_prompt_safety_async, set_object_value, ) +from opentelemetry.instrumentation.fortifyroot.retry_registry import ( + clear_for_thread as clear_framework_attempts_for_thread, + is_framework_owned, + register_framework_attempt, + unregister_framework_attempt, +) from opentelemetry.instrumentation.fortifyroot.text_streaming import ( CompletionTextStreamGroup, ) @@ -40,6 +46,10 @@ __all__ = [ "SAFETY_EVENT_NAME", "HANDLER_LOCK", + "clear_framework_attempts_for_thread", + "is_framework_owned", + "register_framework_attempt", + "unregister_framework_attempt", "SafetyContext", "SafetyDecision", "SafetyFinding", diff --git a/packages/opentelemetry-instrumentation-fortifyroot/opentelemetry/instrumentation/fortifyroot/retry_registry.py b/packages/opentelemetry-instrumentation-fortifyroot/opentelemetry/instrumentation/fortifyroot/retry_registry.py new file mode 100644 index 0000000000..a548991ac4 --- /dev/null +++ b/packages/opentelemetry-instrumentation-fortifyroot/opentelemetry/instrumentation/fortifyroot/retry_registry.py @@ -0,0 +1,284 @@ +"""§4.7.1 token-based framework-attempt registry. + +Shared across all FR fork instrumentations. The flow is: + + - Framework wrappers (LiteLLM / LangChain / LlamaIndex) call + ``register_framework_attempt()`` at the start of EACH HTTP attempt + they observe (i.e. once per ``log_pre_api_call`` / ``on_chat_model_start`` + / ``LLMChatStartEvent``). They get back an opaque token. + + - The same wrappers call ``unregister_framework_attempt(token)`` when + that attempt finishes (success or failure callback). + + - Direct-SDK wrappers (OpenAI / Anthropic / Bedrock) consult + ``is_framework_owned()`` at every wrapped ``send()`` / event-hook + invocation. If True, they SUPPRESS their own ``retry_attempt`` span + emission so the framework's emitter remains the single source of + truth for that logical call. + +The "owned" state is tracked PER OS THREAD because LiteLLM's +``log_pre_api_call`` and the underlying ``httpx.send()`` it triggers +run on the same thread synchronously. Async paths use the same +underlying httpx client, so the thread-local model still applies for +the duration of one attempt. (Re-entrancy across threads is handled by +the per-TID dict.) + +Design references: + - RETRY_LOOP.md §4.7 — suppression discipline rationale. + - RETRY_LOOP.md §4.7.1 — registry shape, eviction policy, required + tests (re-entrancy, stale-cleanup, + parent-end cleanup, cap-eviction, + thread-ID reuse). + - phase_st10_retryloop.txt round-3 disposition — replaces an + earlier ``set[int]`` design that didn't handle re-entrancy. +""" + +from __future__ import annotations + +import logging +import threading +import time +import uuid +from typing import Optional + +logger = logging.getLogger(__name__) + + +# Module-level state (per process). Per RETRY_LOOP.md §4.7.1: +# shape: dict[tid, dict[token, started_at_monotonic_seconds]] +# Keying by tid (thread ID) is what makes per-thread ownership work; +# keying by token within a TID is what makes re-entrancy work (one +# thread can own multiple nested attempts and only suppression-stops +# when ALL tokens have been unregistered). +_REGISTRY: "dict[int, dict[str, float]]" = {} + +# Reverse index: token → originating TID. Required so that +# unregister_framework_attempt() works correctly when the framework's +# terminal callback runs on a different OS thread from its start +# callback (e.g. async paths where success/failure callbacks dispatch +# to a worker thread). Without this mapping, unregister would only +# clean up the CURRENT thread's bucket — leaking the original token +# in the original TID's bucket and falsely keeping that TID +# "framework-owned" until TTL eviction. +# (Review-batch-1 Blocker fix 2026-05-10.) +_TOKEN_TO_TID: "dict[str, int]" = {} + +_REGISTRY_LOCK = threading.Lock() + +# Bounded-size + TTL-based eviction caps. Tunable; chosen to match +# the existing _FR_COMPLETION_SAFETY_MARKERS_MAX shape so all FR +# bookkeeping has consistent memory ceilings. +_REGISTRY_MAX = 4096 # max ENTRIES (sum of inner-dict sizes) before cap-eviction kicks in +_REGISTRY_EVICT_BATCH = 1024 # number of entries to drop in one cap-eviction pass +_REGISTRY_STALE_TTL_SEC = 60.0 # TTL after which a token is considered stale and evicted + +# Rate-limit warning emission to at most one per N evictions per bucket +# (per-TID). Stale-eviction can spike noisily if a framework crashes +# with many leaked tokens; logging every one would flood operator +# dashboards. +_WARN_EVERY_N = 32 +_evict_warn_counter: dict[int, int] = {} +_evict_warn_lock = threading.Lock() + + +def _now() -> float: + return time.monotonic() + + +def _evict_stale_for_tid_locked(tid: int, now: float) -> int: + """Remove tokens older than the TTL for the given TID. Returns + number of entries evicted. Caller must hold _REGISTRY_LOCK. + + This runs in-band on EVERY ``is_framework_owned()`` call so a + leaked token (framework crashed without unregistering) cannot + suppress direct-SDK emission indefinitely. The cost is one O(K) + pass per call where K = tokens active on this TID — almost always + 0 or 1 in practice. + """ + bucket = _REGISTRY.get(tid) + if not bucket: + return 0 + cutoff = now - _REGISTRY_STALE_TTL_SEC + stale_tokens = [tok for tok, started_at in bucket.items() if started_at < cutoff] + for tok in stale_tokens: + bucket.pop(tok, None) + _TOKEN_TO_TID.pop(tok, None) + if not bucket: + _REGISTRY.pop(tid, None) + if stale_tokens: + _maybe_warn_eviction(tid, len(stale_tokens)) + return len(stale_tokens) + + +def _enforce_max_locked(now: float) -> int: + """If the total entry count exceeds _REGISTRY_MAX, evict + _REGISTRY_EVICT_BATCH oldest entries (across all TIDs). Returns + number evicted. Caller must hold _REGISTRY_LOCK. + + This is a hard cap defending against pathological leak storms. + Normal operation never triggers it because TTL-eviction keeps the + registry small. + """ + total = sum(len(b) for b in _REGISTRY.values()) + if total <= _REGISTRY_MAX: + return 0 + # Flatten + sort by started_at; drop the oldest _REGISTRY_EVICT_BATCH. + flat: list[tuple[float, int, str]] = [] + for tid, bucket in _REGISTRY.items(): + for tok, started_at in bucket.items(): + flat.append((started_at, tid, tok)) + flat.sort() # ascending by started_at — oldest first + to_evict = flat[:_REGISTRY_EVICT_BATCH] + for _, tid, tok in to_evict: + bucket = _REGISTRY.get(tid) + if bucket is None: + continue + bucket.pop(tok, None) + _TOKEN_TO_TID.pop(tok, None) + if not bucket: + _REGISTRY.pop(tid, None) + logger.warning( + "fortifyroot retry_registry: cap-evicted %d entries (total exceeded %d); " + "this indicates leaked tokens from a misbehaving framework wrapper", + len(to_evict), + _REGISTRY_MAX, + ) + return len(to_evict) + + +def _maybe_warn_eviction(tid: int, evicted: int) -> None: + with _evict_warn_lock: + c = _evict_warn_counter.get(tid, 0) + evicted + if c < _WARN_EVERY_N: + _evict_warn_counter[tid] = c + return + _evict_warn_counter[tid] = 0 + # Outside the lock so logger.warning's I/O can't deadlock with + # registry operations elsewhere. + logger.warning( + "fortifyroot retry_registry: stale-evicted %d+ tokens on tid=%d; " + "framework wrapper may be leaking attempts (TTL=%.0fs)", + evicted, + tid, + _REGISTRY_STALE_TTL_SEC, + ) + + +def register_framework_attempt() -> str: + """Register a new framework-attempt token on the current OS thread. + + Returns the opaque token string. The caller MUST pass this token + to ``unregister_framework_attempt`` when the attempt completes. + + Idempotency: calling this multiple times on the same thread is + safe and supported (re-entrancy / nested attempts) — each call + mints a fresh token. + """ + token = uuid.uuid4().hex + tid = threading.get_ident() + now = _now() + with _REGISTRY_LOCK: + bucket = _REGISTRY.get(tid) + if bucket is None: + bucket = {} + _REGISTRY[tid] = bucket + bucket[token] = now + _TOKEN_TO_TID[token] = tid + _enforce_max_locked(now) + return token + + +def unregister_framework_attempt(token: Optional[str]) -> None: + """Remove the given token from its ORIGINATING TID bucket. No-op + if token is None or unknown (defensive — tolerates double-unregister + when both sync and async callbacks fire). + + The originating TID is recovered via the ``_TOKEN_TO_TID`` reverse + index — we deliberately do NOT use ``threading.get_ident()`` here, + because the framework's terminal callback may run on a different + OS thread from the start callback (e.g. asyncio worker dispatch). + Using the current thread's TID would leak the entry in the + originating TID's bucket and falsely keep that thread + "framework-owned" until TTL eviction. + """ + if not token: + return + with _REGISTRY_LOCK: + original_tid = _TOKEN_TO_TID.pop(token, None) + if original_tid is None: + # Already unregistered (or unknown). Defensive no-op. + return + bucket = _REGISTRY.get(original_tid) + if bucket is None: + return + bucket.pop(token, None) + if not bucket: + _REGISTRY.pop(original_tid, None) + + +def is_framework_owned(tid: Optional[int] = None) -> bool: + """Return True iff the given TID (default: current thread) has + at least one live framework-attempt token registered. + + This is the LOAD-BEARING read path consulted by direct-SDK + wrappers (OpenAI / Anthropic / Bedrock) at every wrapped HTTP + send. If it returns True, the direct-SDK wrapper SUPPRESSES its + own retry_attempt span emission — the framework is the source of + truth for that logical call. + + Performs per-TID stale eviction in-band before answering, so a + leaked token cannot indefinitely suppress emission (see + review-round-4 Q2 fix in phase_st10_retryloop.txt). + """ + if tid is None: + tid = threading.get_ident() + now = _now() + with _REGISTRY_LOCK: + _evict_stale_for_tid_locked(tid, now) + bucket = _REGISTRY.get(tid) + return bool(bucket) + + +def clear_for_thread(tid: Optional[int] = None) -> int: + """Drop all tokens for the given TID (default: current thread). + Returns number of tokens dropped. Used as a parent-span-end + orphan-cleanup primitive: when a framework wrapper's parent span + ends, any tokens still registered for that thread are by + definition orphans (the parent finished without success/failure + callbacks firing for those attempts). + """ + if tid is None: + tid = threading.get_ident() + with _REGISTRY_LOCK: + bucket = _REGISTRY.pop(tid, None) + if bucket: + for tok in bucket: + _TOKEN_TO_TID.pop(tok, None) + return len(bucket) if bucket else 0 + + +def _registry_size_for_test() -> "tuple[int, int]": + """Test-only helper: returns (num_tids, total_tokens). Not part + of the public contract; do not import outside tests. + """ + with _REGISTRY_LOCK: + return (len(_REGISTRY), sum(len(b) for b in _REGISTRY.values())) + + +def _reset_for_test() -> None: + """Test-only helper: wipe registry state. Not part of the public + contract; do not import outside tests. + """ + with _REGISTRY_LOCK: + _REGISTRY.clear() + _TOKEN_TO_TID.clear() + with _evict_warn_lock: + _evict_warn_counter.clear() + + +__all__ = [ + "register_framework_attempt", + "unregister_framework_attempt", + "is_framework_owned", + "clear_for_thread", +] diff --git a/packages/opentelemetry-instrumentation-fortifyroot/tests/test_retry_registry.py b/packages/opentelemetry-instrumentation-fortifyroot/tests/test_retry_registry.py new file mode 100644 index 0000000000..1f146ef045 --- /dev/null +++ b/packages/opentelemetry-instrumentation-fortifyroot/tests/test_retry_registry.py @@ -0,0 +1,304 @@ +"""Tests for the §4.7.1 framework-attempt registry. + +Required tests per RETRY_LOOP.md §4.7.1: + - re-entrancy (one thread can hold multiple tokens, suppression + only stops when ALL are unregistered) + - stale-entry eviction (TTL-based, runs in-band on read path) + - cap-eviction (when total entries exceed _REGISTRY_MAX) + - parent-end cleanup (clear_for_thread drops orphan tokens) + - thread-ID reuse (registry survives a TID being recycled by + the OS for a subsequent thread, by virtue of token uniqueness) +""" + +from __future__ import annotations + +import threading +import time +from concurrent.futures import ThreadPoolExecutor + +import pytest + +from opentelemetry.instrumentation.fortifyroot import ( + clear_framework_attempts_for_thread, + is_framework_owned, + register_framework_attempt, + unregister_framework_attempt, +) +from opentelemetry.instrumentation.fortifyroot import retry_registry + + +@pytest.fixture(autouse=True) +def reset_registry(): + """Each test starts with an empty registry.""" + retry_registry._reset_for_test() + yield + retry_registry._reset_for_test() + + +def test_register_unregister_basic(): + """Registering then unregistering a single token leaves the + registry empty and is_framework_owned() False.""" + assert not is_framework_owned() + token = register_framework_attempt() + assert isinstance(token, str) and len(token) > 0 + assert is_framework_owned() + unregister_framework_attempt(token) + assert not is_framework_owned() + + +def test_register_returns_unique_tokens(): + """Each register call mints a fresh token (uniqueness invariant).""" + tokens = {register_framework_attempt() for _ in range(50)} + assert len(tokens) == 50, "all tokens must be unique" + + +def test_reentrancy_one_thread_multiple_tokens(): + """Re-entrancy: one thread can register multiple tokens (e.g. + nested framework attempts). Suppression remains True until ALL + tokens are unregistered. This is the canonical case the + refcount-via-token design solves vs the rejected set[int] + design (round-3 Blocker disposition).""" + t1 = register_framework_attempt() + t2 = register_framework_attempt() + t3 = register_framework_attempt() + + assert is_framework_owned() + + unregister_framework_attempt(t1) + assert is_framework_owned(), "still owned: t2 + t3 outstanding" + + unregister_framework_attempt(t2) + assert is_framework_owned(), "still owned: t3 outstanding" + + unregister_framework_attempt(t3) + assert not is_framework_owned(), "fully released after all tokens removed" + + +def test_unregister_idempotent_on_unknown_token(): + """Unregistering an unknown / already-removed token is a no-op + (defensive contract — handles double-unregister from sync+async + callbacks both firing).""" + t = register_framework_attempt() + unregister_framework_attempt(t) + # Should not raise, should not affect state. + unregister_framework_attempt(t) + unregister_framework_attempt("nonexistent-token") + unregister_framework_attempt(None) + assert not is_framework_owned() + + +def test_per_thread_isolation(): + """Tokens registered on thread A do NOT make thread B "owned".""" + main_tid = threading.get_ident() + register_framework_attempt() + assert is_framework_owned(main_tid) + + foreign_tid_seen: dict[str, bool] = {} + + def worker(): + # Different thread → fresh TID → no tokens registered. + foreign_tid_seen["owned"] = is_framework_owned() + + t = threading.Thread(target=worker) + t.start() + t.join() + + assert foreign_tid_seen["owned"] is False, ( + "thread B sees its own (empty) bucket, not thread A's" + ) + + +def test_stale_entry_eviction_via_ttl(monkeypatch): + """Tokens older than _REGISTRY_STALE_TTL_SEC are evicted on the + next is_framework_owned() call. This is the load-bearing + in-band eviction (review-round-4 Q2 fix) — without it a leaked + token (framework crashed) would suppress emission forever.""" + # Shrink the TTL so the test runs in real time. + monkeypatch.setattr(retry_registry, "_REGISTRY_STALE_TTL_SEC", 0.1) + + register_framework_attempt() + assert is_framework_owned() + + time.sleep(0.15) # past the TTL + + # is_framework_owned() runs eviction in-band. + assert not is_framework_owned(), "stale token must be evicted on read path" + + +def test_stale_eviction_with_no_intervening_registration(monkeypatch): + """The boundary case from §4.7.1: register, sleep past TTL, then + check is_framework_owned() WITHOUT another registration in + between. This is the case that the in-band eviction design + explicitly fixes — eviction happens on every read, not only on + write paths.""" + monkeypatch.setattr(retry_registry, "_REGISTRY_STALE_TTL_SEC", 0.1) + + register_framework_attempt() + time.sleep(0.15) + + # No further register() calls. Eviction MUST happen in + # is_framework_owned itself. + assert not is_framework_owned() + # Registry should be fully empty after eviction. + n_tids, n_tokens = retry_registry._registry_size_for_test() + assert n_tids == 0 + assert n_tokens == 0 + + +def test_cap_eviction_at_max(monkeypatch): + """When total entries exceed _REGISTRY_MAX, _REGISTRY_EVICT_BATCH + oldest entries are dropped. Verifies the explicit math: + register MAX+1 tokens, expect MAX+1 - EVICT_BATCH remaining.""" + # Use the real cap (4096) but evict in smaller batches to keep + # the test fast. Force eviction by exceeding the cap. + monkeypatch.setattr(retry_registry, "_REGISTRY_MAX", 16) + monkeypatch.setattr(retry_registry, "_REGISTRY_EVICT_BATCH", 4) + + for _ in range(17): + register_framework_attempt() + # On the 17th register, cap is exceeded → 4 oldest dropped. + _, n_tokens = retry_registry._registry_size_for_test() + assert n_tokens == 17 - 4, ( + f"expected 13 tokens after cap-eviction, got {n_tokens}" + ) + + +def test_clear_for_thread_drops_all_tokens(): + """clear_for_thread() drops ALL tokens for the current TID. + Used as a parent-span-end orphan-cleanup primitive: when the + framework wrapper's parent span ends, any tokens still + registered are by definition orphans.""" + t1 = register_framework_attempt() + t2 = register_framework_attempt() + assert is_framework_owned() + + dropped = clear_framework_attempts_for_thread() + assert dropped == 2 + assert not is_framework_owned() + + # Subsequent unregister of an already-cleared token is a safe no-op. + unregister_framework_attempt(t1) + unregister_framework_attempt(t2) + + +def test_thread_id_reuse_does_not_leak_state(): + """When an OS thread terminates and its TID is recycled by the + next thread, the new thread sees a clean bucket (because the + terminated thread's parent-end cleanup or TTL eviction + cleared its tokens before reuse). + + The registry tolerates TID reuse simply because tokens are + UUIDs — a recycled TID with no tokens has an empty bucket + (or is missing from the registry entirely), so + is_framework_owned returns False naturally. + """ + # Drive thread 1 to register + clear, simulating a clean shutdown. + state: dict[str, int] = {} + + def thread_one(): + state["tid"] = threading.get_ident() + register_framework_attempt() + clear_framework_attempts_for_thread() + + t = threading.Thread(target=thread_one) + t.start() + t.join() + + # Now drive thread 2; if the OS happens to give it the same TID + # (we can't deterministically force this, but we can verify the + # invariant holds either way), it must not see thread 1's state. + seen: dict[str, bool] = {} + + def thread_two(): + # If TID was reused, the bucket should still be empty + # because thread 1 cleared it. If TID is fresh, ditto. + seen["owned_at_start"] = is_framework_owned() + register_framework_attempt() + seen["owned_after_register"] = is_framework_owned() + clear_framework_attempts_for_thread() + + t2 = threading.Thread(target=thread_two) + t2.start() + t2.join() + + assert seen["owned_at_start"] is False + assert seen["owned_after_register"] is True + + +def test_unregister_works_from_different_thread_than_register(): + """REGRESSION GUARD (review-batch-1 Blocker fix 2026-05-10): a + framework's start callback and terminal callback can run on + different OS threads (e.g. asyncio dispatching success/failure + callbacks to a worker thread). The original implementation used + ``threading.get_ident()`` in BOTH register and unregister, so a + cross-thread unregister silently failed and the originating + thread stayed "framework-owned" until TTL. + + Fix: token → originating-TID reverse index, used by unregister + to find the right bucket regardless of which thread runs it. + """ + main_tid = threading.get_ident() + + captured: dict[str, str] = {} + + def thread_a_register(): + captured["token"] = register_framework_attempt() + captured["registered_tid"] = str(threading.get_ident()) + + ta = threading.Thread(target=thread_a_register) + ta.start() + ta.join() + + # Sanity: thread A registered. main thread is not framework-owned; + # the originating thread A's TID *is* owned. + assert not is_framework_owned(main_tid) + registered_tid = int(captured["registered_tid"]) + assert is_framework_owned(registered_tid) + + # Now unregister from a DIFFERENT thread (could be main, could be + # a third thread — anywhere except the original registrant). + unregister_framework_attempt(captured["token"]) + + # The originating TID's bucket must now be empty — i.e. NOT + # falsely reported as owned. Without the fix, this assertion + # fails until the 60s TTL kicks in. + assert not is_framework_owned(registered_tid), ( + "cross-thread unregister MUST clean up the originating TID's bucket; " + "otherwise direct-SDK suppression on the original thread persists " + "incorrectly until TTL eviction" + ) + + # Registry should be fully empty. + n_tids, n_tokens = retry_registry._registry_size_for_test() + assert n_tids == 0 and n_tokens == 0 + + +def test_concurrent_register_unregister_thread_safety(): + """Many threads registering/unregistering concurrently must not + crash, lose tokens, or report inconsistent state. Smoke-tests the + locking around _REGISTRY.""" + N_THREADS = 8 + N_OPS_PER_THREAD = 50 + + errors: list[Exception] = [] + + def worker(): + try: + tokens = [] + for _ in range(N_OPS_PER_THREAD): + tokens.append(register_framework_attempt()) + assert is_framework_owned() + for tok in tokens: + unregister_framework_attempt(tok) + assert not is_framework_owned() + except Exception as e: # pragma: no cover + errors.append(e) + + with ThreadPoolExecutor(max_workers=N_THREADS) as pool: + futures = [pool.submit(worker) for _ in range(N_THREADS)] + for f in futures: + f.result() + + assert not errors, f"concurrent ops raised: {errors}" + n_tids, n_tokens = retry_registry._registry_size_for_test() + assert n_tokens == 0, f"all tokens should be unregistered, got {n_tokens}" diff --git a/packages/opentelemetry-instrumentation-langchain/opentelemetry/instrumentation/langchain/__init__.py b/packages/opentelemetry-instrumentation-langchain/opentelemetry/instrumentation/langchain/__init__.py index 423935c015..18a496136a 100644 --- a/packages/opentelemetry-instrumentation-langchain/opentelemetry/instrumentation/langchain/__init__.py +++ b/packages/opentelemetry-instrumentation-langchain/opentelemetry/instrumentation/langchain/__init__.py @@ -1,7 +1,8 @@ """OpenTelemetry Langchain instrumentation""" +import inspect import logging -from typing import Collection +from typing import Any, Collection, Optional from opentelemetry import context as context_api @@ -12,6 +13,9 @@ TraceloopCallbackHandler, ) from opentelemetry.instrumentation.langchain.config import Config +from opentelemetry.instrumentation.langchain.retry_handler import ( + _FortifyRootRetryHandler, +) from opentelemetry.instrumentation.langchain.safety import ( instrument_safety_wrappers, uninstrument_safety_wrappers, @@ -94,10 +98,26 @@ def _instrument(self, **kwargs): traceloopCallbackHandler = TraceloopCallbackHandler( tracer, duration_histogram, token_histogram ) + # ST-10.2: register the FR retry-attempt handler alongside the + # existing Traceloop handler. Per ST-10.0 C2 POC findings, this + # captures per-HTTP-attempt callbacks on framework-layer retry + # paths (e.g. Runnable.with_retry) and emits one + # fortifyroot.langchain.retry_attempt sibling span per attempt + # under the parent_run_id's span. The handler gets the Traceloop + # handler as a CONSTRUCTION-time reference so it can resolve the + # workflow parent via Traceloop's run_id-keyed ``spans`` dict + # without mutating shared state per-callback-manager-init + # (review-round-2 Major 5). + fortifyrootRetryHandler = _FortifyRootRetryHandler( + traceloop_handler=traceloopCallbackHandler, + ) wrap_function_wrapper( "langchain_core.callbacks", "BaseCallbackManager.__init__", - _BaseCallbackManagerInitWrapper(traceloopCallbackHandler), + _BaseCallbackManagerInitWrapper( + traceloopCallbackHandler, + fortifyrootRetryHandler, + ), ) instrument_safety_wrappers() @@ -205,8 +225,13 @@ def _uninstrument(self, **kwargs): class _BaseCallbackManagerInitWrapper: - def __init__(self, callback_handler: "TraceloopCallbackHandler"): + def __init__( + self, + callback_handler: "TraceloopCallbackHandler", + retry_handler: "_FortifyRootRetryHandler" = None, + ): self._callback_handler = callback_handler + self._retry_handler = retry_handler def __call__( self, @@ -216,6 +241,18 @@ def __call__( kwargs, ) -> None: wrapped(*args, **kwargs) + # Register the existing Traceloop handler first (existing + # behavior, preserved as the canonical order). The FR retry + # handler is registered AFTER Traceloop and uses parent_run_id + # to resolve the workflow parent via Traceloop's spans dict + # — see _resolve_parent_span in retry_handler.py. This avoids + # the test-isolation bug observed when the FR retry handler + # was registered FIRST: running before Traceloop caused + # Traceloop's context-attach/detach discipline to break + # downstream (LiteLLM tests in the same pytest session + # observed stale OTel ambient context leaking from LangChain + # — review-batch-1 v6 trace-id-shared-across-tests bug, + # 2026-05-11). for handler in instance.inheritable_handlers: if isinstance(handler, type(self._callback_handler)): break @@ -225,6 +262,20 @@ def __call__( # we need a way to determine the type of CallbackManager being wrapped. self._callback_handler._callback_manager = instance instance.add_handler(self._callback_handler, True) + # ST-10.2: register the FR retry-attempt handler AFTER + # Traceloop. Idempotent registration. The handler already + # holds a CONSTRUCTION-time reference to the Traceloop handler + # (set in LangchainInstrumentor._instrument). We deliberately + # do NOT mutate any shared attribute on the retry handler per + # callback-manager-init — that previously raced when concurrent + # Runnable invocations created BaseCallbackManagers on + # different threads (review-round-2 Major 5). + if self._retry_handler is not None: + for handler in instance.inheritable_handlers: + if isinstance(handler, type(self._retry_handler)): + break + else: + instance.add_handler(self._retry_handler, True) # This class wraps a function call to inject tracing information (trace headers) into @@ -262,14 +313,51 @@ def __call__( ) # In legacy chains like LLMChain, suppressing model instrumentations - # within create_llm_span doesn't work, so this should helps as a fallback + # within create_llm_span doesn't work, so this should helps as a fallback. + # + # ST-10 review-round-2 fix (2026-05-11): capture the attach token + # and detach in ``finally`` so this suppression layer doesn't leak + # into the OTel context stack indefinitely. The pre-fix code + # never detached, accumulating SUPPRESS_LANGUAGE_MODEL_INSTRUMENTATION_KEY + # frames on every wrapped openai call, which compounded across + # LangChain tests in long pytest sessions and risked corrupting + # later instrumentor behaviour. See review-round-2 Blocker 3. + suppression_token: Optional[Any] = None try: - context_api.attach( + suppression_token = context_api.attach( context_api.set_value(SUPPRESS_LANGUAGE_MODEL_INSTRUMENTATION_KEY, True) ) except Exception: # If context setting fails, continue without suppression # This is not critical for core functionality - pass + suppression_token = None + + try: + result = wrapped(*args, **kwargs) + except Exception: + self._detach_suppression(suppression_token) + raise + + if inspect.isawaitable(result): + return self._await_with_suppression_cleanup(result, suppression_token) + + self._detach_suppression(suppression_token) + return result - return wrapped(*args, **kwargs) + async def _await_with_suppression_cleanup(self, awaitable, suppression_token): + try: + return await awaitable + finally: + self._detach_suppression(suppression_token) + + @staticmethod + def _detach_suppression(suppression_token) -> None: + if suppression_token is None: + return + try: + context_api.detach(suppression_token) + except Exception: + # Detach can fail in async/concurrent edge cases — safe to + # ignore because this is a best-effort suppression fallback + # for legacy LangChain OpenAI wrappers. + pass diff --git a/packages/opentelemetry-instrumentation-langchain/opentelemetry/instrumentation/langchain/callback_handler.py b/packages/opentelemetry-instrumentation-langchain/opentelemetry/instrumentation/langchain/callback_handler.py index 0970638229..cd5c17100d 100644 --- a/packages/opentelemetry-instrumentation-langchain/opentelemetry/instrumentation/langchain/callback_handler.py +++ b/packages/opentelemetry-instrumentation-langchain/opentelemetry/instrumentation/langchain/callback_handler.py @@ -186,9 +186,23 @@ def _end_span(self, span: Span, run_id: UUID) -> None: if child_span.end_time is None: # avoid warning on ended spans child_span.end() span.end() - token = self.spans[run_id].token - if token: - self._safe_detach_context(token) + # ST-10 review-round-2 fix (2026-05-11): detach ALL attached + # tokens in LIFO order (reverse of attach). Pre-fix, only the + # single ``token`` field was detached — which for LLM spans + # was the suppression token only, leaving the span-context + # token leaked and the ended LLM span permanently "current" + # in OTel context. See SpanHolder.tokens. + tokens_to_detach = list(self.spans[run_id].tokens) + for tok in reversed(tokens_to_detach): + if tok is not None: + self._safe_detach_context(tok) + # Legacy single-token path retained for any direct external + # readers of SpanHolder.token (unchanged contract). If the + # legacy token isn't already part of ``tokens`` (defensive), + # detach it too — _safe_detach_context is idempotent. + legacy_token = self.spans[run_id].token + if legacy_token is not None and legacy_token not in tokens_to_detach: + self._safe_detach_context(legacy_token) del self.spans[run_id] @@ -250,6 +264,19 @@ def _create_span( entity_path: str = "", metadata: Optional[dict[str, Any]] = None, ) -> Span: + # ST-10 review-round-2 fix (2026-05-11): capture every + # context_api.attach()'s return token and append to the + # SpanHolder's ``tokens`` list. ``_end_span`` detaches them in + # LIFO order. Pre-fix, the metadata-association attach below + # was made WITHOUT capturing its token, and the span-context + # attach later overwrote the SpanHolder's single ``token`` + # field with the suppression token (in ``_create_llm_span``). + # Result: every LangChain LLM span permanently leaked its + # span-context attach, leaving the ended span "current" in OTel + # context past the end of the LangChain test, which polluted + # later LiteLLM / LlamaIndex tests in the same pytest session + # (session-scoped sdk_helper). See review-round-2 Blocker 1. + attached_tokens: list[Any] = [] if metadata is not None: current_association_properties = ( context_api.get_value("association_properties") or {} @@ -261,12 +288,14 @@ def _create_span( if v is not None } try: - context_api.attach( + metadata_token = context_api.attach( context_api.set_value( "association_properties", {**current_association_properties, **sanitized_metadata}, ) ) + if metadata_token is not None: + attached_tokens.append(metadata_token) except Exception: # If setting association properties fails, continue without them # This doesn't affect the core span functionality @@ -281,7 +310,9 @@ def _create_span( else: span = self.tracer.start_span(span_name, kind=kind) - token = self._safe_attach_context(span) + span_context_token = self._safe_attach_context(span) + if span_context_token is not None: + attached_tokens.append(span_context_token) _set_span_attribute(span, SpanAttributes.TRACELOOP_WORKFLOW_NAME, workflow_name) _set_span_attribute(span, SpanAttributes.TRACELOOP_ENTITY_PATH, entity_path) @@ -296,7 +327,8 @@ def _create_span( ) self.spans[run_id] = SpanHolder( - span, token, None, [], workflow_name, entity_name, entity_path + span, span_context_token, None, [], workflow_name, entity_name, entity_path, + tokens=attached_tokens, ) if parent_run_id is not None and parent_run_id in self.spans: @@ -362,17 +394,29 @@ def _create_llm_span( # we already have an LLM span by this point, # so skip any downstream instrumentation from here + # + # ST-10 review-round-2 fix (2026-05-11): APPEND the suppression + # token to the existing SpanHolder.tokens list rather than + # replacing the SpanHolder. The pre-fix code created a new + # SpanHolder with ONLY the suppression token, dropping the + # span-context token that was attached by ``_create_span()`` — + # so ``_end_span()`` only detached the suppression token and + # the LLM span stayed "current" in OTel context indefinitely. + # See SpanHolder.tokens for rationale. + suppression_token: Optional[Any] = None try: - token = context_api.attach( + suppression_token = context_api.attach( context_api.set_value(SUPPRESS_LANGUAGE_MODEL_INSTRUMENTATION_KEY, True) ) except Exception: # If context setting fails, continue without suppression token - token = None + suppression_token = None - self.spans[run_id] = SpanHolder( - span, token, None, [], workflow_name, None, entity_path - ) + if suppression_token is not None: + # ``_create_span`` already put a SpanHolder at + # ``self.spans[run_id]`` with the span-context (and + # optionally metadata) tokens — append, do NOT replace. + self.spans[run_id].tokens.append(suppression_token) return span @@ -460,6 +504,16 @@ def on_chain_end( self._end_span(span, run_id) if parent_run_id is None: + # ST-10 review-round-2 note (2026-05-11): pre-existing leak — + # this attach is not paired with a detach. It writes + # SUPPRESS_LANGUAGE_MODEL_INSTRUMENTATION_KEY=False to a new + # context layer that grows the OTel context stack by one + # frame per chain-end. The pre-existing behaviour is + # preserved here because legacy LLMChain paths may depend on + # it; it's NOT the trace-id-leak root cause (the trace-id + # leak was caused by the LLM-span context-token loss in + # ``_create_llm_span`` / ``_end_span``, fixed above). Tracked + # as a follow-up for upstream Traceloop cleanup. try: context_api.attach( context_api.set_value( diff --git a/packages/opentelemetry-instrumentation-langchain/opentelemetry/instrumentation/langchain/retry_handler.py b/packages/opentelemetry-instrumentation-langchain/opentelemetry/instrumentation/langchain/retry_handler.py new file mode 100644 index 0000000000..9a933317bd --- /dev/null +++ b/packages/opentelemetry-instrumentation-langchain/opentelemetry/instrumentation/langchain/retry_handler.py @@ -0,0 +1,653 @@ +"""ST-10.2 retry-aware emission for LangChain. + +Per RETRY_LOOP.md §4.4 LangChain row + §4.4.2 coverage limitation + +ST-10.0 C2 empirical verification (POC results in +phase_st10_retryloop.txt): + + - Hook BOTH ``on_chat_model_start`` (chat models, F1 finding) + AND ``on_llm_start`` (legacy completion LLMs). + - Per-HTTP-attempt firing is verified on framework-layer retry + paths (e.g. ``Runnable.with_retry``); provider-SDK-internal + retries (e.g. ``ChatOpenAI(max_retries=N)``) fire callbacks + once per logical call → §4.4.2 coverage limitation applies. + - Use ``run_id`` as the per-attempt correlation key. + - Emit ``fortifyroot.langchain.retry_attempt`` sibling spans + under the parent_run_id's span (NOT under the per-LLM + Traceloop span — siblings need to share a parent for + RetryDetectorProc's grouping to work). + - §4.7.1: register/unregister framework-attempt tokens so + direct-SDK wrappers (ST-10.4) suppress their own emission. + - §4.5: set has_retry_attempt_child=true on the parent AFTER + the first qualifying retry_attempt successfully starts. + +This handler is registered ALONGSIDE the existing +TraceloopCallbackHandler (not as a replacement) — it captures +metadata only (model, tokens, status), so its placement is not +safety-critical. The existing handler continues to do prompt / +completion attribute capture; the §4.5 backend dedup correctly +skips its per-attempt spans because they're non-retry siblings of +the retry_attempts emitted here. +""" + +from __future__ import annotations + +import json +import logging +import threading +import time +from typing import Any, Optional +from uuid import UUID + +from langchain_core.callbacks import BaseCallbackHandler +from opentelemetry import trace +from opentelemetry.instrumentation.fortifyroot import ( + register_framework_attempt, + unregister_framework_attempt, +) +from opentelemetry.instrumentation.langchain.span_utils import _message_type_to_role +from opentelemetry.instrumentation.langchain.utils import ( + CallbackFilteredJSONEncoder, + should_send_prompts, +) +from opentelemetry.instrumentation.langchain.version import __version__ +from opentelemetry.semconv._incubating.attributes import ( + gen_ai_attributes as GenAIAttributes, +) +from opentelemetry.trace import SpanKind, Status, StatusCode, set_span_in_context + +logger = logging.getLogger(__name__) + +# ST-10 §4.4: per-attempt sibling span name + role. +_FR_RETRY_ATTEMPT_SPAN_NAME = "fortifyroot.langchain.retry_attempt" +_FR_SPAN_ROLE_KEY = "fortifyroot.span.role" +_FR_SPAN_ROLE_RETRY_ATTEMPT = "retry_attempt" + +# ST-10 §4.5 parent marker. +_FR_HAS_RETRY_ATTEMPT_CHILD_KEY = "fortifyroot.span.has_retry_attempt_child" + + +# ---------------------------------------------------------------------- +# Per-attempt correlation map. Key = LangChain run_id (UUID per attempt +# on framework-layer retry paths — verified ST-10.0 C2 POC). Value = +# {span, started_at_monotonic, framework_token, ended}. Bounded-size + +# TTL eviction defends against framework crashes that leave attempts +# open. Mirrors the LiteLLM map's shape for consistency across wrappers. +# ---------------------------------------------------------------------- + +_FR_RETRY_ATTEMPT_MAP: dict[UUID, dict] = {} +_FR_RETRY_ATTEMPT_MAP_LOCK = threading.Lock() +_FR_RETRY_ATTEMPT_MAP_MAX = 4096 +_FR_RETRY_ATTEMPT_MAP_TTL_SEC = 60.0 +_FR_RETRY_ATTEMPT_EVICT_BATCH = 1024 +_FR_RETRY_ATTEMPT_EVICT_WARN_EVERY = 32 +_FR_RETRY_ATTEMPT_EVICT_WARN_COUNTER = 0 + + +# Per-parent-run-id parent-span memo. The first time we see a given +# parent_run_id, we capture the then-current ambient OTel span as the +# "parent" for retry_attempt sibling-grouping. Subsequent retries +# (which share parent_run_id under Runnable.with_retry) reuse the +# memoised span. This gives RetryDetectorProc a stable parent to group +# by even though each attempt has a fresh run_id. + +_FR_PARENT_SPAN_BY_PARENT_RUN_ID: dict[UUID, "trace.Span"] = {} +_FR_PARENT_SPAN_LOCK = threading.Lock() +# Same TTL + cap as the attempt map (parent memo grows alongside it). +_FR_PARENT_SPAN_MAX = 4096 +_FR_PARENT_SPAN_TTL_SEC = 60.0 +_FR_PARENT_SPAN_INSERT_TIMES: dict[UUID, float] = {} + + +def _evict_stale_attempts_locked(now: float) -> int: + """Drop entries older than the TTL. Caller MUST hold the map lock. + Returns count evicted.""" + cutoff = now - _FR_RETRY_ATTEMPT_MAP_TTL_SEC + stale = [k for k, v in _FR_RETRY_ATTEMPT_MAP.items() if v["started_at"] < cutoff] + for k in stale: + entry = _FR_RETRY_ATTEMPT_MAP.pop(k, None) + if entry is None: + continue + try: + sp = entry.get("span") + if sp is not None and not entry.get("ended"): + sp.set_status(Status(StatusCode.ERROR, "retry_attempt orphaned (framework crashed)")) + sp.end() + except Exception: + pass + try: + unregister_framework_attempt(entry.get("framework_token")) + except Exception: + pass + return len(stale) + + +def _enforce_attempt_max_locked() -> int: + if len(_FR_RETRY_ATTEMPT_MAP) <= _FR_RETRY_ATTEMPT_MAP_MAX: + return 0 + items = sorted(_FR_RETRY_ATTEMPT_MAP.items(), key=lambda kv: kv[1]["started_at"]) + to_drop = items[:_FR_RETRY_ATTEMPT_EVICT_BATCH] + for k, entry in to_drop: + _FR_RETRY_ATTEMPT_MAP.pop(k, None) + try: + sp = entry.get("span") + if sp is not None and not entry.get("ended"): + sp.set_status(Status(StatusCode.ERROR, "retry_attempt cap-evicted")) + sp.end() + except Exception: + pass + try: + unregister_framework_attempt(entry.get("framework_token")) + except Exception: + pass + return len(to_drop) + + +def _maybe_warn_attempt_eviction(evicted: int) -> None: + global _FR_RETRY_ATTEMPT_EVICT_WARN_COUNTER + _FR_RETRY_ATTEMPT_EVICT_WARN_COUNTER += evicted + if _FR_RETRY_ATTEMPT_EVICT_WARN_COUNTER < _FR_RETRY_ATTEMPT_EVICT_WARN_EVERY: + return + _FR_RETRY_ATTEMPT_EVICT_WARN_COUNTER = 0 + logger.warning( + "fortifyroot langchain retry_attempt map: evicted %d+ stale/over-cap entries; " + "framework may be leaking attempts (TTL=%.0fs, max=%d)", + evicted, + _FR_RETRY_ATTEMPT_MAP_TTL_SEC, + _FR_RETRY_ATTEMPT_MAP_MAX, + ) + + +def _evict_stale_parent_spans_locked(now: float) -> int: + """Drop parent-span memos older than the TTL. Caller holds lock.""" + cutoff = now - _FR_PARENT_SPAN_TTL_SEC + stale = [k for k, v in _FR_PARENT_SPAN_INSERT_TIMES.items() if v < cutoff] + for k in stale: + _FR_PARENT_SPAN_BY_PARENT_RUN_ID.pop(k, None) + _FR_PARENT_SPAN_INSERT_TIMES.pop(k, None) + return len(stale) + + +def _enforce_parent_span_max_locked() -> int: + if len(_FR_PARENT_SPAN_BY_PARENT_RUN_ID) <= _FR_PARENT_SPAN_MAX: + return 0 + items = sorted(_FR_PARENT_SPAN_INSERT_TIMES.items(), key=lambda kv: kv[1]) + to_drop = items[:_FR_RETRY_ATTEMPT_EVICT_BATCH] + for k, _ in to_drop: + _FR_PARENT_SPAN_BY_PARENT_RUN_ID.pop(k, None) + _FR_PARENT_SPAN_INSERT_TIMES.pop(k, None) + return len(to_drop) + + +def _resolve_parent_span( + parent_run_id: Optional[UUID], + traceloop_handler: Optional[Any] = None, +) -> Optional["trace.Span"]: + """Resolve the OTel parent span for a retry_attempt under the given + parent_run_id. Multi-attempt retries share parent_run_id, so the + retry_attempts must share one OTel parent → RetryDetectorProc can + group them. + + Resolution strategy (revised 2026-05-11 after review-round-2): + + Strategy A — TRACELOOP SPANS-DICT LOOKUP (load-bearing when + parent_run_id is set): + The TraceloopCallbackHandler maintains a ``spans`` dict keyed + by run_id → SpanHolder. When parent_run_id is set, we look up + the parent's SpanHolder there and use its OTel span as the + workflow parent. This gives the correct answer even when + Traceloop has already attached its per-LLM span as the OTel + ambient (which would otherwise be the visible "current span" + and would break sibling-grouping across retries). + + The ``traceloop_handler`` is passed as a DIRECT reference at + ``_FortifyRootRetryHandler`` construction time (NOT via a + mutable shared back-reference to BaseCallbackManager). This + fixes the review-round-2 Major-5 concern: a single shared + ``_FortifyRootRetryHandler`` instance was being mutated + per-BaseCallbackManager-init, racing concurrent callbacks + onto the wrong manager's spans dict. + + Strategy B — AMBIENT FALLBACK (root invocation OR no-Traceloop + standalone use): + Used when: + (a) parent_run_id is None (root invocation with no enclosing + chain / runnable), OR + (b) no traceloop_handler reference was provided to the + retry handler (standalone / test scenario — production + ALWAYS wires this via LangchainInstrumentor). + In both cases there is no risk of Traceloop having attached + its per-LLM span as the OTel ambient, so ambient is the + correct workflow parent. + + No-emission policy (review-round-2 Blocker 2): + If parent_run_id IS set AND traceloop_handler is provided + (production wiring) BUT Strategy A's Traceloop lookup fails + (parent_run_id not in its spans dict — e.g. evicted, stale), + DO NOT fall back to ambient. Ambient at that moment is likely + Traceloop's per-LLM span (handler-order is Traceloop-first) + and parenting under it would break sibling-grouping. Return + None; the caller skips emission with a debug log. The §4.5 + backend dedup degrades gracefully (no retry_attempt → parent + stays canonical → single LLMUsageEvent per call). + """ + # Strategy A: lookup Traceloop's parent SpanHolder via parent_run_id. + if parent_run_id is not None and traceloop_handler is not None: + try: + span_holder = getattr(traceloop_handler, "spans", {}).get(parent_run_id) + if span_holder is not None: + span = getattr(span_holder, "span", None) + if span is not None: + return span + except Exception: + logger.debug("Traceloop spans-dict lookup failed", exc_info=True) + # If parent_run_id was set AND traceloop_handler was provided + # but lookup failed, refuse to fall back to ambient — see + # "No-emission policy" above. + return None + + # Strategy B: ambient fallback. + # - parent_run_id is None (legitimate root), OR + # - traceloop_handler is None (standalone / test). + current = trace.get_current_span() + ctx = current.get_span_context() if current is not None else None + if ctx is None or not ctx.is_valid: + return None + return current + + +def _resolve_routed_provider(serialized: Optional[dict], invocation_params: Optional[dict]) -> Optional[str]: + """Best-effort: derive the routed provider for the gen_ai.system + attribute. Falls back to inspecting the serialized handler / + invocation_params dicts; returns None if undeterminable. + + LangChain's serialized dict typically contains 'id' = list of + module path components, e.g. ['langchain_openai', 'chat_models', + 'base', 'ChatOpenAI']. The first segment ('langchain_openai', + 'langchain_anthropic', 'langchain_aws') maps to a provider name. + """ + if isinstance(serialized, dict): + ids = serialized.get("id") + if isinstance(ids, list) and ids: + first = str(ids[0]).lower() + if first.startswith("langchain_"): + provider = first[len("langchain_"):] + # Common provider mappings. + mapping = { + "openai": "openai", + "anthropic": "anthropic", + "aws": "AWS", # Bedrock SDK + "google_genai": "google", + "google_vertexai": "google", + } + return mapping.get(provider, provider) + if isinstance(invocation_params, dict): + # Some langchain integrations expose an explicit provider field. + cand = invocation_params.get("_type") or invocation_params.get("model_provider") + if isinstance(cand, str): + return cand + return None + + +def _resolve_model(serialized: Optional[dict], invocation_params: Optional[dict]) -> Optional[str]: + if isinstance(invocation_params, dict): + for key in ("model", "model_name", "deployment_name"): + v = invocation_params.get(key) + if isinstance(v, str) and v: + return v + if isinstance(serialized, dict): + kwargs = serialized.get("kwargs") or {} + for key in ("model", "model_name", "deployment_name"): + v = kwargs.get(key) + if isinstance(v, str) and v: + return v + return None + + +def _content_to_string(content: Any) -> str: + if isinstance(content, str): + return content + try: + return json.dumps(content, cls=CallbackFilteredJSONEncoder) + except Exception: + return str(content) + + +def _add_prompt_attrs( + attrs: dict[str, Any], + *, + messages: Optional[list[list[Any]]] = None, + prompts: Optional[list[str]] = None, +) -> None: + """Copy LangChain's request content onto the retry_attempt span. + + Backend §4.5 makes retry_attempt the canonical LLMUsageEvent span + when it exists. Safety E2E tests and customers looking up the + canonical event therefore still need the same prompt content that + Traceloop's normal LLM span carries. The callback receives prompts + after FR prompt-safety masking, so these attributes preserve the + existing masked/plaintext semantics. + """ + try: + if not should_send_prompts(): + return + except Exception: + return + + if prompts is not None: + for i, prompt in enumerate(prompts): + if not isinstance(prompt, str): + continue + attrs[f"{GenAIAttributes.GEN_AI_PROMPT}.{i}.role"] = "user" + attrs[f"{GenAIAttributes.GEN_AI_PROMPT}.{i}.content"] = prompt + return + + if messages is None: + return + + i = 0 + for message_group in messages: + for msg in message_group: + msg_type = getattr(msg, "type", None) + if isinstance(msg_type, str): + attrs[f"{GenAIAttributes.GEN_AI_PROMPT}.{i}.role"] = _message_type_to_role(msg_type) + content = getattr(msg, "content", None) + if content is not None: + attrs[f"{GenAIAttributes.GEN_AI_PROMPT}.{i}.content"] = _content_to_string(content) + i += 1 + + +def _start_retry_attempt( + run_id: UUID, + parent_run_id: Optional[UUID], + serialized: Optional[dict], + invocation_params: Optional[dict], + traceloop_handler: Optional[Any] = None, + messages: Optional[list[list[Any]]] = None, + prompts: Optional[list[str]] = None, +) -> None: + """Open the retry_attempt sibling span and register state. + Idempotent if called twice for the same run_id (defensive — the + callback contract should fire it exactly once, but we don't crash + on duplicates).""" + if run_id is None: + return + parent_span = _resolve_parent_span(parent_run_id, traceloop_handler=traceloop_handler) + if parent_span is None: + # No ambient parent → orphan retry_attempt would have no place + # in the trace tree. Skip emission. The §4.5 backend dedup + # degrades gracefully when no retry_attempt exists. + logger.debug( + "no ambient parent span for langchain retry_attempt; skipping emission " + "(run_id=%s, parent_run_id=%s)", run_id, parent_run_id, + ) + return + + routed_provider = _resolve_routed_provider(serialized, invocation_params) + model = _resolve_model(serialized, invocation_params) + + attrs: dict[str, Any] = { + _FR_SPAN_ROLE_KEY: _FR_SPAN_ROLE_RETRY_ATTEMPT, + "gen_ai.operation.name": "chat", + } + if routed_provider: + attrs["gen_ai.system"] = routed_provider + if model: + attrs["gen_ai.request.model"] = model + _add_prompt_attrs(attrs, messages=messages, prompts=prompts) + + tracer = trace.get_tracer(__name__, __version__) + parent_ctx = set_span_in_context(parent_span) + span = tracer.start_span( + _FR_RETRY_ATTEMPT_SPAN_NAME, + kind=SpanKind.CLIENT, + attributes=attrs, + context=parent_ctx, + ) + + try: + framework_token = register_framework_attempt() + except Exception: + framework_token = None + logger.debug("failed to register framework attempt token", exc_info=True) + + now = time.monotonic() + with _FR_RETRY_ATTEMPT_MAP_LOCK: + evicted = _evict_stale_attempts_locked(now) + evicted += _enforce_attempt_max_locked() + if evicted: + _maybe_warn_attempt_eviction(evicted) + # Defensive: if a duplicate run_id slips through, end the + # previous span as ERROR (orphaned by the duplicate) and + # replace. + prev = _FR_RETRY_ATTEMPT_MAP.get(run_id) + if prev is not None and not prev.get("ended"): + try: + prev["span"].set_status(Status(StatusCode.ERROR, "duplicate run_id; superseded")) + prev["span"].end() + except Exception: + pass + try: + unregister_framework_attempt(prev.get("framework_token")) + except Exception: + pass + _FR_RETRY_ATTEMPT_MAP[run_id] = { + "span": span, + "started_at": now, + "framework_token": framework_token, + "ended": False, + } + + # §4.5 marker timing: set AFTER the first qualifying retry_attempt + # has successfully started under this parent. Idempotent — setting + # the attribute twice on the same parent is a no-op. + try: + parent_span.set_attribute(_FR_HAS_RETRY_ATTEMPT_CHILD_KEY, True) + except Exception: + logger.debug("failed to set has_retry_attempt_child on parent", exc_info=True) + + +def _finalize_retry_attempt( + run_id: UUID, + *, + success: bool, + response: Any = None, + error: Optional[BaseException] = None, +) -> None: + """End the retry_attempt span associated with run_id. Idempotent.""" + if run_id is None: + return + with _FR_RETRY_ATTEMPT_MAP_LOCK: + entry = _FR_RETRY_ATTEMPT_MAP.pop(run_id, None) + if entry is None or entry.get("ended"): + return + entry["ended"] = True + span = entry["span"] + framework_token = entry.get("framework_token") + + try: + if success: + # Best-effort: pull model + token usage from LLMResult. + try: + llm_output = getattr(response, "llm_output", None) or {} + if isinstance(llm_output, dict): + model_name = llm_output.get("model_name") or llm_output.get("model") + if isinstance(model_name, str) and model_name: + span.set_attribute("gen_ai.response.model", model_name) + usage = llm_output.get("token_usage") or llm_output.get("usage") or {} + if isinstance(usage, dict): + prompt_tokens = usage.get("prompt_tokens") or usage.get("input_tokens") + completion_tokens = usage.get("completion_tokens") or usage.get("output_tokens") + if isinstance(prompt_tokens, int): + span.set_attribute("gen_ai.usage.input_tokens", prompt_tokens) + if isinstance(completion_tokens, int): + span.set_attribute("gen_ai.usage.output_tokens", completion_tokens) + # generations[0][0].generation_info often has model + finish reason. + generations = getattr(response, "generations", None) + if generations and generations[0]: + gen0 = generations[0][0] + gen_info = getattr(gen0, "generation_info", None) or {} + if isinstance(gen_info, dict): + rid = gen_info.get("response_id") or gen_info.get("id") + if isinstance(rid, str) and rid: + span.set_attribute("gen_ai.response.id", rid) + except Exception: + logger.debug("failed extracting response attrs", exc_info=True) + span.set_status(Status(StatusCode.OK)) + else: + if error is not None: + error_type = type(error).__name__ + span.set_attribute("error.type", error_type) + # Best-effort: pull HTTP status from the error if it + # carries one (e.g. langchain-openai wraps openai's + # APIStatusError, which has .status_code). + status_code = getattr(error, "status_code", None) or getattr( + getattr(error, "response", None), "status_code", None + ) + if isinstance(status_code, int): + span.set_attribute("http.status_code", status_code) + try: + span.record_exception(error) + except Exception: + pass + span.set_status(Status(StatusCode.ERROR, str(error))) + else: + span.set_status(Status(StatusCode.ERROR, "retry_attempt failed")) + finally: + try: + span.end() + except Exception: + logger.debug("failed to end langchain retry_attempt span", exc_info=True) + try: + unregister_framework_attempt(framework_token) + except Exception: + pass + + +class _FortifyRootRetryHandler(BaseCallbackHandler): + """LangChain BaseCallbackHandler that emits one + fortifyroot.langchain.retry_attempt sibling span per LLM-start + callback invocation. Hooks BOTH on_chat_model_start (chat models — + F1 finding from ST-10.0 C2 POC) AND on_llm_start (legacy + completion LLMs). + + Key correlation: ``run_id`` (UUID minted per attempt by LangChain). + The retry_attempt span is created as a CHILD of the parent_run_id's + span (memoised on first sighting), so multiple attempts under the + same parent_run_id become SIBLINGS that RetryDetectorProc can + group on (parent_span_id, provider, model). + """ + + # LangChain's BaseCallbackHandler is class-based (not pydantic) + # in langchain-core >=1.0; instance state is fine here. We use a + # MODULE-level map for span tracking so that multiple callback + # manager instances (one per Runnable invocation) all funnel + # through the same correlation table — the LangChain wrapper + # registers ONE _FortifyRootRetryHandler per Runnable thanks to + # _BaseCallbackManagerInitWrapper, but module-level state survives + # any handler-instance churn cleanly. + + raise_error: bool = False + run_inline: bool = True + + # Set by ``_BaseCallbackManagerInitWrapper`` at construction time + # (NOT mutated per-callback-manager-init — see review-round-2 + # Major-5). Direct reference to the sibling Traceloop handler whose + # ``spans`` dict we look up by run_id to resolve the workflow + # parent for sibling-grouping across multi-attempt retries. See + # ``_resolve_parent_span``. + _traceloop_handler: Any = None + + def __init__(self, traceloop_handler: Optional[Any] = None) -> None: + super().__init__() + self._traceloop_handler = traceloop_handler + + def on_chat_model_start( # type: ignore[override] + self, + serialized, + messages, + *, + run_id, + parent_run_id=None, + tags=None, + metadata=None, + invocation_params=None, + **kwargs, + ): + try: + _start_retry_attempt( + run_id, parent_run_id, serialized, invocation_params, + traceloop_handler=self._traceloop_handler, + messages=messages, + ) + except Exception: + logger.debug("on_chat_model_start retry-attempt-start failed", exc_info=True) + + def on_llm_start( # type: ignore[override] + self, + serialized, + prompts, + *, + run_id, + parent_run_id=None, + tags=None, + metadata=None, + invocation_params=None, + **kwargs, + ): + try: + _start_retry_attempt( + run_id, parent_run_id, serialized, invocation_params, + traceloop_handler=self._traceloop_handler, + prompts=prompts, + ) + except Exception: + logger.debug("on_llm_start retry-attempt-start failed", exc_info=True) + + def on_llm_end( # type: ignore[override] + self, + response, + *, + run_id, + parent_run_id=None, + **kwargs, + ): + try: + _finalize_retry_attempt(run_id, success=True, response=response) + except Exception: + logger.debug("on_llm_end retry-attempt-finalize failed", exc_info=True) + + def on_llm_error( # type: ignore[override] + self, + error, + *, + run_id, + parent_run_id=None, + **kwargs, + ): + try: + _finalize_retry_attempt(run_id, success=False, error=error) + except Exception: + logger.debug("on_llm_error retry-attempt-finalize failed", exc_info=True) + + +def _reset_state_for_test() -> None: + """Test-only helper: clear all module state. Not part of the + public contract; do not import outside tests.""" + with _FR_RETRY_ATTEMPT_MAP_LOCK: + for entry in _FR_RETRY_ATTEMPT_MAP.values(): + try: + if not entry.get("ended"): + entry["span"].end() + except Exception: + pass + _FR_RETRY_ATTEMPT_MAP.clear() + with _FR_PARENT_SPAN_LOCK: + _FR_PARENT_SPAN_BY_PARENT_RUN_ID.clear() + _FR_PARENT_SPAN_INSERT_TIMES.clear() + + +__all__ = [ + "_FortifyRootRetryHandler", + "_FR_RETRY_ATTEMPT_SPAN_NAME", + "_FR_HAS_RETRY_ATTEMPT_CHILD_KEY", +] diff --git a/packages/opentelemetry-instrumentation-langchain/opentelemetry/instrumentation/langchain/span_utils.py b/packages/opentelemetry-instrumentation-langchain/opentelemetry/instrumentation/langchain/span_utils.py index f7d71b18ab..49dab634ad 100644 --- a/packages/opentelemetry-instrumentation-langchain/opentelemetry/instrumentation/langchain/span_utils.py +++ b/packages/opentelemetry-instrumentation-langchain/opentelemetry/instrumentation/langchain/span_utils.py @@ -28,7 +28,7 @@ @dataclass class SpanHolder: span: Span - token: Any + token: Any # Retained for backward compat; new code should use ``tokens``. context: Context children: list[UUID] workflow_name: str @@ -36,6 +36,14 @@ class SpanHolder: entity_path: str start_time: float = field(default_factory=time.time) request_model: Optional[str] = None + # ST-10 review-round-2 fix (2026-05-11): every context_api.attach() + # performed for this span — span-context, suppression, metadata + # association_properties — appended here in attach order. ``_end_span`` + # detaches them in REVERSE order (LIFO) so OTel's context stack is + # popped to the correct level. The legacy ``token`` field above is + # kept as-is for code that reads it directly; new code should append + # to ``tokens`` and detach via this list. + tokens: list[Any] = field(default_factory=list) def _message_type_to_role(message_type: str) -> str: diff --git a/packages/opentelemetry-instrumentation-langchain/tests/conftest.py b/packages/opentelemetry-instrumentation-langchain/tests/conftest.py index b41cf016db..ca4fc44b18 100644 --- a/packages/opentelemetry-instrumentation-langchain/tests/conftest.py +++ b/packages/opentelemetry-instrumentation-langchain/tests/conftest.py @@ -26,10 +26,23 @@ pytest_plugins = [] +_FR_SPAN_NAME_PREFIX = "fortifyroot." + + +class _LegacyAssertionSpanExporter(InMemorySpanExporter): + def get_finished_spans(self): + spans = super().get_finished_spans() + return tuple( + span for span in spans if not span.name.startswith(_FR_SPAN_NAME_PREFIX) + ) + @pytest.fixture(scope="session", name="span_exporter") def fixture_span_exporter(): - exporter = InMemorySpanExporter() + # The upstream/non-FR LangChain tests assert the historical Traceloop + # spans exactly. ST-10 adds FR implementation spans; keep these legacy + # assertions focused while dedicated ST-10 retry tests use a raw exporter. + exporter = _LegacyAssertionSpanExporter() yield exporter diff --git a/packages/opentelemetry-instrumentation-langchain/tests/test_retry_attempt_emission.py b/packages/opentelemetry-instrumentation-langchain/tests/test_retry_attempt_emission.py new file mode 100644 index 0000000000..7b4878219a --- /dev/null +++ b/packages/opentelemetry-instrumentation-langchain/tests/test_retry_attempt_emission.py @@ -0,0 +1,638 @@ +"""Tests for ST-10.2 LangChain retry-aware emission. + +Covers: + - Instrumentor symmetry: _FortifyRootRetryHandler is wired up by + LangchainInstrumentor and removed at uninstrument. + - on_chat_model_start path (chat models — F1 finding from + ST-10.0 C2 POC). + - on_llm_start path (legacy completion LLMs). + - Single-attempt happy path: 1 retry_attempt span, parent has + has_retry_attempt_child=true, retry_attempt has gen_ai.system / + gen_ai.request.model / role attributes. + - Multi-attempt retry path: 3 sequential attempts sharing a + parent_run_id → 3 retry_attempt SIBLINGS under one parent + span. This is the structural shape RetryDetectorProc requires + for retry-loop detection (siblings, not nested). + - Marker timing (§4.5): parent gets the marker only AFTER the + first attempt's start callback fires. + - §4.7.1 token registration symmetry. + - No-parent guard: skip emission when no ambient parent exists. +""" + +from __future__ import annotations + +from uuid import uuid4 + +import pytest +from langchain_core.messages import HumanMessage +from opentelemetry import trace +from opentelemetry.instrumentation.fortifyroot import ( + is_framework_owned, + retry_registry, +) +from opentelemetry.instrumentation.langchain import LangchainInstrumentor +from opentelemetry.instrumentation.langchain.retry_handler import ( + _FortifyRootRetryHandler, + _FR_HAS_RETRY_ATTEMPT_CHILD_KEY, + _FR_RETRY_ATTEMPT_SPAN_NAME, + _reset_state_for_test, +) +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import ( + InMemorySpanExporter, +) + + +@pytest.fixture +def fresh_tracer(): + """A fresh TracerProvider + in-memory exporter installed as the + GLOBAL tracer provider (so the retry handler's + ``trace.get_tracer(...)`` lookups route through it). Same pattern + as the LiteLLM ST-10.1 tests.""" + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + try: + trace.set_tracer_provider(provider) + except Exception: + pass + current_provider = trace.get_tracer_provider() + if current_provider is not provider: + # Already set by an earlier test — add OUR exporter to it. + try: + current_provider.add_span_processor(SimpleSpanProcessor(exporter)) + except Exception: + pass + yield trace.get_tracer("test"), exporter, current_provider + + +@pytest.fixture(autouse=True) +def reset_state(): + """Each test starts with empty handler state + clean §4.7.1 registry.""" + retry_registry._reset_for_test() + _reset_state_for_test() + yield + retry_registry._reset_for_test() + _reset_state_for_test() + + +# --------------------------------------------------------------------------- +# Instrumentor symmetry. +# --------------------------------------------------------------------------- + +def test_instrumentor_wires_retry_handler_globally(instrument_legacy): + """When LangchainInstrumentor instruments, every newly-created + BaseCallbackManager has the _FortifyRootRetryHandler installed + via the patched __init__.""" + from langchain_core.callbacks import BaseCallbackManager + + del instrument_legacy # only requested for its side effect + # Create a manager; the patched __init__ should register both + # the existing TraceloopCallbackHandler AND our retry handler. + mgr = BaseCallbackManager(handlers=[]) + retry_handlers = [ + h for h in mgr.inheritable_handlers + if isinstance(h, _FortifyRootRetryHandler) + ] + assert len(retry_handlers) == 1, ( + f"expected exactly 1 retry handler on the manager, " + f"got {len(retry_handlers)}; " + f"handlers={[type(h).__name__ for h in mgr.inheritable_handlers]}" + ) + + +def test_traceloop_handler_registered_before_fr_retry_handler(instrument_legacy): + """REGRESSION GUARD (review-batch-1 v6 fix 2026-05-11): + LangChain dispatches callbacks in registration order. The + Traceloop handler MUST run BEFORE the FR retry handler, so + Traceloop's context-attach/detach discipline (existing, + ST-6-validated behavior) is preserved. + + History: an earlier fix attempted the OPPOSITE order (FR + retry first) so the OTel ambient at retry_attempt creation + would be the workflow span — required for retry-attempt + sibling-grouping. That ordering caused a DIFFERENT bug: + Traceloop's `on_chat_model_start` / `on_llm_end` + context-attach/detach pair lost discipline when interleaved + with the FR retry handler running first, leaving STALE + OTel context attached past the end of LangChain tests. + Subsequent LiteLLM tests in the same pytest session inherited + that stale context, causing all LiteLLM spans to share the + LangChain trace_id and break the per-test trace-isolation + invariant. All 10 LiteLLM tests failed under this scenario. + + Fix: keep Traceloop first; resolve the FR retry-attempt + parent via Traceloop's ``spans`` dict + ``parent_run_id`` + (see ``_resolve_parent_span`` Strategy A in + ``retry_handler.py``). This preserves Traceloop's + discipline AND gives the FR retry handler the right parent. + """ + from langchain_core.callbacks import BaseCallbackManager + from opentelemetry.instrumentation.langchain.callback_handler import ( + TraceloopCallbackHandler, + ) + + del instrument_legacy # only requested for its side effect + mgr = BaseCallbackManager(handlers=[]) + types = [type(h) for h in mgr.inheritable_handlers] + try: + retry_idx = next( + i for i, h in enumerate(mgr.inheritable_handlers) + if isinstance(h, _FortifyRootRetryHandler) + ) + except StopIteration: + retry_idx = None + try: + traceloop_idx = next( + i for i, h in enumerate(mgr.inheritable_handlers) + if isinstance(h, TraceloopCallbackHandler) + ) + except StopIteration: + traceloop_idx = None + assert retry_idx is not None, f"FR retry handler missing; saw {types}" + assert traceloop_idx is not None, f"Traceloop handler missing; saw {types}" + assert traceloop_idx < retry_idx, ( + f"Traceloop handler must precede FR retry handler in " + f"inheritable_handlers — traceloop={traceloop_idx}, " + f"retry={retry_idx}, all={types}. Wrong order breaks " + f"Traceloop's context-attach discipline and leaks OTel " + f"context across pytest test boundaries." + ) + + +def test_no_leaked_ambient_context_after_simulated_workflow(instrument_legacy): + """REGRESSION GUARD (review-round-2 end-to-end regression, 2026-05-13): + + The trace-id-leak bug that motivated this guard: a LangChain + workflow that exercises the patched BaseCallbackManager.__init__ + path could leave OTel context attached past the end of its + callback flow. Subsequent test code (e.g. a LiteLLM test in the + same pytest session) then saw a STILL-VALID ambient span carrying + the LangChain trace_id, causing the LiteLLM safety_wrapper to + parent under it and inherit the leaking trace_id. + + This test runs the LangChain callback flow synchronously + (on_chat_model_start → on_llm_end) under a workflow span, ends + the workflow span explicitly, and asserts that NOTHING is left + attached to the ambient OTel context. If anything stays attached, + a subsequent LiteLLM-style ``trace.get_current_span()`` would + return a leaked span — exactly the failure mode this guard + prevents. + + Uses the conftest ``instrument_legacy`` session-scoped fixture + (NOT a fresh instrument/uninstrument cycle) so the test plays + nicely with the rest of the session's instrumentation — see the + docstring on ``test_uninstrument_invokes_unwrap_for_callback_manager`` + for why ad-hoc instrument/uninstrument breaks subsequent tests + when LangchainInstrumentor is a Singleton. + """ + del instrument_legacy # only requested for its side effect + from langchain_core.callbacks import BaseCallbackManager + + mgr = BaseCallbackManager(handlers=[]) + + # Sanity: at the start of this test we expect no leaked + # ambient. If something earlier in the session leaked, this + # test isn't the right place to surface it — bail clean. + starting_ambient_valid = trace.get_current_span().get_span_context().is_valid + if starting_ambient_valid: + pytest.skip( + "ambient OTel context already non-empty at test start; " + "earlier test in this session may have leaked — this " + "guard test wants a clean baseline." + ) + + # Drive a synchronous chat-model-start → end pair through + # every handler the instrumentor registered (Traceloop + + # FR retry handler). + run_id = uuid4() + parent_run_id = uuid4() + serialized = {"id": ["langchain_openai", "chat_models", "base", "ChatOpenAI"]} + + for h in list(mgr.inheritable_handlers): + try: + h.on_chat_model_start( + serialized=serialized, + messages=[], + run_id=run_id, + parent_run_id=parent_run_id, + invocation_params={"model": "gpt-4o-mini"}, + ) + except Exception: + # Tolerate handler-internal errors — the guard is + # about context discipline, not handler correctness. + pass + + for h in list(mgr.inheritable_handlers): + try: + h.on_llm_end( + _FakeLLMResult(model="gpt-4o-mini"), + run_id=run_id, + parent_run_id=parent_run_id, + ) + except Exception: + pass + + # After the synthetic workflow, the ambient must be clean. + # If any handler attached a context and forgot to detach, we + # would see a still-valid ambient here — the precise failure + # mode that contaminates the subsequent test. + remaining = trace.get_current_span().get_span_context() + assert not remaining.is_valid, ( + f"OTel ambient context leaked after LangChain workflow — " + f"a future test in the same pytest session would inherit " + f"trace_id={remaining.trace_id:032x}, span_id=" + f"{remaining.span_id:016x}. This is exactly the bug the " + f"review-batch-1 v6 trace-id-shared fix was meant to prevent." + ) + + +def test_uninstrument_invokes_unwrap_for_callback_manager(): + """LangchainInstrumentor._uninstrument MUST call + ``unwrap("langchain_core.callbacks", "BaseCallbackManager.__init__")`` + so the retry-handler wrap (and the existing Traceloop wrap) are + removed. We can't assert on the GLOBAL BaseCallbackManager state + because LangchainInstrumentor is a Singleton and the session-scoped + ``instrument_legacy`` fixture in conftest.py installs the + instrumentor before our tests run — interleaving instrument/ + uninstrument cycles produces order-dependent state. + + Instead, we verify the unwrap CALL is made by patching + ``opentelemetry.instrumentation.langchain.unwrap`` and asserting + the BaseCallbackManager.__init__ target is in the call args. + """ + from unittest.mock import patch + + with patch( + "opentelemetry.instrumentation.langchain.unwrap" + ) as mock_unwrap, patch( + "opentelemetry.instrumentation.langchain.uninstrument_safety_wrappers" + ): + instrumentor = object.__new__(LangchainInstrumentor) + instrumentor.disable_trace_context_propagation = True + LangchainInstrumentor._uninstrument(instrumentor) + # First positional arg of unwrap is the module path. + unwrap_targets = [ + (call.args[0], call.args[1]) + for call in mock_unwrap.call_args_list + if len(call.args) >= 2 + ] + assert ("langchain_core.callbacks", "BaseCallbackManager.__init__") in unwrap_targets, ( + f"_uninstrument must unwrap BaseCallbackManager.__init__; " + f"observed unwrap calls: {unwrap_targets}" + ) + + +# --------------------------------------------------------------------------- +# on_chat_model_start vs on_llm_start (F1 finding from POC). +# --------------------------------------------------------------------------- + +def test_on_chat_model_start_emits_retry_attempt(fresh_tracer): + """Chat models fire on_chat_model_start (NOT on_llm_start) — F1 + finding from ST-10.0 C2 POC. The handler MUST handle it.""" + tracer, exporter, _ = fresh_tracer + handler = _FortifyRootRetryHandler() + + parent = tracer.start_span("workflow") + with trace.use_span(parent, end_on_exit=False): + run_id = uuid4() + parent_run_id = uuid4() + handler.on_chat_model_start( + serialized={"id": ["langchain_openai", "chat_models", "base", "ChatOpenAI"]}, + messages=[[HumanMessage(content="hello masked [EMAIL]")]], + run_id=run_id, + parent_run_id=parent_run_id, + invocation_params={"model": "gpt-4o-mini"}, + ) + handler.on_llm_end(_FakeLLMResult(model="gpt-4o-mini"), run_id=run_id) + parent.end() + + retry_spans = [ + s for s in exporter.get_finished_spans() + if s.name == _FR_RETRY_ATTEMPT_SPAN_NAME + ] + assert len(retry_spans) == 1, ( + f"on_chat_model_start MUST produce a retry_attempt span; got {len(retry_spans)}" + ) + rs = retry_spans[0] + assert rs.attributes.get("fortifyroot.span.role") == "retry_attempt" + assert rs.attributes.get("gen_ai.system") == "openai", ( + "gen_ai.system must be the routed provider (openai), NOT 'langchain'" + ) + assert rs.attributes.get("gen_ai.request.model") == "gpt-4o-mini" + assert rs.attributes.get("gen_ai.prompt.0.role") == "user" + assert rs.attributes.get("gen_ai.prompt.0.content") == "hello masked [EMAIL]" + + +def test_on_llm_start_emits_retry_attempt(fresh_tracer): + """Legacy completion LLMs fire on_llm_start. Same handling.""" + tracer, exporter, _ = fresh_tracer + handler = _FortifyRootRetryHandler() + + parent = tracer.start_span("workflow") + with trace.use_span(parent, end_on_exit=False): + run_id = uuid4() + parent_run_id = uuid4() + handler.on_llm_start( + serialized={"id": ["langchain_anthropic", "llms", "Anthropic"]}, + prompts=["hi"], + run_id=run_id, + parent_run_id=parent_run_id, + invocation_params={"model": "claude-haiku-4-5"}, + ) + handler.on_llm_end(_FakeLLMResult(model="claude-haiku-4-5"), run_id=run_id) + parent.end() + + retry_spans = [ + s for s in exporter.get_finished_spans() + if s.name == _FR_RETRY_ATTEMPT_SPAN_NAME + ] + assert len(retry_spans) == 1 + assert retry_spans[0].attributes.get("gen_ai.system") == "anthropic" + assert retry_spans[0].attributes.get("gen_ai.request.model") == "claude-haiku-4-5" + assert retry_spans[0].attributes.get("gen_ai.prompt.0.role") == "user" + assert retry_spans[0].attributes.get("gen_ai.prompt.0.content") == "hi" + + +# --------------------------------------------------------------------------- +# Single-attempt happy path. +# --------------------------------------------------------------------------- + +def test_single_attempt_emits_one_sibling_with_marker(fresh_tracer): + """A single attempt produces ONE retry_attempt span, parent under + the workflow, and the workflow span carries + has_retry_attempt_child=true.""" + tracer, exporter, _ = fresh_tracer + handler = _FortifyRootRetryHandler() + + parent = tracer.start_span("workflow") + with trace.use_span(parent, end_on_exit=False): + run_id = uuid4() + parent_run_id = uuid4() + handler.on_chat_model_start( + serialized={"id": ["langchain_openai", "ChatOpenAI"]}, + messages=[], + run_id=run_id, + parent_run_id=parent_run_id, + invocation_params={"model": "gpt-4o-mini"}, + ) + handler.on_llm_end( + _FakeLLMResult( + model="gpt-4o-mini", + response_id="resp-abc", + input_tokens=10, + output_tokens=5, + ), + run_id=run_id, + ) + parent.end() + + spans = exporter.get_finished_spans() + parent_exported = next(s for s in spans if s.name == "workflow") + retry_span = next(s for s in spans if s.name == _FR_RETRY_ATTEMPT_SPAN_NAME) + + assert retry_span.parent.span_id == parent_exported.context.span_id, ( + "retry_attempt MUST be a child of the workflow span (sibling-grouping requires this)" + ) + assert parent_exported.attributes.get(_FR_HAS_RETRY_ATTEMPT_CHILD_KEY) is True + assert retry_span.attributes.get("gen_ai.response.id") == "resp-abc" + assert retry_span.attributes.get("gen_ai.usage.input_tokens") == 10 + assert retry_span.attributes.get("gen_ai.usage.output_tokens") == 5 + + +# --------------------------------------------------------------------------- +# Multi-attempt retry path — siblings under one parent. +# --------------------------------------------------------------------------- + +def test_three_attempts_share_parent_under_workflow(fresh_tracer): + """When 3 attempts share the SAME parent_run_id (the canonical + LangChain Runnable.with_retry shape), all 3 retry_attempt spans + are SIBLINGS under the same workflow parent — the structural + invariant RetryDetectorProc needs for retry-loop detection.""" + tracer, exporter, _ = fresh_tracer + handler = _FortifyRootRetryHandler() + + parent = tracer.start_span("workflow") + parent_run_id = uuid4() # shared across 3 attempts + serialized = {"id": ["langchain_openai", "ChatOpenAI"]} + + class MockHTTPError(Exception): + def __init__(self, status_code: int): + self.status_code = status_code + super().__init__(f"http {status_code}") + + with trace.use_span(parent, end_on_exit=False): + # Attempt 1 — 429 + rid1 = uuid4() + handler.on_chat_model_start( + serialized=serialized, messages=[], run_id=rid1, + parent_run_id=parent_run_id, + invocation_params={"model": "gpt-4o-mini"}, + ) + handler.on_llm_error(MockHTTPError(429), run_id=rid1) + + # Attempt 2 — 429 + rid2 = uuid4() + handler.on_chat_model_start( + serialized=serialized, messages=[], run_id=rid2, + parent_run_id=parent_run_id, + invocation_params={"model": "gpt-4o-mini"}, + ) + handler.on_llm_error(MockHTTPError(429), run_id=rid2) + + # Attempt 3 — 200 + rid3 = uuid4() + handler.on_chat_model_start( + serialized=serialized, messages=[], run_id=rid3, + parent_run_id=parent_run_id, + invocation_params={"model": "gpt-4o-mini"}, + ) + handler.on_llm_end(_FakeLLMResult(model="gpt-4o-mini"), run_id=rid3) + parent.end() + + spans = exporter.get_finished_spans() + parent_exported = next(s for s in spans if s.name == "workflow") + retry_spans = [s for s in spans if s.name == _FR_RETRY_ATTEMPT_SPAN_NAME] + + assert len(retry_spans) == 3, f"expected 3 retry_attempt spans, got {len(retry_spans)}" + + # All 3 must share the workflow as parent — this is the sibling + # invariant RetryDetectorProc relies on. + parent_span_ids = {s.parent.span_id for s in retry_spans} + assert parent_span_ids == {parent_exported.context.span_id}, ( + f"all 3 retry_attempts MUST be siblings under one parent; " + f"saw distinct parents: {parent_span_ids}" + ) + + # 2 ERROR + 1 OK. + from opentelemetry.trace import StatusCode + error_count = sum(1 for s in retry_spans if s.status.status_code == StatusCode.ERROR) + ok_count = sum(1 for s in retry_spans if s.status.status_code == StatusCode.OK) + assert error_count == 2, f"expected 2 ERROR, got {error_count}" + assert ok_count == 1, f"expected 1 OK, got {ok_count}" + + # The 2 error spans must carry http.status_code=429 + error.type. + for s in retry_spans: + if s.status.status_code == StatusCode.ERROR: + assert s.attributes.get("http.status_code") == 429 + assert s.attributes.get("error.type") == "MockHTTPError" + + +# --------------------------------------------------------------------------- +# Marker timing (§4.5). +# --------------------------------------------------------------------------- + +def test_marker_set_AFTER_first_attempt_not_at_parent_creation(fresh_tracer): + """Per §4.5 marker-timing: parent gets has_retry_attempt_child=true + only AFTER the first attempt's start callback fires, NOT at parent + creation.""" + tracer, exporter, _ = fresh_tracer + handler = _FortifyRootRetryHandler() + + parent = tracer.start_span("workflow") + # Before any retry_attempt: no marker. + assert _FR_HAS_RETRY_ATTEMPT_CHILD_KEY not in dict(parent.attributes or {}) + + with trace.use_span(parent, end_on_exit=False): + run_id = uuid4() + parent_run_id = uuid4() + handler.on_chat_model_start( + serialized={"id": ["langchain_openai", "ChatOpenAI"]}, + messages=[], + run_id=run_id, + parent_run_id=parent_run_id, + invocation_params={"model": "gpt-4o-mini"}, + ) + # After first attempt start: parent has the marker. + assert dict(parent.attributes or {}).get(_FR_HAS_RETRY_ATTEMPT_CHILD_KEY) is True + + handler.on_llm_end(_FakeLLMResult(model="gpt-4o-mini"), run_id=run_id) + parent.end() + + +def test_marker_NOT_set_when_no_attempts_fire(fresh_tracer): + """If no on_chat_model_start / on_llm_start ever fires for a + parent, the parent never gets the marker — graceful degradation + per §4.5.""" + tracer, exporter, _ = fresh_tracer + + parent = tracer.start_span("workflow") + parent.end() + + parent_exported = next( + s for s in exporter.get_finished_spans() if s.name == "workflow" + ) + assert _FR_HAS_RETRY_ATTEMPT_CHILD_KEY not in (parent_exported.attributes or {}) + + +# --------------------------------------------------------------------------- +# §4.7.1 token registration. +# --------------------------------------------------------------------------- + +def test_framework_token_registered_during_attempt_unregistered_after(fresh_tracer): + """While a retry_attempt is in flight, is_framework_owned() is + True for the current thread; after the end/error callback fires, + it returns to False.""" + tracer, _, _ = fresh_tracer + handler = _FortifyRootRetryHandler() + + parent = tracer.start_span("workflow") + with trace.use_span(parent, end_on_exit=False): + assert not is_framework_owned() + run_id = uuid4() + handler.on_chat_model_start( + serialized={"id": ["langchain_openai", "ChatOpenAI"]}, + messages=[], + run_id=run_id, + parent_run_id=uuid4(), + invocation_params={"model": "gpt-4o-mini"}, + ) + assert is_framework_owned(), "during attempt → framework owns the call" + handler.on_llm_end(_FakeLLMResult(model="gpt-4o-mini"), run_id=run_id) + assert not is_framework_owned(), "after attempt → released" + parent.end() + + +# --------------------------------------------------------------------------- +# No-parent guard. +# --------------------------------------------------------------------------- + +def test_no_parent_does_not_emit_orphan_retry_attempt(fresh_tracer): + """If the handler fires with no ambient OTel parent and no + parent_run_id memo, it must NOT emit an orphan retry_attempt + (no place in the trace tree).""" + tracer, exporter, _ = fresh_tracer + handler = _FortifyRootRetryHandler() + + # No parent span attached at all. + handler.on_chat_model_start( + serialized={"id": ["langchain_openai", "ChatOpenAI"]}, + messages=[], + run_id=uuid4(), + parent_run_id=uuid4(), + invocation_params={"model": "gpt-4o-mini"}, + ) + + retry_spans = [ + s for s in exporter.get_finished_spans() + if s.name == _FR_RETRY_ATTEMPT_SPAN_NAME + ] + assert len(retry_spans) == 0 + + +# --------------------------------------------------------------------------- +# Idempotency. +# --------------------------------------------------------------------------- + +def test_double_finalize_does_not_double_end(fresh_tracer): + """Calling on_llm_end twice for the same run_id is a no-op the + second time.""" + tracer, exporter, _ = fresh_tracer + handler = _FortifyRootRetryHandler() + + parent = tracer.start_span("workflow") + with trace.use_span(parent, end_on_exit=False): + run_id = uuid4() + handler.on_chat_model_start( + serialized={"id": ["langchain_openai", "ChatOpenAI"]}, + messages=[], + run_id=run_id, + parent_run_id=uuid4(), + invocation_params={"model": "gpt-4o-mini"}, + ) + handler.on_llm_end(_FakeLLMResult(model="gpt-4o-mini"), run_id=run_id) + handler.on_llm_end(_FakeLLMResult(model="gpt-4o-mini"), run_id=run_id) # no-op + parent.end() + + retry_spans = [ + s for s in exporter.get_finished_spans() + if s.name == _FR_RETRY_ATTEMPT_SPAN_NAME + ] + assert len(retry_spans) == 1, ( + f"expected exactly 1 retry_attempt span (idempotent finalize); got {len(retry_spans)}" + ) + + +# --------------------------------------------------------------------------- +# Helpers. +# --------------------------------------------------------------------------- + +class _FakeGeneration: + def __init__(self, response_id=None): + self.generation_info = {"response_id": response_id} if response_id else {} + + +class _FakeLLMResult: + def __init__(self, model=None, response_id=None, input_tokens=None, output_tokens=None): + self.generations = [[_FakeGeneration(response_id=response_id)]] + self.llm_output = {} + if model: + self.llm_output["model_name"] = model + if input_tokens is not None or output_tokens is not None: + usage = {} + if input_tokens is not None: + usage["prompt_tokens"] = input_tokens + if output_tokens is not None: + usage["completion_tokens"] = output_tokens + self.llm_output["token_usage"] = usage diff --git a/packages/opentelemetry-instrumentation-langchain/uv.lock b/packages/opentelemetry-instrumentation-langchain/uv.lock index 1d1e61c875..ed4c346624 100644 --- a/packages/opentelemetry-instrumentation-langchain/uv.lock +++ b/packages/opentelemetry-instrumentation-langchain/uv.lock @@ -631,7 +631,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/32/6a/33d1702184d94106d3cdd7bfb788e19723206fce152e303473ca3b946c7b/greenlet-3.3.0-cp310-cp310-macosx_11_0_universal2.whl", hash = "sha256:6f8496d434d5cb2dce025773ba5597f71f5410ae499d5dd9533e0653258cdb3d", size = 273658, upload-time = "2025-12-04T14:23:37.494Z" }, { url = "https://files.pythonhosted.org/packages/d6/b7/2b5805bbf1907c26e434f4e448cd8b696a0b71725204fa21a211ff0c04a7/greenlet-3.3.0-cp310-cp310-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:b96dc7eef78fd404e022e165ec55327f935b9b52ff355b067eb4a0267fc1cffb", size = 574810, upload-time = "2025-12-04T14:50:04.154Z" }, { url = "https://files.pythonhosted.org/packages/94/38/343242ec12eddf3d8458c73f555c084359883d4ddc674240d9e61ec51fd6/greenlet-3.3.0-cp310-cp310-manylinux_2_24_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:73631cd5cccbcfe63e3f9492aaa664d278fda0ce5c3d43aeda8e77317e38efbd", size = 586248, upload-time = "2025-12-04T14:57:39.35Z" }, - { url = "https://files.pythonhosted.org/packages/f0/d0/0ae86792fb212e4384041e0ef8e7bc66f59a54912ce407d26a966ed2914d/greenlet-3.3.0-cp310-cp310-manylinux_2_24_s390x.manylinux_2_28_s390x.whl", hash = "sha256:b299a0cb979f5d7197442dccc3aee67fce53500cd88951b7e6c35575701c980b", size = 597403, upload-time = "2025-12-04T15:07:10.831Z" }, { url = "https://files.pythonhosted.org/packages/b6/a8/15d0aa26c0036a15d2659175af00954aaaa5d0d66ba538345bd88013b4d7/greenlet-3.3.0-cp310-cp310-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:7dee147740789a4632cace364816046e43310b59ff8fb79833ab043aefa72fd5", size = 586910, upload-time = "2025-12-04T14:25:59.705Z" }, { url = "https://files.pythonhosted.org/packages/e1/9b/68d5e3b7ccaba3907e5532cf8b9bf16f9ef5056a008f195a367db0ff32db/greenlet-3.3.0-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:39b28e339fc3c348427560494e28d8a6f3561c8d2bcf7d706e1c624ed8d822b9", size = 1547206, upload-time = "2025-12-04T15:04:21.027Z" }, { url = "https://files.pythonhosted.org/packages/66/bd/e3086ccedc61e49f91e2cfb5ffad9d8d62e5dc85e512a6200f096875b60c/greenlet-3.3.0-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:b3c374782c2935cc63b2a27ba8708471de4ad1abaa862ffdb1ef45a643ddbb7d", size = 1613359, upload-time = "2025-12-04T14:27:26.548Z" }, @@ -639,7 +638,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/1f/cb/48e964c452ca2b92175a9b2dca037a553036cb053ba69e284650ce755f13/greenlet-3.3.0-cp311-cp311-macosx_11_0_universal2.whl", hash = "sha256:e29f3018580e8412d6aaf5641bb7745d38c85228dacf51a73bd4e26ddf2a6a8e", size = 274908, upload-time = "2025-12-04T14:23:26.435Z" }, { url = "https://files.pythonhosted.org/packages/28/da/38d7bff4d0277b594ec557f479d65272a893f1f2a716cad91efeb8680953/greenlet-3.3.0-cp311-cp311-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:a687205fb22794e838f947e2194c0566d3812966b41c78709554aa883183fb62", size = 577113, upload-time = "2025-12-04T14:50:05.493Z" }, { url = "https://files.pythonhosted.org/packages/3c/f2/89c5eb0faddc3ff014f1c04467d67dee0d1d334ab81fadbf3744847f8a8a/greenlet-3.3.0-cp311-cp311-manylinux_2_24_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:4243050a88ba61842186cb9e63c7dfa677ec146160b0efd73b855a3d9c7fcf32", size = 590338, upload-time = "2025-12-04T14:57:41.136Z" }, - { url = "https://files.pythonhosted.org/packages/80/d7/db0a5085035d05134f8c089643da2b44cc9b80647c39e93129c5ef170d8f/greenlet-3.3.0-cp311-cp311-manylinux_2_24_s390x.manylinux_2_28_s390x.whl", hash = "sha256:670d0f94cd302d81796e37299bcd04b95d62403883b24225c6b5271466612f45", size = 601098, upload-time = "2025-12-04T15:07:11.898Z" }, { url = "https://files.pythonhosted.org/packages/dc/a6/e959a127b630a58e23529972dbc868c107f9d583b5a9f878fb858c46bc1a/greenlet-3.3.0-cp311-cp311-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:6cb3a8ec3db4a3b0eb8a3c25436c2d49e3505821802074969db017b87bc6a948", size = 590206, upload-time = "2025-12-04T14:26:01.254Z" }, { url = "https://files.pythonhosted.org/packages/48/60/29035719feb91798693023608447283b266b12efc576ed013dd9442364bb/greenlet-3.3.0-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:2de5a0b09eab81fc6a382791b995b1ccf2b172a9fec934747a7a23d2ff291794", size = 1550668, upload-time = "2025-12-04T15:04:22.439Z" }, { url = "https://files.pythonhosted.org/packages/0a/5f/783a23754b691bfa86bd72c3033aa107490deac9b2ef190837b860996c9f/greenlet-3.3.0-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:4449a736606bd30f27f8e1ff4678ee193bc47f6ca810d705981cfffd6ce0d8c5", size = 1615483, upload-time = "2025-12-04T14:27:28.083Z" }, @@ -647,7 +645,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/f8/0a/a3871375c7b9727edaeeea994bfff7c63ff7804c9829c19309ba2e058807/greenlet-3.3.0-cp312-cp312-macosx_11_0_universal2.whl", hash = "sha256:b01548f6e0b9e9784a2c99c5651e5dc89ffcbe870bc5fb2e5ef864e9cc6b5dcb", size = 276379, upload-time = "2025-12-04T14:23:30.498Z" }, { url = "https://files.pythonhosted.org/packages/43/ab/7ebfe34dce8b87be0d11dae91acbf76f7b8246bf9d6b319c741f99fa59c6/greenlet-3.3.0-cp312-cp312-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:349345b770dc88f81506c6861d22a6ccd422207829d2c854ae2af8025af303e3", size = 597294, upload-time = "2025-12-04T14:50:06.847Z" }, { url = "https://files.pythonhosted.org/packages/a4/39/f1c8da50024feecd0793dbd5e08f526809b8ab5609224a2da40aad3a7641/greenlet-3.3.0-cp312-cp312-manylinux_2_24_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:e8e18ed6995e9e2c0b4ed264d2cf89260ab3ac7e13555b8032b25a74c6d18655", size = 607742, upload-time = "2025-12-04T14:57:42.349Z" }, - { url = "https://files.pythonhosted.org/packages/77/cb/43692bcd5f7a0da6ec0ec6d58ee7cddb606d055ce94a62ac9b1aa481e969/greenlet-3.3.0-cp312-cp312-manylinux_2_24_s390x.manylinux_2_28_s390x.whl", hash = "sha256:c024b1e5696626890038e34f76140ed1daf858e37496d33f2af57f06189e70d7", size = 622297, upload-time = "2025-12-04T15:07:13.552Z" }, { url = "https://files.pythonhosted.org/packages/75/b0/6bde0b1011a60782108c01de5913c588cf51a839174538d266de15e4bf4d/greenlet-3.3.0-cp312-cp312-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:047ab3df20ede6a57c35c14bf5200fcf04039d50f908270d3f9a7a82064f543b", size = 609885, upload-time = "2025-12-04T14:26:02.368Z" }, { url = "https://files.pythonhosted.org/packages/49/0e/49b46ac39f931f59f987b7cd9f34bfec8ef81d2a1e6e00682f55be5de9f4/greenlet-3.3.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:2d9ad37fc657b1102ec880e637cccf20191581f75c64087a549e66c57e1ceb53", size = 1567424, upload-time = "2025-12-04T15:04:23.757Z" }, { url = "https://files.pythonhosted.org/packages/05/f5/49a9ac2dff7f10091935def9165c90236d8f175afb27cbed38fb1d61ab6b/greenlet-3.3.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:83cd0e36932e0e7f36a64b732a6f60c2fc2df28c351bae79fbaf4f8092fe7614", size = 1636017, upload-time = "2025-12-04T14:27:29.688Z" }, @@ -655,7 +652,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/02/2f/28592176381b9ab2cafa12829ba7b472d177f3acc35d8fbcf3673d966fff/greenlet-3.3.0-cp313-cp313-macosx_11_0_universal2.whl", hash = "sha256:a1e41a81c7e2825822f4e068c48cb2196002362619e2d70b148f20a831c00739", size = 275140, upload-time = "2025-12-04T14:23:01.282Z" }, { url = "https://files.pythonhosted.org/packages/2c/80/fbe937bf81e9fca98c981fe499e59a3f45df2a04da0baa5c2be0dca0d329/greenlet-3.3.0-cp313-cp313-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:9f515a47d02da4d30caaa85b69474cec77b7929b2e936ff7fb853d42f4bf8808", size = 599219, upload-time = "2025-12-04T14:50:08.309Z" }, { url = "https://files.pythonhosted.org/packages/c2/ff/7c985128f0514271b8268476af89aee6866df5eec04ac17dcfbc676213df/greenlet-3.3.0-cp313-cp313-manylinux_2_24_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:7d2d9fd66bfadf230b385fdc90426fcd6eb64db54b40c495b72ac0feb5766c54", size = 610211, upload-time = "2025-12-04T14:57:43.968Z" }, - { url = "https://files.pythonhosted.org/packages/79/07/c47a82d881319ec18a4510bb30463ed6891f2ad2c1901ed5ec23d3de351f/greenlet-3.3.0-cp313-cp313-manylinux_2_24_s390x.manylinux_2_28_s390x.whl", hash = "sha256:30a6e28487a790417d036088b3bcb3f3ac7d8babaa7d0139edbaddebf3af9492", size = 624311, upload-time = "2025-12-04T15:07:14.697Z" }, { url = "https://files.pythonhosted.org/packages/fd/8e/424b8c6e78bd9837d14ff7df01a9829fc883ba2ab4ea787d4f848435f23f/greenlet-3.3.0-cp313-cp313-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:087ea5e004437321508a8d6f20efc4cfec5e3c30118e1417ea96ed1d93950527", size = 612833, upload-time = "2025-12-04T14:26:03.669Z" }, { url = "https://files.pythonhosted.org/packages/b5/ba/56699ff9b7c76ca12f1cdc27a886d0f81f2189c3455ff9f65246780f713d/greenlet-3.3.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:ab97cf74045343f6c60a39913fa59710e4bd26a536ce7ab2397adf8b27e67c39", size = 1567256, upload-time = "2025-12-04T15:04:25.276Z" }, { url = "https://files.pythonhosted.org/packages/1e/37/f31136132967982d698c71a281a8901daf1a8fbab935dce7c0cf15f942cc/greenlet-3.3.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:5375d2e23184629112ca1ea89a53389dddbffcf417dad40125713d88eb5f96e8", size = 1636483, upload-time = "2025-12-04T14:27:30.804Z" }, @@ -663,7 +659,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/d7/7c/f0a6d0ede2c7bf092d00bc83ad5bafb7e6ec9b4aab2fbdfa6f134dc73327/greenlet-3.3.0-cp314-cp314-macosx_11_0_universal2.whl", hash = "sha256:60c2ef0f578afb3c8d92ea07ad327f9a062547137afe91f38408f08aacab667f", size = 275671, upload-time = "2025-12-04T14:23:05.267Z" }, { url = "https://files.pythonhosted.org/packages/44/06/dac639ae1a50f5969d82d2e3dd9767d30d6dbdbab0e1a54010c8fe90263c/greenlet-3.3.0-cp314-cp314-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:0a5d554d0712ba1de0a6c94c640f7aeba3f85b3a6e1f2899c11c2c0428da9365", size = 646360, upload-time = "2025-12-04T14:50:10.026Z" }, { url = "https://files.pythonhosted.org/packages/e0/94/0fb76fe6c5369fba9bf98529ada6f4c3a1adf19e406a47332245ef0eb357/greenlet-3.3.0-cp314-cp314-manylinux_2_24_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:3a898b1e9c5f7307ebbde4102908e6cbfcb9ea16284a3abe15cab996bee8b9b3", size = 658160, upload-time = "2025-12-04T14:57:45.41Z" }, - { url = "https://files.pythonhosted.org/packages/93/79/d2c70cae6e823fac36c3bbc9077962105052b7ef81db2f01ec3b9bf17e2b/greenlet-3.3.0-cp314-cp314-manylinux_2_24_s390x.manylinux_2_28_s390x.whl", hash = "sha256:dcd2bdbd444ff340e8d6bdf54d2f206ccddbb3ccfdcd3c25bf4afaa7b8f0cf45", size = 671388, upload-time = "2025-12-04T15:07:15.789Z" }, { url = "https://files.pythonhosted.org/packages/b8/14/bab308fc2c1b5228c3224ec2bf928ce2e4d21d8046c161e44a2012b5203e/greenlet-3.3.0-cp314-cp314-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:5773edda4dc00e173820722711d043799d3adb4f01731f40619e07ea2750b955", size = 660166, upload-time = "2025-12-04T14:26:05.099Z" }, { url = "https://files.pythonhosted.org/packages/4b/d2/91465d39164eaa0085177f61983d80ffe746c5a1860f009811d498e7259c/greenlet-3.3.0-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:ac0549373982b36d5fd5d30beb8a7a33ee541ff98d2b502714a09f1169f31b55", size = 1615193, upload-time = "2025-12-04T15:04:27.041Z" }, { url = "https://files.pythonhosted.org/packages/42/1b/83d110a37044b92423084d52d5d5a3b3a73cafb51b547e6d7366ff62eff1/greenlet-3.3.0-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:d198d2d977460358c3b3a4dc844f875d1adb33817f0613f663a656f463764ccc", size = 1683653, upload-time = "2025-12-04T14:27:32.366Z" }, @@ -671,7 +666,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/a0/66/bd6317bc5932accf351fc19f177ffba53712a202f9df10587da8df257c7e/greenlet-3.3.0-cp314-cp314t-macosx_11_0_universal2.whl", hash = "sha256:d6ed6f85fae6cdfdb9ce04c9bf7a08d666cfcfb914e7d006f44f840b46741931", size = 282638, upload-time = "2025-12-04T14:25:20.941Z" }, { url = "https://files.pythonhosted.org/packages/30/cf/cc81cb030b40e738d6e69502ccbd0dd1bced0588e958f9e757945de24404/greenlet-3.3.0-cp314-cp314t-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:d9125050fcf24554e69c4cacb086b87b3b55dc395a8b3ebe6487b045b2614388", size = 651145, upload-time = "2025-12-04T14:50:11.039Z" }, { url = "https://files.pythonhosted.org/packages/9c/ea/1020037b5ecfe95ca7df8d8549959baceb8186031da83d5ecceff8b08cd2/greenlet-3.3.0-cp314-cp314t-manylinux_2_24_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:87e63ccfa13c0a0f6234ed0add552af24cc67dd886731f2261e46e241608bee3", size = 654236, upload-time = "2025-12-04T14:57:47.007Z" }, - { url = "https://files.pythonhosted.org/packages/69/cc/1e4bae2e45ca2fa55299f4e85854606a78ecc37fead20d69322f96000504/greenlet-3.3.0-cp314-cp314t-manylinux_2_24_s390x.manylinux_2_28_s390x.whl", hash = "sha256:2662433acbca297c9153a4023fe2161c8dcfdcc91f10433171cf7e7d94ba2221", size = 662506, upload-time = "2025-12-04T15:07:16.906Z" }, { url = "https://files.pythonhosted.org/packages/57/b9/f8025d71a6085c441a7eaff0fd928bbb275a6633773667023d19179fe815/greenlet-3.3.0-cp314-cp314t-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:3c6e9b9c1527a78520357de498b0e709fb9e2f49c3a513afd5a249007261911b", size = 653783, upload-time = "2025-12-04T14:26:06.225Z" }, { url = "https://files.pythonhosted.org/packages/f6/c7/876a8c7a7485d5d6b5c6821201d542ef28be645aa024cfe1145b35c120c1/greenlet-3.3.0-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:286d093f95ec98fdd92fcb955003b8a3d054b4e2cab3e2707a5039e7b50520fd", size = 1614857, upload-time = "2025-12-04T15:04:28.484Z" }, { url = "https://files.pythonhosted.org/packages/4f/dc/041be1dff9f23dac5f48a43323cd0789cb798342011c19a248d9c9335536/greenlet-3.3.0-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:6c10513330af5b8ae16f023e8ddbfb486ab355d04467c4679c5cfe4659975dd9", size = 1676034, upload-time = "2025-12-04T14:27:33.531Z" }, @@ -1521,13 +1515,37 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/1a/8a/02435bcdc4c45ff5b0816547fa5b4613ccce2671ec1deb23818f6f284346/opentelemetry_instrumentation_bedrock-0.52.4-py3-none-any.whl", hash = "sha256:ef5cc6f0ba78c500476f2781134beac19bc83279aeaa827dd90c5b09c20d18ae", size = 19361, upload-time = "2026-02-19T13:20:58.869Z" }, ] +[[package]] +name = "opentelemetry-instrumentation-fortifyroot" +version = "0.52.6" +source = { editable = "../opentelemetry-instrumentation-fortifyroot" } +dependencies = [ + { name = "opentelemetry-api" }, +] + +[package.metadata] +requires-dist = [{ name = "opentelemetry-api", specifier = ">=1.38.0,<2" }] + +[package.metadata.requires-dev] +dev = [ + { name = "autopep8", specifier = ">=2.2.0,<3" }, + { name = "pytest", specifier = ">=8.2.2,<9" }, + { name = "ruff", specifier = ">=0.4.0" }, +] +test = [ + { name = "opentelemetry-sdk", specifier = ">=1.38.0,<2" }, + { name = "pytest", specifier = ">=8.2.2,<9" }, + { name = "pytest-asyncio", specifier = ">=0.23.7,<0.24.0" }, +] + [[package]] name = "opentelemetry-instrumentation-langchain" -version = "0.52.4" +version = "0.52.6" source = { editable = "." } dependencies = [ { name = "opentelemetry-api" }, { name = "opentelemetry-instrumentation" }, + { name = "opentelemetry-instrumentation-fortifyroot" }, { name = "opentelemetry-semantic-conventions" }, { name = "opentelemetry-semantic-conventions-ai" }, ] @@ -1575,6 +1593,7 @@ requires-dist = [ { name = "langchain", marker = "extra == 'instruments'" }, { name = "opentelemetry-api", specifier = ">=1.38.0,<2" }, { name = "opentelemetry-instrumentation", specifier = ">=0.59b0" }, + { name = "opentelemetry-instrumentation-fortifyroot", editable = "../opentelemetry-instrumentation-fortifyroot" }, { name = "opentelemetry-semantic-conventions", specifier = ">=0.59b0" }, { name = "opentelemetry-semantic-conventions-ai", specifier = ">=0.4.13,<0.5.0" }, ] diff --git a/packages/opentelemetry-instrumentation-litellm/opentelemetry/instrumentation/litellm/__init__.py b/packages/opentelemetry-instrumentation-litellm/opentelemetry/instrumentation/litellm/__init__.py index ec70deae5b..58d2fcd587 100644 --- a/packages/opentelemetry-instrumentation-litellm/opentelemetry/instrumentation/litellm/__init__.py +++ b/packages/opentelemetry-instrumentation-litellm/opentelemetry/instrumentation/litellm/__init__.py @@ -4,11 +4,16 @@ import inspect import logging import threading -from typing import Collection +import time +from typing import Collection, Optional from opentelemetry import context as context_api from opentelemetry import trace -from opentelemetry.instrumentation.fortifyroot import get_object_value +from opentelemetry.instrumentation.fortifyroot import ( + get_object_value, + register_framework_attempt, + unregister_framework_attempt, +) from opentelemetry.instrumentation.instrumentor import BaseInstrumentor from opentelemetry.instrumentation.litellm.safety import ( apply_completion_safety, @@ -52,6 +57,16 @@ # own ResourceSpans and the existing parent→child correlation in # ``buildSafetyWrapperDedupeSet`` cannot observe siblings across batches. _FR_HAS_NATIVE_OTEL_CHILD_KEY = "fortifyroot.span.has_native_otel_child" + +# ST-10 §4.5: marker on parent set AFTER the first qualifying retry_attempt +# child has started, so the FR backend's LLMUsageExtractor can dedup the +# parent + non-retry siblings cross-batch (mirrors has_native_otel_child). +_FR_HAS_RETRY_ATTEMPT_CHILD_KEY = "fortifyroot.span.has_retry_attempt_child" + +# ST-10 §4.4: per-attempt sibling span emitted by _FortifyRootRetryEmitter +# under the safety_wrapper parent. +_FR_RETRY_ATTEMPT_SPAN_NAME = "fortifyroot.litellm.retry_attempt" +_FR_SPAN_ROLE_RETRY_ATTEMPT = "retry_attempt" _FR_COMPLETION_SAFETY_MARKER = "_fortifyroot_completion_safety_applied" _FR_COMPLETION_SAFETY_FALLBACK_MARKERS: set[tuple[int, type]] = set() _FR_COMPLETION_SAFETY_MARKERS_LOCK = threading.Lock() @@ -157,6 +172,508 @@ async def _ensure_completion_safety_applied_async( ] +# Lazy-import LiteLLM's CustomLogger base for inheritance. We MUST inherit +# (not duck-type) because LiteLLM's dispatch loop gates every callback hook +# on ``isinstance(callback, CustomLogger)`` — see +# litellm_logging.py:1015 (log_pre_api_call), :1216, :1267, :2303 +# (log_success_event), :2613 (async_log_success_event), etc. A duck-typed +# class is silently SKIPPED by the dispatch, so the retry emitter never +# fires and no retry_attempt spans get emitted. +# +# The fallback to ``object`` keeps the module importable even when litellm +# isn't installed (the instrumentor's ``instrumentation_dependencies`` check +# guards actual use). +# +# Discovered end-to-end during ST-10 review-batch-1 re-verification 2026-05-10 +# after local vendoring. The pre-existing _FortifyRootCompletionLogger is +# also duck-typed (same latent bug) but its primary safety-masking path is +# synchronous inside _finalize_response, so its callback never firing is +# masked in production. The retry emitter has no such backup — purely +# callback-driven — which made this bug observable. +try: + from litellm.integrations.custom_logger import CustomLogger as _LiteLLMCustomLoggerBase +except ImportError: + _LiteLLMCustomLoggerBase = object # type: ignore[assignment,misc] + + +# ---------------------------------------------------------------------- +# ST-10 §4.4 / §4.3: retry-attempt sibling-span emission via a second +# LiteLLM CustomLogger. +# ---------------------------------------------------------------------- +# +# The emitter opens one ``fortifyroot.litellm.retry_attempt`` sibling span +# per attempt-start callback fired by LiteLLM, and ends it on the matching +# success/failure callback. Per ST-10.0 C1 source-verified findings, this +# fires per-attempt only on the ``completion_with_retries(num_retries=N)`` +# and ``Router(...)`` retry surfaces — see RETRY_LOOP.md §4.4.2 for the +# documented coverage limitation on the ``completion(num_retries=N)`` +# path (where retries delegate to the underlying provider SDK and are +# invisible to LiteLLM's callback layer). + +# Per-attempt correlation map (§4.3). Key = LiteLLM's per-call +# ``litellm_call_id`` (sufficient because each attempt at the +# observable surfaces gets a fresh ID — verified ST-10.0 C2 POC). Value = +# {span, started_at_monotonic, parent_span, framework_token, ended}. +# Bounded-size + TTL eviction defends against framework crashes that +# leave attempts open. +_FR_RETRY_ATTEMPT_MAP: dict[str, dict] = {} +_FR_RETRY_ATTEMPT_MAP_LOCK = threading.Lock() +_FR_RETRY_ATTEMPT_MAP_MAX = 4096 +_FR_RETRY_ATTEMPT_MAP_TTL_SEC = 60.0 +_FR_RETRY_ATTEMPT_EVICT_BATCH = 1024 +# Rate-limit eviction warnings to avoid log floods on pathological leaks. +_FR_RETRY_ATTEMPT_EVICT_WARN_EVERY = 32 +_FR_RETRY_ATTEMPT_EVICT_WARN_COUNTER = 0 + + +def _evict_stale_retry_attempts_locked(now: float) -> int: + """Drop entries older than the TTL. Caller MUST hold the map lock. + Returns count evicted.""" + cutoff = now - _FR_RETRY_ATTEMPT_MAP_TTL_SEC + stale = [k for k, v in _FR_RETRY_ATTEMPT_MAP.items() if v["started_at"] < cutoff] + for k in stale: + entry = _FR_RETRY_ATTEMPT_MAP.pop(k, None) + if entry is None: + continue + # Best-effort: end the leaked span and unregister its framework token. + try: + sp = entry.get("span") + if sp is not None and not entry.get("ended"): + sp.set_status(Status(StatusCode.ERROR, "retry_attempt orphaned (framework crashed)")) + sp.end() + except Exception: + pass + try: + unregister_framework_attempt(entry.get("framework_token")) + except Exception: + pass + return len(stale) + + +def _enforce_retry_attempt_max_locked() -> int: + """Cap-evict the oldest _FR_RETRY_ATTEMPT_EVICT_BATCH entries when + the map exceeds _FR_RETRY_ATTEMPT_MAP_MAX. Caller MUST hold the + map lock. Returns count evicted.""" + if len(_FR_RETRY_ATTEMPT_MAP) <= _FR_RETRY_ATTEMPT_MAP_MAX: + return 0 + items = sorted(_FR_RETRY_ATTEMPT_MAP.items(), key=lambda kv: kv[1]["started_at"]) + to_drop = items[:_FR_RETRY_ATTEMPT_EVICT_BATCH] + for k, entry in to_drop: + _FR_RETRY_ATTEMPT_MAP.pop(k, None) + try: + sp = entry.get("span") + if sp is not None and not entry.get("ended"): + sp.set_status(Status(StatusCode.ERROR, "retry_attempt cap-evicted")) + sp.end() + except Exception: + pass + try: + unregister_framework_attempt(entry.get("framework_token")) + except Exception: + pass + return len(to_drop) + + +def _maybe_warn_retry_attempt_eviction(evicted: int) -> None: + global _FR_RETRY_ATTEMPT_EVICT_WARN_COUNTER + _FR_RETRY_ATTEMPT_EVICT_WARN_COUNTER += evicted + if _FR_RETRY_ATTEMPT_EVICT_WARN_COUNTER < _FR_RETRY_ATTEMPT_EVICT_WARN_EVERY: + return + _FR_RETRY_ATTEMPT_EVICT_WARN_COUNTER = 0 + logger.warning( + "fortifyroot litellm retry_attempt map: evicted %d+ stale/over-cap entries; " + "framework may be leaking attempts (TTL=%.0fs, max=%d)", + evicted, + _FR_RETRY_ATTEMPT_MAP_TTL_SEC, + _FR_RETRY_ATTEMPT_MAP_MAX, + ) + + +def _resolve_routed_provider(kwargs) -> Optional[str]: + """Best-effort: derive the ROUTED provider (e.g. ``openai``) from + LiteLLM kwargs for the retry_attempt span's ``gen_ai.system`` + attribute. Per RETRY_LOOP.md §4.2, this is the routed provider + (NOT the framework name). Falls back to ``litellm`` if undetermined. + """ + candidate = ( + kwargs.get("custom_llm_provider") + or kwargs.get("provider") + or (kwargs.get("model_response_object") and get_object_value(kwargs["model_response_object"], "model")) + ) + raw: Optional[str] = None + if isinstance(candidate, str) and "/" in candidate: + # e.g. "openai/gpt-4o-mini" — take the prefix. + raw = candidate.split("/", 1)[0] + elif isinstance(candidate, str) and candidate: + raw = candidate + else: + model = kwargs.get("model") + if isinstance(model, str): + model_lower = model.lower() + if "/" in model: + raw = model.split("/", 1)[0] + elif model_lower.startswith("claude-"): + # LiteLLM callbacks can pass Anthropic-routed calls as + # bare Claude model ids (for example + # ``claude-4-sonnet-20250514``) even when the public + # call used ``anthropic/``. Without this inference + # retry_attempt spans fall back to gen_ai.system="litellm", + # so the backend stores the canonical ST-10 event under the + # framework rather than the routed provider. + raw = "anthropic" + elif model and "." in model: + # No explicit provider AND no slash, but the model + # carries a Bedrock-style prefix like + # ``amazon.nova-lite-v1:0`` or + # ``anthropic.claude-3-5-sonnet`` — pass to the + # normaliser so it can recognise the prefix and + # map to "AWS". Bare model strings without "." or + # "/" (e.g. ``gpt-4o-mini``) carry no provider + # signal; we return None so the wrapper falls + # back to ``"litellm"`` rather than guessing. + raw = model + if raw is None: + return None + return _normalize_routed_provider(raw) + + +# Normalisation for the ``gen_ai.system`` attribute. +# +# Per RETRY_LOOP.md §4.2, the value MUST be the ROUTED provider, NOT +# the framework, AND it MUST be the canonical OTel-semconv form (e.g. +# Bedrock = ``"AWS"``). LiteLLM's ``custom_llm_provider`` field uses +# its own taxonomy (``"bedrock"``, ``"bedrock_converse"``, +# ``"sagemaker"``, ...), so we map those to the §4.2 canonical +# values and leave already-canonical values untouched. (Review-batch-1 +# Minor 4 fix 2026-05-10 — keeps cross-wrapper consistency with +# LangChain's _resolve_routed_provider which already normalises +# langchain_aws → ``"AWS"``.) +_LITELLM_PROVIDER_NORMALISATION = { + # AWS Bedrock variants → "AWS" + "bedrock": "AWS", + "bedrock_converse": "AWS", + "amazon": "AWS", + "aws": "AWS", + # Google variants → "google" (gemini, vertex, ...) + "gemini": "google", + "vertex_ai": "google", + "vertex": "google", + "google_genai": "google", + "google_generativeai": "google", +} + + +def _normalize_routed_provider(raw: str) -> str: + """Map LiteLLM's provider taxonomy to RETRY_LOOP.md §4.2's + routed-provider form. If no mapping applies, return the raw + value lower-cased (matches OpenAI / Anthropic which already + use the canonical form). + """ + lower = raw.lower() + if lower in _LITELLM_PROVIDER_NORMALISATION: + return _LITELLM_PROVIDER_NORMALISATION[lower] + # Bedrock model-prefix detection: LiteLLM model strings like + # "anthropic.claude-3-5-sonnet-20241022-v2:0" or "amazon.titan..." + # routed via Bedrock surface as ``custom_llm_provider="bedrock"``, + # but defensive pattern matching catches edge cases. + if lower.startswith(("amazon.", "anthropic.", "meta.", "ai21.", "cohere.", "mistral.")): + # These prefixes appear on Bedrock model IDs. + return "AWS" + return lower + + +def _httpx_status_code(response_obj) -> Optional[int]: + """Best-effort extract HTTP status code from a LiteLLM exception or response.""" + if response_obj is None: + return None + for attr in ("status_code", "http_status", "code"): + v = getattr(response_obj, attr, None) + if isinstance(v, int): + return v + response = getattr(response_obj, "response", None) + if response is not None: + return _httpx_status_code(response) + return None + + +def _server_address(kwargs) -> Optional[str]: + api_base = kwargs.get("api_base") + if not isinstance(api_base, str) or not api_base: + return None + # api_base is typically like "https://api.openai.com/v1" — strip scheme/path. + try: + from urllib.parse import urlparse + parsed = urlparse(api_base) + return parsed.hostname + except Exception: + return None + + +def _add_retry_attempt_prompt_attrs( + attrs: dict, + kwargs: dict, + *, + is_text_completion: bool = False, +) -> None: + """Copy request prompt content onto the retry_attempt span. + + Backend §4.5 makes retry_attempt the canonical LLMUsageEvent span + when it exists. Safety correlation and masking assertions therefore + need the same request content on the retry_attempt span that the + safety_wrapper parent carries. + """ + operation_is_text = is_text_completion or kwargs.get("text_completion") + if operation_is_text: + prompt = kwargs.get("prompt") + for index, text in enumerate(extract_prompt_texts(prompt)): + attrs[f"{SpanAttributes.LLM_PROMPTS}.{index}.role"] = "user" + attrs[f"{SpanAttributes.LLM_PROMPTS}.{index}.content"] = text + return + + messages = kwargs.get("messages") + if not isinstance(messages, list): + return + for index, message in enumerate(messages): + role = get_object_value(message, "role") + content = extract_text_content(get_object_value(message, "content")) + if role is not None: + attrs[f"{SpanAttributes.LLM_PROMPTS}.{index}.role"] = str(role) + if content: + attrs[f"{SpanAttributes.LLM_PROMPTS}.{index}.content"] = content + + +def _start_retry_attempt_span(kwargs, *, is_text_completion: bool = False) -> None: + """Open a retry_attempt sibling span under the current ambient FR + parent span and register it in the correlation map. Idempotent: + if a span already exists for this litellm_call_id, replace it + (defensive — shouldn't happen on the per-attempt observable + surfaces, but tolerated to handle edge cases without losing the + parent-marker side effect).""" + call_id = kwargs.get("litellm_call_id") if isinstance(kwargs, dict) else None + if not call_id: + # No correlation key → cannot match success/failure later. Skip. + return + + parent = trace.get_current_span() + if parent is None or not parent.get_span_context().is_valid: + # No ambient FR parent — likely the FR safety_wrapper context + # was never attached (e.g. user invoked LiteLLM's logger + # directly without going through the wrapped completion). + # Skip — parent-orphan retry_attempts have no meaningful place + # in the trace tree. + return + + routed_provider = _resolve_routed_provider(kwargs) or "litellm" + model = kwargs.get("model") + operation = "text_completion" if (is_text_completion or kwargs.get("text_completion")) else "chat" + attrs = { + GenAIAttributes.GEN_AI_SYSTEM: routed_provider, + GenAIAttributes.GEN_AI_OPERATION_NAME: operation, + _FR_SPAN_ROLE_KEY: _FR_SPAN_ROLE_RETRY_ATTEMPT, + } + if model is not None: + attrs[GenAIAttributes.GEN_AI_REQUEST_MODEL] = str(model) + server = _server_address(kwargs) + if server: + # Per RETRY_LOOP.md §4.2, attribute name is the OTel-standard + # ``server.address`` (network namespace) — using the literal + # key string for forward-compat across semconv lib changes. + attrs["server.address"] = server + _add_retry_attempt_prompt_attrs(attrs, kwargs, is_text_completion=is_text_completion) + + tracer = trace.get_tracer(__name__, __version__) + # Set the parent context explicitly so the retry_attempt span is a + # CHILD of the FR safety_wrapper, not a child of whatever happens + # to be ambient (which IS the safety_wrapper here, but explicit is + # safer for future refactoring). + parent_ctx = set_span_in_context(parent) + span = tracer.start_span( + _FR_RETRY_ATTEMPT_SPAN_NAME, + kind=SpanKind.CLIENT, + attributes=attrs, + context=parent_ctx, + ) + + # §4.7.1: register a framework-attempt token so direct-SDK + # wrappers (OpenAI/Anthropic/Bedrock) suppress their own emission + # while this attempt is in flight. + try: + framework_token = register_framework_attempt() + except Exception: + framework_token = None + logger.debug("Failed to register framework attempt token", exc_info=True) + + now = time.monotonic() + with _FR_RETRY_ATTEMPT_MAP_LOCK: + # Defensive: TTL + cap eviction on every insert path. + evicted = _evict_stale_retry_attempts_locked(now) + evicted += _enforce_retry_attempt_max_locked() + if evicted: + _maybe_warn_retry_attempt_eviction(evicted) + _FR_RETRY_ATTEMPT_MAP[call_id] = { + "span": span, + "parent": parent, + "started_at": now, + "framework_token": framework_token, + "ended": False, + } + + # §4.5 marker timing: set has_retry_attempt_child=true on the + # parent ONLY AFTER the first qualifying retry_attempt has + # successfully started. We just succeeded; mark the parent now. + # Idempotent: setting the attribute twice on the same parent is a + # no-op (OTel deduplicates). + try: + parent.set_attribute(_FR_HAS_RETRY_ATTEMPT_CHILD_KEY, True) + except Exception: + # Some span impls (e.g. a NonRecordingSpan during shutdown) + # may reject set_attribute. Non-fatal; the in-batch dedup + # path still fires from the in-batch retry_attempt children. + logger.debug("Failed to set has_retry_attempt_child on parent", exc_info=True) + + +def _finalize_retry_attempt_span( + kwargs, + response_obj, + *, + success: bool, +) -> None: + """End the retry_attempt span associated with this kwargs's + ``litellm_call_id``. Idempotent: a no-op if the entry is already + ended (handles the sync+async race where both + ``log_success_event`` and ``async_log_success_event`` may fire).""" + call_id = kwargs.get("litellm_call_id") if isinstance(kwargs, dict) else None + if not call_id: + return + with _FR_RETRY_ATTEMPT_MAP_LOCK: + entry = _FR_RETRY_ATTEMPT_MAP.get(call_id) + if entry is None or entry.get("ended"): + return + entry["ended"] = True + # Keep the entry around briefly — TTL sweep will clean it up, + # OR the success/failure callback's symmetric counterpart can + # short-circuit on the "ended" flag. + # Actually pop now: the "ended" sentinel is only useful within + # this critical section; outside it, popping is cleaner. + _FR_RETRY_ATTEMPT_MAP.pop(call_id, None) + span = entry["span"] + framework_token = entry.get("framework_token") + + try: + if success: + response_model = get_object_value(response_obj, "model") + if response_model: + span.set_attribute(GenAIAttributes.GEN_AI_RESPONSE_MODEL, str(response_model)) + response_id = get_object_value(response_obj, "id") + if response_id: + span.set_attribute(GenAIAttributes.GEN_AI_RESPONSE_ID, str(response_id)) + usage = get_object_value(response_obj, "usage") + input_tokens = get_object_value(usage, "prompt_tokens") + if input_tokens is None: + input_tokens = get_object_value(usage, "input_tokens") + output_tokens = get_object_value(usage, "completion_tokens") + if output_tokens is None: + output_tokens = get_object_value(usage, "output_tokens") + # Per §4.2 token-usage rule: SET when known, OMIT when unknown. + if input_tokens is not None: + span.set_attribute(GenAIAttributes.GEN_AI_USAGE_INPUT_TOKENS, int(input_tokens)) + if output_tokens is not None: + span.set_attribute(GenAIAttributes.GEN_AI_USAGE_OUTPUT_TOKENS, int(output_tokens)) + span.set_status(Status(StatusCode.OK)) + else: + exception = kwargs.get("exception") if isinstance(kwargs, dict) else None + status_code = _httpx_status_code(exception) or _httpx_status_code(response_obj) + if status_code is not None: + # Per RETRY_LOOP.md §4.2 the attribute name is the + # OTel-standard ``http.status_code`` (legacy semconv; + # backend extractor reads the literal key, not a + # python-binding constant). + span.set_attribute("http.status_code", int(status_code)) + if exception is not None: + error_type = type(exception).__name__ + span.set_attribute("error.type", error_type) + try: + span.record_exception(exception) + except Exception: + pass + span.set_status(Status(StatusCode.ERROR, str(exception))) + else: + span.set_status(Status(StatusCode.ERROR, "retry_attempt failed")) + finally: + try: + span.end() + except Exception: + logger.debug("Failed to end retry_attempt span", exc_info=True) + try: + unregister_framework_attempt(framework_token) + except Exception: + pass + + +class _FortifyRootRetryEmitter(_LiteLLMCustomLoggerBase): + """Second LiteLLM CustomLogger. Emits per-HTTP-attempt sibling + spans (``fortifyroot.litellm.retry_attempt``) under the FR + safety_wrapper parent. + + Registered AFTER ``_FortifyRootCompletionLogger`` in the LiteLLM + callbacks list so that completion-safety masking still runs + BEFORE retry_attempt finalization (the retry_attempt span captures + request prompt content at start, then model/tokens/status at end). + + MUST inherit from ``litellm.integrations.custom_logger.CustomLogger`` + because LiteLLM's dispatch loop gates every callback hook on + ``isinstance(callback, CustomLogger)``. A duck-typed class is + silently skipped — see _LiteLLMCustomLoggerBase docstring above. + """ + + def log_pre_api_call(self, model, messages, kwargs): + try: + _start_retry_attempt_span(kwargs) + except Exception: + logger.debug("retry-attempt span start failed", exc_info=True) + + async def async_log_pre_api_call(self, model, messages, kwargs): + # Defensive: LiteLLM's dispatch can fire either log_pre_api_call + # OR async_log_pre_api_call depending on the call path (sync vs + # async, callback registration timing). The _start helper is + # idempotent on (litellm_call_id) — defensive replace path + # tolerates duplicate fires — so making BOTH hooks the entry + # point avoids missed attempts on async-only paths. + # Review-round-2 Blocker 4 (2026-05-11): we previously only + # implemented the sync hook; some LiteLLM call paths (acompletion + # via certain routers) skipped the sync pre-call hook and our + # retry_attempt span never opened, leaving the failure callback + # with no entry to finalise. + try: + _start_retry_attempt_span(kwargs) + except Exception: + logger.debug("retry-attempt span async start failed", exc_info=True) + + def log_success_event(self, kwargs, response_obj, start_time, end_time): + try: + _finalize_retry_attempt_span(kwargs, response_obj, success=True) + except Exception: + logger.debug("retry-attempt span finalize-success failed", exc_info=True) + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + try: + _finalize_retry_attempt_span(kwargs, response_obj, success=True) + except Exception: + logger.debug("retry-attempt span async finalize-success failed", exc_info=True) + + def log_failure_event(self, kwargs, response_obj, start_time, end_time): + try: + _finalize_retry_attempt_span(kwargs, response_obj, success=False) + except Exception: + logger.debug("retry-attempt span finalize-failure failed", exc_info=True) + + async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): + try: + _finalize_retry_attempt_span(kwargs, response_obj, success=False) + except Exception: + logger.debug("retry-attempt span async finalize-failure failed", exc_info=True) + + class _FortifyRootCompletionLogger: """LiteLLM duck-typed CustomLogger that masks completions before native OTel fires. @@ -220,9 +737,19 @@ def _instrument(self, **kwargs): litellm.callbacks = [] self._fr_logger = _FortifyRootCompletionLogger() litellm.callbacks.insert(0, self._fr_logger) + # ST-10.1: register the retry-attempt emitter immediately + # AFTER the completion logger so completion-safety masking + # still runs first. The retry emitter captures metadata + # only (model, tokens, status), so its placement is not + # safety-critical. Insert at index 1 to keep ordering + # deterministic regardless of customer-registered + # callbacks. + self._fr_retry_emitter = _FortifyRootRetryEmitter() + litellm.callbacks.insert(1, self._fr_retry_emitter) except Exception: - logger.debug("Failed to register _FortifyRootCompletionLogger") + logger.debug("Failed to register FR LiteLLM callbacks") self._fr_logger = None + self._fr_retry_emitter = None for module_name, func_name, is_async, is_text_completion in _WRAPPED_METHODS: wrapper = ( @@ -233,7 +760,19 @@ def _instrument(self, **kwargs): wrap_function_wrapper(module_name, func_name, wrapper) def _uninstrument(self, **kwargs): - # Remove FR's logger from litellm.callbacks + # Remove FR's loggers from litellm.callbacks. Order: retry + # emitter first so it can no longer pin tokens, THEN the + # completion logger. + fr_retry_emitter = getattr(self, "_fr_retry_emitter", None) + if fr_retry_emitter is not None: + try: + import litellm + if isinstance(getattr(litellm, "callbacks", None), list): + litellm.callbacks.remove(fr_retry_emitter) + except Exception: + pass + self._fr_retry_emitter = None + fr_logger = getattr(self, "_fr_logger", None) if fr_logger is not None: try: diff --git a/packages/opentelemetry-instrumentation-litellm/tests/test_logger_integration.py b/packages/opentelemetry-instrumentation-litellm/tests/test_logger_integration.py index 355ff2d25c..4ae0281bbe 100644 --- a/packages/opentelemetry-instrumentation-litellm/tests/test_logger_integration.py +++ b/packages/opentelemetry-instrumentation-litellm/tests/test_logger_integration.py @@ -264,8 +264,11 @@ def test_logger_non_text_content_is_noop(): # --------------------------------------------------------------------------- def test_instrumentor_registers_logger_at_position_zero(): - """_FortifyRootCompletionLogger must be at index 0 after instrument().""" + """_FortifyRootCompletionLogger must be at index 0, and (post-ST-10.1) + _FortifyRootRetryEmitter must be at index 1, with any pre-existing + customer callbacks pushed to index 2+ after instrument().""" import litellm + from opentelemetry.instrumentation.litellm import _FortifyRootRetryEmitter original_callbacks = list(getattr(litellm, "callbacks", [])) try: litellm.callbacks = ["existing_cb"] @@ -274,10 +277,12 @@ def test_instrumentor_registers_logger_at_position_zero(): patch("opentelemetry.instrumentation.litellm.unwrap"): instrumentor._instrument() assert isinstance(litellm.callbacks[0], _FortifyRootCompletionLogger) - assert litellm.callbacks[1] == "existing_cb" + assert isinstance(litellm.callbacks[1], _FortifyRootRetryEmitter) + assert litellm.callbacks[2] == "existing_cb" instrumentor._uninstrument() - # After uninstrument, FR logger removed. + # After uninstrument, BOTH FR callbacks removed. assert not any(isinstance(cb, _FortifyRootCompletionLogger) for cb in litellm.callbacks) + assert not any(isinstance(cb, _FortifyRootRetryEmitter) for cb in litellm.callbacks) finally: litellm.callbacks = original_callbacks diff --git a/packages/opentelemetry-instrumentation-litellm/tests/test_retry_attempt_emission.py b/packages/opentelemetry-instrumentation-litellm/tests/test_retry_attempt_emission.py new file mode 100644 index 0000000000..9e13cf736e --- /dev/null +++ b/packages/opentelemetry-instrumentation-litellm/tests/test_retry_attempt_emission.py @@ -0,0 +1,565 @@ +"""Tests for ST-10.1 LiteLLM retry-aware emission. + +Covers: + - Instrumentor symmetry: _FortifyRootRetryEmitter is registered at + instrument() AND removed at uninstrument() (Fallback B child-emission + proof requirement (ii) per RETRY_LOOP.md §4.4.1). + - Single-attempt happy path: one retry_attempt span emitted under the + safety_wrapper parent, parent carries has_retry_attempt_child=true. + - Multi-attempt retry path: 3 retry_attempt spans emitted (2 ERROR + 1 + OK) for a 429→429→200 sequence, all under one safety_wrapper. + - Marker timing (§4.5): the parent's has_retry_attempt_child marker + is set AFTER the first retry_attempt's start, NOT at parent creation. + - §4.7.1 token registration: framework-attempt tokens are registered + on attempt-start AND unregistered on attempt-end. + - Idempotency: sync + async success callbacks both firing for one + attempt do not double-end the span. + - No-parent guard: invoking the emitter without an ambient FR parent + does not crash and does not emit an orphan retry_attempt span. +""" + +from __future__ import annotations + +from typing import Any + +import pytest +from opentelemetry import trace +from opentelemetry.instrumentation.fortifyroot import retry_registry +from opentelemetry.instrumentation.litellm import ( + LiteLLMInstrumentor, + _FortifyRootCompletionLogger, + _FortifyRootRetryEmitter, + _FR_HAS_RETRY_ATTEMPT_CHILD_KEY, + _FR_RETRY_ATTEMPT_MAP, + _FR_RETRY_ATTEMPT_SPAN_NAME, + _resolve_routed_provider, + _start_retry_attempt_span, + _finalize_retry_attempt_span, +) +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import ( + InMemorySpanExporter, +) + + +# --------------------------------------------------------------------------- +# Fixtures local to this file (intentionally NOT using conftest's +# session-scoped instrument fixture — these tests need fresh +# instrument/uninstrument cycles to validate symmetry). +# --------------------------------------------------------------------------- + +@pytest.fixture +def fresh_tracer(): + """A fresh TracerProvider + in-memory exporter per test, installed + as the GLOBAL tracer provider so the retry emitter's + ``trace.get_tracer(...)`` lookups route through it. After the + test, the global provider is left as-is — OTel's + set_tracer_provider only allows one set per process, so we don't + try to "restore" — but the in-memory exporter is cleared between + tests.""" + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + # Install as global so the retry emitter's tracer lookup uses it. + # Idempotent across tests because OTel allows only one set; second + # call is a no-op (the existing provider stays). To work around + # that for test isolation, we keep clearing the exporter between + # tests and reuse whatever provider was first set. + try: + trace.set_tracer_provider(provider) + except Exception: + # set_tracer_provider raises a warning (not an error) if a + # provider is already set; we tolerate it because the + # InMemorySpanExporter on the FIRST test's provider is what + # all subsequent tests will use, and that's fine since each + # test clears it. + pass + # Always re-fetch the global tracer so the test sees the same + # one the emitter sees. + tracer = trace.get_tracer("test") + # If a previous test set its own provider, fall back to using that + # test's exporter. We can't access it directly, so re-create our + # own and install via add_span_processor on the existing provider. + current_provider = trace.get_tracer_provider() + if current_provider is not provider: + # Existing provider — add OUR exporter to it so we still capture spans. + try: + current_provider.add_span_processor(SimpleSpanProcessor(exporter)) + except Exception: + pass + yield tracer, exporter, current_provider + + +@pytest.fixture(autouse=True) +def reset_registry_and_map(): + """Each test starts with empty retry-emitter state.""" + retry_registry._reset_for_test() + _FR_RETRY_ATTEMPT_MAP.clear() + yield + retry_registry._reset_for_test() + _FR_RETRY_ATTEMPT_MAP.clear() + + +# --------------------------------------------------------------------------- +# Instrumentor symmetry — Fallback B child-emission proof (ii). +# --------------------------------------------------------------------------- + +def test_retry_emitter_inherits_from_litellm_custom_logger(): + """REGRESSION GUARD (review-batch-1 end-to-end discovery 2026-05-10): + LiteLLM's dispatch loop gates every callback hook on + ``isinstance(callback, CustomLogger)`` — see + litellm_logging.py:1015 (log_pre_api_call) and :2303 + (log_success_event). A duck-typed _FortifyRootRetryEmitter is + silently SKIPPED by the dispatch — log_pre_api_call never fires + → no retry_attempt span ever emitted → §4.5 backend dedup has + nothing to dedup → the entire ST-10.1 contract is no-op'd. + + The bug was invisible to fork-side unit tests (which call the + emitter's hooks directly, bypassing LiteLLM's dispatch). End-to-end + Tier 1 with vendored fork surfaced it. + + This test catches the regression at unit-test time so a future + refactor that drops the inheritance gets a fast failure. + """ + from litellm.integrations.custom_logger import CustomLogger + emitter = _FortifyRootRetryEmitter() + assert isinstance(emitter, CustomLogger), ( + "_FortifyRootRetryEmitter MUST inherit from " + "litellm.integrations.custom_logger.CustomLogger so LiteLLM's " + "isinstance-gated callback dispatch fires its hooks. " + "Without inheritance, the emitter is silently skipped and no " + "retry_attempt spans are emitted." + ) + + +def test_instrumentor_registers_retry_emitter_at_instrument(): + """At _instrument() time, exactly one _FortifyRootRetryEmitter + is present in litellm.callbacks.""" + import litellm + + instrumentor = LiteLLMInstrumentor() + try: + instrumentor.instrument() + assert any( + isinstance(cb, _FortifyRootRetryEmitter) for cb in litellm.callbacks + ), "_FortifyRootRetryEmitter must be registered at instrument()" + emitters = [ + cb for cb in litellm.callbacks if isinstance(cb, _FortifyRootRetryEmitter) + ] + assert len(emitters) == 1, f"expected exactly 1 emitter, found {len(emitters)}" + finally: + instrumentor.uninstrument() + + +def test_instrumentor_removes_retry_emitter_at_uninstrument(): + """At _uninstrument(), the _FortifyRootRetryEmitter is removed + from litellm.callbacks.""" + import litellm + + instrumentor = LiteLLMInstrumentor() + instrumentor.instrument() + assert any(isinstance(cb, _FortifyRootRetryEmitter) for cb in litellm.callbacks) + instrumentor.uninstrument() + assert not any( + isinstance(cb, _FortifyRootRetryEmitter) for cb in litellm.callbacks + ), "_FortifyRootRetryEmitter must be unregistered at uninstrument()" + + +def test_completion_logger_fires_before_retry_emitter(): + """Order in litellm.callbacks: _FortifyRootCompletionLogger BEFORE + _FortifyRootRetryEmitter. Completion-safety masking must run + before retry_attempt finalization (which only captures metadata).""" + import litellm + + instrumentor = LiteLLMInstrumentor() + try: + instrumentor.instrument() + types = [type(cb).__name__ for cb in litellm.callbacks] + completion_idx = next( + (i for i, cb in enumerate(litellm.callbacks) + if isinstance(cb, _FortifyRootCompletionLogger)), + None, + ) + retry_idx = next( + (i for i, cb in enumerate(litellm.callbacks) + if isinstance(cb, _FortifyRootRetryEmitter)), + None, + ) + assert completion_idx is not None, f"no completion logger found; callbacks={types}" + assert retry_idx is not None, f"no retry emitter found; callbacks={types}" + assert completion_idx < retry_idx, ( + f"completion logger must precede retry emitter; got " + f"completion={completion_idx} retry={retry_idx} callbacks={types}" + ) + finally: + instrumentor.uninstrument() + + +# --------------------------------------------------------------------------- +# Single-attempt happy path. +# --------------------------------------------------------------------------- + +def test_single_attempt_emits_one_retry_attempt_under_parent(fresh_tracer): + """Drive _start_retry_attempt_span + _finalize_retry_attempt_span + directly with an ambient parent span. Verify: + - exactly 1 retry_attempt span exported + - retry_attempt's parent is the safety_wrapper + - parent has has_retry_attempt_child=true + - retry_attempt has fortifyroot.span.role=retry_attempt + - retry_attempt has gen_ai.system, gen_ai.request.model + """ + tracer, exporter, _ = fresh_tracer + + parent = tracer.start_span("fortifyroot.litellm.safety") + with trace.use_span(parent, end_on_exit=False): + kwargs: dict[str, Any] = { + "litellm_call_id": "test-call-001", + "model": "openai/gpt-4o-mini", + "api_base": "https://api.openai.com/v1", + "custom_llm_provider": "openai", + "messages": [{"role": "user", "content": "hello masked [EMAIL]"}], + } + _start_retry_attempt_span(kwargs) + + # Mock response object with usage/model. + class MockUsage: + prompt_tokens = 10 + completion_tokens = 5 + + class MockResponse: + id = "resp-abc" + model = "gpt-4o-mini" + usage = MockUsage() + + _finalize_retry_attempt_span(kwargs, MockResponse(), success=True) + parent.end() + + spans = exporter.get_finished_spans() + span_names = [s.name for s in spans] + assert _FR_RETRY_ATTEMPT_SPAN_NAME in span_names, ( + f"retry_attempt span missing; got {span_names}" + ) + retry_span = next(s for s in spans if s.name == _FR_RETRY_ATTEMPT_SPAN_NAME) + parent_span_exported = next(s for s in spans if s.name == "fortifyroot.litellm.safety") + + assert retry_span.parent.span_id == parent_span_exported.context.span_id, ( + "retry_attempt MUST be a child of the safety_wrapper parent" + ) + assert retry_span.attributes.get("fortifyroot.span.role") == "retry_attempt" + assert retry_span.attributes.get("gen_ai.system") == "openai", ( + "gen_ai.system MUST be the routed provider, NOT the framework" + ) + assert retry_span.attributes.get("gen_ai.request.model") == "openai/gpt-4o-mini" + assert retry_span.attributes.get("gen_ai.response.model") == "gpt-4o-mini" + assert retry_span.attributes.get("gen_ai.response.id") == "resp-abc" + assert retry_span.attributes.get("gen_ai.usage.input_tokens") == 10 + assert retry_span.attributes.get("gen_ai.usage.output_tokens") == 5 + assert retry_span.attributes.get("gen_ai.prompt.0.role") == "user" + assert retry_span.attributes.get("gen_ai.prompt.0.content") == "hello masked [EMAIL]" + + # §4.5 marker: must be set on parent. + assert parent_span_exported.attributes.get( + _FR_HAS_RETRY_ATTEMPT_CHILD_KEY + ) is True, "parent MUST carry has_retry_attempt_child=true" + + +def test_retry_attempt_accepts_anthropic_usage_token_names(fresh_tracer): + """Anthropic-shaped LiteLLM responses may expose input/output token + names instead of OpenAI-style prompt/completion names.""" + tracer, exporter, _ = fresh_tracer + + parent = tracer.start_span("fortifyroot.litellm.safety") + with trace.use_span(parent, end_on_exit=False): + kwargs: dict[str, Any] = { + "litellm_call_id": "anthropic-usage-001", + "model": "anthropic/claude-4-sonnet-20250514", + "custom_llm_provider": "anthropic", + "messages": [{"role": "user", "content": "hello"}], + } + _start_retry_attempt_span(kwargs) + + class MockUsage: + input_tokens = 11 + output_tokens = 7 + + class MockResponse: + model = "claude-sonnet-4-20250514" + usage = MockUsage() + + _finalize_retry_attempt_span(kwargs, MockResponse(), success=True) + parent.end() + + retry_span = next( + s for s in exporter.get_finished_spans() + if s.name == _FR_RETRY_ATTEMPT_SPAN_NAME + ) + assert retry_span.attributes.get("gen_ai.usage.input_tokens") == 11 + assert retry_span.attributes.get("gen_ai.usage.output_tokens") == 7 + + +# --------------------------------------------------------------------------- +# Multi-attempt retry path. +# --------------------------------------------------------------------------- + +def test_three_attempt_retry_path_emits_three_retry_attempt_spans(fresh_tracer): + """Three sequential attempts (429 → 429 → 200) under one parent. + Each gets its own litellm_call_id. Expect exactly 3 + retry_attempt children, with IsError flags (2 ERROR + 1 OK).""" + tracer, exporter, _ = fresh_tracer + + parent = tracer.start_span("fortifyroot.litellm.safety") + + class MockHTTPError(Exception): + def __init__(self, status_code: int): + self.status_code = status_code + + with trace.use_span(parent, end_on_exit=False): + # Attempt 1 — 429 + k1 = {"litellm_call_id": "call-001", "model": "openai/gpt-4o-mini"} + _start_retry_attempt_span(k1) + k1["exception"] = MockHTTPError(429) + _finalize_retry_attempt_span(k1, None, success=False) + + # Attempt 2 — 429 + k2 = {"litellm_call_id": "call-002", "model": "openai/gpt-4o-mini"} + _start_retry_attempt_span(k2) + k2["exception"] = MockHTTPError(429) + _finalize_retry_attempt_span(k2, None, success=False) + + # Attempt 3 — 200 + k3 = {"litellm_call_id": "call-003", "model": "openai/gpt-4o-mini"} + _start_retry_attempt_span(k3) + + class MockSuccess: + id = "ok-id" + model = "gpt-4o-mini" + usage = type("U", (), {"prompt_tokens": 10, "completion_tokens": 5})() + + _finalize_retry_attempt_span(k3, MockSuccess(), success=True) + parent.end() + + retry_spans = [ + s for s in exporter.get_finished_spans() + if s.name == _FR_RETRY_ATTEMPT_SPAN_NAME + ] + assert len(retry_spans) == 3, f"expected 3 retry_attempt spans, got {len(retry_spans)}" + + # 2 of them must be ERROR status, 1 OK. + from opentelemetry.trace import StatusCode + error_spans = [s for s in retry_spans if s.status.status_code == StatusCode.ERROR] + ok_spans = [s for s in retry_spans if s.status.status_code == StatusCode.OK] + assert len(error_spans) == 2, f"expected 2 ERROR spans, got {len(error_spans)}" + assert len(ok_spans) == 1, f"expected 1 OK span, got {len(ok_spans)}" + + # The 2 error spans carry http.status_code=429 + error.type. + for s in error_spans: + assert s.attributes.get("http.status_code") == 429 + assert s.attributes.get("error.type") == "MockHTTPError" + + +# --------------------------------------------------------------------------- +# Marker timing — §4.5 "marker timing" boundary case. +# --------------------------------------------------------------------------- + +def test_marker_set_AFTER_first_attempt_starts_not_at_parent_creation(fresh_tracer): + """Per §4.5 marker-timing paragraph: the marker is set AFTER the + first retry_attempt successfully starts, NOT at parent creation. + + This test verifies: a parent span observed BEFORE any + _start_retry_attempt_span call has NO marker. After the first + _start_retry_attempt_span call, the parent has the marker. + """ + tracer, exporter, _ = fresh_tracer + + parent = tracer.start_span("fortifyroot.litellm.safety") + + # Before any retry_attempt: no marker. + # (Use the live span object since it hasn't been exported yet.) + parent_attrs_before = dict(parent.attributes or {}) + assert _FR_HAS_RETRY_ATTEMPT_CHILD_KEY not in parent_attrs_before, ( + "parent MUST NOT have the marker before any retry_attempt starts" + ) + + with trace.use_span(parent, end_on_exit=False): + kwargs = {"litellm_call_id": "marker-test", "model": "openai/gpt-4o-mini"} + _start_retry_attempt_span(kwargs) + # After: parent has the marker. + parent_attrs_after = dict(parent.attributes or {}) + assert parent_attrs_after.get(_FR_HAS_RETRY_ATTEMPT_CHILD_KEY) is True, ( + "parent MUST have the marker after first retry_attempt starts" + ) + + # Cleanup. + with trace.use_span(parent, end_on_exit=False): + _finalize_retry_attempt_span(kwargs, None, success=False) + parent.end() + + +def test_marker_NOT_set_when_no_retry_attempt_starts(fresh_tracer): + """If _start_retry_attempt_span is never called for a parent, the + parent NEVER gets the marker. This is the "framework returned + without invoking the retry-emitter callback at all" scenario from + the §4.5 marker-timing paragraph — telemetry-loss is avoided + because the §4.5 backend dedup degrades gracefully (parent stays + canonical when the marker is unset).""" + tracer, exporter, _ = fresh_tracer + + parent = tracer.start_span("fortifyroot.litellm.safety") + parent.end() + + spans = exporter.get_finished_spans() + parent_exported = next(s for s in spans if s.name == "fortifyroot.litellm.safety") + assert _FR_HAS_RETRY_ATTEMPT_CHILD_KEY not in (parent_exported.attributes or {}), ( + "no retry_attempt → no marker → parent stays canonical for §4.5 dedup" + ) + + +# --------------------------------------------------------------------------- +# §4.7.1 token registration. +# --------------------------------------------------------------------------- + +def test_framework_token_registered_during_attempt_unregistered_after(fresh_tracer): + """While a retry_attempt span is open, is_framework_owned() is + True for the current thread. After the success/failure callback + fires, it returns to False.""" + tracer, _, _ = fresh_tracer + from opentelemetry.instrumentation.fortifyroot import is_framework_owned + + parent = tracer.start_span("fortifyroot.litellm.safety") + with trace.use_span(parent, end_on_exit=False): + assert not is_framework_owned(), "no attempt yet → not owned" + + kwargs = {"litellm_call_id": "tok-test", "model": "openai/gpt-4o-mini"} + _start_retry_attempt_span(kwargs) + + assert is_framework_owned(), ( + "during attempt → framework owns the call → direct-SDK wrappers " + "would suppress emission" + ) + + _finalize_retry_attempt_span(kwargs, None, success=False) + + assert not is_framework_owned(), ( + "after attempt → token unregistered → direct-SDK wrappers " + "may emit again" + ) + parent.end() + + +# --------------------------------------------------------------------------- +# Idempotency. +# --------------------------------------------------------------------------- + +def test_double_finalize_does_not_double_end(fresh_tracer): + """LiteLLM's sync log_success_event AND async_log_success_event + can both fire for the same attempt (the worker may run late). The + second finalize MUST be a no-op.""" + tracer, exporter, _ = fresh_tracer + + parent = tracer.start_span("fortifyroot.litellm.safety") + with trace.use_span(parent, end_on_exit=False): + kwargs = {"litellm_call_id": "idem-001", "model": "openai/gpt-4o-mini"} + _start_retry_attempt_span(kwargs) + + class MockSuccess: + id = "x" + model = "gpt-4o-mini" + usage = type("U", (), {"prompt_tokens": 1, "completion_tokens": 1})() + + # First finalize. + _finalize_retry_attempt_span(kwargs, MockSuccess(), success=True) + # Second finalize — must be a no-op. + _finalize_retry_attempt_span(kwargs, MockSuccess(), success=True) + parent.end() + + retry_spans = [ + s for s in exporter.get_finished_spans() + if s.name == _FR_RETRY_ATTEMPT_SPAN_NAME + ] + assert len(retry_spans) == 1, ( + f"expected exactly 1 retry_attempt span (idempotent finalize); " + f"got {len(retry_spans)}" + ) + + +# --------------------------------------------------------------------------- +# No-parent guard. +# --------------------------------------------------------------------------- + +def test_no_parent_span_does_not_emit_orphan_retry_attempt(fresh_tracer): + """If _start_retry_attempt_span fires without an ambient FR + parent (e.g. user invoked LiteLLM's logger directly), the + emitter MUST NOT create an orphan retry_attempt span — they + have no meaningful place in the trace tree, and the §4.5 + backend dedup expects retry_attempts to have a parent.""" + tracer, exporter, _ = fresh_tracer + # Note: NO parent span attached. + + kwargs = {"litellm_call_id": "no-parent-test", "model": "openai/gpt-4o-mini"} + _start_retry_attempt_span(kwargs) + + spans = exporter.get_finished_spans() + retry_spans = [s for s in spans if s.name == _FR_RETRY_ATTEMPT_SPAN_NAME] + assert len(retry_spans) == 0, ( + f"orphan retry_attempt MUST NOT be emitted; got {len(retry_spans)}" + ) + + # Map should also be empty (no entry registered for orphan calls). + assert "no-parent-test" not in _FR_RETRY_ATTEMPT_MAP + + +def test_resolve_routed_provider_normalisation(): + """REGRESSION GUARD (review-batch-1 Minor 4 fix 2026-05-10): + LiteLLM's ``custom_llm_provider`` taxonomy uses values like + ``"bedrock"``, ``"bedrock_converse"``, ``"vertex_ai"`` that + diverge from RETRY_LOOP.md §4.2's canonical routed-provider + form (``"AWS"``, ``"google"``, etc.). Cross-wrapper drift + here would mean LiteLLM-routed-Bedrock spans carry + ``gen_ai.system="bedrock"`` while LangChain/LlamaIndex + Bedrock spans carry ``gen_ai.system="AWS"``, breaking the + ST-10.0 §4.2 cross-wrapper consistency contract. + """ + # AWS Bedrock — all variants must normalise to "AWS". + assert _resolve_routed_provider({"custom_llm_provider": "bedrock"}) == "AWS" + assert _resolve_routed_provider({"custom_llm_provider": "bedrock_converse"}) == "AWS" + assert _resolve_routed_provider({"custom_llm_provider": "aws"}) == "AWS" + # Bedrock model-prefix detection (LiteLLM may pass the + # provider via the model string only). + assert _resolve_routed_provider({"model": "bedrock/anthropic.claude-3-5-sonnet"}) == "AWS" + assert _resolve_routed_provider({"model": "amazon.nova-lite-v1:0"}) == "AWS" + + # Google Gemini / Vertex variants → "google". + assert _resolve_routed_provider({"custom_llm_provider": "gemini"}) == "google" + assert _resolve_routed_provider({"custom_llm_provider": "vertex_ai"}) == "google" + assert _resolve_routed_provider({"custom_llm_provider": "google_genai"}) == "google" + + # Already-canonical values pass through (lower-cased). + assert _resolve_routed_provider({"custom_llm_provider": "openai"}) == "openai" + assert _resolve_routed_provider({"custom_llm_provider": "anthropic"}) == "anthropic" + assert _resolve_routed_provider({"model": "claude-4-sonnet-20250514"}) == "anthropic" + assert _resolve_routed_provider({"model": "claude-sonnet-4-20250514"}) == "anthropic" + + # Empty / undeterminable → None. + assert _resolve_routed_provider({}) is None + assert _resolve_routed_provider({"custom_llm_provider": ""}) is None + + +def test_no_litellm_call_id_does_not_register_map_entry(fresh_tracer): + """If kwargs lacks litellm_call_id, the emitter has no + correlation key for success/failure callbacks → silently skips + span emission rather than emitting an unmatched span.""" + tracer, exporter, _ = fresh_tracer + + parent = tracer.start_span("fortifyroot.litellm.safety") + with trace.use_span(parent, end_on_exit=False): + kwargs = {"model": "openai/gpt-4o-mini"} # no litellm_call_id + _start_retry_attempt_span(kwargs) + parent.end() + + retry_spans = [ + s for s in exporter.get_finished_spans() + if s.name == _FR_RETRY_ATTEMPT_SPAN_NAME + ] + assert len(retry_spans) == 0, "missing litellm_call_id → no emission" + assert len(_FR_RETRY_ATTEMPT_MAP) == 0 diff --git a/packages/opentelemetry-instrumentation-litellm/uv.lock b/packages/opentelemetry-instrumentation-litellm/uv.lock index 6105c6dc71..a8b244c008 100644 --- a/packages/opentelemetry-instrumentation-litellm/uv.lock +++ b/packages/opentelemetry-instrumentation-litellm/uv.lock @@ -1222,6 +1222,7 @@ dev = [ test = [ { name = "opentelemetry-sdk", specifier = ">=1.38.0,<2" }, { name = "pytest", specifier = ">=8.2.2,<9" }, + { name = "pytest-asyncio", specifier = ">=0.23.7,<0.24.0" }, ] [[package]] @@ -1253,12 +1254,14 @@ test = [ { name = "opentelemetry-sdk" }, { name = "pytest" }, { name = "pytest-asyncio" }, + { name = "pytest-recording" }, { name = "pytest-sugar" }, + { name = "vcrpy" }, ] [package.metadata] requires-dist = [ - { name = "litellm", marker = "extra == 'instruments'", specifier = ">=1.71.2,<2" }, + { name = "litellm", marker = "extra == 'instruments'", specifier = ">=1.71.2,!=1.82.7,!=1.82.8,<2" }, { name = "opentelemetry-api", specifier = ">=1.38.0,<2" }, { name = "opentelemetry-instrumentation", specifier = ">=0.59b0" }, { name = "opentelemetry-instrumentation-fortifyroot", editable = "../opentelemetry-instrumentation-fortifyroot" }, @@ -1275,11 +1278,13 @@ dev = [ { name = "ruff", specifier = ">=0.4.0" }, ] test = [ - { name = "litellm", specifier = ">=1.71.2,<2" }, + { name = "litellm", specifier = ">=1.71.2,!=1.82.7,!=1.82.8,<2" }, { name = "opentelemetry-sdk", specifier = ">=1.38.0,<2" }, { name = "pytest", specifier = ">=8.2.2,<9" }, { name = "pytest-asyncio", specifier = ">=0.23.7,<0.24.0" }, + { name = "pytest-recording", specifier = ">=0.13.2,<0.14.0" }, { name = "pytest-sugar", specifier = "==1.0.0" }, + { name = "vcrpy", specifier = ">=8.0.0,<9" }, ] [[package]] @@ -1653,6 +1658,19 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/ee/82/62e2d63639ecb0fbe8a7ee59ef0bc69a4669ec50f6d3459f74ad4e4189a2/pytest_asyncio-0.23.8-py3-none-any.whl", hash = "sha256:50265d892689a5faefb84df80819d1ecef566eb3549cf915dfb33569359d1ce2", size = 17663, upload-time = "2024-07-17T17:39:32.478Z" }, ] +[[package]] +name = "pytest-recording" +version = "0.13.4" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "pytest" }, + { name = "vcrpy" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/32/9c/f4027c5f1693847b06d11caf4b4f6bb09f22c1581ada4663877ec166b8c6/pytest_recording-0.13.4.tar.gz", hash = "sha256:568d64b2a85992eec4ae0a419c855d5fd96782c5fb016784d86f18053792768c", size = 26576, upload-time = "2025-05-08T10:41:11.231Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/42/c2/ce34735972cc42d912173e79f200fe66530225190c06655c5632a9d88f1e/pytest_recording-0.13.4-py3-none-any.whl", hash = "sha256:ad49a434b51b1c4f78e85b1e6b74fdcc2a0a581ca16e52c798c6ace971f7f439", size = 13723, upload-time = "2025-05-08T10:41:09.684Z" }, +] + [[package]] name = "pytest-sugar" version = "1.0.0" @@ -2318,6 +2336,19 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/39/08/aaaad47bc4e9dc8c725e68f9d04865dbcb2052843ff09c97b08904852d84/urllib3-2.6.3-py3-none-any.whl", hash = "sha256:bf272323e553dfb2e87d9bfd225ca7b0f467b919d7bbd355436d3fd37cb0acd4", size = 131584, upload-time = "2026-01-07T16:24:42.685Z" }, ] +[[package]] +name = "vcrpy" +version = "8.1.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "pyyaml" }, + { name = "wrapt" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/b3/07/bcfd5ebd7cb308026ab78a353e091bd699593358be49197d39d004e5ad83/vcrpy-8.1.1.tar.gz", hash = "sha256:58e3053e33b423f3594031cb758c3f4d1df931307f1e67928e30cf352df7709f", size = 85770, upload-time = "2026-01-04T19:22:03.886Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/3a/d7/f79b05a5d728f8786876a7d75dfb0c5cae27e428081b2d60152fb52f155f/vcrpy-8.1.1-py3-none-any.whl", hash = "sha256:2d16f31ad56493efb6165182dd99767207031b0da3f68b18f975545ede8ac4b9", size = 42445, upload-time = "2026-01-04T19:22:02.532Z" }, +] + [[package]] name = "wrapt" version = "1.17.3" diff --git a/packages/opentelemetry-instrumentation-llamaindex/opentelemetry/instrumentation/llamaindex/dispatcher_wrapper.py b/packages/opentelemetry-instrumentation-llamaindex/opentelemetry/instrumentation/llamaindex/dispatcher_wrapper.py index 1daadbc76d..e21bd81c66 100644 --- a/packages/opentelemetry-instrumentation-llamaindex/opentelemetry/instrumentation/llamaindex/dispatcher_wrapper.py +++ b/packages/opentelemetry-instrumentation-llamaindex/opentelemetry/instrumentation/llamaindex/dispatcher_wrapper.py @@ -152,6 +152,21 @@ def _stamp_llm_response_model_for_safety(event: BaseEvent, span) -> None: def instrument_with_dispatcher(tracer: Tracer): instrument_llm_safety_wrappers() dispatcher = get_dispatcher() + # ST-10.3: register the FR retry-attempt handler FIRST, before + # OpenLLMetrySpanHandler. Order matters because span handlers + # fire in registration order — and our handler reads the + # ambient OTel context to decide the parent of the + # retry_attempt span. If OpenLLMetrySpanHandler ran first, it + # would have already swapped the OTel ambient context to its + # per-call SpanHolder span, meaning each retry attempt would + # land under a DIFFERENT parent → siblings invariant broken → + # RetryDetectorProc grouping fails. Running FR's handler first + # preserves the user's enclosing OTel span as the shared parent + # for all retry_attempts under one logical retry loop. + from opentelemetry.instrumentation.llamaindex.retry_handler import ( + _FortifyRootRetryHandler, + ) + dispatcher.add_span_handler(_FortifyRootRetryHandler()) openllmetry_span_handler = OpenLLMetrySpanHandler(tracer) dispatcher.add_span_handler(openllmetry_span_handler) dispatcher.add_event_handler(OpenLLMetryEventHandler(openllmetry_span_handler)) diff --git a/packages/opentelemetry-instrumentation-llamaindex/opentelemetry/instrumentation/llamaindex/retry_handler.py b/packages/opentelemetry-instrumentation-llamaindex/opentelemetry/instrumentation/llamaindex/retry_handler.py new file mode 100644 index 0000000000..9048b2154c --- /dev/null +++ b/packages/opentelemetry-instrumentation-llamaindex/opentelemetry/instrumentation/llamaindex/retry_handler.py @@ -0,0 +1,488 @@ +"""ST-10.3 retry-aware emission for LlamaIndex. + +Per RETRY_LOOP.md §4.4 LlamaIndex row + §4.4.2 coverage limitation + +ST-10.0 C3 empirical verification (POC findings F4-F6 in +phase_st10_retryloop.txt): + + - LlamaIndex's ``OpenAI.chat()`` (and equivalents) fires the + dispatcher span TWICE per HTTP attempt — once on the public + ``chat()`` method and once on the inner ``_chat()`` (F4). + A naive event-handler approach emits 2x retry_attempt children + per attempt. ST-10.3 hooks the SPAN HANDLER (not the event + handler) and filters to OUTER public methods only via a + method-name whitelist + ``BaseLLM`` instance check. + - Per-attempt firing is verified empirically on framework-layer + retry paths (e.g. tenacity wrapping at the application layer). + Provider-SDK-internal retries (e.g. + ``OpenAI(max_retries=N).chat(...)``) fire dispatcher spans ONCE + per logical call → §4.4.2 coverage limitation applies (same + pattern as LiteLLM C1 / LangChain C2). + - ``span_enter``/``span_exit``/``span_drop`` lifecycle hooks + correlate cleanly: each enter has a matching exit OR drop, so + we don't need separate event-side de-dup logic. + - Use ambient OTel span at first attempt as the parent for + sibling-grouping (RetryDetectorProc requires retry_attempts to + be siblings under one parent_span_id). + +This handler is registered alongside the existing +``OpenLLMetrySpanHandler`` — they capture orthogonal data +(OpenLLMetry: full per-call telemetry; FR retry-handler: per-attempt +sibling spans for retry detection only). +""" + +from __future__ import annotations + +import logging +import re +import threading +import time +from typing import Any, Optional + +from llama_index.core.base.llms.base import BaseLLM +from llama_index.core.instrumentation.span_handlers.base import BaseSpanHandler +from opentelemetry import trace +from opentelemetry.instrumentation.fortifyroot import ( + register_framework_attempt, + unregister_framework_attempt, +) +from opentelemetry.instrumentation.llamaindex.version import __version__ +from opentelemetry.trace import SpanKind, Status, StatusCode, set_span_in_context + +logger = logging.getLogger(__name__) + +# ST-10 §4.4: per-attempt sibling span name + role. +_FR_RETRY_ATTEMPT_SPAN_NAME = "fortifyroot.llamaindex.retry_attempt" +_FR_SPAN_ROLE_KEY = "fortifyroot.span.role" +_FR_SPAN_ROLE_RETRY_ATTEMPT = "retry_attempt" + +# ST-10 §4.5 parent marker. +_FR_HAS_RETRY_ATTEMPT_CHILD_KEY = "fortifyroot.span.has_retry_attempt_child" + +# Dispatcher span IDs follow the pattern "ClassName.method-uuid". +# Capture the method name so we can filter outer (public) calls +# from inner (underscore-prefixed) ones — F4 finding from POC. +_DISPATCHER_ID_RE = re.compile(r"^([A-Za-z][A-Za-z0-9_]*)\.([A-Za-z_][A-Za-z0-9_]*)-") + +# Whitelist of OUTER public LlamaIndex BaseLLM methods we want to +# track. The inner private methods (``_chat``, ``_complete``, etc.) +# do NOT appear here, so they're filtered out — the F4 de-dup. +# Sourced from llama_index.core.llms.llm.LLM public surface: +# chat / achat / stream_chat / astream_chat +# complete / acomplete / stream_complete / astream_complete +# predict / apredict +# structured_predict / astructured_predict / +# stream_structured_predict / astream_structured_predict +_OUTER_LLM_METHODS = { + "chat", "achat", "stream_chat", "astream_chat", + "complete", "acomplete", "stream_complete", "astream_complete", + "predict", "apredict", + "structured_predict", "astructured_predict", + "stream_structured_predict", "astream_structured_predict", +} + + +# ---------------------------------------------------------------------- +# Per-attempt correlation map. Key = dispatcher span id_ (the OUTER +# span's id_, post-filter). Value = {span, started_at, framework_token, +# ended}. Bounded-size + TTL eviction defends against framework +# crashes. +# ---------------------------------------------------------------------- + +_FR_RETRY_ATTEMPT_MAP: dict[str, dict] = {} +_FR_RETRY_ATTEMPT_MAP_LOCK = threading.Lock() +_FR_RETRY_ATTEMPT_MAP_MAX = 4096 +_FR_RETRY_ATTEMPT_MAP_TTL_SEC = 60.0 +_FR_RETRY_ATTEMPT_EVICT_BATCH = 1024 +_FR_RETRY_ATTEMPT_EVICT_WARN_EVERY = 32 +_FR_RETRY_ATTEMPT_EVICT_WARN_COUNTER = 0 + + +def _evict_stale_attempts_locked(now: float) -> int: + cutoff = now - _FR_RETRY_ATTEMPT_MAP_TTL_SEC + stale = [k for k, v in _FR_RETRY_ATTEMPT_MAP.items() if v["started_at"] < cutoff] + for k in stale: + entry = _FR_RETRY_ATTEMPT_MAP.pop(k, None) + if entry is None: + continue + try: + sp = entry.get("span") + if sp is not None and not entry.get("ended"): + sp.set_status(Status(StatusCode.ERROR, "retry_attempt orphaned (framework crashed)")) + sp.end() + except Exception: + pass + try: + unregister_framework_attempt(entry.get("framework_token")) + except Exception: + pass + return len(stale) + + +def _enforce_attempt_max_locked() -> int: + if len(_FR_RETRY_ATTEMPT_MAP) <= _FR_RETRY_ATTEMPT_MAP_MAX: + return 0 + items = sorted(_FR_RETRY_ATTEMPT_MAP.items(), key=lambda kv: kv[1]["started_at"]) + to_drop = items[:_FR_RETRY_ATTEMPT_EVICT_BATCH] + for k, entry in to_drop: + _FR_RETRY_ATTEMPT_MAP.pop(k, None) + try: + sp = entry.get("span") + if sp is not None and not entry.get("ended"): + sp.set_status(Status(StatusCode.ERROR, "retry_attempt cap-evicted")) + sp.end() + except Exception: + pass + try: + unregister_framework_attempt(entry.get("framework_token")) + except Exception: + pass + return len(to_drop) + + +def _maybe_warn_eviction(evicted: int) -> None: + global _FR_RETRY_ATTEMPT_EVICT_WARN_COUNTER + _FR_RETRY_ATTEMPT_EVICT_WARN_COUNTER += evicted + if _FR_RETRY_ATTEMPT_EVICT_WARN_COUNTER < _FR_RETRY_ATTEMPT_EVICT_WARN_EVERY: + return + _FR_RETRY_ATTEMPT_EVICT_WARN_COUNTER = 0 + logger.warning( + "fortifyroot llamaindex retry_attempt map: evicted %d+ stale/over-cap entries; " + "framework may be leaking attempts (TTL=%.0fs, max=%d)", + evicted, + _FR_RETRY_ATTEMPT_MAP_TTL_SEC, + _FR_RETRY_ATTEMPT_MAP_MAX, + ) + + +def _is_outer_llm_method(id_: str, instance: Any) -> bool: + """F4 de-dup filter: only emit retry_attempt for OUTER public + LlamaIndex LLM methods. Returns False for: + - non-LLM dispatcher spans (e.g. workflow / chain spans) + - inner private methods (``_chat``, ``_complete``, etc.) + - non-BaseLLM instances + """ + if not isinstance(instance, BaseLLM): + return False + m = _DISPATCHER_ID_RE.match(id_) + if m is None: + return False + method = m.group(2) + return method in _OUTER_LLM_METHODS + + +def _resolve_routed_provider(instance: Any) -> Optional[str]: + """Best-effort: derive the routed provider for gen_ai.system from + the instance's class module path.""" + cls = type(instance) + module_name = (cls.__module__ or "").lower() + # Common shapes: ``llama_index.llms.openai.base``, + # ``llama_index_llms_openai`` (newer), etc. + if "openai" in module_name: + return "openai" + if "anthropic" in module_name: + return "anthropic" + if "bedrock" in module_name or "aws" in module_name: + return "AWS" + if "google" in module_name or "gemini" in module_name or "vertex" in module_name: + return "google" + return None + + +def _resolve_model(instance: Any) -> Optional[str]: + for attr in ("model", "model_name", "deployment_name"): + v = getattr(instance, attr, None) + if isinstance(v, str) and v: + return v + return None + + +def _content_to_string(content: Any) -> str: + if isinstance(content, str): + return content + try: + import json + return json.dumps(content, default=str) + except Exception: + return str(content) + + +def _add_prompt_attrs(attrs: dict[str, Any], bound_args: Any) -> None: + """Copy request prompt content onto the retry_attempt span. + + Backend §4.5 makes retry_attempt the canonical LLMUsageEvent span + when it exists, so safety correlation still needs prompt content on + this span. LlamaIndex safety wrappers have already processed the + bound arguments by the time dispatcher span handlers see them. + """ + arguments = getattr(bound_args, "arguments", None) + if not isinstance(arguments, dict): + return + + prompt = arguments.get("prompt") + if isinstance(prompt, str): + attrs["gen_ai.prompt.0.role"] = "user" + attrs["gen_ai.prompt.0.content"] = prompt + return + + messages = arguments.get("messages") + if not isinstance(messages, list): + return + for i, msg in enumerate(messages): + role = getattr(msg, "role", None) + if role is None and isinstance(msg, dict): + role = msg.get("role") + content = getattr(msg, "content", None) + if content is None and isinstance(msg, dict): + content = msg.get("content") + if role is not None: + attrs[f"gen_ai.prompt.{i}.role"] = str(role) + if content is not None: + attrs[f"gen_ai.prompt.{i}.content"] = _content_to_string(content) + + +def _start_retry_attempt(id_: str, instance: Any, bound_args: Any = None) -> None: + parent_span = trace.get_current_span() + if parent_span is None or not parent_span.get_span_context().is_valid: + # No ambient parent → orphan retry_attempt would have no place + # in the trace tree. Skip emission. Same graceful-degradation + # contract as LiteLLM/LangChain. + logger.debug( + "no ambient parent span for llamaindex retry_attempt; skipping (id_=%s)", + id_, + ) + return + + routed_provider = _resolve_routed_provider(instance) + model = _resolve_model(instance) + operation = "chat" + m = _DISPATCHER_ID_RE.match(id_) + if m is not None: + method = m.group(2) + if "complete" in method: + operation = "text_completion" + elif "predict" in method: + operation = "chat" # treat as chat-shaped + + attrs: dict[str, Any] = { + _FR_SPAN_ROLE_KEY: _FR_SPAN_ROLE_RETRY_ATTEMPT, + "gen_ai.operation.name": operation, + } + if routed_provider: + attrs["gen_ai.system"] = routed_provider + if model: + attrs["gen_ai.request.model"] = model + _add_prompt_attrs(attrs, bound_args) + + tracer = trace.get_tracer(__name__, __version__) + parent_ctx = set_span_in_context(parent_span) + span = tracer.start_span( + _FR_RETRY_ATTEMPT_SPAN_NAME, + kind=SpanKind.CLIENT, + attributes=attrs, + context=parent_ctx, + ) + + try: + framework_token = register_framework_attempt() + except Exception: + framework_token = None + logger.debug("failed to register framework attempt token", exc_info=True) + + now = time.monotonic() + with _FR_RETRY_ATTEMPT_MAP_LOCK: + evicted = _evict_stale_attempts_locked(now) + evicted += _enforce_attempt_max_locked() + if evicted: + _maybe_warn_eviction(evicted) + # Defensive: replace any prior entry with same id_ (shouldn't + # happen — dispatcher span ids are UUID-suffixed — but + # tolerated to avoid leaks if it does). + prev = _FR_RETRY_ATTEMPT_MAP.get(id_) + if prev is not None and not prev.get("ended"): + try: + prev["span"].set_status(Status(StatusCode.ERROR, "duplicate id_; superseded")) + prev["span"].end() + except Exception: + pass + try: + unregister_framework_attempt(prev.get("framework_token")) + except Exception: + pass + _FR_RETRY_ATTEMPT_MAP[id_] = { + "span": span, + "started_at": now, + "framework_token": framework_token, + "ended": False, + } + + # §4.5 marker timing: set on parent AFTER the first qualifying + # retry_attempt has successfully started under it. + try: + parent_span.set_attribute(_FR_HAS_RETRY_ATTEMPT_CHILD_KEY, True) + except Exception: + logger.debug("failed to set has_retry_attempt_child on parent", exc_info=True) + + +def _finalize_retry_attempt( + id_: str, + *, + success: bool, + result: Any = None, + err: Optional[BaseException] = None, +) -> None: + with _FR_RETRY_ATTEMPT_MAP_LOCK: + entry = _FR_RETRY_ATTEMPT_MAP.pop(id_, None) + if entry is None or entry.get("ended"): + return + entry["ended"] = True + span = entry["span"] + framework_token = entry.get("framework_token") + + try: + if success: + try: + # Best-effort attribute extraction. ChatResponse / CompletionResponse + # both expose .raw (provider-native), .message, .usage. + raw = getattr(result, "raw", None) + if raw is not None: + rid = getattr(raw, "id", None) + if isinstance(rid, str) and rid: + span.set_attribute("gen_ai.response.id", rid) + rmodel = getattr(raw, "model", None) + if isinstance(rmodel, str) and rmodel: + span.set_attribute("gen_ai.response.model", rmodel) + # Usage may live on .raw.usage or .additional_kwargs. + usage = getattr(raw, "usage", None) if raw is not None else None + if usage is not None: + pt = getattr(usage, "prompt_tokens", None) or getattr(usage, "input_tokens", None) + ct = getattr(usage, "completion_tokens", None) or getattr(usage, "output_tokens", None) + if isinstance(pt, int): + span.set_attribute("gen_ai.usage.input_tokens", pt) + if isinstance(ct, int): + span.set_attribute("gen_ai.usage.output_tokens", ct) + except Exception: + logger.debug("failed extracting llamaindex result attrs", exc_info=True) + span.set_status(Status(StatusCode.OK)) + else: + if err is not None: + error_type = type(err).__name__ + span.set_attribute("error.type", error_type) + status_code = getattr(err, "status_code", None) or getattr( + getattr(err, "response", None), "status_code", None + ) + if isinstance(status_code, int): + span.set_attribute("http.status_code", status_code) + try: + span.record_exception(err) + except Exception: + pass + span.set_status(Status(StatusCode.ERROR, str(err))) + else: + span.set_status(Status(StatusCode.ERROR, "retry_attempt failed")) + finally: + try: + span.end() + except Exception: + logger.debug("failed to end llamaindex retry_attempt span", exc_info=True) + try: + unregister_framework_attempt(framework_token) + except Exception: + pass + + +class _FortifyRootRetryHandler(BaseSpanHandler): + """LlamaIndex SpanHandler that emits one + ``fortifyroot.llamaindex.retry_attempt`` sibling span per OUTER + public LLM method invocation (chat/achat/complete/acomplete/...). + + De-dup vs the inner ``_chat``/``_complete`` spans is enforced by + a method-name whitelist in ``_is_outer_llm_method`` (F4 finding + from ST-10.0 C3 POC). So one HTTP attempt = exactly one + retry_attempt span — even though LlamaIndex internally fires + dispatcher spans on both the outer and inner methods. + + Registered ALONGSIDE the existing ``OpenLLMetrySpanHandler`` (not + as a replacement). Order matters: this handler MUST be registered + BEFORE OpenLLMetrySpanHandler so its ``span_enter`` fires while + the OTel ambient context is still the user's enclosing span (and + NOT yet the SpanHolder's per-call OTel span). This ensures all + retry_attempts under one logical user-call become SIBLINGS under + one OTel parent — the structural invariant RetryDetectorProc + needs. + """ + + @classmethod + def class_name(cls) -> str: + return "FortifyRootRetryHandler" + + def new_span( # type: ignore[override] + self, + id_: str, + bound_args, + instance: Optional[Any] = None, + parent_span_id: Optional[str] = None, + tags: Optional[dict] = None, + **kwargs, + ) -> None: + # We don't store SpanHolder objects (the existing + # OpenLLMetrySpanHandler does). We just open a retry_attempt + # OTel span and remember the id_ → otel_span mapping in + # _FR_RETRY_ATTEMPT_MAP. Returning None signals BaseSpanHandler + # to not store the result. + try: + if not _is_outer_llm_method(id_, instance): + return None + _start_retry_attempt(id_, instance, bound_args) + except Exception: + logger.debug("new_span retry-attempt-start failed", exc_info=True) + return None + + def prepare_to_exit_span( # type: ignore[override] + self, + id_: str, + bound_args, + instance: Optional[Any] = None, + result: Optional[Any] = None, + **kwargs, + ) -> None: + try: + if id_ in _FR_RETRY_ATTEMPT_MAP: + _finalize_retry_attempt(id_, success=True, result=result) + except Exception: + logger.debug("prepare_to_exit_span retry-attempt-finalize failed", exc_info=True) + return None + + def prepare_to_drop_span( # type: ignore[override] + self, + id_: str, + bound_args, + instance: Optional[Any] = None, + err: Optional[BaseException] = None, + **kwargs, + ) -> None: + try: + if id_ in _FR_RETRY_ATTEMPT_MAP: + _finalize_retry_attempt(id_, success=False, err=err) + except Exception: + logger.debug("prepare_to_drop_span retry-attempt-finalize failed", exc_info=True) + return None + + +def _reset_state_for_test() -> None: + """Test-only helper: clear all module state.""" + with _FR_RETRY_ATTEMPT_MAP_LOCK: + for entry in _FR_RETRY_ATTEMPT_MAP.values(): + try: + if not entry.get("ended"): + entry["span"].end() + except Exception: + pass + _FR_RETRY_ATTEMPT_MAP.clear() + + +__all__ = [ + "_FortifyRootRetryHandler", + "_FR_RETRY_ATTEMPT_SPAN_NAME", + "_FR_HAS_RETRY_ATTEMPT_CHILD_KEY", + "_OUTER_LLM_METHODS", +] diff --git a/packages/opentelemetry-instrumentation-llamaindex/tests/test_retry_attempt_emission.py b/packages/opentelemetry-instrumentation-llamaindex/tests/test_retry_attempt_emission.py new file mode 100644 index 0000000000..be25a165b7 --- /dev/null +++ b/packages/opentelemetry-instrumentation-llamaindex/tests/test_retry_attempt_emission.py @@ -0,0 +1,487 @@ +"""Tests for ST-10.3 LlamaIndex retry-aware emission. + +Covers: + - F4 de-dup contract: outer chat() AND inner _chat() both fire + dispatcher spans, but only ONE retry_attempt span is emitted + (the outer's) — the load-bearing test. + - Method-name whitelist: non-LLM dispatcher spans don't trigger + retry_attempt; outer-only filter correctly rejects "_chat". + - BaseLLM instance check: dispatcher spans on non-LLM classes + are ignored. + - Single-attempt happy path under tenacity-style wrapping. + - Multi-attempt retry path: 3 attempts share parent → 3 + retry_attempt SIBLINGS under one OTel parent. + - Marker timing (§4.5). + - §4.7.1 token registration symmetry. + - No-parent guard. + - Idempotent finalize. +""" + +from __future__ import annotations + +import inspect + +import pytest +from llama_index.core.base.llms.base import BaseLLM +from opentelemetry import trace +from opentelemetry.instrumentation.fortifyroot import ( + is_framework_owned, + retry_registry, +) +from opentelemetry.instrumentation.llamaindex.retry_handler import ( + _FortifyRootRetryHandler, + _FR_HAS_RETRY_ATTEMPT_CHILD_KEY, + _FR_RETRY_ATTEMPT_MAP, + _FR_RETRY_ATTEMPT_SPAN_NAME, + _is_outer_llm_method, + _OUTER_LLM_METHODS, + _reset_state_for_test, +) +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import ( + InMemorySpanExporter, +) + + +@pytest.fixture +def fresh_tracer(): + """Fresh TracerProvider + in-memory exporter installed as global, + matching the ST-10.1 / ST-10.2 fixture pattern.""" + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + try: + trace.set_tracer_provider(provider) + except Exception: + pass + current_provider = trace.get_tracer_provider() + if current_provider is not provider: + try: + current_provider.add_span_processor(SimpleSpanProcessor(exporter)) + except Exception: + pass + yield trace.get_tracer("test"), exporter, current_provider + + +@pytest.fixture(autouse=True) +def reset_state(): + retry_registry._reset_for_test() + _reset_state_for_test() + yield + retry_registry._reset_for_test() + _reset_state_for_test() + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +class _FakeLLM(BaseLLM): + """A minimal BaseLLM-shaped object for testing. Just satisfies the + isinstance(instance, BaseLLM) check; doesn't implement any + abstract methods.""" + + model: str = "gpt-4o-mini" + + @classmethod + def class_name(cls) -> str: + return "FakeLLM" + + @property + def metadata(self): + from llama_index.core.base.llms.types import LLMMetadata + return LLMMetadata() + + def chat(self, messages, **kwargs): # pragma: no cover + raise NotImplementedError + + async def achat(self, messages, **kwargs): # pragma: no cover + raise NotImplementedError + + def stream_chat(self, messages, **kwargs): # pragma: no cover + raise NotImplementedError + + async def astream_chat(self, messages, **kwargs): # pragma: no cover + raise NotImplementedError + + def complete(self, prompt, formatted=False, **kwargs): # pragma: no cover + raise NotImplementedError + + async def acomplete(self, prompt, formatted=False, **kwargs): # pragma: no cover + raise NotImplementedError + + def stream_complete(self, prompt, formatted=False, **kwargs): # pragma: no cover + raise NotImplementedError + + async def astream_complete(self, prompt, formatted=False, **kwargs): # pragma: no cover + raise NotImplementedError + + +def _empty_bound_args(): + """Return an empty inspect.BoundArguments for tests that don't + care about the arg payload.""" + sig = inspect.signature(lambda: None) + return sig.bind() + + +def _prompt_bound_args(prompt: str): + sig = inspect.signature(lambda prompt: None) + return sig.bind(prompt) + + +def _make_id(class_name: str, method: str) -> str: + """Mimic LlamaIndex's dispatcher span id format + ``ClassName.method-uuid``.""" + return f"{class_name}.{method}-test-uuid-1234" + + +class _FakeChatResponse: + def __init__(self, response_id=None, model=None, prompt_tokens=None, completion_tokens=None): + class _Raw: + pass + self.raw = _Raw() + if response_id: + self.raw.id = response_id + if model: + self.raw.model = model + if prompt_tokens is not None or completion_tokens is not None: + class _Usage: + pass + self.raw.usage = _Usage() + if prompt_tokens is not None: + self.raw.usage.prompt_tokens = prompt_tokens + if completion_tokens is not None: + self.raw.usage.completion_tokens = completion_tokens + + +# --------------------------------------------------------------------------- +# F4 de-dup contract — the LOAD-BEARING test. +# --------------------------------------------------------------------------- + +def test_outer_method_emits_inner_method_does_not(fresh_tracer): + """LlamaIndex dispatcher fires spans on BOTH the public ``chat()`` + AND the inner ``_chat()`` (F4 finding from POC). ST-10.3 hooks + SpanHandler.new_span and filters via + ``_is_outer_llm_method`` → only the outer method emits a + retry_attempt span. + + Net: ONE HTTP attempt = ONE retry_attempt span (NOT TWO), + even though the dispatcher emits two spans. + """ + tracer, exporter, _ = fresh_tracer + handler = _FortifyRootRetryHandler() + instance = _FakeLLM() + + parent = tracer.start_span("workflow") + with trace.use_span(parent, end_on_exit=False): + # Outer chat() span fires → emit retry_attempt. + outer_id = _make_id("FakeLLM", "chat") + handler.new_span(id_=outer_id, bound_args=_empty_bound_args(), instance=instance) + + # Inner _chat() span fires → DO NOT emit retry_attempt. + inner_id = _make_id("FakeLLM", "_chat") + handler.new_span(id_=inner_id, bound_args=_empty_bound_args(), instance=instance) + + # Inner _chat() exits → no-op (no entry to finalize). + handler.prepare_to_exit_span( + id_=inner_id, bound_args=_empty_bound_args(), + instance=instance, result=_FakeChatResponse(model="gpt-4o-mini"), + ) + # Outer chat() exits → finalize the retry_attempt. + handler.prepare_to_exit_span( + id_=outer_id, bound_args=_empty_bound_args(), + instance=instance, result=_FakeChatResponse(model="gpt-4o-mini"), + ) + parent.end() + + retry_spans = [ + s for s in exporter.get_finished_spans() + if s.name == _FR_RETRY_ATTEMPT_SPAN_NAME + ] + assert len(retry_spans) == 1, ( + f"F4 de-dup invariant: outer chat() emits ONE retry_attempt; " + f"inner _chat() must NOT emit. Got {len(retry_spans)} spans." + ) + + +def test_is_outer_llm_method_filter(): + """Direct unit test of the filter function. Outer public methods + pass, inner private methods fail, non-LLM instances fail.""" + instance = _FakeLLM() + not_instance = "not-an-llm" + + # Outer public methods — all pass. + assert _is_outer_llm_method(_make_id("FakeLLM", "chat"), instance) + assert _is_outer_llm_method(_make_id("FakeLLM", "achat"), instance) + assert _is_outer_llm_method(_make_id("FakeLLM", "complete"), instance) + assert _is_outer_llm_method(_make_id("FakeLLM", "stream_chat"), instance) + assert _is_outer_llm_method(_make_id("FakeLLM", "astream_complete"), instance) + + # Inner private methods — all fail. + assert not _is_outer_llm_method(_make_id("FakeLLM", "_chat"), instance) + assert not _is_outer_llm_method(_make_id("FakeLLM", "_complete"), instance) + assert not _is_outer_llm_method(_make_id("FakeLLM", "_predict"), instance) + + # Not a BaseLLM instance — fail. + assert not _is_outer_llm_method(_make_id("FakeLLM", "chat"), not_instance) + assert not _is_outer_llm_method(_make_id("FakeLLM", "chat"), None) + + # Malformed dispatcher id — fail. + assert not _is_outer_llm_method("malformed", instance) + + +def test_outer_method_whitelist_is_complete(): + """Sanity-check the _OUTER_LLM_METHODS whitelist matches the + BaseLLM public surface. If LlamaIndex adds new public methods + in a future version, the whitelist may need updating — this + test alerts on that condition.""" + expected_at_minimum = {"chat", "achat", "complete", "acomplete"} + assert expected_at_minimum.issubset(_OUTER_LLM_METHODS) + + +# --------------------------------------------------------------------------- +# Single-attempt happy path. +# --------------------------------------------------------------------------- + +def test_single_attempt_emits_one_retry_attempt(fresh_tracer): + tracer, exporter, _ = fresh_tracer + handler = _FortifyRootRetryHandler() + instance = _FakeLLM() + + parent = tracer.start_span("workflow") + with trace.use_span(parent, end_on_exit=False): + id_ = _make_id("FakeLLM", "complete") + handler.new_span( + id_=id_, + bound_args=_prompt_bound_args("hello masked [EMAIL]"), + instance=instance, + ) + handler.prepare_to_exit_span( + id_=id_, bound_args=_empty_bound_args(), + instance=instance, + result=_FakeChatResponse( + response_id="resp-001", model="gpt-4o-mini", + prompt_tokens=10, completion_tokens=5, + ), + ) + parent.end() + + spans = exporter.get_finished_spans() + parent_exported = next(s for s in spans if s.name == "workflow") + retry_span = next(s for s in spans if s.name == _FR_RETRY_ATTEMPT_SPAN_NAME) + + assert retry_span.parent.span_id == parent_exported.context.span_id + assert parent_exported.attributes.get(_FR_HAS_RETRY_ATTEMPT_CHILD_KEY) is True + assert retry_span.attributes.get("fortifyroot.span.role") == "retry_attempt" + assert retry_span.attributes.get("gen_ai.request.model") == "gpt-4o-mini" + assert retry_span.attributes.get("gen_ai.response.id") == "resp-001" + assert retry_span.attributes.get("gen_ai.usage.input_tokens") == 10 + assert retry_span.attributes.get("gen_ai.usage.output_tokens") == 5 + assert retry_span.attributes.get("gen_ai.prompt.0.role") == "user" + assert retry_span.attributes.get("gen_ai.prompt.0.content") == "hello masked [EMAIL]" + + +# --------------------------------------------------------------------------- +# Multi-attempt retry path — siblings under one parent. +# --------------------------------------------------------------------------- + +def test_three_attempts_share_one_otel_parent(fresh_tracer): + """Under tenacity wrapping at the application layer, each retry + creates a fresh dispatcher span tree (different id_ per attempt), + BUT they all run under the same enclosing OTel span as ambient. + The retry_handler uses ambient OTel context as parent → all 3 + retry_attempts become SIBLINGS under one OTel parent.""" + tracer, exporter, _ = fresh_tracer + handler = _FortifyRootRetryHandler() + instance = _FakeLLM() + + class MockHTTPError(Exception): + def __init__(self, status_code): + self.status_code = status_code + super().__init__(f"http {status_code}") + + parent = tracer.start_span("workflow") + with trace.use_span(parent, end_on_exit=False): + # Attempt 1 — 429 + id1 = _make_id("FakeLLM", "chat") + "-attempt1" + handler.new_span(id_=id1, bound_args=_empty_bound_args(), instance=instance) + # Inner _chat span also fires — but is filtered. + inner1 = _make_id("FakeLLM", "_chat") + "-attempt1" + handler.new_span( + id_=inner1, bound_args=_empty_bound_args(), instance=instance, + ) + handler.prepare_to_drop_span( + id_=inner1, bound_args=_empty_bound_args(), + instance=instance, err=MockHTTPError(429), + ) + handler.prepare_to_drop_span( + id_=id1, bound_args=_empty_bound_args(), + instance=instance, err=MockHTTPError(429), + ) + + # Attempt 2 — 429 + id2 = _make_id("FakeLLM", "chat") + "-attempt2" + handler.new_span(id_=id2, bound_args=_empty_bound_args(), instance=instance) + handler.prepare_to_drop_span(id_=id2, bound_args=_empty_bound_args(), instance=instance, err=MockHTTPError(429)) + + # Attempt 3 — 200 + id3 = _make_id("FakeLLM", "chat") + "-attempt3" + handler.new_span(id_=id3, bound_args=_empty_bound_args(), instance=instance) + handler.prepare_to_exit_span( + id_=id3, bound_args=_empty_bound_args(), instance=instance, + result=_FakeChatResponse(model="gpt-4o-mini"), + ) + parent.end() + + spans = exporter.get_finished_spans() + parent_exported = next(s for s in spans if s.name == "workflow") + retry_spans = [s for s in spans if s.name == _FR_RETRY_ATTEMPT_SPAN_NAME] + + assert len(retry_spans) == 3, ( + f"3 outer-method spans should produce 3 retry_attempts (inner _chat filtered out); " + f"got {len(retry_spans)}" + ) + parent_ids = {s.parent.span_id for s in retry_spans} + assert parent_ids == {parent_exported.context.span_id}, ( + f"all 3 retry_attempts must be SIBLINGS under one OTel parent; " + f"saw distinct parents: {parent_ids}" + ) + + from opentelemetry.trace import StatusCode + error_count = sum(1 for s in retry_spans if s.status.status_code == StatusCode.ERROR) + ok_count = sum(1 for s in retry_spans if s.status.status_code == StatusCode.OK) + assert error_count == 2 and ok_count == 1 + for s in retry_spans: + if s.status.status_code == StatusCode.ERROR: + assert s.attributes.get("http.status_code") == 429 + assert s.attributes.get("error.type") == "MockHTTPError" + + +# --------------------------------------------------------------------------- +# Marker timing (§4.5). +# --------------------------------------------------------------------------- + +def test_marker_set_AFTER_first_attempt(fresh_tracer): + tracer, _, _ = fresh_tracer + handler = _FortifyRootRetryHandler() + instance = _FakeLLM() + + parent = tracer.start_span("workflow") + assert _FR_HAS_RETRY_ATTEMPT_CHILD_KEY not in dict(parent.attributes or {}) + + with trace.use_span(parent, end_on_exit=False): + handler.new_span( + id_=_make_id("FakeLLM", "chat"), bound_args=_empty_bound_args(), + instance=instance, + ) + assert dict(parent.attributes or {}).get(_FR_HAS_RETRY_ATTEMPT_CHILD_KEY) is True + + handler.prepare_to_exit_span( + id_=_make_id("FakeLLM", "chat"), bound_args=_empty_bound_args(), + instance=instance, result=_FakeChatResponse(model="gpt-4o-mini"), + ) + parent.end() + + +def test_marker_NOT_set_when_inner_only_fires(fresh_tracer): + """If only inner methods fire (no outer chat ever), no + retry_attempt is emitted and the parent stays unmarked.""" + tracer, exporter, _ = fresh_tracer + handler = _FortifyRootRetryHandler() + instance = _FakeLLM() + + parent = tracer.start_span("workflow") + with trace.use_span(parent, end_on_exit=False): + # Only inner method fires. + handler.new_span( + id_=_make_id("FakeLLM", "_chat"), bound_args=_empty_bound_args(), + instance=instance, + ) + handler.prepare_to_exit_span( + id_=_make_id("FakeLLM", "_chat"), bound_args=_empty_bound_args(), + instance=instance, result=_FakeChatResponse(model="gpt-4o-mini"), + ) + parent.end() + + parent_exported = next( + s for s in exporter.get_finished_spans() if s.name == "workflow" + ) + assert _FR_HAS_RETRY_ATTEMPT_CHILD_KEY not in (parent_exported.attributes or {}) + + +# --------------------------------------------------------------------------- +# §4.7.1 token registration. +# --------------------------------------------------------------------------- + +def test_framework_token_lifecycle(fresh_tracer): + tracer, _, _ = fresh_tracer + handler = _FortifyRootRetryHandler() + instance = _FakeLLM() + + parent = tracer.start_span("workflow") + with trace.use_span(parent, end_on_exit=False): + assert not is_framework_owned() + id_ = _make_id("FakeLLM", "chat") + handler.new_span(id_=id_, bound_args=_empty_bound_args(), instance=instance) + assert is_framework_owned(), "during attempt → framework owns the call" + handler.prepare_to_exit_span( + id_=id_, bound_args=_empty_bound_args(), + instance=instance, result=_FakeChatResponse(model="gpt-4o-mini"), + ) + assert not is_framework_owned(), "after attempt → released" + parent.end() + + +# --------------------------------------------------------------------------- +# No-parent guard. +# --------------------------------------------------------------------------- + +def test_no_ambient_parent_no_emission(fresh_tracer): + tracer, exporter, _ = fresh_tracer + handler = _FortifyRootRetryHandler() + instance = _FakeLLM() + + # No parent span attached. + handler.new_span( + id_=_make_id("FakeLLM", "chat"), bound_args=_empty_bound_args(), + instance=instance, + ) + + retry_spans = [ + s for s in exporter.get_finished_spans() + if s.name == _FR_RETRY_ATTEMPT_SPAN_NAME + ] + assert len(retry_spans) == 0 + assert len(_FR_RETRY_ATTEMPT_MAP) == 0 + + +# --------------------------------------------------------------------------- +# Idempotency. +# --------------------------------------------------------------------------- + +def test_double_finalize_is_idempotent(fresh_tracer): + tracer, exporter, _ = fresh_tracer + handler = _FortifyRootRetryHandler() + instance = _FakeLLM() + + parent = tracer.start_span("workflow") + with trace.use_span(parent, end_on_exit=False): + id_ = _make_id("FakeLLM", "chat") + handler.new_span(id_=id_, bound_args=_empty_bound_args(), instance=instance) + handler.prepare_to_exit_span( + id_=id_, bound_args=_empty_bound_args(), + instance=instance, result=_FakeChatResponse(model="gpt-4o-mini"), + ) + # Second finalize — no-op. + handler.prepare_to_exit_span( + id_=id_, bound_args=_empty_bound_args(), + instance=instance, result=_FakeChatResponse(model="gpt-4o-mini"), + ) + parent.end() + + retry_spans = [ + s for s in exporter.get_finished_spans() + if s.name == _FR_RETRY_ATTEMPT_SPAN_NAME + ] + assert len(retry_spans) == 1 diff --git a/packages/opentelemetry-instrumentation-openai/opentelemetry/instrumentation/openai/retry_handler.py b/packages/opentelemetry-instrumentation-openai/opentelemetry/instrumentation/openai/retry_handler.py new file mode 100644 index 0000000000..0afba46e5e --- /dev/null +++ b/packages/opentelemetry-instrumentation-openai/opentelemetry/instrumentation/openai/retry_handler.py @@ -0,0 +1,691 @@ +"""ST-10.4 retry-aware emission for the OpenAI direct SDK. + +Per RETRY_LOOP.md §4.4 OpenAI row + §4.7 suppression discipline + +ST-10.0 hook-table addendum (in phase_st10_retryloop.txt): + + - Hook the SDK-internal private httpx wrapper classes + ``openai._base_client.SyncHttpxClientWrapper.send`` and + ``AsyncHttpxClientWrapper.send``. These run ONCE per HTTP attempt + inside the SDK's internal retry loop, giving us a clean per-attempt + boundary without globally monkey-patching ``httpx.Client`` (which + would emit retry_attempt spans for every httpx caller in the + process — including non-LLM ones). + - The wrap is INSTANCE-level via ``wrapt.wrap_function_wrapper`` + against the class, so newly-constructed clients automatically pick + it up. Unwrap is symmetric (``opentelemetry.instrumentation.utils.unwrap``). + - Private-symbol guard: ``openai._base_client`` is private to the + SDK and may change between versions. We wrap defensively via + ``getattr`` + a try/except around the wrap call; if either + ``SyncHttpxClientWrapper.send`` or ``AsyncHttpxClientWrapper.send`` + is missing or incompatibly typed, we log a warning and skip + retry-attempt emission for that variant. Normal openai + instrumentation continues unaffected. + - Endpoint allow-list: only LLM endpoints emit retry_attempt spans + (``/v1/chat/completions``, ``/v1/completions``, ``/v1/embeddings``, + ``/v1/responses``, plus Azure ``/openai/deployments/.../chat/completions``). + Non-LLM SDK traffic (e.g. ``/v1/models`` for token-refresh / model + listing) does NOT emit retry_attempt — per §4.4.1 allow-listing + requirement. + - Suppression discipline (§4.7): before emitting, check BOTH the + OTel context ``SUPPRESS_LANGUAGE_MODEL_INSTRUMENTATION_KEY`` AND + the shared ``is_framework_owned()`` registry. If either says + "suppress", we skip emission. This is the key invariant that + prevents LiteLLM/LangChain/LlamaIndex framework retries from + DOUBLE-emitting when both framework wrappers and direct-SDK + wrappers are active. + - Parent resolution: use the current OTel ambient span. If invalid + or absent, gracefully skip (no orphan retry_attempt). The §4.5 + backend dedup degrades gracefully. + - Per-attempt span lifetime: open on wrap entry, end on wrap exit + (success OR exception). On non-2xx HTTP, set + ``http.status_code`` + ``error.type`` and ERROR status; on 2xx, + set best-effort response attrs + OK status. + +Tests: see ``tests/test_retry_attempt_emission.py``. +""" + +from __future__ import annotations + +import logging +import threading +import time +from typing import Any, Optional + +from opentelemetry import context as context_api +from opentelemetry import trace +from opentelemetry.instrumentation.fortifyroot import ( + is_framework_owned, +) +from opentelemetry.instrumentation.openai.version import __version__ +from opentelemetry.instrumentation.utils import unwrap +from opentelemetry.semconv_ai import SUPPRESS_LANGUAGE_MODEL_INSTRUMENTATION_KEY +from opentelemetry.trace import SpanKind, Status, StatusCode +from wrapt import wrap_function_wrapper + +logger = logging.getLogger(__name__) + + +# ST-10 §4.4 / §4.5 constants. +_FR_RETRY_ATTEMPT_SPAN_NAME = "fortifyroot.openai.retry_attempt" +_FR_SPAN_ROLE_KEY = "fortifyroot.span.role" +_FR_SPAN_ROLE_RETRY_ATTEMPT = "retry_attempt" +_FR_HAS_RETRY_ATTEMPT_CHILD_KEY = "fortifyroot.span.has_retry_attempt_child" + +# Package-local context key set by the OpenAI logical-call wrappers +# (currently `chat_wrappers.chat_wrapper` / `achat_wrapper`) BEFORE +# they invoke the wrapped OpenAI SDK call. Used by ``_is_suppressed()`` +# to distinguish: +# +# - "openai wrapper is making its OWN internal HTTP attempt under +# its OWN logical span" — retry_attempt MUST emit (the openai +# wrapper sets ``SUPPRESS_LANGUAGE_MODEL_INSTRUMENTATION_KEY`` +# to protect against OTHER instrumentors double-counting; we are +# NOT one of those — we are part of the openai instrumentor). +# +# - "external suppression" (user set the key globally, OR a framework +# wrapper set it to disable LLM instrumentation for a scope) — +# retry_attempt MUST be suppressed. +# +# The override key is read via ``context_api.get_value`` and is True +# only inside the chat_wrapper's `wrapped(*args, **kwargs)` call. See +# ``opentelemetry.instrumentation.openai.shared.chat_wrappers``. +OPENAI_DIRECT_RETRY_PARENT_ACTIVE_KEY = "fortifyroot.openai.direct_retry_parent_active" + +# Private symbols we wrap. Recorded here as module-level constants so the +# instrument/uninstrument code paths share one source of truth and the +# private-symbol guard tests can introspect them. +_OPENAI_BASE_CLIENT_MODULE = "openai._base_client" +_SYNC_WRAPPER_CLASS = "SyncHttpxClientWrapper" +_ASYNC_WRAPPER_CLASS = "AsyncHttpxClientWrapper" +_WRAPPED_METHOD = "send" + +# §4.4.1 endpoint allow-list: matched against the request URL's path +# suffix. Hits any of these → LLM endpoint → emit retry_attempt. +# Misses → non-LLM SDK traffic (e.g. /v1/models, auth refresh) → +# skip emission. +_LLM_PATH_SUFFIXES = ( + "/chat/completions", # /v1/chat/completions, Azure /openai/deployments/.../chat/completions + "/completions", # /v1/completions (legacy completion); also matches /chat/completions + "/embeddings", # /v1/embeddings + "/responses", # /v1/responses (Responses API) + "/messages", # /v1/messages (Azure-style messages endpoints, OpenAI parity layer) +) + + +# Tracks whether instrumentation is currently installed, so the +# private-symbol guard test path can verify symmetry without double-wrap +# / double-unwrap explosions when called by ``instrument()`` machinery +# (the BaseInstrumentor singleton has its own state, but this module +# also exposes a direct install/uninstall for tests). +_state_lock = threading.Lock() +_installed = False + +# Tracer provider the wrap was installed with. ``None`` → fall back to +# the global ``TracerProvider`` (matches behaviour pre-2026-05-16). Set +# via ``instrument_retry_emitter(tracer_provider=...)`` so retry spans +# go to the SAME provider as the openai logical span emitted by +# ``chat_wrapper`` etc. Otherwise a consumer that passes an explicit +# provider to ``OpenAIInstrumentor().instrument(tracer_provider=...)`` +# gets the logical openai span on their provider but the retry_attempt +# span sent into a no-op tracer (when there is no global provider). +_tracer_provider = None + + +def _request_path(request: Any) -> str: + """Return the URL path of an httpx.Request, or '' if unavailable. + + The OpenAI SDK passes an ``httpx.Request`` as the first positional + argument to ``SyncHttpxClientWrapper.send``. ``httpx.URL.path`` is + the percent-decoded path. + """ + try: + url = getattr(request, "url", None) + if url is None: + return "" + path = getattr(url, "path", None) + if isinstance(path, str): + return path + # Older httpx versions: stringify and rely on the .path attribute + # being populated. If not present, fall back to str(url) which + # includes the path. + return str(url) + except Exception: + return "" + + +def _is_llm_endpoint(path: str) -> bool: + if not path: + return False + for suffix in _LLM_PATH_SUFFIXES: + if path.endswith(suffix): + return True + return False + + +def _operation_for_path(path: str) -> str: + """Map a URL path to a ``gen_ai.operation.name`` value.""" + if path.endswith("/embeddings"): + return "embeddings" + if path.endswith("/chat/completions") or path.endswith("/messages"): + return "chat" + if path.endswith("/responses"): + return "chat" # Responses API is chat-shaped + if path.endswith("/completions"): + return "text_completion" + return "chat" + + +def _resolve_model_from_request(request: Any) -> Optional[str]: + """Best-effort: pull the ``model`` field from the request body. The + OpenAI SDK serialises the JSON payload before constructing the + httpx.Request, so ``request.content`` is bytes carrying the JSON. + + On failure (binary streamed body, non-JSON, parse error), return + None — we omit the attribute rather than guessing. + """ + try: + content = getattr(request, "content", None) + if content is None or not isinstance(content, (bytes, bytearray)): + return None + # Cheap, defensive JSON decode. + import json + body = json.loads(content.decode("utf-8")) + if isinstance(body, dict): + model = body.get("model") + if isinstance(model, str) and model: + return model + except Exception: + return None + return None + + +def _server_attrs_from_request(request: Any) -> dict[str, Any]: + """Extract ``server.address`` and ``server.port`` from an httpx.Request.""" + out: dict[str, Any] = {} + try: + url = getattr(request, "url", None) + if url is None: + return out + host = getattr(url, "host", None) + if isinstance(host, str) and host: + out["server.address"] = host + port = getattr(url, "port", None) + if port is None: + scheme = getattr(url, "scheme", None) + port = 443 if scheme == "https" else 80 if scheme == "http" else None + if isinstance(port, int): + out["server.port"] = port + except Exception: + pass + return out + + +def _is_suppressed() -> bool: + """§4.7: skip emission if EITHER suppression signal is active. + + Subtle: the OpenAI ``chat_wrapper`` (and ``achat_wrapper``) sets + ``SUPPRESS_LANGUAGE_MODEL_INSTRUMENTATION_KEY`` around its own + wrapped SDK call to protect against OTHER (external) LLM + instrumentors double-counting the same call. That is NOT the + contract our retry handler should respect — we are part of the + openai instrumentor itself and want to emit retry_attempt under + the openai logical span. The chat_wrapper signals this case by + ALSO setting ``OPENAI_DIRECT_RETRY_PARENT_ACTIVE_KEY`` in the + same context; we treat that as an override. + + For all OTHER scenarios where ``SUPPRESS_LANGUAGE_MODEL_INSTRUMENTATION_KEY`` + is True without the override key (user explicitly disabled LLM + instrumentation, framework wrapper suppressing direct-SDK, etc.), + suppression still applies. + """ + try: + if context_api.get_value(SUPPRESS_LANGUAGE_MODEL_INSTRUMENTATION_KEY): + # Override: openai-wrapper-internal suppression — emit anyway. + if not context_api.get_value(OPENAI_DIRECT_RETRY_PARENT_ACTIVE_KEY): + return True + except Exception: + pass + try: + if is_framework_owned(): + return True + except Exception: + pass + return False + + +def _resolve_parent_span() -> Optional["trace.Span"]: + current = trace.get_current_span() + if current is None: + return None + ctx = current.get_span_context() + if ctx is None or not ctx.is_valid: + return None + return current + + +def _set_parent_marker(parent_span: "trace.Span") -> None: + """§4.5: mark the parent as 'has retry_attempt child' AFTER the + first child has successfully started. Idempotent — setting twice + is a no-op.""" + try: + parent_span.set_attribute(_FR_HAS_RETRY_ATTEMPT_CHILD_KEY, True) + except Exception: + logger.debug("failed to set has_retry_attempt_child on parent", exc_info=True) + + +def _start_attempt_span(request: Any, parent_span: "trace.Span") -> "trace.Span": + """Open the retry_attempt sibling span under the given parent. + Caller is responsible for adding response attrs + ending it. + """ + path = _request_path(request) + attrs: dict[str, Any] = { + _FR_SPAN_ROLE_KEY: _FR_SPAN_ROLE_RETRY_ATTEMPT, + "gen_ai.system": "openai", + "gen_ai.operation.name": _operation_for_path(path), + } + model = _resolve_model_from_request(request) + if model: + attrs["gen_ai.request.model"] = model + attrs.update(_server_attrs_from_request(request)) + + tracer = trace.get_tracer(__name__, __version__, _tracer_provider) + parent_ctx = trace.set_span_in_context(parent_span) + span = tracer.start_span( + _FR_RETRY_ATTEMPT_SPAN_NAME, + kind=SpanKind.CLIENT, + attributes=attrs, + context=parent_ctx, + ) + return span + + +def _finalize_success(span: "trace.Span", response: Any, *, is_streaming: bool = False) -> None: + """Apply response-side attributes + status. Treats non-2xx HTTP + responses as errors even when no exception is raised — httpx does + NOT raise on non-2xx by default; the OpenAI SDK's high-level layers + raise above us, but at the wrap layer we see the raw response. + + For non-streaming responses (regardless of 2xx vs non-2xx), parse + the JSON body and extract usage tokens + response id + response + model. Per RETRY_LOOP.md §4.4 token-usage rule (around line 164): + wrappers MUST extract usage from the response body whenever it's + present, regardless of whether the attempt succeeded — some + failures (context-length-exceeded errors etc.) DO consume tokens + and the provider returns usage in the error body. Backend §4.5 + dedup makes a qualifying retry_attempt canonical (even single- + attempt) and reads token attrs from it + (``fr-backend/internal/processing/proc_llm_extractor.go`` + ``isRetryAttempt`` block) — so usage attribution must live here. + + For streaming responses (``stream=True`` passed to ``send()``), + body reading would consume the SSE stream before the SDK can + iterate it. Usage attrs are therefore OMITTED on streaming + retry_attempts in this MVP; full streaming usage capture would + require intercepting the SSE chunk stream and is tracked as a + follow-up. + """ + try: + status_code = getattr(response, "status_code", None) + if isinstance(status_code, int): + span.set_attribute("http.status_code", status_code) + + # Header-based response id (works for streaming too, both success and error). + try: + headers = getattr(response, "headers", None) or {} + rid = None + if hasattr(headers, "get"): + rid = headers.get("x-request-id") or headers.get("openai-request-id") + if isinstance(rid, str) and rid: + span.set_attribute("gen_ai.response.id", rid) + except Exception: + pass + + # Non-streaming body parse — applies to BOTH 2xx and non-2xx per the + # §4.4 token-usage rule. ``_extract_usage_from_body`` is fully + # defensive: missing fields, parse failure, and bodies without + # ``usage`` all degrade silently to "no attrs set". + if not is_streaming: + _extract_usage_from_body(span, response) + + # Status + error attribution. + if isinstance(status_code, int) and 200 <= status_code < 300: + span.set_status(Status(StatusCode.OK)) + else: + err_type = "openai.HTTPStatusError" + if isinstance(status_code, int): + if status_code == 429: + err_type = "openai.RateLimitError" + elif 500 <= status_code < 600: + err_type = "openai.InternalServerError" + elif status_code == 401: + err_type = "openai.AuthenticationError" + elif status_code == 403: + err_type = "openai.PermissionDeniedError" + span.set_attribute("error.type", err_type) + span.set_status(Status(StatusCode.ERROR, f"http {status_code}")) + except Exception: + logger.debug("failed to set response attrs on openai retry_attempt", exc_info=True) + + +def _extract_usage_from_body(span: "trace.Span", response: Any) -> None: + """Parse a non-streaming OpenAI response body and copy usage / + response id / response model attrs to the retry_attempt span. + + All operations are wrapped in try/except — body parsing is + best-effort. If anything fails (non-JSON body, malformed schema, + missing fields), the span is left without those attrs rather than + raising. + """ + try: + # ``.json()`` triggers ``.read()`` if the body isn't already + # buffered, then parses. Subsequent SDK reads hit the cached + # ``_content``. + body = response.json() + except Exception: + return + if not isinstance(body, dict): + return + try: + rid = body.get("id") + if isinstance(rid, str) and rid: + span.set_attribute("gen_ai.response.id", rid) + rmodel = body.get("model") + if isinstance(rmodel, str) and rmodel: + span.set_attribute("gen_ai.response.model", rmodel) + usage = body.get("usage") + if isinstance(usage, dict): + # OpenAI chat / completions schema: prompt_tokens / completion_tokens. + # Newer Responses API uses input_tokens / output_tokens — tolerate both. + pt = usage.get("prompt_tokens") + if not isinstance(pt, int): + pt = usage.get("input_tokens") + ct = usage.get("completion_tokens") + if not isinstance(ct, int): + ct = usage.get("output_tokens") + if isinstance(pt, int): + span.set_attribute("gen_ai.usage.input_tokens", pt) + if isinstance(ct, int): + span.set_attribute("gen_ai.usage.output_tokens", ct) + except Exception: + logger.debug("failed to extract usage from openai response body", exc_info=True) + + +def _finalize_error(span: "trace.Span", error: BaseException) -> None: + """Apply error-side attributes + status.""" + try: + error_type = type(error).__name__ + # Fully-qualified type name when available. + mod = type(error).__module__ + if isinstance(mod, str) and mod and mod != "builtins": + error_type = f"{mod.split('.')[0]}.{error_type}" + span.set_attribute("error.type", error_type) + status_code = getattr(error, "status_code", None) or getattr( + getattr(error, "response", None), "status_code", None + ) + if isinstance(status_code, int): + span.set_attribute("http.status_code", status_code) + try: + span.record_exception(error) + except Exception: + pass + span.set_status(Status(StatusCode.ERROR, str(error))) + except Exception: + logger.debug("failed to set error attrs on openai retry_attempt", exc_info=True) + + +def _should_emit_for(request: Any) -> bool: + """Combine all skip-emission guards into one decision. Returns + True iff a retry_attempt span SHOULD be emitted for this request. + """ + if _is_suppressed(): + return False + path = _request_path(request) + if not _is_llm_endpoint(path): + return False + return True + + +def _sync_send_wrapper(wrapped, instance, args, kwargs): + """Wraps ``openai._base_client.SyncHttpxClientWrapper.send``. + + The ``send(request, **kwargs)`` signature is the standard httpx one; + ``request`` is always the first positional argument; ``stream`` (bool) + is passed by the OpenAI SDK when the high-level call is a streaming + request. + + Framework-attempt registry: this wrapper does NOT register a token. + The §4.7.1 registry's documented contract (see ``retry_registry.py`` + docstring) reserves registration for FRAMEWORK wrappers (LiteLLM / + LangChain / LlamaIndex). Direct-SDK wrappers only CONSULT via + ``is_framework_owned()``. Self-registering would falsely suppress + other concurrent direct-SDK calls on the same OS thread — a real + issue under asyncio where two tasks share a thread. + + Streaming skip (ST-10.4 review-driven 2026-05-17): when + ``stream=True`` is passed to send, this wrapper SKIPS retry_attempt + emission entirely. Rationale: streaming retry_attempts cannot + carry token usage (the body is the SSE stream and reading it would + break the SDK), but backend §4.5 dedup makes any qualifying + retry_attempt the canonical LLMUsageEvent — producing a zero-token + canonical event for streaming calls. By not emitting at all, the + parent ``openai.chat`` span (which DOES get full usage attribution + from ``ChatStream``'s stream-completion callback) remains the + canonical event. Streaming retry-loop detection is the documented + deferred follow-up ``ST-10.4-FOLLOWUP-streaming-usage``; this + explicit skip is part of that deferral. + """ + request = args[0] if args else kwargs.get("request") + if request is None or not _should_emit_for(request): + return wrapped(*args, **kwargs) + + # ST-10.4: skip retry_attempt emission for streaming calls. + if bool(kwargs.get("stream", False)): + return wrapped(*args, **kwargs) + + parent = _resolve_parent_span() + if parent is None: + return wrapped(*args, **kwargs) + + is_streaming = False # only reached on non-streaming path + span = _start_attempt_span(request, parent) + _set_parent_marker(parent) + + try: + response = wrapped(*args, **kwargs) + except BaseException as exc: + try: + _finalize_error(span, exc) + finally: + try: + span.end() + except Exception: + pass + raise + try: + _finalize_success(span, response, is_streaming=is_streaming) + finally: + try: + span.end() + except Exception: + pass + return response + + +async def _async_send_wrapper(wrapped, instance, args, kwargs): + """Wraps ``openai._base_client.AsyncHttpxClientWrapper.send``. + + Same no-self-registration contract and streaming-skip behaviour as + ``_sync_send_wrapper``. + """ + request = args[0] if args else kwargs.get("request") + if request is None or not _should_emit_for(request): + return await wrapped(*args, **kwargs) + + # ST-10.4: skip retry_attempt emission for streaming (see + # _sync_send_wrapper docstring for rationale). + if bool(kwargs.get("stream", False)): + return await wrapped(*args, **kwargs) + + parent = _resolve_parent_span() + if parent is None: + return await wrapped(*args, **kwargs) + + is_streaming = False + span = _start_attempt_span(request, parent) + _set_parent_marker(parent) + + try: + response = await wrapped(*args, **kwargs) + except BaseException as exc: + try: + _finalize_error(span, exc) + finally: + try: + span.end() + except Exception: + pass + raise + try: + _finalize_success(span, response, is_streaming=is_streaming) + finally: + try: + span.end() + except Exception: + pass + return response + + +def _has_wrappable_symbol(module_name: str, class_name: str, method_name: str) -> bool: + """Import-time wrappability check. Returns True iff the named symbol + exists, is a class, and has the named method as a function-like + attribute. Used by ``instrument_retry_emitter`` to skip wrapping + when private OpenAI symbols are missing or incompatible (without + breaking normal instrumentation). + """ + try: + import importlib + module = importlib.import_module(module_name) + except Exception: + return False + cls = getattr(module, class_name, None) + if cls is None or not isinstance(cls, type): + return False + method = getattr(cls, method_name, None) + if method is None: + return False + if not callable(method): + return False + return True + + +def instrument_retry_emitter(tracer_provider=None) -> None: + """Install the OpenAI retry_attempt emitter. Idempotent — calling + twice without an intervening uninstrument is a no-op. + + Wraps: + - openai._base_client.SyncHttpxClientWrapper.send + - openai._base_client.AsyncHttpxClientWrapper.send + + ``tracer_provider`` is the TracerProvider retry spans should be + emitted through. Caller should pass the same provider the parent + ``OpenAIInstrumentor`` was configured with so retry_attempt spans + land in the same exporter chain as the openai logical span. If + ``None``, falls back to the global TracerProvider (pre-2026-05-16 + behaviour; correct when the consumer has set a global provider). + + If a private symbol is missing or incompatible (newer/older openai + SDK has refactored ``_base_client``), logs a warning and skips + wrapping that variant. Normal openai instrumentation is NOT + affected. + """ + global _installed, _tracer_provider + with _state_lock: + if _installed: + return + _tracer_provider = tracer_provider + sync_ok = _has_wrappable_symbol( + _OPENAI_BASE_CLIENT_MODULE, _SYNC_WRAPPER_CLASS, _WRAPPED_METHOD, + ) + async_ok = _has_wrappable_symbol( + _OPENAI_BASE_CLIENT_MODULE, _ASYNC_WRAPPER_CLASS, _WRAPPED_METHOD, + ) + if not sync_ok: + logger.warning( + "ST-10.4: openai._base_client.%s.%s missing/incompatible; " + "skipping sync retry_attempt emission. Normal openai " + "instrumentation is unaffected.", + _SYNC_WRAPPER_CLASS, _WRAPPED_METHOD, + ) + else: + try: + wrap_function_wrapper( + _OPENAI_BASE_CLIENT_MODULE, + f"{_SYNC_WRAPPER_CLASS}.{_WRAPPED_METHOD}", + _sync_send_wrapper, + ) + except Exception as e: + logger.warning( + "ST-10.4: failed to wrap openai sync httpx send (%s); " + "retry_attempt emission disabled for sync path", + e, + ) + if not async_ok: + logger.warning( + "ST-10.4: openai._base_client.%s.%s missing/incompatible; " + "skipping async retry_attempt emission. Normal openai " + "instrumentation is unaffected.", + _ASYNC_WRAPPER_CLASS, _WRAPPED_METHOD, + ) + else: + try: + wrap_function_wrapper( + _OPENAI_BASE_CLIENT_MODULE, + f"{_ASYNC_WRAPPER_CLASS}.{_WRAPPED_METHOD}", + _async_send_wrapper, + ) + except Exception as e: + logger.warning( + "ST-10.4: failed to wrap openai async httpx send (%s); " + "retry_attempt emission disabled for async path", + e, + ) + _installed = True + + +def uninstrument_retry_emitter() -> None: + """Remove the OpenAI retry_attempt wraps. Idempotent — calling + twice or before install is a no-op. + """ + global _installed, _tracer_provider + with _state_lock: + if not _installed: + return + _tracer_provider = None + for cls in (_SYNC_WRAPPER_CLASS, _ASYNC_WRAPPER_CLASS): + try: + unwrap(f"{_OPENAI_BASE_CLIENT_MODULE}.{cls}", _WRAPPED_METHOD) + except Exception: + logger.debug( + "ST-10.4: openai unwrap of %s.%s failed (likely " + "wrap was never installed for this variant)", + cls, _WRAPPED_METHOD, + exc_info=True, + ) + _installed = False + + +def _is_installed_for_test() -> bool: + """Test-only helper: True iff retry-emitter wraps are currently + installed. Not part of the public contract; do not import outside + tests.""" + with _state_lock: + return _installed + + +__all__ = [ + "instrument_retry_emitter", + "uninstrument_retry_emitter", + "_FR_RETRY_ATTEMPT_SPAN_NAME", + "_FR_HAS_RETRY_ATTEMPT_CHILD_KEY", + "_LLM_PATH_SUFFIXES", +] diff --git a/packages/opentelemetry-instrumentation-openai/opentelemetry/instrumentation/openai/shared/chat_wrappers.py b/packages/opentelemetry-instrumentation-openai/opentelemetry/instrumentation/openai/shared/chat_wrappers.py index 6f00a08a21..a9938afe0d 100644 --- a/packages/opentelemetry-instrumentation-openai/opentelemetry/instrumentation/openai/shared/chat_wrappers.py +++ b/packages/opentelemetry-instrumentation-openai/opentelemetry/instrumentation/openai/shared/chat_wrappers.py @@ -58,6 +58,9 @@ LLMRequestTypeValues, SpanAttributes, ) +from opentelemetry.instrumentation.openai.retry_handler import ( + OPENAI_DIRECT_RETRY_PARENT_ACTIVE_KEY, +) from opentelemetry.trace import SpanKind, Tracer from opentelemetry import trace from opentelemetry.trace.status import Status, StatusCode @@ -104,12 +107,22 @@ def chat_wrapper( run_async(_handle_request(span, kwargs, instance)) try: start_time = time.time() - token = context_api.attach( - context_api.set_value( - SUPPRESS_LANGUAGE_MODEL_INSTRUMENTATION_KEY, - True, - ) + # ST-10.4 (review-driven 2026-05-16): set + # OPENAI_DIRECT_RETRY_PARENT_ACTIVE_KEY alongside the + # existing SUPPRESS_LANGUAGE_MODEL_INSTRUMENTATION_KEY so + # the FortifyRoot retry handler can distinguish this + # "openai wrapper's own internal HTTP send" case from + # external user/framework suppression and emit + # retry_attempt spans under the ``openai.chat`` parent. + # The SUPPRESS key still protects against OTHER LLM + # instrumentors double-counting this call. + attempt_ctx = context_api.set_value( + SUPPRESS_LANGUAGE_MODEL_INSTRUMENTATION_KEY, True, + ) + attempt_ctx = context_api.set_value( + OPENAI_DIRECT_RETRY_PARENT_ACTIVE_KEY, True, attempt_ctx, ) + token = context_api.attach(attempt_ctx) try: response = wrapped(*args, **kwargs) finally: @@ -214,12 +227,15 @@ async def achat_wrapper( try: start_time = time.time() - token = context_api.attach( - context_api.set_value( - SUPPRESS_LANGUAGE_MODEL_INSTRUMENTATION_KEY, - True, - ) + # ST-10.4 (review-driven 2026-05-16): see sync chat_wrapper + # above for the rationale on OPENAI_DIRECT_RETRY_PARENT_ACTIVE_KEY. + attempt_ctx = context_api.set_value( + SUPPRESS_LANGUAGE_MODEL_INSTRUMENTATION_KEY, True, + ) + attempt_ctx = context_api.set_value( + OPENAI_DIRECT_RETRY_PARENT_ACTIVE_KEY, True, attempt_ctx, ) + token = context_api.attach(attempt_ctx) try: response = await wrapped(*args, **kwargs) finally: diff --git a/packages/opentelemetry-instrumentation-openai/opentelemetry/instrumentation/openai/v1/__init__.py b/packages/opentelemetry-instrumentation-openai/opentelemetry/instrumentation/openai/v1/__init__.py index 22ddfb9fa9..f1ad96344e 100644 --- a/packages/opentelemetry-instrumentation-openai/opentelemetry/instrumentation/openai/v1/__init__.py +++ b/packages/opentelemetry-instrumentation-openai/opentelemetry/instrumentation/openai/v1/__init__.py @@ -42,6 +42,10 @@ from opentelemetry.instrumentation.openai.v1.realtime_wrappers import ( realtime_connect_wrapper, ) +from opentelemetry.instrumentation.openai.retry_handler import ( + instrument_retry_emitter, + uninstrument_retry_emitter, +) from opentelemetry.instrumentation.openai.version import __version__ from opentelemetry.metrics import get_meter @@ -350,7 +354,19 @@ def _instrument(self, **kwargs): realtime_connect_wrapper(tracer), ) + # ST-10.4: per-attempt retry_attempt emission via private + # ``openai._base_client`` httpx wrapper classes. Guarded against + # missing private symbols (logs warning + skips emission). Pass + # the same tracer_provider the rest of the instrumentor uses so + # retry_attempt spans land in the same exporter as the openai + # logical span (e.g. ``openai.chat``) — without this, a + # consumer passing an explicit provider would get the parent + # span on their provider but the retry_attempt span lost to the + # global no-op tracer. + instrument_retry_emitter(tracer_provider=tracer_provider) + def _uninstrument(self, **kwargs): + uninstrument_retry_emitter() # ST-10.4 symmetry unwrap_dotted_method("openai.resources.chat.completions", "Completions.create") unwrap_dotted_method("openai.resources.completions", "Completions.create") unwrap_dotted_method("openai.resources.embeddings", "Embeddings.create") diff --git a/packages/opentelemetry-instrumentation-openai/tests/conftest.py b/packages/opentelemetry-instrumentation-openai/tests/conftest.py index 347a2691d0..9f06622ad0 100644 --- a/packages/opentelemetry-instrumentation-openai/tests/conftest.py +++ b/packages/opentelemetry-instrumentation-openai/tests/conftest.py @@ -27,6 +27,35 @@ pytest_plugins = [] +class _NoFortifyRootSpanExporter(InMemorySpanExporter): + """ST-10.4: filter out any span whose name starts with + ``fortifyroot.`` from the upstream-test span exporter. + + ST-10.4 added per-attempt sibling spans (e.g. + ``fortifyroot.openai.retry_attempt``) under every openai logical + call. Upstream/legacy OpenAI tests assert exact span-name lists + (e.g. ``[span.name for span in spans] == ["openai.chat"]``); without + a filter, every such assertion would now fail because the + retry_attempt sibling is also exported. + + Mirrors the same pattern used by the LangChain test conftest from + the 2026-05-15 CI-hardening addendum. ST-10.4 unit tests in + ``tests/test_retry_attempt_emission.py`` use their own ``fresh_tracer`` + fixture (not this one), so they continue to see retry_attempt + spans and aren't affected by the filter. + """ + + def get_finished_spans(self): # type: ignore[override] + # Filter by role rather than name prefix so legitimate + # fortifyroot.*.safety / .llm_wrapper / .has_native_otel_child + # spans remain visible to tests that inspect them. Only + # ST-10.4 retry_attempt siblings carry role=retry_attempt. + return tuple( + s for s in super().get_finished_spans() + if (s.attributes or {}).get("fortifyroot.span.role") != "retry_attempt" + ) + + @pytest.fixture(autouse=True) def environment(): if not os.getenv("OPENAI_API_KEY"): @@ -89,7 +118,7 @@ def async_vllm_openai_client(): @pytest.fixture(scope="session", name="span_exporter") def fixture_span_exporter(): - exporter = InMemorySpanExporter() + exporter = _NoFortifyRootSpanExporter() yield exporter diff --git a/packages/opentelemetry-instrumentation-openai/tests/test_retry_attempt_emission.py b/packages/opentelemetry-instrumentation-openai/tests/test_retry_attempt_emission.py new file mode 100644 index 0000000000..6333e81254 --- /dev/null +++ b/packages/opentelemetry-instrumentation-openai/tests/test_retry_attempt_emission.py @@ -0,0 +1,936 @@ +"""Tests for ST-10.4 OpenAI direct-SDK retry-attempt emission. + +Covers (per RETRY_LOOP.md §4.4 OpenAI row + §4.7 suppression): + - Instrumentor symmetry: install/uninstall flips state cleanly and is + idempotent. + - Single-attempt happy path: ONE retry_attempt span under the active + parent; parent gets has_retry_attempt_child=true; span carries + role / gen_ai.system=openai / gen_ai.request.model / http.status_code. + - Multi-attempt retry path: N siblings under one parent (the + structural shape RetryDetectorProc relies on). + - Error-attempt: status=ERROR, http.status_code set, error.type set. + - No-parent guard: skip when no ambient OTel parent. + - §4.7 suppression — BOTH paths: + * OTel context SUPPRESS_LANGUAGE_MODEL_INSTRUMENTATION_KEY active → skip + * framework-attempt registry says owned (LiteLLM/LangChain/LlamaIndex + wrapping in flight) → skip + - §4.4.1 endpoint allow-list: non-LLM SDK traffic (e.g. /v1/models) + does NOT emit a retry_attempt span. + - Private-symbol missing guard: import-time wrappability check — + if the private symbol is absent, log warning + skip, do NOT crash. + +These are UNIT tests — we drive ``_sync_send_wrapper`` / +``_async_send_wrapper`` directly with fake httpx Request / Response +objects so the test doesn't depend on real openai SDK behaviour or +network calls. +""" + +from __future__ import annotations + +import asyncio +import threading +from types import SimpleNamespace +from typing import Any, Optional +from unittest.mock import patch + +import pytest +from opentelemetry import context as context_api +from opentelemetry import trace +from opentelemetry.instrumentation.fortifyroot import ( + is_framework_owned, + register_framework_attempt, + retry_registry, + unregister_framework_attempt, +) +from opentelemetry.instrumentation.openai.retry_handler import ( + _FR_HAS_RETRY_ATTEMPT_CHILD_KEY, + _FR_RETRY_ATTEMPT_SPAN_NAME, + OPENAI_DIRECT_RETRY_PARENT_ACTIVE_KEY, + _async_send_wrapper, + _has_wrappable_symbol, + _is_installed_for_test, + _sync_send_wrapper, + instrument_retry_emitter, + uninstrument_retry_emitter, +) +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import ( + InMemorySpanExporter, +) +from opentelemetry.semconv_ai import SUPPRESS_LANGUAGE_MODEL_INSTRUMENTATION_KEY + + +# --------------------------------------------------------------------------- +# Test infrastructure. +# --------------------------------------------------------------------------- + +@pytest.fixture +def fresh_tracer(): + """Fresh TracerProvider + in-memory exporter as the global provider, + so the retry handler's ``trace.get_tracer(...)`` lookups route here. + """ + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + try: + trace.set_tracer_provider(provider) + except Exception: + pass + current = trace.get_tracer_provider() + if current is not provider: + try: + current.add_span_processor(SimpleSpanProcessor(exporter)) + except Exception: + pass + yield trace.get_tracer("test"), exporter, current + + +@pytest.fixture(autouse=True) +def reset_registry_and_state(): + """Each test starts with empty registry + emitter uninstalled.""" + retry_registry._reset_for_test() + # Defensive uninstall in case a prior test left it installed. + try: + uninstrument_retry_emitter() + except Exception: + pass + yield + retry_registry._reset_for_test() + try: + uninstrument_retry_emitter() + except Exception: + pass + + +def _make_request(path: str = "/v1/chat/completions", host: str = "api.openai.com", + port: Optional[int] = None, scheme: str = "https", + model: Optional[str] = "gpt-4o-mini") -> SimpleNamespace: + """Build a fake ``httpx.Request``-shaped object for the wrapper. + + The wrapper reads .url.path / .url.host / .url.port / .url.scheme + plus .content (JSON-encoded body) — all we need. + """ + import json + body = {"model": model, "messages": [{"role": "user", "content": "hi"}]} if model else {} + content = json.dumps(body).encode("utf-8") + url = SimpleNamespace(path=path, host=host, port=port, scheme=scheme) + return SimpleNamespace(url=url, content=content) + + +def _make_response(status_code: int = 200, request_id: str = "req-abc", + body: Optional[dict] = None) -> SimpleNamespace: + """Build a fake ``httpx.Response``-shaped object. + + The retry handler's non-streaming body-parse path calls + ``response.json()`` directly, so we expose ``json()`` as a method + returning the supplied dict. The default body matches the OpenAI + chat-completion schema with usage so the success path's usage + extraction has something to read. + """ + if body is None: + body = { + "id": f"chatcmpl-{request_id}", + "model": "gpt-4o-mini-2024-07-18", + "usage": {"prompt_tokens": 7, "completion_tokens": 3, "total_tokens": 10}, + } + headers = {"x-request-id": request_id, "openai-request-id": request_id} + return SimpleNamespace( + status_code=status_code, + headers=headers, + json=lambda: body, + ) + + +def _retry_spans(exporter: InMemorySpanExporter): + return [s for s in exporter.get_finished_spans() + if s.name == _FR_RETRY_ATTEMPT_SPAN_NAME] + + +# --------------------------------------------------------------------------- +# Instrumentor symmetry. +# --------------------------------------------------------------------------- + +def test_install_then_uninstall_is_idempotent_and_symmetric(): + """install → installed; uninstall → not installed; double-install + and double-uninstall are no-ops.""" + assert not _is_installed_for_test() + + instrument_retry_emitter() + assert _is_installed_for_test() + + # Idempotent install: state stays installed; no exception. + instrument_retry_emitter() + assert _is_installed_for_test() + + uninstrument_retry_emitter() + assert not _is_installed_for_test() + + # Idempotent uninstall. + uninstrument_retry_emitter() + assert not _is_installed_for_test() + + +def test_install_actually_wraps_the_openai_private_send(): + """REGRESSION GUARD: install must put a wrapt wrapper on + ``openai._base_client.SyncHttpxClientWrapper.send``. Without this + check the install path could silently no-op and leave retry_attempt + emission disabled in production. + """ + pytest.importorskip("openai") + from openai import _base_client as base + + if not hasattr(base, "SyncHttpxClientWrapper") or not hasattr( + base.SyncHttpxClientWrapper, "send" + ): + pytest.skip("openai._base_client.SyncHttpxClientWrapper.send not present") + + instrument_retry_emitter() + try: + send = base.SyncHttpxClientWrapper.send + # wrapt sets __wrapped__ on the wrapped descriptor. + assert hasattr(send, "__wrapped__"), ( + "after instrument, SyncHttpxClientWrapper.send must be a wrapt wrapper" + ) + finally: + uninstrument_retry_emitter() + + # After uninstall, the wrap must be gone. + send_after = base.SyncHttpxClientWrapper.send + assert not hasattr(send_after, "__wrapped__"), ( + "after uninstall, SyncHttpxClientWrapper.send must be the original (no __wrapped__)" + ) + + +# --------------------------------------------------------------------------- +# Private-symbol-missing guard. +# --------------------------------------------------------------------------- + +def test_missing_private_symbol_logs_warning_and_does_not_crash(caplog): + """If ``openai._base_client.SyncHttpxClientWrapper.send`` is missing, + install must log a warning and skip emission for that variant — + NOT raise. Normal openai instrumentation continues unaffected. + """ + import opentelemetry.instrumentation.openai.retry_handler as rh + + with patch.object(rh, "_has_wrappable_symbol", return_value=False): + with caplog.at_level("WARNING"): + # Must not raise. + instrument_retry_emitter() + # State still flips to installed (a no-op install is still an + # install — uninstall is symmetric). + assert _is_installed_for_test() + # Warning emitted for at least one of the two variants. + msgs = " ".join(r.getMessage() for r in caplog.records) + assert "missing/incompatible" in msgs + uninstrument_retry_emitter() + + +def test_has_wrappable_symbol_returns_true_for_real_openai(): + """Sanity: the import-time wrappability check works against the + real installed openai SDK. Catches the case where the private API + is renamed between SDK versions (the regression this guard + protects against).""" + pytest.importorskip("openai") + assert _has_wrappable_symbol("openai._base_client", "SyncHttpxClientWrapper", "send") + assert _has_wrappable_symbol("openai._base_client", "AsyncHttpxClientWrapper", "send") + + +def test_has_wrappable_symbol_returns_false_for_unknown(): + """A bogus class name returns False (does not raise).""" + assert not _has_wrappable_symbol("openai._base_client", "NoSuchClass", "send") + assert not _has_wrappable_symbol("openai._base_client", "SyncHttpxClientWrapper", "no_such_method") + + +# --------------------------------------------------------------------------- +# Single-attempt happy path. +# --------------------------------------------------------------------------- + +def test_single_attempt_emits_one_span_with_marker(fresh_tracer): + tracer, exporter, _ = fresh_tracer + + request = _make_request(model="gpt-4o-mini", path="/v1/chat/completions") + response = _make_response(status_code=200, request_id="req-1") + + parent = tracer.start_span("openai.chat") # mimics outer span + with trace.use_span(parent, end_on_exit=False): + # _sync_send_wrapper(wrapped, instance, args, kwargs) + result = _sync_send_wrapper(lambda *a, **kw: response, None, (request,), {}) + parent.end() + + assert result is response + + spans = exporter.get_finished_spans() + parent_exported = next(s for s in spans if s.name == "openai.chat") + retry_spans = _retry_spans(exporter) + assert len(retry_spans) == 1 + rs = retry_spans[0] + assert rs.parent.span_id == parent_exported.context.span_id, ( + "retry_attempt must be a child of the active parent" + ) + assert rs.attributes.get("fortifyroot.span.role") == "retry_attempt" + assert rs.attributes.get("gen_ai.system") == "openai" + assert rs.attributes.get("gen_ai.request.model") == "gpt-4o-mini" + assert rs.attributes.get("gen_ai.operation.name") == "chat" + assert rs.attributes.get("http.status_code") == 200 + # Response-body-derived attrs (overrides the header-only response id + # with the body's chatcmpl-* id; adds response.model and usage tokens). + assert rs.attributes.get("gen_ai.response.id") == "chatcmpl-req-1" + assert rs.attributes.get("gen_ai.response.model") == "gpt-4o-mini-2024-07-18" + assert rs.attributes.get("gen_ai.usage.input_tokens") == 7 + assert rs.attributes.get("gen_ai.usage.output_tokens") == 3 + assert rs.attributes.get("server.address") == "api.openai.com" + assert rs.attributes.get("server.port") == 443 + assert parent_exported.attributes.get(_FR_HAS_RETRY_ATTEMPT_CHILD_KEY) is True + + +# --------------------------------------------------------------------------- +# Multi-attempt retry path — siblings under one parent. +# --------------------------------------------------------------------------- + +def test_three_attempts_share_parent(fresh_tracer): + """429 → 429 → 200 chain produces 3 sibling retry_attempt spans + under one parent. RetryDetectorProc relies on this shape.""" + tracer, exporter, _ = fresh_tracer + + request = _make_request(model="gpt-4o-mini", path="/v1/chat/completions") + resp_429 = _make_response(status_code=429, request_id="req-429a") + resp_200 = _make_response(status_code=200, request_id="req-200") + + parent = tracer.start_span("openai.chat") + with trace.use_span(parent, end_on_exit=False): + # attempt 1 — 429 + _sync_send_wrapper(lambda *a, **kw: resp_429, None, (request,), {}) + # attempt 2 — 429 + _sync_send_wrapper(lambda *a, **kw: resp_429, None, (request,), {}) + # attempt 3 — 200 + _sync_send_wrapper(lambda *a, **kw: resp_200, None, (request,), {}) + parent.end() + + spans = exporter.get_finished_spans() + parent_exported = next(s for s in spans if s.name == "openai.chat") + retry_spans = _retry_spans(exporter) + assert len(retry_spans) == 3, f"expected 3 retry_attempts, got {len(retry_spans)}" + + parent_ids = {s.parent.span_id for s in retry_spans} + assert parent_ids == {parent_exported.context.span_id}, ( + f"all 3 retry_attempts must be siblings under one parent; got {parent_ids}" + ) + + from opentelemetry.trace import StatusCode + error_count = sum(1 for s in retry_spans if s.status.status_code == StatusCode.ERROR) + ok_count = sum(1 for s in retry_spans if s.status.status_code == StatusCode.OK) + assert error_count == 2 and ok_count == 1, ( + f"expected 2 ERROR + 1 OK, got error={error_count}, ok={ok_count}" + ) + + for s in retry_spans: + if s.status.status_code == StatusCode.ERROR: + assert s.attributes.get("http.status_code") == 429 + assert s.attributes.get("error.type") == "openai.RateLimitError" + + +# --------------------------------------------------------------------------- +# Error attempt: exception path. +# --------------------------------------------------------------------------- + +def test_exception_path_records_error_and_reraises(fresh_tracer): + """If wrapped(send) raises, the wrapper finalises the span as + ERROR with error.type set, then re-raises the original exception.""" + tracer, exporter, _ = fresh_tracer + + request = _make_request(model="gpt-4o-mini") + + class TimeoutLikeError(Exception): + pass + + err = TimeoutLikeError("connect timeout") + + def raising(*a, **kw): + raise err + + parent = tracer.start_span("openai.chat") + with trace.use_span(parent, end_on_exit=False): + with pytest.raises(TimeoutLikeError): + _sync_send_wrapper(raising, None, (request,), {}) + parent.end() + + from opentelemetry.trace import StatusCode + retry_spans = _retry_spans(exporter) + assert len(retry_spans) == 1 + rs = retry_spans[0] + assert rs.status.status_code == StatusCode.ERROR + # error.type carries the namespaced exception class. + assert "TimeoutLikeError" in (rs.attributes.get("error.type") or "") + + +def test_http_status_code_on_response_exception(fresh_tracer): + """If the SDK raises something that carries .status_code (e.g. + APIStatusError-shaped), the wrapper extracts and records it.""" + tracer, exporter, _ = fresh_tracer + + request = _make_request(model="gpt-4o-mini") + + class FakeStatusError(Exception): + def __init__(self, status_code): + self.status_code = status_code + super().__init__(f"http {status_code}") + + err = FakeStatusError(503) + + parent = tracer.start_span("openai.chat") + with trace.use_span(parent, end_on_exit=False): + with pytest.raises(FakeStatusError): + _sync_send_wrapper(lambda *a, **kw: (_ for _ in ()).throw(err), None, (request,), {}) + parent.end() + + retry_spans = _retry_spans(exporter) + assert len(retry_spans) == 1 + rs = retry_spans[0] + assert rs.attributes.get("http.status_code") == 503 + + +# --------------------------------------------------------------------------- +# Marker timing (§4.5). +# --------------------------------------------------------------------------- + +def test_marker_set_AFTER_first_attempt_not_at_parent_creation(fresh_tracer): + tracer, exporter, _ = fresh_tracer + request = _make_request() + + parent = tracer.start_span("openai.chat") + # Before any retry_attempt: no marker. + assert _FR_HAS_RETRY_ATTEMPT_CHILD_KEY not in dict(parent.attributes or {}) + + with trace.use_span(parent, end_on_exit=False): + _sync_send_wrapper(lambda *a, **kw: _make_response(200), None, (request,), {}) + + # After first attempt: marker present on the still-open parent. + assert dict(parent.attributes or {}).get(_FR_HAS_RETRY_ATTEMPT_CHILD_KEY) is True + parent.end() + + +def test_marker_NOT_set_when_no_attempts_fire(fresh_tracer): + tracer, exporter, _ = fresh_tracer + parent = tracer.start_span("openai.chat") + parent.end() + parent_exported = next( + s for s in exporter.get_finished_spans() if s.name == "openai.chat" + ) + assert _FR_HAS_RETRY_ATTEMPT_CHILD_KEY not in (parent_exported.attributes or {}) + + +# --------------------------------------------------------------------------- +# No-parent guard. +# --------------------------------------------------------------------------- + +def test_no_parent_does_not_emit(fresh_tracer): + tracer, exporter, _ = fresh_tracer + request = _make_request() + # No active span at all. + result = _sync_send_wrapper(lambda *a, **kw: _make_response(200), None, (request,), {}) + assert result.status_code == 200 + assert len(_retry_spans(exporter)) == 0 + + +# --------------------------------------------------------------------------- +# §4.7 suppression — context API. +# --------------------------------------------------------------------------- + +def test_context_suppression_skips_emission(fresh_tracer): + """If SUPPRESS_LANGUAGE_MODEL_INSTRUMENTATION_KEY is True in the + OTel context, no retry_attempt span emits — wrapped(send) still + runs normally.""" + tracer, exporter, _ = fresh_tracer + request = _make_request() + + parent = tracer.start_span("openai.chat") + suppress_token = context_api.attach( + context_api.set_value(SUPPRESS_LANGUAGE_MODEL_INSTRUMENTATION_KEY, True) + ) + try: + with trace.use_span(parent, end_on_exit=False): + result = _sync_send_wrapper(lambda *a, **kw: _make_response(200), None, (request,), {}) + finally: + context_api.detach(suppress_token) + parent.end() + + assert result.status_code == 200 + assert len(_retry_spans(exporter)) == 0, ( + "SUPPRESS_LANGUAGE_MODEL_INSTRUMENTATION_KEY active → no retry_attempt" + ) + + +# --------------------------------------------------------------------------- +# §4.7 suppression — framework-attempt registry. +# --------------------------------------------------------------------------- + +def test_framework_registry_suppression_skips_emission(fresh_tracer): + """If is_framework_owned() is True (a framework wrapper has an + in-flight attempt on this thread), the direct-SDK wrapper must + SKIP emission. This is the §4.7 invariant that prevents double- + emission when both framework + direct-SDK wrappers are active.""" + tracer, exporter, _ = fresh_tracer + request = _make_request() + + parent = tracer.start_span("openai.chat") + token = register_framework_attempt() + assert is_framework_owned() + try: + with trace.use_span(parent, end_on_exit=False): + result = _sync_send_wrapper(lambda *a, **kw: _make_response(200), None, (request,), {}) + finally: + unregister_framework_attempt(token) + parent.end() + + assert result.status_code == 200 + assert len(_retry_spans(exporter)) == 0, ( + "is_framework_owned() active → no retry_attempt" + ) + + +# --------------------------------------------------------------------------- +# Framework registry token: registered DURING our attempt, released AFTER. +# --------------------------------------------------------------------------- + +def test_direct_sdk_wrapper_does_not_self_register_in_framework_registry(fresh_tracer): + """REGRESSION GUARD (review-driven fix 2026-05-13): direct-SDK + wrappers MUST NOT register tokens in the §4.7.1 framework registry. + The registry's contract reserves registration for FRAMEWORK wrappers + (LiteLLM / LangChain / LlamaIndex); direct-SDK wrappers only CONSULT + via ``is_framework_owned()``. + + If this contract is violated, two concurrent asyncio tasks on the + same OS thread will see each other's token via the thread-keyed + registry, and the second task's emission gets suppressed — + silently dropping a retry_attempt span. See + ``test_two_concurrent_async_sends_each_emit_a_retry_attempt`` for + the end-to-end repro.""" + tracer, _, _ = fresh_tracer + request = _make_request() + + seen_during: list[bool] = [] + + def wrapped(*a, **kw): + seen_during.append(is_framework_owned()) + return _make_response(200) + + parent = tracer.start_span("openai.chat") + with trace.use_span(parent, end_on_exit=False): + assert not is_framework_owned() + _sync_send_wrapper(wrapped, None, (request,), {}) + assert not is_framework_owned(), "after attempt → still not registered" + parent.end() + + assert seen_during == [False], ( + "the direct-SDK wrapper MUST NOT self-register a framework " + "token; is_framework_owned() must remain False during the wrapped " + "send. Got True → wrapper is incorrectly registering." + ) + + +def test_two_concurrent_async_sends_each_emit_a_retry_attempt(fresh_tracer): + """REGRESSION GUARD (review-driven fix 2026-05-13): two asyncio + tasks sharing one OS thread MUST each emit their own retry_attempt + span. Pre-fix, the first task's self-registered framework token + would suppress the second task via the thread-keyed registry — + silently dropping one of the spans.""" + tracer, exporter, _ = fresh_tracer + + request = _make_request(path="/v1/chat/completions", model="gpt-4o-mini") + parent = tracer.start_span("openai.chat") + + async def one_send(): + # Yield to the event loop after entering the wrap, so the two + # tasks' wrap bodies interleave (otherwise they could run + # sequentially and the bug would be hidden). + async def wrapped(*a, **kw): + await asyncio.sleep(0) + return _make_response(200) + with trace.use_span(parent, end_on_exit=False): + return await _async_send_wrapper(wrapped, None, (request,), {}) + + async def run(): + return await asyncio.gather(one_send(), one_send()) + + loop = asyncio.new_event_loop() + try: + results = loop.run_until_complete(run()) + finally: + loop.close() + parent.end() + + assert all(r.status_code == 200 for r in results) + retry_spans = _retry_spans(exporter) + assert len(retry_spans) == 2, ( + f"expected 2 retry_attempt spans (one per concurrent task); got {len(retry_spans)}. " + f"This is the §4.7 self-registration regression — the second task is being " + f"suppressed by the first task's lingering framework token." + ) + + +# --------------------------------------------------------------------------- +# §4.4.1 endpoint allow-list. +# --------------------------------------------------------------------------- + +@pytest.mark.parametrize("path", [ + "/v1/models", + "/v1/files", + "/v1/files/file-123", + "/v1/fine_tuning/jobs", + "/v1/threads/thread-abc", + "/", + "", +]) +def test_non_llm_endpoints_do_not_emit(fresh_tracer, path): + """Non-LLM SDK traffic (model listing, file ops, auth refresh, etc.) + MUST NOT emit retry_attempt spans — per §4.4.1 allow-list.""" + tracer, exporter, _ = fresh_tracer + request = _make_request(path=path, model=None) + parent = tracer.start_span("openai.chat") + with trace.use_span(parent, end_on_exit=False): + _sync_send_wrapper(lambda *a, **kw: _make_response(200), None, (request,), {}) + parent.end() + assert len(_retry_spans(exporter)) == 0, ( + f"non-LLM endpoint {path!r} must not emit retry_attempt" + ) + + +@pytest.mark.parametrize("path,expected_op", [ + ("/v1/chat/completions", "chat"), + ("/v1/completions", "text_completion"), + ("/v1/embeddings", "embeddings"), + ("/v1/responses", "chat"), + ("/v1/messages", "chat"), # Azure-OpenAI messages parity layer +]) +def test_llm_endpoints_emit_with_correct_operation(fresh_tracer, path, expected_op): + tracer, exporter, _ = fresh_tracer + request = _make_request(path=path, model="gpt-4o-mini") + parent = tracer.start_span("openai.chat") + with trace.use_span(parent, end_on_exit=False): + _sync_send_wrapper(lambda *a, **kw: _make_response(200), None, (request,), {}) + parent.end() + retry_spans = _retry_spans(exporter) + assert len(retry_spans) == 1 + assert retry_spans[0].attributes.get("gen_ai.operation.name") == expected_op + + +# --------------------------------------------------------------------------- +# §4.5-driven usage-extraction policy (review-driven fix 2026-05-13). +# --------------------------------------------------------------------------- + +def test_non_streaming_response_body_populates_usage_id_and_model(fresh_tracer): + """REGRESSION GUARD: backend dedup makes the retry_attempt span the + canonical LLMUsageEvent even for single-attempt calls. If the wrap + omits usage tokens, the canonical event has zero cost. For + non-streaming responses we MUST parse the body and copy usage onto + the retry_attempt span. See review finding C2 in the 2026-05-13 + addendum. + """ + tracer, exporter, _ = fresh_tracer + request = _make_request(model="gpt-4o-mini") + response = _make_response( + status_code=200, + request_id="r-1", + body={ + "id": "chatcmpl-XYZ", + "model": "gpt-4o-mini-2024-07-18", + "usage": {"prompt_tokens": 42, "completion_tokens": 11, "total_tokens": 53}, + }, + ) + parent = tracer.start_span("openai.chat") + with trace.use_span(parent, end_on_exit=False): + _sync_send_wrapper(lambda *a, **kw: response, None, (request,), {}) + parent.end() + + retry_spans = _retry_spans(exporter) + assert len(retry_spans) == 1 + rs = retry_spans[0] + assert rs.attributes.get("gen_ai.usage.input_tokens") == 42 + assert rs.attributes.get("gen_ai.usage.output_tokens") == 11 + assert rs.attributes.get("gen_ai.response.id") == "chatcmpl-XYZ" + assert rs.attributes.get("gen_ai.response.model") == "gpt-4o-mini-2024-07-18" + + +def test_streaming_request_skips_emission_entirely(fresh_tracer): + """ST-10.4 (review-driven 2026-05-17): when ``stream=True`` is + passed to send, the wrap SKIPS retry_attempt emission entirely — + no span is created, and ``response.json()`` is never called. + + Rationale: streaming retry_attempts cannot carry token usage at + attempt-end (SSE stream consumption would break the SDK), but + backend §4.5 dedup would promote them to canonical + LLMUsageEvents → zero-token events → fr-system-tests + ``prompt_tokens > 0`` assertions fail. The parent + ``openai.chat`` span (which gets full usage from ``ChatStream``'s + stream-completion callback) stays canonical. Streaming + retry-loop detection is the deferred follow-up + ``ST-10.4-FOLLOWUP-streaming-usage``. + """ + tracer, exporter, _ = fresh_tracer + request = _make_request(model="gpt-4o-mini") + + json_called = {"n": 0} + + def explode(): + json_called["n"] += 1 + raise RuntimeError("body must not be read on streaming response") + + headers = {"x-request-id": "req-stream-1"} + response = SimpleNamespace(status_code=200, headers=headers, json=explode) + + parent = tracer.start_span("openai.chat") + with trace.use_span(parent, end_on_exit=False): + result = _sync_send_wrapper( + lambda *a, **kw: response, + None, + (request,), + {"stream": True}, + ) + parent.end() + + # wrap returns the response unchanged. + assert result is response + # response.json() never invoked. + assert json_called["n"] == 0 + # No retry_attempt span emitted at all. + retry_spans = _retry_spans(exporter) + assert len(retry_spans) == 0, ( + f"streaming attempts MUST skip retry_attempt emission entirely; " + f"got {len(retry_spans)} span(s). The parent span stays the " + f"canonical LLM event." + ) + # Parent MUST NOT receive the has_retry_attempt_child marker + # either (no child was emitted). + parent_exported = next( + s for s in exporter.get_finished_spans() if s.name == "openai.chat" + ) + assert _FR_HAS_RETRY_ATTEMPT_CHILD_KEY not in (parent_exported.attributes or {}) + + +def test_non_2xx_response_with_usage_in_body_populates_usage_tokens(fresh_tracer): + """REGRESSION GUARD (review-driven follow-up 2026-05-13): per + RETRY_LOOP.md §4.4 token-usage rule (around line 164), wrappers + MUST extract usage from the response body whenever it's present, + REGARDLESS of whether the attempt succeeded. Some failures (e.g. + context-length-exceeded errors) consume tokens and the provider + returns usage in the error body. The retry_attempt span MUST carry + those tokens so backend §4.5 dedup doesn't promote a zero-cost + canonical event. + """ + tracer, exporter, _ = fresh_tracer + request = _make_request(model="gpt-4o-mini") + # Simulated context-length-exceeded shape: 400 + usage in body. + response = _make_response( + status_code=400, + request_id="r-ctx-exceeded", + body={ + "error": { + "type": "invalid_request_error", + "code": "context_length_exceeded", + "message": "This model's max context...", + }, + "id": "chatcmpl-ERR", + "model": "gpt-4o-mini-2024-07-18", + "usage": {"prompt_tokens": 128000, "completion_tokens": 0, "total_tokens": 128000}, + }, + ) + parent = tracer.start_span("openai.chat") + with trace.use_span(parent, end_on_exit=False): + _sync_send_wrapper(lambda *a, **kw: response, None, (request,), {}) + parent.end() + + from opentelemetry.trace import StatusCode + retry_spans = _retry_spans(exporter) + assert len(retry_spans) == 1 + rs = retry_spans[0] + assert rs.status.status_code == StatusCode.ERROR + assert rs.attributes.get("http.status_code") == 400 + assert rs.attributes.get("error.type") == "openai.HTTPStatusError" + # Usage was in the error body → must be on the span (the load-bearing assertion). + assert rs.attributes.get("gen_ai.usage.input_tokens") == 128000 + assert rs.attributes.get("gen_ai.usage.output_tokens") == 0 + assert rs.attributes.get("gen_ai.response.id") == "chatcmpl-ERR" + + +def test_non_2xx_response_without_usage_in_body_omits_usage_tokens(fresh_tracer): + """Counterpart to the above: when the error body has no ``usage`` + field, the wrap must NOT set usage attrs (per §4.4 rule: omit when + unknown; setting 0 risks under-counting cost on other paths).""" + tracer, exporter, _ = fresh_tracer + request = _make_request(model="gpt-4o-mini") + response = _make_response( + status_code=500, + request_id="r-server-err", + body={"error": {"type": "internal_server_error", "message": "boom"}}, + ) + parent = tracer.start_span("openai.chat") + with trace.use_span(parent, end_on_exit=False): + _sync_send_wrapper(lambda *a, **kw: response, None, (request,), {}) + parent.end() + + from opentelemetry.trace import StatusCode + retry_spans = _retry_spans(exporter) + assert len(retry_spans) == 1 + rs = retry_spans[0] + assert rs.status.status_code == StatusCode.ERROR + assert rs.attributes.get("http.status_code") == 500 + assert rs.attributes.get("error.type") == "openai.InternalServerError" + assert rs.attributes.get("gen_ai.usage.input_tokens") is None + assert rs.attributes.get("gen_ai.usage.output_tokens") is None + + +def test_non_streaming_response_with_malformed_body_does_not_crash(fresh_tracer): + """If ``response.json()`` raises (non-JSON body, e.g. error HTML), + the wrap must record the basic status attrs without raising and + without leaving the span in an unended state.""" + tracer, exporter, _ = fresh_tracer + request = _make_request(model="gpt-4o-mini") + + def explode(): + raise ValueError("malformed JSON") + + headers = {"x-request-id": "req-bad-body"} + response = SimpleNamespace(status_code=200, headers=headers, json=explode) + + parent = tracer.start_span("openai.chat") + with trace.use_span(parent, end_on_exit=False): + # MUST NOT raise. + _sync_send_wrapper(lambda *a, **kw: response, None, (request,), {}) + parent.end() + + retry_spans = _retry_spans(exporter) + assert len(retry_spans) == 1 + from opentelemetry.trace import StatusCode + assert retry_spans[0].status.status_code == StatusCode.OK + assert retry_spans[0].attributes.get("http.status_code") == 200 + # Header-based response id survives even when body parsing fails. + assert retry_spans[0].attributes.get("gen_ai.response.id") == "req-bad-body" + + +# --------------------------------------------------------------------------- +# Async path mirrors sync. +# --------------------------------------------------------------------------- + +def test_external_suppression_with_no_override_skips_emission(fresh_tracer): + """REGRESSION GUARD (review-driven 2026-05-16, Issue 3B counter-proof): + when SUPPRESS_LANGUAGE_MODEL_INSTRUMENTATION_KEY is set in the OTel + context WITHOUT the OpenAI override key, the retry handler MUST + suppress emission. This is the user-explicit "disable LLM + instrumentation for this scope" path, distinct from the + openai-wrapper-internal use of the same key.""" + tracer, exporter, _ = fresh_tracer + request = _make_request() + parent = tracer.start_span("some.workflow") + suppress_token = context_api.attach( + context_api.set_value(SUPPRESS_LANGUAGE_MODEL_INSTRUMENTATION_KEY, True) + ) + try: + with trace.use_span(parent, end_on_exit=False): + _sync_send_wrapper(lambda *a, **kw: _make_response(200), None, (request,), {}) + finally: + context_api.detach(suppress_token) + parent.end() + + assert len(_retry_spans(exporter)) == 0, ( + "external SUPPRESS_LANGUAGE_MODEL_INSTRUMENTATION_KEY without the " + "openai override key MUST skip emission" + ) + + +def test_openai_wrapper_override_key_unblocks_emission_under_suppression(fresh_tracer): + """REGRESSION GUARD (review-driven 2026-05-16, Issue 3B): + the openai chat_wrapper sets SUPPRESS_LANGUAGE_MODEL_INSTRUMENTATION_KEY + around its wrapped SDK call (to protect against OTHER LLM + instrumentors double-counting). To preserve retry_attempt emission + under that scope, the chat_wrapper ALSO sets + OPENAI_DIRECT_RETRY_PARENT_ACTIVE_KEY in the same context. This + test mimics that contract: when BOTH keys are set, emission MUST + proceed (the suppression is openai-wrapper-internal, not external).""" + tracer, exporter, _ = fresh_tracer + request = _make_request() + parent = tracer.start_span("openai.chat") + ctx = context_api.set_value(SUPPRESS_LANGUAGE_MODEL_INSTRUMENTATION_KEY, True) + ctx = context_api.set_value(OPENAI_DIRECT_RETRY_PARENT_ACTIVE_KEY, True, ctx) + token = context_api.attach(ctx) + try: + with trace.use_span(parent, end_on_exit=False): + _sync_send_wrapper(lambda *a, **kw: _make_response(200), None, (request,), {}) + finally: + context_api.detach(token) + parent.end() + + retry_spans = _retry_spans(exporter) + assert len(retry_spans) == 1, ( + "retry_attempt MUST emit when the openai override key is set " + "alongside SUPPRESS_LANGUAGE_MODEL_INSTRUMENTATION_KEY (this is the " + "chat_wrapper's own internal HTTP-send scope)" + ) + + +def test_tracer_provider_plumbed_through_instrument_retry_emitter(): + """REGRESSION GUARD (review-driven 2026-05-16, Issue 3A): + ``instrument_retry_emitter(tracer_provider=provider)`` MUST cause + the wrapper to emit retry_attempt spans through that explicit + provider, not the global default. Without this plumbing, a + consumer who passes a tracer_provider to + ``OpenAIInstrumentor.instrument(...)`` would get the parent openai + span on their provider but the retry_attempt sent into the global + (no-op when global is unset).""" + from opentelemetry.instrumentation.fortifyroot import retry_registry as _rr + _rr._reset_for_test() + try: + uninstrument_retry_emitter() + except Exception: + pass + + # Build a NON-global provider with its own exporter. + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + + instrument_retry_emitter(tracer_provider=provider) + try: + request = _make_request(path="/v1/chat/completions", model="gpt-4o-mini") + # Create the parent via the SAME local provider so both the + # parent and the retry_attempt share one trace tree on `provider`. + parent_tracer = provider.get_tracer("test") + parent = parent_tracer.start_span("openai.chat") + with trace.use_span(parent, end_on_exit=False): + _sync_send_wrapper(lambda *a, **kw: _make_response(200), None, (request,), {}) + parent.end() + finally: + uninstrument_retry_emitter() + + spans = exporter.get_finished_spans() + retry_spans = [s for s in spans if s.name == _FR_RETRY_ATTEMPT_SPAN_NAME] + assert len(retry_spans) == 1, ( + f"retry_attempt MUST be exported via the explicit tracer_provider " + f"passed to instrument_retry_emitter; saw {len(retry_spans)} in this exporter " + f"(all span names: {[s.name for s in spans]})" + ) + + +def test_async_send_wrapper_emits(fresh_tracer): + tracer, exporter, _ = fresh_tracer + request = _make_request() + + async def wrapped(*a, **kw): + return _make_response(200) + + async def run(): + parent = tracer.start_span("openai.chat") + try: + with trace.use_span(parent, end_on_exit=False): + return await _async_send_wrapper(wrapped, None, (request,), {}) + finally: + parent.end() + + loop = asyncio.new_event_loop() + try: + result = loop.run_until_complete(run()) + finally: + loop.close() + assert result.status_code == 200 + assert len(_retry_spans(exporter)) == 1 + assert _retry_spans(exporter)[0].attributes.get("gen_ai.system") == "openai" diff --git a/packages/opentelemetry-instrumentation-openai/tests/traces/test_chat.py b/packages/opentelemetry-instrumentation-openai/tests/traces/test_chat.py index f5441c2ef7..629d15e63c 100644 --- a/packages/opentelemetry-instrumentation-openai/tests/traces/test_chat.py +++ b/packages/opentelemetry-instrumentation-openai/tests/traces/test_chat.py @@ -3,6 +3,7 @@ from unittest.mock import patch import httpx +from openai import _base_client # ST-10.4: spy target for retry-aware httpx wrap import pytest from openai.resources.chat.completions import Completions from openai.types.chat.chat_completion_message_tool_call import ( @@ -1119,8 +1120,8 @@ def test_with_asyncio_run_with_events_with_no_content( def test_chat_context_propagation( instrument_legacy, span_exporter, log_exporter, vllm_openai_client ): - send_spy = spy_decorator(httpx.Client.send) - with patch.object(httpx.Client, "send", send_spy): + send_spy = spy_decorator(_base_client.SyncHttpxClientWrapper.send) + with patch.object(_base_client.SyncHttpxClientWrapper, "send", send_spy): vllm_openai_client.chat.completions.create( model="meta-llama/Llama-3.2-1B-Instruct", messages=[ @@ -1153,8 +1154,8 @@ def test_chat_context_propagation( def test_chat_context_propagation_with_events_with_content( instrument_with_content, span_exporter, log_exporter, vllm_openai_client ): - send_spy = spy_decorator(httpx.Client.send) - with patch.object(httpx.Client, "send", send_spy): + send_spy = spy_decorator(_base_client.SyncHttpxClientWrapper.send) + with patch.object(_base_client.SyncHttpxClientWrapper, "send", send_spy): vllm_openai_client.chat.completions.create( model="meta-llama/Llama-3.2-1B-Instruct", messages=[ @@ -1206,8 +1207,8 @@ def test_chat_context_propagation_with_events_with_content( def test_chat_context_propagation_with_events_with_no_content( instrument_with_no_content, span_exporter, log_exporter, vllm_openai_client ): - send_spy = spy_decorator(httpx.Client.send) - with patch.object(httpx.Client, "send", send_spy): + send_spy = spy_decorator(_base_client.SyncHttpxClientWrapper.send) + with patch.object(_base_client.SyncHttpxClientWrapper, "send", send_spy): vllm_openai_client.chat.completions.create( model="meta-llama/Llama-3.2-1B-Instruct", messages=[ @@ -1247,8 +1248,8 @@ def test_chat_context_propagation_with_events_with_no_content( async def test_chat_async_context_propagation( instrument_legacy, span_exporter, log_exporter, async_vllm_openai_client ): - send_spy = spy_decorator(httpx.AsyncClient.send) - with patch.object(httpx.AsyncClient, "send", send_spy): + send_spy = spy_decorator(_base_client.AsyncHttpxClientWrapper.send) + with patch.object(_base_client.AsyncHttpxClientWrapper, "send", send_spy): await async_vllm_openai_client.chat.completions.create( model="meta-llama/Llama-3.2-1B-Instruct", messages=[ @@ -1282,8 +1283,8 @@ async def test_chat_async_context_propagation( async def test_chat_async_context_propagation_with_events_with_content( instrument_with_content, span_exporter, log_exporter, async_vllm_openai_client ): - send_spy = spy_decorator(httpx.AsyncClient.send) - with patch.object(httpx.AsyncClient, "send", send_spy): + send_spy = spy_decorator(_base_client.AsyncHttpxClientWrapper.send) + with patch.object(_base_client.AsyncHttpxClientWrapper, "send", send_spy): await async_vllm_openai_client.chat.completions.create( model="meta-llama/Llama-3.2-1B-Instruct", messages=[ @@ -1336,8 +1337,8 @@ async def test_chat_async_context_propagation_with_events_with_content( async def test_chat_async_context_propagation_with_events_with_no_content( instrument_with_no_content, span_exporter, log_exporter, async_vllm_openai_client ): - send_spy = spy_decorator(httpx.AsyncClient.send) - with patch.object(httpx.AsyncClient, "send", send_spy): + send_spy = spy_decorator(_base_client.AsyncHttpxClientWrapper.send) + with patch.object(_base_client.AsyncHttpxClientWrapper, "send", send_spy): await async_vllm_openai_client.chat.completions.create( model="meta-llama/Llama-3.2-1B-Instruct", messages=[ diff --git a/packages/opentelemetry-instrumentation-openai/tests/traces/test_completions.py b/packages/opentelemetry-instrumentation-openai/tests/traces/test_completions.py index 5a3fed82a0..9c67dfeb2c 100644 --- a/packages/opentelemetry-instrumentation-openai/tests/traces/test_completions.py +++ b/packages/opentelemetry-instrumentation-openai/tests/traces/test_completions.py @@ -2,6 +2,7 @@ import httpx import pytest +from openai import _base_client # ST-10.4: spy target for retry-aware httpx wrap from opentelemetry.sdk._logs import ReadableLogRecord from opentelemetry.semconv._incubating.attributes import ( gen_ai_attributes as GenAIAttributes, @@ -642,8 +643,8 @@ async def test_async_completion_streaming_with_events_with_no_content( def test_completion_context_propagation( instrument_legacy, span_exporter, log_exporter, vllm_openai_client ): - send_spy = spy_decorator(httpx.Client.send) - with patch.object(httpx.Client, "send", send_spy): + send_spy = spy_decorator(_base_client.SyncHttpxClientWrapper.send) + with patch.object(_base_client.SyncHttpxClientWrapper, "send", send_spy): vllm_openai_client.completions.create( # model="davinci-002", model="meta-llama/Llama-3.2-1B-Instruct", @@ -676,8 +677,8 @@ def test_completion_context_propagation( def test_completion_context_propagation_with_events_with_content( instrument_with_content, span_exporter, log_exporter, vllm_openai_client ): - send_spy = spy_decorator(httpx.Client.send) - with patch.object(httpx.Client, "send", send_spy): + send_spy = spy_decorator(_base_client.SyncHttpxClientWrapper.send) + with patch.object(_base_client.SyncHttpxClientWrapper, "send", send_spy): vllm_openai_client.completions.create( # model="davinci-002", model="meta-llama/Llama-3.2-1B-Instruct", @@ -726,8 +727,8 @@ def test_completion_context_propagation_with_events_with_content( def test_completion_context_propagation_with_events_with_no_content( instrument_with_no_content, span_exporter, log_exporter, vllm_openai_client ): - send_spy = spy_decorator(httpx.Client.send) - with patch.object(httpx.Client, "send", send_spy): + send_spy = spy_decorator(_base_client.SyncHttpxClientWrapper.send) + with patch.object(_base_client.SyncHttpxClientWrapper, "send", send_spy): vllm_openai_client.completions.create( # model="davinci-002", model="meta-llama/Llama-3.2-1B-Instruct", @@ -767,8 +768,8 @@ def test_completion_context_propagation_with_events_with_no_content( async def test_async_completion_context_propagation( instrument_legacy, span_exporter, log_exporter, async_vllm_openai_client ): - send_spy = spy_decorator(httpx.AsyncClient.send) - with patch.object(httpx.AsyncClient, "send", send_spy): + send_spy = spy_decorator(_base_client.AsyncHttpxClientWrapper.send) + with patch.object(_base_client.AsyncHttpxClientWrapper, "send", send_spy): await async_vllm_openai_client.completions.create( model="meta-llama/Llama-3.2-1B-Instruct", prompt="Tell me a joke about opentelemetry", @@ -801,8 +802,8 @@ async def test_async_completion_context_propagation( async def test_async_completion_context_propagation_with_events_with_content( instrument_with_content, span_exporter, log_exporter, async_vllm_openai_client ): - send_spy = spy_decorator(httpx.AsyncClient.send) - with patch.object(httpx.AsyncClient, "send", send_spy): + send_spy = spy_decorator(_base_client.AsyncHttpxClientWrapper.send) + with patch.object(_base_client.AsyncHttpxClientWrapper, "send", send_spy): await async_vllm_openai_client.completions.create( model="meta-llama/Llama-3.2-1B-Instruct", prompt="Tell me a joke about opentelemetry", @@ -851,8 +852,8 @@ async def test_async_completion_context_propagation_with_events_with_content( async def test_async_completion_context_propagation_with_events_with_no_content( instrument_with_no_content, span_exporter, log_exporter, async_vllm_openai_client ): - send_spy = spy_decorator(httpx.AsyncClient.send) - with patch.object(httpx.AsyncClient, "send", send_spy): + send_spy = spy_decorator(_base_client.AsyncHttpxClientWrapper.send) + with patch.object(_base_client.AsyncHttpxClientWrapper, "send", send_spy): await async_vllm_openai_client.completions.create( model="meta-llama/Llama-3.2-1B-Instruct", prompt="Tell me a joke about opentelemetry", diff --git a/packages/opentelemetry-instrumentation-openai/tests/traces/test_embeddings.py b/packages/opentelemetry-instrumentation-openai/tests/traces/test_embeddings.py index 2aab558d8e..9932982a59 100644 --- a/packages/opentelemetry-instrumentation-openai/tests/traces/test_embeddings.py +++ b/packages/opentelemetry-instrumentation-openai/tests/traces/test_embeddings.py @@ -3,6 +3,7 @@ import httpx import openai import pytest +from openai import _base_client # ST-10.4: spy target for retry-aware httpx wrap from opentelemetry.sdk._logs import ReadableLogRecord from opentelemetry.semconv._incubating.attributes import ( gen_ai_attributes as GenAIAttributes, @@ -401,8 +402,8 @@ def test_azure_openai_embeddings_with_events_with_no_content( def test_embeddings_context_propagation( instrument_legacy, span_exporter, log_exporter, vllm_openai_client ): - send_spy = spy_decorator(httpx.Client.send) - with patch.object(httpx.Client, "send", send_spy): + send_spy = spy_decorator(_base_client.SyncHttpxClientWrapper.send) + with patch.object(_base_client.SyncHttpxClientWrapper, "send", send_spy): vllm_openai_client.embeddings.create( input="Tell me a joke about opentelemetry", model="intfloat/e5-mistral-7b-instruct", @@ -428,8 +429,8 @@ def test_embeddings_context_propagation( def test_embeddings_context_propagation_with_events_with_content( instrument_with_content, span_exporter, log_exporter, vllm_openai_client ): - send_spy = spy_decorator(httpx.Client.send) - with patch.object(httpx.Client, "send", send_spy): + send_spy = spy_decorator(_base_client.SyncHttpxClientWrapper.send) + with patch.object(_base_client.SyncHttpxClientWrapper, "send", send_spy): response = vllm_openai_client.embeddings.create( input="Tell me a joke about opentelemetry", model="intfloat/e5-mistral-7b-instruct", @@ -469,8 +470,8 @@ def test_embeddings_context_propagation_with_events_with_content( def test_embeddings_context_propagation_with_events_with_no_content( instrument_with_no_content, span_exporter, log_exporter, vllm_openai_client ): - send_spy = spy_decorator(httpx.Client.send) - with patch.object(httpx.Client, "send", send_spy): + send_spy = spy_decorator(_base_client.SyncHttpxClientWrapper.send) + with patch.object(_base_client.SyncHttpxClientWrapper, "send", send_spy): vllm_openai_client.embeddings.create( input="Tell me a joke about opentelemetry", model="intfloat/e5-mistral-7b-instruct", @@ -503,8 +504,8 @@ def test_embeddings_context_propagation_with_events_with_no_content( async def test_async_embeddings_context_propagation( instrument_legacy, span_exporter, log_exporter, async_vllm_openai_client ): - send_spy = spy_decorator(httpx.AsyncClient.send) - with patch.object(httpx.AsyncClient, "send", send_spy): + send_spy = spy_decorator(_base_client.AsyncHttpxClientWrapper.send) + with patch.object(_base_client.AsyncHttpxClientWrapper, "send", send_spy): await async_vllm_openai_client.embeddings.create( input="Tell me a joke about opentelemetry", model="intfloat/e5-mistral-7b-instruct", @@ -531,8 +532,8 @@ async def test_async_embeddings_context_propagation( async def test_async_embeddings_context_propagation_with_events_with_content( instrument_with_content, span_exporter, log_exporter, async_vllm_openai_client ): - send_spy = spy_decorator(httpx.AsyncClient.send) - with patch.object(httpx.AsyncClient, "send", send_spy): + send_spy = spy_decorator(_base_client.AsyncHttpxClientWrapper.send) + with patch.object(_base_client.AsyncHttpxClientWrapper, "send", send_spy): response = await async_vllm_openai_client.embeddings.create( input="Tell me a joke about opentelemetry", model="intfloat/e5-mistral-7b-instruct", @@ -573,8 +574,8 @@ async def test_async_embeddings_context_propagation_with_events_with_content( async def test_async_embeddings_context_propagation_with_events_with_no_content( instrument_with_no_content, span_exporter, log_exporter, async_vllm_openai_client ): - send_spy = spy_decorator(httpx.AsyncClient.send) - with patch.object(httpx.AsyncClient, "send", send_spy): + send_spy = spy_decorator(_base_client.AsyncHttpxClientWrapper.send) + with patch.object(_base_client.AsyncHttpxClientWrapper, "send", send_spy): await async_vllm_openai_client.embeddings.create( input="Tell me a joke about opentelemetry", model="intfloat/e5-mistral-7b-instruct", diff --git a/packages/opentelemetry-instrumentation-openai/tests/traces/test_responses.py b/packages/opentelemetry-instrumentation-openai/tests/traces/test_responses.py index 6f9c1ee036..3e758ed8b1 100644 --- a/packages/opentelemetry-instrumentation-openai/tests/traces/test_responses.py +++ b/packages/opentelemetry-instrumentation-openai/tests/traces/test_responses.py @@ -529,14 +529,24 @@ def test_responses_trace_context_propagation_unit(): from opentelemetry.instrumentation.openai.v1.responses_wrappers import TracedData import time - # Set up tracing + # Set up tracing. Use the LOCAL provider directly via + # ``provider.get_tracer(...)`` and DO NOT call + # ``trace.set_tracer_provider(provider)`` — OTel only allows the + # global TracerProvider to be set once per process. The ST-10.4 + # retry-attempt test suite (which runs before this test under + # ``--all``) already sets a global provider via its ``fresh_tracer`` + # fixture; a second ``set_tracer_provider`` call would silently + # no-op (only logging "Overriding of current TracerProvider is not + # allowed"), leaving ``trace.get_tracer(__name__)`` bound to the + # earlier global — so this test's local ``exporter`` would never + # receive spans and the assertion `len(parent_spans) == 1` would + # fail in full-suite order (review-driven 2026-05-16 fix). provider = TracerProvider() exporter = InMemorySpanExporter() from opentelemetry.sdk.trace.export import SimpleSpanProcessor provider.add_span_processor(SimpleSpanProcessor(exporter)) - trace.set_tracer_provider(provider) - tracer = trace.get_tracer(__name__) + tracer = provider.get_tracer(__name__) # Create a parent span and capture its trace context with tracer.start_as_current_span("parent-span") as parent_span: diff --git a/packages/traceloop-sdk/tests/test_association_properties.py b/packages/traceloop-sdk/tests/test_association_properties.py index 989925b922..a51aa6b3de 100644 --- a/packages/traceloop-sdk/tests/test_association_properties.py +++ b/packages/traceloop-sdk/tests/test_association_properties.py @@ -5,6 +5,24 @@ from traceloop.sdk import Traceloop from traceloop.sdk.decorators import task, workflow +# ST-10.4 (review-driven 2026-05-16): generalized from name-based +# LangChain-only filter to role-based. Drops every provider's +# retry_attempt sibling (openai / anthropic / bedrock / langchain / +# llamaindex / litellm) uniformly. A new provider's retry_attempt +# automatically participates without touching this helper. +_FR_SPAN_ROLE_KEY = "fortifyroot.span.role" +_FR_SPAN_ROLE_RETRY_ATTEMPT = "retry_attempt" + + +def _without_fr_langchain_retry_attempt_spans(spans): + # Kept under the original LangChain-only name for callsite stability + # in this file. New tests should use a locally-defined role-based + # filter (see e.g. test_workflows.py, test_prompt_management.py). + return [ + s for s in spans + if (s.attributes or {}).get(_FR_SPAN_ROLE_KEY) != _FR_SPAN_ROLE_RETRY_ATTEMPT + ] + def test_association_properties(exporter): @workflow(name="test_workflow") @@ -87,7 +105,7 @@ def test_langchain_association_properties(exporter): {"metadata": {"user_id": "1234", "session_id": 456}}, ) - spans = exporter.get_finished_spans() + spans = _without_fr_langchain_retry_attempt_spans(exporter.get_finished_spans()) assert [ "ChatPromptTemplate.task", @@ -158,7 +176,7 @@ def test_workflow_external(): test_workflow_external() - spans = exporter.get_finished_spans() + spans = _without_fr_langchain_retry_attempt_spans(exporter.get_finished_spans()) assert [ "ChatPromptTemplate.task", diff --git a/packages/traceloop-sdk/tests/test_privacy_no_prompts.py b/packages/traceloop-sdk/tests/test_privacy_no_prompts.py index c25275137b..bd31251ca6 100644 --- a/packages/traceloop-sdk/tests/test_privacy_no_prompts.py +++ b/packages/traceloop-sdk/tests/test_privacy_no_prompts.py @@ -8,6 +8,24 @@ from traceloop.sdk.decorators import workflow, task +# ST-10.4: ``fortifyroot.openai.retry_attempt`` (and equivalents for +# anthropic / bedrock / framework wrappers) lands in this test's +# session-scoped exporter ahead of the logical ``openai.chat`` span. +# Filter by the canonical ``fortifyroot.span.role`` attribute so every +# provider's retry_attempt is dropped uniformly. See +# fr-system-tests/docs/development/ai-logs/st_phase_10.txt addendum +# 2026-05-16 for context. +_FR_SPAN_ROLE_KEY = "fortifyroot.span.role" +_FR_SPAN_ROLE_RETRY_ATTEMPT = "retry_attempt" + + +def _without_retry_attempt_spans(spans): + return [ + s for s in spans + if (s.attributes or {}).get(_FR_SPAN_ROLE_KEY) != _FR_SPAN_ROLE_RETRY_ATTEMPT + ] + + @pytest.fixture(autouse=True) def disable_trace_content(): os.environ["TRACELOOP_TRACE_CONTENT"] = "false" @@ -38,13 +56,13 @@ def joke_workflow(): joke_workflow() - spans = exporter.get_finished_spans() + spans = _without_retry_attempt_spans(exporter.get_finished_spans()) assert [span.name for span in spans] == [ "openai.chat", "joke_creation.task", "pirate_joke_generator.workflow", ] - open_ai_span = spans[0] + open_ai_span = next(s for s in spans if s.name == "openai.chat") assert open_ai_span.attributes[GenAIAttributes.GEN_AI_USAGE_INPUT_TOKENS] == 15 assert not open_ai_span.attributes.get(f"{GenAIAttributes.GEN_AI_PROMPT}.0.content") assert not open_ai_span.attributes.get( diff --git a/packages/traceloop-sdk/tests/test_prompt_management.py b/packages/traceloop-sdk/tests/test_prompt_management.py index ff2aa16ad2..57ac4f71a6 100644 --- a/packages/traceloop-sdk/tests/test_prompt_management.py +++ b/packages/traceloop-sdk/tests/test_prompt_management.py @@ -8,6 +8,21 @@ from traceloop.sdk.prompts import get_prompt from traceloop.sdk.prompts.client import PromptRegistryClient + +# ST-10.4: filter ``fortifyroot.*.retry_attempt`` sibling spans out of +# legacy exact-span-list / ``spans[0]`` assertions. Role-based filter +# so every provider's retry_attempt is dropped uniformly. See addendum +# 2026-05-16 in st_phase_10.txt for context. +_FR_SPAN_ROLE_KEY = "fortifyroot.span.role" +_FR_SPAN_ROLE_RETRY_ATTEMPT = "retry_attempt" + + +def _without_retry_attempt_spans(spans): + return [ + s for s in spans + if (s.attributes or {}).get(_FR_SPAN_ROLE_KEY) != _FR_SPAN_ROLE_RETRY_ATTEMPT + ] + prompts_json = """ { "prompts": [ @@ -229,11 +244,11 @@ def test_prompt_management(exporter, openai_client): prompt_args = get_prompt(key="joke_generator", variables={"style": "pirate"}) openai_client.chat.completions.create(**prompt_args) - spans = exporter.get_finished_spans() + spans = _without_retry_attempt_spans(exporter.get_finished_spans()) assert [span.name for span in spans] == [ "openai.chat", ] - open_ai_span = spans[0] + open_ai_span = next(s for s in spans if s.name == "openai.chat") assert ( open_ai_span.attributes[f"{GenAIAttributes.GEN_AI_PROMPT}.0.content"] == "Tell me a joke about OpenTelemetry, pirate style" @@ -250,8 +265,8 @@ def test_prompt_management_with_tools(exporter, openai_client): prompt_args = get_prompt(key="joke_generator", variables={"style": "pirate"}) openai_client.chat.completions.create(**prompt_args) - spans = exporter.get_finished_spans() - open_ai_span = spans[0] + spans = _without_retry_attempt_spans(exporter.get_finished_spans()) + open_ai_span = next(s for s in spans if s.name == "openai.chat") completion = open_ai_span.attributes.get( f"{GenAIAttributes.GEN_AI_COMPLETION}.0.tool_calls.0.name" ) @@ -266,8 +281,8 @@ def test_prompt_management_with_response_format(exporter, openai_client): prompt_args = get_prompt(key="joke_generator", variables={"style": "pirate"}) openai_client.chat.completions.create(**prompt_args) - spans = exporter.get_finished_spans() - open_ai_span = spans[0] + spans = _without_retry_attempt_spans(exporter.get_finished_spans()) + open_ai_span = next(s for s in spans if s.name == "openai.chat") completion = open_ai_span.attributes.get( f"{GenAIAttributes.GEN_AI_COMPLETION}.0.content" ) diff --git a/packages/traceloop-sdk/tests/test_sdk_initialization.py b/packages/traceloop-sdk/tests/test_sdk_initialization.py index d1475938da..89876daacc 100644 --- a/packages/traceloop-sdk/tests/test_sdk_initialization.py +++ b/packages/traceloop-sdk/tests/test_sdk_initialization.py @@ -4,6 +4,24 @@ from traceloop.sdk.decorators import workflow +# ST-10.4: ``fortifyroot.openai.retry_attempt`` lands in the exporter +# ahead of the logical ``openai.chat`` span; filter by role so the +# legacy ``spans[0] == openai.chat`` assumption survives. The user's +# custom span_postprocess_callback is intentionally still invoked on +# the retry_attempt span (product-code contract: callback sees every +# exported span); only the test's assertion is filtered. See addendum +# 2026-05-16 in st_phase_10.txt for context. +_FR_SPAN_ROLE_KEY = "fortifyroot.span.role" +_FR_SPAN_ROLE_RETRY_ATTEMPT = "retry_attempt" + + +def _without_retry_attempt_spans(spans): + return [ + s for s in spans + if (s.attributes or {}).get(_FR_SPAN_ROLE_KEY) != _FR_SPAN_ROLE_RETRY_ATTEMPT + ] + + @pytest.fixture def openai_client(): return OpenAI() @@ -112,8 +130,10 @@ def test_span_postprocess_callback(exporter_with_custom_span_postprocess_callbac messages=[{"role": "user", "content": "Tell me a joke about opentelemetry"}], ) - spans = exporter_with_custom_span_postprocess_callback.get_finished_spans() - open_ai_span = spans[0] + spans = _without_retry_attempt_spans( + exporter_with_custom_span_postprocess_callback.get_finished_spans() + ) + open_ai_span = next(s for s in spans if s.name == "openai.chat") assert open_ai_span.attributes["gen_ai.prompt.0.content"] == "REDACTED" assert open_ai_span.attributes["gen_ai.completion.0.content"] == "REDACTED" diff --git a/packages/traceloop-sdk/tests/test_workflows.py b/packages/traceloop-sdk/tests/test_workflows.py index daa7107449..eaadaa269f 100644 --- a/packages/traceloop-sdk/tests/test_workflows.py +++ b/packages/traceloop-sdk/tests/test_workflows.py @@ -13,6 +13,22 @@ from traceloop.sdk.decorators import workflow, task +# ST-10.4: filter ``fortifyroot.*.retry_attempt`` sibling spans out +# of legacy exact-span-list / set assertions. Role-based filter so +# every provider's retry_attempt (openai, anthropic, bedrock, +# langchain, llamaindex, litellm) is dropped uniformly. See addendum +# 2026-05-16 in st_phase_10.txt for context. +_FR_SPAN_ROLE_KEY = "fortifyroot.span.role" +_FR_SPAN_ROLE_RETRY_ATTEMPT = "retry_attempt" + + +def _without_retry_attempt_spans(spans): + return [ + s for s in spans + if (s.attributes or {}).get(_FR_SPAN_ROLE_KEY) != _FR_SPAN_ROLE_RETRY_ATTEMPT + ] + + @pytest.fixture def openai_client(): return OpenAI() @@ -40,7 +56,7 @@ def joke_workflow(): joke = joke_workflow() - spans = exporter.get_finished_spans() + spans = _without_retry_attempt_spans(exporter.get_finished_spans()) assert [span.name for span in spans] == [ "openai.chat", "something_creator.task", @@ -86,7 +102,7 @@ async def joke_workflow(): joke = await joke_workflow() - spans = exporter.get_finished_spans() + spans = _without_retry_attempt_spans(exporter.get_finished_spans()) assert [span.name for span in spans] == [ "openai.chat", "something_creator.task", @@ -139,7 +155,7 @@ def joke_workflow(): joke_workflow() - spans = exporter.get_finished_spans() + spans = _without_retry_attempt_spans(exporter.get_finished_spans()) assert set([span.name for span in spans]) == set( [ "openai.chat",