diff --git a/src/xagent/core/model/chat/__init__.py b/src/xagent/core/model/chat/__init__.py index fa9e092989..4d4b465269 100644 --- a/src/xagent/core/model/chat/__init__.py +++ b/src/xagent/core/model/chat/__init__.py @@ -4,9 +4,15 @@ from .basic.base import BaseLLM from .timeout_config import TimeoutConfig from .token_context import ( + MediaCallType, + MediaUnit, TokenContextManager, TokenUsage, + add_media_usage, add_token_usage, + aggregate_media_usage_by_model, + aggregate_token_usage_by_model, + estimate_media_tokens, get_and_reset_token_usage, get_token_usage, reset_token_usage, @@ -25,7 +31,13 @@ # Token tracking "TokenUsage", "TokenContextManager", + "MediaUnit", + "MediaCallType", "add_token_usage", + "add_media_usage", + "aggregate_token_usage_by_model", + "aggregate_media_usage_by_model", + "estimate_media_tokens", "get_token_usage", "reset_token_usage", "get_and_reset_token_usage", diff --git a/src/xagent/core/model/chat/token_context.py b/src/xagent/core/model/chat/token_context.py index 7e9b59a53f..d47c8a1e5d 100644 --- a/src/xagent/core/model/chat/token_context.py +++ b/src/xagent/core/model/chat/token_context.py @@ -1,17 +1,136 @@ -"""Token usage tracking using contextvars. +"""Usage tracking using contextvars, for LLM tokens and non-LLM media calls. -This module provides a thread-safe way to track token usage across LLM calls -without modifying function signatures. Using contextvars allows the token -statistics to be automatically collected during task execution. +Tracks usage across calls without threading it through function signatures: +contextvars let statistics be collected automatically during task execution. + +Two dimensions live here. LLM tokens use ``add_token_usage`` and aggregate via +``aggregate_token_usage_by_model``. Non-LLM media calls — image, video, TTS, +ASR, music, sound effect, embedding, rerank — use the ``MediaUnit`` / +``MediaCallType`` vocabulary, record through ``add_media_usage``, and aggregate +via ``aggregate_media_usage_by_model``. Both write into the same +``TokenUsage.details`` list, discriminated by each row's ``type``, so existing +persistence and quota paths carry media rows with no schema change. """ import contextvars import logging +import math +import threading from dataclasses import dataclass, field +from enum import Enum from typing import Any, Dict, List, Optional logger = logging.getLogger(__name__) +# Sentinel: "argument not supplied", distinct from any value a caller could pass. +_UNSET: Any = object() + + +class MediaUnit(str, Enum): + """Billable dimension of a non-LLM media call. + + One unit per modality, chosen so a given (model, call_type) always reports + the same unit regardless of how complete the provider's response was — a + price table keyed on (model, unit) is only usable if the unit is stable. + ``REQUESTS`` means exactly one provider call and always carries + ``quantity=1``; use it only when the modality genuinely has no finer + billable dimension, never as a degraded fallback for a missing measurement. + """ + + IMAGES = "images" + SECONDS = "seconds" + CHARACTERS = "characters" + TEXTS = "texts" + REQUESTS = "requests" + + +class MediaCallType(str, Enum): + """Modality/operation that produced a media usage entry.""" + + GENERATE_IMAGE = "generate_image" + EDIT_IMAGE = "edit_image" + VIDEO = "video" + TTS = "tts" + ASR = "asr" + MUSIC = "music" + SOUND_EFFECT = "sound_effect" + EMBEDDING = "embedding" + RERANK = "rerank" + + +#: The billable unit each modality reports. This is what makes a price table +#: keyed on (model, unit) usable: the unit is a property of the modality, never +#: of how complete a particular response happened to be. A duration-billed call +#: whose length is unknown records ``seconds`` with ``quantity=0``, it does not +#: fall back to ``requests``. +#: +#: Enforced at the write boundary rather than only documented — the invariant was +#: stated in three docstrings and still violated by six call sites in this +#: module's own tests (``images`` paired with ``video``). +MEDIA_UNIT_BY_CALL_TYPE: Dict[str, "MediaUnit"] = { + MediaCallType.GENERATE_IMAGE.value: MediaUnit.IMAGES, + MediaCallType.EDIT_IMAGE.value: MediaUnit.IMAGES, + MediaCallType.VIDEO.value: MediaUnit.SECONDS, + MediaCallType.ASR.value: MediaUnit.SECONDS, + MediaCallType.MUSIC.value: MediaUnit.SECONDS, + MediaCallType.SOUND_EFFECT.value: MediaUnit.SECONDS, + MediaCallType.TTS.value: MediaUnit.CHARACTERS, + MediaCallType.EMBEDDING.value: MediaUnit.TEXTS, + MediaCallType.RERANK.value: MediaUnit.REQUESTS, +} + +_MEDIA_UNIT_VALUES = frozenset(member.value for member in MediaUnit) +_MEDIA_CALL_TYPE_VALUES = frozenset(member.value for member in MediaCallType) + + +def _validated_media_unit(unit: "MediaUnit | str | None") -> str: + """Normalise ``unit`` to a known :class:`MediaUnit` value. + + A typo'd unit silently mints a new billing dimension that + :func:`aggregate_media_usage_by_model` will happily key off, and a usage + record cannot be repaired retroactively once written. Rejecting at the + write boundary is the only point where the error is still fixable. + """ + if unit is None: + raise ValueError("Media unit cannot be None") + value = unit.value if isinstance(unit, MediaUnit) else str(unit) + if value not in _MEDIA_UNIT_VALUES: + raise ValueError( + f"Unknown media unit {value!r}; expected one of " + f"{sorted(_MEDIA_UNIT_VALUES)}" + ) + return value + + +def _validated_media_call_type(call_type: "MediaCallType | str | None") -> str: + """Normalise ``call_type`` to a known :class:`MediaCallType` value. + + Empty is allowed: ``call_type`` is optional metadata rather than a billing + dimension on its own, and omitting it is a legitimate caller choice. + """ + if call_type is None: + return "" + value = call_type.value if isinstance(call_type, MediaCallType) else str(call_type) + if value and value not in _MEDIA_CALL_TYPE_VALUES: + raise ValueError( + f"Unknown media call type {value!r}; expected one of " + f"{sorted(_MEDIA_CALL_TYPE_VALUES)}" + ) + return value + + +def copy_detail_rows(raw_details: Any) -> List[Dict]: + """Detached copies of the dict rows in ``raw_details``, non-dicts dropped. + + One helper for every place a details list is handed across a boundary: + sharing the inner dicts lets a consumer mutate live usage state, and the + ``isinstance`` filter keeps a malformed legacy row (``details`` is persisted + as free-form JSON) from raising in an accounting path. + """ + if not isinstance(raw_details, list): + return [] + return [dict(item) for item in raw_details if isinstance(item, dict)] + @dataclass class TokenUsage: @@ -21,6 +140,8 @@ class TokenUsage: input_tokens: Number of tokens in prompts sent to LLM output_tokens: Number of tokens generated by LLM llm_calls: Number of LLM API calls made + media_calls: Number of non-LLM media model calls made (image/video/ + tts/asr/embedding/rerank/...) details: Detailed breakdown by model/call type """ @@ -29,11 +150,82 @@ class TokenUsage: llm_calls: int = 0 tool_calls: int = 0 details: List[Dict] = field(default_factory=list) + # Appended AFTER the legacy fields, and keyword-only, so the historical + # positional signature — TokenUsage(input, output, llm_calls, tool_calls, + # details) — keeps binding the way existing callers expect. Inserting it + # before tool_calls silently rebound their 4th and 5th arguments (tool + # count became media count, details became tool count) with no error. + media_calls: int = field(default=0, kw_only=True) + + def __post_init__(self) -> None: + # Counter updates are read-modify-write, and one TokenUsage is routinely + # shared across worker threads (RAG ingestion pools, and the executor hops + # the producer PRs add); without it a concurrent read can observe a + # counter that disagrees with its own detail rows. + # + # Deliberately NOT a dataclass field. As a field it would sit in + # ``__dataclass_fields__``, and ``dataclasses.asdict`` walks those + # directly (ignoring ``__getstate__``), so it would raise + # ``TypeError: cannot pickle '_thread.lock' object`` for any generic + # consumer. Keeping it in ``__dict__`` also leaves the public + # positional signature — ``TokenUsage(input, output, llm_calls, + # tool_calls, details)`` — free of a private slot. + self.__dict__["_lock"] = threading.RLock() + # Normalise once here so the constructor and from_dict both hand this + # object a private, dict-only list. Storing a caller's list by reference + # would let them mutate rows outside the lock, and a non-dict row would + # misalign the index-based delta slice. + self.details = copy_detail_rows(self.details) + + @property + def _lock(self) -> "threading._RLock": + """The per-instance mutation lock (see :meth:`__post_init__`). + + Reentrant so a future nested acquisition cannot self-deadlock: the + current no-nesting discipline (``to_dict`` inlines the token total rather + than calling ``total_tokens``, ``merge`` takes the two locks + sequentially) is invisible to a later editor. + """ + # threading.RLock is a factory, not a class, so the annotation names the + # object it returns (typeshed exposes it as threading._RLock). + lock: threading._RLock = self.__dict__["_lock"] + return lock + + def __getstate__(self) -> Dict[str, Any]: + """Drop the lock: it is not picklable and is per-instance state.""" + state = self.__dict__.copy() + state.pop("_lock", None) + return state + + def __setstate__(self, state: Dict[str, Any]) -> None: + """Restore, giving the revived instance its own fresh lock.""" + self.__dict__.update(state) + self.__dict__["_lock"] = threading.RLock() + + def __copy__(self) -> "TokenUsage": + """Shallow copy is unsupported; delegate to :meth:`snapshot`. + + The default shallow copy routes through the ``__getstate__`` pair above, + so the new instance gets a *different* lock while still referencing the + *same* ``details`` list. Two objects each believing they hold exclusive + access would then mutate one shared list under two different locks — + exactly the mutual exclusion the lock exists to provide. ``snapshot()`` + is the supported way to fork a usage object, so ``copy.copy`` is routed + there rather than left as a trap. + """ + return self.snapshot() + + def __deepcopy__(self, memo: Dict[int, Any]) -> "TokenUsage": + """Deep copy via :meth:`snapshot` so the copy is taken under the lock.""" + copied = self.snapshot() + memo[id(self)] = copied + return copied @property def total_tokens(self) -> int: """Total tokens used (input + output).""" - return self.input_tokens + self.output_tokens + with self._lock: + return self.input_tokens + self.output_tokens def add_input_tokens( self, @@ -51,62 +243,228 @@ def add_input_tokens( ``cache_write_tokens`` is the subset of ``tokens`` written to the cache this call (Claude bills these at a premium); 0 when unknown. """ - self.input_tokens += tokens - if model or call_type: - self.details.append( - { - "type": "input", - "tokens": tokens, - "cached_tokens": cached_tokens, - "cache_write_tokens": cache_write_tokens, - "model": model, - "model_id": model_id, - "call_type": call_type, - } - ) + with self._lock: + self.input_tokens += tokens + if model or call_type: + self.details.append( + { + "type": "input", + "tokens": tokens, + "cached_tokens": cached_tokens, + "cache_write_tokens": cache_write_tokens, + "model": model, + "model_id": model_id, + "call_type": call_type, + } + ) def add_output_tokens( self, tokens: int, model: str = "", call_type: str = "", model_id: str = "" ) -> None: """Add output tokens from a completion.""" - self.output_tokens += tokens - if model or call_type: + with self._lock: + self.output_tokens += tokens + if model or call_type: + self.details.append( + { + "type": "output", + "tokens": tokens, + "model": model, + "model_id": model_id, + "call_type": call_type, + } + ) + + def record_media_call( + self, + unit: "MediaUnit | str | None", + quantity: float, + model: str = "", + call_type: "MediaCallType | str | None" = "", + model_id: str = "", + input_tokens: int = 0, + output_tokens: int = 0, + resolution: str = "", + tokens_estimated: bool = False, + _raw_quantity: Any = _UNSET, + ) -> None: + """Count the call and append its detail under one lock acquisition. + + The only media write path. Bumping the counter and appending the row + separately would take the lock twice, letting a concurrent + ``to_dict``/``merge``/``snapshot`` land between them and observe a + counter with no matching detail row (or the reverse). Consumers derive + per-model rows from ``details`` and the call count from the scalar, so a + torn pair reports mutually inconsistent billing. + """ + unit_value = _validated_media_unit(unit) + call_type_value = _validated_media_call_type(call_type) + # Validated here rather than only in the module-level + # ``add_media_usage``: this is the actual write boundary, so a direct + # caller holding a TokenUsage gets the same guarantees instead of + # persisting a boolean, negative, or non-finite quantity that cannot be + # repaired later. + # _raw_quantity lets the module-level add_media_usage, which coerces one + # layer up, report what its own caller actually passed. + raw_quantity = quantity if _raw_quantity is _UNSET else _raw_quantity + quantity = _coerce_float(quantity) + # REQUESTS is defined as exactly one provider call, so its quantity is + # not a free variable. Letting 0/0.5/2 through would make the billable + # quantity disagree with the media_calls/aggregate `calls` count derived + # from the same row — quota and pricing would then read different + # numbers off one record. Raised rather than clamped: a caller passing + # something else has a real bug, and silently rewriting it hides that. + # Producers route through ``media_usage.record_media_usage``, which + # swallows this so an accounting bug still cannot break a media call. + # Rejecting before the lock is taken means a bad call leaves no state at + # all — neither a counter bump nor an orphan row. + expected_unit = MEDIA_UNIT_BY_CALL_TYPE.get(call_type_value) + if expected_unit is not None and unit_value != expected_unit.value: + raise ValueError( + f"call_type {call_type_value!r} bills in " + f"{expected_unit.value!r}, not {unit_value!r}; the unit is a " + f"property of the modality, so a (model, unit) price table " + f"breaks if one modality reports two units" + ) + if unit_value == MediaUnit.REQUESTS.value and quantity != 1.0: + raise ValueError( + f"MediaUnit.REQUESTS means exactly one call, so quantity must " + f"be 1; got {raw_quantity!r}" + ) + input_tokens = _coerce_int(input_tokens) + output_tokens = _coerce_int(output_tokens) + with self._lock: + self.media_calls += 1 self.details.append( { - "type": "output", - "tokens": tokens, + "type": "media", + "unit": unit_value, + "quantity": quantity, + "provider_tokens": input_tokens + output_tokens, + "provider_input_tokens": input_tokens, + "provider_output_tokens": output_tokens, + "tokens_estimated": tokens_estimated, "model": model, "model_id": model_id, - "call_type": call_type, + "call_type": call_type_value, + "resolution": resolution, } ) + def detail_tail(self, start: int = 0) -> tuple[List[Dict], int]: + """Detail rows from ``start`` onward plus the tool-call count, atomically. + + The narrow read that per-turn delta computation actually needs. Both + values come from one lock acquisition, so they describe the same instant + — a concurrent ``record_media_call``/``increment_tool_calls`` cannot land + between them and yield a pair from two different logical points in time. + + Prefer this over ``snapshot()`` whenever only the tail is wanted. + ``snapshot()`` deep-copies the whole cumulative ``details`` list, which + grows monotonically across turns (seeds are restored from the persisted + list), so using it to read a small tail costs O(total usage) instead of + O(delta) — and holds the lock that serialises every LLM adapter's token + write for the duration of that copy. On the quota-gate path, polled once + per agent step *and* once per streamed LLM chunk, that difference is + three orders of magnitude at a few thousand accumulated rows. + + Rows are copied so the caller cannot mutate live state through them. + """ + with self._lock: + return ( + copy_detail_rows(self.details[start:]), + self.tool_calls, + ) + + def snapshot(self) -> "TokenUsage": + """A detached copy taken under the lock. + + Reading the fields individually — as an external copier must — can + interleave with a concurrent ``record_media_call``/``merge`` and produce + a torn snapshot: a counter without its matching detail row, or the + reverse. Consumers derive per-model rows from ``details`` and the call + count from the scalar, so a torn pair reports mutually inconsistent + billing. Taking the lock once here is the only way a copier can get a + self-consistent view. + """ + with self._lock: + copy = TokenUsage( + input_tokens=self.input_tokens, + output_tokens=self.output_tokens, + llm_calls=self.llm_calls, + tool_calls=self.tool_calls, + media_calls=self.media_calls, + ) + # Assigned after construction, not passed in: __post_init__ would + # copy the list a second time, doubling the cost of a copy that is + # already O(total details). + copy.details = copy_detail_rows(self.details) + return copy + def increment_llm_calls(self) -> None: """Increment the LLM call counter.""" - self.llm_calls += 1 + with self._lock: + self.llm_calls += 1 def increment_tool_calls(self, count: int = 1) -> None: - """Increment the tool-call counter (one per tool invocation).""" - self.tool_calls += count + """Increment the tool-call counter (one per tool invocation). + + Negative counts are ignored: they would drive the counter below the + number of matching detail rows, and no caller has a reason to decrement + a monotonic usage counter. + """ + if count < 0: + logger.warning("Ignoring negative tool call increment: %r", count) + return + with self._lock: + self.tool_calls += count def merge(self, other: "TokenUsage") -> None: """Merge another TokenUsage into this one.""" - self.input_tokens += other.input_tokens - self.output_tokens += other.output_tokens - self.llm_calls += other.llm_calls - self.tool_calls += other.tool_calls - self.details.extend(other.details) + # Snapshot ``other`` under its own lock and release it before taking + # ours: holding both at once would deadlock a concurrent ``b.merge(a)``. + with other._lock: + input_tokens = other.input_tokens + output_tokens = other.output_tokens + llm_calls = other.llm_calls + media_calls = other.media_calls + tool_calls = other.tool_calls + # dict(item), not list(...): sharing the inner dicts would leave the + # merged rows aliased to the source's, so mutating one usage object + # would silently rewrite the other's billing rows. Same reason + # to_dict and snapshot copy each entry. + details = copy_detail_rows(other.details) + with self._lock: + self.input_tokens += input_tokens + self.output_tokens += output_tokens + self.llm_calls += llm_calls + self.media_calls += media_calls + self.tool_calls += tool_calls + self.details.extend(details) def to_dict(self) -> Dict: """Convert to dictionary for serialization.""" - return { - "input_tokens": self.input_tokens, - "output_tokens": self.output_tokens, - "total_tokens": self.total_tokens, - "llm_calls": self.llm_calls, - "tool_calls": self.tool_calls, - "details": self.details, - } + with self._lock: + input_tokens = self.input_tokens + output_tokens = self.output_tokens + return { + "input_tokens": input_tokens, + "output_tokens": output_tokens, + "total_tokens": input_tokens + output_tokens, + "llm_calls": self.llm_calls, + "media_calls": self.media_calls, + "tool_calls": self.tool_calls, + # dict(item), not just list(...): a new outer list around the + # same inner dicts lets a caller mutate the live usage object + # through the returned payload, entirely outside this lock. + # Matches snapshot(), which already does this. + # copy_detail_rows also drops non-dict rows, which the + # previous inline copy here did not. Deliberate: it aligns this + # with snapshot/merge, and a malformed legacy row passed through + # here would fail later at JSON serialisation rather than being + # skipped at the boundary. + "details": copy_detail_rows(self.details), + } @classmethod def from_dict(cls, data: Dict) -> "TokenUsage": @@ -115,6 +473,7 @@ def from_dict(cls, data: Dict) -> "TokenUsage": input_tokens=data.get("input_tokens", 0), output_tokens=data.get("output_tokens", 0), llm_calls=data.get("llm_calls", 0), + media_calls=data.get("media_calls", 0), tool_calls=data.get("tool_calls", 0), details=data.get("details", []), ) @@ -196,18 +555,131 @@ def get_token_usage() -> TokenUsage: def _coerce_int(value: Any) -> int: - """Best-effort int; 0 if the value isn't a usable number (e.g. None/mock).""" + """A finite, non-negative int; 0 if the value isn't a usable token count. + + Token counts share the quantity guarantees documented on + :func:`_coerce_float`, for the same reason: this is the boundary before the + value is persisted into task and quota details, where it cannot be + repaired. Specifically: + + * ``bool`` — ``True``/``False`` are a caller bug, not a count of 1 or 0. + * negatives — a negative token count would subtract from a bill. + * ``NaN``/``inf`` — ``int(float("inf"))`` raises ``OverflowError``, which + would propagate out of the accounting path and drop the whole (billable) + media row rather than just the bad field. + """ if isinstance(value, bool): - return int(value) + logger.warning("Discarding boolean token count: %r", value) + return 0 try: - return int(value) - except (TypeError, ValueError): + result = int(value) + except (TypeError, ValueError, OverflowError): # None/absent is expected (provider omitted the field) — stay quiet. # A present-but-malformed value signals a provider-adapter bug worth # surfacing rather than silently billing it as zero. if value is not None: logger.warning("Discarding non-numeric token count: %r", value) return 0 + if result < 0: + logger.warning("Discarding negative token count: %r", value) + return 0 + return result + + +def estimate_media_tokens(text: Any) -> int: + """Language-aware token estimate for providers that report no usage. + + CJK characters are roughly one token each, while Latin script averages + about four characters per token; a flat chars/4 heuristic therefore + undercounts Chinese text by close to 4x. Accepts a string or an iterable of + strings and ignores anything else, so a malformed input can never raise in + an accounting path. Callers must pass ``tokens_estimated=True`` alongside + the result so billing can tell this apart from a measured count. + """ + if isinstance(text, str): + items: List[str] = [text] + else: + # Any iterable of strings, including generators: the docstring promises + # that, and an embedding producer passing a generator would otherwise + # silently estimate 0 tokens for a real batch. Non-iterables and + # non-string members are ignored rather than raising, since this runs in + # an accounting path. + try: + items = [item for item in text if isinstance(item, str)] + except Exception as e: # noqa: BLE001 + # Not just TypeError: a custom iterable's __iter__/__next__ can + # raise anything, and this runs in an accounting path whose + # documented contract is that malformed input never breaks the call + # being measured. Estimate 0 and say so. + logger.warning("Token estimation failed for %r: %s", type(text), e) + return 0 + + cjk = 0 + other = 0 + for item in items: + for char in item: + # CJK Unified Ideographs, Japanese kana, and Hangul syllables — + # all roughly one token per character. + code = ord(char) + if ( + 0x4E00 <= code <= 0x9FFF # CJK Unified Ideographs + or 0x3400 <= code <= 0x4DBF # CJK Ext-A + or 0x3000 <= code <= 0x303F # CJK punctuation (、。「」etc.) + or 0x3040 <= code <= 0x30FF # Japanese kana + or 0xAC00 <= code <= 0xD7AF # Hangul syllables + or 0xFF00 <= code <= 0xFFEF # Fullwidth / halfwidth forms + ): + cjk += 1 + else: + other += 1 + # Round up rather than truncate: `other // 4` alone estimates zero for any + # 1-3 character Latin string, so short non-empty text would bill nothing. + # A non-empty input must never estimate 0 tokens. + latin_tokens = -(-other // 4) # ceil division + return cjk + latin_tokens + + +def _coerce_float(value: Any) -> float: + """A finite, non-negative float; 0.0 if the value isn't a usable quantity. + + Media quantities can be fractional (e.g. audio seconds), so quantity uses + this rather than ``_coerce_int``. + + Rejects what a billable quantity can never be, because this is the last + boundary before the value is persisted and it cannot be repaired + afterwards: + + * ``bool`` — ``True``/``False`` are almost certainly a caller bug, not a + quantity of 1 or 0. + * negatives — no modality can consume a negative amount, and a negative + would subtract from a bill. + * ``NaN``/``inf`` — these are not JSON-serialisable, so they would produce + literal ``NaN``/``Infinity`` tokens in the persisted task and quota + details that strict parsers reject. + + A rejected value records 0.0, matching the "call happened but is + unmeasured" convention rather than dropping the record entirely. + """ + if isinstance(value, bool): + logger.warning("Discarding boolean media quantity: %r", value) + return 0.0 + try: + result = float(value) + except (TypeError, ValueError, OverflowError): + # OverflowError, not just TypeError/ValueError: float(10**400) raises it, + # and an uncaught raise here would propagate out and drop the whole + # billing row — the opposite of this function's reject-to-0.0 contract. + # _coerce_int already catches it; these two must stay in step. + if value is not None: + logger.warning("Discarding non-numeric media quantity: %r", value) + return 0.0 + if not math.isfinite(result): + logger.warning("Discarding non-finite media quantity: %r", value) + return 0.0 + if result < 0: + logger.warning("Discarding negative media quantity: %r", value) + return 0.0 + return result def _usage_field(usage: Any, name: str) -> Any: @@ -248,6 +720,12 @@ def aggregate_token_usage_by_model(details: Any) -> List[Dict[str, Any]]: merged into an id-backed group only when that name identifies exactly one configured model. Entries without either an id or name are retained as an unattributed group rather than silently dropping tokens from the breakdown. + + Pass a detached list. This iterates ``details`` without any lock, so handing + in a live ``TokenUsage.details`` while another thread appends risks a + "list changed size during iteration" RuntimeError. Use + ``TokenUsage.snapshot().details`` (or a value read from the DB column, as + the API path does) rather than the live list. """ if not isinstance(details, list): return [] @@ -338,6 +816,95 @@ def aggregate_token_usage_by_model(details: Any) -> List[Dict[str, Any]]: ] +def aggregate_media_usage_by_model(details: Any) -> List[Dict[str, Any]]: + """Aggregate ``type:"media"`` detail entries by model/unit/call_type/resolution. + + Companion to :func:`aggregate_token_usage_by_model`, which only keeps + input/output token entries. Media entries are billed per non-token unit + (images, seconds, ...), so they are grouped separately by their unit and + modality rather than summed into a single token total. Resolution is part of + the key because an image model's price varies by resolution, so different + resolutions of the same model surface as separate billable line items. + Returns one entry per (model, unit, call_type, resolution) combination with + the summed quantity, call count and provider-reported tokens. A group is + marked ``tokens_estimated`` when any entry in it carried estimated tokens, + so a consumer never prices a mixed group as if it were measured. + + Pass a detached list. This iterates ``details`` without any lock, so handing + in a live ``TokenUsage.details`` while another thread appends risks a + "list changed size during iteration" RuntimeError. Use + ``TokenUsage.snapshot().details`` (or a value read from the DB column, as + the API path does) rather than the live list. + """ + if not isinstance(details, list): + return [] + + grouped: Dict[tuple[str, str, str, str], Dict[str, Any]] = {} + for detail in details: + if not isinstance(detail, dict): + continue + if detail.get("type") != "media": + continue + quantity = max(0.0, _coerce_float(detail.get("quantity"))) + # Zero-quantity entries are deliberately KEPT. Unlike a zero-token LLM + # entry, a zero-quantity media entry is meaningful: the duration-billed + # tools record quantity=0 precisely to say "this provider call happened + # but its size is unknown" (an async video with no duration yet). + # Dropping it would hide the whole media popover and report + # media_calls=0 for a task that really did make billable calls. + tokens = max(0, _coerce_int(detail.get("provider_tokens"))) + + raw_model_id = detail.get("model_id") + raw_model_name = detail.get("model") + raw_unit = detail.get("unit") + raw_call_type = detail.get("call_type") + raw_resolution = detail.get("resolution") + model_id = raw_model_id.strip() if isinstance(raw_model_id, str) else "" + model_name = raw_model_name.strip() if isinstance(raw_model_name, str) else "" + unit = raw_unit if isinstance(raw_unit, str) else "" + call_type = raw_call_type if isinstance(raw_call_type, str) else "" + resolution = raw_resolution if isinstance(raw_resolution, str) else "" + # model_id is redundant in the key: it equals identity when set, and is + # constant "" otherwise. + identity = model_id or model_name + key = (identity, unit, call_type, resolution) + + aggregate = grouped.setdefault( + key, + { + "model_id": model_id, + "model_name": model_name, + "unit": unit, + "call_type": call_type, + "resolution": resolution, + "quantity": 0.0, + "calls": 0, + "provider_tokens": 0, + "tokens_estimated": False, + }, + ) + if not aggregate["model_name"] and model_name: + aggregate["model_name"] = model_name + if not aggregate["model_id"] and model_id: + aggregate["model_id"] = model_id + aggregate["quantity"] += quantity + aggregate["calls"] += 1 + aggregate["provider_tokens"] += tokens + if detail.get("tokens_estimated"): + aggregate["tokens_estimated"] = True + + return sorted( + grouped.values(), + key=lambda item: ( + -item["quantity"], + str(item["model_name"]).casefold(), + str(item["unit"]).casefold(), + str(item["call_type"]).casefold(), + str(item["resolution"]).casefold(), + ), + ) + + def add_token_usage( input_tokens: int = 0, output_tokens: int = 0, @@ -389,6 +956,99 @@ def add_token_usage( ) +def add_media_usage( + unit: "MediaUnit | str | None", + quantity: float, + model: str = "", + call_type: "MediaCallType | str | None" = "", + model_id: str = "", + input_tokens: int = 0, + output_tokens: int = 0, + resolution: str = "", + tokens_estimated: bool = False, +) -> None: + """Record non-LLM media model usage on the current context. + + Mirrors ``add_token_usage`` for image/video/tts/asr/embedding/rerank and + any other non-chat modality. The resulting ``type:"media"`` detail entry + flows through the same ``TokenUsage.details`` list into DB persistence and + the quota ``delta_details`` contract, so callers need only this one call. + + Args: + unit: Billable dimension; see :class:`MediaUnit`. Must be stable for a + given (model, call_type) — never vary it by response completeness. + quantity: Amount of ``unit`` consumed (may be fractional, e.g. seconds). + Always 1 for ``MediaUnit.REQUESTS``. + model: Model name for tracking. + call_type: Modality/operation; see :class:`MediaCallType`. + model_id: Unique model id (disambiguates identically-named models). + input_tokens: Provider-reported input tokens; 0 if none. + output_tokens: Provider-reported output tokens; 0 if none. + resolution: Size tier ("1K"/"2K"/"4K" or "1024x1024") for image models + whose price varies by resolution; "" when not applicable. + Token-reporting providers (Gemini / OpenAI gpt-image) also fill + input/output_tokens, recorded as raw ``provider_tokens`` for a future + consumer. Deliberately NOT claimed as a pricing rule: nothing here + expresses or enforces "price by tokens instead of by unit", the row + schema carries no such discriminator, and the aggregate groups purely + by (model, unit, call_type, resolution). Defining that precedence is + tracked in #1461, with the first pricing consumer. + tokens_estimated: True when the token counts are a local heuristic + rather than provider-reported, so billing can refuse to price them. + + Raises: + ValueError: If ``unit`` or ``call_type`` is not a known + :class:`MediaUnit` / :class:`MediaCallType` value, or if ``unit`` is + ``MediaUnit.REQUESTS`` and ``quantity`` is not exactly 1. The first + two are checked here, before the context is touched; the REQUESTS + constraint is enforced inside :meth:`TokenUsage.record_media_call`, + which still rejects before taking its lock, so no partial state is + left either way. Producers route through + ``media_usage.record_media_usage``, which swallows all three so an + accounting bug can never break the underlying media call. + """ + # Coerce defensively so a provider returning a malformed count can never + # crash the underlying media call over accounting. + # Keep the caller's raw value for the REQUESTS error message below: + # record_media_call reports what it was handed, which by then is already + # coerced, so -1 would surface as "got 0.0" without this. + raw_quantity = quantity + quantity = _coerce_float(quantity) + input_tokens = _coerce_int(input_tokens) + output_tokens = _coerce_int(output_tokens) + # Validate before touching the context: a rejected unit must not leave + # media_calls incremented with no matching detail entry behind it. + # Plain strings keep the details list JSON-serialisable whether the caller + # passed an enum member or a bare string. + unit_value = _validated_media_unit(unit) + call_type_value = _validated_media_call_type(call_type) + + usage = get_token_usage() + # One locked operation: see TokenUsage.record_media_call for why the count + # and its detail row must not be observable apart. + usage.record_media_call( + _raw_quantity=raw_quantity, + unit=unit_value, + quantity=quantity, + model=model, + call_type=call_type_value, + model_id=model_id, + input_tokens=input_tokens, + output_tokens=output_tokens, + resolution=resolution, + tokens_estimated=tokens_estimated, + ) + + logger.debug( + f"Media usage added: unit={unit_value}, quantity={quantity}, " + f"resolution={resolution}, model={model}, model_id={model_id}, " + f"call_type={call_type_value}, " + f"provider_tokens={input_tokens + output_tokens}" + f"{' (estimated)' if tokens_estimated else ''}, " + f"total_media_calls={usage.media_calls}" + ) + + def add_tool_call_usage(count: int = 1) -> None: """Record one (or more) tool invocations on the current context.""" get_token_usage().increment_tool_calls(count) diff --git a/src/xagent/core/tools/core/media_usage.py b/src/xagent/core/tools/core/media_usage.py new file mode 100644 index 0000000000..78711762ea --- /dev/null +++ b/src/xagent/core/tools/core/media_usage.py @@ -0,0 +1,191 @@ +"""Best-effort media-usage recording for media generation tools. + +TTS/ASR/video/music/sound-effect models don't return a normalised usage +payload the way image/chat models do, and their adapters are factory-only, so +the natural metering point is the tool call site where the request params and +result are both in scope. This helper wraps ``add_media_usage`` so a failure in +accounting can never break the underlying media call. + +Metering invariants +------------------- +A usage record is only worth what its identity and unit are worth, so every +producer must satisfy all of these: + +1. **Metering must survive adapter unwrapping.** Record on the object callers + actually hold. Reaching past an adapter to its inner provider silently drops + the metering — this is how rerank shipped entirely unbilled. +2. **Metering must survive thread boundaries.** ``ThreadPoolExecutor`` does not + copy contextvars, so a worker gets a fresh empty ``TokenUsage`` unless the + caller's is bound explicitly. +3. **The unit is a property of the modality, never of the response.** A + duration-billed modality always reports seconds, recording ``quantity=0`` + when unmeasured rather than switching units. +4. **Never bill a placeholder identity.** ``"default"``, ``"None"`` and ``""`` + are not models; resolve through :func:`resolve_billing_model`. + +Model identity convention +------------------------- +``model`` carries the human-readable **name**, ``model_id`` the configured +**id**. Populate both when known: the aggregator groups on +``model_id or model``, so a producer that leaves ``model_id`` empty while +another sets it splits one physical model into two un-mergeable billing rows. +""" + +from __future__ import annotations + +import logging +import math +from typing import Any, Optional, TypeGuard + +from ...model.chat.token_context import MediaCallType, MediaUnit, add_media_usage + +logger = logging.getLogger(__name__) + +# Placeholders that must never reach a usage record as a model identity. +_PLACEHOLDER_MODEL_NAMES = {"", "none", "null", "default"} + + +def _usable_model_name(value: Any) -> TypeGuard[str]: + """A real model identity, not a placeholder. TypeGuard so callers narrow.""" + return ( + isinstance(value, str) and value.strip().lower() not in _PLACEHOLDER_MODEL_NAMES + ) + + +def resolve_billing_model( + configured_id: Optional[str], + model: Any = None, + *, + fallback: str = "default", +) -> str: + """Best available identity for a model, never a placeholder string. + + ``_configured_model_id``-style lookups return ``Optional[str]``, and passing + that through ``str()`` records a model literally named ``"None"``. Prefer + the configured id, fall back to the provider's own ``model_name``/``model`` + attribute, and only then to ``fallback``. + """ + + # The placeholder filter applies to the configured id too: a config that + # literally names the model "default" or "none" must not be billed as one. + if _usable_model_name(configured_id): + return configured_id + for attr in ("model_name", "model"): + # Guarded individually: this runs *before* record_media_usage's own + # try/except, so a provider whose model_name is a property that raises + # would break the user's media call over an accounting lookup. + try: + value = getattr(model, attr, None) + except Exception as e: # noqa: BLE001 + logger.warning("Reading %s for billing identity failed: %s", attr, e) + continue + if _usable_model_name(value): + return value + return fallback + + +def coerce_duration(value: object) -> Optional[float]: + """A positive duration in seconds, or None when unusable. + + Distinct from ``token_context._coerce_float``, which folds bad input to + ``0.0``: here the caller must be able to tell "provider reported no + duration" apart from "provider reported zero", because those take + different metering branches. + """ + if value is None or isinstance(value, bool): + return None + try: + seconds = float(value) # type: ignore[arg-type] + except (TypeError, ValueError): + return None + # ``inf > 0`` is True, so non-finite values would otherwise pass as a + # usable duration and reach the record as a non-JSON-serialisable + # quantity. Treat them as "no duration reported" instead. + if not math.isfinite(seconds): + logger.warning("Ignoring non-finite duration: %r", value) + return None + return seconds if seconds > 0 else None + + +def record_media_usage( + unit: MediaUnit | str | None, + quantity: float, + *, + model: str = "", + model_id: str = "", + call_type: MediaCallType | str | None = "", + resolution: str = "", + input_tokens: int = 0, + output_tokens: int = 0, + tokens_estimated: bool = False, +) -> None: + """Record one media model call; swallow any error. + + Includes the ``ValueError`` ``add_media_usage`` raises for an unknown + ``unit``/``call_type``: a metering bug must never break the media call the + user actually asked for. The record is dropped and logged rather than + written under a bogus billing dimension, which is unrepairable once + persisted — so a typo surfaces as a missing row plus this warning, not as a + silently mis-billed one. + """ + try: + add_media_usage( + unit=unit, + quantity=quantity, + model=model, + model_id=model_id, + call_type=call_type, + resolution=resolution, + input_tokens=input_tokens, + output_tokens=output_tokens, + tokens_estimated=tokens_estimated, + ) + except Exception as e: # noqa: BLE001 + # exc_info: the expected case here is the deliberate validation + # ValueError, but anything genuinely unexpected would otherwise surface + # as one context-free line. Persisted billing rows cannot be repaired + # after the fact, so a metering bug needs to be diagnosable from the log + # alone. + logger.warning( + "Failed to record media usage: call_type=%r unit=%r quantity=%r " + "model=%r model_id=%r: %s", + call_type, + unit, + quantity, + model, + model_id, + e, + exc_info=True, + ) + + +def record_media_seconds( + seconds: Optional[float], + *, + model: str = "", + model_id: str = "", + call_type: MediaCallType | str | None = "", +) -> None: + """Record a duration-billed media call, keeping the unit stable. + + Duration-billed modalities (video/ASR/music/sound effect) must always + report ``MediaUnit.SECONDS``: a price table keyed on (model, unit) breaks + if the same model sometimes reports "requests" just because the provider + omitted a duration. When the duration is unknown the call is still recorded + — as ``seconds`` with ``quantity=0`` and a warning — so the event is + visible to billing as unmeasured rather than silently mis-dimensioned. + """ + if seconds is None: + logger.warning( + "No duration reported for %s call on model %r; recording 0 seconds " + "(call happened but is unmeasured)", + call_type, + model, + ) + record_media_usage( + MediaUnit.SECONDS, + seconds or 0.0, + model=model, + model_id=model_id, + call_type=call_type, + ) diff --git a/src/xagent/web/tracking/task_tracker.py b/src/xagent/web/tracking/task_tracker.py index 841e47b523..6827132928 100644 --- a/src/xagent/web/tracking/task_tracker.py +++ b/src/xagent/web/tracking/task_tracker.py @@ -11,6 +11,7 @@ from ...core.model.chat.token_context import ( TokenUsage, + copy_detail_rows, get_token_usage, set_token_usage, ) @@ -50,21 +51,14 @@ def _optional_int(value: Any) -> int | None: return None -def _copy_details(raw_details: Any) -> list[dict[str, Any]]: - if not isinstance(raw_details, list): - return [] - return [dict(item) for item in raw_details if isinstance(item, dict)] - - def _copy_usage(usage: TokenUsage) -> TokenUsage: - """Detach a stable snapshot before yielding to a database worker.""" - return TokenUsage( - input_tokens=usage.input_tokens, - output_tokens=usage.output_tokens, - llm_calls=usage.llm_calls, - tool_calls=usage.tool_calls, - details=_copy_details(usage.details), - ) + """Detach a stable snapshot before yielding to a database worker. + + Delegates to ``TokenUsage.snapshot`` so the copy is taken under the usage + object's own lock. Reading the fields here field-by-field could interleave + with a concurrent media write and yield a counter without its detail row. + """ + return usage.snapshot() def _task_for_run( @@ -100,13 +94,22 @@ def _task_seed_from_session( f" for run {expected_run_id}" if expected_run_id is not None else "" ) raise ValueError(f"Task {task_id}{run_suffix} not found") + seed_details = copy_detail_rows(getattr(task, "token_usage_details", None)) return _TaskTrackingSeed( user_id=_optional_int(getattr(task, "user_id", None)), usage=TokenUsage( input_tokens=_safe_int(getattr(task, "input_tokens", 0)), output_tokens=_safe_int(getattr(task, "output_tokens", 0)), llm_calls=_safe_int(getattr(task, "llm_calls", 0)), - details=_copy_details(getattr(task, "token_usage_details", None)), + # Derived from the surviving detail rows rather than read from a + # column: Task has no media_calls column. This only keeps the + # in-memory running total consistent with the rows it summarises — + # the delta path reads `details` and `tool_calls`, never this + # counter, so seeding zero would not re-report prior-turn calls. + media_calls=sum( + 1 for detail in seed_details if detail.get("type") == "media" + ), + details=seed_details, ), ) @@ -160,7 +163,7 @@ def _commit_task_usage_if_owned( Task.output_tokens: usage.output_tokens, Task.total_tokens: usage.total_tokens, Task.llm_calls: usage.llm_calls, - Task.token_usage_details: _copy_details(usage.details), + Task.token_usage_details: copy_detail_rows(usage.details), }, synchronize_session=False, ) @@ -528,16 +531,44 @@ async def stop_periodic_updates(self) -> None: self._update_task = None logger.info(f"Stopped periodic token updates for task {self.task_id}") - def _turn_delta(self, usage: TokenUsage | None = None) -> tuple[list, int]: + def _turn_delta( + self, usage: TokenUsage | None = None + ) -> tuple[list[dict[str, Any]], int]: """This turn's (detail entries, tool-call count) over the baseline seeded in start_tracking. Single source for the completion meter and the mid-run - gate so they can't disagree on what 'this run's usage' means.""" + gate so they can't disagree on what 'this run's usage' means. + + Always snapshots first. The two reads below — the details slice and the + tool-call counter — must describe the same instant: a concurrent + ``record_media_call``/``increment_tool_calls`` landing between them + yields a pair from two different logical points in time (a row visible + with its counter bump missing, or the reverse), and that inconsistent + pair is what gets fed to quota gating. ``interrupt_reason_for_quota`` + polls this at every safe point, so it is a hot concurrent path rather + than a one-shot teardown read. + + Relies on an invariant worth stating: every row in ``details`` is a + dict. ``_initial_details_len`` indexes the same filtered list the live + usage object holds (``TokenUsage.__post_init__`` normalises on + construction and every mutation path appends a literal dict), so the + slice below stays aligned. A future non-dict append would be filtered + out of the tail and shift the boundary, re-metering prior-turn rows. + + Reads through ``detail_tail``, which copies only the rows past the + baseline. ``snapshot()`` would deep-copy the entire cumulative list and + then discard everything before ``_initial_details_len`` — the list grows + monotonically across turns, so that costs O(total usage) per poll on a + path hit once per agent step and once per streamed LLM chunk, while + holding the lock every LLM adapter needs to write tokens. + """ if usage is None: usage = get_token_usage() - return ( - usage.details[self._initial_details_len :], - max(0, usage.tool_calls - self._initial_tool_calls), - ) + # Rows come back already detached, so callers need no further copying + # before handing them to the quota hooks. tool_calls is the cumulative + # counter, so the baseline still has to be subtracted here — detail_tail + # slices details for us but cannot know this tracker's baseline. + tail, tool_calls = usage.detail_tail(self._initial_details_len) + return (tail, max(0, tool_calls - self._initial_tool_calls)) async def interrupt_reason_for_quota(self) -> str | None: """Per-step interrupt-checker: return a reason when this run's live-so-far @@ -558,9 +589,10 @@ async def interrupt_reason_for_quota(self) -> str | None: return None try: delta_details, delta_actions = self._turn_delta() + # No extra copy: _turn_delta already returns detached rows. reason = _check_quota_on_event_loop( self._user_id, - _copy_details(delta_details), + delta_details, delta_actions, ) if reason is not None: @@ -613,9 +645,10 @@ async def _complete_tracking_once(self, usage: TokenUsage) -> TokenUsage: # the task. The remaining blocking risk is tracked separately from the # database-lifecycle changes in this PR. try: + # No extra copy: _turn_delta already returns detached rows. _record_usage_on_event_loop( self._user_id, - _copy_details(delta_details), + delta_details, delta_actions, ) except Exception as e: # noqa: BLE001 diff --git a/tests/core/model/chat/test_media_usage.py b/tests/core/model/chat/test_media_usage.py new file mode 100644 index 0000000000..7083cfd6c3 --- /dev/null +++ b/tests/core/model/chat/test_media_usage.py @@ -0,0 +1,961 @@ +"""Media (non-LLM) usage tracking: image/video/tts/asr/embedding/rerank. + +These modalities record usage via ``add_media_usage`` into the same +``TokenUsage.details`` list that LLM tokens use, so they flow through DB +persistence and the quota ``delta_details`` contract without special-casing. +""" + +import sys +import threading + +import pytest + +from xagent.core.model.chat.token_context import ( + MediaCallType, + MediaUnit, + TokenContextManager, + TokenUsage, + add_media_usage, + add_token_usage, + aggregate_media_usage_by_model, + aggregate_token_usage_by_model, +) + + +def test_add_media_usage_appends_media_entry_and_counts_call() -> None: + with TokenContextManager() as manager: + add_media_usage( + unit="images", + quantity=2, + model="sd-xl", + model_id="m1", + call_type="generate_image", + resolution="1K", + ) + usage = manager.get_usage() + + assert usage.media_calls == 1 + # Media does not count as an LLM call or add tokens. + assert usage.llm_calls == 0 + assert usage.input_tokens == 0 + assert usage.output_tokens == 0 + assert len(usage.details) == 1 + entry = usage.details[0] + assert entry["type"] == "media" + assert entry["unit"] == "images" + assert entry["quantity"] == 2.0 + assert entry["model"] == "sd-xl" + assert entry["model_id"] == "m1" + assert entry["call_type"] == "generate_image" + assert entry["resolution"] == "1K" + + +def test_add_media_usage_carries_accompanying_tokens() -> None: + with TokenContextManager() as manager: + add_media_usage( + unit="images", + quantity=1, + model="gemini-image", + call_type="generate_image", + input_tokens=5, + output_tokens=3, + ) + usage = manager.get_usage() + + entry = usage.details[0] + # Stored under provider_tokens, never "tokens": a consumer that sums the + # "tokens" key across all entries must not pick up media counts. + assert "tokens" not in entry + assert entry["provider_tokens"] == 8 + assert entry["provider_input_tokens"] == 5 + assert entry["provider_output_tokens"] == 3 + assert entry["tokens_estimated"] is False + # Media token passthrough must NOT inflate the LLM token totals. + assert usage.input_tokens == 0 + assert usage.output_tokens == 0 + + +def test_estimated_tokens_are_flagged() -> None: + with TokenContextManager() as manager: + add_media_usage( + unit="texts", + quantity=2, + model="embed", + call_type="embedding", + input_tokens=12, + tokens_estimated=True, + ) + details = manager.get_usage().details + + assert details[0]["tokens_estimated"] is True + # The flag survives aggregation so billing can refuse to price an estimate. + assert aggregate_media_usage_by_model(details)[0]["tokens_estimated"] is True + + +def test_dirty_quantity_is_coerced_on_free_quantity_units() -> None: + # A malformed quantity records 0 rather than raising: the provider call + # happened, so the row must survive as "unmeasured". Uses a duration-billed + # unit, because REQUESTS additionally pins its quantity to exactly 1 (see + # test_requests_unit_rejects_any_quantity_but_one) and would reject these. + with TokenContextManager() as manager: + add_media_usage(unit="seconds", quantity=None, model="x", call_type="asr") # type: ignore[arg-type] + add_media_usage(unit="seconds", quantity="oops", model="x", call_type="asr") # type: ignore[arg-type] + usage = manager.get_usage() + + assert usage.media_calls == 2 + assert all(entry["quantity"] == 0.0 for entry in usage.details) + + +def test_dirty_quantity_on_requests_unit_is_rejected_not_coerced() -> None: + # For REQUESTS the permissive path is wrong: coercing a malformed quantity + # to 0 would record a call whose billable quantity contradicts its own + # call count. Reject instead, leaving no state behind. + with TokenContextManager() as manager: + for bad in (None, "oops", 0): + with pytest.raises(ValueError, match="exactly one call"): + add_media_usage( + unit="requests", quantity=bad, model="x", call_type="rerank" + ) # type: ignore[arg-type] + usage = manager.get_usage() + + assert usage.media_calls == 0 + assert usage.details == [] + + +def test_to_dict_from_dict_roundtrip_preserves_media() -> None: + with TokenContextManager() as manager: + add_token_usage(input_tokens=10, output_tokens=4, model="gpt", model_id="g1") + add_media_usage(unit="seconds", quantity=3.5, model="tts", call_type="asr") + usage = manager.get_usage() + + data = usage.to_dict() + assert data["media_calls"] == 1 + assert data["llm_calls"] == 1 + + restored = TokenUsage.from_dict(data) + assert restored.media_calls == 1 + assert restored.llm_calls == 1 + assert restored.input_tokens == 10 + # 2 token entries (one input, one output) + 1 media entry. + assert len(restored.details) == 3 + assert sum(1 for d in restored.details if d["type"] == "media") == 1 + + +def test_merge_combines_media_calls_and_details() -> None: + a = TokenUsage() + a.record_media_call( + unit="images", quantity=1, model="x", call_type="generate_image" + ) + + b = TokenUsage() + b.record_media_call(unit="seconds", quantity=2, model="y", call_type="asr") + + a.merge(b) + assert a.media_calls == 2 + assert len(a.details) == 2 + + +def test_token_aggregation_ignores_media_entries() -> None: + with TokenContextManager() as manager: + add_token_usage(input_tokens=10, output_tokens=5, model="gpt", model_id="g1") + add_media_usage( + unit="images", quantity=2, model="sd", call_type="generate_image" + ) + details = manager.get_usage().details + + token_groups = aggregate_token_usage_by_model(details) + assert len(token_groups) == 1 + assert token_groups[0]["model_name"] == "gpt" + assert token_groups[0]["input_tokens"] == 10 + assert token_groups[0]["output_tokens"] == 5 + + +def test_media_aggregation_groups_by_model_unit_and_call_type() -> None: + with TokenContextManager() as manager: + add_media_usage( + unit="images", quantity=2, model="sd", call_type="generate_image" + ) + add_media_usage( + unit="images", quantity=3, model="sd", call_type="generate_image" + ) + add_media_usage(unit="seconds", quantity=4, model="tts", call_type="asr") + # LLM tokens must never appear in the media aggregation. + add_token_usage(input_tokens=7, output_tokens=2, model="gpt", model_id="g1") + details = manager.get_usage().details + + media_groups = aggregate_media_usage_by_model(details) + assert len(media_groups) == 2 + + by_unit = {group["unit"]: group for group in media_groups} + assert by_unit["images"]["quantity"] == 5.0 + assert by_unit["images"]["calls"] == 2 + assert by_unit["images"]["call_type"] == "generate_image" + assert by_unit["seconds"]["quantity"] == 4.0 + assert by_unit["seconds"]["calls"] == 1 + + +def test_media_aggregation_splits_by_resolution() -> None: + # Same model+call_type at different resolutions must bill as separate line + # items, since an image model's price varies by resolution. + with TokenContextManager() as manager: + add_media_usage( + unit="images", + quantity=1, + model="gemini-image", + call_type="generate_image", + resolution="1K", + ) + add_media_usage( + unit="images", + quantity=1, + model="gemini-image", + call_type="generate_image", + resolution="4K", + ) + details = manager.get_usage().details + + groups = aggregate_media_usage_by_model(details) + assert len(groups) == 2 + by_res = {group["resolution"]: group for group in groups} + assert set(by_res) == {"1K", "4K"} + assert by_res["1K"]["calls"] == 1 + assert by_res["4K"]["calls"] == 1 + + +def test_aggregations_tolerate_non_list_and_dirty_entries() -> None: + assert aggregate_media_usage_by_model(None) == [] + assert aggregate_media_usage_by_model("nope") == [] + # Non-dict junk is ignored; a bare media entry still counts as a call, since + # a media row's existence is itself the billing signal. + groups = aggregate_media_usage_by_model([{"type": "media"}, 42, "junk"]) + assert len(groups) == 1 + assert groups[0]["calls"] == 1 + assert groups[0]["quantity"] == 0.0 + + +def test_zero_quantity_media_entries_stay_visible() -> None: + # A duration-billed call the provider never measured records 0 seconds. + # That entry must survive aggregation: it is the only evidence the task + # made a billable provider call, and dropping it would report + # media_calls=0 (and hide the whole popover) for a task that did. + with TokenContextManager() as manager: + add_media_usage(unit="seconds", quantity=0, model="tts", call_type="asr") + add_media_usage(unit="seconds", quantity=5, model="tts", call_type="asr") + details = manager.get_usage().details + + groups = aggregate_media_usage_by_model(details) + assert len(groups) == 1 + assert groups[0]["quantity"] == 5.0 + assert groups[0]["calls"] == 2 # both calls counted, including the unmeasured one + + +def test_only_unmeasured_calls_still_surface() -> None: + # The async-video case: no duration is available yet for any call, so the + # whole group is zero-quantity. It must still be reported. + with TokenContextManager() as manager: + for _ in range(3): + add_media_usage(unit="seconds", quantity=0, model="veo", call_type="video") + groups = aggregate_media_usage_by_model(manager.get_usage().details) + + assert len(groups) == 1 + assert groups[0]["calls"] == 3 + assert groups[0]["quantity"] == 0.0 + + +@pytest.mark.parametrize("bad_unit", ["image", "second", "tokens", "IMAGES", ""]) +def test_unknown_unit_is_rejected(bad_unit: str) -> None: + # A typo'd unit mints a new billing dimension that the aggregator will + # happily key off, and a written usage record cannot be repaired + # retroactively. The write boundary is the last point it is still fixable. + with TokenContextManager() as manager: + with pytest.raises(ValueError, match="Unknown media unit"): + add_media_usage(unit=bad_unit, quantity=1, model="m", call_type="tts") + usage = manager.get_usage() + + # A rejected call must leave no partial state behind: media_calls is + # incremented before the detail entry is appended, so validating late would + # record a call with no matching entry. + assert usage.media_calls == 0 + assert usage.details == [] + + +def test_unknown_call_type_is_rejected() -> None: + with TokenContextManager() as manager: + with pytest.raises(ValueError, match="Unknown media call type"): + add_media_usage(unit="seconds", quantity=1, model="m", call_type="speech") + usage = manager.get_usage() + + assert usage.media_calls == 0 + assert usage.details == [] + + +def test_empty_call_type_is_allowed() -> None: + # call_type is optional metadata rather than a billing dimension of its + # own, so omitting it stays legal while a typo does not. + with TokenContextManager() as manager: + add_media_usage(unit="requests", quantity=1, model="m") + usage = manager.get_usage() + + assert usage.media_calls == 1 + assert usage.details[0]["call_type"] == "" + + +def test_none_call_type_is_normalised_to_empty() -> None: + with TokenContextManager() as manager: + add_media_usage(unit="requests", quantity=1, model="m", call_type=None) + usage = manager.get_usage() + + assert usage.media_calls == 1 + assert usage.details[0]["call_type"] == "" + + +def test_none_unit_is_rejected_with_clear_error() -> None: + with TokenContextManager() as manager: + with pytest.raises(ValueError, match="Media unit cannot be None"): + add_media_usage(unit=None, quantity=1, model="m") + usage = manager.get_usage() + + assert usage.media_calls == 0 + assert usage.details == [] + + +def test_enum_members_are_accepted() -> None: + with TokenContextManager() as manager: + add_media_usage( + unit=MediaUnit.SECONDS, + quantity=3, + model="whisper", + call_type=MediaCallType.ASR, + ) + usage = manager.get_usage() + + entry = usage.details[0] + # Stored as plain strings so details stays JSON-serialisable. + assert entry["unit"] == "seconds" + assert entry["call_type"] == "asr" + assert isinstance(entry["unit"], str) + + +def test_concurrent_snapshots_never_see_a_torn_counter_and_rows_pair() -> None: + # What the lock actually protects is the PAIRING of the counter with its + # detail rows, not `+=` itself. On CPython 3.12 a bare `self.n += 1` does not + # measurably lose increments even at 8x200000 threads-iterations, so a test + # that only counts increments passes with the lock entirely removed — an + # earlier version of this test did exactly that. + # + # record_media_call bumps the counter and appends the row as one operation. + # Without the lock a concurrent reader lands between them and observes + # media_calls disagreeing with the number of media rows. Mutation-verified: + # replacing _lock with contextlib.nullcontext() makes this fail in 30/30 + # runs, typically within the first few hundred writes. + usage = TokenUsage() + torn: list[tuple[int, int]] = [] + stop = threading.Event() + + # Required, not decoration: at the default switch interval the window is too + # narrow to hit, and this test passes with the lock removed. At 1e-9 the + # unlocked variant tears in 10/10 runs. + previous_interval = sys.getswitchinterval() + sys.setswitchinterval(1e-9) + + def writer() -> None: + for _ in range(3000): + usage.record_media_call( + unit="images", quantity=1, call_type="generate_image" + ) + stop.set() + + def reader() -> None: + while not stop.is_set(): + snap = usage.to_dict() + media = sum(1 for d in snap["details"] if d.get("type") == "media") + if snap["media_calls"] != media: + torn.append((snap["media_calls"], media)) + return + + threads = [threading.Thread(target=writer), threading.Thread(target=reader)] + try: + for thread in threads: + thread.start() + for thread in threads: + thread.join(timeout=60) + finally: + sys.setswitchinterval(previous_interval) + + assert not torn, f"observed torn counter/rows pairs: {torn[:5]}" + assert usage.media_calls == 3000 + assert sum(1 for d in usage.details if d.get("type") == "media") == 3000 + + +def test_concurrent_merge_loses_no_counts() -> None: + # Each thread merges repeatedly and all start together at a barrier, so the + # read-modify-write in merge() genuinely overlaps. One merge per thread of a + # one-row source would very likely pass against an unlocked merge() too. + target = TokenUsage() + workers, merges_per_worker = 8, 100 + + source = TokenUsage() + source.record_media_call(unit="seconds", quantity=2, call_type="asr") + + barrier = threading.Barrier(workers) + + def merge_many() -> None: + barrier.wait(timeout=10) + for _ in range(merges_per_worker): + target.merge(source) + + threads = [threading.Thread(target=merge_many) for _ in range(workers)] + for thread in threads: + thread.start() + for thread in threads: + thread.join(timeout=30) + + for thread in threads: + assert not thread.is_alive(), "merge deadlocked under contention" + + expected = workers * merges_per_worker + assert target.media_calls == expected + media_rows = [d for d in target.details if d.get("type") == "media"] + assert len(media_rows) == expected + # The counter and its rows must agree exactly — a lost += shows up here. + assert target.media_calls == len(media_rows) + + +def test_usage_survives_asdict_deepcopy_and_pickle() -> None: + # The mutation lock is not picklable, so as a dataclass field it would break + # every generic consumer of this object (dataclasses.asdict walks + # __dataclass_fields__ and ignores __getstate__). + import copy + import dataclasses + import pickle + + usage = TokenUsage() + usage.record_media_call(unit="images", quantity=2, call_type="generate_image") + + assert dataclasses.asdict(usage)["media_calls"] == 1 + assert copy.deepcopy(usage).media_calls == 1 + revived = pickle.loads(pickle.dumps(usage)) + assert revived.media_calls == 1 + # The revived object gets a working lock of its own, not a shared one. + revived.record_media_call(unit="images", quantity=1, call_type="generate_image") + assert revived.media_calls == 2 + + +def test_legacy_positional_construction_still_binds_the_same_fields() -> None: + # media_calls must NOT be inserted into the historical positional sequence. + # An existing TokenUsage(input, output, llm_calls, tool_calls, details) call + # would otherwise bind its 4th argument to media_calls and shift the rest — + # tool count read as media count, details read as tool count — silently, + # with no error. The field is appended and keyword-only instead. + details = [{"type": "input", "tokens": 7}] + usage = TokenUsage(10, 5, 2, 3, details) + + assert usage.input_tokens == 10 + assert usage.output_tokens == 5 + assert usage.llm_calls == 2 + assert usage.tool_calls == 3 + assert usage.details == details + assert usage.media_calls == 0 + + # And it is reachable by keyword, where it cannot be confused with anything. + assert TokenUsage(media_calls=4).media_calls == 4 + + # The private lock must not occupy a positional slot either: six positional + # arguments is the documented maximum (five fields plus self). + with pytest.raises(TypeError): + TokenUsage(1, 2, 3, 4, [], 5) # type: ignore[misc] + + +def test_concurrent_merges_in_both_directions_do_not_deadlock() -> None: + # merge() snapshots the source under its own lock and releases it before + # taking the target's; holding both at once would deadlock on opposing + # merges. A plain thread start/join can be scheduled so the two never + # actually overlap, which would pass even against a both-locks-held + # implementation — so a Barrier forces them into the window. + a = TokenUsage() + b = TokenUsage() + per_side = 50 + for _ in range(per_side): + a.record_media_call(unit="images", quantity=1, call_type="generate_image") + b.record_media_call(unit="seconds", quantity=1, call_type="asr") + + barrier = threading.Barrier(2) + + def merge_after_barrier(target: TokenUsage, source: TokenUsage) -> None: + barrier.wait(timeout=10) + target.merge(source) + + threads = [ + threading.Thread(target=merge_after_barrier, args=(a, b)), + threading.Thread(target=merge_after_barrier, args=(b, a)), + ] + for thread in threads: + thread.start() + for thread in threads: + thread.join(timeout=10) + + for thread in threads: + assert not thread.is_alive(), "merge deadlocked in opposing directions" + + # Exact counts, not lower bounds: a lower bound would pass while rows were + # dropped or a counter drifted from its detail rows. Each side ends with its + # own 50 plus whatever snapshot it read of the other — between 50 and 100 + # extra depending on interleaving — but the scalar and the rows must always + # agree with each other, and neither side may lose its own rows. + for usage, own_unit in ((a, "images"), (b, "seconds")): + media_rows = [d for d in usage.details if d.get("type") == "media"] + assert usage.media_calls == len(media_rows), "counter drifted from rows" + assert len(media_rows) >= 2 * per_side + assert len(media_rows) <= 3 * per_side + own_rows = [d for d in media_rows if d["unit"] == own_unit] + assert len(own_rows) >= per_side, "a side lost its own rows" + + +def test_estimate_media_tokens_accepts_any_iterable_of_strings() -> None: + # The docstring promises an iterable; a generator previously estimated 0, + # so an embedding producer passing one would have billed nothing. + from xagent.core.model.chat.token_context import estimate_media_tokens + + assert estimate_media_tokens(x for x in ["abcd", "efgh"]) == 2 + assert estimate_media_tokens(["abcd", "efgh"]) == 2 + # Non-iterables and non-string members are ignored, never raised on. + assert estimate_media_tokens(42) == 0 + assert estimate_media_tokens(["abcd", 42, None]) == 1 + + +@pytest.mark.parametrize( + ("text", "expected"), + [("", 0), ("a", 1), ("ab", 1), ("abc", 1), ("abcd", 1), ("abcde", 2)], +) +def test_estimate_media_tokens_never_undercounts_short_text(text, expected) -> None: + # Truncating division estimated 0 for any 1-3 character Latin string, so + # short non-empty text billed nothing. Empty still estimates 0. + from xagent.core.model.chat.token_context import estimate_media_tokens + + assert estimate_media_tokens(text) == expected + + +@pytest.mark.parametrize("bad_quantity", [0, 0.5, 2, 7, -1]) +def test_requests_unit_rejects_any_quantity_but_one(bad_quantity) -> None: + # MediaUnit.REQUESTS is defined as exactly one provider call, so its + # quantity is not a free variable. Letting another value through would make + # the billable quantity disagree with the media_calls / aggregate `calls` + # count derived from the same row, so quota and pricing would read + # different numbers off one record. + usage = TokenUsage() + with pytest.raises(ValueError, match="exactly one call"): + usage.record_media_call( + unit=MediaUnit.REQUESTS, quantity=bad_quantity, call_type="rerank" + ) + + # Rejected before any mutation: no counter bump, no orphan detail row. + assert usage.media_calls == 0 + assert usage.details == [] + + +def test_requests_unit_accepts_exactly_one() -> None: + usage = TokenUsage() + usage.record_media_call(unit=MediaUnit.REQUESTS, quantity=1, call_type="rerank") + + assert usage.media_calls == 1 + assert usage.details[0]["quantity"] == 1.0 + + +def test_other_units_keep_free_quantities() -> None: + # The REQUESTS constraint must not leak into duration/count-billed units. + usage = TokenUsage() + usage.record_media_call(unit=MediaUnit.SECONDS, quantity=2.5, call_type="asr") + usage.record_media_call(unit=MediaUnit.IMAGES, quantity=4, call_type="edit_image") + + assert [d["quantity"] for d in usage.details] == [2.5, 4.0] + + +def test_requests_rejection_is_swallowed_by_the_wrapper() -> None: + # Producers route through record_media_usage, whose contract is that an + # accounting bug never breaks the media call the user asked for. + from xagent.core.tools.core.media_usage import record_media_usage + + with TokenContextManager() as manager: + record_media_usage( + MediaUnit.REQUESTS, 3, model="m", call_type=MediaCallType.RERANK + ) + usage = manager.get_usage() + + assert usage.media_calls == 0 + assert usage.details == [] + + +@pytest.mark.parametrize( + "bad_tokens", + [True, False, -100, float("inf"), float("-inf"), float("nan"), "abc"], +) +def test_malformed_provider_tokens_sanitize_without_losing_the_row(bad_tokens) -> None: + # int(float("inf")) raises OverflowError, which was not caught — so a + # malformed provider payload propagated out of the accounting path and + # dropped the whole billable media row. Negatives and booleans were + # persisted as-is, letting a provider bug subtract from a bill. + usage = TokenUsage() + usage.record_media_call( + unit="images", + quantity=1, + call_type="generate_image", + input_tokens=bad_tokens, + output_tokens=bad_tokens, + ) + + # The row survives — the call happened and is billable by quantity ... + assert usage.media_calls == 1 + entry = usage.details[0] + assert entry["quantity"] == 1.0 + # ... with only the unusable token fields zeroed. + assert entry["provider_input_tokens"] == 0 + assert entry["provider_output_tokens"] == 0 + assert entry["provider_tokens"] == 0 + + +def test_valid_provider_tokens_survive() -> None: + usage = TokenUsage() + usage.record_media_call( + unit="images", + quantity=1, + call_type="generate_image", + input_tokens=11, + output_tokens=5, + ) + entry = usage.details[0] + + assert (entry["provider_input_tokens"], entry["provider_output_tokens"]) == (11, 5) + assert entry["provider_tokens"] == 16 + + +def test_snapshot_is_taken_under_the_lock() -> None: + # An external copier (TaskTracker._copy_usage) reading fields one by one can + # interleave with a concurrent write and see a counter without its matching + # detail row. snapshot() takes the lock once so the pair always agrees. + usage = TokenUsage() + torn: list[tuple[int, int]] = [] + stop = threading.Event() + + def writer() -> None: + for _ in range(400): + usage.record_media_call( + unit="images", quantity=1, call_type="generate_image" + ) + stop.set() + + def reader() -> None: + while not stop.is_set(): + snap = usage.snapshot() + media = sum(1 for d in snap.details if d.get("type") == "media") + if snap.media_calls != media: + torn.append((snap.media_calls, media)) + + threads = [threading.Thread(target=writer), threading.Thread(target=reader)] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + + assert not torn, f"snapshot() returned torn state: {torn[:5]}" + assert usage.media_calls == 400 + + +def test_snapshot_detaches_from_the_source() -> None: + usage = TokenUsage() + usage.record_media_call(unit="images", quantity=1, call_type="generate_image") + snap = usage.snapshot() + + usage.record_media_call(unit="images", quantity=1, call_type="generate_image") + + # The snapshot must not see writes that landed after it was taken. + assert snap.media_calls == 1 + assert len(snap.details) == 1 + assert snap.details is not usage.details + + +def test_estimate_media_tokens_never_raises_from_a_hostile_iterable() -> None: + # The docstring promises malformed input cannot raise, but only TypeError + # was caught: a custom iterable raising anything else escaped into the + # accounting path and would break the call being measured. + from xagent.core.model.chat.token_context import estimate_media_tokens + + class _RaisesMidIteration: + def __iter__(self): + yield "abcd" + raise RuntimeError("iterator exploded") + + class _RaisesImmediately: + def __iter__(self): + raise ValueError("cannot iterate") + + assert estimate_media_tokens(_RaisesMidIteration()) == 0 + assert estimate_media_tokens(_RaisesImmediately()) == 0 + + +def test_copy_copy_does_not_share_details_under_a_different_lock() -> None: + # The default shallow copy routes through __getstate__/__setstate__, which + # gives the new instance a fresh lock while still referencing the SAME + # details list — two objects each believing they hold exclusive access, + # mutating one list under two different locks. __copy__ routes to snapshot() + # so copy.copy is a supported fork rather than a trap. + import copy + + usage = TokenUsage() + usage.record_media_call(unit="images", quantity=1, call_type="generate_image") + clone = copy.copy(usage) + + assert clone.details is not usage.details + assert clone._lock is not usage._lock + + clone.record_media_call(unit="images", quantity=1, call_type="generate_image") + assert len(usage.details) == 1, "the original saw a write through its copy" + assert len(clone.details) == 2 + assert usage.media_calls == 1 + + +def test_deepcopy_also_routes_through_snapshot() -> None: + import copy + + usage = TokenUsage() + usage.record_media_call(unit="seconds", quantity=3, call_type="asr") + clone = copy.deepcopy(usage) + + assert clone.details is not usage.details + assert clone.media_calls == 1 + clone.record_media_call(unit="seconds", quantity=1, call_type="asr") + assert usage.media_calls == 1 + + +def test_to_dict_does_not_leak_the_live_detail_dicts() -> None: + # to_dict takes the lock, but a new outer list around the same inner dicts + # lets a caller mutate the live usage object through the returned payload, + # entirely outside that lock. + usage = TokenUsage() + usage.record_media_call(unit="images", quantity=2, call_type="generate_image") + + payload = usage.to_dict() + payload["details"][0]["quantity"] = 999 + payload["details"].append({"type": "media", "unit": "images"}) + + assert usage.details[0]["quantity"] == 2.0 + assert len(usage.details) == 1 + + +@pytest.mark.parametrize("negative", [-1, -5, -100]) +def test_counters_ignore_negative_increments(negative) -> None: + # A negative count would drive the counter below the number of matching + # detail rows — the same counter/rows divergence the REQUESTS constraint + # exists to prevent. + usage = TokenUsage() + usage.increment_tool_calls(negative) + + assert usage.tool_calls == 0 + + # Positive increments still work. + usage.increment_tool_calls(2) + assert usage.tool_calls == 2 + + +def test_merge_does_not_alias_the_source_detail_dicts() -> None: + # Sharing the inner dicts would leave merged rows aliased to the source's, + # so mutating one usage object silently rewrites the other's billing rows. + # Found by self-review as the sibling of the to_dict leak. + target = TokenUsage() + source = TokenUsage() + source.record_media_call(unit="images", quantity=1, call_type="generate_image") + + target.merge(source) + target.details[0]["quantity"] = 999 + + assert source.details[0]["quantity"] == 1.0 + assert target.media_calls == 1 + + +def test_detail_tail_returns_only_the_tail_atomically() -> None: + # The narrow read the delta path needs: copying the whole cumulative list to + # use a small tail costs O(total usage) on a per-chunk quota path. + usage = TokenUsage() + for _ in range(5): + usage.record_media_call(unit="images", quantity=1, call_type="generate_image") + usage.increment_tool_calls(4) + + tail, tool_calls = usage.detail_tail(3) + + assert len(tail) == 2 + assert tool_calls == 4 + # Rows are detached, so a caller cannot mutate live state through them. + tail[0]["quantity"] = 999 + assert usage.details[3]["quantity"] == 1.0 + + # A start past the end is empty rather than an error. + assert usage.detail_tail(99) == ([], 4) + + +def test_constructor_and_from_dict_detach_and_filter_details() -> None: + # details is persisted as free-form JSON, so a legacy row may not be a dict; + # and storing the caller's list by reference would let them mutate rows + # outside the lock. Both are normalised at construction. + caller_list = [{"type": "media", "unit": "images"}, "junk", None, 42] + + usage = TokenUsage(details=caller_list) # type: ignore[arg-type] + assert len(usage.details) == 1 + + caller_list[0]["unit"] = "MUTATED" # type: ignore[index] + assert usage.details[0]["unit"] == "images" + + revived = TokenUsage.from_dict({"details": [{"type": "media"}, "bad"]}) + assert len(revived.details) == 1 + + +def test_requests_error_reports_the_value_the_caller_passed() -> None: + # quantity is coerced before the guard runs, so interpolating the coerced + # value would report "got 0.0" for a caller who passed -1. + usage = TokenUsage() + with pytest.raises(ValueError, match=r"got -1"): + usage.record_media_call(unit="requests", quantity=-1, call_type="rerank") + + +def test_lock_is_reentrant() -> None: + # A plain Lock makes any future nested acquisition a self-deadlock, resting + # on manual discipline; RLock removes the trap. + usage = TokenUsage() + with usage._lock: + with usage._lock: + usage.increment_tool_calls(1) + assert usage.tool_calls == 1 + + +@pytest.mark.parametrize( + ("text", "expected"), + [ + # CJK is ~1 token per character; a flat chars/4 heuristic undercounts + # Chinese by close to 4x. Punctuation and fullwidth forms count too — + # they appear in essentially every Chinese sentence. + ("中文", 2), + ("中文abcd", 3), + ("中文，你好。", 6), + ("（全角）", 4), + ], +) +def test_estimate_media_tokens_counts_cjk(text, expected) -> None: + from xagent.core.model.chat.token_context import estimate_media_tokens + + assert estimate_media_tokens(text) == expected + + +def test_coerce_int_change_applies_to_the_llm_token_path() -> None: + # _coerce_int is shared with add_token_usage, so its hardening (bool -> 0, + # negatives -> 0) changes behaviour for every LLM adapter, not just media. + with TokenContextManager() as manager: + add_token_usage( + input_tokens=True, # type: ignore[arg-type] + output_tokens=-5, + model="m", + call_type="chat", + ) + usage = manager.get_usage() + + assert usage.input_tokens == 0 + assert usage.output_tokens == 0 + + +def test_media_aggregation_uses_model_id_when_present() -> None: + # identity = model_id or model_name — the model_id branch had no coverage, + # so rows keyed by id always fell through to the name branch in tests. + with TokenContextManager() as manager: + add_media_usage( + unit="images", + quantity=1, + model="display-name", + model_id="img-1", + call_type="generate_image", + ) + add_media_usage( + unit="images", + quantity=2, + model="", + model_id="img-1", + call_type="generate_image", + ) + groups = aggregate_media_usage_by_model(manager.get_usage().details) + + # Both rows share one identity via model_id even though one lacks a name... + assert len(groups) == 1 + assert groups[0]["model_id"] == "img-1" + assert groups[0]["quantity"] == 3.0 + # ...and the name is backfilled from whichever row carried it. + assert groups[0]["model_name"] == "display-name" + + +@pytest.mark.parametrize("huge", [10**400, -(10**400), "1e400"]) +def test_huge_quantity_records_zero_rather_than_dropping_the_row(huge) -> None: + # float(10**400) raises OverflowError, which _coerce_float did not catch. + # An uncaught raise here propagates out and the error-swallowing wrapper + # drops the whole billing row — the opposite of the reject-to-0.0 contract. + usage = TokenUsage() + usage.record_media_call(unit="images", quantity=huge, call_type="generate_image") + + assert usage.media_calls == 1 + assert usage.details[0]["quantity"] == 0.0 + + +def test_add_media_usage_reports_the_raw_quantity_in_the_requests_error() -> None: + # The wrapper coerces one layer above record_media_call, so without + # threading the raw value through, a caller passing -1 saw "got 0.0". + with TokenContextManager(): + with pytest.raises(ValueError, match=r"got -1"): + add_media_usage(unit="requests", quantity=-1, model="m", call_type="rerank") + + +@pytest.mark.parametrize( + ("unit", "call_type"), + [ + ("images", "video"), # video bills in seconds + ("seconds", "tts"), # tts bills in characters + ("images", "asr"), + ("requests", "embedding"), # embedding bills per text + ], +) +def test_unit_must_match_the_modality(unit, call_type) -> None: + # "The unit is a property of the modality" was stated in three docstrings and + # violated by six call sites in this file. MEDIA_UNIT_BY_CALL_TYPE turns the + # invariant into an enforced constraint: a (model, unit) price table is only + # usable if one modality never reports two units. + usage = TokenUsage() + with pytest.raises(ValueError, match="bills in"): + usage.record_media_call(unit=unit, quantity=1, call_type=call_type) + + assert usage.media_calls == 0 + assert usage.details == [] + + +@pytest.mark.parametrize( + ("call_type", "unit"), + [ + ("generate_image", "images"), + ("edit_image", "images"), + ("video", "seconds"), + ("asr", "seconds"), + ("music", "seconds"), + ("sound_effect", "seconds"), + ("tts", "characters"), + ("embedding", "texts"), + ("rerank", "requests"), + ], +) +def test_every_call_type_accepts_its_own_unit(call_type, unit) -> None: + # The mapping must cover every MediaCallType member, or a legitimate + # producer would be rejected. + usage = TokenUsage() + usage.record_media_call(unit=unit, quantity=1, call_type=call_type) + + assert usage.details[0]["unit"] == unit + + +def test_unit_is_unconstrained_when_call_type_is_omitted() -> None: + # call_type is optional metadata; with none given there is no modality to + # check the unit against. + usage = TokenUsage() + usage.record_media_call(unit="images", quantity=1) + + assert usage.details[0]["unit"] == "images" diff --git a/tests/core/tools/core/test_media_usage_helpers.py b/tests/core/tools/core/test_media_usage_helpers.py new file mode 100644 index 0000000000..fa745ad57f --- /dev/null +++ b/tests/core/tools/core/test_media_usage_helpers.py @@ -0,0 +1,198 @@ +"""Unit stability for duration-billed media tools. + +The billed unit must depend on the modality, never on how complete the +provider's response happened to be — a price table keyed on (model, unit) is +unusable if the same model sometimes reports "seconds" and sometimes +"requests". +""" + +import pytest + +from xagent.core.model.chat.token_context import ( + MediaCallType, + MediaUnit, + TokenContextManager, + aggregate_media_usage_by_model, +) +from xagent.core.tools.core.media_usage import ( + coerce_duration, + record_media_seconds, + record_media_usage, + resolve_billing_model, +) + + +@pytest.mark.parametrize( + "value,expected", + [ + (12.5, 12.5), + ("3", 3.0), + (0, None), # zero is "no usable duration", not a real measurement + (-5, None), + (None, None), + (True, None), # bool must not sneak through as 1.0 + ("nonsense", None), + ], +) +def test_coerce_duration(value, expected) -> None: + assert coerce_duration(value) == expected + + +def test_record_media_seconds_keeps_unit_stable_when_duration_missing() -> None: + # Same model, one call with a duration and one without: both must report + # seconds so billing sees a single line item, not two different units. + with TokenContextManager() as manager: + record_media_seconds(30.0, model="veo", call_type="video") + record_media_seconds(None, model="veo", call_type="video") + details = manager.get_usage().details + + assert [entry["unit"] for entry in details] == ["seconds", "seconds"] + # The unmeasured call records 0 seconds and is deliberately KEPT in the + # rollup: it is the only evidence that a billable provider call happened. + # It contributes nothing to quantity but still counts toward calls. + assert [entry["quantity"] for entry in details] == [30.0, 0.0] + groups = aggregate_media_usage_by_model(details) + assert len(groups) == 1 + assert groups[0]["unit"] == "seconds" + assert groups[0]["quantity"] == 30.0 + assert groups[0]["calls"] == 2 + + +def test_record_media_seconds_warns_when_unmeasured(caplog) -> None: + with caplog.at_level("WARNING"): + with TokenContextManager(): + record_media_seconds(None, model="veo", call_type="video") + assert "unmeasured" in caplog.text + + +def test_record_media_usage_never_raises() -> None: + # Accounting must never break the underlying media call. + with TokenContextManager() as manager: + record_media_usage("seconds", None, model="m", call_type="video") # type: ignore[arg-type] + assert manager.get_usage().media_calls == 1 + + +def test_resolve_billing_model_never_returns_a_placeholder() -> None: + """`_configured_model_id`-style lookups return Optional[str]; passing that + through str() records a model literally named "None".""" + + class _Model: + model_name = "elevenlabs-music-v1" + + # None id -> falls back to the provider's own name, not "None". + assert resolve_billing_model(None, _Model()) == "elevenlabs-music-v1" + assert resolve_billing_model("", _Model()) == "elevenlabs-music-v1" + # A real configured id always wins. + assert resolve_billing_model("cfg-id", _Model()) == "cfg-id" + # Nothing identifies the model: an explicit fallback, never None/"None". + assert resolve_billing_model(None, None) == "default" + + # Placeholder names on the model are not treated as identities. + class _Placeholder: + model_name = "None" + + assert resolve_billing_model(None, _Placeholder()) == "default" + + +def test_record_media_usage_swallows_recording_errors(monkeypatch) -> None: + # The best-effort guarantee is the whole point of this wrapper: a metering + # bug must never break the media call the user asked for. The existing + # quantity=None case coerces to 0 and records fine, so it would still pass + # with the try/except deleted — this one would not. + import xagent.core.tools.core.media_usage as mu + + def _boom(**kwargs: object) -> None: + raise RuntimeError("recording backend exploded") + + monkeypatch.setattr(mu, "add_media_usage", _boom) + + with TokenContextManager() as manager: + # Must not raise. + mu.record_media_usage( + MediaUnit.IMAGES, 1, model="m", call_type=MediaCallType.GENERATE_IMAGE + ) + usage = manager.get_usage() + + # And must leave no partial state behind. + assert usage.media_calls == 0 + assert usage.details == [] + + +def test_record_media_usage_swallows_invalid_unit(monkeypatch) -> None: + # A typo'd unit raises ValueError from the validator; the wrapper drops the + # record with a warning rather than propagating into the media call. + import xagent.core.tools.core.media_usage as mu + + with TokenContextManager() as manager: + mu.record_media_usage("not-a-unit", 1, model="m", call_type="tts") + usage = manager.get_usage() + + assert usage.media_calls == 0 + assert usage.details == [] + + +def test_record_media_seconds_swallows_invalid_call_type() -> None: + import xagent.core.tools.core.media_usage as mu + + with TokenContextManager() as manager: + mu.record_media_seconds(3.0, model="m", call_type="not-a-call-type") + usage = manager.get_usage() + + assert usage.media_calls == 0 + assert usage.details == [] + + +def test_record_media_seconds_ignores_non_finite_duration() -> None: + # inf > 0 is True, so without the finiteness guard this would record a + # non-JSON-serialisable quantity. + import xagent.core.tools.core.media_usage as mu + + with TokenContextManager() as manager: + mu.record_media_seconds( + mu.coerce_duration(float("inf")), model="m", call_type=MediaCallType.ASR + ) + entry = manager.get_usage().details[0] + + # Treated as "no duration reported": recorded as 0 seconds, unit unchanged. + assert entry["unit"] == "seconds" + assert entry["quantity"] == 0.0 + + +def test_record_media_usage_forwards_optional_billing_metadata() -> None: + # The wrapper must mirror the core primitive's optional fields: passing them + # previously raised TypeError during argument binding, before the try block. + import xagent.core.tools.core.media_usage as mu + + with TokenContextManager() as manager: + mu.record_media_usage( + MediaUnit.IMAGES, + 1, + model="m", + call_type=MediaCallType.GENERATE_IMAGE, + resolution="2K", + input_tokens=7, + output_tokens=3, + tokens_estimated=True, + ) + entry = manager.get_usage().details[0] + + assert entry["resolution"] == "2K" + assert entry["provider_tokens"] == 10 + assert entry["tokens_estimated"] is True + + +def test_resolve_billing_model_survives_raising_descriptor() -> None: + # Identity resolution runs before record_media_usage's try/except, so a + # provider whose model_name property raises would otherwise break the call. + import xagent.core.tools.core.media_usage as mu + + class _Hostile: + @property + def model_name(self) -> str: + raise RuntimeError("descriptor exploded") + + @property + def model(self) -> str: + return "real-name" + + assert mu.resolve_billing_model(None, _Hostile()) == "real-name" diff --git a/tests/web/tracking/test_task_tracker.py b/tests/web/tracking/test_task_tracker.py index 93c2ca173d..543d5d0339 100644 --- a/tests/web/tracking/test_task_tracker.py +++ b/tests/web/tracking/test_task_tracker.py @@ -1116,6 +1116,180 @@ def _hook(db, user_id, delta_details, delta_actions): assert len(captured["details"]) == 2 assert sorted(d["tokens"] for d in captured["details"]) == [5, 10] + @pytest.mark.asyncio + async def test_turn_delta_pairs_details_and_tool_calls_atomically(self, db_session): + """The details tail and the tool-call count must describe one instant. + + ``interrupt_reason_for_quota`` polls ``_turn_delta`` at every safe point + — once per agent step and once per streamed LLM chunk — against the live + shared TokenUsage. Reading the two fields in separate lock acquisitions + lets a concurrent write land between them, yielding a pair from two + different logical points in time, which is then fed to quota gating. + + The interleaving is forced from *inside* the locked read, by patching + ``copy_detail_rows`` — which ``detail_tail`` calls after taking the lock + and reading ``details``, but before reading ``tool_calls``. Two earlier + versions of this test wrapped ``detail_tail`` itself and released the + writer before the real accessor ran, so the write always completed + before the locked read began and the test passed against a deliberately + torn implementation. Mutation-verified: splitting ``detail_tail`` into + two lock acquisitions makes this fail with media=0, tool_calls=1. + """ + import threading + + from xagent.core.model.chat import token_context as tc_module + from xagent.core.model.chat.token_context import get_token_usage + + task = db_session.query.return_value.filter.return_value.first.return_value + task.user_id = 42 + task.input_tokens = 0 + task.output_tokens = 0 + task.llm_calls = 1 + task.token_usage_details = [] + + tracker = TaskTracker(task_id=123, db_session=db_session) + await tracker.start_tracking() + usage = get_token_usage() + + inside_locked_read = threading.Event() + writer_done = threading.Event() + real_copy = tc_module.copy_detail_rows + + def copy_then_pause(raw_details): + rows = real_copy(raw_details) + # Inside detail_tail's lock: details已读, tool_calls 还没读. + inside_locked_read.set() + writer_done.wait(timeout=5) + return rows + + results = {} + + def read_delta() -> None: + # usage passed explicitly — a bare thread inherits no contextvars. + results["delta"] = tracker._turn_delta(usage) + + def write_during_read() -> None: + assert inside_locked_read.wait(timeout=5), ( + "reader never entered the locked read" + ) + # These block until the reader releases the lock, when detail_tail is + # atomic. Against a torn implementation they land between its two + # acquisitions. + usage.record_media_call( + unit="images", quantity=1, call_type="generate_image" + ) + usage.increment_tool_calls(1) + writer_done.set() + + with patch.object(tc_module, "copy_detail_rows", copy_then_pause): + threads = [ + threading.Thread(target=read_delta), + threading.Thread(target=write_during_read), + ] + for thread in threads: + thread.start() + # The writer must not be released by a silent timeout: give it a + # nudge so an atomic implementation (where the writer blocks on the + # lock) still completes instead of both sides waiting. + threads[0].join(timeout=10) + writer_done.set() + for thread in threads: + thread.join(timeout=10) + + assert inside_locked_read.is_set(), "the interleaving never happened" + assert "delta" in results, "the reader never returned" + + delta_details, delta_actions = results["delta"] + media_rows = sum(1 for d in delta_details if d.get("type") == "media") + # Whichever side of the write the locked read landed on, the pair agrees: + # either both saw it (1 row, 1 action) or neither did (0 and 0). + assert media_rows == delta_actions, ( + f"torn pair: {media_rows} media rows vs {delta_actions} tool calls" + ) + + @pytest.mark.asyncio + async def test_media_rows_report_only_the_current_turn_delta(self, db_session): + """Media rows must behave like token rows across the TaskTracker seam. + + The seed derives ``media_calls`` from persisted media rows rather than a + column, so a multi-turn task must not re-report a prior turn's media + calls; and both the mid-run gate and the completion hook must see only + this turn's rows. + """ + from xagent.core.model.chat.token_context import ( + MediaCallType, + MediaUnit, + add_media_usage, + ) + from xagent.web.services import quota_hooks + + task = db_session.query.return_value.filter.return_value.first.return_value + task.user_id = 42 + task.input_tokens = 0 + task.output_tokens = 0 + task.llm_calls = 1 + # A prior turn already recorded two media calls, persisted as rows. + task.token_usage_details = [ + { + "type": "media", + "unit": "images", + "quantity": 2.0, + "model": "sd", + "call_type": "generate_image", + }, + { + "type": "media", + "unit": "seconds", + "quantity": 9.0, + "model": "whisper", + "call_type": "asr", + }, + ] + + recorded = {} + gated = {} + + def _usage_hook(db, user_id, delta_details, delta_actions): + recorded.update(details=delta_details, actions=delta_actions) + + def _gate_hook(db, user_id, delta_details, delta_actions): + gated.update(details=delta_details) + return None + + quota_hooks.set_usage_record_hook(_usage_hook) + quota_hooks.set_run_progress_gate_hook(_gate_hook) + try: + tracker = TaskTracker(task_id=123, db_session=db_session) + await tracker.start_tracking() + + # The seed counted the prior rows, so the baseline is not zero. + seeded = get_token_usage() + assert seeded.media_calls == 2 + + # This turn records one more media call. + add_media_usage( + unit=MediaUnit.SECONDS, + quantity=4.5, + model="whisper", + call_type=MediaCallType.ASR, + ) + + # A periodic snapshot must carry the running total, not drop it. + snapshot = get_token_usage().snapshot() + assert snapshot.media_calls == 3 + + # The mid-run gate sees only this turn's row. + await tracker.interrupt_reason_for_quota() + await tracker.complete_tracking() + finally: + quota_hooks.set_usage_record_hook(None) + quota_hooks.set_run_progress_gate_hook(None) + + for label, payload in (("gate", gated), ("completion", recorded)): + media = [d for d in payload["details"] if d.get("type") == "media"] + assert len(media) == 1, f"{label} hook saw {len(media)} media rows" + assert media[0]["quantity"] == 4.5, f"{label} hook saw a prior-turn row" + @pytest.mark.asyncio async def test_interrupt_reason_for_quota_passes_turn_delta(self, db_session): """The per-step quota gate must see the same this-turn delta the metering