UsageCallbackHandler shares token/cost state across instances (class-level _model_usage)

[chuenchen309]

Bug

UsageCallbackHandler._model_usage (and price_map / llm_output) are class-level mutable attributes, so token/cost statistics are shared across every UsageCallbackHandler instance instead of being per-instance. A freshly constructed handler inherits the accumulated totals of any earlier one.

Reproduce

from unittest.mock import MagicMock
from langchain_core.outputs import LLMResult
from langchain_nvidia_ai_endpoints.callbacks import UsageCallbackHandler

result = LLMResult(
    generations=[],
    llm_output={"token_usage": {"total_tokens": 10, "prompt_tokens": 4, "completion_tokens": 6},
                "model_name": "meta/llama-3.1-8b-instruct"},
)

h1 = UsageCallbackHandler()
h1.on_llm_end(result)
print(h1.total_tokens)   # 10

h2 = UsageCallbackHandler()
print(h2.total_tokens)   # 10  ← expected 0; inherits h1's totals

Two independent trackers — e.g. two sequential get_usage_callback() contexts — therefore contaminate each other's token/cost counts.

Why this looks like a bug (not "just call reset()")

• Each handler creates its own self._lock, but that lock guards the class-shared _model_usage. A per-instance lock protecting shared state provides no real mutual exclusion — which only makes sense if _model_usage was intended to be per-instance too.
• The handler advertises itself as compatible with OpenAICallbackHandler / get_openai_callback, both of which are per-instance and need no reset between separate tracking sessions.

Fix direction

Construct _model_usage in __init__ (and give price_map / llm_output per-instance copies) so each handler owns its state. Happy to open a PR with a regression test if this direction looks right.

[chuenchen309]

Two closely related issues in the same UsageCallbackHandler / get_usage_callback area, for completeness — same "per-instance state kept at class/module scope" theme:

1. price_map leaks across contexts. UsageCallbackHandler.price_map is a mutable class attribute, and get_usage_callback does callback.price_map.update(price_map), mutating the shared dict in place. So a custom price_map passed to one context is inherited by every later one:

with get_usage_callback(price_map={"m": 1.0}):
    ...
with get_usage_callback() as cb:
    assert "m" not in cb.price_map   # fails — leaked from the first context

Fix: assign a fresh dict on the instance instead of mutating the class one — callback.price_map = {**callback.price_map, **price_map}.

2. Nested get_usage_callback() breaks outer tracking. On exit the manager runs usage_callback_var.set(None) unconditionally instead of restoring the previous handler, so leaving an inner context stops the outer one from tracking:

with get_usage_callback() as outer:
    with get_usage_callback():
        ...
    assert usage_callback_var.get() is outer   # fails — it's None

Fix: token = usage_callback_var.set(callback); try: yield callback; finally: usage_callback_var.reset(token) (also makes teardown exception-safe).

All three (this issue + the two above) share the same root pattern. Happy to bundle them into a single PR with regression tests if that's the preferred shape.

[chuenchen309]

Richard (@richardchen874-sys) Good question — looked at the source to answer accurately rather than guess: UsageCallbackHandler tracks total_cost/price_map (a per-1k-token cost table) and its docstring says it's "compatible with OpenAICallbackHandler" — LangChain's own well-known cost-tracking utility (the one people use inside with get_openai_callback() as cb: for real spend accounting). So this isn't debug-only tooling; it's designed to feed actual cost/billing numbers if you wire it up that way.

That makes the class-level leak worse than a cosmetic bug for your multi-model use case — with OpenAI-compatible endpoints mixed with provider-specific adapters like you're testing, a shared _model_usage/price_map across instances means costs from one model/tenant can silently bleed into another's totals, exactly the cross-contamination risk you flagged. Your suggested test (two providers/models running in parallel, asserting no cross-talk) is a good one — I'll fold that scenario into the regression tests when I get to implementing the fix.

[chuenchen309]

Thanks for the offer, but I'd rather not take credentials or move this off GitHub — no reflection on you, it's just a line I keep for anything involving someone else's API keys.

It shouldn't be needed here either. The defect is that _model_usage is class-level state, so the sharing is observable without any provider access at all: construct two handlers in one process, mutate the usage on one, read it back from the other. No network, no keys, no price maps. That's also the form a regression test would want to take, since a test that needs live multi-model access wouldn't be runnable in CI.

If your private-beta setup does surface something the in-process reproduction doesn't — for example the accounting going wrong only when the same model name is served by two different base URLs — a description of the shape here would be more useful than access, and it stays on the record for whoever picks this up.