feat(localFileStore): Add Cache Layer to LocalFileStore to Avoid I/O Overhead from Read Operations

openhands-sdk/openhands/sdk/io/local.py

@@ -1,6 +1,8 @@
 import os
 import shutil
 
+from cachetools import LRUCache
+
 from openhands.sdk.logger import get_logger
 from openhands.sdk.observability.laminar import observe
 
@@ -12,13 +14,15 @@
 
 class LocalFileStore(FileStore):
     root: str
+    cache: LRUCache
 
-    def __init__(self, root: str):
+    def __init__(self, root: str, cache_size: int = 100) -> None:
         if root.startswith("~"):
             root = os.path.expanduser(root)
         root = os.path.abspath(os.path.normpath(root))
         self.root = root
         os.makedirs(self.root, exist_ok=True)
+        self.cache = LRUCache(maxsize=cache_size)
 
     def get_full_path(self, path: str) -> str:
         # strip leading slash to keep relative under root
@@ -32,6 +36,7 @@ def get_full_path(self, path: str) -> str:
         # ensure sandboxing
         if os.path.commonpath([self.root, full]) != self.root:
             raise ValueError(f"path escapes filestore root: {path}")
+
         return full
 
     @observe(name="LocalFileStore.write", span_type="TOOL")
@@ -44,11 +49,33 @@ def write(self, path: str, contents: str | bytes) -> None:
         else:
             with open(full_path, "wb") as f:
                 f.write(contents)
+        cache_content = (
+            contents.decode("utf-8", errors="replace")
+            if isinstance(contents, bytes)
+            else contents
+        )
+
+        self.cache[full_path] = cache_content
 
     def read(self, path: str) -> str:
         full_path = self.get_full_path(path)
-        with open(full_path, encoding="utf-8") as f:
-            return f.read()
+
+        if full_path in self.cache:
+            return self.cache[full_path]
+
+        if not os.path.exists(full_path):
+            raise FileNotFoundError(path)
+        result: str
+        try:
+            with open(full_path, encoding="utf-8") as f:
+                result = f.read()
+        except UnicodeDecodeError:
+            logger.debug(f"File {full_path} is binary, reading as bytes")
+            with open(full_path, "rb") as f:
+                result = f.read().decode("utf-8", errors="replace")
+
+        self.cache[full_path] = result
+        return result
 
     @observe(name="LocalFileStore.list", span_type="TOOL")
     def list(self, path: str) -> list[str]:
@@ -67,9 +94,12 @@ def list(self, path: str) -> list[str]:
 
     @observe(name="LocalFileStore.delete", span_type="TOOL")
     def delete(self, path: str) -> None:
+        has_exist: bool = True
+        full_path: str | None = None
         try:
             full_path = self.get_full_path(path)
             if not os.path.exists(full_path):
+                has_exist = False
                 logger.debug(f"Local path does not exist: {full_path}")
                 return
             if os.path.isfile(full_path):
@@ -80,3 +110,15 @@ def delete(self, path: str) -> None:
                 logger.debug(f"Removed local directory: {full_path}")
         except Exception as e:
             logger.error(f"Error clearing local file store: {str(e)}")
+        finally:
+            if has_exist and full_path is not None:
+                self._cache_delete(full_path)
+
+    def _cache_delete(self, path: str) -> None:
+        try:
+            keys_to_delete = [key for key in self.cache.keys() if key.startswith(path)]
+            for key in keys_to_delete:
+                del self.cache[key]
+            logger.debug(f"Cleared LRU cache: {path}")
+        except Exception as e:
+            logger.error(f"Error clearing LRU cache: {str(e)}")