Replace diskcache with custom file-based storage system to resolve memory issues (#82)

* Fix memory issue in file_editor and add large history insert test

* Replace diskcache with custom file-based storage system

* Remove diskcache dependency from pyproject.toml

* Simplify file_editor function by removing chunk serialization

* Implement FileCache class and refactor FileHistoryManager to use it

* Minimize changes while maintaining new functionality

* Implement key hashing in FileCache to handle filepath issues

* Update clear method in FileCache to handle nested directories

* Fix FileCache issues: size limit functionality and empty string keys

* Add test case for size limit with empty key

* Update poetry.lock and pyproject.toml

* Improve FileHistoryManager with consistent key generation and add get_all_history method

* Fix FileCache and FileHistoryManager issues, update tests

* Apply linter fixes

* Fix FileCache size limit functionality and edge cases

* Apply linting fixes

* Fix JSON formatting in file_editor function

* Apply linting fixes

* Add detailed logging to FileCache for debugging

* Apply linting fixes

* Optimize memory usage in file_editor function

* Apply ruff-format changes

* Add FileCache import to __init__.py

* Add FileCache import to openhands_aci/__init__.py

* Update FileCache import in test_file_cache.py

* Export FileCache in editor/__init__.py

* Apply ruff linter fixes

* use mtime

* bump to 0.2.4

* use utf8 measurement

---------

Co-authored-by: openhands <openhands@all-hands.dev>
Co-authored-by: Hoang Tran <descience.thh10@gmail.com>

Author: 3 people

.openhands/microagents/repo.md

@@ -43,4 +43,4 @@ CI/CD:
   - `lint.yml`: Code linting checks
   - `py-unit-tests.yml`: Unit tests
   - `py-intg-tests.yml`: Integration tests
-  - `pypi-release.yml`: Package publishing to PyPI
+  - `pypi-release.yml`: Package publishing to PyPI

openhands_aci/__init__.py

@@ -1,3 +1,4 @@
 from .editor import file_editor
+from .editor.file_cache import FileCache
 
-__all__ = ['file_editor']
+__all__ = ['file_editor', 'FileCache']

openhands_aci/editor/__init__.py

@@ -3,10 +3,13 @@
 
 from .editor import Command, OHEditor
 from .exceptions import ToolError
+from .file_cache import FileCache
 from .results import ToolResult
 
 _GLOBAL_EDITOR = OHEditor()
 
+__all__ = ['Command', 'OHEditor', 'ToolError', 'ToolResult', 'FileCache', 'file_editor']
+
 
 def _make_api_tool_result(tool_result: ToolResult) -> str:
     """Convert an agent ToolResult to an API ToolResultBlockParam."""
@@ -44,6 +47,20 @@ def file_editor(
 
     formatted_output_and_error = _make_api_tool_result(result)
     marker_id = uuid.uuid4().hex
-    return f"""<oh_aci_output_{marker_id}>
-{json.dumps(result.to_dict(extra_field={'formatted_output_and_error': formatted_output_and_error}), indent=2)}
-</oh_aci_output_{marker_id}>"""
+
+    def json_generator():
+        yield '{'
+        first = True
+        for key, value in result.to_dict().items():
+            if not first:
+                yield ','
+            first = False
+            yield f'"{key}": {json.dumps(value)}'
+        yield f', "formatted_output_and_error": {json.dumps(formatted_output_and_error)}'
+        yield '}'
+
+    return (
+        f'<oh_aci_output_{marker_id}>\n'
+        + ''.join(json_generator())
+        + f'\n</oh_aci_output_{marker_id}>'
+    )

openhands_aci/editor/file_cache.py

@@ -0,0 +1,146 @@
+import hashlib
+import json
+import logging
+import os
+import time
+from pathlib import Path
+from typing import Any, Optional
+
+logging.basicConfig(level=logging.DEBUG)
+logger = logging.getLogger(__name__)
+
+
+class FileCache:
+    def __init__(self, directory: str, size_limit: Optional[int] = None):
+        self.directory = Path(directory)
+        self.directory.mkdir(parents=True, exist_ok=True)
+        self.size_limit = size_limit
+        self.current_size = 0
+        self._update_current_size()
+        logger.debug(
+            f'FileCache initialized with directory: {self.directory}, size_limit: {self.size_limit}, current_size: {self.current_size}'
+        )
+
+    def _get_file_path(self, key: str) -> Path:
+        hashed_key = hashlib.sha256(key.encode()).hexdigest()
+        return self.directory / f'{hashed_key}.json'
+
+    def _update_current_size(self):
+        self.current_size = sum(
+            f.stat().st_size for f in self.directory.glob('*.json') if f.is_file()
+        )
+        logger.debug(f'Current size updated: {self.current_size}')
+
+    def set(self, key: str, value: Any) -> None:
+        file_path = self._get_file_path(key)
+        content = json.dumps({'key': key, 'value': value})
+        content_size = len(content.encode('utf-8'))
+        logger.debug(f'Setting key: {key}, content_size: {content_size}')
+
+        if self.size_limit is not None:
+            if file_path.exists():
+                old_size = file_path.stat().st_size
+                size_diff = content_size - old_size
+                logger.debug(
+                    f'Existing file: old_size: {old_size}, size_diff: {size_diff}'
+                )
+                if size_diff > 0:
+                    while (
+                        self.current_size + size_diff > self.size_limit
+                        and len(self) > 1
+                    ):
+                        logger.debug(
+                            f'Evicting oldest (existing file case): current_size: {self.current_size}, size_limit: {self.size_limit}'
+                        )
+                        self._evict_oldest(file_path)
+            else:
+                while (
+                    self.current_size + content_size > self.size_limit and len(self) > 1
+                ):
+                    logger.debug(
+                        f'Evicting oldest (new file case): current_size: {self.current_size}, size_limit: {self.size_limit}'
+                    )
+                    self._evict_oldest(file_path)
+
+        if file_path.exists():
+            self.current_size -= file_path.stat().st_size
+            logger.debug(
+                f'Existing file removed from current_size: {self.current_size}'
+            )
+
+        with open(file_path, 'w') as f:
+            f.write(content)
+
+        self.current_size += content_size
+        logger.debug(f'File written, new current_size: {self.current_size}')
+        os.utime(
+            file_path, (time.time(), time.time())
+        )  # Update access and modification time
+
+    def _evict_oldest(self, exclude_path: Optional[Path] = None):
+        oldest_file = min(
+            (
+                f
+                for f in self.directory.glob('*.json')
+                if f.is_file() and f != exclude_path
+            ),
+            key=os.path.getmtime,
+        )
+        evicted_size = oldest_file.stat().st_size
+        self.current_size -= evicted_size
+        os.remove(oldest_file)
+        logger.debug(
+            f'Evicted file: {oldest_file}, size: {evicted_size}, new current_size: {self.current_size}'
+        )
+
+    def get(self, key: str, default: Any = None) -> Any:
+        file_path = self._get_file_path(key)
+        if not file_path.exists():
+            logger.debug(f'Get: Key not found: {key}')
+            return default
+        with open(file_path, 'r') as f:
+            data = json.load(f)
+            os.utime(file_path, (time.time(), time.time()))  # Update access time
+            logger.debug(f'Get: Key found: {key}')
+            return data['value']
+
+    def delete(self, key: str) -> None:
+        file_path = self._get_file_path(key)
+        if file_path.exists():
+            deleted_size = file_path.stat().st_size
+            self.current_size -= deleted_size
+            os.remove(file_path)
+            logger.debug(
+                f'Deleted key: {key}, size: {deleted_size}, new current_size: {self.current_size}'
+            )
+
+    def clear(self) -> None:
+        for item in self.directory.glob('*.json'):
+            if item.is_file():
+                os.remove(item)
+        self.current_size = 0
+        logger.debug('Cache cleared')
+
+    def __contains__(self, key: str) -> bool:
+        exists = self._get_file_path(key).exists()
+        logger.debug(f'Contains check: {key}, result: {exists}')
+        return exists
+
+    def __len__(self) -> int:
+        length = sum(1 for _ in self.directory.glob('*.json') if _.is_file())
+        logger.debug(f'Cache length: {length}')
+        return length
+
+    def __iter__(self):
+        for file in self.directory.glob('*.json'):
+            if file.is_file():
+                with open(file, 'r') as f:
+                    data = json.load(f)
+                    logger.debug(f"Yielding key: {data['key']}")
+                    yield data['key']
+
+    def __getitem__(self, key: str) -> Any:
+        return self.get(key)
+
+    def __setitem__(self, key: str, value: Any) -> None:
+        self.set(key, value)

openhands_aci/editor/history.py

@@ -1,10 +1,11 @@
 """History management for file edits with disk-based storage and memory constraints."""
 
+import logging
 import tempfile
 from pathlib import Path
-from typing import Optional
+from typing import List, Optional
 
-from diskcache import Cache
+from .file_cache import FileCache
 
 
 class FileHistoryManager:
@@ -21,67 +22,91 @@ def __init__(
 
         Notes:
             - Each file's history is limited to the last N entries to conserve memory
-            - The disk cache is limited to 500MB total to prevent excessive disk usage
+            - The file cache is limited to prevent excessive disk usage
             - Older entries are automatically removed when limits are exceeded
         """
         self.max_history_per_file = max_history_per_file
         if history_dir is None:
             history_dir = Path(tempfile.mkdtemp(prefix='oh_editor_history_'))
-        self.cache = Cache(str(history_dir), size_limit=5e8)  # 500MB size limit
+        self.cache = FileCache(str(history_dir))
+        self.logger = logging.getLogger(__name__)
+
+    def _get_metadata_key(self, file_path: Path) -> str:
+        return f'{file_path}.metadata'
+
+    def _get_history_key(self, file_path: Path, counter: int) -> str:
+        return f'{file_path}.{counter}'
 
     def add_history(self, file_path: Path, content: str):
         """Add a new history entry for a file."""
-        key = str(file_path)
-        # Get list of entry indices and counter for this file
-        entries_key = f'{key}:entries'
-        counter_key = f'{key}:counter'
-        entries = self.cache.get(entries_key, [])
-        counter = self.cache.get(counter_key, 0)
-
-        # Add new entry with monotonically increasing counter
-        entry_key = f'{key}:{counter}'
-        self.cache.set(entry_key, content)
-        entries.append(entry_key)
-        counter += 1
+        metadata_key = self._get_metadata_key(file_path)
+        metadata = self.cache.get(metadata_key, {'entries': [], 'counter': 0})
+        counter = metadata['counter']
+
+        # Add new entry
+        history_key = self._get_history_key(file_path, counter)
+        self.cache.set(history_key, content)
+
+        metadata['entries'].append(counter)
+        metadata['counter'] += 1
 
         # Keep only last N entries
-        if len(entries) > self.max_history_per_file:
-            old_key = entries.pop(0)
-            self.cache.delete(old_key)
+        while len(metadata['entries']) > self.max_history_per_file:
+            old_counter = metadata['entries'].pop(0)
+            old_history_key = self._get_history_key(file_path, old_counter)
+            self.cache.delete(old_history_key)
 
-        # Update entries list and counter
-        self.cache.set(entries_key, entries)
-        self.cache.set(counter_key, counter)
+        self.cache.set(metadata_key, metadata)
 
     def get_last_history(self, file_path: Path) -> Optional[str]:
         """Get the most recent history entry for a file."""
-        key = str(file_path)
-        entries_key = f'{key}:entries'
-        entries = self.cache.get(entries_key, [])
+        metadata_key = self._get_metadata_key(file_path)
+        metadata = self.cache.get(metadata_key, {'entries': [], 'counter': 0})
+        entries = metadata['entries']
 
         if not entries:
             return None
 
-        # Get and remove last entry
-        last_key = entries.pop()
-        content = self.cache.get(last_key)
-        self.cache.delete(last_key)
+        # Get the last entry without removing it
+        last_counter = entries[-1]
+        history_key = self._get_history_key(file_path, last_counter)
+        content = self.cache.get(history_key)
+
+        if content is None:
+            self.logger.warning(f'History entry not found for {file_path}')
+            return None
 
-        # Update entries list
-        self.cache.set(entries_key, entries)
         return content
 
     def clear_history(self, file_path: Path):
         """Clear history for a given file."""
-        key = str(file_path)
-        entries_key = f'{key}:entries'
-        counter_key = f'{key}:counter'
-        entries = self.cache.get(entries_key, [])
-
-        # Delete all entries
-        for entry_key in entries:
-            self.cache.delete(entry_key)
-
-        # Delete entries list and counter
-        self.cache.delete(entries_key)
-        self.cache.delete(counter_key)
+        metadata_key = self._get_metadata_key(file_path)
+        metadata = self.cache.get(metadata_key, {'entries': [], 'counter': 0})
+
+        # Delete all history entries
+        for counter in metadata['entries']:
+            history_key = self._get_history_key(file_path, counter)
+            self.cache.delete(history_key)
+
+        # Clear metadata
+        self.cache.set(metadata_key, {'entries': [], 'counter': 0})
+
+    def get_all_history(self, file_path: Path) -> List[str]:
+        """Get all history entries for a file."""
+        metadata_key = self._get_metadata_key(file_path)
+        metadata = self.cache.get(metadata_key, {'entries': [], 'counter': 0})
+        entries = metadata['entries']
+
+        history = []
+        for counter in entries:
+            history_key = self._get_history_key(file_path, counter)
+            content = self.cache.get(history_key)
+            if content is not None:
+                history.append(content)
+
+        return history
+
+    def get_metadata(self, file_path: Path):
+        """Get metadata for a file (for testing purposes)."""
+        metadata_key = self._get_metadata_key(file_path)
+        return self.cache.get(metadata_key, {'entries': [], 'counter': 0})

openhands_aci/linter/impl/treesitter_compat.py

@@ -1,6 +1,7 @@
 """Compatibility layer for tree-sitter 0.24.0."""
 
 import importlib
+
 from tree_sitter import Language, Parser
 
 # Cache of loaded languages
@@ -11,13 +12,13 @@ def get_parser(language):
     """Get a Parser object for the given language name."""
     if language not in _language_cache:
         # Try to import the language module
-        module_name = f"tree_sitter_{language}"
+        module_name = f'tree_sitter_{language}'
         try:
             module = importlib.import_module(module_name)
             _language_cache[language] = Language(module.language())
         except ImportError:
             raise ValueError(
-                f"Language {language} is not supported. Please install {module_name} package."
+                f'Language {language} is not supported. Please install {module_name} package.'
             )
 
-    return Parser(_language_cache[language])
+    return Parser(_language_cache[language])