feat: add media usage billing primitives - #1422
Conversation
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.
There was a problem hiding this comment.
Code Review
This pull request introduces tracking, validation, and aggregation for non-LLM media usage (such as images, video, TTS, ASR, embeddings, and reranking) alongside existing LLM token tracking. It also adds thread-safety to the TokenUsage class using a lock to prevent race conditions during concurrent updates. The review feedback suggests explicitly handling None values in the validation helpers _validated_media_call_type and _validated_media_unit to prevent them from being converted to the literal string 'None', which would trigger validation errors.
rogercloud
left a comment
There was a problem hiding this comment.
This PR adds the MediaUnit/MediaCallType vocabulary, media-detail aggregation, a billing-model resolver, and best-effort recording wrappers for future image/audio/embedding/rerank producers. It extends TokenUsage with media_calls and locking and sends media entries through TokenUsage.details, allowing the existing TaskTracker persistence and quota delta_details path to carry later producer records. That moves the billing boundary from token-only entries to a mixed detail stream, so numeric validity, union compatibility, identity, and snapshot state must be closed before those producers rely on it.
Blocking: yes — recommended event: REQUEST_CHANGES
Approach / design verdict
acceptable-with-reservations. Reusing TokenUsage.details and the existing TaskTracker/quota path is the right integration seam, and the shared primitive is a reasonable foundation for follow-up producers. The design is not closed yet: the numeric write boundary accepts invalid billing values, the mixed media/token detail union is not compatible with existing quota consumers, unresolved model identity can fall back to a placeholder and the resolver API does not preserve name-vs-ID semantics, and the new counter/detail state is not atomically snapshotted or propagated through TaskTracker.
Prior review status
- P1 — FIXED:
call_type=Noneis normalized to an empty string before stringification. Original inline comment · review - P2 — FIXED:
unit=Noneis rejected explicitly without mutating usage. Original inline comment · review
These prior roots are not repeated as new findings below.
New findings — major
-
N3 — major —
src/xagent/core/model/chat/token_context.py:726._coerce_floataccepts booleans, negative values,NaN, infinities, and float-coercible strings, and the write path stores the result;record_media_secondscan also forward negative or non-finite values throughseconds or 0.0. A malformed quantity from a future producer can therefore persist invalid billing data, withInfinityreaching serialized task/quota details. Validate finite, non-negative, unit-specific values before incrementing or appending (while keeping an absent duration asseconds=0), and add boundary tests for booleans, negatives,NaN, infinities, bad strings, and duration forwarding. -
N6 — major —
src/xagent/core/model/chat/token_context.py:204. Media entries intentionally omit the legacytokenskey but are mixed into the same details list; existing quota hooks assume token entries and indexdetail["tokens"]. Any task whose delta contains a media entry can make progress handling fail open and cause completion recording to log/drop usage. Define and migrate a discriminateddelta_detailsunion, or filter unknown types/capability-gate legacy hooks before mixed entries reach them, and add an integration test containing token and media entries that proves both progress and completion accounting. -
N8a — major —
src/xagent/core/tools/core/media_usage.py:78.resolve_billing_modelreturns the hard-coded"default"fallback even though the module invariant says placeholders must never be recorded. When the configured ID and provider attributes are absent or placeholders, a future producer creates a fakedefaultbilling row instead of representing an unresolved identity. Return an optional/structured unresolved result and let producers skip/log it; if a fallback is retained, validate it and prohibit placeholders, and update the test that currently institutionalizes"default".
New findings — minor
The following contract risks are latent where noted because this PR wires no production media producers; they become observable when follow-up producers consume these helpers.
-
N1 — minor —
src/xagent/core/model/chat/token_context.py:106. Addingmedia_callsdoes not updateTaskTracker._copy_usage(src/xagent/web/tracking/task_tracker.py:59-66) or_task_seed_from_session(:105-110). After a media call, periodic/final snapshots and the next-turn seed reportmedia_calls=0even though the detail rows survive. Thread the scalar through copy/seed/persistence paths and add a regression test covering periodic, final, and next-turn usage. -
N2 — minor —
src/xagent/core/model/chat/token_context.py:737.increment_media_calls()andadd_media_usage()acquire separate locks, so another thread can snapshot or merge the object between them and observe a count/detail mismatch; TaskTracker/delta readers also read fields outside the lock. Replace these calls with one atomicrecord_media_usageoperation and use a lock-protected snapshot for tracker/delta consumers. Add an interleaving test that asserts the counter equals the number of media details in every snapshot. -
N5 — minor —
src/xagent/core/model/chat/token_context.py:579. Raw media details retainprovider_input_tokensandprovider_output_tokens, but the public aggregate only sumsprovider_tokens. Once a consumer has distinct input/output rates, it cannot apply them through this aggregate and may under- or over-bill. Preserve both split fields (or an explicit split structure) in aggregate output and test unequal input/output counts; there is no current pricing consumer in this PR. -
N7 — minor —
src/xagent/core/model/chat/token_context.py:593.identity = model_id or model_nameuses one untagged key, so an ID-backed row withmodel_id="foo"can merge with an unrelated name-only row withmodel="foo". When future producers mix these forms, distinct physical models can be billed as one. Use a tagged/structured identity key and add a collision test; no producer currently exercises the path. -
N8b — minor —
src/xagent/core/tools/core/media_usage.py:61. The resolver returns one untaggedstr, but that value is a configured ID in one branch and a provider name in another, while the recording contract requiresmodel=nameandmodel_id=id. A natural future caller can put an ID inmodeland leavemodel_idempty, splitting or colliding aggregate rows. Return a structured pair with explicit name and ID fields (or two explicit APIs), and test both configured-ID and provider-name fallback cases. -
N9 — minor —
src/xagent/core/tools/core/media_usage.py:98.record_media_usagedoes not accept or forwardresolution,input_tokens,output_tokens, ortokens_estimated, although the core primitive and its documentation expose them. A future producer either loses billing metadata or passes one of these arguments and gets aTypeErrorduring argument binding before the wrapper'stryblock. Mirror the core optional fields, forward them inside the guarded call, and add a forwarding test; no production caller exists yet. -
N11 — minor —
src/xagent/core/model/chat/token_context.py:106. Insertingmedia_callsbefore legacytool_calls/detailssilently remaps positionalTokenUsage(input, output, llm, tool, details)calls, while the init-visible_lockat:114-116can break genericdataclasses.asdict,deepcopy, or pickle consumers. Append the public field or make it keyword-only, and keep lock state out of the public constructor/serialization with explicit state handling; add a compatibility test or document the intentional public break. -
N12 — minor —
src/xagent/core/tools/core/media_usage.py:75.getattr(model, attr, None)is evaluated beforerecord_media_usage's best-efforttryblock. A provider object whosemodel_nameormodeldescriptor raises can therefore break the user call despite the module promise. Guard each attribute access, continue to the next candidate, and add a raising-descriptor test; no production resolver caller exists yet. -
N13 — minor test gap —
tests/core/tools/core/test_media_usage_helpers.py:66.quantity=Noneis coerced to0with valid unit/call-type metadata, so the test never enters the wrapper's exception path and would still pass if thetry/exceptwere removed. Add invalid-unit/call-type cases and a monkeypatchedadd_media_usagethat raises, asserting no propagation and no state mutation; those tests would catch a regression in the best-effort guarantee. -
N14 — minor test gap —
tests/core/model/chat/test_media_usage.py:340. The fan-in test only exercisestarget.merge(source)and has no timeout or opposite-directiona.merge(b)/b.merge(a)scenario, so it cannot catch a deadlock or snapshot-loss regression despite the implementation's stated contract. Add concurrent reverse-direction merges with a timeout and count/detail assertions; the test should fail on lock-order or incomplete-snapshot regressions. -
N15 — minor —
src/xagent/core/model/chat/token_context.py:382.estimate_tokensdocuments an iterable of strings but accepts only list/tuple/set, so a generator returns0; additionally,other // 4makes non-empty one-to-three-character Latin input estimate to zero. A future estimator-based producer can therefore silently undercount. Support the documented iterable (including generators) or narrow the contract, define the minimum/zero semantics, and add generator and short-text boundary tests; no current producer calls this helper.
No Simplification opportunities section is included because no Spark finding was confirmed.
Verification
No local tests were run under the PR-review workflow. CI preflight had 14 completed successful checks and no failures; this review does not claim that tests passed locally.
Blocking status & recommended decision
Blocking: yes
src/xagent/core/model/chat/token_context.py:726— major invalid/non-finite quantities can enter persisted billing details [new]src/xagent/core/model/chat/token_context.py:204— major mixed media/token details are incompatible with existing quota hooks [new]src/xagent/core/tools/core/media_usage.py:78— major unresolved model identity is billed under a placeholder fallback [new]
Recommended event: REQUEST_CHANGES
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"]`.
|
Worked through all 15 findings (review body + inline). Pushed in 50afcc4. FixedN3 — invalid quantities (major). Validation also moved to the real write boundary. It lived only in the N2 — non-atomic counter/detail write. New N1 — N11 — lock breaks asdict/deepcopy/pickle. Confirmed all three raised N9 — wrapper drops optional metadata. N12 — raising descriptor escapes the guard. Each provider attribute read N15 — N13, N14 — test gaps. Both were fair. N13: added cases that monkeypatch Deferred with issuesN6 — mixed detail union (raised as blocking). Deferred to #1461 with
N5, N7, N8a, N8b — contract shape. Deferred to #1460. All four are VerificationTests were written to fail on revert. Local test runs are not usable in my |
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.
rogercloud
left a comment
There was a problem hiding this comment.
Summary
This PR introduces reusable billing primitives for non-LLM media calls: MediaUnit/MediaCallType, public recording APIs, provider-token metadata, aggregation, and best-effort wrappers. It routes media rows through TokenUsage.details so TaskTracker persistence, per-turn deltas, and the existing quota boundary can consume them without a separate storage path, while adding locking and serialization support for shared usage. No producers are wired in this PR, so the boundary and compatibility contracts here will be inherited by the upcoming image, audio, embedding, and rerank integrations.
Blocking: yes — recommended event: REQUEST_CHANGES
Since the last review
Two commits landed since the previous review at 979c0515. 50afcc49 closes numeric, snapshot, and serialization gaps: it sanitizes quantities, unifies direct and module-level writes under the atomic recording path, carries media_calls through TaskTracker seeds/copies, makes the lock safe for asdict/deepcopy/pickle, forwards wrapper metadata, and adds regression coverage. 832c5273 is a typing-only follow-up for the media write path and does not change runtime behavior.
Approach / design verdict
Verdict: acceptable-with-reservations. Reusing the existing TokenUsage.details → TaskTracker persistence → delta/quota flow is the right architectural direction: it avoids a parallel media store, keeps LLM token totals isolated, and gives each modality an explicit unit and call type. The lock-aware write/merge/snapshot work is also aimed at the real shared-state boundary rather than masking races at individual producers.
C2 — minor design-level scalability risk
src/xagent/web/tracking/task_tracker.py:53-67, 166-173, 432 still clones the complete unbounded details list and rewrites the complete JSON value on each periodic/final update. The new public API appends one row per call, so once the planned embedding, rerank, or segment producers are high-frequency, retained task JSON grows with every call and cumulative copy/serialization work can become O(n²). This is a latent minor scalability risk, not an immediate regression: no producer is wired in this PR and the copy/write mechanism predates it. The same root was called out as NEW-D2 in the linked #997 discussion. Before those producers land, prefer bounded/incremental keyed media persistence or a persisted delta cursor.
New findings
C1 — major — REQUESTS quantities are not enforced
src/xagent/core/model/chat/token_context.py:817 documents that MediaUnit.REQUESTS always has quantity 1, but both public write paths pass the value through generic _coerce_float and persist 0, fractional, or greater-than-one quantities. A requests record can therefore contribute one media_calls/aggregate calls entry while its billable quantity says something else, making quota and pricing disagree. Enforce exactly 1.0 for REQUESTS at the shared record_media_call boundary before any mutation; the best-effort wrapper may catch and log the direct-API error, while direct APIs should fail. Replace the permissive dirty-quantity test with rejection and no-state-change cases.
C4 — major — provider token counts accept malformed values
src/xagent/core/model/chat/token_context.py:839 sends provider input/output counts through _coerce_int, which accepts booleans, preserves negatives, and lets int(float("inf")) raise OverflowError because that exception is not caught. A malformed provider payload can persist invalid raw counts, or the best-effort wrapper can drop an otherwise valid media row when coercion raises. Use finite, non-negative integral coercion at this boundary, catch overflow, sanitize invalid fields to zero while preserving the media row, and add both direct-API and wrapper-path tests.
C5 — minor — TaskTracker media seam is not tested
src/xagent/web/tracking/task_tracker.py:104 changes seed/persistence/delta behavior, but the changed tests do not pass media rows through TaskTracker seeding, periodic/final persistence, and both progress/completion hooks. Add one seam-level integration test with prior-turn and current-turn media rows, asserting that the prior rows are not re-reported and that each hook receives only the current delta.
Prior findings checklist
The PR author is OliverBryant, so no qinxuye rebuttal waiver applies. Statuses below are canonical roots, not duplicate new findings.
- P1 —
FIXED—call_type=Noneis normalized to the empty optional value at the write boundary. Original comment (review4949744795). - P2 —
FIXED—unit=Noneis rejected before mutation. Original comment (review4949744795). - P3 —
FIXED— finite/non-negative quantity handling and the duration fallback are now present. Original comment. - P4 —
DROPPED(tracked, not fixed) — the discriminateddelta_details/external quota-hook contract is explicitly deferred to #1461; see the author's reply to the original finding. The stock tree has no hook registration or in-tree consumer, so the current concern is out-of-tree and is not re-reported here. - P5 —
DROPPED(tracked, not fixed) — the unresolved/placeholder identity fallback remains a producer/contract follow-up in #1460; see the author's reply to the original finding. The current fallback is not being re-raised while that issue is tracked. - P6 —
FIXED— TaskTracker snapshots preservemedia_calls, and seeds derive it from persisted media rows. Original comment. - P7 —
PARTIAL(minor) —src/xagent/core/model/chat/token_context.py:242andsrc/xagent/web/tracking/task_tracker.py:65: the canonical recording path is atomic, but the public two-stepadd_media_usage(count_call=False)plusincrement_media_callsremains, and TaskTracker still copies fields/details without the usage lock. Use a lock-aware snapshot API and add an interleaving test. Original comment. - P8 —
DROPPED(tracked, not fixed) — raw input/output provider-token fields are preserved, there is no in-tree pricing consumer, and the aggregate output contract is tracked in #1460; see the author's reply to the original finding. - P9 —
DROPPED(tracked, not fixed) — the identity collision is not reachable with current producers and tagged identity is tracked in #1460; see the author's reply to the original finding. - P10 —
DROPPED(tracked, not fixed) — no current resolver producer exercises the unstructured return shape, which is tracked in #1460; see the author's reply to the original finding. - P11 —
FIXED— the best-effort wrapper accepts and forwards all optional media metadata inside its guarded call. Original comment. - P12 —
PARTIAL(minor) —src/xagent/core/model/chat/token_context.py:107insertsmedia_callsbefore the legacytool_calls/detailspositional fields, so old five-positionalTokenUsage(...)callers can silently bind the wrong fields. Preserve the legacy positional order (or make the new field keyword-only/custom-init) and add a five-positional compatibility test. Original comment. - P13 —
FIXED— descriptor reads are individually guarded against provider properties that raise. Original comment. - P14 —
FIXED— exception and invalid-metadata paths now have exercising tests. Original comment. - P15 —
PARTIAL(minor) —tests/core/model/chat/test_media_usage.py:461checks opposing-merge timeouts and lower bounds, but has no deterministic overlap barrier or exactmedia_calls/detail consistency assertions. Add those assertions and force the interleaving. Original comment. - P16 —
PARTIAL(minor) —src/xagent/core/model/chat/token_context.py:467catches onlyTypeErroreven though the estimator docstring promises malformed iterables never raise. Narrow the promise or catch the intended iterator errors, and add a custom raising-iterable test. Original comment.
Verification
CI preflight reports 14 completed checks, all successful. No local tests were run under this review.
Blocking status & recommended decision
Blocking: yes — recommended event: REQUEST_CHANGES
src/xagent/core/model/chat/token_context.py:817— major —REQUESTSquantities are not constrained to exactly one, so persisted quantity and call count can diverge.[new]src/xagent/core/model/chat/token_context.py:839— major — malformed provider token counts can persist invalid values or drop an otherwise valid media row.[new]
C2, C5, P7, P12, P15, and P16 are minor and do not block; the tracked P4/P5/P8/P9/P10 roots are dropped under the follow-up history rule.
…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.
|
Round 2 addressed in 0ccbd35. Verified each finding against the code before FixedC1 — C4 — provider token counts (major). Confirmed all three parts: Fixed on P12 — positional ABI (was PARTIAL). You're right, and my previous test P7 — lock-aware snapshot (was PARTIAL). Added P16 — estimator exception scope (was PARTIAL). Confirmed a custom C5 — TaskTracker seam (minor). Added P15 — merge test (was PARTIAL). Agreed the old assertions were too Deferred with an issueC2 — unbounded P4/P5/P8/P9/P10 stay tracked in #1460 / #1461 as you recorded. Self-review notesTwo things my own pass caught, neither reported:
Verification
|
…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.
rogercloud
left a comment
There was a problem hiding this comment.
Summary
This PR adds the shared vocabulary (MediaUnit, MediaCallType) and recording primitives for media (non-LLM) usage billing, threading media rows through TokenUsage.details and the existing TaskTracker / quota path rather than inventing a parallel meter. Since the last review, two commits landed that close both previously blocking findings: REQUESTS quantity is now constrained at the single shared write boundary, provider token counts are sanitized instead of trusted, the dataclass positional ABI is restored, and the flaky-by-construction concurrency test now uses a real barrier. One new major issue surfaced this pass: _turn_delta() still reads the live shared TokenUsage without the lock this PR introduced, on the hot quota-gating path — the sibling call site was fixed in the newest commit, this one was missed.
Blocking: yes — recommended event: REQUEST_CHANGES
Since the last review
Two commits landed on top of 832c527:
0ccbd35— "fix: constrain REQUESTS quantity, sanitize token counts, restore positional ABI"- C1 (was blocking) — FIXED.
record_media_callnow raisesValueErrorbefore any mutation whenunit == MediaUnit.REQUESTSandquantity != 1.0(src/xagent/core/model/chat/token_context.py:296-300). Placing it at the single shared write boundary means both public entry points inherit the constraint rather than each re-implementing it. - C4 (was blocking) — FIXED.
_coerce_int(src/xagent/core/model/chat/token_context.py:481-508) rejectsbool, floors negatives to 0, and catchesOverflowErrorfromint(float("inf")). Importantly it sanitizes to 0 while preserving the row, so a malformed provider payload degrades the number rather than silently dropping the billing record. - P12 — FIXED.
media_callsis nowfield(default=0, kw_only=True)declared aftertool_calls/details(src/xagent/core/model/chat/token_context.py:104-112), restoring the legacy 5-positional-argument binding. - P16 — FIXED.
estimate_tokensnow catches broadExceptionand logs vialogger.warningwith the exception message, instead of onlyTypeError. - P7 Part A — FIXED.
TaskTracker._copy_usagenow delegates toTokenUsage.snapshot(), which takes the lock and deep-copies detail dicts.
- C1 (was blocking) — FIXED.
43e019e— "test: split the dirty-quantity case by unit after the REQUESTS constraint"- Splits the dirty-quantity coverage by unit so the
REQUESTSrejection path and the sanitize-to-0.0path for other units are asserted separately, with no-partial-state assertions on the rejection branch and coverage of the wrapper's swallow path. - P15 — FIXED. The opposing-merge test now uses
threading.Barrier(2)to force genuine overlap and asserts exactmedia_calls == len(media_rows)consistency plus bounded row counts, rather than only "did not deadlock". - C5 — FIXED.
tests/web/tracking/test_task_tracker.py:1119(test_media_rows_report_only_the_current_turn_delta) seeds prior-turn media rows, records a current-turn call, and asserts both the run-progress gate hook and the completion/usage-record hook observe only the current-turn delta.
- Splits the dirty-quantity coverage by unit so the
(All of the above verified by reading the code and the test assertions, not by executing a CI-equivalent suite.)
Approach / design verdict
Acceptable with reservations. Reusing TokenUsage.details and the existing TaskTracker/quota pipeline is the right call — it avoids a second, divergent meter — and MediaUnit/MediaCallType is a reasonable vocabulary to build pricing on. Three reservations keep this from a clean "sound", none of them newly blocking:
- No
MediaUnitpricing consumer exists yet, so media prices at a literal $0 until one is built (tracked in #1460 / #1461). resolve_billing_model's documented invariant ("never return a placeholder") is contradicted by its own default fallback returning"default", which is a member of its own placeholder set (prior finding P5, tracked in #1460).- The
_lock/__getstate__/__setstate__/snapshot()machinery onTokenUsageis wider in scope than "shared vocabulary and recording helpers" — it changes concurrency and serialization semantics of a class used across the codebase, and the stated pickling motivation has no current in-tree caller. Worth naming as scope creep; not blocking on its own. See Simplification opportunities.
New findings (this pass)
NEW-A — MAJOR — unlocked read of live shared TokenUsage on the quota-gate path
src/xagent/web/tracking/task_tracker.py:538-547 (_turn_delta), reached from interrupt_reason_for_quota at src/xagent/web/tracking/task_tracker.py:567.
def _turn_delta(self, usage: TokenUsage | None = None) -> tuple[list, int]:
if usage is None:
usage = get_token_usage()
return (
usage.details[self._initial_details_len :],
max(0, usage.tool_calls - self._initial_tool_calls),
)On the interrupt_reason_for_quota path, usage is the live, shared TokenUsage from get_token_usage(), and both the details slice and the tool_calls read happen outside TokenUsage._lock — the very lock this PR introduced to prevent this class of torn read. The two reads are separate un-synchronized operations, so a concurrent record_media_call or increment_tool_calls interleaving between them yields a (details, tool_calls) pair reflecting two different logical points in time: the slice may already contain a just-appended media/tool row whose counter bump has not landed, or the counter may have advanced past the rows visible in the slice. That internally inconsistent pair is fed straight into billing/quota gating.
This matters more than a typical racy read because interrupt_reason_for_quota is documented as polled at every safe point (each LLM reply / tool call) — it is a hot, concurrent path, not a one-shot teardown read. It is also specifically the sibling of _copy_usage, which this PR's newest commit correctly re-pointed through usage.snapshot() for exactly this reason; _turn_delta was missed by that fix. Since TokenUsage._lock did not exist before this PR, this inconsistency window is introduced by this PR.
Suggested fix: have _turn_delta take a usage.snapshot() (or hold usage._lock across both reads) before computing the pair — at minimum on the interrupt_reason_for_quota call path, ideally unconditionally so the two call sites cannot diverge again.
NEW-B — MINOR — copy.copy(usage) shares details but not the lock
src/xagent/core/model/chat/token_context.py:137-146 (__getstate__ / __setstate__, with no __copy__).
__getstate__/__setstate__ correctly give a pickled or deep-copied instance its own fresh lock, and that path is tested. copy.copy() is not covered: Python's shallow copy goes through the same __getstate__/__setstate__ pair, so the new instance gets a different threading.Lock while still referencing the same details list object. Two objects that each believe they hold exclusive access then mutate one shared list under two different locks — precisely the mutual exclusion the lock exists to provide. Confirmed empirically.
No in-tree caller does copy.copy() on a TokenUsage today and snapshot() is already the right alternative, so this is latent rather than exploited. Suggest either defining __copy__ to delegate to snapshot(), or a short comment on __getstate__ warning that copy.copy is unsafe and snapshot() is the supported way to fork a usage object.
NEW-C — MINOR — to_dict() leaks the live inner detail dicts
src/xagent/core/model/chat/token_context.py:390 — "details": list(self.details).
to_dict() correctly takes the lock, but list(self.details) creates a new outer list around the same inner dict objects. A caller that mutates an entry of the returned payload mutates the live usage object directly, entirely outside the lock. Confirmed empirically. This is inconsistent with snapshot() in the same class, which already establishes the correct pattern ([dict(item) for item in self.details ...]).
No in-tree caller mutates the returned dict today, so this is an encapsulation break rather than an active bug. Suggest matching snapshot(): "details": [dict(item) for item in self.details].
NEW-D — MINOR — counter increments accept negative counts
src/xagent/core/model/chat/token_context.py:349-357 (increment_media_calls, increment_tool_calls).
Neither validates that count is non-negative, while _coerce_int / _coerce_float in the same module explicitly reject or floor negatives. A negative count would drive the counter below the number of matching detail rows, producing the same counter/rows divergence the C1 fix was written to prevent. The only in-tree caller (add_tool_call_usage) always passes count=1, so this is a hardening gap, not an active bug. Suggest rejecting or clamping count < 0, consistent with the coercion helpers.
NEW-E — MINOR — swallowed accounting errors lose the traceback
src/xagent/core/tools/core/media_usage.py:145-146 (record_media_usage).
To be clear, the broad except Exception here is correct and deliberate — a metering bug must never break the media call the user asked for, and that intent is documented and tested. The narrower gap is diagnosability: the handler logs logger.warning("Failed to record %s media usage: %s", call_type, e) with no exc_info, so a genuinely unexpected bug (as opposed to the intended validation ValueError) surfaces as a single line with no stack trace. Given this PR's own framing that persisted billing rows are unrepairable after the fact, that is the wrong trade. Suggest logger.warning(..., exc_info=True) or logger.exception(...).
NEW-F — MINOR — zero quantity is overloaded (observability, not correctness)
_coerce_float folds negative / NaN / inf / non-numeric quantities into the same 0.0 the design deliberately uses to mean "call happened, unmeasured". aggregate_media_usage_by_model then sums quantity and counts calls per group with no field distinguishing measured from unmeasured or corrupt contributors — unlike token estimation, which does carry a separate tokens_estimated boolean and thereby sets the precedent this dimension doesn't follow.
This is currently tested and accepted behavior (test_zero_quantity_media_entries_stay_visible asserts exactly the blended count), so it is a refinement suggestion, not a defect: consider an analogous unmeasured_calls / has_unmeasured field in a follow-up so aggregate rows can be audited.
Checked, no issue
Placing media_usage.py under src/xagent/core/tools/core/ is fine. That directory already holds several non-tool helper modules (e.g. command_path_guard.py), and there is no directory-scanning tool auto-discovery in this codebase — tools are wired via explicit named imports only, so the module cannot be mistaken for a registered tool.
Prior findings checklist
| # | Finding | Status |
|---|---|---|
| C1 | REQUESTS quantity not enforced |
FIXED — validated pre-mutation at the shared write boundary |
| C2 | Unbounded details clone on every persist |
STILL OPEN — _commit_task_usage_if_owned full-column rewrite (src/xagent/web/tracking/task_tracker.py:144-178) untouched; previously agreed out of scope for a helpers-only PR, non-blocking |
| C4 | Provider token counts accept malformed values | FIXED — _coerce_int sanitizes bool/negative/inf/nan/string, row preserved |
| C5 | TaskTracker media seam untested |
FIXED — tests/web/tracking/test_task_tracker.py:1119 |
| P1 | call_type=None normalization |
STILL FIXED — untouched by the two new commits |
| P2 | unit=None rejection before mutation |
STILL FIXED — untouched by the two new commits |
| P3 | Finite/non-negative quantity + duration fallback | STILL FIXED — untouched by the two new commits |
| P4 | Discriminated delta_details/quota-hook contract |
STILL TRACKED — no change; #1460 / #1461 confirmed open |
| P5 | resolve_billing_model returns the "default" placeholder it forbids |
STILL TRACKED — no change; tracked in #1460 |
| P6 | TaskTracker snapshots preserve media_calls |
STILL FIXED — untouched by the two new commits |
| P7 | Two-step API + TaskTracker copies without lock |
PARTIAL — Part A (lock safety) FIXED via snapshot(); Part B (two-step API) not fixed, see below |
| P8 | Raw provider-token fields, no in-tree pricing consumer | STILL TRACKED — no change; #1461 |
| P9 | Model identity collision in aggregation | STILL TRACKED — no change; #1460 / #1461 |
| P10 | Unstructured resolver return shape | STILL TRACKED — no change; #1460 / #1461 |
| P11 | Best-effort wrapper forwards optional metadata | STILL FIXED — untouched by the two new commits |
| P12 | media_calls positional ABI break |
FIXED — kw_only=True after tool_calls/details |
| P13 | Descriptor reads guarded against raising providers | STILL FIXED — untouched by the two new commits |
| P14 | Exception/invalid-metadata paths exercised by tests | STILL FIXED — untouched by the two new commits |
| P15 | No deterministic overlap barrier in the merge test | FIXED — threading.Barrier(2) + exact consistency assertion |
| P16 | estimate_tokens only caught TypeError |
FIXED — broad catch, warning log, non-TypeError test |
The two new commits touch only REQUESTS validation, _coerce_int, dataclass field ordering, snapshot(), and estimate_tokens's exception handling — disjoint from the code paths behind P1/P2/P3/P6/P11/P13/P14, which were re-read to confirm no regression. resolve_billing_model, the quota_hooks.py contract docs, aggregate_media_usage_by_model's identity key, and the provider-token-field consumer are all unmodified, so P4/P5/P8/P9/P10 keep their tracked disposition — not re-raised here.
P7 Part B (remaining, minor, non-blocking). TokenUsage.add_media_usage(count_call=False) and increment_media_calls() remain public and are still documented as kept only for backwards compatibility. In that two-step sequence the caller bumps media_calls first, and the subsequent record_media_call performs its validation before taking its lock — so a rejected call (e.g. a bad REQUESTS quantity) leaves media_calls incremented with no matching detail row. That is exactly the torn state C1 was fixed to prevent, still reachable through the legacy door. There are no in-tree production callers of this path (grep-confirmed); only tests exercise it. Recommendation: delete it rather than harden it — see below.
Simplification opportunities
Two clusters, both reductions rather than new abstractions:
-
Delete the two-step media API (same root cause as P7 Part B above — one recommendation, not two).
TokenUsage.add_media_usageas an instance method, thecount_callparameter onrecord_media_call(src/xagent/core/model/chat/token_context.py:250), andincrement_media_calls(src/xagent/core/model/chat/token_context.py:349) exist solely to preserve a legacy two-step flow with zero production callers, whose only exercise is this PR's own tests, and which is documented as inherently capable of torn state. Keepingrecord_media_call(always counting) plus the module-leveladd_media_usage/record_media_usageentry points covers every real caller. Deleting it removes the P7 Part B defect and shrinks the diff at the same time. -
Reconsider the pickle/lock compatibility layer (
__post_init__, the_lockproperty,__getstate__,__setstate__atsrc/xagent/core/model/chat/token_context.py:116-146). This machinery exists only to keepTokenUsagepicklable anddataclasses.asdict()-safe after introducing a lock, but nothing in the codebase pickles, deep-copies, orasdict()s aTokenUsagein production (grep-confirmed). This is the concrete form of the scope-creep reservation in the design verdict above — and it is also what makes NEW-B possible. A plain private attribute assigned outside__init__, or simply not supporting pickle, would be less surface for no lost capability.
Net: roughly 80-110 lines removable — the two-step API and its dedicated tests account for ~50, and the pickle/lock compatibility layer plus its tests for most of the remainder. Estimated from the size of the methods and test cases named above, not measured by applying the change.
Blocking status & recommended decision
Blocking: yes — recommended event: REQUEST_CHANGES
Blocking issues:
- NEW-A
[new]— unlocked, non-atomic read of the live sharedTokenUsagein_turn_deltaon theinterrupt_reason_for_quotaquota-gating path (src/xagent/web/tracking/task_tracker.py:538-547). Major, and newly introduced by this PR; the sibling_copy_usagecall site was fixed in0ccbd35but this one was not.
No longer blocking:
- C1
[fixed], C4[fixed]— both previously blocking findings are resolved by0ccbd35/43e019e.
Non-blocking (minor, can land as follow-ups): P7 Part B (two-step API deletion), C2 (persist-time clone, out of scope), and NEW-B / NEW-C / NEW-D / NEW-E / NEW-F.
Findings in this review were verified by reading the implementation and the test assertions and by targeted experiments on the copy/serialization semantics; no full CI-equivalent suite run was performed locally.
…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.
|
Round 3 addressed in 650663f. Reproduced each finding before fixing. NEW-A (blocking) — fixed, and it was my regressionYou're right, and the diagnosis is exact: this is the sibling of Reproduced with a forced interleaving (barrier between the two reads, write in Randomly-scheduled concurrency showed 0/4000 — the GIL makes each individual Fixed by snapshotting unconditionally, including when the caller already Simplifications — both takenTwo-step API deleted. Pickle/lock layer narrowed rather than removed. I kept NEW-B / NEW-C / NEW-D / NEW-E — fixedAll four confirmed empirically first:
Self-review found one more
DeferredNEW-F → #1495. Agreed it's a real observability gap and that C2 stays in #1466. P4/P5/P8/P9/P10 keep their tracked disposition. Verification
|
rogercloud
left a comment
There was a problem hiding this comment.
Summary
This PR adds the shared vocabulary (MediaUnit, MediaCallType) and recording primitives (record_media_call, add_media_usage, aggregate_media_usage_by_model) for non-LLM media usage — image, video, TTS, ASR, music, embedding, rerank — threading type:"media" rows through the existing TokenUsage.details list so TaskTracker persistence and the quota delta_details path carry them with no DB schema change. It also introduces the first threading.Lock on TokenUsage, which previously had bare unlocked += on every counter. No producers are wired up yet; this is the shared primitive layer only, first of five PRs splitting #997.
Blocking: yes — recommended event: REQUEST_CHANGES
Since the last review
One commit landed on top of 43e019e8: 650663f6 — "fix: snapshot before pairing turn-delta reads, and drop the two-step media API" (231 insertions / 91 deletions across token_context.py, media_usage.py, task_tracker.py and two test files). It closes six prior findings in one pass:
- NEW-A (was the sole blocking issue) — FIXED.
_turn_delta(src/xagent/web/tracking/task_tracker.py:556-563) now callsusage.snapshot()once and derives both the details slice and the tool-call delta from that single detached snapshot;snapshot()(src/xagent/core/model/chat/token_context.py:296-305) reads both fields under onewith self._lock:hold. The torn-read window is genuinely closed. - NEW-B (
copy.copylock aliasing) — FIXED.__copy__now delegates tosnapshot(). - NEW-C (
to_dictleaked live inner dicts) — FIXED. Now builds[dict(item) for item in self.details]. - NEW-D (negative counter increments) — FIXED.
increment_media_callsdeleted outright;increment_tool_callsrejects negatives with a warning. - NEW-E (swallowed errors lost the traceback) — FIXED.
exc_info=Trueadded inmedia_usage.py. - P7 Part B (two-step media API) — FIXED. The instance-method
TokenUsage.add_media_usage,record_media_call'scount_callparameter, andincrement_media_callsare all gone (repo-wide grep forincrement_media_callsreturns zero matches). Only the module-leveladd_media_usageremains, so the torn-state sequence is no longer reachable.
That is a clean round, and deleting the legacy path rather than hardening it was the right call.
Approach / design verdict
Acceptable with reservations. Reusing TokenUsage.details and the existing TaskTracker/quota pipeline remains the right integration seam, and six prior findings closed correctly. Three reservations keep it from "sound": the NEW-A fix traded a correctness bug for a hot-path cost regression (N1 below); the only regression test written for that fix cannot fail (N2); and the aggregate-row identity and delta_details typing questions stay deferred to #1460 / #1461 as previously agreed.
New findings this pass
N1 — MAJOR — the NEW-A fix introduces an O(n²) deep copy on the quota hot path
src/xagent/web/tracking/task_tracker.py:556-563 (_turn_delta) with src/xagent/core/model/chat/token_context.py:296-305 (snapshot).
_turn_delta unconditionally calls usage.snapshot(), which deep-copies the entire cumulative details list under the lock, and then discards everything before _initial_details_len. At the base commit _turn_delta was a plain slice with no copy at all, so this cost is new to this PR.
Three facts make it matter rather than being a micro-optimization:
- The list grows across turns.
_task_seed_from_session(src/xagent/web/tracking/task_tracker.py:94-121) seedsdetailsfrom the persisted cumulative list, so its length is monotonically non-decreasing over a run. The copy is proportional to total accumulated usage, not to the delta being read. - The path is hotter than its own docstring says.
interrupt_reason_for_quotais wired as theinterrupt_checker(src/xagent/web/api/chat.py:3243-3245,src/xagent/web/api/websocket.py:3365-3366) and reached fromruntime.pyviashould_interrupt()both once per ReAct/DAG/auto step (react.py:3009,dag.py:1812,auto.py:633) and once per streamed chunk inside a single LLM call (src/xagent/core/agent/runtime.py:291). One long streamed response triggers dozens of full-list copies on its own. - There is no cheap short-circuit.
check_run_progress_gate's early return (src/xagent/web/services/quota_hooks.py:104) runs after_turn_deltahas already paid the cost, so disabling the quota hook does not avoid it. The lock held during the copy is also the one serializingadd_input_tokens/add_output_tokens/record_media_callfrom the LLM adapters.
Suggested fix: add a narrow locked accessor that returns only the tail plus the counter in one lock acquisition, e.g. TokenUsage.detail_tail(start) -> tuple[list[dict], int] returning ([dict(d) for d in self.details[start:]], self.tool_calls). That preserves exactly the atomicity guarantee the snapshot was added for, at O(delta) instead of O(total). The _copy_details(delta_details) calls at the two call sites can then go as well — though note that second copy is already O(delta), so it is a cleanup, not the main win.
N2 — MAJOR (test quality) — the regression test for the fix just landed cannot fail
tests/web/tracking/test_task_tracker.py:1120 (test_turn_delta_pairs_details_and_tool_calls_atomically).
Verified by direct experiment. The test passes identically with and without the fix, for two independent reasons:
token_contextis acontextvars.ContextVar.read_delta()callstracker._turn_delta()withusage=None, so_turn_deltacallsget_token_usage()from inside a plainthreading.Thread— which does not inherit the calling context. It lazily creates a brand-new emptyTokenUsage, entirely disconnected from theusageobject the writer thread mutates via closure. The read side therefore observes([], 0)no matter what the implementation does.- Against the pre-fix implementation,
snapshot()was never called at all, so the patchedslow_snapshotnever fires,enteredis never set, and the writer'sentered.wait(timeout=5)times out silently — its boolean return is never checked — again yielding(0, 0).
The test contains exactly one assertion, assert media_rows == delta_actions, which degenerates to 0 == 0 on both sides of the change; nothing asserts that the interleave actually occurred. Since this is the only regression test for the race that was just fixed, the fix is currently unguarded.
Suggested fix: pass usage explicitly into _turn_delta() (or propagate the context with contextvars.copy_context().run(...)) so the reader observes the same live object the writer mutates, and assert entered.is_set() so a silent timeout fails rather than passes.
Minor findings
- N3 — the lock's detachment invariant is not enforced at the constructor /
from_dict.__post_init__(token_context.py:116-129) creates the lock but does not copydetails, andfrom_dict(:368-377) stores the caller's list by reference. No exploitable path exists today —from_dicthas no production caller, and the one production constructor call passingdetails=(task_tracker.py:102-117) hands over a fresh_copy_detailsresult nothing else retains. Preventive fix: normalize once in__post_init__(self.details = [dict(d) for d in self.details if isinstance(d, dict)]), covering the constructor andfrom_dicttogether. - N4 — string fields reach a JSON column unsanitized while numbers are hardened.
record_media_call(token_context.py:221-283) writesmodel,model_id,resolutionwith no type check while the numeric fields go through_coerce_float/_coerce_int.detailslands inTask.token_usage_details, a plainColumn(JSON). A non-serializable value would not raise at the call site —record_media_usage'stry/exceptonly covers write-time validation — but far away at commit time. Zero blast radius today, but this is the boundary the numeric-coercion docstrings exist to protect. - N5 — non-reentrant
threading.Lockis a latent trap._lockis a plainLock(token_context.py:129). No self-deadlock exists today:to_dictdeliberately inlines the total instead of callingtotal_tokens(),mergetakes the two locks sequentially and never nested, andsnapshotbuilds a new object with its own lock. But that rests on manual discipline, and nothing documents the non-reentrancy as deliberate.threading.RLock()removes the trap for free. - N6 —
_coerce_int's changed semantics are untested on the pre-existing LLM path. The change frombool -> int(value)/ negatives-passed-through tobool -> 0/ negatives-> 0(token_context.py:455-481) affectsadd_token_usage,extract_cached_input_tokens, andaggregate_token_usage_by_model, which every adapter calls. To be clear, this was explicitly disclosed in0ccbd35b's message and is intentional, and real provider payloads are unlikely to be bool or negative — the only gap is that no test exercises the new behavior throughadd_token_usage. One small case would close it. - N7 — the
REQUESTSerror message reports a value the caller never passed. See the inline comment ontoken_context.py:262. - N8 —
add_media_usage's "Raises" section omits theREQUESTSquantityValueError. The docstring (token_context.py:872-877) lists only theunit/call_typeerrors, and the "Validate before touching the context" comment at:884-885is true only for those; theREQUESTSconstraint is enforced insiderecord_media_call(:260), which runs afterget_token_usage()at:891. Worth documenting the third raise. (The duplicate coercion between the two functions is not a finding — it is deliberate defense-in-depth, explained in-code at:244-247, becauserecord_media_callis independently public.) - N9 —
estimate_tokensis unexported and collides in name with an established helper.token_context.py:487is the only new top-level symbol this PR adds that is not exported fromsrc/xagent/core/model/chat/__init__.py(MediaUnit,MediaCallType,add_media_usage,aggregate_media_usage_by_model,aggregate_token_usage_by_modelall are). MeanwhileCompactUtils.estimate_tokensalready exists (src/xagent/core/agent/utils/compact.py:37) with a different signature and a different algorithm (chars // 4), with 13 call sites plus aMessageUtils.estimate_tokenswrapper. Suggest renaming (e.g.estimate_media_tokens) and exporting it, or deferring it to the producer PR that consumes it. - N10 —
estimate_tokens' CJK ranges omit CJK punctuation and fullwidth forms.token_context.py:519-527covers U+4E00–9FFF, U+3040–30FF and U+AC00–D7AF but not U+3000–303F or U+FF00–FFEF, which appear in essentially every Chinese sentence (nor the rarer Ext-A/Ext-B blocks). For typical prose the undercount is a few percent rather than the 4x a worst case would suggest, but it is real and invisible to the current tests, whose only CJK inputs are two-character Han strings (with and without an ASCII suffix), both inside the covered range. - N11 — a false comment justifies the
media_callsseed. See the inline comment ontask_tracker.py:110. - N12 — an implicit, undocumented "details are always dicts" invariant.
snapshot()filters non-dicts (token_context.py:302) while_turn_deltaslices the filtered list by_initial_details_len. This is currently safe and unreachable —_initial_details_lenis computed instart_trackingfrom_copy_usage(seed.usage)->snapshot(), i.e. an already-filtered list which then becomes the live usage object; every mutation path appends a literal dict; andfrom_dict, the one unfiltered path, has no production caller. Still worth documenting the invariant (or indexing by a monotonic sequence number) so a future non-dict append cannot silently misalign the slice and re-meter prior-turn rows. - N13 — the module docstring no longer describes the module.
src/xagent/core/model/chat/token_context.py:1-6still reads "Token usage tracking using contextvars … across LLM calls"; the file now also definesMediaUnit(:19-34),MediaCallType(:37-48),record_media_call,add_media_usage, andaggregate_media_usage_by_model. - N14 — the aggregators' "pass a detached list" contract is undocumented.
aggregate_token_usage_by_model(:606-616) andaggregate_media_usage_by_model(:704-717) iterate the passed list with no lock and no docstring warning. Safe today — the only real caller (src/xagent/web/api/chat.py:4536) passes a DB-read column value — but nothing stops a future caller from handing in a liveusage.detailsand iterating while another thread appends. - N15 —
test_concurrent_merge_loses_no_countshas little regression power. See the inline comment ontest_media_usage.py:369. - N16 — a comment contradicts the assertion below it. See the inline comment on
test_media_usage_helpers.py:50. - N17 — the
model_idbranch ofaggregate_media_usage_by_modelis untested. No test passes a media detail with a non-emptymodel_idinto that aggregator: the two tests that setmodel_idon a media entry (test_media_usage.py:30,47) never aggregate, and themodel_iduses at:126,157,180feed the LLM-token aggregator instead. Soidentity = model_id or model_name(:748) always takes the name branch and the backfill at:753-767never runs. Since this identity logic is already slated for rework in #1460 / #1461, a one-line test can fold into that follow-up rather than block here. - N18 — stale test counts in the PR description. The body says "22 and 11 cases". Actual:
tests/core/model/chat/test_media_usage.pyhas 45def test_functions (74 with parametrize expansion), andtests/core/tools/core/test_media_usage_helpers.pyhas 11 functions (17 expanded). They undercount rather than overstate, but please correct or drop the numbers. - N19 (nit) —
media_callsstill has no consumer. Verified it has no reader outsidetoken_context.pyand the seed attask_tracker.py:113, is not persisted, does not enterdelta_detailsor any API/frontend surface, and is fully derivable fromdetails. To be fair about scope: an earlier draft of this finding claimed the PR's whole locking design existed to protect this field — that is refuted. The base commit had no lock onTokenUsageat all, and the lock fixes a real read-modify-write race affectingtool_callsand the token counters just as much (the PR reports ~85% count loss under 8x200 increments without it). #997 also names the follow-up consumers. So this is only a "please confirm the wiring lands in the follow-up" nit.
Prior findings checklist
| Finding | Status |
|---|---|
NEW-A — torn read of details vs tool_calls on the quota path |
FIXED |
NEW-B — copy.copy shares details but not the lock |
FIXED |
NEW-C — to_dict() leaks live inner detail dicts |
FIXED |
| NEW-D — counter increments accept negative counts | FIXED |
| NEW-E — swallowed accounting errors lose the traceback | FIXED |
| P7 Part B — two-step media API allows torn state | FIXED |
| NEW-F — zero quantity overloaded (unmeasured vs corrupt) | NOT FIXED — unchanged; recorded last round as a follow-up refinement, non-blocking |
C2 — _commit_task_usage_if_owned rewrites the whole column on every persist |
NOT FIXED — byte-for-byte unchanged; previously agreed out of scope, non-blocking |
P4 / P5 / P8 / P9 / P10 — aggregate-row identity, provider-token split, delta_details union |
Still tracked in #1460 and #1461, both confirmed open; not re-raised here |
Simplification opportunities
src/xagent/core/model/chat/token_context.py:302: shrink: the same detail-copy expression [dict(item) for item in ... if isinstance(item, dict)] is reimplemented four times — snapshot() (:302), merge() (:339), to_dict() (:364, minus the isinstance filter, so a non-dict entry would raise there but be skipped in the other two), and _copy_details (src/xagent/web/tracking/task_tracker.py:53-56). Extract one shared helper and call it from all four; the filter asymmetry disappears with it.
net: -15 lines possible
Two other simplification candidates were considered and deliberately rejected, so you know they were not overlooked: deleting the __getstate__ / __setstate__ / __copy__ / __deepcopy__ set (__copy__/__deepcopy__ were added specifically to fix NEW-B at this reviewer's request), and narrowing estimate_tokens' generic-iterable / broad-except handling (added specifically to fix P16 at this reviewer's request). Both should stay as they are.
Blocking status & recommended decision
Blocking: yes — recommended event: REQUEST_CHANGES
Blocking issues:
src/xagent/web/tracking/task_tracker.py:559— major[new]— unconditional full-list deep copy in_turn_deltaon the per-step/per-chunk quota path, held under the usage lock.tests/web/tracking/test_task_tracker.py:1120— major[new]— the sole regression test for the NEW-A fix cannot fail, because a barethreading.Threaddoes not inherit thecontextvarscontext and the reader observes a differentTokenUsageobject.
No longer blocking: NEW-A is confirmed fixed, along with five other prior findings.
Findings were verified by reading the implementation and the test assertions, and by targeted experiments on the contextvars/threading behavior behind N2. No local test suite was run; CI on this PR is green.
|
Round 4 addressed in a20c098. Both blocking findings were consequences of how N1 (blocking) — confirmed and measured
Also confirmed the two facts that make it matter: Fixed with One thing I had to keep: N2 (blocking) — confirmed; my test was emptyVerified the mechanism directly: So The distinction that let it through: I mutation-verified the fix in a side
Minor findings — fixedN3 ( Simplification taken: Self-review found two things I introduced this round
Also swept for the sibling-call-site question that caused NEW-A: no other Accepted for this layerN4 (string fields reaching the JSON column uncoerced) and N19 N18: dropped the stale counts from the PR body rather than restate numbers NEW-F stays in #1495, C2 in #1466, P4/P5/P8/P9/P10 in #1460/#1461. Verification
|
a20c098 to
88b30d5
Compare
…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.
88b30d5 to
a489258
Compare
|
Follow-up to my round-4 reply, after a further self-review pass. Nothing here Head is now Refinements since the round-4 pushNone of these came from your findings; they came from re-reading my own diff. 1.
2. 3. The 4. Two redundant test pairs merged (49 → 47 functions): the three CJK cases An incidental correctness gain from the N1 fixWorth flagging because it was not the point of the change. The old path was Not reachable today (that is your N12 invariant), but the new ordering is Also verified while re-reading, no change made
The scope question — your call, not mineYou raised this in the round-4 design verdict, and I want to put numbers on it
Every new symbol traces to a specific request — Taken together, though, a PR whose stated job is "shared vocabulary and The concrete candidate is the pickle/copy layer: I did not remove it, because you wrote that Two coherent end states, and I am happy with either:
Say which and I will do it in one commit. Disclosure: cancelled CI runs are my fault, not test failuresWhile re-chaining the five-PR stack I force-pushed repeatedly, which tripped the I also briefly wiped #1457's and #1463's own commits during that re-chain (a Verification
|
rogercloud
left a comment
There was a problem hiding this comment.
Summary
This is PR 1 of a 5-PR series splitting up #997 (2841 lines, previously too large to review). It introduces only the shared vocabulary and recording primitives for non-LLM media usage billing: MediaUnit/MediaCallType enums, TokenUsage.record_media_call/add_media_usage, aggregate_media_usage_by_model, resolve_billing_model, estimate_media_tokens, record_media_seconds/coerce_duration, and error-swallowing wrapper record_media_usage, plus locking/serialization hardening on TokenUsage (snapshot, to_dict/from_dict, merge, __copy__/__deepcopy__) and an atomicity fix in task_tracker.py's detail_tail. No producers are wired up in this PR — confirmed by grep, zero production callers exist for any new symbol — which matches the PR's stated scope for this slice.
Blocking: yes — recommended event: REQUEST_CHANGES
Round 0 verdict: acceptable-with-reservations
The direction — reusing TokenUsage.details as the substrate for media billing rows rather than a parallel structure — is sound, and this slice is genuinely inert in production today. However, real design gaps remain that should be resolved before downstream PRs build on this foundation:
- [Design] Two unenforced billing dimensions per row. The PR's central stability claim is that "unit is a property of the modality, so a price table keyed on
(model, unit)stays usable." ButMediaUnithas notokensmember, while token-reporting providers (embedding, gpt-image) also populateprovider_tokens/provider_input_tokens/provider_output_tokens(token_context.py:299-317,token_context.py:951). The rule "token-based price takes precedence over unit-based price" exists only as a code comment — never enforced or represented in the row schema, andaggregate_media_usage_by_modelgroups purely by(model, unit, call_type, resolution)with no field distinguishing "price by tokens instead." Zero consumers ofprovider_tokensexist anywhere outside this file, confirming the precedence rule is aspirational, not implemented. See inline comment ontoken_context.py:951. - [Design, non-blocking]
TokenUsage.media_callsis a redundant field (no DB column; recomputed by row-counting on every session-resume seed) that is the direct cause of avoidable complexity (kw_onlyABI-preservation dance, extra handling inmerge/to_dict/from_dict/snapshot). Note: the RLock and most of its hardening would be needed regardless ofmedia_calls, since the base commit had zero locking on the pre-existingtool_calls/token counters — that is a real, independently-justified fix, so please don't read this as "the locking design exists to protect media_calls." Suggestion: replacemedia_callswith a computed property (sum(1 for d in details if d["type"]=="media"), mirroringtotal_tokens), eliminating several touch-points. - [Informational]
merge()'s snapshot-then-lock rewrite has one call site insrc/(TokenContextManager.__init__,token_context.py:475), but no instantiation in the repo passesparent_usage, so the path is functionally unreachable today. Also worth documenting: snapshot-then-lock meansmerge()capturesotherat one instant, not a live drain — rows appended between snapshot and extend are not included. Expected semantics, but the namemergecould mislead a future caller into expecting a full drain.
Findings
Major
-
src/xagent/core/model/chat/token_context.py:633—_coerce_floatcatches onlyTypeError/ValueError, missingOverflowErrorthat its sibling_coerce_intwas already fixed for._coerce_float(10**400)raises uncaught. Effect: a huge/malformed quantity from a producer, routed throughrecord_media_usage's bareexcept Exception, causes the entire billing row to be silently dropped rather than recorded asquantity=0.0— violating_coerce_float's own documented "reject-to-0.0, never drop" contract. Any direct caller ofadd_media_usage/TokenUsage.record_media_callbypassing the wrapper crashes uncaught. No test covers a huge-int quantity.
Fix: addOverflowErrorto the except clause, mirroring_coerce_int; add a huge-int regression test. -
tests/core/model/chat/test_media_usage.py:336— Test-quality, major. The PR description claims "Fixes are verified by mutation — reverting a guard must fail its test," but this is false for the locking fix. Verified: monkey-patchingTokenUsage._locktocontextlib.nullcontext()(removing all mutual exclusion) and running the full concurrency test suite — all 6 concurrency-named tests still pass (test_concurrent_media_records_lose_no_counts,test_concurrent_merge_loses_no_counts,test_concurrent_merges_in_both_directions_do_not_deadlock,test_snapshot_is_taken_under_the_lock,test_copy_copy_does_not_share_details_under_a_different_lock,test_lock_is_reentrant). The 8-thread × 200-iteration loops are too short to reliably lose a+=race under CPython's GIL switch granularity, so these tests pass by luck, not because the lock is verified necessary.
Fix: increase iteration counts substantially and/or force switches viasys.setswitchinterval, and assert on a materially larger observed-loss threshold. -
tests/web/tracking/test_task_tracker.py:1120— Test-quality, major, possible regression on a previously-claimed fix.test_turn_delta_pairs_details_and_tool_calls_atomicallystill does not testdetail_tail's atomicity, despite the author previously replying that this exact concern (from an earlier review round) was fixed. Root cause: theslow_detail_tailmock setsentered/waits onreleasedbefore delegating to the realdetail_tailaccessor, so the concurrent writer thread always finishes its write before the real locked read even begins — the interleaving the test claims to check for never occurs. Confirmed by mutation: temporarily splittingdetail_tail's singlewith self._lock:block into two separate lock acquisitions with a 50ms sleep between readingdetailsandtool_calls(a deliberately torn implementation) — the test still passes. This is the second review round this exact test has been shown not to test what it claims; the previous fix addressed a different bug (silent mock timeout) without closing the real gap. -
src/xagent/core/model/chat/token_context.py:951— See "Two unenforced billing dimensions per row" in the Round 0 verdict above. Treating as blocking because the PR's own stability claim for the pricing key is unenforced and contradicted by the schema. -
tests/core/model/chat/test_media_usage.py:145— The module's own stated invariant ("unit is a property of the modality, never of the response" —media_usage.py:20-22,token_context.py:29-31) is not enforced anywhere.add_media_usage/record_media_callvalidateunitandcall_typeindependently against their own enums, with noMEDIA_UNIT_BY_CALL_TYPE-style mapping. Proof: this PR's own tests violate the invariant on day one —unit="images"paired withcall_type="video"appears at lines 145, 610, 632, 635, 671, 677 (six occurrences), even thoughrecord_media_seconds's own docstring says video must always reportseconds. This directly contradicts the PR's central pricing-stability claim, in its own test suite.
Minor (non-blocking, confirmed)
-
src/xagent/core/model/chat/token_context.py:968—add_media_usage's REQUESTS-quantity validation error reports the already-coerced value, not the caller's raw input, on the public wrapper path (quantity is coerced before callingrecord_media_call, discarding the true raw value one layer up). Previously flagged; the prior fix covers only the directTokenUsage.record_media_call()path (token_context.py:282-298). The only test (test_requests_error_reports_the_value_the_caller_passed) covers just the direct path, giving false confidence. Low severity (debuggability only, not billing correctness), but the fix is incomplete. -
src/xagent/core/model/chat/token_context.py:170—__getstate__does a shallow, unlocked copy (self.__dict__.copy()), leavingstate["details"]as a live-list reference — the one gap in an otherwise fully lock-hardened serialization surface (snapshot(),to_dict(),merge(),__copy__,__deepcopy__all correctly lock and defensively copy).dataclasses.asdict()also walks the livedetailslist unlocked (_lockis deliberately excluded from__dataclass_fields__, soasdictbypasses__getstate__entirely). No crash occurs (list iteration has no torn-iteration guard, unlike dict/set), but a torn/inconsistent snapshot is possible. -
src/xagent/core/model/chat/token_context.py:387—TokenUsage.merge(self)(self-merge) silently doubles all counters and rows — no identity guard exists. Confirmed: on an object withinput=10,output=5,llm_calls=1,details=1,a.merge(a)producesinput=20,output=10,llm_calls=2,details=2. No test covers self-merge. -
src/xagent/core/model/chat/token_context.py:859—aggregate_media_usage_by_modelsorts output by-quantitydescending, pushing all-zero-quantity groups (deliberately kept, per the PR's own design, as "evidence an unmeasured call happened") to the end. Not an active bug today (no consumer truncates the list), but the sort order sends the opposite signal from the stated design intent — a future "top N" consumer would silently drop exactly the evidence rows this design exists to preserve. -
src/xagent/core/tools/core/media_usage.py:112— Inconsistent positional-vs-keyword-only API discipline:record_media_call/add_media_usage(token_context.py) acceptmodel/call_type/model_idpositional-or-keyword in that order;record_media_usagehere makes the same three keyword-only, and in a different order (model,model_id,call_type). Minor consistency nit across three sibling entry points. -
src/xagent/core/model/chat/token_context.py:854—tokens_estimated: bool(and alsomodel/model_id/resolution) are not coerced/sanitized at the write boundary, unlikequantity/input_tokens/output_tokens. Sinceaggregate_media_usage_by_modelteststokens_estimatedfor truthiness, a stray non-bool value (e.g. string"no") would be misread asTrue— uniquely dangerous among the uncoerced fields due to this truthiness test. -
src/xagent/core/tools/core/media_usage.py:152—record_media_usage's warning log includes onlycall_type, omittingmodel/model_id/unit/quantity— hard to identify which specific producer/call failed when multiple media calls share a call_type. (Note: the bareexcept Exceptionitself was independently re-checked and is fine as designed —exc_info=Truealready surfaces the real exception type in the traceback.) -
tests/core/tools/core/test_media_usage_helpers.py:121— No test usescaplogto verifyrecord_media_usage's warning log actually fires with useful content; existing tests only assert on state (no row written). Also this test has an unusedmonkeypatchfixture parameter — dead test fixture, please remove. -
src/xagent/web/tracking/task_tracker.py:105—_task_seed_from_sessionnever seeds atool_callsbaseline (the Task ORM model has no persistedtool_callscolumn at all), so_initial_tool_callsis always 0 for a resumed multi-turn task — while_turn_delta's docstring (line 557) states the details-baseline and tool_calls-baseline "must describe the same instant." Pre-existing gap made more visible by this PR's new atomicity language. -
src/xagent/web/tracking/task_tracker.py:539—_turn_deltareturns an unparameterizedtuple[list, int]rather thantuple[list[dict[str, Any]], int]. -
src/xagent/core/model/chat/token_context.py:920—add_media_usage/record_media_usageboth always returnNone, success and swallowed-failure alike — a producer has no programmatic signal that its metering silently failed, only the (already under-informative, see comment 12) log line. -
src/xagent/core/model/chat/token_context.py:139andtests/core/model/chat/test_media_usage.py:338— Both a lock-justifying comment and a test docstring cite a helperbind_usage_to_threadas the original motivating use case for the lock — this symbol does not exist anywhere in the codebase (confirmed by full-repo grep). Stale/copy-pasted justification from a design that changed. -
src/xagent/core/model/chat/token_context.py:590—estimate_media_tokens's CJK character-range table was previously flagged for missing CJK punctuation/fullwidth-forms handling; that specific gap was fixed. But still incomplete: several Unicode blocks remain unhandled (CJK Extension B+, CJK Compatibility Ideographs, Bopomofo, Hangul Jamo), and fullwidth Latin characters (U+FF01–FF5E) are counted at the CJK 1-token/char rate instead of the ~4-chars/token Latin rate — a ~4x overcount for fullwidth-Latin content. No test covers these remaining gaps.
Previously-tracked findings (not newly blocking, referenced for context)
- Provider-token pricing precedence /
quota_hooks.pycontract not updated for media rows lacking a "tokens" key — deferred to #1461, confirmed still open through 3 prior review rounds. (Directly related to finding #4 above.) aggregate_media_usage_by_modellacks id/name reconciliation (unlikeaggregate_token_usage_by_model), so rows for the same model can incorrectly merge/split when only one carriesmodel_id— confirmed real, deferred to #1460.resolve_billing_model's default fallback returns"default", itself a forbidden placeholder per the module's own invariant — partially mitigated in #1425, deferred to #1460. Note: this PR's owntest_resolve_billing_model_never_returns_a_placeholderliterally asserts the placeholder is returned — worth folding this observation into the tracked issue.quantity=0.0conflates unmeasured/measured-zero/discarded-malformed states with no disambiguating flag — deferred to #1495.model/model_id/resolutionstring fields not normalized at the write boundary — accepted for this layer, producers land in #1424/#1425/#1457.media_callshas no DB-persisted column / no consumer yet — expected, tracked, consumer named in #997.
Simplification opportunities
native: src/xagent/core/tools/core/media_usage.py:40 — imports TypeGuard from typing_extensions, but the project requires Python >=3.11 (stdlib typing.TypeGuard since 3.10); other modules (type_check.py, jwt_validation.py) already import it natively. Use `from typing import TypeGuard`.
yagni: src/xagent/core/tools/core/media_usage.py:57 — resolve_billing_model has no production caller anywhere in the repo (only its own test file). Drop until a real caller needs it, or inline when it lands.
shrink: src/xagent/core/model/chat/token_context.py:98 — task_tracker.py's _copy_details (lines 53-56, 102, 171) duplicates this new shared copy_detail_rows helper, which this same PR added specifically to consolidate this duplication (previously flagged as reimplemented four times: snapshot(), merge(), to_dict(), and task_tracker._copy_details). Import and use copy_detail_rows in task_tracker.py; delete _copy_details.
yagni: src/xagent/core/tools/core/media_usage.py:156 — record_media_seconds (and coerce_duration) has no production caller anywhere in the repo (only its own test file). Drop until a real video/ASR/music caller needs it.
net: -10 lines possible (the import fix and dedup are immediate wins; the two yagni removals are better deferred to the producer PRs that will actually need them, rather than counted as lines removed today).
Blocking status & recommended decision
Blocking: yes — Recommended event: REQUEST_CHANGES
Blocking issues:
src/xagent/core/model/chat/token_context.py:633— major —_coerce_floatmissingOverflowErrorguard causes silent row-drop on huge/malformed quantity, violating the module's own contract[new]tests/core/model/chat/test_media_usage.py:336— major — locking fix is not actually verified by its own tests (mutation-tested: all concurrency tests pass with the lock fully removed)[new]tests/web/tracking/test_task_tracker.py:1120— major —detail_tailatomicity test still doesn't exercise the interleaving it claims to test (second round this exact gap surfaces)[new]src/xagent/core/model/chat/token_context.py:951— major (design) — token-vs-unit pricing precedence is unenforced and unrepresented in the row schema, undermining the PR's central pricing-key stability claim[new]tests/core/model/chat/test_media_usage.py:145— major — the module's own "unit is a property of the modality" invariant is unenforced and contradicted by this PR's own test data (unit="images"+call_type="video")[new]
| # all roughly one token per character. | ||
| code = ord(char) | ||
| if ( | ||
| 0x4E00 <= code <= 0x9FFF # CJK Unified Ideographs |
There was a problem hiding this comment.
Minor (claimed-fixed-but-incomplete). The previously-flagged gap (missing CJK punctuation/fullwidth-forms handling) was fixed, but several Unicode blocks remain unhandled: CJK Extension B+, CJK Compatibility Ideographs, Bopomofo, Hangul Jamo. Also, fullwidth Latin characters (U+FF01–FF5E) are counted at the CJK 1-token/char rate instead of the ~4-chars/token Latin rate — a ~4x overcount for fullwidth-Latin content. No test covers these remaining gaps.
Suggest: extend the range table for the listed blocks and add a fullwidth-Latin-specific branch, with test coverage for both.
There was a problem hiding this comment.
Partially addressed: CJK punctuation, fullwidth forms and Ext-A are now covered, with tests. You are right that Ext-B+, Compatibility Ideographs, Bopomofo and Hangul Jamo remain unhandled, and that fullwidth Latin (U+FF01-FF5E) is over-counted at the CJK rate. Both are real; the estimator is a heuristic feeding tokens_estimated=True rows that billing is documented to refuse to price, so I stopped at the blocks that appear in ordinary prose rather than chasing completeness in a PR with no estimator consumer. Leaving open.
| ) | ||
|
|
||
|
|
||
| def resolve_billing_model( |
There was a problem hiding this comment.
Simplification (yagni). resolve_billing_model has no production caller anywhere in the codebase (only its own test file). Suggest dropping until a real caller needs it, or inlining it into that caller when it lands, rather than shipping unused surface area now.
There was a problem hiding this comment.
True in this PR read alone, but resolve_billing_model is called by audio_tool and the Telegram ASR path in #1425 and by the memory store in #1457 — same series. Removing it here would move the diff rather than shrink it, and would make #1425 reintroduce an identical helper. Leaving open in case you would still rather it land with its first caller.
| ) | ||
|
|
||
|
|
||
| def record_media_seconds( |
There was a problem hiding this comment.
Simplification (yagni). record_media_seconds (and the coerce_duration helper it wraps) has no production caller anywhere in the codebase (only its own test file). Suggest dropping until a real video/ASR/music caller needs it, consistent with the "no producers wired up yet" scope of this PR.
There was a problem hiding this comment.
Same as resolve_billing_model: record_media_seconds/coerce_duration are the metering entry point for ASR, video, music and sound effects in #1425, and the duration-billed invariant they encode (always seconds, quantity=0 when unmeasured) is exactly what the new MEDIA_UNIT_BY_CALL_TYPE check now enforces. Deferring the helper without deferring the invariant would split one idea across two PRs. Leaving open.
… 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.
|
Round 5 addressed in f9eb91e. I reproduced all five blocking findings before Finding 2 — the locking tests, and a correction to my own numbersConfirmed: replacing Investigating why turned up something worse than the test being weak. My On CPython 3.12 The interval is restored in a Finding 3 —
|
rogercloud
left a comment
There was a problem hiding this comment.
Summary
This is PR 1 of 5 splitting #997 into reviewable chunks. It adds only the foundational vocabulary/data model for billing non-LLM media usage (images, video, TTS, ASR, music, embedding, rerank): MediaUnit/MediaCallType enums, TokenUsage.record_media_call/add_media_usage/record_media_usage, aggregate_media_usage_by_model, and thread-safety (locking) plus merge() on TokenUsage. No producers are wired up yet — those land in the 4 follow-up PRs. Media entries are appended as type:"media" dicts into the existing TokenUsage.details list to reuse existing DB persistence and the quota delta_details contract without a schema migration.
Round 0 (design) verdict: ACCEPTABLE-WITH-RESERVATIONS
Reusing details as a discriminated union and avoiding a schema migration is the right call, and the lock ordering in merge() has no deadlock. However, some of the concurrency machinery is justified by comments that assert facts about the codebase which are independently verified to be false today, and one primitive (merge()) has a real, if currently unreachable, correctness gap. Since this PR's entire purpose is to hand 4 follow-up PRs a trustworthy set of concurrency primitives and documentation to build on, factual accuracy of the justifying comments and actual verification of the concurrency tests both matter more than usual here.
Design-level findings
[Major, test-quality, [prior] PARTIAL fix] Three of the four concurrency tests in tests/core/model/chat/test_media_usage.py (test_concurrent_merge_loses_no_counts, test_concurrent_merges_in_both_directions_do_not_deadlock, test_snapshot_is_taken_under_the_lock) still pass even with TokenUsage._lock fully replaced by contextlib.nullcontext() (reproduced by patching the _lock property and running each test 3x). Only test_concurrent_snapshots_never_see_a_torn_counter_and_rows_pair, which forces the race window open with sys.setswitchinterval(1e-9), actually fails without the lock. This exact gap was raised in a prior round; the author's reply (round 5) confirmed via their own mutation testing that removing the lock left "all six concurrency tests passing" on CPython due to the GIL, and fixed only the one torn-counter test with the switch-interval technique, calling it "required, not decoration." The other three tests were not given equivalent treatment and still name/document themselves as verifying locking behavior they cannot detect the absence of. This is a partial fix, not a complete one.
- Recommend: apply the same forced-interleaving technique (
sys.setswitchinterval, or athreading.Barrierat the critical section) to the remaining three tests, or explicitly document/rename them as deadlock/smoke tests rather than lock-verification tests.
[Medium] merge() can double-count rows under genuinely concurrent, opposing merges. For a.merge(b) racing b.merge(a), if b.merge(a) completes first, a's later snapshot of b already contains a's own rows, so a merges its own rows back into itself. test_concurrent_merges_in_both_directions_do_not_deadlock (tests/core/model/chat/test_media_usage.py:508-511) tolerates this with a loose upper bound (<= 3 * per_side) rather than asserting the exact expected count, so it doesn't catch the defect. merge() has exactly one production caller (TokenContextManager.__init__), always invoked with zero args, so there is zero live risk today — but the primitive itself is not correct for concurrent bidirectional use.
- Recommend tightening the assertion to the exact expected count (so a future regression is caught) and/or documenting that
merge()is not safe for concurrent bidirectional merging.
[Minor, live but zero production impact, [prior] reminder] The documented delta_details contract in src/xagent/web/services/quota_hooks.py (a list of {"type","tokens","model"}) was not updated for media rows, which instead carry unit/quantity. A hook written to the documented contract doing d["tokens"] would KeyError on a media row; one doing .get("tokens", 0) would silently zero-bill media usage. No in-repo hook implementation exists yet and no producer emits media rows in this PR, so there's no live impact today. Already raised in a prior round and tracked in follow-up issue #1461 — not re-raising as a fresh blocking item, just flagging that it's still open.
[Minor, pre-existing, not introduced by this PR] This PR documents an invariant ("TokenUsage must be kept alive across thread hops for accounting to work") that has no enforcement mechanism and is already violated today by an unrelated, pre-existing code path: FunctionTool.run_json_async (src/xagent/core/tools/adapters/vibe/function.py) uses a bare run_in_executor without copy_context(), so translate_json.py's sync LLM-calling path silently loses token accounting (a fresh, empty TokenUsage is created and discarded in the executor thread). Confirmed pre-existing at the base commit. Worth a non-blocking follow-up issue since this PR is the one that writes the invariant down.
Line-level findings (inline comments posted separately, summarized here)
- [Major] Two design-justifying comments in
src/xagent/core/model/chat/token_context.pyassert facts about the codebase that are independently verified false: (a) the claim thatTokenUsageis "routinely shared across worker threads (RAG ingestion pools...)" — the RAG ingestion thread pool does not usecopy_context(), so no liveTokenUsagereference crosses into those threads, and the pooled function never calls any token-accounting function; all actual production writers ofTokenUsagerun on the asyncio event loop whererecord_media_callhas noawaitand is already atomic without a lock; (b) the claim that "existing callers depend on positional argument binding" forTokenUsage(...)— every existing constructor call in the repo at the base commit uses keyword args (or none); the only positional call anywhere is a new test this PR itself adds. The underlying design decisions (adding a lock; makingmedia_callskw_only+trailing) may still be reasonable defensive choices for the producer PRs to come, but the comments should say "will be needed once producer PRs add executor hops," not assert it's already true. - [Medium,
[prior], tracked as #1460]resolve_billing_model's module docstring states placeholder identities ("default"/"None"/"") must never be recorded as a model name, yet its ownfallbackparameter defaults to the literal string"default", andresolve_billing_model(None, None) == "default"is asserted as expected behavior in the tests. No live bug (zero production callers in this PR), contradiction already tracked in #1460. - [Minor]
test_snapshot_is_taken_under_the_lock's reader thread is joined with no timeout at all (sibling tests in the same file usejoin(timeout=...)), so a writer failure beforestop.set()would hang CI rather than fail cleanly. - [Minor]
MEDIA_UNIT_BY_CALL_TYPEcompleteness is verified via a hand-written 9-tuple parametrize list rather than iteratingMediaCallTypeitself, so a future new enum member added without a table entry would pass silently instead of failing the test. - [Minor, style] A private
_UNSET/_raw_quantitysentinel leaks throughrecord_media_call's public signature solely to preserve a pre-coercion value for one error message. - [Minor, style] One Chinese-language comment (line 1160) in
tests/web/tracking/test_task_tracker.py, in an otherwise English codebase.
Other minor observations (non-blocking, no action required)
record_media_usage→add_media_usage→TokenUsage.record_media_callre-runs validation/coercion at each layer (harmless, idempotent) and the parameter order for shared concepts differs slightly between layers (low risk since the differing layer is keyword-only).- New media-billing vocabulary lives in
token_context.py, a "chat/token"-named module — costs little to move, worth a note for a follow-up PR rather than this one. aggregate_media_usage_by_modeluses a flatter model-identity policy (model_id or model_name) than its sibling token aggregator's careful reconciliation — an unenforced docstring convention with no producers to validate it yet.TokenUsage.__getstate__readsdetailswithout acquiring the lock, inconsistent with every other reader in the class. Flagged in a prior round, still unfixed, low severity (no in-repo pickle-of-live-object caller today).- Test coverage gaps worth picking up opportunistically: the warning log on a swallowed recording error is never asserted;
quantity=True(bool) on a non-REQUESTS unit is untested; pickle/deepcopy tests only assertmedia_callssurvives, notdetailscontent;resolve_billing_model's"null"/whitespace/configured_id="default"variants are untested;_validated_media_unitreceiving a non-str/non-enum value is untested. - All new
src/code in this PR has zero production callers by design (this is the primitives-only PR) — already acknowledged in prior rounds, not new.
Verified but intentionally dropped (not defects)
Empty call_type bypassing the unit-consistency check, _coerce_float/_coerce_int collapsing invalid input to 0 colliding with the unmeasured-call sentinel, MEDIA_UNIT_BY_CALL_TYPE lacking a version tag, _coerce_int hardening also touching the pre-existing LLM token path, and RERANK forced to quantity=1 were all independently checked and confirmed to be intentional, already-discussed, and/or tested design decisions from prior rounds — not re-raising any of these.
Blocking status & recommended decision
Blocking: yes — REQUEST_CHANGES.
Two findings drive this: (1) three of the four concurrency tests do not detect TokenUsage's lock being removed, which contradicts the PR's own stated mutation-testing methodology and is only a partial fix of a previously-raised issue; (2) two core design-justifying comments assert facts about the codebase that are demonstrably false. Both go to the integrity of the exact thing this PR exists to deliver — trustworthy concurrency primitives and documentation for 4 follow-up PRs to build on — so they should be fixed before merge rather than deferred. The merge() double-count gap (Medium) is real but currently unreachable in production; recommend fixing the test assertion at minimum, tightening merge() itself if easy. All other findings (resolve_billing_model default contradiction #1460, quota_hooks contract drift #1461, __getstate__ lock gap, style/coverage notes) are already tracked or prior-known and should not block this PR on their own.
| # positional signature — TokenUsage(input, output, llm_calls, tool_calls, | ||
| # details) — keeps binding the way existing callers expect. Inserting it | ||
| # before tool_calls silently rebound their 4th and 5th arguments (tool | ||
| # count became media count, details became tool count) with no error. |
There was a problem hiding this comment.
This comment ("Inserting it before tool_calls silently rebound their 4th and 5th arguments... with no error") justifies the kw_only+trailing placement by claiming "existing callers depend on positional argument binding." That's not accurate at the base commit: every existing TokenUsage(...) constructor call in the repo uses keyword args (or none). The only positional call anywhere is the new test this PR adds (tests/core/model/chat/test_media_usage.py:451). The placement choice is still fine as a defensive default for future callers, but please reword this to say it protects future positional callers rather than asserting current callers rely on it.
| def __post_init__(self) -> None: | ||
| # Counter updates are read-modify-write, and one TokenUsage is routinely | ||
| # shared across worker threads (RAG ingestion pools, and the executor hops | ||
| # the producer PRs add); without it a concurrent read can observe a |
There was a problem hiding this comment.
"one TokenUsage is routinely shared across worker threads (RAG ingestion pools...)" is not accurate today: the RAG ingestion ThreadPoolExecutor usage does not use copy_context(), so no live TokenUsage reference actually crosses into those worker threads, and the pooled function itself never calls any token-accounting method. All current production writers of TokenUsage run on the asyncio event loop, where record_media_call's body has no await and is therefore already atomic without a lock. Adding the lock now as a forward-looking safeguard for the executor hops the producer PRs will add is reasonable, but please reword this to say "will be needed once producer PRs add executor hops" rather than asserting it's already the case.
| configured_id: Optional[str], | ||
| model: Any = None, | ||
| *, | ||
| fallback: str = "default", |
There was a problem hiding this comment.
The module docstring (around line 24) states that placeholder identities like "default"/"None"/"" must never be recorded as a model name, but this parameter's own default is the literal string "default", and resolve_billing_model(None, None) == "default" is asserted as expected in tests/core/tools/core/test_media_usage_helpers.py:88. No live impact yet since this function has zero production callers in this PR, but the contradiction is real. Already raised previously and tracked in follow-up issue #1460 — flagging as a reminder, not a new blocking item.
| for thread in threads: | ||
| thread.start() | ||
| for thread in threads: | ||
| thread.join() |
There was a problem hiding this comment.
This reader thread is joined with no timeout at all, unlike sibling tests in this file which use join(timeout=...). If the writer ever failed before calling stop.set(), this would hang CI instead of failing cleanly. Consider a timeout here and a try/finally around stop.set() for consistency with the other concurrency tests in this file.
|
|
||
|
|
||
| @pytest.mark.parametrize( | ||
| ("call_type", "unit"), |
There was a problem hiding this comment.
This hand-written 9-tuple list is a duplicate, maintained-by-hand encoding of MediaCallType's members rather than an iteration over the enum itself. A future new MediaCallType member added without a corresponding entry in MEDIA_UNIT_BY_CALL_TYPE would pass this test silently instead of being caught. Consider for ct in MediaCallType: assert ct.value in MEDIA_UNIT_BY_CALL_TYPE (or a set-equality check) instead.
| output_tokens: int = 0, | ||
| resolution: str = "", | ||
| tokens_estimated: bool = False, | ||
| _raw_quantity: Any = _UNSET, |
There was a problem hiding this comment.
_raw_quantity is a private, module-level _UNSET sentinel that leaks into this public method's signature solely so the module-level add_media_usage can preserve the pre-coercion value for one error message. It works, but it's a bit of an odd surface for a public method to carry. Not blocking, just a style note for a follow-up cleanup.
|
|
||
| def copy_then_pause(raw_details): | ||
| rows = real_copy(raw_details) | ||
| # Inside detail_tail's lock: details已读, tool_calls 还没读. |
There was a problem hiding this comment.
This inline comment is written in Chinese, in an otherwise English codebase/comment style. Consider translating to English for consistency.
|
Closing in favour of #1527, with the scope and a post-mortem recorded in Five review rounds without converging is the signal. The direction was never
#1527 carries the same billing vocabulary at 1496 lines across 5 files, Thanks for the five rounds. Findings 2, 3 and 5 of the last round in |
`test_malformed_provider_tokens_sanitize_without_losing_the_row` was carried over from xorbitsai#1422 and asserts the hardened `_coerce_int` behaviour — booleans and negatives coerced to 0, `OverflowError` caught. That hardening is part of the TokenUsage work deliberately split out of this PR, so on `main`'s `_coerce_int` the boolean and negative cases assert the wrong value and the infinity cases raise. Narrowed to the cases `_coerce_int` handles today (NaN, non-numeric, None), with a comment naming what is not yet sanitised and why: `_coerce_int` is shared with every LLM adapter, so changing it is not a media-billing change. Tracked in xorbitsai#1526. I dropped the sibling `test_coerce_int_change_applies_to_the_llm_token_path` when assembling this branch and missed this one — my importlib harness passed because it only exercises the media path, where these values never reach a real provider payload. CI ran the parametrised cases and caught it.
Self-review against the xorbitsai#1422 finding list, before this PR gets a review round. Three things it turned up. `_coerce_int` still raised `OverflowError` on `int(float("inf"))`. I had classified this as part of the TokenUsage hardening split out of xorbitsai#1422 and scoped a test away from it — that was wrong. Image providers pass `prompt_tokens` straight from provider JSON into this boundary, so an infinity in a malformed payload propagates out and the error-swallowing wrapper drops the whole billing row: a call that happened, lost to salvage one bad field. Reproduced through both `record_media_call` and `record_media_usage`. Adding `OverflowError` to the existing except clause is a one-word change with no behaviour difference for any current caller. Huge integers (`10**400`) are deliberately left passing through: unlike `inf` they serialise to valid JSON, so they cannot break persistence, and clamping them is a policy decision for the pricing consumer. Two guards had no test that failed when reverted — the same shape as the xorbitsai#1422 findings where a guard existed and nothing verified it. Mutation testing each guard against the committed tests found the `_coerce_int` overflow path and the bool/negative quantity rejection both uncovered. Both now have cases, and all ten guards in this PR are verified: reverting each one fails at least one test.
…tsai#1527) * feat: add media usage billing vocabulary and recording helpers Non-LLM media calls — image, video, TTS, ASR, music, sound effect, embedding, rerank — currently produce no usage records at all: providers bill and nothing is metered. This adds the shared vocabulary those producers record through. No producers are wired up here; they follow in separate changes. Replaces xorbitsai#1422, which grew to 2296 lines across five review rounds by absorbing a concurrency and serialization overhaul of TokenUsage alongside the billing work. See xorbitsai#1526 for the scope and the post-mortem. This slice is the billing vocabulary only, and deliberately leaves TokenUsage's threading semantics exactly as it found them. - `MediaUnit` / `MediaCallType` name the billable dimension and the modality that produced a row. `MEDIA_UNIT_BY_CALL_TYPE` pairs them and is enforced at the write boundary: the unit is a property of the modality, never of the response, so a duration-billed call whose length is unknown records `seconds` with `quantity=0` rather than degrading to `requests`. A (model, unit) price table is unusable otherwise. Stating this in a docstring was not enough last time — three docstrings said it and ten call sites in the test suite violated it. - `add_media_usage` writes a `type:"media"` row into the existing `TokenUsage.details` list, so media flows through the current TaskTracker persistence and quota `delta_details` paths with no schema change. Provider tokens are stored under `provider_tokens`, not `tokens`: the latter means "billable LLM tokens", and a consumer summing it must not pick up media counts. No pricing precedence between tokens and units is claimed or implied — that design is tracked in xorbitsai#1461. - `aggregate_media_usage_by_model` groups by (model, unit, call_type, resolution). Zero-quantity rows are kept deliberately: they are the only evidence that an unmeasured provider call happened, and dropping them would report `media_calls=0` for a task that really did make billable calls. - `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 use — a metering bug must never break the media call the user asked for. Quantities and provider token counts are validated at the write boundary, which is the last point they can be repaired: booleans, negatives, NaN, infinities and overflowing values all record 0 with a warning while the row survives, because the provider call happened and is billable even when its size is unknown. `MediaUnit.REQUESTS` additionally rejects any quantity but 1, since it is defined as exactly one call and a different value would make the billable quantity contradict the call count derived from the same row. Not included, and tracked separately: making the counter/rows pairing atomic under concurrency. TokenUsage has never been thread-safe — not for its pre-existing `tool_calls` or token counters either — and fixing that touches a class used well beyond media billing. Bundling it is what made xorbitsai#1422 unreviewable. * test: scope the provider-token test to what this PR actually guarantees `test_malformed_provider_tokens_sanitize_without_losing_the_row` was carried over from xorbitsai#1422 and asserts the hardened `_coerce_int` behaviour — booleans and negatives coerced to 0, `OverflowError` caught. That hardening is part of the TokenUsage work deliberately split out of this PR, so on `main`'s `_coerce_int` the boolean and negative cases assert the wrong value and the infinity cases raise. Narrowed to the cases `_coerce_int` handles today (NaN, non-numeric, None), with a comment naming what is not yet sanitised and why: `_coerce_int` is shared with every LLM adapter, so changing it is not a media-billing change. Tracked in xorbitsai#1526. I dropped the sibling `test_coerce_int_change_applies_to_the_llm_token_path` when assembling this branch and missed this one — my importlib harness passed because it only exercises the media path, where these values never reach a real provider payload. CI ran the parametrised cases and caught it. * fix: catch OverflowError in _coerce_int, and cover two unverified guards Self-review against the xorbitsai#1422 finding list, before this PR gets a review round. Three things it turned up. `_coerce_int` still raised `OverflowError` on `int(float("inf"))`. I had classified this as part of the TokenUsage hardening split out of xorbitsai#1422 and scoped a test away from it — that was wrong. Image providers pass `prompt_tokens` straight from provider JSON into this boundary, so an infinity in a malformed payload propagates out and the error-swallowing wrapper drops the whole billing row: a call that happened, lost to salvage one bad field. Reproduced through both `record_media_call` and `record_media_usage`. Adding `OverflowError` to the existing except clause is a one-word change with no behaviour difference for any current caller. Huge integers (`10**400`) are deliberately left passing through: unlike `inf` they serialise to valid JSON, so they cannot break persistence, and clamping them is a policy decision for the pricing consumer. Two guards had no test that failed when reverted — the same shape as the xorbitsai#1422 findings where a guard existed and nothing verified it. Mutation testing each guard against the committed tests found the `_coerce_int` overflow path and the bool/negative quantity rejection both uncovered. Both now have cases, and all ten guards in this PR are verified: reverting each one fails at least one test. * ref: derive the media billing unit from the modality Follows the review's option A: make a wrong (unit, modality) pair unrepresentable rather than validated. Each MediaCallType member now carries the unit it bills in, reachable as `.unit`, so the unit is a property of the modality and never of the response. `unit` is gone as a parameter from `record_media_call`, `add_media_usage` and `record_media_usage`, and `call_type` is required, which closes the omitted-call_type bypass structurally instead of by convention. `MEDIA_UNIT_BY_CALL_TYPE` and `_validated_media_unit` are no longer needed. Note this does change the previously-intentional permissive path: an omitted or empty `call_type` is now rejected rather than recorded with an unconstrained unit. `record_media_seconds` now refuses modalities that do not bill in seconds instead of writing a row whose quantity is a duration but whose unit says otherwise -- a silently mispriced row. `media_calls` becomes a property derived from the detail rows, so the counter cannot drift from the rows it counts. This also matters at the quota seam: `_copy_usage` rebuilds TokenUsage from an explicit kwarg list that does not mention `media_calls`, which would have zeroed a stored field while carrying the rows. Also in this change: - `tokens_estimated` requires a real `True`; `bool()` let `bool("no")` mark measured counts as estimates. - Provider token counts are floored at the media write boundary only, so a negative count cannot credit back tokens that were never spent, while the LLM path's `if input_tokens:` keeps seeing real values. - Model identities are stripped, so " sd " and "sd" bill as one model rather than two invoice lines. - `coerce_duration` catches OverflowError: `float()` of an oversized int raises it, and that is what json.loads yields for a large JSON number. - The fullwidth Unicode range is split on the character rather than the block. Fullwidth letters and digits tokenize like Latin (billing "ABCD" per character over-counted ~4x) while fullwidth punctuation appears in CJK text and tokenizes per character. - `_raw_quantity`/`_UNSET` are dropped; with no layer coercing above the write boundary they had no callers. - `resolve_billing_model`'s docstring states the `fallback="default"` hole honestly, tracked in xorbitsai#1460. Tests: adds coverage at the TaskTracker/quota seam, which is where this PR's claim that media rides the existing `details` list is actually cashed. Every guard above is mutation-tested -- reverting it fails a test -- after finding that several previously-claimed fixes had no test that could detect a regression. * fix: migrate two wrapper tests left on the old signature Both were invisible to my local check because the ad-hoc runner I was using skipped any test taking a fixture, so the two monkeypatch tests were never executed. Runner fixed to run them and to report anything it still cannot execute instead of passing over it silently. `test_record_media_usage_swallows_invalid_unit` also needed its premise changed, not just its call: it passed "tts" as the invalid value, but with the unit derived from the modality "tts" is now a *valid* call_type, so the row records correctly and the test was asserting the opposite of the truth. It now uses a real typo and is renamed for what it checks. * fix: three defects found reviewing my own option-A changes **Making media_calls derived broke `details: null`.** The property iterates `details`, so a persisted null -- harmless while the counter was a stored field, since nothing read the list -- now raises TypeError on every read of `media_calls`, `to_dict` and `merge`. `from_dict` passed `data.get("details", [])` straight through, so it constructed the broken object happily. Coerced at that boundary; non-list shapes get the same treatment, since `details` is append-only downstream and a str or dict would fail later and further from the cause. Not reachable in production today (`TokenUsage.from_dict` has no production callers and nothing assigns None to a live `details`), but it is a public classmethod and the fragility is new. **record_media_seconds raised into the caller's own except.** Producers call it from inside `try/except Exception` (`music_tool` lines 120/205) whose handler returns `success: False`. A wrong `call_type` would therefore convert a media call that had already generated audio and already been billed into a reported failure -- destroying the user's result to report an accounting mistake, which is the exact outcome the swallow boundary exists to prevent. Now logs at error level and drops the row. Verified across 16 hostile inputs that it never raises and that every degenerate duration lands as `quantity=0.0` in a JSON-serialisable row. **The fullwidth range still had a two-character hole.** U+FF5F and U+FF60, fullwidth white parentheses, fell between `<= 0xFF5E` and `>= 0xFF61` and billed as Latin. Same class of bug as the previous attempt at this line. The test now asserts over the whole U+FF01-FF5E block rather than by example, so a boundary edit cannot silently reclassify part of it. Local runner also gained caplog support. It had been silently skipping `test_record_media_seconds_warns_when_unmeasured` all session -- the same blind spot that put two stale monkeypatch tests into CI one commit ago. It now runs every test in these files with no skips. * fix: address the non-blocking minors from the approving review Five of the eight minors; the other three are tracked as follow-ups because they are separate units of work rather than small corrections. **Stale counter prose (finding 1).** Removing the lock wording last round left the surrounding claims describing a scalar counter that no longer exists: the docstring said this method bumps a counter and appends a row and that a reader can see one without the other, and the rejection comment still said "counter bump". With `media_calls` derived from `details`, the count and the per-model breakdown are two readings of one list and cannot drift. Rewritten around the real caveat -- a concurrent reader gets a non-snapshot view of `details` -- keeping xorbitsai#1526 as the tracking reference. This is the same defect as the original Major 1, reintroduced by my own change. **Wrong list-iteration failure mode (finding 8).** The aggregate docstring claimed appending during iteration risks a "list changed size during iteration" RuntimeError. It does not -- that is dict/set behaviour; a list silently yields a moving, non-snapshot view. Verified before rewriting. **Boolean provider tokens billed as one (finding 6).** `_coerce_int(True)` is 1, since bool subclasses int, so a provider returning a JSON boolean for a token count billed one token -- while `quantity` at the same boundary already rejected booleans. Added `_coerce_media_tokens`, which rejects bools then delegates. Deliberately not changed inside `_coerce_int`: its bool handling is pre-existing and intentional on the LLM path. **Resolution not normalized (finding 7).** `resolution` was stored verbatim beside two stripped fields while being part of the aggregate key, so ' 1K ' and '1K' billed as two line items for one tier and the padded row missed an exact price-table join. This is the sibling site I missed when stripping model/model_id. **Backfill test could not fail (finding 2).** The named row was written first, so setdefault seeded `model_name` and the backfill branch never ran; the test passed with that branch deleted. Rows reversed. The three new guards are mutation-checked: reverting each fails a test.
First of five changes splitting #997, which grew to 2841 lines across 42
files and became impractical to review. #997 is now a draft and will be
closed once this series lands.
This change adds the shared vocabulary and recording helpers for non-LLM
media usage. No producers are wired up yet — image, audio,
embedding/rerank and the frontend follow in separate PRs that each build
on this one.
What lands here
MediaUnit/MediaCallTypename the billable dimension and themodality that produced an entry. The unit is a property of the
modality, never of the response: a duration-billed call records
secondswithquantity=0when unmeasured rather than switching torequests, so a price table keyed on (model, unit) stays usable.add_media_usagewritestype:"media"entries into the existingTokenUsage.detailslist, so media flows through DB persistence andthe quota
delta_detailscontract without special-casing. Providertokens are stored under
provider_tokens, nottokens, so a consumersumming billable LLM tokens cannot pick them up.
aggregate_media_usage_by_modelgroups 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.pyaddsresolve_billing_model, which never recordsa placeholder identity such as
"default"or"None", plus theerror-swallowing wrappers producers will use.
Review feedback addressed from #997
unit/call_typeare 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.
record_media_usagestill swallows the resultingValueError— ametering bug must not break the user's media call — so a typo surfaces
as a missing row plus a warning, never as a mis-billed one.
TokenUsageis shared across workerthreads, where
+=is a read-modify-write that silently loses counts.All counter updates now hold a lock;
merge()snapshots the sourceunder its own lock before taking the target's, so concurrent merges in
both directions cannot deadlock.
Testing
tests/core/model/chat/test_media_usage.pyandtests/core/tools/core/test_media_usage_helpers.py. Fixes are verified bymutation — reverting a guard must fail its test — rather than by a case count
that drifts as the suite grows.