From 78931f4222b934a1be396fbf7d66cf1d0b4d40fa Mon Sep 17 00:00:00 2001 From: OliverBryant <2713999266@qq.com> Date: Mon, 17 Aug 2026 17:33:18 +0800 Subject: [PATCH 1/6] feat: meter audio and video generation MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Wire the audio and video modalities into the media usage primitives: ASR billed by transcribed seconds, TTS by input characters, and music, sound effects and video by duration. Video shares the duration-billed path because it shares the same unit invariants, not to pad the change. Also lands `web/tracking/standalone_usage.py`, which `/speech/transcribe` needs here: it binds a TokenUsage for work that is not a tracked agent task and reports it to the quota hook on exit. Without it that endpoint records into a throwaway object nothing reads — the provider bills and the usage evaporates. `bind_usage_to_thread` runs each call inside `copy_context()` + `ctx.run` so a binding cannot outlive the job on a pooled executor thread, and `_report` checks `has_usage_record_hook()` before checking out a DB session, since with no hook installed the record call is a guaranteed no-op. Fix ASR recording 0 seconds on every call outside audio_tool. `/speech/transcribe` and the Telegram voice handler both called `transcribe()` without `verbose=True`, so the provider returned a bare string with no timings and every record was an unbillable 0-second entry. Both now pass it; each already read the text via `getattr(result, "text", result)`, so the richer result needs no other handling. The Telegram test fakes accepted no `verbose` argument at all, unlike every real provider, so they were widened to match the real signature — a test double narrower than the interface it stands in for is what let this ship. Unify ASR billing identity on the model name. The three entry points previously wrote three different identities for one physical model: audio_tool wrote a name into `model_id` (its registries are keyed by `model_name`, so the membership test was a name test), `/speech/ transcribe` wrote the real DB id, and Telegram read a bare `.model` attribute with no placeholder filter. The aggregator groups on `model_id or model`, so one model produced up to three rows that could never be reconciled once written. Only the name is available on all three paths, so `model_id` is now left unset everywhere and Telegram goes through the shared resolver. The same name-into-`model_id` mistake is fixed on both TTS paths. Bill provider calls that already happened. Music and sound effects raised on an empty or malformed result *before* recording, so a call that succeeded at the HTTP level but returned nothing went unmetered even though the provider charged for it. Recording now precedes validation in both, matching TTS, and the policy is stated in a comment at each site — the four otherwise-symmetric tools previously had three different implicit rules. Delete the duplicate billing helpers. `record_asr_seconds` is now a thin wrapper over the shared `record_media_seconds`, and audio_tool's private `_resolve_billing_model` copy delegates to the shared `resolve_billing_model` instead of reimplementing it without the placeholder filter. That copy billed the Xinference default ASR/TTS model — which exposes no `model_name` and whose default-getter builds a separate instance, so identity matching never hit — under the literal string "default". It now falls back to the provider class name, which attributes cost to something real. --- src/xagent/core/model/asr/usage.py | 115 +++++++++++++ src/xagent/core/tools/core/audio_tool.py | 88 +++++++++- src/xagent/core/tools/core/music_tool.py | 32 ++++ .../core/tools/core/sound_effect_tool.py | 29 ++++ src/xagent/core/tools/core/video_tool.py | 20 +++ src/xagent/web/api/model.py | 41 ++++- src/xagent/web/channels/telegram/bot.py | 17 ++ src/xagent/web/services/quota_hooks.py | 42 ++++- src/xagent/web/tracking/standalone_usage.py | 142 ++++++++++++++++ .../core/test_media_tool_billing_policy.py | 156 ++++++++++++++++++ tests/core/tools/test_audio_tool.py | 121 ++++++++++++++ .../test_connector_runtime_entrypoints_e2e.py | 9 +- tests/web/test_telegram_message_queue.py | 35 +++- 13 files changed, 825 insertions(+), 22 deletions(-) create mode 100644 src/xagent/core/model/asr/usage.py create mode 100644 src/xagent/web/tracking/standalone_usage.py create mode 100644 tests/core/tools/core/test_media_tool_billing_policy.py diff --git a/src/xagent/core/model/asr/usage.py b/src/xagent/core/model/asr/usage.py new file mode 100644 index 0000000000..00dde17144 --- /dev/null +++ b/src/xagent/core/model/asr/usage.py @@ -0,0 +1,115 @@ +"""Media-usage recording for ASR. + +ASR is reached from several entry points that do not go through ``audio_tool`` +— the Telegram channel and the ``/speech/transcribe`` API both call +``transcribe`` directly — so this module is the single place that knows how to +turn a transcription result into a usage record. Every ASR caller routes here, +including ``audio_tool``: keeping one implementation is what stops the copies +from drifting (an earlier pair differed in whether they logged at all). +""" + +from __future__ import annotations + +import logging +from typing import Any, Optional, Sequence, Union + +from ...tools.core.media_usage import coerce_duration, record_media_seconds +from ..chat.token_context import MediaCallType +from .base import ASRResult + +logger = logging.getLogger(__name__) + +# Keys providers use for the total length of the submitted audio. +_DURATION_KEYS = ("duration", "audio_duration", "duration_seconds") + + +def duration_from_raw_response(raw_response: Any) -> Optional[float]: + """Provider-reported total audio duration, if it exposed one. + + Preferred over segment timings: the end of the last spoken segment + undercounts a recording with trailing silence, which the provider still + processed and billed for. + """ + if not isinstance(raw_response, dict): + return None + for key in _DURATION_KEYS: + seconds = coerce_duration(raw_response.get(key)) + if seconds is not None: + return seconds + return None + + +def duration_from_segments(segments: Optional[Sequence[Any]]) -> Optional[float]: + """Transcribed duration inferred from the end of the last timed segment. + + Accepts both ``ASRSegment`` objects and the plain dicts ``audio_tool`` + builds from them. + """ + if not segments: + return None + last_end = 0.0 + for segment in segments: + end = ( + segment.get("end") + if isinstance(segment, dict) + else getattr(segment, "end", None) + ) + seconds = coerce_duration(end) + if seconds is not None: + last_end = max(last_end, seconds) + return last_end if last_end > 0 else None + + +def resolve_asr_seconds( + raw_response: Any = None, + segments: Optional[Sequence[Any]] = None, +) -> Optional[float]: + """Best available transcribed-audio duration, provider field first.""" + return duration_from_raw_response(raw_response) or duration_from_segments(segments) + + +def record_asr_seconds( + seconds: Optional[float], + *, + model_name: str = "", + model_id: str = "", +) -> None: + """Record one transcription from an already-resolved duration. + + Thin wrapper over the shared :func:`record_media_seconds`: ASR is one of + several duration-billed modalities, and keeping a second implementation is + how the two drifted before. The shared helper owns the invariants — unit + always seconds (a call whose duration is unknown records 0 rather than + switching units, which would make a (model, unit) price table unusable), + the unmeasured-call warning, and swallowing errors so accounting never + breaks a transcription. + """ + record_media_seconds( + seconds, + model=model_name, + model_id=model_id, + call_type=MediaCallType.ASR, + ) + + +def record_asr_usage( + result: Union[str, ASRResult], + *, + model_name: str = "", + model_id: str = "", +) -> None: + """Record one transcription from a provider result. + + A non-verbose call returns a bare string with no timing information, so it + is recorded as unmeasured rather than silently skipped. + """ + raw_response = None + segments = None + if isinstance(result, ASRResult): + raw_response = result.raw_response + segments = result.segments + record_asr_seconds( + resolve_asr_seconds(raw_response, segments), + model_name=model_name, + model_id=model_id, + ) diff --git a/src/xagent/core/tools/core/audio_tool.py b/src/xagent/core/tools/core/audio_tool.py index 168b92e38e..94c331fe2a 100644 --- a/src/xagent/core/tools/core/audio_tool.py +++ b/src/xagent/core/tools/core/audio_tool.py @@ -17,6 +17,8 @@ from ...file_ref import build_workspace_file_ref from ...model.asr.base import ASRResult, BaseASR +from ...model.asr.usage import record_asr_seconds, resolve_asr_seconds +from ...model.chat.token_context import MediaCallType from ...model.tts.base import BaseTTS, TTSResult from ...workspace import TaskWorkspace from .audio_tool_descriptions import ( @@ -27,6 +29,7 @@ SYNTHESIZE_SPEECH_JSON_DESCRIPTION, TRANSCRIBE_AUDIO_DESCRIPTION, ) +from .media_usage import record_media_usage, resolve_billing_model logger = logging.getLogger(__name__) @@ -572,6 +575,33 @@ def _get_tts_model_id(self, tts_model: BaseTTS) -> str: return model_id return "default" + @staticmethod + def _resolve_billing_model( + models: Dict[str, Any], model: Any, model_id: Optional[str] + ) -> str: + """Identify the model that actually served a call, for usage records. + + Resolves by configured key, then by object identity (the batch path + hands back a different instance than the registry dict holds), then + delegates to the shared :func:`resolve_billing_model` for the + provider-attribute fallback and the placeholder filter — so the + Xinference default model, which exposes no ``model_name``, is never + billed under the literal string ``"default"``. + """ + if model_id and model_id in models: + return model_id + if model is not None: + for configured_id, configured_model in models.items(): + if configured_model is model: + return configured_id + # The provider class name is the last identity that still attributes + # cost to something real. The shared helper's own fallback is the + # literal "default", which the metering invariants forbid as a billing + # identity — the Xinference default model reaches exactly that case, + # exposing no model_name and never matching by identity above. + fallback = type(model).__name__ if model is not None else "default" + return resolve_billing_model(None, model, fallback=fallback) + def _get_provider_tts_model( self, *, @@ -713,14 +743,15 @@ async def transcribe_audio( ) # Determine the actual model used - actual_model_id = ( - model_id if model_id and model_id in self._asr_models else "default" + actual_model_id = self._resolve_billing_model( + self._asr_models, asr_model, model_id ) # Handle different result types text = None raw_segments = None language_detected = None + raw_response: dict[str, Any] = {} if isinstance(result, str): text = result @@ -741,6 +772,13 @@ async def transcribe_audio( else None ) language_detected = result.language + if isinstance(result.raw_response, dict): + raw_response = result.raw_response + + # Shared with /speech/transcribe and the Telegram channel so the + # duration rule (provider's own total first, segment timings only + # as a fallback) has exactly one implementation. + audio_seconds = resolve_asr_seconds(raw_response, raw_segments) segment_view = "raw" if verbose else "processed" segments = raw_segments @@ -754,6 +792,23 @@ async def transcribe_audio( ) segments_merged = len(segments) < len(raw_segments) + # Recorded only after every step that can still fail, so a call that + # errors out is not billed. Routed through the shared ASR recorder + # so every entry point meters identically. + record_asr_seconds( + audio_seconds, + # model_id is deliberately left unset. This tool is handed + # name-keyed registries (model_service keys _asr_models by + # model_name), so the only identity in scope here is a name — + # writing it into model_id would persist a name under an id + # field and still not match the real DB model_id that + # /speech/transcribe records. The aggregator groups on + # `model_id or model`, so leaving it empty lets these rows key + # on the name consistently instead of inventing a third + # identity that can never be reconciled. + model_name=str(actual_model_id), + ) + # Save transcription to JSON file if workspace is available file_id: Optional[str] = None file_ref: Optional[dict[str, Any]] = None @@ -916,8 +971,19 @@ async def synthesize_speech( ) # Determine the actual model used - actual_model_id = ( - model_id if model_id and model_id in self._tts_models else "default" + actual_model_id = self._resolve_billing_model( + self._tts_models, tts_model, model_id + ) + + # Meter TTS by input characters (how providers like ElevenLabs + # bill). Recorded before the result is unpacked below: the provider + # call already happened and is billable regardless of what fails + # afterwards. model_id stays unset — see synthesize/ASR above, the + # registries here are name-keyed so only a name is ever in scope. + record_media_usage( + MediaCallType.TTS, + len(text or ""), + model=str(actual_model_id), ) audio_data: Optional[bytes] = None @@ -1873,6 +1939,20 @@ async def _synthesize_single_segment( **kwargs, ) + # Meter TTS by input characters, matching synthesize_speech so batch + # synthesis is metered the same way as single-shot synthesis. + # _resolve_billing_model rather than _get_tts_model_id: the latter + # bottoms out at the literal "default" with no model_name fallback, + # so an implicitly-resolved model would bill a phantom name here. + batch_model_id = self._resolve_billing_model( + self._tts_models, tts_model, None + ) + record_media_usage( + MediaCallType.TTS, + len(text or ""), + model=batch_model_id, + ) + # Handle result if isinstance(audio_data, bytes): audio_binary = audio_data diff --git a/src/xagent/core/tools/core/music_tool.py b/src/xagent/core/tools/core/music_tool.py index 1ea0baf387..21383fcf36 100644 --- a/src/xagent/core/tools/core/music_tool.py +++ b/src/xagent/core/tools/core/music_tool.py @@ -8,8 +8,10 @@ from typing import Any, Optional from ...file_ref import build_workspace_file_ref +from ...model.chat.token_context import MediaCallType from ...model.music import BaseMusicModel, MusicResult from ...workspace import TaskWorkspace +from .media_usage import coerce_duration, record_media_seconds, resolve_billing_model logger = logging.getLogger(__name__) @@ -137,6 +139,36 @@ async def generate_music( force_instrumental=force_instrumental, output_format=output_format, ) + # Billing policy for the media tools: a provider call that has + # already happened is billable regardless of what fails afterwards, + # so usage is recorded before the response is validated. Metering + # after the checks below would silently drop every call that + # succeeded at the HTTP level but came back empty or malformed — + # which the provider still charges for. + # + # Music is duration-billed: always meter in seconds, even when the + # length was auto-selected and no duration came back. `result` is + # not yet known to be a MusicResult at this point, so read through + # getattr rather than attribute access. + raw_response = getattr(result, "raw_response", None) + reported_length = ( + raw_response.get("music_length_seconds") + if isinstance(raw_response, dict) + else None + ) + seconds = coerce_duration(reported_length) or coerce_duration( + music_length_seconds + ) + record_media_seconds( + seconds, + # Never str(None): _configured_model_id returns Optional[str] + # and an inactive shared default resolves to None, which would + # bill against a phantom model literally named "None". + model=resolve_billing_model(configured_model_id, model), + model_id=configured_model_id or "", + call_type=MediaCallType.MUSIC, + ) + if not isinstance(result, MusicResult): raise RuntimeError(f"Unexpected music response: {type(result)}") if not result.audio: diff --git a/src/xagent/core/tools/core/sound_effect_tool.py b/src/xagent/core/tools/core/sound_effect_tool.py index d404551007..7109655124 100644 --- a/src/xagent/core/tools/core/sound_effect_tool.py +++ b/src/xagent/core/tools/core/sound_effect_tool.py @@ -8,8 +8,10 @@ from typing import Any, Optional from ...file_ref import build_workspace_file_ref +from ...model.chat.token_context import MediaCallType from ...model.sound_effect import BaseSoundEffectModel, SoundEffectResult from ...workspace import TaskWorkspace +from .media_usage import coerce_duration, record_media_seconds, resolve_billing_model logger = logging.getLogger(__name__) @@ -151,6 +153,33 @@ async def generate_sound_effect( loop=loop, output_format=output_format, ) + # Billing policy for the media tools: a provider call that has + # already happened is billable regardless of what fails afterwards, + # so usage is recorded before the response is validated — a call + # that returned empty or malformed audio was still charged for. + # + # ElevenLabs prices sound effects by duration, so seconds is the + # only meaningful unit here — an auto-length call with no reported + # duration records 0 seconds rather than switching to characters, + # which would be wrong in kind, not just in magnitude. `result` is + # not yet known to be a SoundEffectResult, so read via getattr. + raw_response = getattr(result, "raw_response", None) + reported_duration = ( + raw_response.get("duration_seconds") + if isinstance(raw_response, dict) + else None + ) + seconds = coerce_duration(reported_duration) or coerce_duration( + duration_seconds + ) + record_media_seconds( + seconds, + # Never str(None) — see resolve_billing_model. + model=resolve_billing_model(configured_model_id, model), + model_id=configured_model_id or "", + call_type=MediaCallType.SOUND_EFFECT, + ) + if not isinstance(result, SoundEffectResult): raise RuntimeError(f"Unexpected sound effect response: {type(result)}") if not result.audio: diff --git a/src/xagent/core/tools/core/video_tool.py b/src/xagent/core/tools/core/video_tool.py index b995d74771..0c69e66297 100644 --- a/src/xagent/core/tools/core/video_tool.py +++ b/src/xagent/core/tools/core/video_tool.py @@ -22,10 +22,12 @@ from pydantic import Field from ...file_ref import build_workspace_file_ref, guess_mime_type, parse_file_id_ref +from ...model.chat.token_context import MediaCallType from ...model.video.ark import ArkVideoModel from ...model.video.base import BaseVideoModel from ...model.video.xinference import XinferenceVideoModel from ...workspace import TaskWorkspace +from .media_usage import coerce_duration, record_media_seconds logger = logging.getLogger(__name__) @@ -721,6 +723,24 @@ async def generate_video( result = await video_model.generate_video(**generate_params) + # Video is duration-billed, so always meter in seconds — never + # switch units when the provider omits a duration (async tasks + # started with wait_for_result=False have none yet). + # + # The provider reports one duration but generates n videos and + # bills for all of them, so the billable total is duration * n. + # Ark rejects n>1; Xinference does not, and previously only the + # first video's duration was recorded. + per_video_seconds = coerce_duration(result.get("duration")) + billable_count = max(1, int(n or 1)) + record_media_seconds( + per_video_seconds * billable_count + if per_video_seconds is not None + else None, + model=str(actual_model_id), + call_type=MediaCallType.VIDEO, + ) + video_url = result.get("video_url") video_path = None video_file_id: Optional[str] = None diff --git a/src/xagent/web/api/model.py b/src/xagent/web/api/model.py index 77a162503a..2d8082256e 100644 --- a/src/xagent/web/api/model.py +++ b/src/xagent/web/api/model.py @@ -980,6 +980,9 @@ async def transcribe_speech_input( """Transcribe a short UI voice input clip with an accessible ASR model.""" from xagent.core.model.asr.adapter import get_asr_model_instance + from xagent.core.model.asr.usage import record_asr_usage + + from ..tracking.standalone_usage import usage_scope try: audio_bytes = await _read_transcribe_upload_with_size_limit(file) @@ -991,14 +994,36 @@ async def transcribe_speech_input( db_model = _resolve_asr_model_for_transcription(db, user, model_id) asr_model = get_asr_model_instance(db_model) try: - result = await asyncio.wait_for( - asr_model.transcribe( - audio_bytes, - language=language, - format=_format_from_upload(file), - ), - timeout=180.0, - ) + # usage_scope binds a TokenUsage and reports it to the quota hook on + # exit. Recording alone is not enough here: this endpoint has no + # TaskTracker, so without a bound context the usage would land in a + # throwaway object and the transcription would bill nothing. + with usage_scope(int(user.id) if user and user.id is not None else None): + result = await asyncio.wait_for( + asr_model.transcribe( + audio_bytes, + language=language, + format=_format_from_upload(file), + # Without this the provider returns a bare string carrying + # no timings, so the usage record below would be a + # 0-second — unbillable — entry on every call. The extra + # fields are discarded; only `text` is returned to the UI. + verbose=True, + ), + timeout=180.0, + ) + # model_id is deliberately left unset even though the real DB id is + # in scope here. The aggregator groups on `model_id or model`, and + # the other two ASR entry points (audio_tool, the Telegram channel) + # only ever see name-keyed model registries — model_service keys + # them by model_name — so no DB id is available there. Setting it + # on this path alone would split one physical model into two + # billing rows that can never be reconciled once written. Name is + # the one identity all three share. + record_asr_usage( + result, + model_name=str(db_model.model_name), + ) except asyncio.TimeoutError as exc: raise HTTPException(status_code=504, detail="Transcription timed out") from exc except Exception as exc: diff --git a/src/xagent/web/channels/telegram/bot.py b/src/xagent/web/channels/telegram/bot.py index 55b8f335dc..1599e3c3e9 100644 --- a/src/xagent/web/channels/telegram/bot.py +++ b/src/xagent/web/channels/telegram/bot.py @@ -1368,6 +1368,9 @@ async def _transcribe_uploaded_voice_files( uploaded_info: list[dict[str, Any]], asr_model: Any, ) -> dict[str, str]: + from ....core.model.asr.usage import record_asr_usage + from ....core.tools.core.media_usage import resolve_billing_model + uploaded_by_source_id = { str(info.get("telegram_file_id")): info for info in uploaded_info @@ -1387,9 +1390,23 @@ async def _transcribe_uploaded_voice_files( asr_model.transcribe( audio=str(file_info["path"]), format=self._audio_format_from_file_info(file_info), + # Without this the provider returns a bare string + # carrying no timings, making every voice message a + # 0-second — unbillable — usage record. Only `text` is + # used below; the extra fields are discarded. + verbose=True, ), timeout=self.voice_transcription_timeout_seconds, ) + # Telegram calls the ASR provider directly rather than going + # through audio_tool, so meter here too. Identity goes through + # the shared resolver so a provider exposing no model_name is + # never billed under a placeholder, and all three ASR entry + # points agree on the name the aggregator groups by. + record_asr_usage( + result, + model_name=resolve_billing_model(None, asr_model), + ) except asyncio.TimeoutError as exc: raise TelegramVoiceTranscriptionError( "Telegram voice transcription timed out" diff --git a/src/xagent/web/services/quota_hooks.py b/src/xagent/web/services/quota_hooks.py index a215b5c7e5..f4eb8f3e4d 100644 --- a/src/xagent/web/services/quota_hooks.py +++ b/src/xagent/web/services/quota_hooks.py @@ -23,9 +23,30 @@ # callbacks may rely on that affinity and must remain non-blocking. _run_gate_hook: Callable[[Any, Any], str | Mapping[str, Any] | None] | None = None # (db, user_id, delta_details, delta_actions) -> None; best-effort post-run -# metering. delta_details is this turn's per-model token breakdown (list of -# {"type","tokens","model"}) for cost-based credits; delta_actions counts tool -# calls (one billable action per tool invocation). +# metering. delta_details is this turn's per-model usage breakdown for +# cost-based credits; delta_actions counts tool calls (one billable action per +# tool invocation). Entry shapes in delta_details: +# - LLM tokens: {"type":"input"|"output", "tokens", "model", "model_id", +# "call_type", ...cache fields} +# - Non-LLM media (image/video/tts/asr/embedding/rerank/...): +# {"type":"media", "unit":"images"|"seconds"|"characters"|"texts"| +# "requests", "quantity", "provider_tokens", "provider_input_tokens", +# "provider_output_tokens", "tokens_estimated", "model", "model_id", +# "call_type", "resolution"} +# The app layer should price media entries by their "unit"/"quantity". The unit +# is stable for a given (model, call_type) — a duration-billed modality always +# reports "seconds", recording quantity=0 when the provider gave no duration, +# rather than switching units. "requests" always carries quantity=1; a batch of +# N embedded texts reports unit="texts" with quantity=N and calls=1. +# For image models whose price varies by resolution, "resolution" ("1K"/"2K"/ +# "4K" or "1024x1024") keys a per-(model, resolution) price table. Providers +# that report real image tokens (Gemini, OpenAI gpt-image) also fill +# "provider_tokens" so a token-based price ($/1M tokens) can take precedence +# over the resolution table — but only when "tokens_estimated" is False; +# embedding/rerank counts are local heuristics and must not be priced as +# measured tokens. Media token counts are deliberately NOT under the "tokens" +# key, so a consumer summing "tokens" across entries cannot double-count them. +# Unknown entry types must be ignored, not summed as tokens. # # TRANSACTION CONTRACT: the hook is invoked from TaskTracker.complete_tracking # only after the run/runner-fenced token-usage update commits. The hook owns @@ -41,7 +62,8 @@ # in-flight run's live-so-far usage would push the team over a run-gated quota, # else None. Polled per step (each LLM reply / tool call) during a run so a # single long/expensive run is stopped mid-flight instead of only being metered -# at completion. +# at completion. delta_details carries the same entry shapes documented on the +# metering hook above (LLM token entries plus type:"media" entries). # # CONTRACT: invoked SYNCHRONOUSLY on the event loop once per step. It MUST NOT # block (no synchronous network/DB round-trips per call) — blocking work stalls @@ -90,6 +112,18 @@ def check_run_gate(db: Any, user_id: Any) -> str | Mapping[str, Any] | None: return _run_gate_hook(db, user_id) +def has_usage_record_hook() -> bool: + """Whether a usage-record hook is installed. + + Lets a caller skip the work of preparing a report — notably checking out a + DB session — when :func:`record_usage` would be a no-op anyway. In the + stock open-source configuration no hook is registered, so every ingest and + transcription would otherwise pay a pool checkout, transaction and close + for nothing. + """ + return _usage_record_hook is not None + + def record_usage( db: Any, user_id: Any, delta_details: list, delta_actions: int ) -> None: diff --git a/src/xagent/web/tracking/standalone_usage.py b/src/xagent/web/tracking/standalone_usage.py new file mode 100644 index 0000000000..a530101285 --- /dev/null +++ b/src/xagent/web/tracking/standalone_usage.py @@ -0,0 +1,142 @@ +"""Usage metering for work that is not a tracked agent task. + +``TaskTracker`` binds a ``TokenUsage`` for chat/agent runs and reports the +delta to the quota hook when the run completes. Everything else — KB ingestion +over HTTP or Celery, ``/speech/transcribe``, Telegram voice — records usage +with no context bound, so ``get_token_usage()`` lazily creates a throwaway +object that nothing ever reads. The calls succeed, the provider bills, and the +usage silently evaporates. + +This module is the equivalent sink for those paths: bind a context around the +work, then report whatever was recorded. + +Why a shared helper rather than a `TokenContextManager` at each call site: the +quota hook has a transaction contract that is easy to violate by accident. It +must not be handed a caller's request Session (it manages its own durability +and must not commit or leave writes pending on someone else's session), so the +report step opens and disposes a short-lived compatibility Session of its own. +Reproducing that at four call sites would mean four chances to get it wrong. +""" + +from __future__ import annotations + +import functools +import logging +from contextlib import contextmanager +from contextvars import copy_context +from typing import Any, Callable, Iterator, Optional, TypeVar + +from ...core.model.chat.token_context import TokenUsage, set_token_usage + +logger = logging.getLogger(__name__) + +T = TypeVar("T") + + +def _report(user_id: Optional[int], usage: TokenUsage) -> None: + """Hand this unit of work's usage to the quota hook, best-effort.""" + details = [dict(item) for item in usage.details if isinstance(item, dict)] + if not details: + return + try: + from ..models.database import get_session_local + from ..services.quota_hooks import has_usage_record_hook, record_usage + + # Check the hook before checking out a session: with no hook installed + # (the stock configuration) record_usage is a guaranteed no-op, and a + # pool checkout + transaction + close per ingest/transcription is pure + # overhead. + if not has_usage_record_hook(): + return + + db_session = get_session_local()() + try: + # delta_actions=0: these paths make provider calls, not agent tool + # calls, and tool invocations are what that counter bills for. + record_usage(db_session, user_id, details, 0) + finally: + # The hook owns its own durability and must not leave work pending + # on this compatibility Session. + if db_session.in_transaction(): + db_session.rollback() + db_session.close() + except Exception as e: # noqa: BLE001 + # Metering must never break the work it is measuring. + logger.warning("Standalone usage recording failed: %s", e) + + +@contextmanager +def usage_scope(user_id: Optional[int]) -> Iterator[TokenUsage]: + """Bind a usage context for one unit of non-task work and report it after. + + Usage is reported even when the body raises: a provider call that already + happened is billable regardless of what fails afterwards. + + Note the body must not cross a thread boundary that drops contextvars — + see :func:`bind_usage_to_thread` for `run_in_executor`-style hops. + """ + from ...core.model.chat.token_context import token_context + + usage = TokenUsage() + # Restore whatever was bound before (usually None) so a nested scope cannot + # leak its usage object into the caller's. + previous = token_context.get(None) + set_token_usage(usage) + try: + yield usage + finally: + try: + # Only restore if we are still the bound context. A TaskTracker + # started *inside* this scope also calls set_token_usage, and + # clobbering that would silently detach it from its own run. + # Unreachable at today's call sites, but cheap to get right before + # a sixth one appears. + if token_context.get(None) is usage: + set_token_usage(previous) # type: ignore[arg-type] + else: + logger.debug( + "usage_scope exiting with a different context bound; " + "leaving it in place rather than detaching its owner" + ) + except Exception: # noqa: BLE001 + pass + _report(user_id, usage) + + +def bind_usage_to_thread(fn: Callable[..., T]) -> Callable[..., T]: + """Wrap a callable so it records into the *calling* thread's usage. + + ``loop.run_in_executor(None, fn)`` and bare ``ThreadPoolExecutor`` do not + propagate contextvars, so without this the worker records into a fresh + ``TokenUsage`` that is discarded on return. ``asyncio.to_thread`` already + copies the context and does not need this. + + Captures at wrap time, on the calling thread — wrapping must therefore + happen before the hop, not inside the worker. + + The binding runs inside a copied context (``copy_context()`` + + ``ctx.run``), the same idiom used at + ``core/tools/adapters/vibe/file_ingestion_tool.py``. That confines the + contextvar write to this one call instead of mutating the worker thread: + ``run_in_executor(None, ...)`` uses the loop's long-lived default executor, + so a plain ``set_token_usage`` would outlive the job and leak the caller's + ``TokenUsage`` into the next unrelated task that reuses the thread — + cross-tenant usage misattribution, plus a strong reference pinned for the + thread's lifetime. A copied context also needs no ``finally`` restore: + there is nothing in the caller's context left to restore. + """ + from ...core.model.chat.token_context import get_token_usage + + caller_usage = get_token_usage() + + @functools.wraps(fn) + def _bound(*args: Any, **kwargs: Any) -> T: + def _run() -> T: + set_token_usage(caller_usage) + return fn(*args, **kwargs) + + # Fresh copy per invocation: a wrapped callable may be submitted more + # than once, and each run must get its own isolated context. + return copy_context().run(_run) + + return _bound diff --git a/tests/core/tools/core/test_media_tool_billing_policy.py b/tests/core/tools/core/test_media_tool_billing_policy.py new file mode 100644 index 0000000000..8d8804aeac --- /dev/null +++ b/tests/core/tools/core/test_media_tool_billing_policy.py @@ -0,0 +1,156 @@ +"""The media tools bill a provider call that happened, even if it then fails. + +Music, sound effect, TTS and ASR are otherwise-symmetric tools that had three +different implicit "is this call billable" rules. The policy is now uniform: +usage is recorded as soon as the provider call returns, before the response is +validated, because a call that succeeded at the HTTP level but came back empty +or malformed was still charged for. +""" + +from typing import Any, Optional + +import pytest + +from xagent.core.model.chat.token_context import TokenContextManager +from xagent.core.model.music.base import MusicResult +from xagent.core.tools.core.music_tool import MusicToolCore +from xagent.core.tools.core.sound_effect_tool import SoundEffectToolCore + + +class _EmptyMusicModel: + """Provider that returns a well-formed-but-empty result. Still billed.""" + + model_name = "music-a" + + async def generate_music(self, **kwargs: Any) -> MusicResult: + _ = kwargs + return MusicResult(audio=b"", format="mp3", raw_response={}) + + +class _GarbageMusicModel: + """Provider that returns the wrong type entirely. Still billed.""" + + model_name = "music-b" + + async def generate_music(self, **kwargs: Any) -> Any: + _ = kwargs + return {"not": "a MusicResult"} + + +def _media_entries(manager: TokenContextManager) -> list[dict]: + return [d for d in manager.get_usage().details if d.get("type") == "media"] + + +@pytest.mark.asyncio +async def test_music_bills_empty_audio_response() -> None: + tool = MusicToolCore(models={"music-a": _EmptyMusicModel()}) + + with TokenContextManager() as manager: + result = await tool.generate_music(prompt="p", music_length_seconds=30) + entries = _media_entries(manager) + + # The tool still reports failure to the caller... + assert result["success"] is False + # ...but the provider call happened and was billed, so it is metered. + assert len(entries) == 1 + assert entries[0]["unit"] == "seconds" + assert entries[0]["quantity"] == 30.0 + assert entries[0]["call_type"] == "music" + + +@pytest.mark.asyncio +async def test_music_bills_malformed_response() -> None: + tool = MusicToolCore(models={"music-b": _GarbageMusicModel()}) + + with TokenContextManager() as manager: + result = await tool.generate_music(prompt="p", music_length_seconds=12) + entries = _media_entries(manager) + + assert result["success"] is False + assert len(entries) == 1 + # Unit stays seconds even though nothing usable came back: a (model, unit) + # price table breaks if the unit varies with response completeness. + assert entries[0]["unit"] == "seconds" + assert entries[0]["quantity"] == 12.0 + + +class _EmptySoundEffectModel: + model_name = "sfx-a" + + async def generate_sound_effect(self, **kwargs: Any) -> Any: + _ = kwargs + return {"not": "a SoundEffectResult"} + + +@pytest.mark.asyncio +async def test_sound_effect_bills_malformed_response() -> None: + tool = SoundEffectToolCore(models={"sfx-a": _EmptySoundEffectModel()}) + + with TokenContextManager() as manager: + result = await tool.generate_sound_effect(text="p", duration_seconds=4) + entries = _media_entries(manager) + + assert result["success"] is False + assert len(entries) == 1 + assert entries[0]["unit"] == "seconds" + assert entries[0]["quantity"] == 4.0 + assert entries[0]["call_type"] == "sound_effect" + + +class _RecordingASR: + """Captures the verbose flag the caller passed.""" + + model_name = "asr-a" + + def __init__(self) -> None: + self.verbose_seen: Optional[bool] = None + + async def transcribe(self, *args: Any, verbose: bool = False, **kwargs: Any) -> Any: + _ = args, kwargs + self.verbose_seen = verbose + if not verbose: + return "bare string, no timings" + from xagent.core.model.asr.base import ASRResult, ASRSegment + + return ASRResult( + text="hi", + segments=[ASRSegment(text="hi", start=0.0, end=2.5, confidence=1.0)], + language="eng", + ) + + +@pytest.mark.asyncio +async def test_asr_usage_is_unmeasured_without_verbose() -> None: + # Guards the reason /speech/transcribe and the Telegram channel now pass + # verbose=True: without it the provider returns a bare string with no + # timings, and every call is metered as an unbillable 0 seconds. + from xagent.core.model.asr.usage import record_asr_usage + + with TokenContextManager() as manager: + record_asr_usage("bare string, no timings", model_name="asr-a") + entries = _media_entries(manager) + + assert entries[0]["unit"] == "seconds" + assert entries[0]["quantity"] == 0.0 + + +@pytest.mark.asyncio +async def test_asr_usage_meters_duration_with_verbose() -> None: + from xagent.core.model.asr.base import ASRResult, ASRSegment + from xagent.core.model.asr.usage import record_asr_usage + + result = ASRResult( + text="hi", + segments=[ASRSegment(text="hi", start=0.0, end=2.5, confidence=1.0)], + language="eng", + ) + + with TokenContextManager() as manager: + record_asr_usage(result, model_name="asr-a") + entries = _media_entries(manager) + + assert entries[0]["unit"] == "seconds" + assert entries[0]["quantity"] == 2.5 + # Name-keyed identity across all three ASR entry points; see PR body. + assert entries[0]["model"] == "asr-a" + assert entries[0]["model_id"] == "" diff --git a/tests/core/tools/test_audio_tool.py b/tests/core/tools/test_audio_tool.py index 80200ec0d5..4716980338 100644 --- a/tests/core/tools/test_audio_tool.py +++ b/tests/core/tools/test_audio_tool.py @@ -949,6 +949,71 @@ async def test_synthesize_speech_json_merges_default_and_segment_options() -> No ] +async def test_synthesize_speech_json_records_media_usage_per_segment() -> None: + from xagent.core.model.chat.token_context import TokenContextManager + + tts = FakeTTS() + tool = AudioToolCore(tts_models={"fake": tts}) + + with TokenContextManager() as manager: + result = await tool.synthesize_speech_json( + json_data={ + "segments": [ + {"text": "First line"}, + {"text": "Second longer line"}, + ], + "default_voice": "voice-1", + }, + model_id="fake", + ) + usage = manager.get_usage() + + assert result["success"] is True + media_entries = [d for d in usage.details if d.get("type") == "media"] + # One TTS media entry per synthesized segment, metered by input characters. + assert usage.media_calls == 2 + assert [(d["unit"], d["quantity"], d["call_type"]) for d in media_entries] == [ + ("characters", len("First line"), "tts"), + ("characters", len("Second longer line"), "tts"), + ] + + +async def test_tts_usage_records_real_model_not_default() -> None: + """Omitting model_id is the documented common case, so the usage entry must + still name the model — billing the literal string "default" would destroy + per-model cost attribution.""" + from xagent.core.model.chat.token_context import TokenContextManager + + tool = AudioToolCore(tts_models={"tts-a": FakeTTS()}) + + with TokenContextManager() as manager: + await tool.synthesize_speech(text="hello") + entries = [d for d in manager.get_usage().details if d.get("type") == "media"] + + assert len(entries) == 1 + assert entries[0]["model"] == "tts-a" + assert entries[0]["model"] != "default" + + +async def test_asr_usage_meters_seconds_from_segments() -> None: + """ASR is duration-billed, so the unit is seconds and the quantity comes + from the transcribed audio's timing.""" + from xagent.core.model.chat.token_context import TokenContextManager + + tool = AudioToolCore(asr_models={"asr-a": FakeASR()}) + + with TokenContextManager() as manager: + result = await tool.transcribe_audio(audio_file_path="x.wav") + entries = [d for d in manager.get_usage().details if d.get("type") == "media"] + + assert result["success"] is True + assert len(entries) == 1 + assert entries[0]["unit"] == "seconds" + assert entries[0]["call_type"] == "asr" + assert entries[0]["quantity"] > 0 + assert entries[0]["model"] == "asr-a" + + async def test_synthesize_speech_json_rejects_non_object_json() -> None: tool = AudioToolCore(tts_models={"fake": FakeTTS()}) @@ -963,3 +1028,59 @@ async def test_synthesize_speech_json_rejects_non_object_json() -> None: "failed": 0, "errors": ["JSON data must be an object"], } + + +class _NamelessASR(FakeASR): + """The Xinference default-model shape: no ``model_name`` attribute. + + Its default-getter builds a separate instance from the one in the registry + dict, so identity-by-``is`` never matches either. + """ + + +async def test_asr_default_model_is_not_billed_as_placeholder() -> None: + # Previously fell through to the literal string "default" — the exact + # placeholder the metering invariants forbid as a billing identity. + from xagent.core.model.chat.token_context import TokenContextManager + + tool = AudioToolCore( + asr_models={"asr-a": _NamelessASR()}, + default_asr_model=_NamelessASR(), + ) + + with TokenContextManager() as manager: + await tool.transcribe_audio(audio_file_path="x.wav") + entries = [d for d in manager.get_usage().details if d.get("type") == "media"] + + assert len(entries) == 1 + assert entries[0]["model"] not in {"default", "none", "null", ""} + + +async def test_asr_usage_leaves_model_id_unset() -> None: + # The tool only ever sees name-keyed registries, so writing the name into + # model_id would persist a name under an id field and still not match the + # real DB model_id /speech/transcribe records. The aggregator groups on + # `model_id or model`, so all three ASR entry points must key on the name. + from xagent.core.model.chat.token_context import TokenContextManager + + tool = AudioToolCore(asr_models={"asr-a": FakeASR()}) + + with TokenContextManager() as manager: + await tool.transcribe_audio(audio_file_path="x.wav") + entries = [d for d in manager.get_usage().details if d.get("type") == "media"] + + assert entries[0]["model"] == "asr-a" + assert entries[0]["model_id"] == "" + + +async def test_tts_usage_leaves_model_id_unset() -> None: + from xagent.core.model.chat.token_context import TokenContextManager + + tool = AudioToolCore(tts_models={"tts-a": FakeTTS()}) + + with TokenContextManager() as manager: + await tool.synthesize_speech(text="hello") + entries = [d for d in manager.get_usage().details if d.get("type") == "media"] + + assert entries[0]["model"] == "tts-a" + assert entries[0]["model_id"] == "" diff --git a/tests/web/test_connector_runtime_entrypoints_e2e.py b/tests/web/test_connector_runtime_entrypoints_e2e.py index 84f4cd3145..c2b85feb2d 100644 --- a/tests/web/test_connector_runtime_entrypoints_e2e.py +++ b/tests/web/test_connector_runtime_entrypoints_e2e.py @@ -1101,7 +1101,14 @@ class _FakeASR: def __init__(self) -> None: self.closed = False - async def transcribe(self, *, audio: str, format: str | None = None) -> str: + async def transcribe( + self, + *, + audio: str, + format: str | None = None, + verbose: bool = False, + **kwargs: object, + ) -> str: assert audio == "/workspace/input/voice.oga" assert format == "ogg" return "今晚有世界杯比赛吗？" diff --git a/tests/web/test_telegram_message_queue.py b/tests/web/test_telegram_message_queue.py index 562e65c404..543a8a79e1 100644 --- a/tests/web/test_telegram_message_queue.py +++ b/tests/web/test_telegram_message_queue.py @@ -166,8 +166,15 @@ class FakeASR: def __init__(self) -> None: self.calls: list[dict[str, str | None]] = [] - async def transcribe(self, *, audio: str, format: str | None = None) -> str: - self.calls.append({"audio": audio, "format": format}) + async def transcribe( + self, + *, + audio: str, + format: str | None = None, + verbose: bool = False, + **kwargs: object, + ) -> str: + self.calls.append({"audio": audio, "format": format, "verbose": verbose}) return "今晚有世界杯比赛吗？" asr = FakeASR() @@ -189,8 +196,14 @@ async def transcribe(self, *, audio: str, format: str | None = None) -> str: ) assert transcripts == {"voice-file-id": "今晚有世界杯比赛吗？"} + # verbose=True is required: without it the provider returns a bare string + # with no timings and the ASR usage record is an unbillable 0 seconds. assert asr.calls == [ - {"audio": "/workspace/input/voice-file-id.oga", "format": "ogg"} + { + "audio": "/workspace/input/voice-file-id.oga", + "format": "ogg", + "verbose": True, + } ] assert uploaded_info[0]["file_id"] == "workspace-file-id" @@ -201,7 +214,12 @@ async def test_transcribe_uploaded_voice_files_extracts_result_text() -> None: class ResultASR: async def transcribe( - self, *, audio: str, format: str | None = None + self, + *, + audio: str, + format: str | None = None, + verbose: bool = False, + **kwargs: object, ) -> SimpleNamespace: return SimpleNamespace(text="今晚有世界杯比赛吗？") @@ -242,7 +260,14 @@ async def test_transcribe_uploaded_voice_files_rejects_empty_result() -> None: bot = make_bot() class EmptyASR: - async def transcribe(self, *, audio: str, format: str | None = None) -> str: + async def transcribe( + self, + *, + audio: str, + format: str | None = None, + verbose: bool = False, + **kwargs: object, + ) -> str: return " " with pytest.raises( From 910b8f11a8b8e2c7ef9b865b23e92f27353aa5ce Mon Sep 17 00:00:00 2001 From: OliverBryant <2713999266@qq.com> Date: Tue, 18 Aug 2026 16:24:49 +0800 Subject: [PATCH 2/6] fix: widen the /speech/transcribe test fake to the real ASR signature MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `FakeASR.transcribe` in test_model_api took only (audio, language, format), so the endpoint's new `verbose=True` raised TypeError and the request returned 500. The real `BaseASR.transcribe` takes `verbose` and `**kwargs`; the double was narrower than the interface it stood in for, which is the same reason the missing `verbose=True` went unnoticed in the first place. The fake now mirrors the real signature and records the flag, and the call assertion checks `verbose=True` is actually passed — so a revert of the endpoint change fails here instead of silently recording 0-second, unbillable ASR usage. --- tests/web/test_model_api.py | 24 +++++++++++++++++++++--- 1 file changed, 21 insertions(+), 3 deletions(-) diff --git a/tests/web/test_model_api.py b/tests/web/test_model_api.py index 129d9b1097..d88a505a96 100644 --- a/tests/web/test_model_api.py +++ b/tests/web/test_model_api.py @@ -1010,9 +1010,19 @@ def __init__(self, model_name: str): self.calls = [] self.closed = False - async def transcribe(self, audio, language=None, format=None): + async def transcribe( + self, audio, language=None, format=None, verbose=False, **kwargs + ): + # Mirrors the real BaseASR signature, which takes verbose and + # **kwargs. A double narrower than the interface it stands in + # for is how the missing verbose=True went unnoticed. self.calls.append( - {"audio": audio, "language": language, "format": format} + { + "audio": audio, + "language": language, + "format": format, + "verbose": verbose, + } ) return f"text from {self.model_name}" @@ -1042,8 +1052,16 @@ def create_fake_asr(db_model): assert data["text"] == "text from asr-second" assert data["model_id"] == "test-asr-second" assert len(created_models) == 1 + # verbose=True is required: without it the provider returns a bare + # string with no timings and the ASR usage record is an unbillable + # 0 seconds. assert created_models[0].calls == [ - {"audio": b"audio-bytes", "language": "en", "format": "webm"} + { + "audio": b"audio-bytes", + "language": "en", + "format": "webm", + "verbose": True, + } ] assert created_models[0].closed is True From 2e24aa0fe7e90d287864479e821fe8e8ffcd3105 Mon Sep 17 00:00:00 2001 From: OliverBryant Date: Fri, 21 Aug 2026 12:17:28 +0800 Subject: [PATCH 3/6] test: cover the standalone usage sink MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit standalone_usage.py sits on the quota path and shipped with no tests: the original split lost this file, which exists in #997 but landed in neither this branch nor main. Recovered and migrated to the merged media API (#1527), where the unit is derived from call_type rather than passed in. Two changes beyond the mechanical migration: - Stub has_usage_record_hook. _report gained a short-circuit that returns before checking out a session when no hook is installed, which the #997 tests predate — dropped in verbatim they would pass vacuously, asserting nothing. - Scope the entry-point wiring guard to transcribe_speech_input. The KB ingest and Telegram entry points the original asserted on bind their scopes in a later PR in this series. Added a case for the short-circuit itself, so the no-hook fast path cannot regress into a pool checkout per transcription. --- tests/web/tracking/test_standalone_usage.py | 206 ++++++++++++++++++++ 1 file changed, 206 insertions(+) create mode 100644 tests/web/tracking/test_standalone_usage.py diff --git a/tests/web/tracking/test_standalone_usage.py b/tests/web/tracking/test_standalone_usage.py new file mode 100644 index 0000000000..0111100585 --- /dev/null +++ b/tests/web/tracking/test_standalone_usage.py @@ -0,0 +1,206 @@ +"""The usage sink for work that has no TaskTracker. + +Recording usage is only half of metering: something has to bind a TokenUsage +and hand it to the quota hook. Entry points such as ``/speech/transcribe`` have +no TaskTracker, so without this their recorded usage lands in a throwaway +object. +""" + +import concurrent.futures +from typing import Any, Optional + +import pytest + +from xagent.core.model.chat.token_context import ( + MediaCallType, + add_media_usage, + get_token_usage, +) +from xagent.web.tracking.standalone_usage import bind_usage_to_thread, usage_scope + + +class _FakeSession: + def __init__(self) -> None: + self.rolled_back = False + self.closed = False + + def in_transaction(self) -> bool: + return False + + def rollback(self) -> None: + self.rolled_back = True + + def close(self) -> None: + self.closed = True + + +@pytest.fixture +def captured(monkeypatch: pytest.MonkeyPatch): + """Capture what reaches quota_hooks.record_usage.""" + calls: list[dict[str, Any]] = [] + session = _FakeSession() + + def _record(db: Any, user_id: Any, details: list, actions: int) -> None: + calls.append( + {"db": db, "user_id": user_id, "details": details, "actions": actions} + ) + + from xagent.web.models import database + from xagent.web.services import quota_hooks + + monkeypatch.setattr(quota_hooks, "record_usage", _record) + # _report short-circuits before touching a session when no hook is + # installed, so the stock configuration would otherwise report nothing. + monkeypatch.setattr(quota_hooks, "has_usage_record_hook", lambda: True) + monkeypatch.setattr(database, "get_session_local", lambda: lambda: session) + return calls, session + + +def test_usage_recorded_in_scope_reaches_the_quota_hook(captured) -> None: + calls, _ = captured + with usage_scope(42): + add_media_usage(MediaCallType.ASR, 3, model="asr-1") + + assert len(calls) == 1 + assert calls[0]["user_id"] == 42 + assert calls[0]["details"][0]["call_type"] == "asr" + # The unit is derived from the call type, never passed in. + assert calls[0]["details"][0]["unit"] == "seconds" + assert calls[0]["details"][0]["quantity"] == 3.0 + # These paths make provider calls, not agent tool calls. + assert calls[0]["actions"] == 0 + + +def test_scope_reports_even_when_the_body_raises(captured) -> None: + """A provider call that already happened is billable regardless of what + fails afterwards.""" + calls, _ = captured + with pytest.raises(RuntimeError): + with usage_scope(7): + add_media_usage(MediaCallType.TTS, 12, model="t") + raise RuntimeError("downstream failure") + + assert len(calls) == 1 + assert calls[0]["details"][0]["call_type"] == "tts" + assert calls[0]["details"][0]["unit"] == "characters" + + +def test_no_usage_means_no_hook_call(captured) -> None: + calls, _ = captured + with usage_scope(1): + pass + assert calls == [] + + +def test_compatibility_session_is_disposed(captured) -> None: + """The hook manages its own durability; the session handed to it must be + closed by us and never left holding a transaction.""" + calls, session = captured + with usage_scope(5): + add_media_usage(MediaCallType.ASR, 2, model="e") + + assert session.closed is True + assert calls[0]["db"] is session + + +def test_bind_usage_to_thread_carries_usage_across_the_hop(captured) -> None: + """run_in_executor / ThreadPoolExecutor drop contextvars; the wrapper is + what keeps the worker's records attached to the caller.""" + calls, _ = captured + + with usage_scope(9): + + def work() -> None: + add_media_usage(MediaCallType.ASR, 4, model="e") + + bound = bind_usage_to_thread(work) + with concurrent.futures.ThreadPoolExecutor(max_workers=2) as pool: + pool.submit(bound).result() + + assert len(calls) == 1 + assert calls[0]["details"][0]["quantity"] == 4.0 + + +def test_unbound_thread_hop_loses_usage(captured) -> None: + """Documents why the wrapper exists: the same code without it records + nothing, which is the defect this module fixes.""" + calls, _ = captured + + with usage_scope(9): + + def work() -> None: + add_media_usage(MediaCallType.ASR, 4, model="e") + + with concurrent.futures.ThreadPoolExecutor(max_workers=2) as pool: + pool.submit(work).result() # not bound + + assert calls == [] + + +def test_scope_restores_the_previous_context(captured) -> None: + """A nested scope must not leak its usage object into the outer one.""" + _calls, _ = captured + with usage_scope(1) as outer: + add_media_usage(MediaCallType.TTS, 1, model="a") + with usage_scope(2): + add_media_usage(MediaCallType.TTS, 1, model="b") + # Back on the outer usage, and the inner call did not land here. + assert get_token_usage() is outer + assert len(outer.details) == 1 + + +def test_none_user_id_is_tolerated(captured) -> None: + """record_usage no-ops on a missing user; the scope must not raise.""" + calls, _ = captured + user_id: Optional[int] = None + with usage_scope(user_id): + add_media_usage(MediaCallType.ASR, 1, model="m") + assert calls[0]["user_id"] is None + + +def test_hook_failure_does_not_break_the_work(monkeypatch: pytest.MonkeyPatch) -> None: + """Metering must never break the operation it measures.""" + from xagent.web.models import database + from xagent.web.services import quota_hooks + + def _boom(*_a: Any, **_k: Any) -> None: + raise RuntimeError("quota backend down") + + monkeypatch.setattr(quota_hooks, "record_usage", _boom) + monkeypatch.setattr(quota_hooks, "has_usage_record_hook", lambda: True) + monkeypatch.setattr(database, "get_session_local", lambda: _FakeSession) + + with usage_scope(1): + add_media_usage(MediaCallType.ASR, 1, model="m") + # No exception escaped. + + +def test_no_hook_installed_skips_the_session_checkout( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """The stock configuration has no hook, so _report must not pay for a pool + checkout per transcription.""" + from xagent.web.models import database + from xagent.web.services import quota_hooks + + checkouts: list[int] = [] + + monkeypatch.setattr(quota_hooks, "has_usage_record_hook", lambda: False) + monkeypatch.setattr( + database, "get_session_local", lambda: checkouts.append(1) or _FakeSession + ) + + with usage_scope(1): + add_media_usage(MediaCallType.ASR, 1, model="m") + + assert checkouts == [] + + +def test_patched_entry_points_bind_a_scope() -> None: + """Guard the wiring itself: if a future refactor drops the scope from an + entry point, its usage silently stops being billed again.""" + import inspect + + from xagent.web.api import model + + assert "usage_scope(" in inspect.getsource(model.transcribe_speech_input) From 62c4ca3a0af78d3819074739029a5fe16dc979d9 Mon Sep 17 00:00:00 2001 From: OliverBryant Date: Fri, 21 Aug 2026 15:52:04 +0800 Subject: [PATCH 4/6] fix: address media metering review feedback Carries the provider name and the configured id in the fields they belong in, drops a pricing promise the authoritative contract does not make, and closes the test gaps that let those defects stay green. Identity fields (`model` = provider name, `model_id` = configured id): - music/sound_effect passed the configured id as resolve_billing_model's first argument, which returns it unchanged, so the same id was written into both fields and the provider name was lost. Pass None instead so the resolver falls through to the provider's model_name, with the class name as the last resort rather than the placeholder "default". - video wrote the configured id into `model` and left `model_id` empty. Populate both; the aggregator groups on `model_id or model`, so row identity is unchanged while the name is no longer discarded. Pricing contract: quota_hooks claimed provider-reported image tokens can take precedence over resolution pricing, but add_media_usage explicitly declines to define that rule -- there is no discriminator in the row schema to express it, and the aggregate groups purely by (model, unit, call_type, resolution). Removed the promise and pointed at the authoritative contract and #1461. Tests: - The transcribe wiring guard only asserted that the literal "usage_scope(" appeared in the endpoint source, which would pass with the wrong user, the wrong scope boundary, or no recording at all. Replaced with a behavioural test that drives the endpoint and asserts what reaches the quota hook: owner, unit, quantity, identity, exactly once. - _FakeSession.in_transaction() was hard-coded False, leaving _report's rollback branch unreachable -- deleting the rollback kept the suite green. Made the transaction state configurable and added cases for the success path, the raising-hook path and the inactive case, asserting rollback happens before close. - resolve_asr_seconds had two covered branches out of several. Added tests for each provider total field, the fallback to segment ends, unsorted segments, unusable totals (non-finite, negative, boolean, non-numeric) and the no-timing case. - Added distinct-name/distinct-id regression tests for the music and sound-effect identity fix. Telegram ASR lifecycle (#1582) and async video reconciliation (#1583) are tracked separately; both need a new reporting/settlement mechanism rather than a change to the metering call. --- src/xagent/core/tools/core/music_tool.py | 17 +- .../core/tools/core/sound_effect_tool.py | 10 +- src/xagent/core/tools/core/video_tool.py | 17 +- src/xagent/web/services/quota_hooks.py | 16 +- .../core/test_media_tool_billing_policy.py | 157 ++++++++++++++++ tests/web/tracking/test_standalone_usage.py | 171 +++++++++++++++++- 6 files changed, 367 insertions(+), 21 deletions(-) diff --git a/src/xagent/core/tools/core/music_tool.py b/src/xagent/core/tools/core/music_tool.py index 21383fcf36..0cf9cc23d1 100644 --- a/src/xagent/core/tools/core/music_tool.py +++ b/src/xagent/core/tools/core/music_tool.py @@ -161,10 +161,19 @@ async def generate_music( ) record_media_seconds( seconds, - # Never str(None): _configured_model_id returns Optional[str] - # and an inactive shared default resolves to None, which would - # bill against a phantom model literally named "None". - model=resolve_billing_model(configured_model_id, model), + # `model` carries the provider's name, `model_id` the + # configured id. Passing configured_model_id as the first + # argument here would return it unchanged, writing the same id + # into both fields and losing the provider name entirely. + # None is passed instead so the resolver falls through to the + # provider's model_name. The fallback keeps the configured id + # ahead of the class name -- a provider exposing no model_name + # (Xinference's default) would otherwise be billed under a + # Python class name while its real id sat unused in scope -- + # and never the placeholder "default". + model=resolve_billing_model( + None, model, fallback=configured_model_id or type(model).__name__ + ), model_id=configured_model_id or "", call_type=MediaCallType.MUSIC, ) diff --git a/src/xagent/core/tools/core/sound_effect_tool.py b/src/xagent/core/tools/core/sound_effect_tool.py index 7109655124..40cabfbc9d 100644 --- a/src/xagent/core/tools/core/sound_effect_tool.py +++ b/src/xagent/core/tools/core/sound_effect_tool.py @@ -174,8 +174,14 @@ async def generate_sound_effect( ) record_media_seconds( seconds, - # Never str(None) — see resolve_billing_model. - model=resolve_billing_model(configured_model_id, model), + # `model` is the provider name, `model_id` the configured id — + # see music_tool: passing configured_model_id as the first + # argument would put the same id in both fields and drop the + # provider name, while the fallback keeps it ahead of the + # class name for providers exposing no model_name. + model=resolve_billing_model( + None, model, fallback=configured_model_id or type(model).__name__ + ), model_id=configured_model_id or "", call_type=MediaCallType.SOUND_EFFECT, ) diff --git a/src/xagent/core/tools/core/video_tool.py b/src/xagent/core/tools/core/video_tool.py index 0c69e66297..41f4521ddf 100644 --- a/src/xagent/core/tools/core/video_tool.py +++ b/src/xagent/core/tools/core/video_tool.py @@ -27,7 +27,11 @@ from ...model.video.base import BaseVideoModel from ...model.video.xinference import XinferenceVideoModel from ...workspace import TaskWorkspace -from .media_usage import coerce_duration, record_media_seconds +from .media_usage import ( + coerce_duration, + record_media_seconds, + resolve_billing_model, +) logger = logging.getLogger(__name__) @@ -737,7 +741,16 @@ async def generate_video( per_video_seconds * billable_count if per_video_seconds is not None else None, - model=str(actual_model_id), + # `model` is the provider's name and `model_id` the configured + # id. Writing the configured id into `model` and leaving + # `model_id` empty loses the canonical name for display and + # external consumers; the aggregator groups on + # `model_id or model`, so populating both keeps the same row + # identity while carrying the name. + model=resolve_billing_model( + None, video_model, fallback=str(actual_model_id) + ), + model_id=str(actual_model_id), call_type=MediaCallType.VIDEO, ) diff --git a/src/xagent/web/services/quota_hooks.py b/src/xagent/web/services/quota_hooks.py index f4eb8f3e4d..e6dbf5a92e 100644 --- a/src/xagent/web/services/quota_hooks.py +++ b/src/xagent/web/services/quota_hooks.py @@ -41,11 +41,17 @@ # For image models whose price varies by resolution, "resolution" ("1K"/"2K"/ # "4K" or "1024x1024") keys a per-(model, resolution) price table. Providers # that report real image tokens (Gemini, OpenAI gpt-image) also fill -# "provider_tokens" so a token-based price ($/1M tokens) can take precedence -# over the resolution table — but only when "tokens_estimated" is False; -# embedding/rerank counts are local heuristics and must not be priced as -# measured tokens. Media token counts are deliberately NOT under the "tokens" -# key, so a consumer summing "tokens" across entries cannot double-count them. +# "provider_tokens", recorded raw for a future consumer. That is NOT a pricing +# rule: no precedence between token-based and unit/resolution-based pricing is +# defined here, the row schema carries no discriminator to express one, and the +# aggregate groups purely by (model, unit, call_type, resolution). See the +# authoritative contract on ``add_media_usage`` in +# ``core/model/chat/token_context.py``; defining that precedence is tracked in +# #1461. When a hook does price "provider_tokens", note "tokens_estimated" +# marks counts that are local heuristics (embedding/rerank) rather than +# provider-reported, so they must not be priced as measured tokens. +# Media token counts are deliberately NOT under the "tokens" key, so a consumer +# summing "tokens" across entries cannot double-count them. # Unknown entry types must be ignored, not summed as tokens. # # TRANSACTION CONTRACT: the hook is invoked from TaskTracker.complete_tracking diff --git a/tests/core/tools/core/test_media_tool_billing_policy.py b/tests/core/tools/core/test_media_tool_billing_policy.py index 8d8804aeac..b96186cd7e 100644 --- a/tests/core/tools/core/test_media_tool_billing_policy.py +++ b/tests/core/tools/core/test_media_tool_billing_policy.py @@ -154,3 +154,160 @@ async def test_asr_usage_meters_duration_with_verbose() -> None: # Name-keyed identity across all three ASR entry points; see PR body. assert entries[0]["model"] == "asr-a" assert entries[0]["model_id"] == "" + + +# --- ASR duration precedence (resolve_asr_seconds) --------------------------- +# +# The producer tests above cover a bare string and one ordered segment. The +# resolver has several more branches — each provider total field, the fallback +# to segment ends, unsorted segments and unusable values — and a regression in +# any of them would silently change the billed quantity while leaving the +# suite green. + + +@pytest.mark.parametrize("field", ["duration", "audio_duration", "duration_seconds"]) +def test_provider_total_duration_wins_over_segments(field: str) -> None: + """A provider total covers trailing silence the last segment end misses.""" + from xagent.core.model.asr.usage import resolve_asr_seconds + + seconds = resolve_asr_seconds( + {field: 30.0}, + [{"start": 0.0, "end": 2.5}], + ) + assert seconds == 30.0 + + +def test_segment_ends_are_used_when_no_provider_total() -> None: + from xagent.core.model.asr.usage import resolve_asr_seconds + + assert resolve_asr_seconds({}, [{"start": 0.0, "end": 7.5}]) == 7.5 + + +def test_unsorted_segments_take_the_maximum_end() -> None: + """Segment order is not guaranteed, so the last element is not the end.""" + from xagent.core.model.asr.usage import resolve_asr_seconds + + seconds = resolve_asr_seconds( + None, + [{"start": 5.0, "end": 9.0}, {"start": 0.0, "end": 3.0}], + ) + assert seconds == 9.0 + + +@pytest.mark.parametrize("bad", [None, "abc", float("inf"), float("nan"), -1.0, True]) +def test_unusable_provider_total_falls_back_to_segments(bad: Any) -> None: + """A non-finite/negative/boolean total must not be billed as a duration.""" + from xagent.core.model.asr.usage import resolve_asr_seconds + + seconds = resolve_asr_seconds({"duration": bad}, [{"start": 0.0, "end": 4.0}]) + assert seconds == 4.0 + + +def test_no_usable_timing_reports_none() -> None: + """None (not 0.0) so the caller can tell "unmeasured" from "zero-length"; + record_media_seconds turns it into a 0-second row plus a warning.""" + from xagent.core.model.asr.usage import resolve_asr_seconds + + assert resolve_asr_seconds({}, []) is None + assert resolve_asr_seconds(None, None) is None + assert resolve_asr_seconds({"duration": None}, [{"start": 0.0, "end": 0.0}]) is None + + +# --- Media identity field shape --------------------------------------------- +# +# Convention: `model` carries the human-readable provider name, `model_id` the +# configured id. Writing the configured id into both (or leaving model_id +# empty) loses the canonical name for display and external consumers. Each +# case below uses a provider name and a configured id that differ, so a +# resolver that returns the id for both fields cannot pass. + + +class _NamedMusicModel: + model_name = "music-provider-name" + + async def generate_music(self, **kwargs: Any) -> MusicResult: + _ = kwargs + return MusicResult(audio=b"x", format="mp3", raw_response={}) + + +@pytest.mark.asyncio +async def test_music_records_provider_name_and_configured_id_separately() -> None: + tool = MusicToolCore(models={"configured-music-id": _NamedMusicModel()}) + + with TokenContextManager() as manager: + await tool.generate_music(prompt="p", music_length_seconds=10) + entries = _media_entries(manager) + + assert len(entries) == 1 + assert entries[0]["model"] == "music-provider-name" + assert entries[0]["model_id"] == "configured-music-id" + + +class _NamedSoundEffectModel: + model_name = "sfx-provider-name" + + async def generate_sound_effect(self, **kwargs: Any) -> Any: + _ = kwargs + return {"not": "a SoundEffectResult"} + + +@pytest.mark.asyncio +async def test_sound_effect_records_provider_name_and_configured_id_separately() -> ( + None +): + tool = SoundEffectToolCore(models={"configured-sfx-id": _NamedSoundEffectModel()}) + + with TokenContextManager() as manager: + await tool.generate_sound_effect(text="p", duration_seconds=4) + entries = _media_entries(manager) + + assert len(entries) == 1 + assert entries[0]["model"] == "sfx-provider-name" + assert entries[0]["model_id"] == "configured-sfx-id" + + +class _UnnamedMusicModel: + """Provider exposing no model_name — Xinference's default behaves this way.""" + + async def generate_music(self, **kwargs: Any) -> MusicResult: + _ = kwargs + return MusicResult(audio=b"x", format="mp3", raw_response={}) + + +@pytest.mark.asyncio +async def test_music_falls_back_to_configured_id_not_the_class_name() -> None: + """A provider with no model_name must still bill under its configured id. + + The resolver's fallback is what decides this, and a Python class name is + not a billing identity: it is not configurable, not unique across + providers, and means nothing to a price table. + """ + tool = MusicToolCore(models={"configured-music-id": _UnnamedMusicModel()}) + + with TokenContextManager() as manager: + await tool.generate_music(prompt="p", music_length_seconds=10) + entries = _media_entries(manager) + + assert len(entries) == 1 + assert entries[0]["model"] == "configured-music-id" + assert entries[0]["model"] != "_UnnamedMusicModel" + assert entries[0]["model_id"] == "configured-music-id" + + +class _UnnamedSoundEffectModel: + async def generate_sound_effect(self, **kwargs: Any) -> Any: + _ = kwargs + return {"not": "a SoundEffectResult"} + + +@pytest.mark.asyncio +async def test_sound_effect_falls_back_to_configured_id_not_the_class_name() -> None: + tool = SoundEffectToolCore(models={"configured-sfx-id": _UnnamedSoundEffectModel()}) + + with TokenContextManager() as manager: + await tool.generate_sound_effect(text="p", duration_seconds=4) + entries = _media_entries(manager) + + assert len(entries) == 1 + assert entries[0]["model"] == "configured-sfx-id" + assert entries[0]["model"] != "_UnnamedSoundEffectModel" diff --git a/tests/web/tracking/test_standalone_usage.py b/tests/web/tracking/test_standalone_usage.py index 0111100585..a606053fa0 100644 --- a/tests/web/tracking/test_standalone_usage.py +++ b/tests/web/tracking/test_standalone_usage.py @@ -6,6 +6,7 @@ object. """ +import asyncio import concurrent.futures from typing import Any, Optional @@ -19,19 +20,40 @@ from xagent.web.tracking.standalone_usage import bind_usage_to_thread, usage_scope +def _async_return(value: Any): + """An async callable ignoring its args and returning ``value``.""" + + async def _fn(*_a: Any, **_k: Any) -> Any: + return value + + return _fn + + class _FakeSession: - def __init__(self) -> None: + """Stand-in for the short-lived compatibility Session ``_report`` opens. + + ``in_transaction`` is configurable and the calls are recorded in order: + with it hard-coded to False the rollback branch is unreachable, so + deleting the rollback from _report would leave every test green. + """ + + def __init__(self, in_transaction: bool = False) -> None: + self._in_transaction = in_transaction self.rolled_back = False self.closed = False + self.calls: list[str] = [] def in_transaction(self) -> bool: - return False + return self._in_transaction def rollback(self) -> None: self.rolled_back = True + self._in_transaction = False + self.calls.append("rollback") def close(self) -> None: self.closed = True + self.calls.append("close") @pytest.fixture @@ -103,6 +125,63 @@ def test_compatibility_session_is_disposed(captured) -> None: assert calls[0]["db"] is session +def test_active_transaction_is_rolled_back_before_close( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """The hook owns its own durability and must not be left holding one. + + With in_transaction() hard-coded False the rollback branch is unreachable, + so removing the rollback from _report would leave the suite green. + """ + from xagent.web.models import database + from xagent.web.services import quota_hooks + + session = _FakeSession(in_transaction=True) + monkeypatch.setattr(quota_hooks, "record_usage", lambda *a, **k: None) + monkeypatch.setattr(quota_hooks, "has_usage_record_hook", lambda: True) + monkeypatch.setattr(database, "get_session_local", lambda: lambda: session) + + with usage_scope(1): + add_media_usage(MediaCallType.ASR, 1, model="m") + + assert session.rolled_back is True + # Ordering matters: rolling back after close would raise on a real Session. + assert session.calls == ["rollback", "close"] + + +def test_active_transaction_is_rolled_back_when_the_hook_raises( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """The finally-branch must dispose the session on the failure path too.""" + from xagent.web.models import database + from xagent.web.services import quota_hooks + + session = _FakeSession(in_transaction=True) + + def _boom(*_a: Any, **_k: Any) -> None: + raise RuntimeError("quota backend down") + + monkeypatch.setattr(quota_hooks, "record_usage", _boom) + monkeypatch.setattr(quota_hooks, "has_usage_record_hook", lambda: True) + monkeypatch.setattr(database, "get_session_local", lambda: lambda: session) + + with usage_scope(1): + add_media_usage(MediaCallType.ASR, 1, model="m") + + assert session.rolled_back is True + assert session.calls == ["rollback", "close"] + + +def test_inactive_transaction_is_closed_without_rollback(captured) -> None: + """The common case: nothing pending, so no rollback is issued.""" + _calls, session = captured + with usage_scope(1): + add_media_usage(MediaCallType.ASR, 1, model="m") + + assert session.rolled_back is False + assert session.calls == ["close"] + + def test_bind_usage_to_thread_carries_usage_across_the_hop(captured) -> None: """run_in_executor / ThreadPoolExecutor drop contextvars; the wrapper is what keeps the worker's records attached to the caller.""" @@ -196,11 +275,87 @@ def test_no_hook_installed_skips_the_session_checkout( assert checkouts == [] -def test_patched_entry_points_bind_a_scope() -> None: - """Guard the wiring itself: if a future refactor drops the scope from an - entry point, its usage silently stops being billed again.""" - import inspect +def test_transcribe_endpoint_reports_asr_usage_for_the_authenticated_owner( + captured, monkeypatch: pytest.MonkeyPatch +) -> None: + """Drive the real endpoint and observe the quota hook. + + Asserting that the literal "usage_scope(" appears in the source would pass + with the wrong user, the wrong scope boundary, or no recording at all. This + observes what actually reaches the hook: owner, quantity, unit and that it + is reported exactly once. + """ + calls, _ = captured + + from xagent.core.model.asr import adapter + from xagent.core.model.asr.base import ASRResult + from xagent.web.api import model as model_api + + class _FakeASR: + model_name = "asr-provider" + + async def transcribe(self, audio: Any, **kwargs: Any) -> ASRResult: + _ = (audio, kwargs) + # A real ASRResult, not a dict: record_asr_usage reads + # raw_response/segments off ASRResult, so a dict meters as 0s. + return ASRResult(text="hello", raw_response={"duration": 12.5}) + + class _DBModel: + # Mirrors the real DBModel: `id` is the integer PK and `model_id` the + # configured string id. The endpoint returns both model_id and + # model_name, so a fake carrying only one of them fails on the + # response build rather than on anything this test is about. + id = 1 + model_id = "configured-asr-id" + model_name = "asr-provider" + category = "speech" + + class _User: + id = 4242 + + class _Upload: + filename = "clip.wav" + content_type = "audio/wav" + + async def read(self, *_a: Any, **_k: Any) -> bytes: + return b"RIFFfake" + + async def close(self) -> None: + return None - from xagent.web.api import model + monkeypatch.setattr( + model_api, "_resolve_asr_model_for_transcription", lambda *a, **k: _DBModel() + ) + monkeypatch.setattr(adapter, "get_asr_model_instance", lambda *a, **k: _FakeASR()) + monkeypatch.setattr( + model_api, + "_read_transcribe_upload_with_size_limit", + _async_return(b"RIFFfake"), + ) + + response = asyncio.run( + model_api.transcribe_speech_input( + file=_Upload(), + language=None, + model_id=None, + db=object(), + user=_User(), + ) + ) - assert "usage_scope(" in inspect.getsource(model.transcribe_speech_input) + # Assert the response too: it is built from db_model after the scope + # closes, so a fake that diverges from the real schema shows up here + # instead of as an AttributeError deep in the endpoint. + assert response["text"] == "hello" + assert response["model_id"] == "configured-asr-id" + assert response["model_name"] == "asr-provider" + + # Exactly one report, for the authenticated owner, with the ASR row. + assert len(calls) == 1, calls + assert calls[0]["user_id"] == 4242 + media = [d for d in calls[0]["details"] if d.get("type") == "media"] + assert len(media) == 1, media + assert media[0]["call_type"] == "asr" + assert media[0]["unit"] == "seconds" + assert media[0]["quantity"] == 12.5 + assert media[0]["model"] == "asr-provider" From cdde2360333376333b5a21b7bef37eac6ac03ce5 Mon Sep 17 00:00:00 2001 From: OliverBryant Date: Fri, 21 Aug 2026 18:25:20 +0800 Subject: [PATCH 5/6] fix: apply the record-before-validation policy uniformly The stated invariant -- record as soon as the provider call returns, before the response is validated -- held for TTS, music, sound effect and video but not for ASR, and the identity rules were applied at four of five call sites. Both gaps are the bug class this metering exists to remove, reappearing inside the PR that removes it. ASR recording order (M1): transcribe_audio resolved the duration and then recorded it only after _aggregate_segments had run. That function raises ValueError on a segment missing start/end, and the tool's broad handler turns any raise into success:False -- so a provider call that succeeded and was billed went entirely unmetered. Reproduced end to end: provider returns 42s, post-processing raises, zero rows recorded. The record now happens immediately after the duration is resolved, matching TTS in the same file. This also corrects the comment at the site, which asserted the opposite policy, and the stale model_id note flagged separately (N5). Telegram identity fallback (M2): resolve_billing_model(None, asr_model) fell through to the helper's own default of "default" -- the exact placeholder the module invariants forbid as a billing identity, while the adjacent comment claimed the opposite. Now passes an explicit fallback like the other four call sites. Video billing coverage (M3): no test exercised video's metering at all. Added the duration * n multiplication, the provider-name/configured-id split shipped earlier in this PR, and the async no-duration path that records seconds with quantity 0. Video placeholder id (N1): _model_id_for_model bottoms out at the literal "default" when a model exposes neither id nor name, and that was written straight into model_id. Dropped to "" instead, matching music/sound_effect, so the aggregator falls through to the resolved name rather than inventing a phantom model shared by every unidentifiable provider. Dead scaffolding (N7): _RecordingASR captured a verbose flag no test ever asserted on. Completed the intended assertion rather than deleting it -- verbose=True reaching the provider is what makes ASR billable at all, since without it the provider returns a bare string and every call meters as an unbillable 0 seconds. bind_usage_to_thread (N3): removed along with its tests. It had no callers anywhere in the repo, every thread hop in this diff uses asyncio.to_thread which copies the context already, and shipping an unused concurrency helper built on a documented non-thread-safe object is a trap. It can return with its first real consumer, where the locking contract can be settled against actual usage. Every fix is mutation-checked: reverting the ASR record past _aggregate_segments, or the video/music fallbacks, turns the new tests red. --- src/xagent/core/tools/core/audio_tool.py | 39 ++-- src/xagent/core/tools/core/video_tool.py | 10 +- src/xagent/web/api/model.py | 13 +- src/xagent/web/channels/telegram/bot.py | 15 +- src/xagent/web/tracking/standalone_usage.py | 81 ++------ .../core/test_media_tool_billing_policy.py | 183 ++++++++++++++++++ tests/web/tracking/test_standalone_usage.py | 105 ++++++---- 7 files changed, 327 insertions(+), 119 deletions(-) diff --git a/src/xagent/core/tools/core/audio_tool.py b/src/xagent/core/tools/core/audio_tool.py index 94c331fe2a..66e9604695 100644 --- a/src/xagent/core/tools/core/audio_tool.py +++ b/src/xagent/core/tools/core/audio_tool.py @@ -780,6 +780,28 @@ async def transcribe_audio( # as a fallback) has exactly one implementation. audio_seconds = resolve_asr_seconds(raw_response, raw_segments) + # Recorded here, before _aggregate_segments and the rest of the + # post-processing below. Those steps can raise (a segment missing + # start/end raises ValueError), and the broad handler at the bottom + # turns any raise into success:False -- so recording afterwards + # meant a provider call that succeeded and was billed went + # unmetered. That is the bug class this metering exists to remove, + # and TTS in this same file already records before unpacking. + # Routed through the shared ASR recorder so every entry point + # meters identically. + record_asr_seconds( + audio_seconds, + # model_id is deliberately left unset. This tool is handed + # name-keyed registries (model_service keys _asr_models by + # model_name), so the only identity in scope here is a name — + # writing a name into model_id would persist it under an id + # field. The aggregator groups on `model_id or model`, so + # leaving it empty lets these rows key on the name + # consistently instead of inventing a third identity that can + # never be reconciled. + model_name=str(actual_model_id), + ) + segment_view = "raw" if verbose else "processed" segments = raw_segments segments_merged = False @@ -792,23 +814,6 @@ async def transcribe_audio( ) segments_merged = len(segments) < len(raw_segments) - # Recorded only after every step that can still fail, so a call that - # errors out is not billed. Routed through the shared ASR recorder - # so every entry point meters identically. - record_asr_seconds( - audio_seconds, - # model_id is deliberately left unset. This tool is handed - # name-keyed registries (model_service keys _asr_models by - # model_name), so the only identity in scope here is a name — - # writing it into model_id would persist a name under an id - # field and still not match the real DB model_id that - # /speech/transcribe records. The aggregator groups on - # `model_id or model`, so leaving it empty lets these rows key - # on the name consistently instead of inventing a third - # identity that can never be reconciled. - model_name=str(actual_model_id), - ) - # Save transcription to JSON file if workspace is available file_id: Optional[str] = None file_ref: Optional[dict[str, Any]] = None diff --git a/src/xagent/core/tools/core/video_tool.py b/src/xagent/core/tools/core/video_tool.py index 41f4521ddf..443c4edf50 100644 --- a/src/xagent/core/tools/core/video_tool.py +++ b/src/xagent/core/tools/core/video_tool.py @@ -737,6 +737,7 @@ async def generate_video( # first video's duration was recorded. per_video_seconds = coerce_duration(result.get("duration")) billable_count = max(1, int(n or 1)) + billing_model_id = str(actual_model_id) record_media_seconds( per_video_seconds * billable_count if per_video_seconds is not None @@ -750,7 +751,14 @@ async def generate_video( model=resolve_billing_model( None, video_model, fallback=str(actual_model_id) ), - model_id=str(actual_model_id), + # _model_id_for_model bottoms out at the literal "default" when + # a model exposes neither an id nor a name, and that string is + # a forbidden billing identity. Drop it rather than persist it, + # matching music/sound_effect which pass `configured_id or ""`. + # The aggregator groups on `model_id or model`, so an empty id + # falls through to the resolved name instead of inventing a + # phantom model shared by every unidentifiable provider. + model_id=("" if billing_model_id == "default" else billing_model_id), call_type=MediaCallType.VIDEO, ) diff --git a/src/xagent/web/api/model.py b/src/xagent/web/api/model.py index 2d8082256e..da96b9e4bf 100644 --- a/src/xagent/web/api/model.py +++ b/src/xagent/web/api/model.py @@ -981,6 +981,7 @@ async def transcribe_speech_input( from xagent.core.model.asr.adapter import get_asr_model_instance from xagent.core.model.asr.usage import record_asr_usage + from xagent.core.tools.core.media_usage import resolve_billing_model from ..tracking.standalone_usage import usage_scope @@ -1022,7 +1023,17 @@ async def transcribe_speech_input( # the one identity all three share. record_asr_usage( result, - model_name=str(db_model.model_name), + # Routed through the shared resolver rather than str() on the + # DB column: model_name is nullable=False but not constrained + # to be non-placeholder, so an empty or literally "default" + # name would otherwise be billed as a model identity, which + # the module invariants forbid. The configured id is the + # fallback here -- it is the one real identity in scope -- and + # it only ever reaches `model`, never `model_id`, so the + # single-identity rule above still holds. + model_name=resolve_billing_model( + None, db_model, fallback=str(db_model.model_id) + ), ) except asyncio.TimeoutError as exc: raise HTTPException(status_code=504, detail="Transcription timed out") from exc diff --git a/src/xagent/web/channels/telegram/bot.py b/src/xagent/web/channels/telegram/bot.py index 1599e3c3e9..80e47dd51c 100644 --- a/src/xagent/web/channels/telegram/bot.py +++ b/src/xagent/web/channels/telegram/bot.py @@ -1400,12 +1400,19 @@ async def _transcribe_uploaded_voice_files( ) # Telegram calls the ASR provider directly rather than going # through audio_tool, so meter here too. Identity goes through - # the shared resolver so a provider exposing no model_name is - # never billed under a placeholder, and all three ASR entry - # points agree on the name the aggregator groups by. + # the shared resolver so all three ASR entry points agree on + # the name the aggregator groups by. The explicit fallback + # matters: resolve_billing_model's own default is the literal + # "default", which the module invariants forbid as a billing + # identity, so a provider exposing no model_name would + # otherwise be billed under exactly the placeholder this is + # meant to avoid. The provider class name is the last identity + # that still attributes cost to something real. record_asr_usage( result, - model_name=resolve_billing_model(None, asr_model), + model_name=resolve_billing_model( + None, asr_model, fallback=type(asr_model).__name__ + ), ) except asyncio.TimeoutError as exc: raise TelegramVoiceTranscriptionError( diff --git a/src/xagent/web/tracking/standalone_usage.py b/src/xagent/web/tracking/standalone_usage.py index a530101285..eff592b2fa 100644 --- a/src/xagent/web/tracking/standalone_usage.py +++ b/src/xagent/web/tracking/standalone_usage.py @@ -20,18 +20,14 @@ from __future__ import annotations -import functools import logging from contextlib import contextmanager -from contextvars import copy_context -from typing import Any, Callable, Iterator, Optional, TypeVar +from typing import Iterator, Optional from ...core.model.chat.token_context import TokenUsage, set_token_usage logger = logging.getLogger(__name__) -T = TypeVar("T") - def _report(user_id: Optional[int], usage: TokenUsage) -> None: """Hand this unit of work's usage to the quota hook, best-effort.""" @@ -39,27 +35,29 @@ def _report(user_id: Optional[int], usage: TokenUsage) -> None: if not details: return try: - from ..models.database import get_session_local - from ..services.quota_hooks import has_usage_record_hook, record_usage + from ..services.quota_hooks import has_usage_record_hook + + # Imported lazily, like the hook above: this module is deliberately + # light so importing it costs nothing at startup, and task_tracker + # pulls in the whole tracking chain. + from .task_tracker import _record_usage_on_event_loop # Check the hook before checking out a session: with no hook installed # (the stock configuration) record_usage is a guaranteed no-op, and a # pool checkout + transaction + close per ingest/transcription is pure - # overhead. + # overhead. The sibling path in task_tracker predates this check and + # still pays it; extending the short-circuit there is a separate change + # to a hot path this PR does not otherwise touch. if not has_usage_record_hook(): return - db_session = get_session_local()() - try: - # delta_actions=0: these paths make provider calls, not agent tool - # calls, and tool invocations are what that counter bills for. - record_usage(db_session, user_id, details, 0) - finally: - # The hook owns its own durability and must not leave work pending - # on this compatibility Session. - if db_session.in_transaction(): - db_session.rollback() - db_session.close() + # Reuses task_tracker's helper rather than repeating its session + # lifecycle: it already owns the "hand the hook a short-lived + # compatibility Session, never leave it holding a transaction" + # contract, and two copies of that is how they drift. + # delta_actions=0: these paths make provider calls, not agent tool + # calls, and tool invocations are what that counter bills for. + _record_usage_on_event_loop(user_id, details, 0) except Exception as e: # noqa: BLE001 # Metering must never break the work it is measuring. logger.warning("Standalone usage recording failed: %s", e) @@ -72,8 +70,10 @@ def usage_scope(user_id: Optional[int]) -> Iterator[TokenUsage]: Usage is reported even when the body raises: a provider call that already happened is billable regardless of what fails afterwards. - Note the body must not cross a thread boundary that drops contextvars — - see :func:`bind_usage_to_thread` for `run_in_executor`-style hops. + Note the body must not cross a thread boundary that drops contextvars. + ``asyncio.to_thread`` copies the context and is safe; a bare + ``ThreadPoolExecutor`` or ``run_in_executor`` is not, and would need the + caller's usage bound explicitly inside the worker. """ from ...core.model.chat.token_context import token_context @@ -101,42 +101,3 @@ def usage_scope(user_id: Optional[int]) -> Iterator[TokenUsage]: except Exception: # noqa: BLE001 pass _report(user_id, usage) - - -def bind_usage_to_thread(fn: Callable[..., T]) -> Callable[..., T]: - """Wrap a callable so it records into the *calling* thread's usage. - - ``loop.run_in_executor(None, fn)`` and bare ``ThreadPoolExecutor`` do not - propagate contextvars, so without this the worker records into a fresh - ``TokenUsage`` that is discarded on return. ``asyncio.to_thread`` already - copies the context and does not need this. - - Captures at wrap time, on the calling thread — wrapping must therefore - happen before the hop, not inside the worker. - - The binding runs inside a copied context (``copy_context()`` + - ``ctx.run``), the same idiom used at - ``core/tools/adapters/vibe/file_ingestion_tool.py``. That confines the - contextvar write to this one call instead of mutating the worker thread: - ``run_in_executor(None, ...)`` uses the loop's long-lived default executor, - so a plain ``set_token_usage`` would outlive the job and leak the caller's - ``TokenUsage`` into the next unrelated task that reuses the thread — - cross-tenant usage misattribution, plus a strong reference pinned for the - thread's lifetime. A copied context also needs no ``finally`` restore: - there is nothing in the caller's context left to restore. - """ - from ...core.model.chat.token_context import get_token_usage - - caller_usage = get_token_usage() - - @functools.wraps(fn) - def _bound(*args: Any, **kwargs: Any) -> T: - def _run() -> T: - set_token_usage(caller_usage) - return fn(*args, **kwargs) - - # Fresh copy per invocation: a wrapped callable may be submitted more - # than once, and each run must get its own isolated context. - return copy_context().run(_run) - - return _bound diff --git a/tests/core/tools/core/test_media_tool_billing_policy.py b/tests/core/tools/core/test_media_tool_billing_policy.py index b96186cd7e..f37d44f539 100644 --- a/tests/core/tools/core/test_media_tool_billing_policy.py +++ b/tests/core/tools/core/test_media_tool_billing_policy.py @@ -311,3 +311,186 @@ async def test_sound_effect_falls_back_to_configured_id_not_the_class_name() -> assert len(entries) == 1 assert entries[0]["model"] == "configured-sfx-id" assert entries[0]["model"] != "_UnnamedSoundEffectModel" + + +# --- Video billing -------------------------------------------------------- +# +# video_tool is duration-billed like music/sound effect, but bills +# `duration * n` because the provider reports one duration while generating +# and charging for n videos. These cover the multiplication, the identity +# fields, and the async path that reports no duration. + + +def _video_model(duration: Any, *, name: str = "video-provider") -> Any: + from unittest.mock import AsyncMock, Mock + + from xagent.core.model.video.base import BaseVideoModel + + model = Mock(spec=BaseVideoModel) + model.has_ability = Mock(return_value=True) + model.abilities = ["generate"] + model.model_name = name + result = {"task_id": "t-1", "status": "succeeded", "video_url": "", "ratio": "16:9"} + if duration is not None: + result["duration"] = duration + model.generate_video = AsyncMock(return_value=result) + return model + + +@pytest.mark.asyncio +async def test_video_bills_duration_times_count() -> None: + """The provider reports one duration but generates and bills for n videos.""" + from xagent.core.tools.core.video_tool import VideoGenerationToolCore + + tool = VideoGenerationToolCore(video_models={"cfg-video-id": _video_model(5)}) + + with TokenContextManager() as manager: + await tool.generate_video(prompt="p", n=3) + entries = _media_entries(manager) + + assert len(entries) == 1 + assert entries[0]["unit"] == "seconds" + assert entries[0]["quantity"] == 15.0 + assert entries[0]["call_type"] == "video" + + +@pytest.mark.asyncio +async def test_video_records_provider_name_and_configured_id_separately() -> None: + """`model` is the provider name, `model_id` the configured id.""" + from xagent.core.tools.core.video_tool import VideoGenerationToolCore + + tool = VideoGenerationToolCore( + video_models={"cfg-video-id": _video_model(5, name="video-provider")} + ) + + with TokenContextManager() as manager: + await tool.generate_video(prompt="p") + entries = _media_entries(manager) + + assert len(entries) == 1 + assert entries[0]["model"] == "video-provider" + assert entries[0]["model_id"] == "cfg-video-id" + + +@pytest.mark.asyncio +async def test_video_without_duration_is_recorded_as_unmeasured_seconds() -> None: + """An async task reports no duration yet. The unit must stay seconds with + quantity 0 rather than degrading to another unit -- reconciling that row is + tracked in #1583.""" + from xagent.core.tools.core.video_tool import VideoGenerationToolCore + + tool = VideoGenerationToolCore(video_models={"cfg-video-id": _video_model(None)}) + + with TokenContextManager() as manager: + await tool.generate_video(prompt="p") + entries = _media_entries(manager) + + assert len(entries) == 1 + assert entries[0]["unit"] == "seconds" + assert entries[0]["quantity"] == 0.0 + + +# --- ASR records before post-processing ----------------------------------- + + +class _MalformedSegmentASR: + """Provider call succeeds and is billed; a segment lacks an end time. + + _aggregate_segments raises ValueError on this input, which the tool's + broad handler turns into success:False. The provider still ran and + charged, so the row must already be recorded by then. + """ + + model_name = "asr-provider" + + @property + def abilities(self) -> list: + return ["asr", "timestamps"] + + async def transcribe(self, audio: Any, **kwargs: Any) -> Any: + _ = (audio, kwargs) + from xagent.core.model.asr.base import ASRResult, ASRSegment + + return ASRResult( + text="hello world", + raw_response={"duration": 42.0}, + segments=[ + ASRSegment(text="hello", start=0.0, end=1.0, confidence=1.0), + ASRSegment(text="world", start=1.0, end=None, confidence=1.0), + ], + ) + + +@pytest.mark.asyncio +async def test_asr_bills_when_post_processing_raises() -> None: + """Metering must not depend on post-processing succeeding. + + Recording after _aggregate_segments meant a provider call that succeeded + and was billed went entirely unmetered whenever segment data was + malformed -- the exact bug class this module's policy exists to prevent. + """ + from xagent.core.tools.core.audio_tool import AudioToolCore + + tool = AudioToolCore(asr_models={"asr-provider": _MalformedSegmentASR()}) + + with TokenContextManager() as manager: + result = await tool.transcribe_audio( + audio_file_path="/tmp/does-not-exist.wav", verbose=False + ) + entries = _media_entries(manager) + + # The tool still reports failure to the caller... + assert result["success"] is False + # ...but the provider call happened and was billed, so it is metered. + assert len(entries) == 1 + assert entries[0]["unit"] == "seconds" + assert entries[0]["quantity"] == 42.0 + assert entries[0]["call_type"] == "asr" + + +# --- verbose=True reaches the provider ------------------------------------ +# +# _RecordingASR exists to capture the flag. Without verbose=True the provider +# returns a bare string with no timings and every call meters as an +# unbillable 0 seconds, so the flag reaching the provider is what makes ASR +# billable at all. + + +@pytest.mark.asyncio +async def test_telegram_passes_verbose_and_bills_a_real_duration(tmp_path) -> None: + """Drive the real Telegram method, not a re-implementation of it. + + Asserting on a locally-issued transcribe() call would pass with + verbose=True deleted from bot.py, which is the only thing worth guarding + here: without it the provider returns a bare string with no timings and + every voice message meters as an unbillable 0 seconds. + """ + from xagent.web.channels.telegram.bot import TelegramBotInstance + + provider = _RecordingASR() + # __new__ rather than the real constructor: the method only touches a + # class-level timeout and a staticmethod, so a fully wired bot instance + # would add setup without adding coverage. + bot = TelegramBotInstance.__new__(TelegramBotInstance) + + audio = tmp_path / "voice.ogg" + audio.write_bytes(b"fake-ogg") + + with TokenContextManager() as manager: + transcripts = await bot._transcribe_uploaded_voice_files( + ["file-1"], + [{"telegram_file_id": "file-1", "path": str(audio), "name": "voice.ogg"}], + provider, + ) + entries = _media_entries(manager) + + # bot.py passed verbose through to the provider... + assert provider.verbose_seen is True + assert transcripts == {"file-1": "hi"} + # ...so the row carries a real duration instead of an unbillable zero. + assert len(entries) == 1 + assert entries[0]["unit"] == "seconds" + assert entries[0]["quantity"] == 2.5 + # and the identity is the provider name, never the forbidden placeholder. + assert entries[0]["model"] == "asr-a" + assert entries[0]["model"] != "default" diff --git a/tests/web/tracking/test_standalone_usage.py b/tests/web/tracking/test_standalone_usage.py index a606053fa0..dd7777dce1 100644 --- a/tests/web/tracking/test_standalone_usage.py +++ b/tests/web/tracking/test_standalone_usage.py @@ -7,7 +7,6 @@ """ import asyncio -import concurrent.futures from typing import Any, Optional import pytest @@ -17,7 +16,7 @@ add_media_usage, get_token_usage, ) -from xagent.web.tracking.standalone_usage import bind_usage_to_thread, usage_scope +from xagent.web.tracking.standalone_usage import usage_scope def _async_return(value: Any): @@ -182,40 +181,6 @@ def test_inactive_transaction_is_closed_without_rollback(captured) -> None: assert session.calls == ["close"] -def test_bind_usage_to_thread_carries_usage_across_the_hop(captured) -> None: - """run_in_executor / ThreadPoolExecutor drop contextvars; the wrapper is - what keeps the worker's records attached to the caller.""" - calls, _ = captured - - with usage_scope(9): - - def work() -> None: - add_media_usage(MediaCallType.ASR, 4, model="e") - - bound = bind_usage_to_thread(work) - with concurrent.futures.ThreadPoolExecutor(max_workers=2) as pool: - pool.submit(bound).result() - - assert len(calls) == 1 - assert calls[0]["details"][0]["quantity"] == 4.0 - - -def test_unbound_thread_hop_loses_usage(captured) -> None: - """Documents why the wrapper exists: the same code without it records - nothing, which is the defect this module fixes.""" - calls, _ = captured - - with usage_scope(9): - - def work() -> None: - add_media_usage(MediaCallType.ASR, 4, model="e") - - with concurrent.futures.ThreadPoolExecutor(max_workers=2) as pool: - pool.submit(work).result() # not bound - - assert calls == [] - - def test_scope_restores_the_previous_context(captured) -> None: """A nested scope must not leak its usage object into the outer one.""" _calls, _ = captured @@ -359,3 +324,71 @@ async def close(self) -> None: assert media[0]["unit"] == "seconds" assert media[0]["quantity"] == 12.5 assert media[0]["model"] == "asr-provider" + + +def test_transcribe_falls_back_to_configured_id_for_a_placeholder_db_name( + captured, monkeypatch: pytest.MonkeyPatch +) -> None: + """model_name is nullable=False but not constrained to be meaningful. + + A row whose name is empty or literally "default" must not be billed under + that string -- the module invariants forbid placeholder identities -- so + the resolver falls back to the configured id. + """ + calls, _ = captured + + from xagent.core.model.asr import adapter + from xagent.core.model.asr.base import ASRResult + from xagent.web.api import model as model_api + + class _FakeASR: + model_name = "asr-provider" + + async def transcribe(self, audio: Any, **kwargs: Any) -> ASRResult: + _ = (audio, kwargs) + return ASRResult(text="hello", raw_response={"duration": 3.0}) + + class _PlaceholderNameDBModel: + id = 1 + model_id = "configured-asr-id" + model_name = "default" + category = "speech" + + class _User: + id = 7 + + class _Upload: + filename = "clip.wav" + content_type = "audio/wav" + + async def read(self, *_a: Any, **_k: Any) -> bytes: + return b"RIFFfake" + + async def close(self) -> None: + return None + + monkeypatch.setattr( + model_api, + "_resolve_asr_model_for_transcription", + lambda *a, **k: _PlaceholderNameDBModel(), + ) + monkeypatch.setattr(adapter, "get_asr_model_instance", lambda *a, **k: _FakeASR()) + monkeypatch.setattr( + model_api, + "_read_transcribe_upload_with_size_limit", + _async_return(b"RIFFfake"), + ) + + asyncio.run( + model_api.transcribe_speech_input( + file=_Upload(), + language=None, + model_id=None, + db=object(), + user=_User(), + ) + ) + + media = [d for d in calls[0]["details"] if d.get("type") == "media"] + assert media[0]["model"] == "configured-asr-id" + assert media[0]["model"] != "default" From 7c01b3db983f3679c78b10c7ac0ec1351825451e Mon Sep 17 00:00:00 2001 From: OliverBryant Date: Wed, 26 Aug 2026 15:36:02 +0800 Subject: [PATCH 6/6] test: collapse duplicated media-billing fakes into factories Seven near-identical fake provider classes differed along only two axes: whether they expose a model_name, and what the call returns. Replaced with two factories taking those as parameters, plus a _NO_NAME sentinel that distinguishes "no model_name attribute at all" (Xinference's default behaves this way) from "the name is empty". The two music tests that assert a failed-but-billed call are now one parametrized test with named cases, since they differed only in which way the response was unusable and in the billed quantity. Both failure modes stay visible as case ids. Kept deliberately separate: - The two identity tests per modality assert opposite outcomes (provider name wins vs configured id wins), so merging them would hide which invariant broke. - Music and sound effect stay separate tests rather than one parametrized pair. A modality parameter would need four lambdas per case to carry the differing constructors and call signatures, which is harder to review than the duplication it removes. No production code touched, no assertion weakened, no scenario dropped. --- .../core/test_media_tool_billing_policy.py | 144 ++++++++---------- 1 file changed, 66 insertions(+), 78 deletions(-) diff --git a/tests/core/tools/core/test_media_tool_billing_policy.py b/tests/core/tools/core/test_media_tool_billing_policy.py index f37d44f539..db3730d2c8 100644 --- a/tests/core/tools/core/test_media_tool_billing_policy.py +++ b/tests/core/tools/core/test_media_tool_billing_policy.py @@ -7,6 +7,7 @@ or malformed was still charged for. """ +from pathlib import Path from typing import Any, Optional import pytest @@ -16,37 +17,67 @@ from xagent.core.tools.core.music_tool import MusicToolCore from xagent.core.tools.core.sound_effect_tool import SoundEffectToolCore +# Sentinel distinguishing "this fake has no model_name attribute at all" +# (Xinference's default model behaves this way) from "its name is empty". +_NO_NAME = object() -class _EmptyMusicModel: - """Provider that returns a well-formed-but-empty result. Still billed.""" +# The two ways a provider call can succeed at the HTTP level and still yield +# nothing usable. Both are billed, which is the policy this module guards. +_EMPTY_MUSIC = MusicResult(audio=b"", format="mp3", raw_response={}) +_GARBAGE = {"not": "a result object"} +_PLAYABLE_MUSIC = MusicResult(audio=b"x", format="mp3", raw_response={}) - model_name = "music-a" - async def generate_music(self, **kwargs: Any) -> MusicResult: - _ = kwargs - return MusicResult(audio=b"", format="mp3", raw_response={}) +def _music_model(result: Any = _GARBAGE, *, name: Any = "music-a") -> Any: + """A music provider returning ``result``; ``name=_NO_NAME`` exposes none.""" + class _FakeMusicModel: + async def generate_music(self, **kwargs: Any) -> Any: + _ = kwargs + return result -class _GarbageMusicModel: - """Provider that returns the wrong type entirely. Still billed.""" + if name is not _NO_NAME: + _FakeMusicModel.model_name = name # type: ignore[attr-defined] + return _FakeMusicModel() - model_name = "music-b" - async def generate_music(self, **kwargs: Any) -> Any: - _ = kwargs - return {"not": "a MusicResult"} +def _sound_effect_model(*, name: Any = "sfx-a") -> Any: + """A sound-effect provider returning the wrong type; still billed.""" + + class _FakeSoundEffectModel: + async def generate_sound_effect(self, **kwargs: Any) -> Any: + _ = kwargs + return _GARBAGE + + if name is not _NO_NAME: + _FakeSoundEffectModel.model_name = name # type: ignore[attr-defined] + return _FakeSoundEffectModel() def _media_entries(manager: TokenContextManager) -> list[dict]: return [d for d in manager.get_usage().details if d.get("type") == "media"] +@pytest.mark.parametrize( + ("bad_result", "seconds"), + [ + pytest.param(_EMPTY_MUSIC, 30.0, id="well-formed-but-empty"), + pytest.param(_GARBAGE, 12.0, id="wrong-type-entirely"), + ], +) @pytest.mark.asyncio -async def test_music_bills_empty_audio_response() -> None: - tool = MusicToolCore(models={"music-a": _EmptyMusicModel()}) +async def test_music_bills_a_call_that_returned_nothing_usable( + bad_result: Any, seconds: float +) -> None: + """Both ways a call can succeed at the HTTP level and yield nothing. + + The unit stays seconds in both: a (model, unit) price table breaks if the + unit varies with how complete the response happened to be. + """ + tool = MusicToolCore(models={"music-a": _music_model(bad_result)}) with TokenContextManager() as manager: - result = await tool.generate_music(prompt="p", music_length_seconds=30) + result = await tool.generate_music(prompt="p", music_length_seconds=seconds) entries = _media_entries(manager) # The tool still reports failure to the caller... @@ -54,37 +85,13 @@ async def test_music_bills_empty_audio_response() -> None: # ...but the provider call happened and was billed, so it is metered. assert len(entries) == 1 assert entries[0]["unit"] == "seconds" - assert entries[0]["quantity"] == 30.0 + assert entries[0]["quantity"] == seconds assert entries[0]["call_type"] == "music" -@pytest.mark.asyncio -async def test_music_bills_malformed_response() -> None: - tool = MusicToolCore(models={"music-b": _GarbageMusicModel()}) - - with TokenContextManager() as manager: - result = await tool.generate_music(prompt="p", music_length_seconds=12) - entries = _media_entries(manager) - - assert result["success"] is False - assert len(entries) == 1 - # Unit stays seconds even though nothing usable came back: a (model, unit) - # price table breaks if the unit varies with response completeness. - assert entries[0]["unit"] == "seconds" - assert entries[0]["quantity"] == 12.0 - - -class _EmptySoundEffectModel: - model_name = "sfx-a" - - async def generate_sound_effect(self, **kwargs: Any) -> Any: - _ = kwargs - return {"not": "a SoundEffectResult"} - - @pytest.mark.asyncio async def test_sound_effect_bills_malformed_response() -> None: - tool = SoundEffectToolCore(models={"sfx-a": _EmptySoundEffectModel()}) + tool = SoundEffectToolCore(models={"sfx-a": _sound_effect_model()}) with TokenContextManager() as manager: result = await tool.generate_sound_effect(text="p", duration_seconds=4) @@ -222,17 +229,15 @@ def test_no_usable_timing_reports_none() -> None: # resolver that returns the id for both fields cannot pass. -class _NamedMusicModel: - model_name = "music-provider-name" - - async def generate_music(self, **kwargs: Any) -> MusicResult: - _ = kwargs - return MusicResult(audio=b"x", format="mp3", raw_response={}) - - @pytest.mark.asyncio async def test_music_records_provider_name_and_configured_id_separately() -> None: - tool = MusicToolCore(models={"configured-music-id": _NamedMusicModel()}) + """`model` carries the provider's name, `model_id` the configured id. + + Writing the id into both fields loses the canonical name for display and + for external consumers. + """ + model = _music_model(_PLAYABLE_MUSIC, name="music-provider-name") + tool = MusicToolCore(models={"configured-music-id": model}) with TokenContextManager() as manager: await tool.generate_music(prompt="p", music_length_seconds=10) @@ -243,19 +248,12 @@ async def test_music_records_provider_name_and_configured_id_separately() -> Non assert entries[0]["model_id"] == "configured-music-id" -class _NamedSoundEffectModel: - model_name = "sfx-provider-name" - - async def generate_sound_effect(self, **kwargs: Any) -> Any: - _ = kwargs - return {"not": "a SoundEffectResult"} - - @pytest.mark.asyncio async def test_sound_effect_records_provider_name_and_configured_id_separately() -> ( None ): - tool = SoundEffectToolCore(models={"configured-sfx-id": _NamedSoundEffectModel()}) + model = _sound_effect_model(name="sfx-provider-name") + tool = SoundEffectToolCore(models={"configured-sfx-id": model}) with TokenContextManager() as manager: await tool.generate_sound_effect(text="p", duration_seconds=4) @@ -266,14 +264,6 @@ async def test_sound_effect_records_provider_name_and_configured_id_separately() assert entries[0]["model_id"] == "configured-sfx-id" -class _UnnamedMusicModel: - """Provider exposing no model_name — Xinference's default behaves this way.""" - - async def generate_music(self, **kwargs: Any) -> MusicResult: - _ = kwargs - return MusicResult(audio=b"x", format="mp3", raw_response={}) - - @pytest.mark.asyncio async def test_music_falls_back_to_configured_id_not_the_class_name() -> None: """A provider with no model_name must still bill under its configured id. @@ -282,7 +272,8 @@ async def test_music_falls_back_to_configured_id_not_the_class_name() -> None: not a billing identity: it is not configurable, not unique across providers, and means nothing to a price table. """ - tool = MusicToolCore(models={"configured-music-id": _UnnamedMusicModel()}) + model = _music_model(_PLAYABLE_MUSIC, name=_NO_NAME) + tool = MusicToolCore(models={"configured-music-id": model}) with TokenContextManager() as manager: await tool.generate_music(prompt="p", music_length_seconds=10) @@ -290,19 +281,14 @@ async def test_music_falls_back_to_configured_id_not_the_class_name() -> None: assert len(entries) == 1 assert entries[0]["model"] == "configured-music-id" - assert entries[0]["model"] != "_UnnamedMusicModel" + assert entries[0]["model"] != type(model).__name__ assert entries[0]["model_id"] == "configured-music-id" -class _UnnamedSoundEffectModel: - async def generate_sound_effect(self, **kwargs: Any) -> Any: - _ = kwargs - return {"not": "a SoundEffectResult"} - - @pytest.mark.asyncio async def test_sound_effect_falls_back_to_configured_id_not_the_class_name() -> None: - tool = SoundEffectToolCore(models={"configured-sfx-id": _UnnamedSoundEffectModel()}) + model = _sound_effect_model(name=_NO_NAME) + tool = SoundEffectToolCore(models={"configured-sfx-id": model}) with TokenContextManager() as manager: await tool.generate_sound_effect(text="p", duration_seconds=4) @@ -310,7 +296,7 @@ async def test_sound_effect_falls_back_to_configured_id_not_the_class_name() -> assert len(entries) == 1 assert entries[0]["model"] == "configured-sfx-id" - assert entries[0]["model"] != "_UnnamedSoundEffectModel" + assert entries[0]["model"] != type(model).__name__ # --- Video billing -------------------------------------------------------- @@ -457,7 +443,9 @@ async def test_asr_bills_when_post_processing_raises() -> None: @pytest.mark.asyncio -async def test_telegram_passes_verbose_and_bills_a_real_duration(tmp_path) -> None: +async def test_telegram_passes_verbose_and_bills_a_real_duration( + tmp_path: Path, +) -> None: """Drive the real Telegram method, not a re-implementation of it. Asserting on a locally-issued transcribe() call would pass with