diff --git a/libs/core/langchain_core/callbacks/usage.py b/libs/core/langchain_core/callbacks/usage.py index 9fbcbb798d..17ad7781af 100644 --- a/libs/core/langchain_core/callbacks/usage.py +++ b/libs/core/langchain_core/callbacks/usage.py @@ -114,6 +114,11 @@ def get_usage_metadata_callback( ) register_configure_hook(usage_metadata_callback_var, inheritable=True) cb = UsageMetadataCallbackHandler() - usage_metadata_callback_var.set(cb) - yield cb - usage_metadata_callback_var.set(None) + token = usage_metadata_callback_var.set(cb) + try: + yield cb + finally: + # Always clear the context var, including when the with-block raises. + # Without finally, post-block model calls keep accumulating into cb + # (see #38989). + usage_metadata_callback_var.reset(token) diff --git a/libs/core/tests/unit_tests/callbacks/test_usage_callback.py b/libs/core/tests/unit_tests/callbacks/test_usage_callback.py index b583faecff..9d709dc079 100644 --- a/libs/core/tests/unit_tests/callbacks/test_usage_callback.py +++ b/libs/core/tests/unit_tests/callbacks/test_usage_callback.py @@ -1,3 +1,4 @@ +import contextlib from typing import Any from langchain_core.callbacks import ( @@ -120,3 +121,15 @@ async def test_usage_callback_async() -> None: callback = UsageMetadataCallbackHandler() _ = await llm.abatch(["Message 1", "Message 2"], config={"callbacks": [callback]}) assert callback.usage_metadata == {"test_model": total_1_2} + + +def test_usage_callback_clears_on_exception() -> None: + """Callback must stop tracking after with-block exits via exception (#38989).""" + llm = FakeChatModelWithResponseMetadata(messages=iter(messages), model_name="fake") + with contextlib.suppress(RuntimeError), get_usage_metadata_callback() as cb: + _ = llm.invoke("in block") + raise RuntimeError + + # Calls after the block must not accumulate into the previous callback. + _ = llm.invoke("outside block") + assert cb.usage_metadata == {"fake": usage1}