fix: prevent memory leak by checking existing hooks

Avoid duplicate hook registration by checking if a ContextVar with
the same name is already registered in _configure_hooks. This prevents
the memory leak while maintaining the original API and import patterns.

Changes:
- Check _configure_hooks before registering new hooks
- Maintain original function signature and behavior
- Use proper token-based context variable management
- Keep imports inside function to satisfy linting

Author: 20ns

libs/core/langchain_core/callbacks/usage.py

@@ -88,16 +88,9 @@ def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
                     )
 
 
-# Module-level ContextVar that gets registered only once when the module is loaded
-_usage_metadata_callback_var: ContextVar[Optional[UsageMetadataCallbackHandler]] = (
-    ContextVar("usage_metadata_callback", default=None)
-)
-_hook_registered = False
-
-
 @contextmanager
 def get_usage_metadata_callback(
-    name: str = "usage_metadata_callback",  # Parameter kept for backward compatibility but not used
+    name: str = "usage_metadata_callback",
 ) -> Generator[UsageMetadataCallbackHandler, None, None]:
     """Get usage metadata callback.
 
@@ -106,9 +99,7 @@ def get_usage_metadata_callback(
 
     Args:
         name (str): The name of the context variable. Defaults to
-            ``'usage_metadata_callback'``. Note: This parameter is kept for
-            backward compatibility but is no longer used, as the context variable
-            is now defined at module level to prevent memory leaks.
+            ``'usage_metadata_callback'``.
 
     Example:
         .. code-block:: python
@@ -139,16 +130,25 @@ def get_usage_metadata_callback(
     .. versionadded:: 0.3.49
 
     """
-    # Register hook only once to prevent memory leak
-    global _hook_registered
-    if not _hook_registered:
-        from langchain_core.tracers.context import register_configure_hook
-        register_configure_hook(_usage_metadata_callback_var, inheritable=True)
-        _hook_registered = True
+    from langchain_core.tracers.context import register_configure_hook, _configure_hooks
+
+    # Check if this ContextVar is already registered to prevent memory leak
+    usage_metadata_callback_var: ContextVar[Optional[UsageMetadataCallbackHandler]] = (
+        ContextVar(name, default=None)
+    )
+    
+    # Only register if not already in the hooks list
+    already_registered = any(
+        hook[0].name == name for hook in _configure_hooks 
+        if hasattr(hook[0], 'name')
+    )
+    
+    if not already_registered:
+        register_configure_hook(usage_metadata_callback_var, inheritable=True)
 
     cb = UsageMetadataCallbackHandler()
-    token = _usage_metadata_callback_var.set(cb)
+    token = usage_metadata_callback_var.set(cb)
     try:
         yield cb
     finally:
-        _usage_metadata_callback_var.reset(token)
+        usage_metadata_callback_var.reset(token)