From f54ef596c8d5c8dea5bdb555af49538e9da7b0cd Mon Sep 17 00:00:00 2001 From: Alexis Georges Date: Tue, 11 Aug 2026 16:11:36 -0400 Subject: [PATCH 1/5] feat(client): cache-aware token accounting and a reusable span lifecycle Six handlers each hand-write their own usage attributes today, which is how they came to disagree. This adds the shared writers they will move onto, so the numbers have one author. parse_usage grows cache handling. It read three key-pair spellings and dropped every cache field, so Anthropic undercounted input by the whole cached portion of a call: a turn that reads 19,971 tokens from cache and writes 3,580 more reports input_tokens: 3, and 3 is what we billed a dashboard on. It now folds every accepted spelling, including Bedrock Converse's cacheWriteInputTokens, and reports the breakdown in input_details. The fold stays provider-blind, which is the contract handlers must respect: return raw usage with the cache fields intact, or an input figure that already includes cache with the fields omitted. Returning a pre-folded input alongside the fields double-counts. add_cached_tokens_to_input applies the Anthropic-shaped rule at the call site, and lang_chain_span_usage the LangChain-shaped one, because the direction differs per provider and centralising it would silently double-count for two providers out of three. SpanUsage is the type that means "the folding is already done". set_usage_span_attributes writes all seven attributes every time, zeros included, because an absent attribute drops a span from every query that groups on usage, which reads as "no cached tokens" rather than "this handler forgot to say". It also owns the two OpenLLMetry aliases, which previously lived beside the completion text and were computed off Anthropic's cache-excluding input field, so they disagreed with the canonical numbers on the same span. RunUsage counts whether any turn reported usage rather than testing the total for zero, so a failed run can put its partial spend on the root while a run that never completed a call correctly says nothing. All-zero attributes would assert the run cost nothing. number_or_zero replaces bare int(...), which raised on None. An emitted NaN is worse than an emitted 0, because the metric guard tests `> 0` and that is false for NaN, so the metric vanishes instead of reading low. end_span_once makes the streaming teardown idempotent and marks an abandoned stream without failing it. Stopping early is a normal thing for a consumer to do, and LaunchDarkly's own metrics record neither success nor error for it, so ERROR would put two dashboards in disagreement about one run. It tracks id(span) because an OTel span is not guaranteed hashable. UsageDict gains input_details, which broke graph.py's UsageDict(**dict) splat. Now built with named arguments, so the next member added here cannot silently arrive from a dict with no business filling it. --- .../src/launchdarkly_ai_server/__init__.py | 20 + .../src/launchdarkly_ai_server/graph.py | 10 +- .../src/launchdarkly_ai_server/types.py | 15 + .../src/launchdarkly_ai_server/utils.py | 308 +++++++++++++- packages/client/tests/test_span_usage.py | 382 ++++++++++++++++++ 5 files changed, 721 insertions(+), 14 deletions(-) create mode 100644 packages/client/tests/test_span_usage.py diff --git a/packages/client/src/launchdarkly_ai_server/__init__.py b/packages/client/src/launchdarkly_ai_server/__init__.py index 5d1f2ee..feaeaec 100644 --- a/packages/client/src/launchdarkly_ai_server/__init__.py +++ b/packages/client/src/launchdarkly_ai_server/__init__.py @@ -34,6 +34,7 @@ HandlerResult, HandlerStreamEvent, InitClientOptions, + InputTokenDetails, JudgeResult, JudgeRunResult, JudgeTask, @@ -56,15 +57,24 @@ ) from .types_validation import parse_ai_config from .utils import ( + RunUsage, + SpanUsage, + add_cached_tokens_to_input, create_handler, + create_run_usage, + end_span_once, + lang_chain_span_usage, make_track_data, normalize_mode, + number_or_zero, parse_json_with_possible_fences, parse_template, parse_usage, set_ld_span_attributes, + set_model_identity_attributes, set_openllmetry_completion, set_openllmetry_prompt, + set_usage_span_attributes, to_ld_context, ) @@ -100,7 +110,17 @@ "StreamDoneEvent", "StreamEvent", "TrackData", + "InputTokenDetails", + "RunUsage", + "SpanUsage", "UsageDict", + "add_cached_tokens_to_input", + "create_run_usage", + "end_span_once", + "lang_chain_span_usage", + "number_or_zero", + "set_model_identity_attributes", + "set_usage_span_attributes", "VariationMeta", # utils "create_handler", diff --git a/packages/client/src/launchdarkly_ai_server/graph.py b/packages/client/src/launchdarkly_ai_server/graph.py index 0a397cf..46d606f 100644 --- a/packages/client/src/launchdarkly_ai_server/graph.py +++ b/packages/client/src/launchdarkly_ai_server/graph.py @@ -668,7 +668,15 @@ async def invoke( return ProviderGraphResponse( response=final_response, - usage=UsageDict(**total_usage), + # Named rather than splatted, so a new UsageDict member cannot silently arrive + # here from a dict that has no business filling it. Graph totals carry no cache + # breakdown: they are a sum across nodes, and the per-node detail is on the node's + # own spans. + usage=UsageDict( + input=total_usage["input"], + output=total_usage["output"], + total=total_usage["total"], + ), judge_results=judge_results, ) diff --git a/packages/client/src/launchdarkly_ai_server/types.py b/packages/client/src/launchdarkly_ai_server/types.py index 55e863f..57f8134 100644 --- a/packages/client/src/launchdarkly_ai_server/types.py +++ b/packages/client/src/launchdarkly_ai_server/types.py @@ -192,11 +192,26 @@ def has_stream(self) -> bool: # --------------------------------------------------------------------------- +@dataclass +class InputTokenDetails: + """The cache breakdown behind an inclusive ``input`` figure. + + Present only when the provider reported at least one cache field. ``uncached + cache_read + + cache_creation`` equals :attr:`UsageDict.input`. + """ + + uncached: int = 0 + cache_read: int = 0 + cache_creation: int = 0 + + @dataclass class UsageDict: input: int = 0 output: int = 0 total: int = 0 + #: Cache breakdown, when the provider reported one. See :class:`InputTokenDetails`. + input_details: InputTokenDetails | None = None @dataclass diff --git a/packages/client/src/launchdarkly_ai_server/utils.py b/packages/client/src/launchdarkly_ai_server/utils.py index 64648dc..81e83fd 100644 --- a/packages/client/src/launchdarkly_ai_server/utils.py +++ b/packages/client/src/launchdarkly_ai_server/utils.py @@ -2,6 +2,7 @@ import json import re +from dataclasses import dataclass, field from typing import Any, Literal from .types import ( @@ -49,27 +50,308 @@ def normalize_mode(mode: str | None) -> Literal["agent", "messages"]: return "agent" if mode == "agent" else "messages" -_USAGE_KEY_PAIRS = [ - ("input_tokens", "output_tokens"), - ("inputTokens", "outputTokens"), - ("input", "output"), -] +def number_or_zero(value: Any) -> int: + """Coerces a provider-reported token count to a finite number, defaulting to 0. + + Provider SDKs report usage as loosely typed bags where a field may be absent, ``None``, or (in + streaming paths) a partially populated value. Every count that reaches a span or a usage total + should pass through here. + An emitted 0 is bad. An emitted ``NaN`` is worse, because the metric guard tests whether the + total is greater than zero, and that test is false for ``NaN``, so the metric is dropped + silently rather than reported low. -def parse_usage(usage: dict[str, Any]) -> dict[str, int]: + Replaces the bare ``int(...)`` this module used to do, which raised on ``None``. """ - Normalises token counts from three possible key-pair shapes. - Always computes ``total`` from ``input + output``; never trusts a ``total`` - field from the raw object. + if value is None or isinstance(value, bool): + return 0 + try: + parsed = float(value) + except (TypeError, ValueError): + return 0 + if parsed != parsed or parsed in (float("inf"), float("-inf")): + return 0 + return int(parsed) + + +#: Provider usage field names, tried in order; the first row whose input *and* output keys are both +#: present wins. Cache fields list every spelling we accept, because providers disagree (Anthropic +#: ``cache_creation_input_tokens``, Bedrock Converse ``cacheWriteInputTokens``) and an unlisted +#: spelling is silently dropped from the total. +#: +#: Contract for handlers: this fold is provider-blind, so it always adds the cache fields on top of +#: the input field. A handler must therefore return *either* raw usage with the cache fields intact +#: (Anthropic-style reporting, where cache is genuinely additional), *or* an input figure that +#: already includes cache with the cache fields omitted (OpenAI- and LangChain-style reporting). +#: Returning a pre-folded input *alongside* the cache fields double-counts the cached portion. +_USAGE_KEY_PAIRS: list[tuple[str, str, tuple[str, ...], tuple[str, ...]]] = [ + ( + "input_tokens", + "output_tokens", + ("cache_read_input_tokens",), + ("cache_creation_input_tokens",), + ), + ( + "inputTokens", + "outputTokens", + ("cacheReadInputTokens",), + ("cacheCreationInputTokens", "cacheWriteInputTokens"), + ), + ("input", "output", (), ()), +] + + +def _read_cache_field(usage: dict[str, Any], keys: tuple[str, ...]) -> int: + """Sums every accepted spelling of a cache field, so an alias never silently reads as 0.""" + return sum(number_or_zero(usage.get(key)) for key in keys) + + +def parse_usage(usage: dict[str, Any]) -> dict[str, Any]: + """Normalises token counts from three possible key-pair shapes, folding cache tokens in. + + Always computes ``total`` from ``input + output``; never trusts a ``total`` field from the raw + object. A provider-supplied total can include tokens that appear in neither input nor output, + which would make the figure derivable on five handlers and not on the sixth. + + ``input`` is the inclusive total: every accepted cache spelling is added on top of the reported + input figure. See the contract note on ``_USAGE_KEY_PAIRS``. + + When the matched row defines cache keys and at least one of them is present, the result also + carries ``input_details`` with ``uncached``, ``cache_read`` and ``cache_creation``. + + An unrecognised bag returns all zeros rather than raising. """ - for input_key, output_key in _USAGE_KEY_PAIRS: + for input_key, output_key, cache_read_keys, cache_creation_keys in _USAGE_KEY_PAIRS: if input_key in usage and output_key in usage: - inp = int(usage[input_key]) - out = int(usage[output_key]) - return {"input": inp, "output": out, "total": inp + out} + cache_read = _read_cache_field(usage, cache_read_keys) + cache_creation = _read_cache_field(usage, cache_creation_keys) + uncached = number_or_zero(usage[input_key]) + inp = uncached + cache_read + cache_creation + out = number_or_zero(usage[output_key]) + result: dict[str, Any] = {"input": inp, "output": out, "total": inp + out} + has_details = bool(cache_read_keys or cache_creation_keys) and any( + key in usage for key in (*cache_read_keys, *cache_creation_keys) + ) + if has_details: + result["input_details"] = { + "uncached": uncached, + "cache_read": cache_read, + "cache_creation": cache_creation, + } + return result return {"input": 0, "output": 0, "total": 0} +@dataclass +class SpanUsage: + """The provider-neutral token counts a span reports, after the caller applied its cache rule. + + ``input`` is the *inclusive* total: whether cache tokens were already counted inside the + provider's input figure (OpenAI, LangChain) or reported alongside it and folded in by the caller + (Anthropic), by the time a value reaches this type the folding is done. + + That is the whole point of the type existing. See :func:`add_cached_tokens_to_input`. + """ + + input: int = 0 + output: int = 0 + cache_read: int = 0 + cache_creation: int = 0 + + +def add_cached_tokens_to_input(raw_usage: dict[str, Any]) -> SpanUsage: + """Adds a provider's cached-token counts into its input total. + + Providers describe the same call two different ways. Anthropic reports cache reads and cache + writes as buckets *beside* ``input_tokens``, so ``input_tokens`` counts only the new tokens: a + turn that read 19,971 tokens from cache and wrote 3,580 more reports ``input_tokens: 3``, when + the model actually processed 23,554. OpenAI and LangChain report one input figure that already + contains the cached tokens, with a subset breakdown alongside. + + This is for the first kind. Applying it to the second would count the cached tokens twice, which + is why the rule lives at the call site and not inside :func:`set_usage_span_attributes`. + + Not named for Anthropic: Bedrock Converse reports the same way, and the shape is the reason it + applies, not the vendor. + """ + cache_read = number_or_zero(raw_usage.get("cache_read_input_tokens")) + cache_creation = number_or_zero(raw_usage.get("cache_creation_input_tokens")) + return SpanUsage( + input=number_or_zero(raw_usage.get("input_tokens")) + + cache_read + + cache_creation, + output=number_or_zero(raw_usage.get("output_tokens")), + cache_read=cache_read, + cache_creation=cache_creation, + ) + + +@dataclass +class RunUsage: + """A run's accumulated token spend, for the ``invoke_agent`` root. + + The root is the only span carrying ``launchdarkly.*`` and the ``feature_flag`` event, so it is + the span a config-scoped query finds. Without a run total on it, that query returns nothing: + summing the children requires having already found them. + + Provider-blind on purpose. It sums four numbers and remembers whether anything was added. The + cache-folding rule that differs per provider has already been applied by the time a + :class:`SpanUsage` exists, so each handler maps its own bag first and this stays shared. + + The two Anthropic handlers deliberately do not use it. Their run total has to stay in Anthropic's + own field names with the cache buckets *unfolded*, because it is also the handler's return value + and :func:`parse_usage` folds it exactly once; handing back a :class:`SpanUsage` there would + count the cache twice. + """ + + total: SpanUsage = field(default_factory=SpanUsage) + _turns: int = 0 + + @property + def reported(self) -> bool: + """Whether any turn reported usage at all. + + The failure path needs this to tell "no call completed" from "a call completed and reported + zero". Only the second may be written to a span: all-zero attributes assert the run cost + nothing, which a run whose first call died mid-flight cannot claim, whereas an absent + attribute correctly says "unknown". + + Counted rather than derived by testing the total for zero, so a provider reporting a + genuinely empty bag stays distinguishable. + """ + return self._turns > 0 + + def add(self, turn: SpanUsage | None) -> None: + """Adds one turn. ``None`` is a no-op and does not count as reported.""" + if turn is None: + return + self._turns += 1 + self.total.input += turn.input + self.total.output += turn.output + self.total.cache_read += turn.cache_read + self.total.cache_creation += turn.cache_creation + + +def create_run_usage() -> RunUsage: + """A fresh run accumulator. Present so call sites read the same as the TypeScript SDK's.""" + return RunUsage() + + +def lang_chain_span_usage(usage: dict[str, Any] | None) -> SpanUsage | None: + """One LangChain ``usage_metadata`` bag as :class:`SpanUsage`, or ``None`` when it reports nothing. + + LangChain already includes cached tokens in ``input_tokens``, so nothing is folded here. Shared + because both LangChain handlers read the identical shape, and because ``None`` for an empty bag + is what keeps a turn the provider said nothing about from registering as reported: the callback + path hands over ``{}`` rather than nothing when a provider omits usage entirely. + """ + if not usage: + return None + if usage.get("input_tokens") is None and usage.get("output_tokens") is None: + return None + details = usage.get("input_token_details") or {} + return SpanUsage( + input=number_or_zero(usage.get("input_tokens")), + output=number_or_zero(usage.get("output_tokens")), + cache_read=number_or_zero(details.get("cache_read")), + cache_creation=number_or_zero(details.get("cache_creation")), + ) + + +def set_usage_span_attributes(span: Any, usage: SpanUsage) -> None: + """Writes the OpenTelemetry ``gen_ai.usage.*`` token attributes onto a span. + + Always writes all seven attributes, every time, including zeros. Consumers group and aggregate + on the complete set, and an *absent* attribute drops a span from those queries entirely, which + reads as "this call had no cached tokens" when it actually means "this handler forgot to say". A + provider with no cache-creation concept still reports 0. + + ``total`` is always ``input + output``. Providers that report their own total are deliberately + not trusted here. + + This helper does **not** know each provider's cache accounting, and must not learn it. Anthropic + reports cache tokens alongside input, so its callers fold them in; OpenAI and LangChain already + count cache tokens inside input, so their callers pass the reported figure through untouched. + Both arrive here as an inclusive ``input``. Centralising that rule would silently double-count + for two providers out of three. + + The last two are OpenLLMetry aliases for the same two numbers. They belong here rather than + beside the completion text so they cannot disagree with the canonical attributes above: + computed at a call site off the raw input field, the alias undercounted Anthropic by every + cached token, because on Anthropic that field excludes the cache. One writer, one number. + """ + inp = number_or_zero(usage.input) + out = number_or_zero(usage.output) + span.set_attribute("gen_ai.usage.input_tokens", inp) + span.set_attribute("gen_ai.usage.output_tokens", out) + span.set_attribute("gen_ai.usage.total_tokens", inp + out) + span.set_attribute( + "gen_ai.usage.cache_read.input_tokens", number_or_zero(usage.cache_read) + ) + span.set_attribute( + "gen_ai.usage.cache_creation.input_tokens", number_or_zero(usage.cache_creation) + ) + span.set_attribute("gen_ai.usage.prompt_tokens", inp) + span.set_attribute("gen_ai.usage.completion_tokens", out) + + +def set_model_identity_attributes( + span: Any, + provider_name: str, + request_model: str, + legacy_system: str | None = None, +) -> None: + """Writes the model identity attributes that every LLM span carries. + + Both spellings of the provider key are emitted on purpose. ``gen_ai.system`` is the pre-1.37 + semconv name and is what handlers shipped before the span hierarchy landed; + ``gen_ai.provider.name`` is the current name. Emitting only the new key would silently break + dashboards written against the old one, and emitting only the old one leaves us off-spec, so + both go out until the next major. + + ``legacy_system`` exists because the two keys do not always want the same value. The LangChain + handlers ship ``gen_ai.system = 'langchain'``, but ``gen_ai.provider.name`` means *who served + the model* and its semconv enum has no ``langchain`` member, so those handlers pass the real + provider for the new key and keep the framework name on the old one. + """ + span.set_attribute( + "gen_ai.system", provider_name if legacy_system is None else legacy_system + ) + span.set_attribute("gen_ai.provider.name", provider_name) + span.set_attribute("gen_ai.request.model", request_model) + + +def end_span_once(span: Any, tracker: set[int], abandoned: bool = False) -> None: + """Ends a span exactly once, even when the caller cannot know whether an earlier path ended it. + + The streaming handlers need this: a consumer that breaks out of ``async for``, or throws inside + the loop body, makes the generator run its ``finally`` without ever entering ``except``, so the + cleanup path and the success path can both reach the same span. Ending twice is silently ignored + by the OTel SDK but recorded as a diagnostic error, and would also hide a genuine leak, so the + guard is explicit. + + *tracker* holds ``id(span)`` rather than the span itself, because an OTel span is not guaranteed + hashable across implementations. + + An abandoned span is marked with ``launchdarkly.stream.abandoned`` and deliberately left at + ``UNSET`` rather than ``ERROR``. Stopping early is a normal thing for a consumer to do, such as + rendering enough of a response and moving on, and nothing failed. Marking it ``ERROR`` would + also put the trace at odds with LaunchDarkly's own metrics, which record neither a success nor + an error for an abandoned stream: two dashboards would disagree about the same run. The + attribute keeps abandonment findable without asserting a failure. + + Ending the span is not always enough. Two handlers also hold a vendor generator or run that must + be closed or cancelled in the same ``finally``. See TELEMETRY-CONTRACT.md section 6. + """ + key = id(span) + if key in tracker: + return + tracker.add(key) + if abandoned: + span.set_attribute("launchdarkly.stream.abandoned", True) + span.end() + + def parse_template(template: str, variables: dict[str, Any]) -> str: """ Replaces ``{{variable}}`` placeholders in *template* with values from diff --git a/packages/client/tests/test_span_usage.py b/packages/client/tests/test_span_usage.py new file mode 100644 index 0000000..605a21d --- /dev/null +++ b/packages/client/tests/test_span_usage.py @@ -0,0 +1,382 @@ +"""Tests for the span usage and identity helpers. + +Covers TELEMETRY-CONTRACT.md sections 6 (span lifecycle), 8 (token accounting) and 9 (model +identity). + +The load-bearing tests here are the cache-direction ones. Folding cache tokens for a provider that +already includes them, or failing to fold for a provider that does not, produces a number that is +wrong by exactly the cached portion of every call, and nothing else catches it. +""" + +from __future__ import annotations + +from typing import Any + +from launchdarkly_ai_server.utils import ( + SpanUsage, + add_cached_tokens_to_input, + create_run_usage, + end_span_once, + lang_chain_span_usage, + number_or_zero, + parse_usage, + set_model_identity_attributes, + set_usage_span_attributes, +) + + +class FakeSpan: + def __init__(self) -> None: + self.attributes: dict[str, Any] = {} + self.ended = 0 + + def set_attribute(self, key: str, value: Any) -> None: + self.attributes[key] = value + + def end(self) -> None: + self.ended += 1 + + +# ─── number_or_zero ────────────────────────────────────────────────────────── + + +class TestNumberOrZero: + def test_passes_a_number_through(self) -> None: + assert number_or_zero(42) == 42 + + def test_none_becomes_zero_rather_than_raising(self) -> None: + # The old bare int(...) raised here, taking the whole call down with it. + assert number_or_zero(None) == 0 + + def test_a_non_numeric_string_becomes_zero(self) -> None: + assert number_or_zero("abc") == 0 + + def test_a_numeric_string_is_read(self) -> None: + assert number_or_zero("17") == 17 + + def test_nan_becomes_zero(self) -> None: + # NaN is worse than 0: the metric guard tests `> 0`, which is false for NaN, so the metric + # would be dropped silently rather than reported low. + assert number_or_zero(float("nan")) == 0 + + def test_infinity_becomes_zero(self) -> None: + assert number_or_zero(float("inf")) == 0 + assert number_or_zero(float("-inf")) == 0 + + def test_a_bool_is_not_a_token_count(self) -> None: + assert number_or_zero(True) == 0 + + def test_a_float_truncates(self) -> None: + assert number_or_zero(3.7) == 3 + + +# ─── parse_usage ───────────────────────────────────────────────────────────── + + +class TestParseUsage: + def test_reads_the_snake_case_pair(self) -> None: + assert parse_usage({"input_tokens": 10, "output_tokens": 5}) == { + "input": 10, + "output": 5, + "total": 15, + } + + def test_reads_the_camel_case_pair(self) -> None: + assert parse_usage({"inputTokens": 10, "outputTokens": 5})["total"] == 15 + + def test_reads_the_bare_pair(self) -> None: + assert parse_usage({"input": 10, "output": 5})["total"] == 15 + + def test_first_matching_pair_wins(self) -> None: + result = parse_usage( + {"input_tokens": 1, "output_tokens": 2, "input": 100, "output": 200} + ) + assert result["total"] == 3 + + def test_never_trusts_a_provider_total(self) -> None: + result = parse_usage( + {"input_tokens": 3, "output_tokens": 7, "total_tokens": 999} + ) + assert result["total"] == 10 + + def test_an_unrecognised_bag_returns_zeros(self) -> None: + assert parse_usage({"foo": 1}) == {"input": 0, "output": 0, "total": 0} + + def test_a_none_field_does_not_raise(self) -> None: + assert parse_usage({"input_tokens": None, "output_tokens": 5})["input"] == 0 + + def test_folds_anthropic_cache_tokens_into_input(self) -> None: + # Anthropic reports cache beside input, so the real input is the sum of all three. + result = parse_usage( + { + "input_tokens": 3, + "output_tokens": 10, + "cache_read_input_tokens": 19971, + "cache_creation_input_tokens": 3580, + } + ) + assert result["input"] == 23554 + assert result["total"] == 23564 + + def test_reports_the_cache_breakdown_when_present(self) -> None: + result = parse_usage( + { + "input_tokens": 3, + "output_tokens": 1, + "cache_read_input_tokens": 10, + "cache_creation_input_tokens": 20, + } + ) + assert result["input_details"] == { + "uncached": 3, + "cache_read": 10, + "cache_creation": 20, + } + + def test_omits_the_breakdown_when_no_cache_field_is_present(self) -> None: + assert "input_details" not in parse_usage( + {"input_tokens": 1, "output_tokens": 2} + ) + + def test_omits_the_breakdown_for_the_bare_pair_which_has_no_cache_keys( + self, + ) -> None: + assert "input_details" not in parse_usage({"input": 1, "output": 2}) + + def test_accepts_the_bedrock_cache_write_alias(self) -> None: + # An unlisted spelling is silently dropped from the total, so every alias must be summed. + result = parse_usage( + {"inputTokens": 5, "outputTokens": 1, "cacheWriteInputTokens": 100} + ) + assert result["input"] == 105 + + def test_sums_both_camel_case_creation_spellings(self) -> None: + result = parse_usage( + { + "inputTokens": 0, + "outputTokens": 0, + "cacheCreationInputTokens": 10, + "cacheWriteInputTokens": 5, + } + ) + assert result["input"] == 15 + + +# ─── add_cached_tokens_to_input ────────────────────────────────────────────── + + +class TestAddCachedTokensToInput: + def test_adds_both_cache_buckets_on_top_of_input(self) -> None: + usage = add_cached_tokens_to_input( + { + "input_tokens": 3, + "output_tokens": 10, + "cache_read_input_tokens": 19971, + "cache_creation_input_tokens": 3580, + } + ) + assert usage.input == 23554 + assert usage.output == 10 + assert usage.cache_read == 19971 + assert usage.cache_creation == 3580 + + def test_an_absent_cache_field_is_zero_not_an_error(self) -> None: + usage = add_cached_tokens_to_input({"input_tokens": 5, "output_tokens": 2}) + assert (usage.input, usage.cache_read, usage.cache_creation) == (5, 0, 0) + + def test_an_absent_usage_bag_is_all_zeros(self) -> None: + usage = add_cached_tokens_to_input({}) + assert (usage.input, usage.output) == (0, 0) + + +# ─── lang_chain_span_usage ─────────────────────────────────────────────────── + + +class TestLangChainSpanUsage: + def test_does_not_add_cache_on_top_of_input(self) -> None: + # LangChain already counts cached tokens inside input_tokens. Adding them would double-count. + usage = lang_chain_span_usage( + { + "input_tokens": 100, + "output_tokens": 10, + "input_token_details": {"cache_read": 80, "cache_creation": 5}, + } + ) + assert usage is not None + assert usage.input == 100 + assert usage.cache_read == 80 + assert usage.cache_creation == 5 + + def test_an_empty_bag_is_none_not_zeros(self) -> None: + # None keeps a turn the provider said nothing about from registering as reported. + assert lang_chain_span_usage({}) is None + assert lang_chain_span_usage(None) is None + + def test_a_bag_with_no_token_keys_is_none(self) -> None: + assert lang_chain_span_usage({"input_token_details": {"cache_read": 5}}) is None + + def test_a_zero_count_is_still_reported(self) -> None: + usage = lang_chain_span_usage({"input_tokens": 0, "output_tokens": 0}) + assert usage is not None + assert usage.input == 0 + + def test_missing_details_default_to_zero(self) -> None: + usage = lang_chain_span_usage({"input_tokens": 5, "output_tokens": 1}) + assert usage is not None + assert usage.cache_read == 0 + + +# ─── The run accumulator ───────────────────────────────────────────────────── + + +class TestRunUsage: + def test_starts_unreported_and_at_zero(self) -> None: + run = create_run_usage() + assert run.reported is False + assert run.total.input == 0 + + def test_sums_across_turns(self) -> None: + run = create_run_usage() + run.add(SpanUsage(input=10, output=1, cache_read=2, cache_creation=3)) + run.add(SpanUsage(input=20, output=2, cache_read=4, cache_creation=5)) + assert run.total.input == 30 + assert run.total.output == 3 + assert run.total.cache_read == 6 + assert run.total.cache_creation == 8 + + def test_none_is_a_no_op_and_does_not_count_as_reported(self) -> None: + run = create_run_usage() + run.add(None) + assert run.reported is False + + def test_a_genuinely_empty_turn_still_counts_as_reported(self) -> None: + # This is the distinction the failure path needs: "reported zero" is not "reported nothing". + run = create_run_usage() + run.add(SpanUsage()) + assert run.reported is True + assert run.total.input == 0 + + def test_two_accumulators_do_not_share_state(self) -> None: + # A mutable default on the dataclass would make this fail. + first = create_run_usage() + first.add(SpanUsage(input=5)) + assert create_run_usage().total.input == 0 + + +# ─── set_usage_span_attributes ─────────────────────────────────────────────── + + +class TestSetUsageSpanAttributes: + def test_writes_all_seven_attributes(self) -> None: + span = FakeSpan() + set_usage_span_attributes( + span, SpanUsage(input=10, output=4, cache_read=6, cache_creation=1) + ) + assert span.attributes == { + "gen_ai.usage.input_tokens": 10, + "gen_ai.usage.output_tokens": 4, + "gen_ai.usage.total_tokens": 14, + "gen_ai.usage.cache_read.input_tokens": 6, + "gen_ai.usage.cache_creation.input_tokens": 1, + "gen_ai.usage.prompt_tokens": 10, + "gen_ai.usage.completion_tokens": 4, + } + + def test_writes_zeros_rather_than_omitting_them(self) -> None: + # An absent attribute drops the span from every query that groups on usage, which reads as + # "no cached tokens" when it means "this handler forgot to say". + span = FakeSpan() + set_usage_span_attributes(span, SpanUsage()) + assert len(span.attributes) == 7 + assert span.attributes["gen_ai.usage.cache_creation.input_tokens"] == 0 + + def test_total_is_derived_not_taken(self) -> None: + span = FakeSpan() + set_usage_span_attributes(span, SpanUsage(input=7, output=3)) + assert span.attributes["gen_ai.usage.total_tokens"] == 10 + + def test_the_openllmetry_aliases_agree_with_the_canonical_keys(self) -> None: + # They disagreed once, because the alias was computed at a call site off Anthropic's + # cache-excluding input field. One writer, one number. + span = FakeSpan() + set_usage_span_attributes(span, SpanUsage(input=100, output=5, cache_read=80)) + assert ( + span.attributes["gen_ai.usage.prompt_tokens"] + == span.attributes["gen_ai.usage.input_tokens"] + ) + assert ( + span.attributes["gen_ai.usage.completion_tokens"] + == span.attributes["gen_ai.usage.output_tokens"] + ) + + +# ─── set_model_identity_attributes ─────────────────────────────────────────── + + +class TestSetModelIdentityAttributes: + def test_writes_both_provider_keys_with_the_same_value_by_default(self) -> None: + span = FakeSpan() + set_model_identity_attributes(span, "anthropic", "claude-3-5-sonnet") + assert span.attributes["gen_ai.system"] == "anthropic" + assert span.attributes["gen_ai.provider.name"] == "anthropic" + assert span.attributes["gen_ai.request.model"] == "claude-3-5-sonnet" + + def test_the_legacy_key_can_differ_from_the_current_one(self) -> None: + # The LangChain handlers keep the framework name on the old key, because the new key's + # semconv enum has no 'langchain' member. + span = FakeSpan() + set_model_identity_attributes( + span, "anthropic", "claude-3-5-sonnet", "langchain" + ) + assert span.attributes["gen_ai.system"] == "langchain" + assert span.attributes["gen_ai.provider.name"] == "anthropic" + + +# ─── end_span_once ─────────────────────────────────────────────────────────── + + +class TestEndSpanOnce: + def test_ends_the_span(self) -> None: + span: FakeSpan = FakeSpan() + tracker: set[int] = set() + end_span_once(span, tracker) + assert span.ended == 1 + + def test_a_second_call_is_ignored(self) -> None: + # The streaming finally and the success path can both reach the same span. + span: FakeSpan = FakeSpan() + tracker: set[int] = set() + end_span_once(span, tracker) + end_span_once(span, tracker) + assert span.ended == 1 + + def test_marks_abandonment_without_asserting_failure(self) -> None: + span: FakeSpan = FakeSpan() + tracker: set[int] = set() + end_span_once(span, tracker, abandoned=True) + assert span.attributes["launchdarkly.stream.abandoned"] is True + assert span.ended == 1 + + def test_does_not_mark_a_normal_end(self) -> None: + span: FakeSpan = FakeSpan() + tracker: set[int] = set() + end_span_once(span, tracker) + assert "launchdarkly.stream.abandoned" not in span.attributes + + def test_tracks_each_span_separately(self) -> None: + first, second = FakeSpan(), FakeSpan() + tracker: set[int] = set() + end_span_once(first, tracker) + end_span_once(second, tracker) + assert (first.ended, second.ended) == (1, 1) + + def test_an_unhashable_span_is_still_tracked(self) -> None: + # The tracker holds id(span), because an OTel span is not guaranteed hashable. + class Unhashable(FakeSpan): + __hash__ = None # type: ignore[assignment] + + span = Unhashable() + tracker: set[int] = set() + end_span_once(span, tracker) + end_span_once(span, tracker) + assert span.ended == 1 From c60fa7c6ade65542cbbbaefd440eca13e56088a4 Mon Sep 17 00:00:00 2001 From: Alexis Georges Date: Tue, 11 Aug 2026 16:27:56 -0400 Subject: [PATCH 2/5] fix(client): carry the cache breakdown into the public UsageDict parse_usage started reporting a cache breakdown, but invoke and the judge runner both built UsageDict by hand from three keys, so input_details was always None on the blocking path while the streaming path handed back the nested dict. The two paths disagreed about the same run. Both now go through to_usage_dict, so the mapping has one author and cannot drift again. graph.py keeps building its own: a graph total is a sum across nodes and carries no per-call breakdown, and it says so. Found by Bugbot on #28. --- .../src/launchdarkly_ai_server/__init__.py | 2 + .../src/launchdarkly_ai_server/client.py | 8 +-- .../src/launchdarkly_ai_server/judges.py | 8 +-- .../src/launchdarkly_ai_server/utils.py | 27 +++++++++ packages/client/tests/test_span_usage.py | 58 +++++++++++++++++++ 5 files changed, 91 insertions(+), 12 deletions(-) diff --git a/packages/client/src/launchdarkly_ai_server/__init__.py b/packages/client/src/launchdarkly_ai_server/__init__.py index feaeaec..039a521 100644 --- a/packages/client/src/launchdarkly_ai_server/__init__.py +++ b/packages/client/src/launchdarkly_ai_server/__init__.py @@ -76,6 +76,7 @@ set_openllmetry_prompt, set_usage_span_attributes, to_ld_context, + to_usage_dict, ) __all__ = [ # noqa: RUF022 @@ -121,6 +122,7 @@ "number_or_zero", "set_model_identity_attributes", "set_usage_span_attributes", + "to_usage_dict", "VariationMeta", # utils "create_handler", diff --git a/packages/client/src/launchdarkly_ai_server/client.py b/packages/client/src/launchdarkly_ai_server/client.py index 9e6ac4c..570ac0d 100644 --- a/packages/client/src/launchdarkly_ai_server/client.py +++ b/packages/client/src/launchdarkly_ai_server/client.py @@ -15,12 +15,12 @@ ProviderHandler, ProviderResponse, StreamEvent, - UsageDict, VariationMeta, ) from .utils import ( parse_json_with_possible_fences, select_handler, + to_usage_dict, ) @@ -112,11 +112,7 @@ async def invoke( else json.dumps(parsed_response) ) - usage_obj = UsageDict( - input=usage.get("input", 0), - output=usage.get("output", 0), - total=usage.get("total", 0), - ) + usage_obj = to_usage_dict(usage) if self._skip_judges: judge_tasks = await build_judge_tasks( diff --git a/packages/client/src/launchdarkly_ai_server/judges.py b/packages/client/src/launchdarkly_ai_server/judges.py index 9e644c2..898f864 100644 --- a/packages/client/src/launchdarkly_ai_server/judges.py +++ b/packages/client/src/launchdarkly_ai_server/judges.py @@ -13,7 +13,6 @@ NativeTool, ProviderHandler, TrackData, - UsageDict, ) from .utils import ( collapse_messages_to_instructions as _collapse_messages_to_instructions, @@ -22,6 +21,7 @@ normalize_mode, parse_json_with_possible_fences, to_ld_context, + to_usage_dict, ) @@ -409,11 +409,7 @@ def _matches(h: ProviderHandler) -> bool: reasoning = parsed.get("reasoning", "") raw_usage = result["usage"] - usage = UsageDict( - input=raw_usage.get("input", 0), - output=raw_usage.get("output", 0), - total=raw_usage.get("total", 0), - ) + usage = to_usage_dict(raw_usage) merged_track_data: TrackData = { **task.parent_track_data, diff --git a/packages/client/src/launchdarkly_ai_server/utils.py b/packages/client/src/launchdarkly_ai_server/utils.py index 81e83fd..bf14b92 100644 --- a/packages/client/src/launchdarkly_ai_server/utils.py +++ b/packages/client/src/launchdarkly_ai_server/utils.py @@ -8,7 +8,9 @@ from .types import ( AiConfigRep, GraphNode, + InputTokenDetails, ProviderHandler, + UsageDict, VariationMeta, _HandlerFn, _StreamFn, @@ -142,6 +144,31 @@ def parse_usage(usage: dict[str, Any]) -> dict[str, Any]: return {"input": 0, "output": 0, "total": 0} +def to_usage_dict(usage: dict[str, Any]) -> UsageDict: + """Builds the public :class:`UsageDict` from a :func:`parse_usage` result. + + Shared because ``invoke`` and the judge runner both need it and both used to build the dataclass + by hand from three keys, which silently dropped the cache breakdown the moment ``parse_usage`` + started reporting one. A caller reading ``input_details`` off a blocking call got ``None`` while + the streaming path handed back the nested dict, so the two paths disagreed about the same run. + """ + details = usage.get("input_details") + return UsageDict( + input=usage.get("input", 0), + output=usage.get("output", 0), + total=usage.get("total", 0), + input_details=( + InputTokenDetails( + uncached=details.get("uncached", 0), + cache_read=details.get("cache_read", 0), + cache_creation=details.get("cache_creation", 0), + ) + if isinstance(details, dict) + else None + ), + ) + + @dataclass class SpanUsage: """The provider-neutral token counts a span reports, after the caller applied its cache rule. diff --git a/packages/client/tests/test_span_usage.py b/packages/client/tests/test_span_usage.py index 605a21d..dc3df9c 100644 --- a/packages/client/tests/test_span_usage.py +++ b/packages/client/tests/test_span_usage.py @@ -22,6 +22,7 @@ parse_usage, set_model_identity_attributes, set_usage_span_attributes, + to_usage_dict, ) @@ -380,3 +381,60 @@ class Unhashable(FakeSpan): end_span_once(span, tracker) end_span_once(span, tracker) assert span.ended == 1 + + +# ─── to_usage_dict ─────────────────────────────────────────────────────────── + + +class TestToUsageDict: + """The public UsageDict must carry everything parse_usage reported. + + Both call sites used to build the dataclass by hand from three keys, so the cache breakdown + vanished the moment parse_usage started reporting one: a caller reading input_details off a + blocking call got None while the streaming path handed back the nested dict. + """ + + def test_carries_the_three_totals(self) -> None: + usage = to_usage_dict({"input": 10, "output": 5, "total": 15}) + assert (usage.input, usage.output, usage.total) == (10, 5, 15) + + def test_carries_the_cache_breakdown_when_present(self) -> None: + usage = to_usage_dict( + { + "input": 23554, + "output": 10, + "total": 23564, + "input_details": { + "uncached": 3, + "cache_read": 19971, + "cache_creation": 3580, + }, + } + ) + assert usage.input_details is not None + assert usage.input_details.uncached == 3 + assert usage.input_details.cache_read == 19971 + assert usage.input_details.cache_creation == 3580 + + def test_is_none_when_the_provider_reported_no_cache(self) -> None: + assert ( + to_usage_dict({"input": 1, "output": 2, "total": 3}).input_details is None + ) + + def test_round_trips_a_parse_usage_result(self) -> None: + # The two functions are used together, so pin them together. + raw = { + "input_tokens": 3, + "output_tokens": 10, + "cache_read_input_tokens": 19971, + "cache_creation_input_tokens": 3580, + } + usage = to_usage_dict(parse_usage(raw)) + assert usage.input == 23554 + assert usage.input_details is not None + assert usage.input_details.cache_read == 19971 + + def test_a_missing_field_defaults_rather_than_raising(self) -> None: + usage = to_usage_dict({}) + assert (usage.input, usage.output, usage.total) == (0, 0, 0) + assert usage.input_details is None From a1772c1a59c4f45e322d9805986ad494c2ef5bcc Mon Sep 17 00:00:00 2001 From: Alexis Georges Date: Tue, 11 Aug 2026 16:29:08 -0400 Subject: [PATCH 3/5] fix(client): make end_span_once a no-op on a None span Handlers hold None for every span whenever the OpenTelemetry SDK is absent, and every other helper in this family already no-ops on it. This one did not, so a streaming cleanup path in a finally crashed with AttributeError on an install without the otel extra: the one place that should be hardest to break was the one place that was not guarded. Fixed in the shared helper rather than at six call sites, because all six handlers use it the same way and the next handler would hit the same edge. Found by Bugbot on #30. --- packages/client/src/launchdarkly_ai_server/utils.py | 6 ++++++ packages/client/tests/test_span_usage.py | 6 ++++++ 2 files changed, 12 insertions(+) diff --git a/packages/client/src/launchdarkly_ai_server/utils.py b/packages/client/src/launchdarkly_ai_server/utils.py index bf14b92..da27fc1 100644 --- a/packages/client/src/launchdarkly_ai_server/utils.py +++ b/packages/client/src/launchdarkly_ai_server/utils.py @@ -369,7 +369,13 @@ def end_span_once(span: Any, tracker: set[int], abandoned: bool = False) -> None Ending the span is not always enough. Two handlers also hold a vendor generator or run that must be closed or cancelled in the same ``finally``. See TELEMETRY-CONTRACT.md section 6. + + A ``None`` span is a no-op, matching every other helper in this family. Handlers hold ``None`` + whenever the OpenTelemetry SDK is absent, and a cleanup path in a ``finally`` is the last place + that should have to remember it. """ + if span is None: + return key = id(span) if key in tracker: return diff --git a/packages/client/tests/test_span_usage.py b/packages/client/tests/test_span_usage.py index dc3df9c..bde037b 100644 --- a/packages/client/tests/test_span_usage.py +++ b/packages/client/tests/test_span_usage.py @@ -371,6 +371,12 @@ def test_tracks_each_span_separately(self) -> None: end_span_once(second, tracker) assert (first.ended, second.ended) == (1, 1) + def test_a_none_span_is_a_no_op(self) -> None: + # Handlers hold None whenever the OTel SDK is absent, and a cleanup path in a `finally` is + # the last place that should have to remember it. Every sibling helper already no-ops. + end_span_once(None, set()) + end_span_once(None, set(), abandoned=True) + def test_an_unhashable_span_is_still_tracked(self) -> None: # The tracker holds id(span), because an OTel span is not guaranteed hashable. class Unhashable(FakeSpan): From dff27f60915f6ea6fb24e0351c9eb5faec83ff27 Mon Sep 17 00:00:00 2001 From: Alexis Georges Date: Fri, 14 Aug 2026 13:21:46 -0400 Subject: [PATCH 4/5] feat(client): end spans a cancelled run would otherwise strand asyncio.CancelledError inherits from BaseException, deliberately, so a timeout or a task.cancel() walks past every except Exception a handler writes. The blocking paths end their spans only from those clauses, so a cancelled run exports nothing at all. Not a wrong attribute: no span. The root carries the feature_flag event and every launchdarkly.* attribute, so a stranded root means the whole run never reaches AI Config Monitoring. end_unfinished_spans belongs in a finally rather than an except, because the paths an except cannot see are the entire point. Spans are left at UNSET and marked launchdarkly.run.cancelled. Nothing failed: the caller went away. That is the reasoning end_span_once already applies to an abandoned stream, and it keeps the trace agreeing with LaunchDarkly's own metrics, which record neither a success nor an error for a run that never finished. A span another path already ended is skipped. is_recording() is False once end() has run, so this is a no-op on the success path. The OTel SDK makes a second end idempotent but logs it, and that log is what would bury a real leak in noise. This has no TypeScript counterpart and needs none. An aborted request there rejects its promise and the catch catches it. Only Python routes cancellation around a handler's guards. Four tests against a real exporter, since a mock span cannot tell an ended span from an open one. Removing the guard, or marking ERROR instead of UNSET, each fails one. Found by Bugbot on the langchain-messages layer, then found in five of six handlers by auditing the rest. --- .../src/launchdarkly_ai_server/__init__.py | 2 + .../src/launchdarkly_ai_server/utils.py | 32 ++++++++ packages/client/tests/test_span_usage.py | 79 +++++++++++++++++++ 3 files changed, 113 insertions(+) diff --git a/packages/client/src/launchdarkly_ai_server/__init__.py b/packages/client/src/launchdarkly_ai_server/__init__.py index 039a521..414a985 100644 --- a/packages/client/src/launchdarkly_ai_server/__init__.py +++ b/packages/client/src/launchdarkly_ai_server/__init__.py @@ -63,6 +63,7 @@ create_handler, create_run_usage, end_span_once, + end_unfinished_spans, lang_chain_span_usage, make_track_data, normalize_mode, @@ -118,6 +119,7 @@ "add_cached_tokens_to_input", "create_run_usage", "end_span_once", + "end_unfinished_spans", "lang_chain_span_usage", "number_or_zero", "set_model_identity_attributes", diff --git a/packages/client/src/launchdarkly_ai_server/utils.py b/packages/client/src/launchdarkly_ai_server/utils.py index da27fc1..d8d0751 100644 --- a/packages/client/src/launchdarkly_ai_server/utils.py +++ b/packages/client/src/launchdarkly_ai_server/utils.py @@ -385,6 +385,38 @@ def end_span_once(span: Any, tracker: set[int], abandoned: bool = False) -> None span.end() +def end_unfinished_spans(*spans: Any) -> None: + """Ends every span still open, for an exception no ``except Exception`` can catch. + + ``asyncio.CancelledError`` inherits from ``BaseException``, deliberately, so a timeout or a + ``task.cancel()`` walks straight past every ``except Exception`` a handler writes. The blocking + paths ended their spans only from those clauses, so a cancelled run exported nothing at all: not a + wrong attribute, no span. The root carries the ``feature_flag`` event and every ``launchdarkly.*`` + attribute, so a stranded root means the whole run never reaches AI Config Monitoring. + + Call this from a ``finally``, not from an ``except``. The point is the paths an ``except`` cannot + see. + + Spans are left at ``UNSET`` and marked ``launchdarkly.run.cancelled``. Nothing failed: the caller + went away. That is the same reasoning :func:`end_span_once` applies to an abandoned stream, and it + keeps the trace agreeing with LaunchDarkly's own metrics, which record neither a success nor an + error for a run that never finished. + + A span another path already ended is skipped. ``is_recording()`` is ``False`` once ``end()`` has + run, so this is a no-op on the success path and cannot end a span twice. Setting an attribute after + ``end()`` logs a warning, which is why the check comes first. A ``None`` span is a no-op, matching + every other helper in this family. + + This has no TypeScript counterpart, and needs none. An aborted request there rejects its promise + and the ``catch`` catches it. Only Python routes cancellation around a handler's guards. + """ + for span in spans: + if span is None or not span.is_recording(): + continue + span.set_attribute("launchdarkly.run.cancelled", True) + span.end() + + def parse_template(template: str, variables: dict[str, Any]) -> str: """ Replaces ``{{variable}}`` placeholders in *template* with values from diff --git a/packages/client/tests/test_span_usage.py b/packages/client/tests/test_span_usage.py index bde037b..98f6fcd 100644 --- a/packages/client/tests/test_span_usage.py +++ b/packages/client/tests/test_span_usage.py @@ -12,6 +12,8 @@ from typing import Any +import pytest + from launchdarkly_ai_server.utils import ( SpanUsage, add_cached_tokens_to_input, @@ -444,3 +446,80 @@ def test_a_missing_field_defaults_rather_than_raising(self) -> None: usage = to_usage_dict({}) assert (usage.input, usage.output, usage.total) == (0, 0, 0) assert usage.input_details is None + + +@pytest.fixture +def tracer_and_exporter(): + """A real tracer and exporter, not a mock. + + A mock span cannot tell a span that was ended from one that was not, which is the only thing these + tests are about. + """ + from opentelemetry.sdk.trace import TracerProvider + from opentelemetry.sdk.trace.export import SimpleSpanProcessor + from opentelemetry.sdk.trace.export.in_memory_span_exporter import ( + InMemorySpanExporter, + ) + + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + yield provider.get_tracer("test"), exporter + provider.shutdown() + + +class TestEndUnfinishedSpans: + """TELEMETRY-CONTRACT.md section 6: a `finally` owns every end.""" + + def test_ends_a_span_a_cancelled_run_left_open(self, tracer_and_exporter) -> None: + # asyncio.CancelledError is a BaseException, so `except Exception` never sees it. Without this + # helper a cancelled run exported no span at all, and the root carries the feature_flag event. + from launchdarkly_ai_server import end_unfinished_spans + + tracer, exporter = tracer_and_exporter + span = tracer.start_span("invoke_agent") + + end_unfinished_spans(span) + + [finished] = exporter.get_finished_spans() + assert finished.name == "invoke_agent" + assert finished.attributes["launchdarkly.run.cancelled"] is True + + def test_leaves_the_span_unset_rather_than_error(self, tracer_and_exporter) -> None: + # Nothing failed. The caller went away. Marking ERROR would disagree with LaunchDarkly's own + # metrics, which record neither a success nor an error for a run that never finished. + from opentelemetry.trace import StatusCode + + from launchdarkly_ai_server import end_unfinished_spans + + tracer, exporter = tracer_and_exporter + end_unfinished_spans(tracer.start_span("chat")) + + [finished] = exporter.get_finished_spans() + assert finished.status.status_code is StatusCode.UNSET + + def test_skips_a_span_another_path_already_ended( + self, tracer_and_exporter, caplog + ) -> None: + # This runs on the success path too, where every span is already closed. The OTel SDK makes a + # second end idempotent but logs it, and that log is the only observable difference: it is how + # a real double-end shows up in production, and it would bury a genuine leak in noise. + import logging + + from launchdarkly_ai_server import end_unfinished_spans + + tracer, exporter = tracer_and_exporter + span = tracer.start_span("chat") + span.end() + + with caplog.at_level(logging.WARNING, logger="opentelemetry.sdk.trace"): + end_unfinished_spans(span) + + assert caplog.records == [] + [finished] = exporter.get_finished_spans() + assert "launchdarkly.run.cancelled" not in (finished.attributes or {}) + + def test_ignores_none_so_a_finally_need_not_check(self) -> None: + from launchdarkly_ai_server import end_unfinished_spans + + end_unfinished_spans(None, None) From 14ed0668708bac31f110451cdf25d748e4076063 Mon Sep 17 00:00:00 2001 From: Alexis Georges Date: Fri, 14 Aug 2026 14:30:35 -0400 Subject: [PATCH 5/5] feat(client): tell a cancelled stream apart from an abandoned one end_span_once marked every unfinished streaming span launchdarkly.stream.abandoned, whatever ended it. Both facts can be true of the same unwind and only one of them is the reason. A consumer that stops reading abandoned the stream, and that is normal. A CancelledError means something cancelled the run, usually a timeout, and the consumer never chose anything. Reporting the second as the first made a timed-out stream and a timed-out blocking call disagree in AI Config Monitoring about why they stopped, since the blocking paths already report launchdarkly.run.cancelled. cancelled wins over abandoned when both are passed, because the cancellation is the cause and the abandonment is a consequence of it. No new attribute. Both keys already exist and are already in the vocabulary lock. Two tests. Removing the branch fails one. Found by Bugbot on the openai-agents layer. --- .../src/launchdarkly_ai_server/utils.py | 14 ++++++-- packages/client/tests/test_span_usage.py | 32 +++++++++++++++++++ 2 files changed, 44 insertions(+), 2 deletions(-) diff --git a/packages/client/src/launchdarkly_ai_server/utils.py b/packages/client/src/launchdarkly_ai_server/utils.py index d8d0751..0eb2c21 100644 --- a/packages/client/src/launchdarkly_ai_server/utils.py +++ b/packages/client/src/launchdarkly_ai_server/utils.py @@ -348,7 +348,9 @@ def set_model_identity_attributes( span.set_attribute("gen_ai.request.model", request_model) -def end_span_once(span: Any, tracker: set[int], abandoned: bool = False) -> None: +def end_span_once( + span: Any, tracker: set[int], abandoned: bool = False, cancelled: bool = False +) -> None: """Ends a span exactly once, even when the caller cannot know whether an earlier path ended it. The streaming handlers need this: a consumer that breaks out of ``async for``, or throws inside @@ -360,6 +362,12 @@ def end_span_once(span: Any, tracker: set[int], abandoned: bool = False) -> None *tracker* holds ``id(span)`` rather than the span itself, because an OTel span is not guaranteed hashable across implementations. + *cancelled* wins over *abandoned*, because both can be true of the same unwind and only one of + them is the reason. A consumer that stops reading abandoned the stream. A ``CancelledError`` + means something cancelled the run, usually a timeout, and the consumer never chose anything. The + blocking paths report that second case as ``launchdarkly.run.cancelled``, so the streaming paths + say the same thing rather than calling a timed-out run an abandoned one. + An abandoned span is marked with ``launchdarkly.stream.abandoned`` and deliberately left at ``UNSET`` rather than ``ERROR``. Stopping early is a normal thing for a consumer to do, such as rendering enough of a response and moving on, and nothing failed. Marking it ``ERROR`` would @@ -380,7 +388,9 @@ def end_span_once(span: Any, tracker: set[int], abandoned: bool = False) -> None if key in tracker: return tracker.add(key) - if abandoned: + if cancelled: + span.set_attribute("launchdarkly.run.cancelled", True) + elif abandoned: span.set_attribute("launchdarkly.stream.abandoned", True) span.end() diff --git a/packages/client/tests/test_span_usage.py b/packages/client/tests/test_span_usage.py index 98f6fcd..197f48d 100644 --- a/packages/client/tests/test_span_usage.py +++ b/packages/client/tests/test_span_usage.py @@ -468,6 +468,38 @@ def tracer_and_exporter(): provider.shutdown() +class TestEndSpanOnceDistinguishesWhyItStopped: + """TELEMETRY-CONTRACT.md section 6.""" + + def test_a_cancelled_stream_is_not_reported_as_abandoned( + self, tracer_and_exporter + ) -> None: + # Both can be true of the same unwind, and only one is the reason. A consumer that stops + # reading abandoned the stream. A CancelledError means something cancelled the run, usually a + # timeout, and the consumer never chose anything. Reporting the second as the first made a + # timed-out stream and a timed-out blocking call disagree about why they stopped. + from launchdarkly_ai_server import end_span_once + + tracer, exporter = tracer_and_exporter + end_span_once(tracer.start_span("chat"), set(), abandoned=True, cancelled=True) + + [finished] = exporter.get_finished_spans() + assert finished.attributes["launchdarkly.run.cancelled"] is True + assert "launchdarkly.stream.abandoned" not in finished.attributes + + def test_a_consumer_that_stops_reading_is_still_abandoned( + self, tracer_and_exporter + ) -> None: + from launchdarkly_ai_server import end_span_once + + tracer, exporter = tracer_and_exporter + end_span_once(tracer.start_span("chat"), set(), abandoned=True) + + [finished] = exporter.get_finished_spans() + assert finished.attributes["launchdarkly.stream.abandoned"] is True + assert "launchdarkly.run.cancelled" not in finished.attributes + + class TestEndUnfinishedSpans: """TELEMETRY-CONTRACT.md section 6: a `finally` owns every end."""