feat: meter embedding and rerank, and bind usage on standalone entry points by OliverBryant · Pull Request #1457 · xorbitsai/xagent

OliverBryant · 2026-08-18T07:33:51Z

Fourth of five changes splitting #997.

Depends on #1422, #1424 and #1425. This branch contains their
commits too, so the diff here shows all four. Review those first; once
they merge this reduces to just the changes below.

Completes the producer side: embeddings billed per text, rerank per
request.

Rerank was previously unbilled end to end. The RAG search pipeline
reached past the metered adapter to its inner provider
(_extract_dashscope_rerank and its Xinference twin), so metering was
structurally unreachable. Those unwrapping helpers are deleted and
compress_with_scores is declared on the base so callers go through the
adapter. retry_methods also gains compress_with_scores — production
search calls that entry point, so listing only compress left the real
path with no retry at all.

Units follow the modality: embeddings are TEXTS (a batch of 32 is one
provider call but 32 billable texts, and REQUESTS is defined as always
1 per call); rerank is REQUESTS (one call is one billable unit whatever
the document count).

KB ingestion binds its usage through usage_scope / bind_usage_to_thread
(landed in #1425, which needed them first for /speech/transcribe).

Review feedback addressed from #997

The RAG ingestion contextvar leak is fixed (flagged
merge-blocking). The embedding ThreadPoolExecutor bound the caller's
TokenUsage onto the worker with no restore. It now captures on the
calling thread and runs each batch inside copy_context() + ctx.run,
so the write is confined to that call and cannot outlive it on a pooled
thread. Done locally rather than importing the web.tracking helper,
since nothing else under RAG_tools/pipelines depends on xagent.web
and metering is not worth introducing that edge.
The memory store no longer switches embedding models (also flagged
merge-blocking). It pins _DEFAULT_MEMORY_EMBEDDING_MODEL for the
embedding call: existing vectors were written with it, and
schema_migration.py's mismatch check compares only vector presence
and width, so a changed model would silently put new vectors in a
different space with nothing triggering a rebuild. The DB row's name is
threaded through as the new billing_model_name field so usage is
still attributed to the configured row rather than the pin.
The unreachable inner-provider fallback in document_search is now
commented as a defensive branch so it is not mistaken for a live path.

Two comments I did not change code for: the rerank/base.py
synthetic-score default was raised as "flagging only, now documented as
an accepted tradeoff", and the code already carries that rationale; and
the model_service.py is_active predicate is unrelated to metering, so
it is split into #1458 with the test coverage that comment asked for.

Testing

CI. Local runs are not usable in my environment — importing
xagent.core.model.chat hangs on this machine, on main as well as on
these branches — so I am relying on CI rather than claiming a local pass
I did not get.

Introduce the shared vocabulary and recording helpers for non-LLM media usage (image/video/audio/embedding/rerank), with no producers wired up yet — those follow in separate changes. - `MediaUnit` / `MediaCallType` name the billable dimension and the modality that produced an entry. The unit is a property of the modality, never of the response: a duration-billed call records `seconds` with `quantity=0` when unmeasured rather than switching to `requests`, so a price table keyed on (model, unit) stays usable. - `TokenUsage.add_media_usage` / `add_media_usage` write `type:"media"` detail entries into the existing `TokenUsage.details` list, so media flows through DB persistence and the quota `delta_details` contract without special-casing. Provider tokens are stored under `provider_tokens`, not `tokens`, so a consumer summing billable LLM tokens cannot pick up media counts. - `aggregate_media_usage_by_model` groups by (model, unit, call_type, resolution). Zero-quantity entries are kept deliberately: they are the only evidence that an unmeasured provider call happened. - `media_usage.py` adds `resolve_billing_model`, which never records a placeholder identity such as "default" or "None", plus the error-swallowing `record_media_usage` / `record_media_seconds` wrappers producers will use. Unknown `unit`/`call_type` values are now rejected at the write boundary. A typo would otherwise mint a new billing dimension that the aggregator keys off, and a usage record cannot be repaired once persisted. Validation runs before any mutation so a rejected call leaves no counter incremented without a matching detail entry. `TokenUsage` counter updates are now guarded by a lock. A single instance is shared across worker threads, where `+=` is a read-modify-write that silently loses counts; a 8x200 concurrent recording test drops ~85% of them without this. `merge()` snapshots the source under its own lock before taking the target's, so concurrent merges in both directions cannot deadlock.

gemini-code-assist

Code Review

This pull request introduces a comprehensive and thread-safe media-usage tracking system for non-LLM modalities such as image, video, TTS, ASR, embedding, and rerank. It adds new enums, thread locks to TokenUsage, and context managers to handle thread boundaries and background tasks. The review feedback suggests enhancing robustness against null values, specifically by safely retrieving usageMetadata from Gemini API responses to prevent AttributeErrors, and avoiding direct str() conversions on potentially null fields in the memory store to prevent them from being coerced into the literal string 'None'.

gemini-code-assist · 2026-08-18T07:37:03Z

+                            id=str(getattr(embedding_model, "model_id", "") or ""),
+                            model_name=_DEFAULT_MEMORY_EMBEDDING_MODEL,
+                            billing_model_name=configured_name or None,
+                            model_provider="dashscope",
+                            api_key=str(embedding_model.api_key),
+                            dimension=int(embedding_model.dimension or 1024),


Avoid converting potentially null values directly to strings using str() without checking for None first, as this converts None to the literal string 'None'. When multiple sites require the same coercion, use a single unified helper function (such as safe_str) to prevent logic drift and ensure consistent behavior.

Suggested change

id=str(getattr(embedding_model, "model_id", "") or ""),

model_name=_DEFAULT_MEMORY_EMBEDDING_MODEL,

billing_model_name=configured_name or None,

model_provider="dashscope",

api_key=str(embedding_model.api_key),

dimension=int(embedding_model.dimension or 1024),

id=safe_str(getattr(embedding_model, "model_id", None)),

model_name=_DEFAULT_MEMORY_EMBEDDING_MODEL,

billing_model_name=configured_name or None,

model_provider="dashscope",

api_key=safe_str(embedding_model.api_key),

dimension=int(embedding_model.dimension or 1024),

References

Avoid converting potentially null values directly to strings using str() without checking for None first, as this converts None to the literal string 'None'. When multiple sites require the same coercion, use a single unified helper function to prevent logic drift and ensure consistent behavior.

gemini-code-assist · 2026-08-18T07:37:03Z

+            # empty parts). Recording after the structural checks below would
+            # drop usage for exactly those billed-but-unusable responses, and
+            # retry_on only matches 429/5xx so they are never retried either.
+            usage_metadata = response_data.get("usageMetadata", {})


If the API response contains "usageMetadata": null, response_data.get("usageMetadata", {}) will return None instead of {} because the key exists. This will subsequently cause an AttributeError when calling .get() on usage_metadata. Use response_data.get("usageMetadata") or {} to safely handle both missing and null values.

Suggested change

usage_metadata = response_data.get("usageMetadata", {})

usage_metadata = response_data.get("usageMetadata") or {}

References

When retrieving values from a dictionary where a key might be present but explicitly set to None, use element.get('key') or 'default' instead of element.get('key', 'default') to ensure the default value is used for None values as well, preventing downstream type errors.

gemini-code-assist · 2026-08-18T07:37:03Z

+            # Meter before validating the response body — see generate_image
+            # for why a billed 200 must be recorded ahead of the structural
+            # checks below.
+            usage_metadata = response_data.get("usageMetadata", {})


If the API response contains "usageMetadata": null, response_data.get("usageMetadata", {}) will return None instead of {} because the key exists. This will subsequently cause an AttributeError when calling .get() on usage_metadata. Use response_data.get("usageMetadata") or {} to safely handle both missing and null values.

Suggested change

usage_metadata = response_data.get("usageMetadata", {})

usage_metadata = response_data.get("usageMetadata") or {}

References

When retrieving values from a dictionary where a key might be present but explicitly set to None, use element.get('key') or 'default' instead of element.get('key', 'default') to ensure the default value is used for None values as well, preventing downstream type errors.

Addresses review findings N1, N2, N3, N9, N11, N12, N13, N14 and N15 on the media billing primitives. Invalid quantities can no longer be persisted. `_coerce_float` rejected only non-numeric input, so booleans, negatives, NaN and infinities all reached the detail row — and NaN/inf are not JSON-serialisable, so they would emit literal `NaN`/`Infinity` tokens into task and quota details that strict parsers reject. All four now record 0.0 with a warning, matching the existing "call happened but is unmeasured" convention. `coerce_duration` gained the same finiteness guard, since `inf > 0` is True and would otherwise pass as a usable duration. Validation moved to the real write boundary. It previously lived only in the module-level `add_media_usage`, so anyone holding a `TokenUsage` bypassed it. Both writes now go through one path: the new atomic `record_media_call`, with the older `add_media_usage` method delegating to it. Self-review caught that these two copies had already drifted — the method appended unvalidated quantities that the function rejected. The counter and its detail row are now written under a single lock acquisition. Taking the lock twice let a concurrent `to_dict`/`merge` observe a count with no matching detail, and consumers derive per-model rows from `details` while reading the call count from the scalar, so a torn snapshot reports mutually inconsistent billing. `media_calls` now survives TaskTracker snapshots. `_copy_usage` and `_task_seed_from_session` both reconstructed `TokenUsage` without it, so periodic and final snapshots reported zero while the detail rows survived. The seed derives the count from the surviving media rows, since tasks have no `media_calls` column and seeding zero would re-report every media call from a prior turn. The mutation lock is no longer a dataclass field. As a field it sat in `__dataclass_fields__`, which `dataclasses.asdict` walks directly while ignoring `__getstate__`, so any generic consumer hit `TypeError: cannot pickle '_thread.lock' object`. It now lives in `__dict__` behind a property, with `__getstate__`/`__setstate__` dropping and recreating it, which also keeps the private slot out of the public positional signature. asdict, deepcopy and pickle all work. `record_media_usage` mirrors the core optional fields (`resolution`, `input_tokens`, `output_tokens`, `tokens_estimated`) and forwards them inside the guarded call; passing them previously raised `TypeError` during argument binding, before the `try` block. `resolve_billing_model` guards each provider attribute read individually, since it runs before that `try` and a `model_name` property that raises would have broken the user's media call over an accounting lookup. `estimate_tokens` accepts any iterable of strings as documented — generators previously estimated 0 — and rounds Latin character counts up, since truncating division estimated 0 tokens for any 1-3 character string. Empty input still estimates 0. Tests cover each fix and were written to fail on revert: invalid-quantity boundaries, a writer/reader interleaving assertion that every snapshot has `media_calls` equal to its media detail count, asdict/deepcopy/pickle round-trips, positional construction, opposing concurrent merges with a timeout, wrapper exception paths with a monkeypatched raising backend, a raising-descriptor resolver case, and generator/short-text token estimation. Deferred with issues rather than guessed at: xorbitsai#1460 covers the aggregate identity tagging, resolver return shape, placeholder fallback and input/output token split (all latent until producers mix the forms), and xorbitsai#1461 covers the discriminated `delta_details` union, which depends on an out-of-tree quota hook whose shape is not visible from this repository — no in-tree consumer indexes `detail["tokens"]`.

Three errors from the previous commit: `record_media_call` now accepts `MediaUnit | str | None` and `MediaCallType | str | None`, matching its callers. It validates and narrows internally, so the narrow `str` annotation rejected the very values `add_media_usage` forwards. The validated results bind to separate `unit_value`/`call_type_value` locals rather than being assigned back over the widened parameters. The `_lock` property binds through a locally annotated variable, since reading it out of `__dict__` yields `Any` and returning that from a function declared to return `LockType` trips `no-any-return`. Verified with mypy 1.19.0 (the pinned pre-commit version) on the changed file: no issues.

…tional ABI Second review round on the media billing primitives: C1, C4, C5, P7, P12, P15 and P16. C1 — MediaUnit.REQUESTS now rejects any quantity but 1. It is defined as exactly one provider call, so its quantity is not a free variable; letting 0/0.5/2 through made the billable quantity disagree with the media_calls and aggregate `calls` count derived from the same row, so quota and pricing would read different numbers off one record. Raised rather than clamped — a caller passing something else has a real bug, and silently rewriting it hides that — and rejected before any mutation, so no counter is bumped without its detail row. `record_media_usage` swallows it, keeping the promise that accounting never breaks a media call. C4 — `_coerce_int` is now a safe token boundary. It accepted booleans, preserved negatives, and let `int(float("inf"))` raise an uncaught `OverflowError` that propagated out of the accounting path and dropped the whole billable media row rather than just the bad field. It now mirrors `_coerce_float`: booleans, negatives and non-finite values sanitize to 0 with a warning while the row survives. Fixed on the shared helper rather than a media-only copy, since the LLM token paths pass provider-reported values through the same function and had the same overflow exposure. P12 — the positional ABI is restored. `media_calls` was inserted before the legacy `tool_calls`/`details` fields, so an existing `TokenUsage(input, output, llm_calls, tool_calls, details)` call bound its fourth argument to `media_calls` and shifted the rest, silently and with no error. The field is now appended after `details` and keyword-only. The test that previously asserted the broken order is rewritten to guard the legacy binding instead — it was institutionalising the bug. P7 — `TokenUsage.snapshot()` takes the lock and returns a detached copy; `TaskTracker._copy_usage` delegates to it. Reading the fields one by one from outside could interleave with a concurrent write and yield a counter without its matching detail row even though the writer is atomic. P16 — `estimate_tokens` catches any exception from iteration, not just `TypeError`. A custom iterable raising `ValueError`/`RuntimeError` escaped into an accounting path whose documented contract is that malformed input never breaks the call being measured. C5 — added a TaskTracker seam test that seeds prior-turn media rows, records a current-turn row, and asserts the seeded baseline is counted, the periodic snapshot carries the running total, and both the mid-run gate and the completion hook receive only this turn's row. P15 — the opposing-merge test now forces overlap with a `Barrier` and asserts exact counter/row agreement plus per-unit row survival, instead of `>= 100` lower bounds that would pass while rows were dropped. C2 (unbounded details persistence, O(n²) as high-frequency producers land) is deferred to xorbitsai#1466: the copy-and-rewrite mechanism predates this work, no producer is wired here, and changing the persistence shape is a much larger change than this PR. Self-review caught two things while making the above: the legacy two-step API leaves an orphan counter when a REQUESTS rejection happens (the caller owns that increment) — documented at the raise site, and unreachable since no in-tree producer uses the two-step path; and a comment describing a decrement that no longer exists.

…aint `test_dirty_quantity_is_coerced_and_does_not_raise` used the `requests` unit with `None`/`"oops"` quantities and asserted both coerce to 0. That is now a contradiction: REQUESTS pins its quantity to exactly 1, so those inputs are rejected rather than coerced. Split into the two behaviours the constraint creates, which is the replacement the review asked for: - free-quantity units (`seconds`) still coerce a malformed quantity to 0 and keep the row, because the provider call happened and must be visible as unmeasured; - `requests` rejects the same inputs and leaves no state behind, because coercing to 0 would record a call whose billable quantity contradicts its own call count. Checked the other `requests` usages in this file: both pass `quantity=1` and are unaffected.

…media API Third review round. NEW-A was blocking and was a regression I introduced last round; the rest are hardening plus the two simplifications the review recommended. NEW-A — `_turn_delta` now snapshots before reading. It read the live shared `TokenUsage` twice — the details slice, then the tool-call counter — with no lock, so a concurrent `record_media_call`/`increment_tool_calls` landing between them returned a pair describing two different instants, and that inconsistent pair fed quota gating. `interrupt_reason_for_quota` polls it at every safe point, so it is a hot concurrent path. This is the sibling of `_copy_usage`, which I re-pointed through `snapshot()` last round for exactly this reason — I fixed one call site and did not check the other in the same file. The snapshot is now taken unconditionally, including when the caller already passed a detached copy, so the two sites cannot diverge again. Reproduced with a forced interleaving: 200/200 inconsistent before, 0/200 after. Deleted the two-step media API — `TokenUsage.add_media_usage`, the `count_call` parameter, and `increment_media_calls`. It existed only to preserve a legacy flow with no production callers (grep-confirmed; only tests used it), and its documented behaviour was that a rejected call leaves `media_calls` incremented with no matching detail row — the exact torn state the REQUESTS constraint was added to prevent. Deleting removes the defect rather than hardening around it, and leaves `record_media_call` plus the module-level `add_media_usage`/`record_media_usage` as the only entry points. Tests migrated to the atomic call. `copy.copy` no longer shares state under a split lock. Shallow copy routed through `__getstate__`/`__setstate__`, giving the clone a fresh lock while still referencing the same `details` list — two objects each believing they held exclusive access to one list. `__copy__`/`__deepcopy__` now delegate to `snapshot()`, which is the supported way to fork a usage object. `to_dict` deep-copies each detail entry. It took the lock but returned a new outer list around the same inner dicts, so a caller mutating the payload mutated the live usage object outside that lock. Counter increments ignore negative counts, consistent with the coercion helpers; a negative would drive a counter below its matching row count. `record_media_usage` logs with `exc_info=True`. The broad catch is deliberate, but an unexpected bug surfaced as one context-free line, which is the wrong trade when persisted billing rows are unrepairable. Self-review found the same aliasing NEW-C describes in `merge`, which copied the outer list but shared inner dicts — mutating one usage object would silently rewrite the other's rows. Fixed with the rest. NEW-F (zero quantity overloading unmeasured and corrupt) is deferred to xorbitsai#1495: the current behaviour is deliberate and asserted by an existing test, and separating the two touches the persisted row shape, so it belongs with the first pricing consumer. C2 remains tracked in xorbitsai#1466.

…ble to fail Fourth review round. Both blocking findings were consequences of how I fixed NEW-A last round. N1 — `_turn_delta` no longer deep-copies the whole details list. `usage.snapshot()` copied the entire cumulative list under the lock and then discarded everything before `_initial_details_len`. The list grows monotonically across turns (seeds restore the persisted list), and the path is polled once per agent step *and* once per streamed LLM chunk (`runtime.py` `_raise_if_interrupted` -> `interrupt_reason_for_quota`), while holding the lock that serialises every LLM adapter's token write. Measured: 10.6us at 100 rows, 103us at 1000, 487us at 5000 — versus 0.3us for the tail. New `TokenUsage.detail_tail(start)` returns the tail plus `tool_calls` in one lock acquisition, preserving exactly the atomicity the snapshot was added for, at O(delta) instead of O(total). The now-redundant `_copy_details(delta_details)` at both quota call sites is gone, since the rows come back detached. N2 — the regression test for that fix could not fail. `read_delta` called `_turn_delta()` with no argument, so it went through `get_token_usage()` inside a bare `threading.Thread`, which does not inherit contextvars: it lazily built a fresh empty `TokenUsage` disconnected from the object the writer mutated, and observed `([], 0)` either way. Against the pre-fix code `snapshot()` was never called at all, so the patch never fired, `entered` was never set, and the writer's `wait()` timed out silently with its return value unchecked. The single assertion degenerated to `0 == 0` on both sides. `usage` is now passed explicitly, and the test asserts `entered.is_set()` and `released.is_set()` so a silent timeout fails. Mutation-verified this time against the committed test rather than a side script: 50/50 failures with two unlocked reads, 0/50 with `detail_tail`. That distinction is the gap that let the empty test through last round. Also addressed from this round: the `REQUESTS` error message now reports the caller's raw value rather than the coerced one (N7); `_lock` is an `RLock` so a future nested acquisition cannot self-deadlock (N5); `__post_init__` normalises `details`, covering the constructor and `from_dict` together (N3); `copy_detail_rows` replaces the same copy expression reimplemented four times, removing the filter asymmetry where `to_dict` would raise on a non-dict row that the others skipped; `estimate_tokens` is renamed `estimate_media_tokens` and exported, since three unrelated `estimate_tokens` already exist with different algorithms (N9); its CJK ranges gain punctuation, fullwidth forms and Ext-A (N10); the module docstring, the `add_media_usage` "Raises" section, the aggregators' detached-list contract, and the details-are-dicts invariant are documented (N13, N8, N14, N12); and three false or contradictory comments are corrected (N11, N16, plus the stale seed rationale). Test gaps closed: `_coerce_int`'s hardening exercised through the LLM `add_token_usage` path (N6), the `model_id` branch of `aggregate_media_usage_by_model` (N17), and `test_concurrent_merge_loses_no_counts` rewritten with a barrier and 8x100 merges — one merge per thread of a one-row source would have passed against an unlocked `merge` (N15). Self-review caught that adding normalisation to `__post_init__` made `snapshot()` copy twice, doubling it from 487us to 1066us at 5000 rows; `snapshot` now assigns `details` after construction, back to 526us. mypy also rejected `threading.RLock` as an annotation (it is a factory, not a class), so the property is typed through a `TYPE_CHECKING` alias. NEW-F stays in xorbitsai#1495, C2 in xorbitsai#1466, and P4/P5/P8/P9/P10 in xorbitsai#1460/xorbitsai#1461. N4 (string fields reaching the JSON column uncoerced) and N19 (`media_calls` has no consumer yet) are noted as accepted for this layer: the producers that supply those strings land in xorbitsai#1424/xorbitsai#1425/xorbitsai#1457, and xorbitsai#997 names the `media_calls` consumer.

… real Fifth review round. Three of the five blocking findings were about claims I made rather than code I wrote, so each is verified by mutation against the committed test this time, not a side script. Locking tests now require the lock (finding 2). Replacing `_lock` with `nullcontext()` left all six concurrency tests passing. Investigating why showed my earlier "~85% count loss" figure was wrong: it came from a mutation that inserted `sleep(0)` between the read and the write, not from the real implementation. On CPython 3.12 even a bare `self.n += 1` loses nothing at 8x200000 threads-iterations, so no increment-counting test can justify the lock. What the lock actually protects is the *pairing* of the counter with its detail rows: `record_media_call` bumps and appends as one operation, and an unlocked reader lands between them. The test now asserts that invariant under `sys.setswitchinterval(1e-9)` (required — at the default interval the window is unhittable). Mutation-verified: 6/6 runs tear without the lock, 0/6 with it. `detail_tail`'s atomicity test now exercises the interleaving (finding 3, the second round this test was shown not to test its claim). Both earlier versions wrapped `detail_tail` itself and released the writer *before* delegating to the real accessor, so the write always completed before the locked read began. The interleaving is now forced from inside the lock, by patching `copy_detail_rows` — called after the lock is taken and `details` read, but before `tool_calls`. Mutation-verified against the reviewer's own torn implementation (two acquisitions with a sleep between): it fails with media=0, tool_calls=1. The unit-per-modality invariant is enforced, not just documented (finding 5). `MEDIA_UNIT_BY_CALL_TYPE` rejects a unit that does not match the call type at the write boundary. The invariant was stated in three docstrings and violated by this PR's own tests; the reviewer found six sites, and the new check found four more (`seconds` paired with `tts`) that a grep for `images`+`video` missed. All ten fixed. Verified every downstream producer in xorbitsai#1424/xorbitsai#1425/xorbitsai#1457 pairs legally, so the new constraint breaks none of them. `_coerce_float` catches `OverflowError` (finding 1). `float(10**400)` raised it uncaught, so the error-swallowing wrapper dropped the entire billing row — the opposite of the documented reject-to-0.0 contract. `_coerce_int` was already fixed for this in an earlier round; the sibling was missed. The token-vs-unit pricing precedence claim is removed rather than implemented (finding 4). It existed only as a docstring sentence: nothing expressed or enforced "price by tokens instead of by unit", the row schema carries no discriminator, and `provider_tokens` has zero consumers outside this file. Implementing it would add a billing dimension to a PR that wires no producers and has no pricing consumer, so the aspirational sentence is gone and the design question stays tracked in xorbitsai#1461. Minor: `add_media_usage` reports the caller's raw quantity in the REQUESTS error (the prior fix covered only the direct path); the wrapper's warning log now includes model/model_id/unit/quantity, not just call_type; the two stale `bind_usage_to_thread` references are dropped (the symbol does not exist); `_turn_delta`'s return type is parameterised. Simplifications taken: `TypeGuard` from stdlib `typing` rather than `typing_extensions` (the project requires >=3.11), and `task_tracker`'s `_copy_details` is deleted in favour of the `copy_detail_rows` helper this PR added for exactly that consolidation. Not changed, with reasons: `resolve_billing_model` and `record_media_seconds` are flagged as having no production caller, but both are consumed by xorbitsai#1425/xorbitsai#1457 in this same series — removing them here would only move the diff. Findings 7-9 and 11-16 are noted as accepted or tracked in the reply.

Wire the image providers into the media usage primitives so image generation and editing are billed alongside LLM tokens. `record_image_usage` normalises each provider's usage payload into a single `add_media_usage` call, billed in `MediaUnit.IMAGES` with the request's `n` as the quantity so multi-image requests are not under-billed as one. The resolution tier is recorded so billing can price by (model, resolution); providers that also report real tokens (Gemini, OpenAI gpt-image) pass them through so a token-based price can take precedence. Gemini and DashScope now record usage immediately after the usage payload is parsed, before the response body is validated. Both providers raise structural RuntimeErrors on a 200 the provider has already billed — a safety-blocked finish reason or a missing image field — and recording after those checks dropped usage for exactly those responses. Gemini's retry_on matches only 429/5xx, so they are never retried either. Applies to both generate_image and edit_image on both providers. The Xinference inpainting branch now forwards n and size. It previously dropped both while usage was still recorded as n, over-billing whenever n>1. Unreachable today because **kwargs is stripped from the exposed tool schema, but it would have gone live silently the moment that changed.

The previous commit applied `image_count=kwargs.get("n", 1)` uniformly to Gemini and DashScope, but the two clients treat `n` differently, so one of them was over-billing. Gemini reverts to a fixed count of 1. The Gemini API has no multi-image parameter, and this client forwards only `temperature` out of `**kwargs` into `generationConfig` — a caller-supplied `n` is silently dropped. The response parser also stops at the first `inlineData` part. Billing `n` there charged for images the provider never generated. DashScope keeps `n`. It spreads `**kwargs` into the request `parameters`, so `n` does reach the provider, which generates and bills for that many images; the recorded count has to match the invoice. Images 2..n are then discarded by the parser, which reads only `content[0]` — real, but a separate defect from metering, and noted in a comment rather than silently under-billed around. Both behaviours are now pinned by tests: the Gemini case asserts `n` never reaches the request body and that the billed quantity stays 1, and the DashScope case asserts `n` does reach `parameters` and that the billed quantity follows it.

Two tests asserted that Gemini bills the caller's `n`. It cannot: both `generate_image` and `edit_image` build `generationConfig` without any multi-image parameter — `generate_image` forwards only `temperature` out of `**kwargs` — so `n` is dropped before the request goes out and the parser returns a single image. Billing it charges for images the provider was never asked to produce. `test_gemini_edit_forwards_image_count` is replaced by `test_gemini_edit_does_not_bill_unsupported_n`, which asserts `n` is absent from the outgoing request body and that `image_count` is therefore not passed at all. The safety-blocked case now expects `quantity == 1` rather than 3, for the same reason. The DashScope tests are unchanged and still expect the caller's `n`: that client spreads `**kwargs` into the request `parameters`, so `n` does reach the provider and does get billed by it.

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.

`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.

…points Completes the producer side: embeddings billed per text, rerank per request, and a shared sink so the paths that are not tracked agent tasks report their usage at all. `standalone_usage.usage_scope` binds a TokenUsage and reports it to the quota hook on exit. Without it, KB ingestion, `/speech/transcribe` and Telegram voice recorded into a throwaway object that nothing read — the provider billed and the usage evaporated. Rerank was previously unbilled end to end. The RAG search pipeline reached past the metered adapter to its inner provider (`_extract_dashscope_rerank` and its Xinference twin), so metering was structurally unreachable; those unwrapping helpers are deleted and `compress_with_scores` is declared on the base so callers go through the adapter. `retry_methods` also gains `compress_with_scores`: production search calls that entry point, so listing only `compress` left the real path with no retry. Embeddings use unit=TEXTS, not REQUESTS — a batch of 32 texts is one provider call but 32 billable texts, and REQUESTS is defined as always 1 per call. Rerank is the opposite: one call is one billable unit whatever the document count, so REQUESTS is right there. Review feedback from xorbitsai#997: - `bind_usage_to_thread` now runs each call inside `copy_context()` + `ctx.run` instead of writing the contextvar directly on the worker. `run_in_executor(None, ...)` uses the loop's long-lived default executor, so the old binding outlived the job and leaked the caller's TokenUsage into the next unrelated task on that thread — cross-tenant misattribution plus a reference pinned for the thread's life. The RAG ingestion pool had the same pattern and is fixed the same way, locally rather than importing the web helper, since nothing else under RAG_tools/pipelines depends on xagent.web. - `_report` checks `has_usage_record_hook()` before checking out a DB session. With no hook installed — the stock configuration — every ingest and transcription paid a pool checkout, transaction and close for a guaranteed no-op. - `usage_scope` only restores the previous context if it still owns the binding, so a TaskTracker started inside the scope is not silently detached on exit. - The memory store no longer switches embedding models. It pins `_DEFAULT_MEMORY_EMBEDDING_MODEL` for the embedding call, because existing memory vectors were written with it and `schema_migration.py`'s mismatch check compares only vector presence and width — a changed model would put new vectors in a different space with nothing triggering a rebuild. The DB row's name is threaded through as `billing_model_name` so usage is still attributed to the configured row. - The unreachable inner-provider fallback in `document_search` is now commented as a defensive branch rather than reading like a live path.

OliverBryant added 2 commits August 17, 2026 16:26

fix: normalize optional media metadata

979c051

XprobeBot added the feature label Aug 18, 2026

gemini-code-assist Bot reviewed Aug 18, 2026

View reviewed changes

OliverBryant force-pushed the feat/embedding-rerank-metering branch 2 times, most recently from 074de16 to aab2e08 Compare August 18, 2026 07:42

OliverBryant mentioned this pull request Aug 18, 2026

fix: skip inactive shared defaults for music and sound-effect models #1458

Merged

OliverBryant force-pushed the feat/embedding-rerank-metering branch from aab2e08 to 3997ac6 Compare August 18, 2026 07:48

OliverBryant mentioned this pull request Aug 18, 2026

Define a discriminated delta_details union before external quota consumers see media entries #1461

Open

OliverBryant added 2 commits August 18, 2026 16:11

OliverBryant force-pushed the feat/embedding-rerank-metering branch from 3997ac6 to e48758f Compare August 18, 2026 08:25

OliverBryant mentioned this pull request Aug 18, 2026

feat: surface media usage in the task API and chat UI #1463

Merged

OliverBryant mentioned this pull request Aug 18, 2026

feat: add media usage billing primitives #1422

Closed

OliverBryant force-pushed the feat/embedding-rerank-metering branch from e48758f to 933d12e Compare August 18, 2026 10:28

OliverBryant force-pushed the feat/embedding-rerank-metering branch from 933d12e to 9c29b0b Compare August 19, 2026 04:37

OliverBryant force-pushed the feat/embedding-rerank-metering branch from 9c29b0b to a9671b0 Compare August 19, 2026 07:41

OliverBryant force-pushed the feat/embedding-rerank-metering branch from a9671b0 to a59361b Compare August 19, 2026 07:48

OliverBryant force-pushed the feat/embedding-rerank-metering branch 2 times, most recently from d2d09ee to 5837e71 Compare August 19, 2026 08:26

OliverBryant added 4 commits August 19, 2026 18:44

fix: correct image usage accounting

f41ee06

OliverBryant added 4 commits August 19, 2026 18:44

OliverBryant force-pushed the feat/embedding-rerank-metering branch from 5837e71 to 26b1d61 Compare August 19, 2026 10:46

OliverBryant mentioned this pull request Aug 20, 2026

Media usage billing primitives: scope, and lessons from the failed #1422 attempt #1526

Open

rogercloud mentioned this pull request Aug 20, 2026

feat: add media usage billing vocabulary and recording helpers #1527

Merged

OliverBryant mentioned this pull request Aug 21, 2026

feat: meter image generation and editing #1424

Open

Navigation Menu

Uh oh!

feat: meter embedding and rerank, and bind usage on standalone entry points - #1457

feat: meter embedding and rerank, and bind usage on standalone entry points#1457
OliverBryant wants to merge 16 commits into
xorbitsai:mainfrom
OliverBryant:feat/embedding-rerank-metering

OliverBryant commented Aug 18, 2026 •

edited

Loading

Uh oh!

gemini-code-assist Bot left a comment

Uh oh!

gemini-code-assist Bot Aug 18, 2026

Uh oh!

gemini-code-assist Bot Aug 18, 2026

Uh oh!

gemini-code-assist Bot Aug 18, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

	usage_metadata = response_data.get("usageMetadata", {})
	usage_metadata = response_data.get("usageMetadata") or {}

Uh oh!

Conversation

OliverBryant commented Aug 18, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Review feedback addressed from #997

Testing

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist Bot Aug 18, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Aug 18, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Aug 18, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

OliverBryant commented Aug 18, 2026 •

edited

Loading