fix: release the token counter when CompactionHook is closed (#12985)

Co-authored-by: anakin87 <[email protected]>

haystack/hooks/compaction/hooks.py

@@ -290,24 +290,24 @@ class CompactionHook:
         """Warm up the token counter and the compactor on the serving event loop."""
         if hasattr(self.token_counter, "warm_up"):
             self.token_counter.warm_up()
-        warm_up_async = getattr(self.compactor, "warm_up_async", None)
-        if warm_up_async is not None:
-            await warm_up_async()
+        if hasattr(self.compactor, "warm_up_async"):
+            await self.compactor.warm_up_async()
         elif hasattr(self.compactor, "warm_up"):
             self.compactor.warm_up()
 
     def close(self) -> None:
-        """Release the compactor's resources."""
-        if hasattr(self.compactor, "close"):
-            self.compactor.close()
+        """Release the token counter's and the compactor's resources."""
+        for resource in (self.token_counter, self.compactor):
+            if hasattr(resource, "close"):
+                resource.close()
 
     async def close_async(self) -> None:
-        """Release the compactor's async resources."""
-        close_async = getattr(self.compactor, "close_async", None)
-        if close_async is not None:
-            await close_async()
-        elif hasattr(self.compactor, "close"):
-            self.compactor.close()
+        """Release the token counter's and the compactor's async resources."""
+        for resource in (self.token_counter, self.compactor):
+            if hasattr(resource, "close_async"):
+                await resource.close_async()
+            elif hasattr(resource, "close"):
+                resource.close()
 
     def to_dict(self) -> dict[str, Any]:
         """

releasenotes/notes/compaction-hook-close-token-counter-0495043806039056.yaml

@@ -0,0 +1,6 @@
+---
+fixes:
+  - |
+    Fixed ``CompactionHook.close()`` and ``close_async()`` to release the token counter's
+    resources as well as the compactor's. Previously, resources such as
+    ``OpenAITokenCounter``'s HTTP client remained open after the hook or Agent was closed.

test/hooks/compaction/test_hooks.py

@@ -4,6 +4,7 @@
 
 import logging
 from typing import Annotated, Any
+from unittest.mock import Mock
 
 import pytest
 
@@ -387,11 +388,14 @@ class TestCompactionHook:
         )
         assert compacted[-2:] == messages[-2:]
 
-    def test_lifecycle_delegates_to_the_compactor(self):
+    def test_lifecycle_delegates_to_the_counter_and_compactor(self):
+        counter = Mock(spec=["warm_up", "close"])
         compactor = _RecordingCompactor()
-        hook = _hook(compactor)
+        hook = _hook(compactor=compactor, token_counter=counter)
         hook.warm_up()
         hook.close()
+        counter.warm_up.assert_called_once_with()
+        counter.close.assert_called_once_with()
         assert compactor.calls == ["warm_up", "close"]
 
 
@@ -422,11 +426,14 @@ class TestCompactionHookAsync:
         assert compactor.calls == ["compact_async"]
 
     @pytest.mark.asyncio
-    async def test_lifecycle_prefers_the_async_methods(self):
+    async def test_lifecycle_prefers_async_methods_with_sync_fallback(self):
+        counter = Mock(spec=["warm_up", "close"])
         compactor = _RecordingCompactor()
-        hook = _hook(compactor)
+        hook = _hook(compactor=compactor, token_counter=counter)
         await hook.warm_up_async()
         await hook.close_async()
+        counter.warm_up.assert_called_once_with()
+        counter.close.assert_called_once_with()
         assert compactor.calls == ["warm_up_async", "close_async"]
 
     @pytest.mark.asyncio