Skip to content

feat: add media usage billing primitives - #1422

Closed
OliverBryant wants to merge 9 commits into
xorbitsai:mainfrom
OliverBryant:feat/media-billing-primitives
Closed

feat: add media usage billing primitives#1422
OliverBryant wants to merge 9 commits into
xorbitsai:mainfrom
OliverBryant:feat/media-billing-primitives

Conversation

@OliverBryant

@OliverBryant OliverBryant commented Aug 17, 2026

Copy link
Copy Markdown
Contributor

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 / 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.
  • add_media_usage writes type:"media" 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 them up.
  • 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 wrappers producers will use.

Review feedback addressed from #997

  • Unknown unit/call_type 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.
    record_media_usage still swallows the resulting ValueError — a
    metering 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.
  • Counter races are fixed. TokenUsage is shared across worker
    threads, where += is a read-modify-write that silently loses counts.
    All counter updates now hold a lock; merge() snapshots the source
    under its own lock before taking the target's, so concurrent merges in
    both directions cannot deadlock.

Testing

tests/core/model/chat/test_media_usage.py and
tests/core/tools/core/test_media_usage_helpers.py. Fixes are verified by
mutation — reverting a guard must fail its test — rather than by a case count
that drifts as the suite grows.

Introduce the shared vocabulary and recording helpers for non-LLM media
usage (image/video/audio/embedding/rerank), with no producers wired up
yet — those follow in separate changes.

- `MediaUnit` / `MediaCallType` name the billable dimension and the
  modality that produced an entry. The unit is a property of the
  modality, never of the response: a duration-billed call records
  `seconds` with `quantity=0` when unmeasured rather than switching to
  `requests`, so a price table keyed on (model, unit) stays usable.
- `TokenUsage.add_media_usage` / `add_media_usage` write `type:"media"`
  detail entries into the existing `TokenUsage.details` list, so media
  flows through DB persistence and the quota `delta_details` contract
  without special-casing. Provider tokens are stored under
  `provider_tokens`, not `tokens`, so a consumer summing billable LLM
  tokens cannot pick up media counts.
- `aggregate_media_usage_by_model` groups by (model, unit, call_type,
  resolution). Zero-quantity entries are kept deliberately: they are the
  only evidence that an unmeasured provider call happened.
- `media_usage.py` adds `resolve_billing_model`, which never records a
  placeholder identity such as "default" or "None", plus the
  error-swallowing `record_media_usage` / `record_media_seconds`
  wrappers producers will use.

Unknown `unit`/`call_type` values are now rejected at the write
boundary. A typo would otherwise mint a new billing dimension that the
aggregator keys off, and a usage record cannot be repaired once
persisted. Validation runs before any mutation so a rejected call leaves
no counter incremented without a matching detail entry.

`TokenUsage` counter updates are now guarded by a lock. A single
instance is shared across worker threads, where `+=` is a
read-modify-write that silently loses counts; a 8x200 concurrent
recording test drops ~85% of them without this. `merge()` snapshots the
source under its own lock before taking the target's, so concurrent
merges in both directions cannot deadlock.

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment thread src/xagent/core/model/chat/token_context.py Outdated
Comment thread src/xagent/core/model/chat/token_context.py Outdated

@rogercloud rogercloud left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

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_float accepts booleans, negative values, NaN, infinities, and float-coercible strings, and the write path stores the result; record_media_seconds can also forward negative or non-finite values through seconds or 0.0. A malformed quantity from a future producer can therefore persist invalid billing data, with Infinity reaching serialized task/quota details. Validate finite, non-negative, unit-specific values before incrementing or appending (while keeping an absent duration as seconds=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 legacy tokens key but are mixed into the same details list; existing quota hooks assume token entries and index detail["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 discriminated delta_details union, 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_model returns 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 fake default billing 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. Adding media_calls does not update TaskTracker._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 report media_calls=0 even 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() and add_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 atomic record_media_usage operation 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 retain provider_input_tokens and provider_output_tokens, but the public aggregate only sums provider_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_name uses one untagged key, so an ID-backed row with model_id="foo" can merge with an unrelated name-only row with model="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 untagged str, but that value is a configured ID in one branch and a provider name in another, while the recording contract requires model=name and model_id=id. A natural future caller can put an ID in model and leave model_id empty, 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_usage does not accept or forward resolution, input_tokens, output_tokens, or tokens_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 a TypeError during argument binding before the wrapper's try block. 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. Inserting media_calls before legacy tool_calls/details silently remaps positional TokenUsage(input, output, llm, tool, details) calls, while the init-visible _lock at :114-116 can break generic dataclasses.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 before record_media_usage's best-effort try block. A provider object whose model_name or model descriptor 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=None is coerced to 0 with valid unit/call-type metadata, so the test never enters the wrapper's exception path and would still pass if the try/except were removed. Add invalid-unit/call-type cases and a monkeypatched add_media_usage that 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 exercises target.merge(source) and has no timeout or opposite-direction a.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_tokens documents an iterable of strings but accepts only list/tuple/set, so a generator returns 0; additionally, other // 4 makes 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:726major invalid/non-finite quantities can enter persisted billing details [new]
  • src/xagent/core/model/chat/token_context.py:204major mixed media/token details are incompatible with existing quota hooks [new]
  • src/xagent/core/tools/core/media_usage.py:78major unresolved model identity is billed under a placeholder fallback [new]

Recommended event: REQUEST_CHANGES

Comment thread src/xagent/core/model/chat/token_context.py
Comment thread src/xagent/core/model/chat/token_context.py
Comment thread src/xagent/core/tools/core/media_usage.py
Comment thread src/xagent/core/model/chat/token_context.py Outdated
Comment thread src/xagent/core/model/chat/token_context.py Outdated
Comment thread src/xagent/core/model/chat/token_context.py Outdated
Comment thread src/xagent/core/tools/core/media_usage.py Outdated
Comment thread tests/core/tools/core/test_media_usage_helpers.py
Comment thread tests/core/model/chat/test_media_usage.py
Comment thread src/xagent/core/model/chat/token_context.py Outdated
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"]`.
@OliverBryant

Copy link
Copy Markdown
Contributor Author

Worked through all 15 findings (review body + inline). Pushed in 50afcc4.
Point by point:

Fixed

N3 — invalid quantities (major). _coerce_float rejected only
non-numeric input, so booleans, negatives, NaN and inf all reached the
detail row. Confirmed NaN/inf are not JSON-serialisable, so they emit
literal NaN/Infinity into persisted task/quota details. All four now
record 0.0 with a warning, matching the existing "call happened but is
unmeasured" convention rather than dropping the record. coerce_duration
got the same finiteness guard, since inf > 0 is True and would otherwise
pass as a usable duration.

Validation also moved to the real write boundary. It lived only in the
module-level add_media_usage, so anyone holding a TokenUsage bypassed
it. My own self-review then caught that the two copies had already drifted:
the TokenUsage.add_media_usage method appended unvalidated quantities
that the module function rejected. There is now exactly one write path —
grep -c '"type": "media"' returns 1.

N2 — non-atomic counter/detail write. New record_media_call does both
under one lock acquisition; add_media_usage delegates to it with
count_call=False (that method never bumped the counter — two-step callers
do it themselves). Verified with a writer/reader interleaving test: 0 torn
snapshots over 500 concurrent writes.

N1 — media_calls lost in snapshots. Confirmed: both _copy_usage and
_task_seed_from_session reconstructed TokenUsage without it. Both now
carry it. The seed derives the count from the surviving media detail rows,
since tasks have no media_calls column. Checked this cannot double-bill:
_turn_delta slices details by _initial_details_len and never reads the
scalar, so the seed value is reporting-only.

N11 — lock breaks asdict/deepcopy/pickle. Confirmed all three raised
TypeError: cannot pickle '_thread.lock' object. init=False alone was not
enough — dataclasses.asdict walks __dataclass_fields__ and ignores
__getstate__. The lock now lives in __dict__ behind a property, with
__getstate__/__setstate__ dropping and recreating it. That also keeps
the private slot out of the public positional signature. All four now pass,
and there is a test.

N9 — wrapper drops optional metadata. record_media_usage now mirrors
resolution, input_tokens, output_tokens, tokens_estimated and
forwards them inside the guarded call. Confirmed the failure mode: passing
them raised TypeError at argument binding, before the try.

N12 — raising descriptor escapes the guard. Each provider attribute read
is now individually guarded, since resolve_billing_model runs before
record_media_usage's try. Test uses a model_name property that raises
and asserts it falls through to model.

N15 — estimate_tokens. Both halves confirmed and fixed: generators
returned 0 (now accepts any iterable of strings, non-iterables still return
0 rather than raising), and other // 4 returned 0 for any 1-3 character
Latin string (now rounds up, so non-empty never estimates 0; empty stays 0).

N13, N14 — test gaps. Both were fair. N13: added cases that monkeypatch
add_media_usage to raise, plus invalid unit/call-type, asserting no
propagation and no partial state — the previous quantity=None case
coerced to 0 and recorded fine, so it passed with the try/except deleted.
N14: added opposing concurrent a.merge(b) / b.merge(a) with a join
timeout and row-count assertions.

Deferred with issues

N6 — mixed detail union (raised as blocking). Deferred to #1461 with
evidence. No consumer in this repo indexes detail["tokens"]:

$ grep -rnE '\["tokens"\]|\[.tokens.\]' --include="*.py" src/xagent
(no matches)

record_usage forwards delta_details verbatim to an optional hook, and
there is no set_usage_record_hook call anywhere in src/ — stock
config registers nothing, so it is a no-op. The risk is real but lives in
an out-of-tree hook whose shape is not visible here; inventing a union
format from this side would mean either guessing a migration for a consumer
I cannot see, or filtering media entries out of the stream this work exists
to deliver. #1461 records what to confirm and how to close it.

N5, N7, N8a, N8b — contract shape. Deferred to #1460. All four are
public-API shape changes (tagged identity key, structured resolver return,
placeholder fallback, input/output token split) that only become observable
once producers mix the forms, and there is no in-tree pricing consumer to
design the output shape against. N8a is partially mitigated in #1425, where
audio_tool passes the provider class name as its fallback so cost attaches
to something real instead of "default".

Verification

Tests were written to fail on revert. Local test runs are not usable in my
environment — importing xagent.core.model.chat hangs on this machine, on
main as well as on this branch — so behavioural checks were done by
exercising the module directly via importlib, and I am relying on CI for
the suite rather than claiming a local pass I did not get.

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 rogercloud left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

  1. P1 — FIXEDcall_type=None is normalized to the empty optional value at the write boundary. Original comment (review 4949744795).
  2. P2 — FIXEDunit=None is rejected before mutation. Original comment (review 4949744795).
  3. P3 — FIXED — finite/non-negative quantity handling and the duration fallback are now present. Original comment.
  4. P4 — DROPPED (tracked, not fixed) — the discriminated delta_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.
  5. 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.
  6. P6 — FIXED — TaskTracker snapshots preserve media_calls, and seeds derive it from persisted media rows. Original comment.
  7. P7 — PARTIAL (minor)src/xagent/core/model/chat/token_context.py:242 and src/xagent/web/tracking/task_tracker.py:65: the canonical recording path is atomic, but the public two-step add_media_usage(count_call=False) plus increment_media_calls remains, and TaskTracker still copies fields/details without the usage lock. Use a lock-aware snapshot API and add an interleaving test. Original comment.
  8. 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.
  9. 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.
  10. 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.
  11. P11 — FIXED — the best-effort wrapper accepts and forwards all optional media metadata inside its guarded call. Original comment.
  12. P12 — PARTIAL (minor)src/xagent/core/model/chat/token_context.py:107 inserts media_calls before the legacy tool_calls/details positional fields, so old five-positional TokenUsage(...) 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.
  13. P13 — FIXED — descriptor reads are individually guarded against provider properties that raise. Original comment.
  14. P14 — FIXED — exception and invalid-metadata paths now have exercising tests. Original comment.
  15. P15 — PARTIAL (minor)tests/core/model/chat/test_media_usage.py:461 checks opposing-merge timeouts and lower bounds, but has no deterministic overlap barrier or exact media_calls/detail consistency assertions. Add those assertions and force the interleaving. Original comment.
  16. P16 — PARTIAL (minor)src/xagent/core/model/chat/token_context.py:467 catches only TypeError even 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:817majorREQUESTS quantities are not constrained to exactly one, so persisted quantity and call count can diverge. [new]
  • src/xagent/core/model/chat/token_context.py:839major — 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.

Comment thread src/xagent/web/tracking/task_tracker.py Outdated
Comment thread src/xagent/web/tracking/task_tracker.py Outdated
Comment thread src/xagent/core/model/chat/token_context.py Outdated
Comment thread tests/core/model/chat/test_media_usage.py
Comment thread src/xagent/core/model/chat/token_context.py
Comment thread src/xagent/core/model/chat/token_context.py
Comment thread src/xagent/core/model/chat/token_context.py
…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.
@OliverBryant

Copy link
Copy Markdown
Contributor Author

Round 2 addressed in 0ccbd35. Verified each finding against the code before
fixing; point by point:

Fixed

C1 — REQUESTS quantity (major). Confirmed: 0, 0.5, 2, 7 all
persisted while media_calls and the aggregate calls still counted one.
record_media_call now raises for any REQUESTS quantity but 1.0, before
any mutation — verified a rejection leaves media_calls == 0 and
details == [], no orphan counter. Raised rather than clamped, as you
suggested: silently rewriting hides the caller bug. record_media_usage
swallows it, so the best-effort contract holds. Other units keep free
quantities (checked seconds=2.5, images=4 still pass).

C4 — provider token counts (major). Confirmed all three parts:
True → 1, -100 → -100, and int(float("inf")) raising OverflowError.
That last one was worse than "drops a field" — it propagated out and
dropped the whole billable media row (details=0, calls=0), and
negatives persisted as a provider_tokens total of -99.

Fixed on _coerce_int itself rather than adding a media-only variant: the
LLM paths at :766-769 pass provider-reported values through the same
helper and had the identical overflow exposure, so a local fix would have
left the class of bug open next door. Booleans, negatives and non-finite
values now sanitize to 0 with a warning while the media row survives
the call happened and is billable by quantity.

P12 — positional ABI (was PARTIAL). You're right, and my previous test
was institutionalising the break: it asserted TokenUsage(1,2,3,4,5,[])
binds media_calls=4. Demonstrated the real damage — an old
TokenUsage(10,5,2,3,[{...}]) call bound media_calls=3,
tool_calls=[{...}], details=[]. The field is now appended after
details and kw_only=True (fine on the >=3.11 floor), so legacy
five-positional calls bind exactly as before. Test rewritten to guard that,
including that a sixth positional argument is still rejected.

P7 — lock-aware snapshot (was PARTIAL). Added
TokenUsage.snapshot(), which takes the lock and returns a detached copy;
_copy_usage now delegates to it. Verified 0 torn snapshots over 400
concurrent writes, and that the copy does not alias details.

P16 — estimator exception scope (was PARTIAL). Confirmed a custom
iterable raising RuntimeError mid-iteration escaped. Now catches any
exception from iteration and logs, since the documented contract is that
malformed input never breaks the call being measured. Tests cover both
raise-immediately and raise-mid-iteration iterables.

C5 — TaskTracker seam (minor). Added
test_media_rows_report_only_the_current_turn_delta, modelled on the
existing test_complete_tracking_reports_only_current_turn_delta: seeds two
prior-turn media rows, asserts the seed counts them (media_calls == 2),
records one current-turn row, asserts the periodic snapshot carries the
running total (== 3), and asserts both the progress gate and the
completion hook receive only this turn's row.

P15 — merge test (was PARTIAL). Agreed the old assertions were too
weak. Now uses a threading.Barrier to force both directions into the
window, and asserts exact counter/row agreement plus per-unit row survival
rather than >= 100. Ran 20 trials locally: no deadlock, no drift.

Deferred with an issue

C2 — unbounded details persistence. Filed as #1466 with the
escalation path (persisted delta cursor, or keyed media persistence) and my
read of when it becomes reachable — the embedding/rerank producers in #1457
are the first high-frequency appenders. Agreed it is a real O(n²) trend;
deferred because the copy-and-rewrite mechanism predates this work, nothing
appends media rows in this PR, and reshaping persistence is a much larger
change than the primitives it was raised against. Not re-litigating the
finding, just its timing.

P4/P5/P8/P9/P10 stay tracked in #1460 / #1461 as you recorded.

Self-review notes

Two things my own pass caught, neither reported:

  1. The legacy two-step API (increment_media_calls + add_media_usage)
    leaves an orphan counter when a REQUESTS rejection happens, because the
    caller owns that increment. Documented at the raise site; unreachable
    today since no in-tree producer uses the two-step path (grep confirms
    only comments reference it).
  2. A comment still described a counter decrement I had already replaced with
    count_call=False. Corrected.

Verification

ruff, isort 6.0.1 and mypy 1.19.0 (the pinned pre-commit versions) are
clean on the changed files. Behavioural checks were run by loading the
module directly via importlib — the pytest suite does not run on this
machine (importing xagent.core.model.chat hangs, on main too), so I am
relying on CI for the suite rather than claiming a local pass I did not get.
Each fix above was confirmed by reproducing the described failure first.

…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 rogercloud left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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_call now raises ValueError before any mutation when unit == MediaUnit.REQUESTS and quantity != 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) rejects bool, floors negatives to 0, and catches OverflowError from int(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_calls is now field(default=0, kw_only=True) declared after tool_calls/details (src/xagent/core/model/chat/token_context.py:104-112), restoring the legacy 5-positional-argument binding.
    • P16 — FIXED. estimate_tokens now catches broad Exception and logs via logger.warning with the exception message, instead of only TypeError.
    • P7 Part A — FIXED. TaskTracker._copy_usage now delegates to TokenUsage.snapshot(), which takes the lock and deep-copies detail dicts.
  • 43e019e — "test: split the dirty-quantity case by unit after the REQUESTS constraint"
    • Splits the dirty-quantity coverage by unit so the REQUESTS rejection path and the sanitize-to-0.0 path 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 exact media_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.

(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:

  1. No MediaUnit pricing consumer exists yet, so media prices at a literal $0 until one is built (tracked in #1460 / #1461).
  2. 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).
  3. The _lock / __getstate__ / __setstate__ / snapshot() machinery on TokenUsage is 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 FIXEDtests/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 FIXEDkw_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 FIXEDthreading.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:

  1. Delete the two-step media API (same root cause as P7 Part B above — one recommendation, not two). TokenUsage.add_media_usage as an instance method, the count_call parameter on record_media_call (src/xagent/core/model/chat/token_context.py:250), and increment_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. Keeping record_media_call (always counting) plus the module-level add_media_usage / record_media_usage entry points covers every real caller. Deleting it removes the P7 Part B defect and shrinks the diff at the same time.

  2. Reconsider the pickle/lock compatibility layer (__post_init__, the _lock property, __getstate__, __setstate__ at src/xagent/core/model/chat/token_context.py:116-146). This machinery exists only to keep TokenUsage picklable and dataclasses.asdict()-safe after introducing a lock, but nothing in the codebase pickles, deep-copies, or asdict()s a TokenUsage in 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 shared TokenUsage in _turn_delta on the interrupt_reason_for_quota quota-gating path (src/xagent/web/tracking/task_tracker.py:538-547). Major, and newly introduced by this PR; the sibling _copy_usage call site was fixed in 0ccbd35 but this one was not.

No longer blocking:

  • C1 [fixed], C4 [fixed] — both previously blocking findings are resolved by 0ccbd35 / 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.

Comment thread src/xagent/core/model/chat/token_context.py Outdated
Comment thread src/xagent/core/model/chat/token_context.py
Comment thread src/xagent/core/model/chat/token_context.py Outdated
Comment thread src/xagent/core/tools/core/media_usage.py
…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.
@OliverBryant

Copy link
Copy Markdown
Contributor Author

Round 3 addressed in 650663f. Reproduced each finding before fixing.

NEW-A (blocking) — fixed, and it was my regression

You're right, and the diagnosis is exact: this is the sibling of _copy_usage,
which I re-pointed through snapshot() last round for precisely this reason. I
fixed one call site and did not check the other in the same file — which is the
"did you check the sibling call sites" question I should have asked myself.

Reproduced with a forced interleaving (barrier between the two reads, write in
the window):

_turn_delta without snapshot: 200/200 inconsistent  (0 media rows vs tool_calls=1)
_turn_delta with snapshot:      0/200

Randomly-scheduled concurrency showed 0/4000 — the GIL makes each individual
read atomic, so the window is narrow enough that it would likely never have been
caught by a stress test. It needed the deterministic barrier to surface, which is
why I'm treating your read of the severity as correct rather than theoretical.

Fixed by snapshotting unconditionally, including when the caller already
passed a detached copy, so the two call sites cannot diverge again. Regression
test added at test_turn_delta_pairs_details_and_tool_calls_atomically, which
patches snapshot to widen the window and asserts the returned pair agrees
whichever side of the write it lands on.

Simplifications — both taken

Two-step API deleted. TokenUsage.add_media_usage, the count_call
parameter, and increment_media_calls are gone. Confirmed no production callers
(only the internal count_call=False delegation), and I'd already documented in
my own comment that the path leaves an orphan counter on rejection — so
"delete rather than harden" is clearly right. record_media_call plus the
module-level add_media_usage/record_media_usage cover every real caller;
tests migrated to the atomic call. That closes P7 Part B.

Pickle/lock layer narrowed rather than removed. I kept __getstate__/
__setstate__ and added __copy__/__deepcopy__ delegating to snapshot().
Reasoning: the part that is genuinely load-bearing is keeping _lock out of
__dataclass_fields__ — without that, dataclasses.asdict() raises
cannot pickle '_thread.lock' object, and putting it back as a field would also
re-break the positional ABI that P12 just fixed. Removing the getstate pair
while keeping the lock out of the fields would leave copy/pickle silently
producing a lock-less object, which is worse than the current surface. Routing
copy.copy to snapshot() is what actually removes NEW-B's failure mode. Happy
to go further and drop pickle support outright if you'd prefer — say so and I'll
do it.

NEW-B / NEW-C / NEW-D / NEW-E — fixed

All four confirmed empirically first:

  • NEW-B: copy.copy gave details is shared, _lock is not shared. Now
    routes through snapshot(); test asserts a write through the clone does not
    reach the original.
  • NEW-C: mutating to_dict()["details"][0] changed the live object to 999.
    Now [dict(item) for item in self.details], matching snapshot().
  • NEW-D: increment_media_calls(-5) produced media_calls == -5. Negative
    counts now ignored with a warning (increment_media_calls itself is deleted,
    so this lands on increment_tool_calls).
  • NEW-E: added exc_info=True. Agreed on the trade — unrepairable persisted
    rows justify the traceback.

Self-review found one more

merge had the same aliasing NEW-C describes in to_dict: it copied the outer
list but shared the inner dicts, so mutating one usage object silently rewrote
the other's billing rows (verified: 999 propagated). Fixed with the rest, test
added. Same sibling-call-site lesson as NEW-A, applied proactively this time.

Deferred

NEW-F#1495. Agreed it's a real observability gap and that
tokens_estimated sets the precedent, but the blended-count behaviour is
deliberate and asserted by test_zero_quantity_media_entries_stay_visible, and
separating "unmeasured" from "corrupt" changes the persisted row shape — so it
belongs with the first pricing consumer alongside #1460/#1461.

C2 stays in #1466. P4/P5/P8/P9/P10 keep their tracked disposition.

Verification

ruff, isort 6.0.1, mypy 1.19.0 clean on the changed files. Behavioural
checks by loading the module via importlib — the pytest suite doesn't run on
this machine (importing xagent.core.model.chat hangs, on main too), so CI is
the suite of record. Every fix above was confirmed by reproducing the failure
first and re-running after.

@rogercloud rogercloud left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 calls usage.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 one with self._lock: hold. The torn-read window is genuinely closed.
  • NEW-B (copy.copy lock aliasing) — FIXED. __copy__ now delegates to snapshot().
  • NEW-C (to_dict leaked live inner dicts) — FIXED. Now builds [dict(item) for item in self.details].
  • NEW-D (negative counter increments) — FIXED. increment_media_calls deleted outright; increment_tool_calls rejects negatives with a warning.
  • NEW-E (swallowed errors lost the traceback) — FIXED. exc_info=True added in media_usage.py.
  • P7 Part B (two-step media API) — FIXED. The instance-method TokenUsage.add_media_usage, record_media_call's count_call parameter, and increment_media_calls are all gone (repo-wide grep for increment_media_calls returns zero matches). Only the module-level add_media_usage remains, 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) seeds details from 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_quota is wired as the interrupt_checker (src/xagent/web/api/chat.py:3243-3245, src/xagent/web/api/websocket.py:3365-3366) and reached from runtime.py via should_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_delta has already paid the cost, so disabling the quota hook does not avoid it. The lock held during the copy is also the one serializing add_input_tokens / add_output_tokens / record_media_call from 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:

  1. token_context is a contextvars.ContextVar. read_delta() calls tracker._turn_delta() with usage=None, so _turn_delta calls get_token_usage() from inside a plain threading.Thread — which does not inherit the calling context. It lazily creates a brand-new empty TokenUsage, entirely disconnected from the usage object the writer thread mutates via closure. The read side therefore observes ([], 0) no matter what the implementation does.
  2. Against the pre-fix implementation, snapshot() was never called at all, so the patched slow_snapshot never fires, entered is never set, and the writer's entered.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 copy details, and from_dict (:368-377) stores the caller's list by reference. No exploitable path exists today — from_dict has no production caller, and the one production constructor call passing details= (task_tracker.py:102-117) hands over a fresh _copy_details result 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 and from_dict together.
  • N4 — string fields reach a JSON column unsanitized while numbers are hardened. record_media_call (token_context.py:221-283) writes model, model_id, resolution with no type check while the numeric fields go through _coerce_float/_coerce_int. details lands in Task.token_usage_details, a plain Column(JSON). A non-serializable value would not raise at the call site — record_media_usage's try/except only 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.Lock is a latent trap. _lock is a plain Lock (token_context.py:129). No self-deadlock exists today: to_dict deliberately inlines the total instead of calling total_tokens(), merge takes the two locks sequentially and never nested, and snapshot builds 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 from bool -> int(value) / negatives-passed-through to bool -> 0 / negatives -> 0 (token_context.py:455-481) affects add_token_usage, extract_cached_input_tokens, and aggregate_token_usage_by_model, which every adapter calls. To be clear, this was explicitly disclosed in 0ccbd35b'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 through add_token_usage. One small case would close it.
  • N7 — the REQUESTS error message reports a value the caller never passed. See the inline comment on token_context.py:262.
  • N8 — add_media_usage's "Raises" section omits the REQUESTS quantity ValueError. The docstring (token_context.py:872-877) lists only the unit/call_type errors, and the "Validate before touching the context" comment at :884-885 is true only for those; the REQUESTS constraint is enforced inside record_media_call (:260), which runs after get_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, because record_media_call is independently public.)
  • N9 — estimate_tokens is unexported and collides in name with an established helper. token_context.py:487 is the only new top-level symbol this PR adds that is not exported from src/xagent/core/model/chat/__init__.py (MediaUnit, MediaCallType, add_media_usage, aggregate_media_usage_by_model, aggregate_token_usage_by_model all are). Meanwhile CompactUtils.estimate_tokens already exists (src/xagent/core/agent/utils/compact.py:37) with a different signature and a different algorithm (chars // 4), with 13 call sites plus a MessageUtils.estimate_tokens wrapper. 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-527 covers 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_calls seed. See the inline comment on task_tracker.py:110.
  • N12 — an implicit, undocumented "details are always dicts" invariant. snapshot() filters non-dicts (token_context.py:302) while _turn_delta slices the filtered list by _initial_details_len. This is currently safe and unreachable — _initial_details_len is computed in start_tracking from _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; and from_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-6 still reads "Token usage tracking using contextvars … across LLM calls"; the file now also defines MediaUnit (:19-34), MediaCallType (:37-48), record_media_call, add_media_usage, and aggregate_media_usage_by_model.
  • N14 — the aggregators' "pass a detached list" contract is undocumented. aggregate_token_usage_by_model (:606-616) and aggregate_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 live usage.details and iterating while another thread appends.
  • N15 — test_concurrent_merge_loses_no_counts has little regression power. See the inline comment on test_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_id branch of aggregate_media_usage_by_model is untested. No test passes a media detail with a non-empty model_id into that aggregator: the two tests that set model_id on a media entry (test_media_usage.py:30,47) never aggregate, and the model_id uses at :126,157,180 feed the LLM-token aggregator instead. So identity = model_id or model_name (:748) always takes the name branch and the backfill at :753-767 never 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.py has 45 def test_ functions (74 with parametrize expansion), and tests/core/tools/core/test_media_usage_helpers.py has 11 functions (17 expanded). They undercount rather than overstate, but please correct or drop the numbers.
  • N19 (nit) — media_calls still has no consumer. Verified it has no reader outside token_context.py and the seed at task_tracker.py:113, is not persisted, does not enter delta_details or any API/frontend surface, and is fully derivable from details. 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 on TokenUsage at all, and the lock fixes a real read-modify-write race affecting tool_calls and 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:559major [new] — unconditional full-list deep copy in _turn_delta on the per-step/per-chunk quota path, held under the usage lock.
  • tests/web/tracking/test_task_tracker.py:1120major [new] — the sole regression test for the NEW-A fix cannot fail, because a bare threading.Thread does not inherit the contextvars context and the reader observes a different TokenUsage object.

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.

Comment thread src/xagent/web/tracking/task_tracker.py Outdated
Comment thread src/xagent/core/model/chat/token_context.py
Comment thread tests/web/tracking/test_task_tracker.py
Comment thread src/xagent/core/model/chat/token_context.py
Comment thread src/xagent/web/tracking/task_tracker.py Outdated
Comment thread tests/core/model/chat/test_media_usage.py Outdated
Comment thread tests/core/tools/core/test_media_usage_helpers.py Outdated
@OliverBryant

Copy link
Copy Markdown
Contributor Author

Round 4 addressed in a20c098. Both blocking findings were consequences of how
I fixed NEW-A last round — reproduced and measured each before fixing.

N1 (blocking) — confirmed and measured

snapshot() copied the whole cumulative list and _turn_delta then discarded
everything before the baseline. Measured cost of copying all versus the tail:

details=  100:  snapshot  10.6µs   tail 0.3µs     38x
details= 1000:  snapshot 103.2µs   tail 0.3µs    375x
details= 5000:  snapshot 487.2µs   tail 0.3µs   1671x

Also confirmed the two facts that make it matter: runtime.py:291
(_raise_if_interrupted) runs inside the streaming loop per chunk, and
interrupt_reason_for_quota is the registered interrupt_checker
(chat.py:3244, websocket.py:3366) — so it is per-chunk, not per-step as my
own docstring claimed. And the list grows monotonically because
_task_seed_from_session restores the persisted cumulative list.

Fixed with TokenUsage.detail_tail(start) exactly as you suggested: tail plus
tool_calls in one lock acquisition, O(delta). The redundant
_copy_details(delta_details) at both quota call sites is gone too.

One thing I had to keep: _turn_delta still subtracts _initial_tool_calls
itself, since detail_tail returns the cumulative counter and cannot know this
tracker's baseline. I dropped that subtraction in my first attempt and caught
it in self-review — it would have reported the whole run's tool calls as this
turn's delta.

N2 (blocking) — confirmed; my test was empty

Verified the mechanism directly:

plain Thread reading a ContextVar -> LookupError (not inherited)

So read_delta built a fresh empty TokenUsage and observed ([], 0)
regardless of the implementation, and against the pre-fix code the
slow_snapshot patch never fired so entered was never set and the writer's
wait() timed out silently with its return value unchecked. One assertion,
0 == 0 on both sides. You're right that this left the race unguarded.

The distinction that let it through: I mutation-verified the fix in a side
script
and reported 200/200 → 0/200, but never mutation-verified the test I
committed
. Those are not the same check. This round I did the latter, against
the committed test's logic:

FIXED   (detail_tail):      0/50 failures
BROKEN  (two unlocked reads): 50/50 failures

usage is now passed explicitly, and the test asserts entered.is_set() and
released.is_set() so a silent timeout fails rather than passes.

Minor findings — fixed

N3 (__post_init__ normalises details, covering the constructor and
from_dict), N5 (RLock), N6 (_coerce_int exercised through
add_token_usage), N7 (raw value in the REQUESTS message — verified it said
"got 0.0" for quantity=-1), N8/N12/N13/N14 (docstrings and the
details-are-dicts invariant), N9 (renamed estimate_media_tokens and exported
— confirmed three unrelated estimate_tokens already exist), N10 (CJK
punctuation, fullwidth, Ext-A: 中文,你好。 was 4, now 6), N11/N16 (both
comments were indeed false; N16 also gains the calls == 2 assertion), N15
(barrier + 8x100 merges), N17 (model_id branch covered).

Simplification taken: copy_detail_rows replaces the same expression
reimplemented four times. That also removes the asymmetry you noted — to_dict
lacked the isinstance filter, so a non-dict row would have raised there while
being skipped in snapshot/merge.

Self-review found two things I introduced this round

  1. Adding normalisation to __post_init__ made snapshot() copy twice —
    487µs → 1066µs at 5000 rows. snapshot now assigns details after
    construction: back to 526µs.
  2. mypy rejected threading.RLock as an annotation (it is a factory, not a
    class). Typed through a TYPE_CHECKING alias.

Also swept for the sibling-call-site question that caused NEW-A: no other
unlocked read of a live shared TokenUsage remains — the two persistence paths
both pass _copy_usage(...), and start_tracking reads initial_usage before
publishing it as the context.

Accepted for this layer

N4 (string fields reaching the JSON column uncoerced) and N19
(media_calls has no consumer): the producers supplying those strings land in
#1424/#1425/#1457, and #997 names the media_calls consumer. Happy to add
string coercion here if you'd rather not wait.

N18: dropped the stale counts from the PR body rather than restate numbers
that will drift again.

NEW-F stays in #1495, C2 in #1466, P4/P5/P8/P9/P10 in #1460/#1461.

Verification

ruff, isort 6.0.1, mypy 1.19.0 clean. Behavioural checks via importlib
(the pytest suite does not run on this machine — import xagent.core.model.chat hangs, on main too), so CI is the suite of record.

@OliverBryant
OliverBryant force-pushed the feat/media-billing-primitives branch from a20c098 to 88b30d5 Compare August 19, 2026 07:47
…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.
@OliverBryant
OliverBryant force-pushed the feat/media-billing-primitives branch from 88b30d5 to a489258 Compare August 19, 2026 08:04
@OliverBryant

Copy link
Copy Markdown
Contributor Author

Follow-up to my round-4 reply, after a further self-review pass. Nothing here
contradicts what you asked for — this is one code change, four refinements, one
scope question I would rather you decide than decide myself, and a disclosure
about a CI artefact I caused.

Head is now a4892581.

Refinements since the round-4 push

None of these came from your findings; they came from re-reading my own diff.
Listing them so the delta between 650663f6 and a4892581 is auditable.

1. snapshot() was copying details twice. Adding the N3 normalisation to
__post_init__ meant snapshot()'s own copy_detail_rows(...) result was
copied again by the constructor. Measured at 5000 rows:

before N3 with N3 (regressed) now
snapshot() 487µs 1066µs 526µs

snapshot now assigns details after construction. This is the persistence and
fork path, not the quota path, but doubling it while fixing a performance
finding would have been a poor trade.

2. _turn_delta briefly lost the tool_calls baseline. My first version
returned detail_tail's cumulative counter directly, dropping
- self._initial_tool_calls. That would have reported the whole run's tool
calls as this turn's delta on every quota poll. Caught before pushing;
_turn_delta keeps the subtraction because detail_tail cannot know a
tracker's baseline.

3. The RLock annotation carried a dead branch. I had written a
TYPE_CHECKING/else alias pair; threading._RLock exists at runtime too, so
the else was unreachable. Simplified to a direct annotation.

4. Two redundant test pairs merged (49 → 47 functions): the three CJK cases
became one parametrised case, and detail_tail_start_beyond_the_end_is_empty
folded into the main detail_tail test as one extra assertion. Net −11 lines —
small, and I am not claiming it offsets the growth discussed below.

An incidental correctness gain from the N1 fix

Worth flagging because it was not the point of the change. The old path was
snapshot().details[baseline:]filter, then slice. detail_tail is
slice, then filter. If a non-dict row ever sat before the baseline, the old
ordering shifted the boundary and dropped a real row:

details = ['JUNK', {'type': 'media', ...}], baseline = 1
old: snapshot().details[1:]  -> []                      # row lost
new: detail_tail(1)          -> [{'type': 'media',...}]  # correct

Not reachable today (that is your N12 invariant), but the new ordering is
robust to it rather than dependent on it.

Also verified while re-reading, no change made

  • Removing _copy_details at the quota call sites is safe. detail_tail
    returns a new outer list of new inner dicts; I confirmed a consumer mutating
    or appending to what it receives cannot reach live state. The two call sites
    each get their own list from separate _turn_delta calls.
  • _coerce_int's hardening is safe on the LLM path (your N6). Of its 13
    call sites, most already wrapped the result in max(0, ...), so negatives
    were never an intended input anywhere; and no adapter passes a bool.
  • __post_init__ normalisation breaks no caller. Only two production
    constructors pass details=, both handing over a freshly built list.
  • No other unlocked read of a live shared TokenUsage remains. This was the
    sibling-call-site sweep I should have done when I fixed _copy_usage: both
    persistence paths pass _copy_usage(...), and start_tracking reads
    initial_usage before publishing it as the context.
  • to_dict now also drops non-dict rows, which the previous inline copy did
    not. Deliberate — it removes the asymmetry you noted — and documented in code.
    I checked for consumers depending on the old passthrough; there are none.

The scope question — your call, not mine

You raised this in the round-4 design verdict, and I want to put numbers on it
rather than leave it as a vibe. Against the merge-base, this PR has grown:

at ffbecfed at a4892581
files 5 7
insertions 991 2120
test_media_usage.py 18 tests 47 tests

Every new symbol traces to a specific request — record_media_call (C1/P7),
detail_tail (N1), snapshot (P7), copy_detail_rows (your simplification
note), estimate_media_tokens (N9), __post_init__/_lock/__getstate__/
__setstate__ (N11), __copy__/__deepcopy__ (NEW-B) — and two were deleted
(P7 Part B). So each step was justified locally.

Taken together, though, a PR whose stated job is "shared vocabulary and
recording helpers, no producers" now also changes the concurrency and
serialization semantics of a class used across the codebase. That is your
reservation, and I think it is correct.

The concrete candidate is the pickle/copy layer:
__getstate__/__setstate__/__copy__/__deepcopy__ plus their tests —
25 source lines and 51 test lines, serving zero production callers (grep
confirms nothing pickles, deep-copies or asdicts a TokenUsage).

I did not remove it, because you wrote that __copy__/__deepcopy__ were
added at your request to fix NEW-B and "should stay as they are." Deleting them
right after that would be me overriding your explicit instruction on a
judgement call that is yours.

Two coherent end states, and I am happy with either:

  • Keep as-is. The lock genuinely needs to stay out of __dataclass_fields__
    (otherwise dataclasses.asdict raises, and moving it back would re-break the
    P12 positional ABI), so __getstate__/__setstate__ follow from that, and
    __copy__/__deepcopy__ close NEW-B.
  • Drop pickle support entirely — remove all four dunders and their tests
    (−76 lines). NEW-B's failure mode disappears with them, since it only exists
    because shallow copy routes through __getstate__. snapshot() remains the
    supported way to fork a usage object.

Say which and I will do it in one commit.

Disclosure: cancelled CI runs are my fault, not test failures

While re-chaining the five-PR stack I force-pushed repeatedly, which tripped the
workflow's cancel-in-progress concurrency rule. Several runs on #1425, #1457
and #1463 show as failed checks but are conclusion: cancelled against the
current head SHA — the giveaway is unexpanded ${{ matrix.name }} job names,
0-second durations and no step marked failure. All have been re-triggered. If
you see a red check while reviewing, please confirm the run conclusion before
treating it as a signal.

I also briefly wiped #1457's and #1463's own commits during that re-chain (a
rebase --onto against a base that no longer existed). Both were recovered from
reflog and verified: 9 and 5 branch-specific files respectively, matching their
pre-incident state. No review-relevant content was lost, but it is the kind of
thing worth stating plainly rather than having you notice a shifted diff.

Verification

ruff, isort 6.0.1 and mypy 1.19.0 (the pinned pre-commit versions) are
clean. Behavioural checks were run by loading the module through importlib
the pytest suite does not run on this machine (import xagent.core.model.chat hangs, on main as well), so CI is the suite of record.
The 47 no-argument and parametrised cases in test_media_usage.py and all 11 in
test_media_usage_helpers.py pass under that harness.

@rogercloud rogercloud left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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." But MediaUnit has no tokens member, while token-reporting providers (embedding, gpt-image) also populate provider_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, and aggregate_media_usage_by_model groups purely by (model, unit, call_type, resolution) with no field distinguishing "price by tokens instead." Zero consumers of provider_tokens exist anywhere outside this file, confirming the precedence rule is aspirational, not implemented. See inline comment on token_context.py:951.
  • [Design, non-blocking] TokenUsage.media_calls is a redundant field (no DB column; recomputed by row-counting on every session-resume seed) that is the direct cause of avoidable complexity (kw_only ABI-preservation dance, extra handling in merge/to_dict/from_dict/snapshot). Note: the RLock and most of its hardening would be needed regardless of media_calls, since the base commit had zero locking on the pre-existing tool_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: replace media_calls with a computed property (sum(1 for d in details if d["type"]=="media"), mirroring total_tokens), eliminating several touch-points.
  • [Informational] merge()'s snapshot-then-lock rewrite has one call site in src/ (TokenContextManager.__init__, token_context.py:475), but no instantiation in the repo passes parent_usage, so the path is functionally unreachable today. Also worth documenting: snapshot-then-lock means merge() captures other at one instant, not a live drain — rows appended between snapshot and extend are not included. Expected semantics, but the name merge could mislead a future caller into expecting a full drain.

Findings

Major

  1. src/xagent/core/model/chat/token_context.py:633_coerce_float catches only TypeError/ValueError, missing OverflowError that its sibling _coerce_int was already fixed for. _coerce_float(10**400) raises uncaught. Effect: a huge/malformed quantity from a producer, routed through record_media_usage's bare except Exception, causes the entire billing row to be silently dropped rather than recorded as quantity=0.0 — violating _coerce_float's own documented "reject-to-0.0, never drop" contract. Any direct caller of add_media_usage/TokenUsage.record_media_call bypassing the wrapper crashes uncaught. No test covers a huge-int quantity.
    Fix: add OverflowError to the except clause, mirroring _coerce_int; add a huge-int regression test.

  2. 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-patching TokenUsage._lock to contextlib.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 via sys.setswitchinterval, and assert on a materially larger observed-loss threshold.

  3. 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_atomically still does not test detail_tail's atomicity, despite the author previously replying that this exact concern (from an earlier review round) was fixed. Root cause: the slow_detail_tail mock sets entered/waits on released before delegating to the real detail_tail accessor, 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 splitting detail_tail's single with self._lock: block into two separate lock acquisitions with a 50ms sleep between reading details and tool_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.

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

  5. 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_call validate unit and call_type independently against their own enums, with no MEDIA_UNIT_BY_CALL_TYPE-style mapping. Proof: this PR's own tests violate the invariant on day one — unit="images" paired with call_type="video" appears at lines 145, 610, 632, 635, 671, 677 (six occurrences), even though record_media_seconds's own docstring says video must always report seconds. This directly contradicts the PR's central pricing-stability claim, in its own test suite.

Minor (non-blocking, confirmed)

  1. src/xagent/core/model/chat/token_context.py:968add_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 calling record_media_call, discarding the true raw value one layer up). Previously flagged; the prior fix covers only the direct TokenUsage.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.

  2. src/xagent/core/model/chat/token_context.py:170__getstate__ does a shallow, unlocked copy (self.__dict__.copy()), leaving state["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 live details list unlocked (_lock is deliberately excluded from __dataclass_fields__, so asdict bypasses __getstate__ entirely). No crash occurs (list iteration has no torn-iteration guard, unlike dict/set), but a torn/inconsistent snapshot is possible.

  3. src/xagent/core/model/chat/token_context.py:387TokenUsage.merge(self) (self-merge) silently doubles all counters and rows — no identity guard exists. Confirmed: on an object with input=10,output=5,llm_calls=1,details=1, a.merge(a) produces input=20,output=10,llm_calls=2,details=2. No test covers self-merge.

  4. src/xagent/core/model/chat/token_context.py:859aggregate_media_usage_by_model sorts output by -quantity descending, 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.

  5. 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) accept model/call_type/model_id positional-or-keyword in that order; record_media_usage here makes the same three keyword-only, and in a different order (model, model_id, call_type). Minor consistency nit across three sibling entry points.

  6. src/xagent/core/model/chat/token_context.py:854tokens_estimated: bool (and also model/model_id/resolution) are not coerced/sanitized at the write boundary, unlike quantity/input_tokens/output_tokens. Since aggregate_media_usage_by_model tests tokens_estimated for truthiness, a stray non-bool value (e.g. string "no") would be misread as True — uniquely dangerous among the uncoerced fields due to this truthiness test.

  7. src/xagent/core/tools/core/media_usage.py:152record_media_usage's warning log includes only call_type, omitting model/model_id/unit/quantity — hard to identify which specific producer/call failed when multiple media calls share a call_type. (Note: the bare except Exception itself was independently re-checked and is fine as designed — exc_info=True already surfaces the real exception type in the traceback.)

  8. tests/core/tools/core/test_media_usage_helpers.py:121 — No test uses caplog to verify record_media_usage's warning log actually fires with useful content; existing tests only assert on state (no row written). Also this test has an unused monkeypatch fixture parameter — dead test fixture, please remove.

  9. src/xagent/web/tracking/task_tracker.py:105_task_seed_from_session never seeds a tool_calls baseline (the Task ORM model has no persisted tool_calls column at all), so _initial_tool_calls is 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.

  10. src/xagent/web/tracking/task_tracker.py:539_turn_delta returns an unparameterized tuple[list, int] rather than tuple[list[dict[str, Any]], int].

  11. src/xagent/core/model/chat/token_context.py:920add_media_usage/record_media_usage both always return None, 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.

  12. src/xagent/core/model/chat/token_context.py:139 and tests/core/model/chat/test_media_usage.py:338 — Both a lock-justifying comment and a test docstring cite a helper bind_usage_to_thread as 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.

  13. src/xagent/core/model/chat/token_context.py:590estimate_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.py contract 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_model lacks id/name reconciliation (unlike aggregate_token_usage_by_model), so rows for the same model can incorrectly merge/split when only one carries model_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 own test_resolve_billing_model_never_returns_a_placeholder literally asserts the placeholder is returned — worth folding this observation into the tracked issue.
  • quantity=0.0 conflates unmeasured/measured-zero/discarded-malformed states with no disambiguating flag — deferred to #1495.
  • model/model_id/resolution string fields not normalized at the write boundary — accepted for this layer, producers land in #1424/#1425/#1457.
  • media_calls has 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_float missing OverflowError guard 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_tail atomicity 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]

Comment thread src/xagent/core/model/chat/token_context.py Outdated
Comment thread tests/core/model/chat/test_media_usage.py Outdated
Comment thread tests/web/tracking/test_task_tracker.py
Comment thread src/xagent/core/model/chat/token_context.py Outdated
Comment thread tests/core/model/chat/test_media_usage.py Outdated
# all roughly one token per character.
code = ord(char)
if (
0x4E00 <= code <= 0x9FFF # CJK Unified Ideographs

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment thread src/xagent/core/tools/core/media_usage.py Outdated
)


def resolve_billing_model(

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment thread src/xagent/core/model/chat/token_context.py
)


def record_media_seconds(

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.
@OliverBryant

Copy link
Copy Markdown
Contributor Author

Round 5 addressed in f9eb91e. I reproduced all five blocking findings before
touching anything — none was a false positive, and three were about claims I
made rather than code I wrote. Details below, including one place where
investigating your finding showed an earlier claim of mine was simply wrong.

Finding 2 — the locking tests, and a correction to my own numbers

Confirmed: replacing _lock with nullcontext() left all six concurrency
tests passing.

Investigating why turned up something worse than the test being weak. My
claim in an earlier round that the lock prevents "~85% count loss under 8x200
increments (231 of 1600)" was wrong. That figure came from a mutation where
I had inserted sleep(0) between the read and the write — not from the real
implementation. Measured on this build:

bare  self.n += 1   8x200000 threads-iterations : 0 lost
increment_tool_calls 8x20000  (nullcontext)     : 0 lost
record_media_call    16x5000  (nullcontext)     : 0 lost
read-sleep-write variant 8x200                  : 1353 of 1600 lost

On CPython 3.12 += is effectively atomic in practice, so no
increment-counting test can justify this lock
, however long you make it. The
lock's real job is the pairing: record_media_call bumps the counter and
appends the row as one operation, and an unlocked reader lands between them.
That is now what the test asserts, under sys.setswitchinterval(1e-9)
required, not decoration:

with lock,    setswitchinterval(1e-9) : 0/6 runs tear
without lock, setswitchinterval(1e-9) : 6/6 runs tear
without lock, default interval        : 0/10 runs tear   <- why the old test passed

The interval is restored in a finally.

Finding 3 — detail_tail atomicity, second time around

You are right, and the diagnosis is exact. Both my earlier versions wrapped
detail_tail itself and set/waited before delegating to the real accessor, so
the writer always finished before the locked read began. My round-4 fix
addressed the silent-timeout bug and I reported "50/50 → 0/50" — but I had
verified my own script's logic, not the mock I committed.

The interleaving is now forced from inside the lock, by patching
copy_detail_rows, which detail_tail calls after taking the lock and reading
details but before reading tool_calls. Mutation-verified against your torn
implementation (two acquisitions, sleep between):

atomic detail_tail : media=0 tool_calls=0  -> passes
torn  detail_tail  : media=0 tool_calls=1  -> fails

Finding 5 — the invariant is now enforced, and there were ten sites

Added MEDIA_UNIT_BY_CALL_TYPE and a check at the write boundary, so a unit
that does not match its modality is rejected rather than merely discouraged by
three docstrings.

Your grep found six images+video sites. Turning the invariant into code
surfaced four more the grep pattern missed — seconds paired with tts at
lines 128, 181, 242 and 243. All ten fixed. This is the clearest evidence for
your point: an invariant that is only documented gets violated, and the
violations are not all findable by pattern.

I verified the new constraint breaks no downstream producer: every
unit/call_type pairing in #1424/#1425/#1457 is legal, including the implicit
SECONDS from record_media_seconds (used with ASR, VIDEO, MUSIC,
SOUND_EFFECT — all mapped to SECONDS).

Finding 1 — _coerce_float and OverflowError

Confirmed: _coerce_float(10**400) raised uncaught, so the wrapper dropped the
whole billing row (details=0, calls=0) — the opposite of the documented
reject-to-0.0 contract. _coerce_int was fixed for exactly this two rounds ago
and I did not check the sibling. Now catches it; regression test covers
10**400, -(10**400) and "1e400".

Finding 4 — taking option B: the claim is removed, not implemented

I agree the precedence rule was aspirational. Rather than add a
price_basis discriminator to the row schema, I deleted the docstring sentence
that promised it. Reasoning: this PR wires no producers and there is no pricing
consumer, so implementing a pricing dimension here would be designing against
nothing — and the sentence was the actual defect, since it described behaviour
that does not exist. The docstring now says plainly that provider_tokens is
recorded raw for a future consumer and that no precedence is expressed or
enforced, pointing at #1461 where the design belongs.

If you would rather have the discriminator now, say so and I will add it.

Minor findings

Fixed: #6 (add_media_usage now reports the caller's raw quantity — the prior
fix covered only the direct path, as you noted), #12 (warning log now carries
model/model_id/unit/quantity), #15 (parameterised return type), #17 (both stale
bind_usage_to_thread references removed — confirmed the symbol does not
exist), and the TypeGuard/_copy_details simplifications.

Accepted or deferred, with reasons:

On the pattern

Three of five blocking findings were "the author said this was verified and it
was not." The common thread is that I verified the right concept against the
wrong object — a side script instead of the committed test, twice. Everything
in this round is mutation-verified against the committed test, with the
before/after numbers above so you can check the claim rather than take it.

Verification

ruff, isort 6.0.1, mypy 1.19.0 clean. 50 cases across
test_media_usage.py and test_media_usage_helpers.py pass under an
importlib harness (the pytest suite does not run on this machine — import xagent.core.model.chat hangs, on main too), so CI remains the suite of
record.

@rogercloud rogercloud left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 a threading.Barrier at 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.py assert facts about the codebase that are independently verified false: (a) the claim that TokenUsage is "routinely shared across worker threads (RAG ingestion pools...)" — the RAG ingestion thread pool does not use copy_context(), so no live TokenUsage reference crosses into those threads, and the pooled function never calls any token-accounting function; all actual production writers of TokenUsage run on the asyncio event loop where record_media_call has no await and is already atomic without a lock; (b) the claim that "existing callers depend on positional argument binding" for TokenUsage(...) — 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; making media_calls kw_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 own fallback parameter defaults to the literal string "default", and resolve_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 use join(timeout=...)), so a writer failure before stop.set() would hang CI rather than fail cleanly.
  • [Minor] MEDIA_UNIT_BY_CALL_TYPE completeness is verified via a hand-written 9-tuple parametrize list rather than iterating MediaCallType itself, 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_quantity sentinel leaks through record_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_usageadd_media_usageTokenUsage.record_media_call re-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_model uses 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__ reads details without 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 assert media_calls survives, not details content; resolve_billing_model's "null"/whitespace/configured_id="default" variants are untested; _validated_media_unit receiving 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.

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

"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",

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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()

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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"),

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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,

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

_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 还没读.

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This inline comment is written in Chinese, in an otherwise English codebase/comment style. Consider translating to English for consistency.

@OliverBryant

Copy link
Copy Markdown
Contributor Author

Closing in favour of #1527, with the scope and a post-mortem recorded in
#1526.

Five review rounds without converging is the signal. The direction was never
the problem — reusing TokenUsage.details rather than building a parallel
meter held up across every round. What failed was the execution, in three
specific ways worth stating plainly since they were mine:

  1. Fix-A-introduces-B, three rounds running. Round 3's _turn_delta race
    came from round 2's snapshot() fix landing on the sibling call site only.
    Round 4's O(n²) hot-path copy came from round 3's fix reusing snapshot()
    for a tail read. Round 5's _coerce_float overflow was the same bug I had
    already fixed in _coerce_int two rounds earlier. Each time I verified the
    target bug was fixed and never asked which sibling had the same shape.

  2. Verification against the wrong object. I mutation-tested fixes in
    scratch scripts, then committed different tests. Two of those committed
    tests passed against deliberately broken implementations — the
    detail_tail atomicity test twice, across two rounds. Worse, my claimed
    "~85% count loss without the lock" was wrong: it came from a mutation
    where I had inserted sleep(0). On CPython 3.12 a bare += loses nothing
    even at 8×200000, so that number never described the real code.

  3. Scope creep, 991 → 2296 lines. Every addition traced to a review
    request, which is precisely how it went unnoticed. A PR whose stated job
    was "shared vocabulary, no producers" ended up rewriting the concurrency
    and serialization semantics of a class used across the codebase, ~184
    lines of which serve zero production callers.

#1527 carries the same billing vocabulary at 1496 lines across 5 files,
doesn't touch task_tracker.py, enforces the unit-per-modality invariant in
code from the first commit rather than in docstrings, and drops the pricing
claim that was never implemented. The TokenUsage concurrency work is real
and will get its own PR with an honest justification — the counter/rows
pairing, not a fabricated loss figure.

Thanks for the five rounds. Findings 2, 3 and 5 of the last round in
particular were things I had asserted were verified, and being shown
otherwise with a reproduction each time is what made the pattern visible.

OliverBryant added a commit to OliverBryant/xagent that referenced this pull request Aug 20, 2026
`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.
OliverBryant added a commit to OliverBryant/xagent that referenced this pull request Aug 20, 2026
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.
OliverBryant added a commit to OliverBryant/xagent that referenced this pull request Aug 21, 2026
…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.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants