Fix file descriptor leak (#6897)

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

openhands/llm/llm.py

@@ -8,6 +8,7 @@
 import requests
 
 from openhands.core.config import LLMConfig
+from openhands.utils.ensure_httpx_close import EnsureHttpxClose
 
 with warnings.catch_warnings():
     warnings.simplefilter('ignore')
@@ -230,9 +231,9 @@ def wrapper(*args, **kwargs):
 
             # Record start time for latency measurement
             start_time = time.time()
-
-            # we don't support streaming here, thus we get a ModelResponse
-            resp: ModelResponse = self._completion_unwrapped(*args, **kwargs)
+            with EnsureHttpxClose():
+                # we don't support streaming here, thus we get a ModelResponse
+                resp: ModelResponse = self._completion_unwrapped(*args, **kwargs)
 
             # Calculate and record latency
             latency = time.time() - start_time
@@ -287,7 +288,11 @@ def wrapper(*args, **kwargs):
                     'messages': messages,
                     'response': resp,
                     'args': args,
-                    'kwargs': {k: v for k, v in kwargs.items() if k != 'messages'},
+                    'kwargs': {
+                        k: v
+                        for k, v in kwargs.items()
+                        if k not in ('messages', 'client')
+                    },
                     'timestamp': time.time(),
                     'cost': cost,
                 }

openhands/utils/ensure_httpx_close.py

@@ -0,0 +1,43 @@
+"""
+LiteLLM currently have an issue where HttpHandlers are being created but not
+closed. We have submitted a PR to them, (https://github.com/BerriAI/litellm/pull/8711)
+and their dev team say they are in the process of a refactor that will fix this, but
+in the meantime, we need to manage the lifecycle of the httpx.Client manually.
+
+We can't simply pass in our own client object, because all the different implementations use
+different types of client object.
+
+So we monkey patch the httpx.Client class to track newly created instances and close these
+when the operations complete. (This is relatively safe, as if the client is reused after this
+then is will transparently reopen)
+
+Hopefully, this will be fixed soon and we can remove this abomination.
+"""
+
+from dataclasses import dataclass, field
+from functools import wraps
+from typing import Callable
+
+from httpx import Client
+
+
+@dataclass
+class EnsureHttpxClose:
+    clients: list[Client] = field(default_factory=list)
+    original_init: Callable | None = None
+
+    def __enter__(self):
+        self.original_init = Client.__init__
+
+        @wraps(Client.__init__)
+        def init_wrapper(*args, **kwargs):
+            self.clients.append(args[0])
+            return self.original_init(*args, **kwargs)  # type: ignore
+
+        Client.__init__ = init_wrapper
+
+    def __exit__(self, type, value, traceback):
+        Client.__init__ = self.original_init
+        while self.clients:
+            client = self.clients.pop()
+            client.close()

tests/unit/test_ensure_httpx_close.py

@@ -0,0 +1,84 @@
+from httpx import Client
+
+from openhands.utils.ensure_httpx_close import EnsureHttpxClose
+
+
+def test_ensure_httpx_close_basic():
+    """Test basic functionality of EnsureHttpxClose."""
+    clients = []
+    ctx = EnsureHttpxClose()
+    with ctx:
+        # Create a client - should be tracked
+        client = Client()
+        assert client in ctx.clients
+        assert len(ctx.clients) == 1
+        clients.append(client)
+
+    # After context exit, client should be closed
+    assert client.is_closed
+
+
+def test_ensure_httpx_close_multiple_clients():
+    """Test EnsureHttpxClose with multiple clients."""
+    ctx = EnsureHttpxClose()
+    with ctx:
+        client1 = Client()
+        client2 = Client()
+        assert len(ctx.clients) == 2
+        assert client1 in ctx.clients
+        assert client2 in ctx.clients
+
+    assert client1.is_closed
+    assert client2.is_closed
+
+
+def test_ensure_httpx_close_nested():
+    """Test nested usage of EnsureHttpxClose."""
+    outer_ctx = EnsureHttpxClose()
+    with outer_ctx:
+        client1 = Client()
+        assert client1 in outer_ctx.clients
+
+        inner_ctx = EnsureHttpxClose()
+        with inner_ctx:
+            client2 = Client()
+            assert client2 in inner_ctx.clients
+            # Since both contexts are using the same monkey-patched __init__,
+            # both contexts will track all clients created while they are active
+            assert client2 in outer_ctx.clients
+
+        # After inner context, client2 should be closed
+        assert client2.is_closed
+        # client1 should still be open since outer context is still active
+        assert not client1.is_closed
+
+    # After outer context, both clients should be closed
+    assert client1.is_closed
+    assert client2.is_closed
+
+
+def test_ensure_httpx_close_exception():
+    """Test EnsureHttpxClose when an exception occurs."""
+    client = None
+    ctx = EnsureHttpxClose()
+    try:
+        with ctx:
+            client = Client()
+            raise ValueError('Test exception')
+    except ValueError:
+        pass
+
+    # Client should be closed even if an exception occurred
+    assert client is not None
+    assert client.is_closed
+
+
+def test_ensure_httpx_close_restore_init():
+    """Test that the original __init__ is restored after context exit."""
+    original_init = Client.__init__
+    ctx = EnsureHttpxClose()
+    with ctx:
+        assert Client.__init__ != original_init
+
+    # Original __init__ should be restored
+    assert Client.__init__ == original_init

tests/unit/test_llm.py

@@ -1,4 +1,6 @@
 import copy
+import tempfile
+from pathlib import Path
 from unittest.mock import MagicMock, patch
 
 import pytest
@@ -489,3 +491,27 @@ def test_llm_token_usage(mock_litellm_completion, default_config):
     assert usage_entry_2['cache_read_tokens'] == 1
     assert usage_entry_2['cache_write_tokens'] == 3
     assert usage_entry_2['response_id'] == 'test-response-usage-2'
+
+
+@patch('openhands.llm.llm.litellm_completion')
+def test_completion_with_log_completions(mock_litellm_completion, default_config):
+    with tempfile.TemporaryDirectory() as temp_dir:
+        default_config.log_completions = True
+        default_config.log_completions_folder = temp_dir
+        mock_response = {
+            'choices': [{'message': {'content': 'This is a mocked response.'}}]
+        }
+        mock_litellm_completion.return_value = mock_response
+
+        test_llm = LLM(config=default_config)
+        response = test_llm.completion(
+            messages=[{'role': 'user', 'content': 'Hello!'}],
+            stream=False,
+            drop_params=True,
+        )
+        assert (
+            response['choices'][0]['message']['content'] == 'This is a mocked response.'
+        )
+        files = list(Path(temp_dir).iterdir())
+        # Expect a log to be generated
+        assert len(files) == 1