Merge branch 'feat/pr-mcp-upstream' of https://github.com/oraichain/OpenHands into feat/pr-mcp-upstream

Author: ducphamle2

openhands/core/setup.py

@@ -119,7 +119,11 @@ def initialize_repository_for_runtime(
     if selected_repository and provider_tokens:
         logger.debug(f'Selected repository {selected_repository}.')
         repo_directory = call_async_from_sync(
-            runtime.clone_repo, GENERAL_TIMEOUT, github_token, selected_repository, None
+            runtime.clone_repo,
+            GENERAL_TIMEOUT,
+            provider_tokens,
+            selected_repository,
+            None,
         )
         # Run setup script if it exists
         runtime.maybe_run_setup_script()

openhands/events/event_store.py

@@ -7,12 +7,37 @@
 from openhands.events.serialization.event import event_from_dict, event_to_dict
 from openhands.storage.files import FileStore
 from openhands.storage.locations import (
+    get_conversation_dir,
     get_conversation_event_filename,
     get_conversation_events_dir,
 )
 from openhands.utils.shutdown_listener import should_continue
 
 
+@dataclass(frozen=True)
+class _CachePage:
+    events: list[dict] | None
+    start: int
+    end: int
+
+    def covers(self, global_index: int) -> bool:
+        if global_index < self.start:
+            return False
+        if global_index >= self.end:
+            return False
+        return True
+
+    def get_event(self, global_index: int) -> Event | None:
+        # If there was not actually a cached page, return None
+        if not self.events:
+            return None
+        local_index = global_index - self.start
+        return event_from_dict(self.events[local_index])
+
+
+_DUMMY_PAGE = _CachePage(None, 1, -1)
+
+
 @dataclass
 class EventStore:
     """
@@ -23,6 +48,7 @@ class EventStore:
     file_store: FileStore
     user_id: str | None
     cur_id: int = -1  # We fix this in post init if it is not specified
+    cache_size: int = 25
 
     def __post_init__(self) -> None:
         if self.cur_id >= 0:
@@ -83,30 +109,33 @@ def should_filter(event: Event) -> bool:
                 return True
             return False
 
+        if end_id is None:
+            end_id = self.cur_id
+        else:
+            end_id += 1  # From inclusive to exclusive
+
         if reverse:
-            if end_id is None:
-                end_id = self.cur_id - 1
-            event_id = end_id
-            while event_id >= start_id:
-                try:
-                    event = self.get_event(event_id)
-                    if not should_filter(event):
-                        yield event
-                except FileNotFoundError:
-                    logger.debug(f'No event found for ID {event_id}')
-                event_id -= 1
+            step = -1
+            start_id, end_id = end_id, start_id
+            start_id -= 1
+            end_id -= 1
         else:
-            event_id = start_id
-            while should_continue():
-                if end_id is not None and event_id > end_id:
-                    break
+            step = 1
+
+        cache_page = _DUMMY_PAGE
+        for index in range(start_id, end_id, step):
+            if not should_continue():
+                return
+            if not cache_page.covers(index):
+                cache_page = self._load_cache_page_for_index(index)
+            event = cache_page.get_event(index)
+            if event is None:
                 try:
-                    event = self.get_event(event_id)
-                    if not should_filter(event):
-                        yield event
+                    event = self.get_event(index)
                 except FileNotFoundError:
-                    break
-                event_id += 1
+                    event = None
+            if event and not should_filter(event):
+                yield event
 
     def get_event(self, id: int) -> Event:
         filename = self._get_filename_for_id(id, self.user_id)
@@ -230,6 +259,25 @@ def get_matching_events(
     def _get_filename_for_id(self, id: int, user_id: str | None) -> str:
         return get_conversation_event_filename(self.sid, id, user_id)
 
+    def _get_filename_for_cache(self, start: int, end: int) -> str:
+        return f'{get_conversation_dir(self.sid, self.user_id)}event_cache/{start}-{end}.json'
+
+    def _load_cache_page(self, start: int, end: int) -> _CachePage:
+        """Read a page from the cache. Reading individual events is slow when there are a lot of them, so we use pages."""
+        cache_filename = self._get_filename_for_cache(start, end)
+        try:
+            content = self.file_store.read(cache_filename)
+            events = json.loads(content)
+        except FileNotFoundError:
+            events = None
+        page = _CachePage(events, start, end)
+        return page
+
+    def _load_cache_page_for_index(self, index: int) -> _CachePage:
+        offset = index % self.cache_size
+        index -= offset
+        return self._load_cache_page(index, index + self.cache_size)
+
     @staticmethod
     def _get_id_from_filename(filename: str) -> int:
         try:

openhands/events/stream.py

@@ -52,6 +52,7 @@ class EventStream(EventStore):
     _queue_loop: asyncio.AbstractEventLoop | None
     _thread_pools: dict[str, dict[str, ThreadPoolExecutor]]
     _thread_loops: dict[str, dict[str, asyncio.AbstractEventLoop]]
+    _write_page_cache: list[dict]
 
     def __init__(self, sid: str, file_store: FileStore, user_id: str | None = None):
         super().__init__(sid, file_store, user_id)
@@ -66,6 +67,7 @@ def __init__(self, sid: str, file_store: FileStore, user_id: str | None = None):
         self._subscribers = {}
         self._lock = threading.Lock()
         self.secrets = {}
+        self._write_page_cache = []
 
     def _init_thread_loop(self, subscriber_id: str, callback_id: str) -> None:
         loop = asyncio.new_event_loop()
@@ -172,8 +174,22 @@ def add_event(self, event: Event, source: EventSource) -> None:
             self.file_store.write(
                 self._get_filename_for_id(event.id, self.user_id), json.dumps(data)
             )
+            self._write_page_cache.append(data)
+            self._store_cache_page()
         self._queue.put(event)
 
+    def _store_cache_page(self):
+        """Store a page in the cache. Reading individual events is slow when there are a lot of them, so we use pages."""
+        current_write_page = self._write_page_cache
+        if len(current_write_page) < self.cache_size:
+            return
+        self._write_page_cache = []
+        start = current_write_page[0]['id']
+        end = start + self.cache_size
+        contents = json.dumps(current_write_page)
+        cache_filename = self._get_filename_for_cache(start, end)
+        self.file_store.write(cache_filename, contents)
+
     def set_secrets(self, secrets: dict[str, str]) -> None:
         self.secrets = secrets.copy()