From ffbecfed6323afadadd6f5de3fc3d337934eb150 Mon Sep 17 00:00:00 2001 From: OliverBryant <2713999266@qq.com> Date: Mon, 17 Aug 2026 16:26:27 +0800 Subject: [PATCH 1/9] feat: add media usage billing primitives MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Introduce the shared vocabulary and recording helpers for non-LLM media usage (image/video/audio/embedding/rerank), with no producers wired up yet — those follow in separate changes. - `MediaUnit` / `MediaCallType` name the billable dimension and the modality that produced an entry. The unit is a property of the modality, never of the response: a duration-billed call records `seconds` with `quantity=0` when unmeasured rather than switching to `requests`, so a price table keyed on (model, unit) stays usable. - `TokenUsage.add_media_usage` / `add_media_usage` write `type:"media"` detail entries into the existing `TokenUsage.details` list, so media flows through DB persistence and the quota `delta_details` contract without special-casing. Provider tokens are stored under `provider_tokens`, not `tokens`, so a consumer summing billable LLM tokens cannot pick up media counts. - `aggregate_media_usage_by_model` groups by (model, unit, call_type, resolution). Zero-quantity entries are kept deliberately: they are the only evidence that an unmeasured provider call happened. - `media_usage.py` adds `resolve_billing_model`, which never records a placeholder identity such as "default" or "None", plus the error-swallowing `record_media_usage` / `record_media_seconds` wrappers producers will use. Unknown `unit`/`call_type` values are now rejected at the write boundary. A typo would otherwise mint a new billing dimension that the aggregator keys off, and a usage record cannot be repaired once persisted. Validation runs before any mutation so a rejected call leaves no counter incremented without a matching detail entry. `TokenUsage` counter updates are now guarded by a lock. A single instance is shared across worker threads, where `+=` is a read-modify-write that silently loses counts; a 8x200 concurrent recording test drops ~85% of them without this. `merge()` snapshots the source under its own lock before taking the target's, so concurrent merges in both directions cannot deadlock. --- src/xagent/core/model/chat/__init__.py | 10 + src/xagent/core/model/chat/token_context.py | 430 ++++++++++++++++-- src/xagent/core/tools/core/media_usage.py | 156 +++++++ tests/core/model/chat/test_media_usage.py | 339 ++++++++++++++ .../tools/core/test_media_usage_helpers.py | 89 ++++ 5 files changed, 991 insertions(+), 33 deletions(-) create mode 100644 src/xagent/core/tools/core/media_usage.py create mode 100644 tests/core/model/chat/test_media_usage.py create mode 100644 tests/core/tools/core/test_media_usage_helpers.py diff --git a/src/xagent/core/model/chat/__init__.py b/src/xagent/core/model/chat/__init__.py index fa9e092989..0545fde72a 100644 --- a/src/xagent/core/model/chat/__init__.py +++ b/src/xagent/core/model/chat/__init__.py @@ -4,9 +4,14 @@ 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, get_and_reset_token_usage, get_token_usage, reset_token_usage, @@ -25,7 +30,12 @@ # Token tracking "TokenUsage", "TokenContextManager", + "MediaUnit", + "MediaCallType", "add_token_usage", + "add_media_usage", + "aggregate_token_usage_by_model", + "aggregate_media_usage_by_model", "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..dd3b8f3966 100644 --- a/src/xagent/core/model/chat/token_context.py +++ b/src/xagent/core/model/chat/token_context.py @@ -7,12 +7,82 @@ import contextvars import logging +import threading from dataclasses import dataclass, field +from enum import Enum from typing import Any, Dict, List, Optional logger = logging.getLogger(__name__) +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" + + +_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") -> 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. + """ + 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") -> 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. + """ + 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 + + @dataclass class TokenUsage: """Token usage statistics for a task or operation. @@ -21,19 +91,31 @@ 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 """ input_tokens: int = 0 output_tokens: int = 0 llm_calls: int = 0 + media_calls: int = 0 tool_calls: int = 0 details: List[Dict] = field(default_factory=list) + # Counter updates are read-modify-write, and one TokenUsage is routinely + # shared across worker threads (RAG ingestion pools, ``bind_usage_to_thread`` + # callers). Without this, concurrent ``+=`` silently loses counts. + # ``repr=False``/``compare=False`` keep the lock out of the dataclass's + # generated ``__repr__``/``__eq__``, which tests compare on. + _lock: threading.Lock = field( + default_factory=threading.Lock, repr=False, compare=False + ) @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 +133,132 @@ 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 add_media_usage( + self, + unit: str, + quantity: float, + model: str = "", + call_type: str = "", + model_id: str = "", + input_tokens: int = 0, + output_tokens: int = 0, + resolution: str = "", + tokens_estimated: bool = False, + ) -> None: + """Record a non-LLM media model call (image/video/tts/asr/...). + + ``unit`` names the billable dimension (see :class:`MediaUnit`); + ``quantity`` is its amount. ``resolution`` records the size tier + ("1K"/"2K"/"4K" or "1024x1024") so a per-(model, resolution) price table + can bill image models whose price varies by resolution; "" otherwise. + + Token passthrough is deliberately **not** stored under ``tokens``: that + key means "billable LLM tokens" on input/output entries, and a consumer + that naively sums it across all entries must not pick up media counts. + ``provider_tokens`` holds the provider-reported count instead, and + ``tokens_estimated`` marks it as a local heuristic (embedding/rerank) + rather than a measurement (Gemini / OpenAI gpt-image), so billing can + refuse to price an estimate. + """ + unit = _validated_media_unit(unit) + call_type = _validated_media_call_type(call_type) + with self._lock: self.details.append( { - "type": "output", - "tokens": tokens, + "type": "media", + "unit": unit, + "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, + "resolution": resolution, } ) def increment_llm_calls(self) -> None: """Increment the LLM call counter.""" - self.llm_calls += 1 + with self._lock: + self.llm_calls += 1 + + def increment_media_calls(self, count: int = 1) -> None: + """Increment the media (non-LLM) model call counter.""" + with self._lock: + self.media_calls += count def increment_tool_calls(self, count: int = 1) -> None: """Increment the tool-call counter (one per tool invocation).""" - self.tool_calls += count + 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 + details = list(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, + "details": list(self.details), + } @classmethod def from_dict(cls, data: Dict) -> "TokenUsage": @@ -115,6 +267,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", []), ) @@ -210,6 +363,57 @@ def _coerce_int(value: Any) -> int: return 0 +def estimate_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] + elif isinstance(text, (list, tuple, set)): + items = [item for item in text if isinstance(item, str)] + else: + 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 + or 0x3040 <= code <= 0x30FF + or 0xAC00 <= code <= 0xD7AF + ): + cjk += 1 + else: + other += 1 + return cjk + other // 4 + + +def _coerce_float(value: Any) -> float: + """Best-effort float; 0.0 if the value isn't a usable number. + + Media quantities can be fractional (e.g. audio seconds), so quantity uses + this rather than ``_coerce_int``. + """ + if isinstance(value, bool): + return float(value) + try: + return float(value) + except (TypeError, ValueError): + if value is not None: + logger.warning("Discarding non-numeric media quantity: %r", value) + return 0.0 + + def _usage_field(usage: Any, name: str) -> Any: """Read a field from a provider usage payload (SDK object or plain dict).""" if isinstance(usage, dict): @@ -338,6 +542,89 @@ 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. + """ + 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 +676,83 @@ def add_token_usage( ) +def add_media_usage( + unit: "MediaUnit | str", + quantity: float, + model: str = "", + call_type: "MediaCallType | str" = "", + 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. Providers + that also report real tokens (Gemini / OpenAI gpt-image) fill + input/output_tokens so a token-based price can take precedence. + 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. Producers route + through ``media_usage.record_media_usage``, which swallows this 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. + 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() + usage.increment_media_calls() + usage.add_media_usage( + 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..47f3a1fade --- /dev/null +++ b/src/xagent/core/tools/core/media_usage.py @@ -0,0 +1,156 @@ +"""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 +from typing import Any, Optional + +from typing_extensions import 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"): + value = getattr(model, attr, None) + 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 + return seconds if seconds > 0 else None + + +def record_media_usage( + unit: MediaUnit | str, + quantity: float, + *, + model: str = "", + model_id: str = "", + call_type: MediaCallType | str = "", +) -> 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, + ) + except Exception as e: # noqa: BLE001 + logger.warning("Failed to record %s media usage: %s", call_type, e) + + +def record_media_seconds( + seconds: Optional[float], + *, + model: str = "", + model_id: str = "", + call_type: MediaCallType | str = "", +) -> 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/tests/core/model/chat/test_media_usage.py b/tests/core/model/chat/test_media_usage.py new file mode 100644 index 0000000000..fbb98654ac --- /dev/null +++ b/tests/core/model/chat/test_media_usage.py @@ -0,0 +1,339 @@ +"""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 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_and_does_not_raise() -> None: + with TokenContextManager() as manager: + add_media_usage(unit="requests", quantity=None, model="x") # type: ignore[arg-type] + add_media_usage(unit="requests", quantity="oops", model="x") # 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_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="tts") + 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.add_media_usage(unit="images", quantity=1, model="x") + a.increment_media_calls() + + b = TokenUsage() + b.add_media_usage(unit="seconds", quantity=2, model="y") + b.increment_media_calls() + + 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="tts") + # 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="tts") + add_media_usage(unit="seconds", quantity=5, model="tts", call_type="tts") + 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_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_media_records_lose_no_counts() -> None: + # One TokenUsage is shared across worker threads (RAG ingestion pools, + # bind_usage_to_thread callers), where ``+=`` is a read-modify-write. + usage = TokenUsage() + workers, per_worker = 8, 200 + + def record() -> None: + for _ in range(per_worker): + usage.increment_media_calls() + usage.add_media_usage(unit="images", quantity=1, call_type="generate_image") + + threads = [threading.Thread(target=record) for _ in range(workers)] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + + expected = workers * per_worker + assert usage.media_calls == expected + assert len(usage.details) == expected + + +def test_concurrent_merge_loses_no_counts() -> None: + # merge() snapshots the source under its own lock before taking the + # target's, so concurrent merges neither deadlock nor drop entries. + target = TokenUsage() + sources = [] + for _ in range(8): + source = TokenUsage() + source.increment_media_calls() + source.add_media_usage(unit="seconds", quantity=2, call_type="asr") + sources.append(source) + + threads = [threading.Thread(target=target.merge, args=(s,)) for s in sources] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + + assert target.media_calls == len(sources) + assert len(target.details) == len(sources) 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..deebc7104b --- /dev/null +++ b/tests/core/tools/core/test_media_usage_helpers.py @@ -0,0 +1,89 @@ +"""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 ( + 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 and is dropped from the billable rollup. + 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 + + +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" From 979c05154555153f212aad8b1f937c60bdeef19b Mon Sep 17 00:00:00 2001 From: OliverBryant <2713999266@qq.com> Date: Mon, 17 Aug 2026 16:44:05 +0800 Subject: [PATCH 2/9] fix: normalize optional media metadata --- src/xagent/core/model/chat/token_context.py | 16 ++++++++++------ src/xagent/core/tools/core/media_usage.py | 6 +++--- tests/core/model/chat/test_media_usage.py | 19 +++++++++++++++++++ 3 files changed, 32 insertions(+), 9 deletions(-) diff --git a/src/xagent/core/model/chat/token_context.py b/src/xagent/core/model/chat/token_context.py index dd3b8f3966..f5e2e9ba23 100644 --- a/src/xagent/core/model/chat/token_context.py +++ b/src/xagent/core/model/chat/token_context.py @@ -51,7 +51,7 @@ class MediaCallType(str, Enum): _MEDIA_CALL_TYPE_VALUES = frozenset(member.value for member in MediaCallType) -def _validated_media_unit(unit: "MediaUnit | str") -> str: +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 @@ -59,6 +59,8 @@ def _validated_media_unit(unit: "MediaUnit | str") -> str: 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( @@ -68,12 +70,14 @@ def _validated_media_unit(unit: "MediaUnit | str") -> str: return value -def _validated_media_call_type(call_type: "MediaCallType | str") -> str: +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( @@ -167,10 +171,10 @@ def add_output_tokens( def add_media_usage( self, - unit: str, + unit: "MediaUnit | str | None", quantity: float, model: str = "", - call_type: str = "", + call_type: "MediaCallType | str | None" = "", model_id: str = "", input_tokens: int = 0, output_tokens: int = 0, @@ -677,10 +681,10 @@ def add_token_usage( def add_media_usage( - unit: "MediaUnit | str", + unit: "MediaUnit | str | None", quantity: float, model: str = "", - call_type: "MediaCallType | str" = "", + call_type: "MediaCallType | str | None" = "", model_id: str = "", input_tokens: int = 0, output_tokens: int = 0, diff --git a/src/xagent/core/tools/core/media_usage.py b/src/xagent/core/tools/core/media_usage.py index 47f3a1fade..4773c9c3ad 100644 --- a/src/xagent/core/tools/core/media_usage.py +++ b/src/xagent/core/tools/core/media_usage.py @@ -96,12 +96,12 @@ def coerce_duration(value: object) -> Optional[float]: def record_media_usage( - unit: MediaUnit | str, + unit: MediaUnit | str | None, quantity: float, *, model: str = "", model_id: str = "", - call_type: MediaCallType | str = "", + call_type: MediaCallType | str | None = "", ) -> None: """Record one media model call; swallow any error. @@ -129,7 +129,7 @@ def record_media_seconds( *, model: str = "", model_id: str = "", - call_type: MediaCallType | str = "", + call_type: MediaCallType | str | None = "", ) -> None: """Record a duration-billed media call, keeping the unit stable. diff --git a/tests/core/model/chat/test_media_usage.py b/tests/core/model/chat/test_media_usage.py index fbb98654ac..52cdabb601 100644 --- a/tests/core/model/chat/test_media_usage.py +++ b/tests/core/model/chat/test_media_usage.py @@ -279,6 +279,25 @@ def test_empty_call_type_is_allowed() -> None: 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( From 50afcc49d7fb617e2e1c781bba0f00d417edb858 Mon Sep 17 00:00:00 2001 From: OliverBryant <2713999266@qq.com> Date: Tue, 18 Aug 2026 16:11:31 +0800 Subject: [PATCH 3/9] fix: close numeric, snapshot and serialisation gaps in media usage MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Addresses review findings N1, N2, N3, N9, N11, N12, N13, N14 and N15 on the media billing primitives. Invalid quantities can no longer be persisted. `_coerce_float` rejected only non-numeric input, so booleans, negatives, NaN and infinities all reached the detail row — and NaN/inf are not JSON-serialisable, so they would emit literal `NaN`/`Infinity` tokens into task and quota details that strict parsers reject. All four now record 0.0 with a warning, matching the existing "call happened but is unmeasured" convention. `coerce_duration` gained the same finiteness guard, since `inf > 0` is True and would otherwise pass as a usable duration. Validation moved to the real write boundary. It previously lived only in the module-level `add_media_usage`, so anyone holding a `TokenUsage` bypassed it. Both writes now go through one path: the new atomic `record_media_call`, with the older `add_media_usage` method delegating to it. Self-review caught that these two copies had already drifted — the method appended unvalidated quantities that the function rejected. The counter and its detail row are now written under a single lock acquisition. Taking the lock twice let a concurrent `to_dict`/`merge` observe a count with no matching detail, and consumers derive per-model rows from `details` while reading the call count from the scalar, so a torn snapshot reports mutually inconsistent billing. `media_calls` now survives TaskTracker snapshots. `_copy_usage` and `_task_seed_from_session` both reconstructed `TokenUsage` without it, so periodic and final snapshots reported zero while the detail rows survived. The seed derives the count from the surviving media rows, since tasks have no `media_calls` column and seeding zero would re-report every media call from a prior turn. The mutation lock is no longer a dataclass field. As a field it sat in `__dataclass_fields__`, which `dataclasses.asdict` walks directly while ignoring `__getstate__`, so any generic consumer hit `TypeError: cannot pickle '_thread.lock' object`. It now lives in `__dict__` behind a property, with `__getstate__`/`__setstate__` dropping and recreating it, which also keeps the private slot out of the public positional signature. asdict, deepcopy and pickle all work. `record_media_usage` mirrors the core optional fields (`resolution`, `input_tokens`, `output_tokens`, `tokens_estimated`) and forwards them inside the guarded call; passing them previously raised `TypeError` during argument binding, before the `try` block. `resolve_billing_model` guards each provider attribute read individually, since it runs before that `try` and a `model_name` property that raises would have broken the user's media call over an accounting lookup. `estimate_tokens` accepts any iterable of strings as documented — generators previously estimated 0 — and rounds Latin character counts up, since truncating division estimated 0 tokens for any 1-3 character string. Empty input still estimates 0. Tests cover each fix and were written to fail on revert: invalid-quantity boundaries, a writer/reader interleaving assertion that every snapshot has `media_calls` equal to its media detail count, asdict/deepcopy/pickle round-trips, positional construction, opposing concurrent merges with a timeout, wrapper exception paths with a monkeypatched raising backend, a raising-descriptor resolver case, and generator/short-text token estimation. Deferred with issues rather than guessed at: #1460 covers the aggregate identity tagging, resolver return shape, placeholder fallback and input/output token split (all latent until producers mix the forms), and #1461 covers the discriminated `delta_details` union, which depends on an out-of-tree quota hook whose shape is not visible from this repository — no in-tree consumer indexes `detail["tokens"]`. --- src/xagent/core/model/chat/token_context.py | 146 ++++++++++++++-- src/xagent/core/tools/core/media_usage.py | 24 ++- src/xagent/web/tracking/task_tracker.py | 11 +- tests/core/model/chat/test_media_usage.py | 158 ++++++++++++++++++ .../tools/core/test_media_usage_helpers.py | 106 ++++++++++++ 5 files changed, 426 insertions(+), 19 deletions(-) diff --git a/src/xagent/core/model/chat/token_context.py b/src/xagent/core/model/chat/token_context.py index f5e2e9ba23..cf7f5fefe3 100644 --- a/src/xagent/core/model/chat/token_context.py +++ b/src/xagent/core/model/chat/token_context.py @@ -7,6 +7,7 @@ import contextvars import logging +import math import threading from dataclasses import dataclass, field from enum import Enum @@ -106,14 +107,37 @@ class TokenUsage: media_calls: int = 0 tool_calls: int = 0 details: List[Dict] = field(default_factory=list) - # Counter updates are read-modify-write, and one TokenUsage is routinely - # shared across worker threads (RAG ingestion pools, ``bind_usage_to_thread`` - # callers). Without this, concurrent ``+=`` silently loses counts. - # ``repr=False``/``compare=False`` keep the lock out of the dataclass's - # generated ``__repr__``/``__eq__``, which tests compare on. - _lock: threading.Lock = field( - default_factory=threading.Lock, repr=False, compare=False - ) + + def __post_init__(self) -> None: + # Counter updates are read-modify-write, and one TokenUsage is routinely + # shared across worker threads (RAG ingestion pools, + # ``bind_usage_to_thread`` callers); without a lock, concurrent ``+=`` + # silently loses counts. + # + # 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, media, tool, + # details)`` — free of a private slot. + self.__dict__["_lock"] = threading.Lock() + + @property + def _lock(self) -> threading.Lock: + """The per-instance mutation lock (see :meth:`__post_init__`).""" + return self.__dict__["_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 unpickled instance its own fresh lock.""" + self.__dict__.update(state) + self.__dict__["_lock"] = threading.Lock() @property def total_tokens(self) -> int: @@ -196,9 +220,63 @@ def add_media_usage( rather than a measurement (Gemini / OpenAI gpt-image), so billing can refuse to price an estimate. """ + # Delegates to record_media_call so there is exactly ONE media write + # path: two copies drifted apart during review, leaving this one + # appending unvalidated quantities (inf/negatives) that the other + # rejected. The only difference is that this method does not bump + # media_calls — callers of the older two-step API bump it themselves via + # increment_media_calls — so undo that here to keep the count honest. + self.record_media_call( + unit=unit, + quantity=quantity, + model=model, + call_type=call_type, + model_id=model_id, + input_tokens=input_tokens, + output_tokens=output_tokens, + resolution=resolution, + tokens_estimated=tokens_estimated, + # This method never bumped the counter: callers of the two-step API + # call increment_media_calls themselves. + count_call=False, + ) + + def record_media_call( + self, + unit: str, + quantity: float, + model: str = "", + call_type: str = "", + model_id: str = "", + input_tokens: int = 0, + output_tokens: int = 0, + resolution: str = "", + tokens_estimated: bool = False, + count_call: bool = True, + ) -> None: + """Count the call and append its detail under one lock acquisition. + + ``count_call=False`` appends the detail without bumping the counter, for + the older two-step API where the caller bumps it separately. + + ``increment_media_calls`` followed by ``add_media_usage`` takes the lock + twice, so a concurrent ``to_dict``/``merge`` can 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 snapshot reports mutually inconsistent billing. + """ unit = _validated_media_unit(unit) call_type = _validated_media_call_type(call_type) + # Validate here, not only in the module-level ``add_media_usage``: this + # is the actual write boundary, and a direct caller (tests, future + # producers holding a TokenUsage) would otherwise persist a boolean, + # negative, or non-finite quantity that cannot be repaired later. + quantity = _coerce_float(quantity) + input_tokens = _coerce_int(input_tokens) + output_tokens = _coerce_int(output_tokens) with self._lock: + if count_call: + self.media_calls += 1 self.details.append( { "type": "media", @@ -379,10 +457,16 @@ def estimate_tokens(text: Any) -> int: """ if isinstance(text, str): items: List[str] = [text] - elif isinstance(text, (list, tuple, set)): - items = [item for item in text if isinstance(item, str)] else: - return 0 + # 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 TypeError: + return 0 cjk = 0 other = 0 @@ -399,23 +483,50 @@ def estimate_tokens(text: Any) -> int: cjk += 1 else: other += 1 - return cjk + other // 4 + # 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: - """Best-effort float; 0.0 if the value isn't a usable number. + """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): - return float(value) + logger.warning("Discarding boolean media quantity: %r", value) + return 0.0 try: - return float(value) + result = float(value) except (TypeError, ValueError): 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: @@ -734,8 +845,9 @@ def add_media_usage( call_type_value = _validated_media_call_type(call_type) usage = get_token_usage() - usage.increment_media_calls() - usage.add_media_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( unit=unit_value, quantity=quantity, model=model, diff --git a/src/xagent/core/tools/core/media_usage.py b/src/xagent/core/tools/core/media_usage.py index 4773c9c3ad..e64c270b47 100644 --- a/src/xagent/core/tools/core/media_usage.py +++ b/src/xagent/core/tools/core/media_usage.py @@ -34,6 +34,7 @@ from __future__ import annotations import logging +import math from typing import Any, Optional from typing_extensions import TypeGuard @@ -72,7 +73,14 @@ def resolve_billing_model( if _usable_model_name(configured_id): return configured_id for attr in ("model_name", "model"): - value = getattr(model, attr, None) + # 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 @@ -92,6 +100,12 @@ def coerce_duration(value: object) -> Optional[float]: 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 @@ -102,6 +116,10 @@ def record_media_usage( 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. @@ -119,6 +137,10 @@ def record_media_usage( 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 logger.warning("Failed to record %s media usage: %s", call_type, e) diff --git a/src/xagent/web/tracking/task_tracker.py b/src/xagent/web/tracking/task_tracker.py index 841e47b523..8f2d87101a 100644 --- a/src/xagent/web/tracking/task_tracker.py +++ b/src/xagent/web/tracking/task_tracker.py @@ -62,6 +62,7 @@ def _copy_usage(usage: TokenUsage) -> TokenUsage: input_tokens=usage.input_tokens, output_tokens=usage.output_tokens, llm_calls=usage.llm_calls, + media_calls=usage.media_calls, tool_calls=usage.tool_calls, details=_copy_details(usage.details), ) @@ -100,13 +101,21 @@ 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_details(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: tasks have no media_calls column, and the next turn's + # delta is computed against this seed, so seeding zero would + # re-report every media call already recorded in a prior turn. + media_calls=sum( + 1 for detail in seed_details if detail.get("type") == "media" + ), + details=seed_details, ), ) diff --git a/tests/core/model/chat/test_media_usage.py b/tests/core/model/chat/test_media_usage.py index 52cdabb601..7e9d5a893b 100644 --- a/tests/core/model/chat/test_media_usage.py +++ b/tests/core/model/chat/test_media_usage.py @@ -356,3 +356,161 @@ def test_concurrent_merge_loses_no_counts() -> None: assert target.media_calls == len(sources) assert len(target.details) == len(sources) + + +@pytest.mark.parametrize( + "bad_quantity", + [True, False, -1, -0.5, float("nan"), float("inf"), float("-inf"), "abc", None], +) +def test_invalid_quantities_record_zero_not_garbage(bad_quantity) -> None: + # This is the last boundary before the value is persisted into task and + # quota details, and a written record cannot be repaired. NaN/inf are not + # JSON-serialisable (they emit literal NaN/Infinity that strict parsers + # reject); negatives would subtract from a bill; booleans are a caller bug. + with TokenContextManager() as manager: + add_media_usage( + unit="images", quantity=bad_quantity, model="m", call_type="generate_image" + ) + entry = manager.get_usage().details[0] + + assert entry["quantity"] == 0.0 + # Recorded, not dropped: the call still happened, it is just unmeasured. + assert manager.get_usage().media_calls == 1 + + +def test_valid_fractional_quantity_survives() -> None: + # The guard must not clamp legitimate fractional durations. + with TokenContextManager() as manager: + add_media_usage(unit="seconds", quantity=2.5, model="m", call_type="asr") + assert manager.get_usage().details[0]["quantity"] == 2.5 + + +def test_direct_record_media_call_also_validates() -> None: + # TokenUsage.record_media_call is the real write boundary; a caller holding + # a TokenUsage directly must get the same guarantees as add_media_usage. + usage = TokenUsage() + usage.record_media_call(unit="images", quantity=float("inf"), call_type="video") + + assert usage.details[0]["quantity"] == 0.0 + + +def test_media_count_and_detail_are_written_atomically() -> None: + # A snapshot taken between the counter bump and the detail append would + # report mutually inconsistent billing. Assert every observed snapshot has + # media_calls equal to the number of media detail rows. + 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.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)) + + threads = [threading.Thread(target=writer), threading.Thread(target=reader)] + for t in threads: + t.start() + for t in threads: + t.join() + + assert not torn, f"observed torn snapshots: {torn[:5]}" + assert usage.media_calls == 400 + + +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_media_calls_is_not_positionally_before_legacy_fields() -> None: + # Guards the public constructor shape: media_calls sits after llm_calls and + # before tool_calls, and the private lock must not occupy a positional slot. + usage = TokenUsage(1, 2, 3, 4, 5, []) + + assert usage.input_tokens == 1 + assert usage.output_tokens == 2 + assert usage.llm_calls == 3 + assert usage.media_calls == 4 + assert usage.tool_calls == 5 + + +def test_concurrent_merges_in_both_directions_do_not_deadlock() -> None: + # merge() snapshots the source under its own lock before taking the + # target's; holding both would deadlock on opposing merges. + a = TokenUsage() + b = TokenUsage() + for _ in range(50): + a.record_media_call(unit="images", quantity=1, call_type="generate_image") + b.record_media_call(unit="seconds", quantity=1, call_type="asr") + + threads = [ + threading.Thread(target=a.merge, args=(b,)), + threading.Thread(target=b.merge, args=(a,)), + ] + for t in threads: + t.start() + for t in threads: + t.join(timeout=10) + + for t in threads: + assert not t.is_alive(), "merge deadlocked in opposing directions" + # Neither side lost rows: each gained at least the other's original 50. + assert len(a.details) >= 100 + assert len(b.details) >= 100 + + +def test_estimate_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_tokens + + assert estimate_tokens(x for x in ["abcd", "efgh"]) == 2 + assert estimate_tokens(["abcd", "efgh"]) == 2 + # Non-iterables and non-string members are ignored, never raised on. + assert estimate_tokens(42) == 0 + assert estimate_tokens(["abcd", 42, None]) == 1 + + +@pytest.mark.parametrize( + ("text", "expected"), + [("", 0), ("a", 1), ("ab", 1), ("abc", 1), ("abcd", 1), ("abcde", 2)], +) +def test_estimate_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_tokens + + assert estimate_tokens(text) == expected + + +def test_estimate_tokens_counts_cjk_per_character() -> None: + from xagent.core.model.chat.token_context import estimate_tokens + + # CJK is roughly one token per character; a flat chars/4 heuristic would + # undercount Chinese by close to 4x. + assert estimate_tokens("中文") == 2 + assert estimate_tokens("中文abcd") == 3 diff --git a/tests/core/tools/core/test_media_usage_helpers.py b/tests/core/tools/core/test_media_usage_helpers.py index deebc7104b..d2a406103d 100644 --- a/tests/core/tools/core/test_media_usage_helpers.py +++ b/tests/core/tools/core/test_media_usage_helpers.py @@ -9,6 +9,8 @@ import pytest from xagent.core.model.chat.token_context import ( + MediaCallType, + MediaUnit, TokenContextManager, aggregate_media_usage_by_model, ) @@ -87,3 +89,107 @@ 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" From 832c52734035e10d944174378ee6356bb594a1e7 Mon Sep 17 00:00:00 2001 From: OliverBryant <2713999266@qq.com> Date: Tue, 18 Aug 2026 16:19:52 +0800 Subject: [PATCH 4/9] fix: satisfy mypy on the media usage write path Three errors from the previous commit: `record_media_call` now accepts `MediaUnit | str | None` and `MediaCallType | str | None`, matching its callers. It validates and narrows internally, so the narrow `str` annotation rejected the very values `add_media_usage` forwards. The validated results bind to separate `unit_value`/`call_type_value` locals rather than being assigned back over the widened parameters. The `_lock` property binds through a locally annotated variable, since reading it out of `__dict__` yields `Any` and returning that from a function declared to return `LockType` trips `no-any-return`. Verified with mypy 1.19.0 (the pinned pre-commit version) on the changed file: no issues. --- src/xagent/core/model/chat/token_context.py | 17 +++++++++-------- 1 file changed, 9 insertions(+), 8 deletions(-) diff --git a/src/xagent/core/model/chat/token_context.py b/src/xagent/core/model/chat/token_context.py index cf7f5fefe3..908a1cc886 100644 --- a/src/xagent/core/model/chat/token_context.py +++ b/src/xagent/core/model/chat/token_context.py @@ -124,9 +124,10 @@ def __post_init__(self) -> None: self.__dict__["_lock"] = threading.Lock() @property - def _lock(self) -> threading.Lock: + def _lock(self) -> "threading.Lock": """The per-instance mutation lock (see :meth:`__post_init__`).""" - return self.__dict__["_lock"] + lock: threading.Lock = self.__dict__["_lock"] + return lock def __getstate__(self) -> Dict[str, Any]: """Drop the lock: it is not picklable and is per-instance state.""" @@ -243,10 +244,10 @@ def add_media_usage( def record_media_call( self, - unit: str, + unit: "MediaUnit | str | None", quantity: float, model: str = "", - call_type: str = "", + call_type: "MediaCallType | str | None" = "", model_id: str = "", input_tokens: int = 0, output_tokens: int = 0, @@ -265,8 +266,8 @@ def record_media_call( Consumers derive per-model rows from ``details`` and the call count from the scalar, so a torn snapshot reports mutually inconsistent billing. """ - unit = _validated_media_unit(unit) - call_type = _validated_media_call_type(call_type) + unit_value = _validated_media_unit(unit) + call_type_value = _validated_media_call_type(call_type) # Validate here, not only in the module-level ``add_media_usage``: this # is the actual write boundary, and a direct caller (tests, future # producers holding a TokenUsage) would otherwise persist a boolean, @@ -280,7 +281,7 @@ def record_media_call( self.details.append( { "type": "media", - "unit": unit, + "unit": unit_value, "quantity": quantity, "provider_tokens": input_tokens + output_tokens, "provider_input_tokens": input_tokens, @@ -288,7 +289,7 @@ def record_media_call( "tokens_estimated": tokens_estimated, "model": model, "model_id": model_id, - "call_type": call_type, + "call_type": call_type_value, "resolution": resolution, } ) From 0ccbd35bd93ed5478cad02c3257de00f0685f68b Mon Sep 17 00:00:00 2001 From: OliverBryant <2713999266@qq.com> Date: Tue, 18 Aug 2026 17:37:49 +0800 Subject: [PATCH 5/9] fix: constrain REQUESTS quantity, sanitize token counts, restore positional ABI MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Second review round on the media billing primitives: C1, C4, C5, P7, P12, P15 and P16. C1 — MediaUnit.REQUESTS now rejects any quantity but 1. It is defined as exactly one provider call, so its quantity is not a free variable; letting 0/0.5/2 through made the billable quantity disagree with the media_calls and aggregate `calls` count derived from the same row, so quota and pricing would 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 — and rejected before any mutation, so no counter is bumped without its detail row. `record_media_usage` swallows it, keeping the promise that accounting never breaks a media call. C4 — `_coerce_int` is now a safe token boundary. It accepted booleans, preserved negatives, and let `int(float("inf"))` raise an uncaught `OverflowError` that propagated out of the accounting path and dropped the whole billable media row rather than just the bad field. It now mirrors `_coerce_float`: booleans, negatives and non-finite values sanitize to 0 with a warning while the row survives. Fixed on the shared helper rather than a media-only copy, since the LLM token paths pass provider-reported values through the same function and had the same overflow exposure. P12 — the positional ABI is restored. `media_calls` was inserted before the legacy `tool_calls`/`details` fields, so an existing `TokenUsage(input, output, llm_calls, tool_calls, details)` call bound its fourth argument to `media_calls` and shifted the rest, silently and with no error. The field is now appended after `details` and keyword-only. The test that previously asserted the broken order is rewritten to guard the legacy binding instead — it was institutionalising the bug. P7 — `TokenUsage.snapshot()` takes the lock and returns a detached copy; `TaskTracker._copy_usage` delegates to it. Reading the fields one by one from outside could interleave with a concurrent write and yield a counter without its matching detail row even though the writer is atomic. P16 — `estimate_tokens` catches any exception from iteration, not just `TypeError`. A custom iterable raising `ValueError`/`RuntimeError` escaped into an accounting path whose documented contract is that malformed input never breaks the call being measured. C5 — added a TaskTracker seam test that seeds prior-turn media rows, records a current-turn row, and asserts the seeded baseline is counted, the periodic snapshot carries the running total, and both the mid-run gate and the completion hook receive only this turn's row. P15 — the opposing-merge test now forces overlap with a `Barrier` and asserts exact counter/row agreement plus per-unit row survival, instead of `>= 100` lower bounds that would pass while rows were dropped. C2 (unbounded details persistence, O(n²) as high-frequency producers land) is deferred to #1466: the copy-and-rewrite mechanism predates this work, no producer is wired here, and changing the persistence shape is a much larger change than this PR. Self-review caught two things while making the above: the legacy two-step API leaves an orphan counter when a REQUESTS rejection happens (the caller owns that increment) — documented at the raise site, and unreachable since no in-tree producer uses the two-step path; and a comment describing a decrement that no longer exists. --- src/xagent/core/model/chat/token_context.py | 87 +++++++- src/xagent/web/tracking/task_tracker.py | 16 +- tests/core/model/chat/test_media_usage.py | 234 ++++++++++++++++++-- tests/web/tracking/test_task_tracker.py | 83 +++++++ 4 files changed, 379 insertions(+), 41 deletions(-) diff --git a/src/xagent/core/model/chat/token_context.py b/src/xagent/core/model/chat/token_context.py index 908a1cc886..4e37e0642f 100644 --- a/src/xagent/core/model/chat/token_context.py +++ b/src/xagent/core/model/chat/token_context.py @@ -104,9 +104,14 @@ class TokenUsage: input_tokens: int = 0 output_tokens: int = 0 llm_calls: int = 0 - media_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 @@ -119,8 +124,8 @@ def __post_init__(self) -> None: # 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, media, tool, - # details)`` — free of a private slot. + # positional signature — ``TokenUsage(input, output, llm_calls, + # tool_calls, details)`` — free of a private slot. self.__dict__["_lock"] = threading.Lock() @property @@ -226,7 +231,7 @@ def add_media_usage( # appending unvalidated quantities (inf/negatives) that the other # rejected. The only difference is that this method does not bump # media_calls — callers of the older two-step API bump it themselves via - # increment_media_calls — so undo that here to keep the count honest. + # increment_media_calls — hence count_call=False below. self.record_media_call( unit=unit, quantity=quantity, @@ -273,6 +278,26 @@ def record_media_call( # producers holding a TokenUsage) would otherwise persist a boolean, # negative, or non-finite quantity that cannot be repaired later. 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. + # + # Note the legacy two-step path (``increment_media_calls`` then + # ``add_media_usage``) bumps the counter in the CALLER, so a rejection + # there leaves a counted call with no detail row. That path is only + # reachable by an explicit two-step caller — every in-tree producer uses + # the single atomic entry point — and is exactly why the two-step API is + # kept only for backwards compatibility rather than used internally. + 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 {quantity!r}" + ) input_tokens = _coerce_int(input_tokens) output_tokens = _coerce_int(output_tokens) with self._lock: @@ -294,6 +319,28 @@ def record_media_call( } ) + 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, + details=[dict(item) for item in self.details if isinstance(item, dict)], + media_calls=self.media_calls, + ) + return copy + def increment_llm_calls(self) -> None: """Increment the LLM call counter.""" with self._lock: @@ -432,18 +479,35 @@ 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_tokens(text: Any) -> int: @@ -466,7 +530,12 @@ def estimate_tokens(text: Any) -> int: # an accounting path. try: items = [item for item in text if isinstance(item, str)] - except TypeError: + 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 diff --git a/src/xagent/web/tracking/task_tracker.py b/src/xagent/web/tracking/task_tracker.py index 8f2d87101a..ef74517602 100644 --- a/src/xagent/web/tracking/task_tracker.py +++ b/src/xagent/web/tracking/task_tracker.py @@ -57,15 +57,13 @@ def _copy_details(raw_details: Any) -> list[dict[str, Any]]: 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, - media_calls=usage.media_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( diff --git a/tests/core/model/chat/test_media_usage.py b/tests/core/model/chat/test_media_usage.py index 7e9d5a893b..f04ff37338 100644 --- a/tests/core/model/chat/test_media_usage.py +++ b/tests/core/model/chat/test_media_usage.py @@ -446,41 +446,74 @@ def test_usage_survives_asdict_deepcopy_and_pickle() -> None: assert revived.media_calls == 2 -def test_media_calls_is_not_positionally_before_legacy_fields() -> None: - # Guards the public constructor shape: media_calls sits after llm_calls and - # before tool_calls, and the private lock must not occupy a positional slot. - usage = TokenUsage(1, 2, 3, 4, 5, []) +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 - assert usage.input_tokens == 1 - assert usage.output_tokens == 2 - assert usage.llm_calls == 3 - assert usage.media_calls == 4 - assert usage.tool_calls == 5 + # 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 before taking the - # target's; holding both would deadlock on opposing merges. + # 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() - for _ in range(50): + 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=a.merge, args=(b,)), - threading.Thread(target=b.merge, args=(a,)), + threading.Thread(target=merge_after_barrier, args=(a, b)), + threading.Thread(target=merge_after_barrier, args=(b, a)), ] - for t in threads: - t.start() - for t in threads: - t.join(timeout=10) + for thread in threads: + thread.start() + for thread in threads: + thread.join(timeout=10) - for t in threads: - assert not t.is_alive(), "merge deadlocked in opposing directions" - # Neither side lost rows: each gained at least the other's original 50. - assert len(a.details) >= 100 - assert len(b.details) >= 100 + 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_tokens_accepts_any_iterable_of_strings() -> None: @@ -514,3 +547,158 @@ def test_estimate_tokens_counts_cjk_per_character() -> None: # undercount Chinese by close to 4x. assert estimate_tokens("中文") == 2 assert estimate_tokens("中文abcd") == 3 + + +@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="video") + 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="video") + snap = usage.snapshot() + + usage.record_media_call(unit="images", quantity=1, call_type="video") + + # 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_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_tokens + + class _RaisesMidIteration: + def __iter__(self): + yield "abcd" + raise RuntimeError("iterator exploded") + + class _RaisesImmediately: + def __iter__(self): + raise ValueError("cannot iterate") + + assert estimate_tokens(_RaisesMidIteration()) == 0 + assert estimate_tokens(_RaisesImmediately()) == 0 diff --git a/tests/web/tracking/test_task_tracker.py b/tests/web/tracking/test_task_tracker.py index 93c2ca173d..aac7ce8e66 100644 --- a/tests/web/tracking/test_task_tracker.py +++ b/tests/web/tracking/test_task_tracker.py @@ -1116,6 +1116,89 @@ 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_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 From 43e019e80a096ae61cac5c673795a5b51e84849e Mon Sep 17 00:00:00 2001 From: OliverBryant <2713999266@qq.com> Date: Tue, 18 Aug 2026 17:55:49 +0800 Subject: [PATCH 6/9] test: split the dirty-quantity case by unit after the REQUESTS constraint `test_dirty_quantity_is_coerced_and_does_not_raise` used the `requests` unit with `None`/`"oops"` quantities and asserted both coerce to 0. That is now a contradiction: REQUESTS pins its quantity to exactly 1, so those inputs are rejected rather than coerced. Split into the two behaviours the constraint creates, which is the replacement the review asked for: - free-quantity units (`seconds`) still coerce a malformed quantity to 0 and keep the row, because the provider call happened and must be visible as unmeasured; - `requests` rejects the same inputs and leaves no state behind, because coercing to 0 would record a call whose billable quantity contradicts its own call count. Checked the other `requests` usages in this file: both pass `quantity=1` and are unaffected. --- tests/core/model/chat/test_media_usage.py | 26 ++++++++++++++++++++--- 1 file changed, 23 insertions(+), 3 deletions(-) diff --git a/tests/core/model/chat/test_media_usage.py b/tests/core/model/chat/test_media_usage.py index f04ff37338..51e3fe2eb8 100644 --- a/tests/core/model/chat/test_media_usage.py +++ b/tests/core/model/chat/test_media_usage.py @@ -91,16 +91,36 @@ def test_estimated_tokens_are_flagged() -> None: assert aggregate_media_usage_by_model(details)[0]["tokens_estimated"] is True -def test_dirty_quantity_is_coerced_and_does_not_raise() -> None: +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="requests", quantity=None, model="x") # type: ignore[arg-type] - add_media_usage(unit="requests", quantity="oops", model="x") # type: ignore[arg-type] + 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") From 650663f6d566579b983a849af5e22cc7985ba0cf Mon Sep 17 00:00:00 2001 From: OliverBryant <2713999266@qq.com> Date: Wed, 19 Aug 2026 12:28:04 +0800 Subject: [PATCH 7/9] fix: snapshot before pairing turn-delta reads, and drop the two-step media API MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Third review round. NEW-A was blocking and was a regression I introduced last round; the rest are hardening plus the two simplifications the review recommended. NEW-A — `_turn_delta` now snapshots before reading. It read the live shared `TokenUsage` twice — the details slice, then the tool-call counter — with no lock, so a concurrent `record_media_call`/`increment_tool_calls` landing between them returned a pair describing two different instants, and that inconsistent pair fed quota gating. `interrupt_reason_for_quota` polls it at every safe point, so it is a hot concurrent path. This is the sibling of `_copy_usage`, which I re-pointed through `snapshot()` last round for exactly this reason — I fixed one call site and did not check the other in the same file. The snapshot is now taken unconditionally, including when the caller already passed a detached copy, so the two sites cannot diverge again. Reproduced with a forced interleaving: 200/200 inconsistent before, 0/200 after. Deleted the two-step media API — `TokenUsage.add_media_usage`, the `count_call` parameter, and `increment_media_calls`. It existed only to preserve a legacy flow with no production callers (grep-confirmed; only tests used it), and its documented behaviour was that a rejected call leaves `media_calls` incremented with no matching detail row — the exact torn state the REQUESTS constraint was added to prevent. Deleting removes the defect rather than hardening around it, and leaves `record_media_call` plus the module-level `add_media_usage`/`record_media_usage` as the only entry points. Tests migrated to the atomic call. `copy.copy` no longer shares state under a split lock. Shallow copy routed through `__getstate__`/`__setstate__`, giving the clone a fresh lock while still referencing the same `details` list — two objects each believing they held exclusive access to one list. `__copy__`/`__deepcopy__` now delegate to `snapshot()`, which is the supported way to fork a usage object. `to_dict` deep-copies each detail entry. It took the lock but returned a new outer list around the same inner dicts, so a caller mutating the payload mutated the live usage object outside that lock. Counter increments ignore negative counts, consistent with the coercion helpers; a negative would drive a counter below its matching row count. `record_media_usage` logs with `exc_info=True`. The broad catch is deliberate, but an unexpected bug surfaced as one context-free line, which is the wrong trade when persisted billing rows are unrepairable. Self-review found the same aliasing NEW-C describes in `merge`, which copied the outer list but shared inner dicts — mutating one usage object would silently rewrite the other's rows. Fixed with the rest. NEW-F (zero quantity overloading unmeasured and corrupt) is deferred to #1495: the current behaviour is deliberate and asserted by an existing test, and separating the two touches the persisted row shape, so it belongs with the first pricing consumer. C2 remains tracked in #1466. --- src/xagent/core/model/chat/token_context.py | 132 ++++++++------------ src/xagent/core/tools/core/media_usage.py | 9 +- src/xagent/web/tracking/task_tracker.py | 22 +++- tests/core/model/chat/test_media_usage.py | 93 ++++++++++++-- tests/web/tracking/test_task_tracker.py | 66 ++++++++++ 5 files changed, 231 insertions(+), 91 deletions(-) diff --git a/src/xagent/core/model/chat/token_context.py b/src/xagent/core/model/chat/token_context.py index 4e37e0642f..06514d477d 100644 --- a/src/xagent/core/model/chat/token_context.py +++ b/src/xagent/core/model/chat/token_context.py @@ -141,10 +141,29 @@ def __getstate__(self) -> Dict[str, Any]: return state def __setstate__(self, state: Dict[str, Any]) -> None: - """Restore, giving the unpickled instance its own fresh lock.""" + """Restore, giving the revived instance its own fresh lock.""" self.__dict__.update(state) self.__dict__["_lock"] = threading.Lock() + 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).""" @@ -199,54 +218,6 @@ def add_output_tokens( } ) - def add_media_usage( - 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, - ) -> None: - """Record a non-LLM media model call (image/video/tts/asr/...). - - ``unit`` names the billable dimension (see :class:`MediaUnit`); - ``quantity`` is its amount. ``resolution`` records the size tier - ("1K"/"2K"/"4K" or "1024x1024") so a per-(model, resolution) price table - can bill image models whose price varies by resolution; "" otherwise. - - Token passthrough is deliberately **not** stored under ``tokens``: that - key means "billable LLM tokens" on input/output entries, and a consumer - that naively sums it across all entries must not pick up media counts. - ``provider_tokens`` holds the provider-reported count instead, and - ``tokens_estimated`` marks it as a local heuristic (embedding/rerank) - rather than a measurement (Gemini / OpenAI gpt-image), so billing can - refuse to price an estimate. - """ - # Delegates to record_media_call so there is exactly ONE media write - # path: two copies drifted apart during review, leaving this one - # appending unvalidated quantities (inf/negatives) that the other - # rejected. The only difference is that this method does not bump - # media_calls — callers of the older two-step API bump it themselves via - # increment_media_calls — hence count_call=False below. - self.record_media_call( - unit=unit, - quantity=quantity, - model=model, - call_type=call_type, - model_id=model_id, - input_tokens=input_tokens, - output_tokens=output_tokens, - resolution=resolution, - tokens_estimated=tokens_estimated, - # This method never bumped the counter: callers of the two-step API - # call increment_media_calls themselves. - count_call=False, - ) - def record_media_call( self, unit: "MediaUnit | str | None", @@ -258,25 +229,23 @@ def record_media_call( output_tokens: int = 0, resolution: str = "", tokens_estimated: bool = False, - count_call: bool = True, ) -> None: """Count the call and append its detail under one lock acquisition. - ``count_call=False`` appends the detail without bumping the counter, for - the older two-step API where the caller bumps it separately. - - ``increment_media_calls`` followed by ``add_media_usage`` takes the lock - twice, so a concurrent ``to_dict``/``merge`` can 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 snapshot reports mutually inconsistent billing. + 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) - # Validate here, not only in the module-level ``add_media_usage``: this - # is the actual write boundary, and a direct caller (tests, future - # producers holding a TokenUsage) would otherwise persist a boolean, - # negative, or non-finite quantity that cannot be repaired later. + # 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. 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 @@ -286,13 +255,8 @@ def record_media_call( # 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. - # - # Note the legacy two-step path (``increment_media_calls`` then - # ``add_media_usage``) bumps the counter in the CALLER, so a rejection - # there leaves a counted call with no detail row. That path is only - # reachable by an explicit two-step caller — every in-tree producer uses - # the single atomic entry point — and is exactly why the two-step API is - # kept only for backwards compatibility rather than used internally. + # Rejecting before the lock is taken means a bad call leaves no state at + # all — neither a counter bump nor an orphan row. if unit_value == MediaUnit.REQUESTS.value and quantity != 1.0: raise ValueError( f"MediaUnit.REQUESTS means exactly one call, so quantity must " @@ -301,8 +265,7 @@ def record_media_call( input_tokens = _coerce_int(input_tokens) output_tokens = _coerce_int(output_tokens) with self._lock: - if count_call: - self.media_calls += 1 + self.media_calls += 1 self.details.append( { "type": "media", @@ -346,13 +309,16 @@ def increment_llm_calls(self) -> None: with self._lock: self.llm_calls += 1 - def increment_media_calls(self, count: int = 1) -> None: - """Increment the media (non-LLM) model call counter.""" - with self._lock: - self.media_calls += count - def increment_tool_calls(self, count: int = 1) -> None: - """Increment the tool-call counter (one per tool invocation).""" + """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 @@ -366,7 +332,11 @@ def merge(self, other: "TokenUsage") -> None: llm_calls = other.llm_calls media_calls = other.media_calls tool_calls = other.tool_calls - details = list(other.details) + # 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 = [dict(item) for item in other.details if isinstance(item, dict)] with self._lock: self.input_tokens += input_tokens self.output_tokens += output_tokens @@ -387,7 +357,11 @@ def to_dict(self) -> Dict: "llm_calls": self.llm_calls, "media_calls": self.media_calls, "tool_calls": self.tool_calls, - "details": list(self.details), + # 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. + "details": [dict(item) for item in self.details], } @classmethod diff --git a/src/xagent/core/tools/core/media_usage.py b/src/xagent/core/tools/core/media_usage.py index e64c270b47..ed37627816 100644 --- a/src/xagent/core/tools/core/media_usage.py +++ b/src/xagent/core/tools/core/media_usage.py @@ -143,7 +143,14 @@ def record_media_usage( tokens_estimated=tokens_estimated, ) except Exception as e: # noqa: BLE001 - logger.warning("Failed to record %s media usage: %s", call_type, e) + # 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 %s media usage: %s", call_type, e, exc_info=True + ) def record_media_seconds( diff --git a/src/xagent/web/tracking/task_tracker.py b/src/xagent/web/tracking/task_tracker.py index ef74517602..476f60ee23 100644 --- a/src/xagent/web/tracking/task_tracker.py +++ b/src/xagent/web/tracking/task_tracker.py @@ -538,12 +538,28 @@ async def stop_periodic_updates(self) -> None: def _turn_delta(self, usage: TokenUsage | None = None) -> tuple[list, 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. + + Snapshotting unconditionally — including when the caller already passed + a detached copy — keeps the two call sites from diverging again; the + cost is one extra shallow copy of a list that is about to be copied + anyway. + """ if usage is None: usage = get_token_usage() + snapshot = usage.snapshot() return ( - usage.details[self._initial_details_len :], - max(0, usage.tool_calls - self._initial_tool_calls), + snapshot.details[self._initial_details_len :], + max(0, snapshot.tool_calls - self._initial_tool_calls), ) async def interrupt_reason_for_quota(self) -> str | None: diff --git a/tests/core/model/chat/test_media_usage.py b/tests/core/model/chat/test_media_usage.py index 51e3fe2eb8..9fd3d861ab 100644 --- a/tests/core/model/chat/test_media_usage.py +++ b/tests/core/model/chat/test_media_usage.py @@ -142,12 +142,10 @@ def test_to_dict_from_dict_roundtrip_preserves_media() -> None: def test_merge_combines_media_calls_and_details() -> None: a = TokenUsage() - a.add_media_usage(unit="images", quantity=1, model="x") - a.increment_media_calls() + a.record_media_call(unit="images", quantity=1, model="x", call_type="video") b = TokenUsage() - b.add_media_usage(unit="seconds", quantity=2, model="y") - b.increment_media_calls() + b.record_media_call(unit="seconds", quantity=2, model="y", call_type="asr") a.merge(b) assert a.media_calls == 2 @@ -343,8 +341,9 @@ def test_concurrent_media_records_lose_no_counts() -> None: def record() -> None: for _ in range(per_worker): - usage.increment_media_calls() - usage.add_media_usage(unit="images", quantity=1, call_type="generate_image") + usage.record_media_call( + unit="images", quantity=1, call_type="generate_image" + ) threads = [threading.Thread(target=record) for _ in range(workers)] for thread in threads: @@ -364,8 +363,7 @@ def test_concurrent_merge_loses_no_counts() -> None: sources = [] for _ in range(8): source = TokenUsage() - source.increment_media_calls() - source.add_media_usage(unit="seconds", quantity=2, call_type="asr") + source.record_media_call(unit="seconds", quantity=2, call_type="asr") sources.append(source) threads = [threading.Thread(target=target.merge, args=(s,)) for s in sources] @@ -722,3 +720,82 @@ def __iter__(self): assert estimate_tokens(_RaisesMidIteration()) == 0 assert estimate_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="video") + 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="video") + 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 diff --git a/tests/web/tracking/test_task_tracker.py b/tests/web/tracking/test_task_tracker.py index aac7ce8e66..be7ee60f71 100644 --- a/tests/web/tracking/test_task_tracker.py +++ b/tests/web/tracking/test_task_tracker.py @@ -1116,6 +1116,72 @@ 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 slice and the tool-call count must describe one instant. + + ``interrupt_reason_for_quota`` polls ``_turn_delta`` at every safe point + against the live shared TokenUsage. Reading the two fields separately + lets a concurrent write land between them, yielding a pair from two + different logical points in time — a row visible with its counter bump + missing, or the reverse — which is then fed to quota gating. + """ + import threading + + 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() + + # Force a write into the window between the two reads _turn_delta makes. + entered = threading.Event() + released = threading.Event() + original_snapshot = type(usage).snapshot + + def slow_snapshot(self): + entered.set() + released.wait(timeout=5) + return original_snapshot(self) + + results = {} + + def read_delta() -> None: + results["delta"] = tracker._turn_delta() + + def write_during_read() -> None: + entered.wait(timeout=5) + usage.record_media_call( + unit="images", quantity=1, call_type="generate_image" + ) + usage.increment_tool_calls(1) + released.set() + + with patch.object(type(usage), "snapshot", slow_snapshot): + threads = [ + threading.Thread(target=read_delta), + threading.Thread(target=write_during_read), + ] + for thread in threads: + thread.start() + for thread in threads: + thread.join(timeout=10) + + 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 snapshot 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. From a48925817163742230044ed80fb037dacfec4065 Mon Sep 17 00:00:00 2001 From: OliverBryant <2713999266@qq.com> Date: Wed, 19 Aug 2026 15:34:18 +0800 Subject: [PATCH 8/9] perf: read only the delta tail on the quota path, and make its test able to fail MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Fourth review round. Both blocking findings were consequences of how I fixed NEW-A last round. N1 — `_turn_delta` no longer deep-copies the whole details list. `usage.snapshot()` copied the entire cumulative list under the lock and then discarded everything before `_initial_details_len`. The list grows monotonically across turns (seeds restore the persisted list), and the path is polled once per agent step *and* once per streamed LLM chunk (`runtime.py` `_raise_if_interrupted` -> `interrupt_reason_for_quota`), while holding the lock that serialises every LLM adapter's token write. Measured: 10.6us at 100 rows, 103us at 1000, 487us at 5000 — versus 0.3us for the tail. New `TokenUsage.detail_tail(start)` returns the tail plus `tool_calls` in one lock acquisition, preserving exactly the atomicity the snapshot was added for, at O(delta) instead of O(total). The now-redundant `_copy_details(delta_details)` at both quota call sites is gone, since the rows come back detached. N2 — the regression test for that fix could not fail. `read_delta` called `_turn_delta()` with no argument, so it went through `get_token_usage()` inside a bare `threading.Thread`, which does not inherit contextvars: it lazily built a fresh empty `TokenUsage` disconnected from the object the writer mutated, and observed `([], 0)` either way. Against the pre-fix code `snapshot()` was never called at all, so the patch never fired, `entered` was never set, and the writer's `wait()` timed out silently with its return value unchecked. The single assertion degenerated to `0 == 0` on both sides. `usage` is now passed explicitly, and the test asserts `entered.is_set()` and `released.is_set()` so a silent timeout fails. Mutation-verified this time against the committed test rather than a side script: 50/50 failures with two unlocked reads, 0/50 with `detail_tail`. That distinction is the gap that let the empty test through last round. Also addressed from this round: the `REQUESTS` error message now reports the caller's raw value rather than the coerced one (N7); `_lock` is an `RLock` so a future nested acquisition cannot self-deadlock (N5); `__post_init__` normalises `details`, covering the constructor and `from_dict` together (N3); `copy_detail_rows` replaces the same copy expression reimplemented four times, removing the filter asymmetry where `to_dict` would raise on a non-dict row that the others skipped; `estimate_tokens` is renamed `estimate_media_tokens` and exported, since three unrelated `estimate_tokens` already exist with different algorithms (N9); its CJK ranges gain punctuation, fullwidth forms and Ext-A (N10); the module docstring, the `add_media_usage` "Raises" section, the aggregators' detached-list contract, and the details-are-dicts invariant are documented (N13, N8, N14, N12); and three false or contradictory comments are corrected (N11, N16, plus the stale seed rationale). Test gaps closed: `_coerce_int`'s hardening exercised through the LLM `add_token_usage` path (N6), the `model_id` branch of `aggregate_media_usage_by_model` (N17), and `test_concurrent_merge_loses_no_counts` rewritten with a barrier and 8x100 merges — one merge per thread of a one-row source would have passed against an unlocked `merge` (N15). Self-review caught that adding normalisation to `__post_init__` made `snapshot()` copy twice, doubling it from 487us to 1066us at 5000 rows; `snapshot` now assigns `details` after construction, back to 526us. mypy also rejected `threading.RLock` as an annotation (it is a factory, not a class), so the property is typed through a `TYPE_CHECKING` alias. NEW-F stays in #1495, C2 in #1466, and P4/P5/P8/P9/P10 in #1460/#1461. N4 (string fields reaching the JSON column uncoerced) and N19 (`media_calls` has no consumer yet) are noted as accepted for this layer: the producers that supply those strings land in #1424/#1425/#1457, and #997 names the `media_calls` consumer. --- src/xagent/core/model/chat/__init__.py | 2 + src/xagent/core/model/chat/token_context.py | 129 +++++++-- src/xagent/web/tracking/task_tracker.py | 41 ++- tests/core/model/chat/test_media_usage.py | 253 +++++++++++------- .../tools/core/test_media_usage_helpers.py | 5 +- tests/web/tracking/test_task_tracker.py | 40 ++- 6 files changed, 322 insertions(+), 148 deletions(-) diff --git a/src/xagent/core/model/chat/__init__.py b/src/xagent/core/model/chat/__init__.py index 0545fde72a..4d4b465269 100644 --- a/src/xagent/core/model/chat/__init__.py +++ b/src/xagent/core/model/chat/__init__.py @@ -12,6 +12,7 @@ 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, @@ -36,6 +37,7 @@ "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 06514d477d..44d8cc0d8c 100644 --- a/src/xagent/core/model/chat/token_context.py +++ b/src/xagent/core/model/chat/token_context.py @@ -1,8 +1,15 @@ -"""Token usage tracking using contextvars. - -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. +"""Usage tracking using contextvars, for LLM tokens and non-LLM media calls. + +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 @@ -88,6 +95,19 @@ def _validated_media_call_type(call_type: "MediaCallType | str | None") -> str: 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: """Token usage statistics for a task or operation. @@ -126,12 +146,25 @@ def __post_init__(self) -> None: # 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.Lock() + 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.Lock": - """The per-instance mutation lock (see :meth:`__post_init__`).""" - lock: threading.Lock = self.__dict__["_lock"] + 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]: @@ -143,7 +176,7 @@ def __getstate__(self) -> Dict[str, Any]: 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.Lock() + self.__dict__["_lock"] = threading.RLock() def __copy__(self) -> "TokenUsage": """Shallow copy is unsupported; delegate to :meth:`snapshot`. @@ -246,6 +279,7 @@ def record_media_call( # 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 = 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 @@ -260,7 +294,7 @@ def record_media_call( 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 {quantity!r}" + f"be 1; got {raw_quantity!r}" ) input_tokens = _coerce_int(input_tokens) output_tokens = _coerce_int(output_tokens) @@ -282,6 +316,31 @@ def record_media_call( } ) + 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. @@ -299,9 +358,12 @@ def snapshot(self) -> "TokenUsage": output_tokens=self.output_tokens, llm_calls=self.llm_calls, tool_calls=self.tool_calls, - details=[dict(item) for item in self.details if isinstance(item, dict)], 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: @@ -336,7 +398,7 @@ def merge(self, other: "TokenUsage") -> None: # 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 = [dict(item) for item in other.details if isinstance(item, dict)] + details = copy_detail_rows(other.details) with self._lock: self.input_tokens += input_tokens self.output_tokens += output_tokens @@ -361,7 +423,12 @@ def to_dict(self) -> Dict: # 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. - "details": [dict(item) for item in self.details], + # 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 @@ -484,7 +551,7 @@ def _coerce_int(value: Any) -> int: return result -def estimate_tokens(text: Any) -> int: +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 @@ -520,9 +587,12 @@ def estimate_tokens(text: Any) -> int: # all roughly one token per character. code = ord(char) if ( - 0x4E00 <= code <= 0x9FFF - or 0x3040 <= code <= 0x30FF - or 0xAC00 <= code <= 0xD7AF + 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: @@ -611,6 +681,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 [] @@ -714,6 +790,12 @@ def aggregate_media_usage_by_model(details: Any) -> List[Dict[str, Any]]: 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 [] @@ -872,9 +954,14 @@ def add_media_usage( Raises: ValueError: If ``unit`` or ``call_type`` is not a known - :class:`MediaUnit` / :class:`MediaCallType` value. Producers route - through ``media_usage.record_media_usage``, which swallows this so - an accounting bug can never break the underlying media call. + :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. diff --git a/src/xagent/web/tracking/task_tracker.py b/src/xagent/web/tracking/task_tracker.py index 476f60ee23..fd9164d921 100644 --- a/src/xagent/web/tracking/task_tracker.py +++ b/src/xagent/web/tracking/task_tracker.py @@ -107,9 +107,10 @@ def _task_seed_from_session( output_tokens=_safe_int(getattr(task, "output_tokens", 0)), llm_calls=_safe_int(getattr(task, "llm_calls", 0)), # Derived from the surviving detail rows rather than read from a - # column: tasks have no media_calls column, and the next turn's - # delta is computed against this seed, so seeding zero would - # re-report every media call already recorded in a prior turn. + # 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" ), @@ -549,18 +550,28 @@ def _turn_delta(self, usage: TokenUsage | None = None) -> tuple[list, int]: polls this at every safe point, so it is a hot concurrent path rather than a one-shot teardown read. - Snapshotting unconditionally — including when the caller already passed - a detached copy — keeps the two call sites from diverging again; the - cost is one extra shallow copy of a list that is about to be copied - anyway. + 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() - snapshot = usage.snapshot() - return ( - snapshot.details[self._initial_details_len :], - max(0, snapshot.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 @@ -581,9 +592,10 @@ async def interrupt_reason_for_quota(self) -> str | None: return None try: delta_details, delta_actions = self._turn_delta() + # No _copy_details: _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: @@ -636,9 +648,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 _copy_details: _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 index 9fd3d861ab..ed2742acbd 100644 --- a/tests/core/model/chat/test_media_usage.py +++ b/tests/core/model/chat/test_media_usage.py @@ -357,91 +357,37 @@ def record() -> None: def test_concurrent_merge_loses_no_counts() -> None: - # merge() snapshots the source under its own lock before taking the - # target's, so concurrent merges neither deadlock nor drop entries. + # 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() - sources = [] - for _ in range(8): - source = TokenUsage() - source.record_media_call(unit="seconds", quantity=2, call_type="asr") - sources.append(source) - - threads = [threading.Thread(target=target.merge, args=(s,)) for s in sources] - for thread in threads: - thread.start() - for thread in threads: - thread.join() - - assert target.media_calls == len(sources) - assert len(target.details) == len(sources) - - -@pytest.mark.parametrize( - "bad_quantity", - [True, False, -1, -0.5, float("nan"), float("inf"), float("-inf"), "abc", None], -) -def test_invalid_quantities_record_zero_not_garbage(bad_quantity) -> None: - # This is the last boundary before the value is persisted into task and - # quota details, and a written record cannot be repaired. NaN/inf are not - # JSON-serialisable (they emit literal NaN/Infinity that strict parsers - # reject); negatives would subtract from a bill; booleans are a caller bug. - with TokenContextManager() as manager: - add_media_usage( - unit="images", quantity=bad_quantity, model="m", call_type="generate_image" - ) - entry = manager.get_usage().details[0] - - assert entry["quantity"] == 0.0 - # Recorded, not dropped: the call still happened, it is just unmeasured. - assert manager.get_usage().media_calls == 1 - - -def test_valid_fractional_quantity_survives() -> None: - # The guard must not clamp legitimate fractional durations. - with TokenContextManager() as manager: - add_media_usage(unit="seconds", quantity=2.5, model="m", call_type="asr") - assert manager.get_usage().details[0]["quantity"] == 2.5 - - -def test_direct_record_media_call_also_validates() -> None: - # TokenUsage.record_media_call is the real write boundary; a caller holding - # a TokenUsage directly must get the same guarantees as add_media_usage. - usage = TokenUsage() - usage.record_media_call(unit="images", quantity=float("inf"), call_type="video") - - assert usage.details[0]["quantity"] == 0.0 + workers, merges_per_worker = 8, 100 + source = TokenUsage() + source.record_media_call(unit="seconds", quantity=2, call_type="asr") -def test_media_count_and_detail_are_written_atomically() -> None: - # A snapshot taken between the counter bump and the detail append would - # report mutually inconsistent billing. Assert every observed snapshot has - # media_calls equal to the number of media detail rows. - usage = TokenUsage() - torn: list[tuple[int, int]] = [] - stop = threading.Event() + barrier = threading.Barrier(workers) - def writer() -> None: - for _ in range(400): - usage.record_media_call( - unit="images", quantity=1, call_type="generate_image" - ) - stop.set() + def merge_many() -> None: + barrier.wait(timeout=10) + for _ in range(merges_per_worker): + target.merge(source) - 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)) + threads = [threading.Thread(target=merge_many) for _ in range(workers)] + for thread in threads: + thread.start() + for thread in threads: + thread.join(timeout=30) - threads = [threading.Thread(target=writer), threading.Thread(target=reader)] - for t in threads: - t.start() - for t in threads: - t.join() + for thread in threads: + assert not thread.is_alive(), "merge deadlocked under contention" - assert not torn, f"observed torn snapshots: {torn[:5]}" - assert usage.media_calls == 400 + 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: @@ -534,37 +480,28 @@ def merge_after_barrier(target: TokenUsage, source: TokenUsage) -> None: assert len(own_rows) >= per_side, "a side lost its own rows" -def test_estimate_tokens_accepts_any_iterable_of_strings() -> None: +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_tokens + from xagent.core.model.chat.token_context import estimate_media_tokens - assert estimate_tokens(x for x in ["abcd", "efgh"]) == 2 - assert estimate_tokens(["abcd", "efgh"]) == 2 + 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_tokens(42) == 0 - assert estimate_tokens(["abcd", 42, None]) == 1 + 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_tokens_never_undercounts_short_text(text, expected) -> None: +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_tokens - - assert estimate_tokens(text) == expected + from xagent.core.model.chat.token_context import estimate_media_tokens - -def test_estimate_tokens_counts_cjk_per_character() -> None: - from xagent.core.model.chat.token_context import estimate_tokens - - # CJK is roughly one token per character; a flat chars/4 heuristic would - # undercount Chinese by close to 4x. - assert estimate_tokens("中文") == 2 - assert estimate_tokens("中文abcd") == 3 + assert estimate_media_tokens(text) == expected @pytest.mark.parametrize("bad_quantity", [0, 0.5, 2, 7, -1]) @@ -703,11 +640,11 @@ def test_snapshot_detaches_from_the_source() -> None: assert snap.details is not usage.details -def test_estimate_tokens_never_raises_from_a_hostile_iterable() -> None: +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_tokens + from xagent.core.model.chat.token_context import estimate_media_tokens class _RaisesMidIteration: def __iter__(self): @@ -718,8 +655,8 @@ class _RaisesImmediately: def __iter__(self): raise ValueError("cannot iterate") - assert estimate_tokens(_RaisesMidIteration()) == 0 - assert estimate_tokens(_RaisesImmediately()) == 0 + 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: @@ -799,3 +736,119 @@ def test_merge_does_not_alias_the_source_detail_dicts() -> None: 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" diff --git a/tests/core/tools/core/test_media_usage_helpers.py b/tests/core/tools/core/test_media_usage_helpers.py index d2a406103d..fa745ad57f 100644 --- a/tests/core/tools/core/test_media_usage_helpers.py +++ b/tests/core/tools/core/test_media_usage_helpers.py @@ -47,12 +47,15 @@ def test_record_media_seconds_keeps_unit_stable_when_duration_missing() -> None: details = manager.get_usage().details assert [entry["unit"] for entry in details] == ["seconds", "seconds"] - # The unmeasured call records 0 and is dropped from the billable rollup. + # 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: diff --git a/tests/web/tracking/test_task_tracker.py b/tests/web/tracking/test_task_tracker.py index be7ee60f71..1ca2ff5de4 100644 --- a/tests/web/tracking/test_task_tracker.py +++ b/tests/web/tracking/test_task_tracker.py @@ -1118,13 +1118,19 @@ def _hook(db, user_id, delta_details, delta_actions): @pytest.mark.asyncio async def test_turn_delta_pairs_details_and_tool_calls_atomically(self, db_session): - """The details slice and the tool-call count must describe one instant. + """The details tail and the tool-call count must describe one instant. ``interrupt_reason_for_quota`` polls ``_turn_delta`` at every safe point - against the live shared TokenUsage. Reading the two fields separately + — 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 — a row visible with its counter bump - missing, or the reverse — which is then fed to quota gating. + different logical points in time, which is then fed to quota gating. + + The reader runs in a plain thread, which does NOT inherit contextvars, so + ``usage`` is passed explicitly: relying on ``get_token_usage()`` there + would lazily create a fresh empty TokenUsage and the test would observe + ``([], 0)`` regardless of the implementation — passing whether or not the + race is fixed. """ import threading @@ -1141,30 +1147,34 @@ async def test_turn_delta_pairs_details_and_tool_calls_atomically(self, db_sessi await tracker.start_tracking() usage = get_token_usage() - # Force a write into the window between the two reads _turn_delta makes. + # Force a write into the window between the two reads, by making the + # locked read itself slow. Patching detail_tail's underlying lock + # acquisition is not possible, so widen the window from the writer side: + # the writer waits until the reader is inside the accessor. entered = threading.Event() released = threading.Event() - original_snapshot = type(usage).snapshot + real_detail_tail = type(usage).detail_tail - def slow_snapshot(self): + def slow_detail_tail(self, start=0): entered.set() released.wait(timeout=5) - return original_snapshot(self) + return real_detail_tail(self, start) results = {} def read_delta() -> None: - results["delta"] = tracker._turn_delta() + # usage passed explicitly — a bare thread inherits no contextvars. + results["delta"] = tracker._turn_delta(usage) def write_during_read() -> None: - entered.wait(timeout=5) + assert entered.wait(timeout=5), "reader never entered detail_tail" usage.record_media_call( unit="images", quantity=1, call_type="generate_image" ) usage.increment_tool_calls(1) released.set() - with patch.object(type(usage), "snapshot", slow_snapshot): + with patch.object(type(usage), "detail_tail", slow_detail_tail): threads = [ threading.Thread(target=read_delta), threading.Thread(target=write_during_read), @@ -1174,9 +1184,15 @@ def write_during_read() -> None: for thread in threads: thread.join(timeout=10) + # A silent timeout must fail rather than pass: without this the writer + # could never run and the assertion below would trivially hold at (0, 0). + assert entered.is_set(), "the interleaving never happened" + assert released.is_set(), "the writer never completed" + 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 snapshot landed on, the pair agrees: + # 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" From f9eb91ee718e0519f0a96a0e4ae733b912f46912 Mon Sep 17 00:00:00 2001 From: OliverBryant <2713999266@qq.com> Date: Wed, 19 Aug 2026 18:44:18 +0800 Subject: [PATCH 9/9] fix: enforce the unit-per-modality invariant, and make the lock tests real MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Fifth review round. Three of the five blocking findings were about claims I made rather than code I wrote, so each is verified by mutation against the committed test this time, not a side script. Locking tests now require the lock (finding 2). Replacing `_lock` with `nullcontext()` left all six concurrency tests passing. Investigating why showed my earlier "~85% count loss" figure was wrong: it came from a mutation that inserted `sleep(0)` between the read and the write, not from the real implementation. On CPython 3.12 even a bare `self.n += 1` loses nothing at 8x200000 threads-iterations, so no increment-counting test can justify the lock. What the lock actually protects is the *pairing* of the counter with its detail rows: `record_media_call` bumps and appends as one operation, and an unlocked reader lands between them. The test now asserts that invariant under `sys.setswitchinterval(1e-9)` (required — at the default interval the window is unhittable). Mutation-verified: 6/6 runs tear without the lock, 0/6 with it. `detail_tail`'s atomicity test now exercises the interleaving (finding 3, the second round this test was shown not to test its claim). Both earlier versions wrapped `detail_tail` itself and released the writer *before* delegating to the real accessor, so the write always completed before the locked read began. The interleaving is now forced from inside the lock, by patching `copy_detail_rows` — called after the lock is taken and `details` read, but before `tool_calls`. Mutation-verified against the reviewer's own torn implementation (two acquisitions with a sleep between): it fails with media=0, tool_calls=1. The unit-per-modality invariant is enforced, not just documented (finding 5). `MEDIA_UNIT_BY_CALL_TYPE` rejects a unit that does not match the call type at the write boundary. The invariant was stated in three docstrings and violated by this PR's own tests; the reviewer found six sites, and the new check found four more (`seconds` paired with `tts`) that a grep for `images`+`video` missed. All ten fixed. Verified every downstream producer in #1424/#1425/#1457 pairs legally, so the new constraint breaks none of them. `_coerce_float` catches `OverflowError` (finding 1). `float(10**400)` raised it uncaught, so the error-swallowing wrapper dropped the entire billing row — the opposite of the documented reject-to-0.0 contract. `_coerce_int` was already fixed for this in an earlier round; the sibling was missed. The token-vs-unit pricing precedence claim is removed rather than implemented (finding 4). It existed only as a docstring sentence: nothing expressed or enforced "price by tokens instead of by unit", the row schema carries no discriminator, and `provider_tokens` has zero consumers outside this file. Implementing it would add a billing dimension to a PR that wires no producers and has no pricing consumer, so the aspirational sentence is gone and the design question stays tracked in #1461. Minor: `add_media_usage` reports the caller's raw quantity in the REQUESTS error (the prior fix covered only the direct path); the wrapper's warning log now includes model/model_id/unit/quantity, not just call_type; the two stale `bind_usage_to_thread` references are dropped (the symbol does not exist); `_turn_delta`'s return type is parameterised. Simplifications taken: `TypeGuard` from stdlib `typing` rather than `typing_extensions` (the project requires >=3.11), and `task_tracker`'s `_copy_details` is deleted in favour of the `copy_detail_rows` helper this PR added for exactly that consolidation. Not changed, with reasons: `resolve_billing_model` and `record_media_seconds` are flagged as having no production caller, but both are consumed by #1425/#1457 in this same series — removing them here would only move the diff. Findings 7-9 and 11-16 are noted as accepted or tracked in the reply. --- src/xagent/core/model/chat/token_context.py | 65 +++++++- src/xagent/core/tools/core/media_usage.py | 14 +- src/xagent/web/tracking/task_tracker.py | 19 +-- tests/core/model/chat/test_media_usage.py | 155 +++++++++++++++++--- tests/web/tracking/test_task_tracker.py | 55 ++++--- 5 files changed, 238 insertions(+), 70 deletions(-) diff --git a/src/xagent/core/model/chat/token_context.py b/src/xagent/core/model/chat/token_context.py index 44d8cc0d8c..d47c8a1e5d 100644 --- a/src/xagent/core/model/chat/token_context.py +++ b/src/xagent/core/model/chat/token_context.py @@ -22,6 +22,9 @@ 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. @@ -55,6 +58,27 @@ class MediaCallType(str, Enum): 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) @@ -135,9 +159,9 @@ class TokenUsage: def __post_init__(self) -> None: # Counter updates are read-modify-write, and one TokenUsage is routinely - # shared across worker threads (RAG ingestion pools, - # ``bind_usage_to_thread`` callers); without a lock, concurrent ``+=`` - # silently loses counts. + # 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 @@ -262,6 +286,7 @@ def record_media_call( 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. @@ -279,7 +304,9 @@ def record_media_call( # 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 = quantity + # _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 @@ -291,6 +318,14 @@ def record_media_call( # 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 " @@ -630,7 +665,11 @@ def _coerce_float(value: Any) -> float: return 0.0 try: result = float(value) - except (TypeError, ValueError): + 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 @@ -946,9 +985,14 @@ def add_media_usage( 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. Providers - that also report real tokens (Gemini / OpenAI gpt-image) fill - input/output_tokens so a token-based price can take precedence. + 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. @@ -965,6 +1009,10 @@ def add_media_usage( """ # 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) @@ -979,6 +1027,7 @@ def add_media_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, diff --git a/src/xagent/core/tools/core/media_usage.py b/src/xagent/core/tools/core/media_usage.py index ed37627816..78711762ea 100644 --- a/src/xagent/core/tools/core/media_usage.py +++ b/src/xagent/core/tools/core/media_usage.py @@ -35,9 +35,7 @@ import logging import math -from typing import Any, Optional - -from typing_extensions import TypeGuard +from typing import Any, Optional, TypeGuard from ...model.chat.token_context import MediaCallType, MediaUnit, add_media_usage @@ -149,7 +147,15 @@ def record_media_usage( # after the fact, so a metering bug needs to be diagnosable from the log # alone. logger.warning( - "Failed to record %s media usage: %s", call_type, e, exc_info=True + "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, ) diff --git a/src/xagent/web/tracking/task_tracker.py b/src/xagent/web/tracking/task_tracker.py index fd9164d921..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,12 +51,6 @@ 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. @@ -99,7 +94,7 @@ 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_details(getattr(task, "token_usage_details", None)) + seed_details = copy_detail_rows(getattr(task, "token_usage_details", None)) return _TaskTrackingSeed( user_id=_optional_int(getattr(task, "user_id", None)), usage=TokenUsage( @@ -168,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, ) @@ -536,7 +531,9 @@ 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. @@ -592,7 +589,7 @@ async def interrupt_reason_for_quota(self) -> str | None: return None try: delta_details, delta_actions = self._turn_delta() - # No _copy_details: _turn_delta already returns detached rows. + # No extra copy: _turn_delta already returns detached rows. reason = _check_quota_on_event_loop( self._user_id, delta_details, @@ -648,7 +645,7 @@ 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 _copy_details: _turn_delta already returns detached rows. + # No extra copy: _turn_delta already returns detached rows. _record_usage_on_event_loop( self._user_id, delta_details, diff --git a/tests/core/model/chat/test_media_usage.py b/tests/core/model/chat/test_media_usage.py index ed2742acbd..7083cfd6c3 100644 --- a/tests/core/model/chat/test_media_usage.py +++ b/tests/core/model/chat/test_media_usage.py @@ -5,6 +5,7 @@ persistence and the quota ``delta_details`` contract without special-casing. """ +import sys import threading import pytest @@ -124,7 +125,7 @@ def test_dirty_quantity_on_requests_unit_is_rejected_not_coerced() -> None: 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="tts") + add_media_usage(unit="seconds", quantity=3.5, model="tts", call_type="asr") usage = manager.get_usage() data = usage.to_dict() @@ -142,7 +143,9 @@ def test_to_dict_from_dict_roundtrip_preserves_media() -> None: def test_merge_combines_media_calls_and_details() -> None: a = TokenUsage() - a.record_media_call(unit="images", quantity=1, model="x", call_type="video") + 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") @@ -175,7 +178,7 @@ def test_media_aggregation_groups_by_model_unit_and_call_type() -> None: add_media_usage( unit="images", quantity=3, model="sd", call_type="generate_image" ) - add_media_usage(unit="seconds", quantity=4, model="tts", call_type="tts") + 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 @@ -236,8 +239,8 @@ def test_zero_quantity_media_entries_stay_visible() -> None: # 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="tts") - add_media_usage(unit="seconds", quantity=5, model="tts", call_type="tts") + 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) @@ -333,27 +336,55 @@ def test_enum_members_are_accepted() -> None: assert isinstance(entry["unit"], str) -def test_concurrent_media_records_lose_no_counts() -> None: - # One TokenUsage is shared across worker threads (RAG ingestion pools, - # bind_usage_to_thread callers), where ``+=`` is a read-modify-write. +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() - workers, per_worker = 8, 200 + torn: list[tuple[int, int]] = [] + stop = threading.Event() - def record() -> None: - for _ in range(per_worker): + # 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() - threads = [threading.Thread(target=record) for _ in range(workers)] - for thread in threads: - thread.start() - for thread in threads: - thread.join() + 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 - expected = workers * per_worker - assert usage.media_calls == expected - assert len(usage.details) == expected + 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: @@ -607,7 +638,9 @@ def test_snapshot_is_taken_under_the_lock() -> None: def writer() -> None: for _ in range(400): - usage.record_media_call(unit="images", quantity=1, call_type="video") + usage.record_media_call( + unit="images", quantity=1, call_type="generate_image" + ) stop.set() def reader() -> None: @@ -629,10 +662,10 @@ def reader() -> None: def test_snapshot_detaches_from_the_source() -> None: usage = TokenUsage() - usage.record_media_call(unit="images", quantity=1, call_type="video") + 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="video") + 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 @@ -668,13 +701,13 @@ def test_copy_copy_does_not_share_details_under_a_different_lock() -> None: import copy usage = TokenUsage() - usage.record_media_call(unit="images", quantity=1, call_type="video") + 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="video") + 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 @@ -852,3 +885,77 @@ def test_media_aggregation_uses_model_id_when_present() -> None: 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/web/tracking/test_task_tracker.py b/tests/web/tracking/test_task_tracker.py index 1ca2ff5de4..543d5d0339 100644 --- a/tests/web/tracking/test_task_tracker.py +++ b/tests/web/tracking/test_task_tracker.py @@ -1126,14 +1126,18 @@ async def test_turn_delta_pairs_details_and_tool_calls_atomically(self, db_sessi 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 reader runs in a plain thread, which does NOT inherit contextvars, so - ``usage`` is passed explicitly: relying on ``get_token_usage()`` there - would lazily create a fresh empty TokenUsage and the test would observe - ``([], 0)`` regardless of the implementation — passing whether or not the - race is fixed. + 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 @@ -1147,18 +1151,16 @@ async def test_turn_delta_pairs_details_and_tool_calls_atomically(self, db_sessi await tracker.start_tracking() usage = get_token_usage() - # Force a write into the window between the two reads, by making the - # locked read itself slow. Patching detail_tail's underlying lock - # acquisition is not possible, so widen the window from the writer side: - # the writer waits until the reader is inside the accessor. - entered = threading.Event() - released = threading.Event() - real_detail_tail = type(usage).detail_tail + inside_locked_read = threading.Event() + writer_done = threading.Event() + real_copy = tc_module.copy_detail_rows - def slow_detail_tail(self, start=0): - entered.set() - released.wait(timeout=5) - return real_detail_tail(self, start) + 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 = {} @@ -1167,27 +1169,34 @@ def read_delta() -> None: results["delta"] = tracker._turn_delta(usage) def write_during_read() -> None: - assert entered.wait(timeout=5), "reader never entered detail_tail" + 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) - released.set() + writer_done.set() - with patch.object(type(usage), "detail_tail", slow_detail_tail): + 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) - # A silent timeout must fail rather than pass: without this the writer - # could never run and the assertion below would trivially hold at (0, 0). - assert entered.is_set(), "the interleaving never happened" - assert released.is_set(), "the writer never completed" + assert inside_locked_read.is_set(), "the interleaving never happened" assert "delta" in results, "the reader never returned" delta_details, delta_actions = results["delta"]