From 9bdc8dda6c14bdbb889fe082a235959dcb2ab576 Mon Sep 17 00:00:00 2001 From: Rohit Malhotra Date: Mon, 10 Feb 2025 11:12:12 -0500 Subject: [PATCH] [Enhancement]: Handle GH token refresh inside runtime (#6632) --- .../github/github_service.py | 11 ++++++++++- .../github/github_types.py | 0 openhands/runtime/base.py | 14 ++++++++++++++ .../action_execution/action_execution_client.py | 2 ++ openhands/runtime/impl/remote/remote_runtime.py | 15 +++++++++++---- openhands/server/config/server_config.py | 2 -- .../standalone_conversation_manager.py | 4 ++-- openhands/server/routes/github.py | 13 +++++-------- openhands/server/routes/manage_conversations.py | 6 +++--- openhands/server/routes/settings.py | 6 ++++-- openhands/server/session/agent_session.py | 9 +++++++++ openhands/server/session/session.py | 5 ++++- tests/unit/test_github_service.py | 4 ++-- 13 files changed, 66 insertions(+), 25 deletions(-) rename openhands/{services => integrations}/github/github_service.py (93%) rename openhands/{services => integrations}/github/github_types.py (100%) diff --git a/openhands/services/github/github_service.py b/openhands/integrations/github/github_service.py similarity index 93% rename from openhands/services/github/github_service.py rename to openhands/integrations/github/github_service.py index 92d14c5653..4d8b731255 100644 --- a/openhands/services/github/github_service.py +++ b/openhands/integrations/github/github_service.py @@ -1,14 +1,16 @@ +import os from typing import Any import httpx from pydantic import SecretStr -from openhands.services.github.github_types import ( +from openhands.integrations.github.github_types import ( GhAuthenticationError, GHUnknownException, GitHubRepository, GitHubUser, ) +from openhands.utils.import_utils import get_impl class GitHubService: @@ -133,3 +135,10 @@ class GitHubService: ] return repos + + +github_service_cls = os.environ.get( + 'OPENHANDS_GITHUB_SERVICE_CLS', + 'openhands.integrations.github.github_service.GitHubService', +) +GithubServiceImpl = get_impl(GitHubService, github_service_cls) diff --git a/openhands/services/github/github_types.py b/openhands/integrations/github/github_types.py similarity index 100% rename from openhands/services/github/github_types.py rename to openhands/integrations/github/github_types.py diff --git a/openhands/runtime/base.py b/openhands/runtime/base.py index f511a8b92f..b7c93eea6d 100644 --- a/openhands/runtime/base.py +++ b/openhands/runtime/base.py @@ -39,6 +39,7 @@ from openhands.events.observation import ( UserRejectObservation, ) from openhands.events.serialization.action import ACTION_TYPE_TO_CLASS +from openhands.integrations.github.github_service import GithubServiceImpl from openhands.microagent import ( BaseMicroAgent, load_microagents_from_dir, @@ -94,6 +95,7 @@ class Runtime(FileEditRuntimeMixin): status_callback: Callable | None = None, attach_to_existing: bool = False, headless_mode: bool = False, + github_user_id: str | None = None, ): self.sid = sid self.event_stream = event_stream @@ -126,6 +128,8 @@ class Runtime(FileEditRuntimeMixin): self, enable_llm_editor=config.get_agent_config().codeact_enable_llm_editor ) + self.github_user_id = github_user_id + def setup_initial_env(self) -> None: if self.attach_to_existing: return @@ -213,6 +217,16 @@ class Runtime(FileEditRuntimeMixin): event.set_hard_timeout(self.config.sandbox.timeout, blocking=False) assert event.timeout is not None try: + if isinstance(event, CmdRunAction): + if self.github_user_id and '$GITHUB_TOKEN' in event.command: + gh_client = GithubServiceImpl(user_id=self.github_user_id) + token = await gh_client.get_latest_token() + if token: + export_cmd = CmdRunAction( + f"export GITHUB_TOKEN='{token.get_secret_value()}'" + ) + await call_sync_from_async(self.run, export_cmd) + observation: Observation = await call_sync_from_async( self.run_action, event ) diff --git a/openhands/runtime/impl/action_execution/action_execution_client.py b/openhands/runtime/impl/action_execution/action_execution_client.py index 38afc54487..7da8c21ff6 100644 --- a/openhands/runtime/impl/action_execution/action_execution_client.py +++ b/openhands/runtime/impl/action_execution/action_execution_client.py @@ -55,6 +55,7 @@ class ActionExecutionClient(Runtime): status_callback: Any | None = None, attach_to_existing: bool = False, headless_mode: bool = True, + github_user_id: str | None = None, ): self.session = HttpSession() self.action_semaphore = threading.Semaphore(1) # Ensure one action at a time @@ -70,6 +71,7 @@ class ActionExecutionClient(Runtime): status_callback, attach_to_existing, headless_mode, + github_user_id, ) @abstractmethod diff --git a/openhands/runtime/impl/remote/remote_runtime.py b/openhands/runtime/impl/remote/remote_runtime.py index f663c1f5af..cb10b2c15b 100644 --- a/openhands/runtime/impl/remote/remote_runtime.py +++ b/openhands/runtime/impl/remote/remote_runtime.py @@ -45,6 +45,7 @@ class RemoteRuntime(ActionExecutionClient): status_callback: Optional[Callable] = None, attach_to_existing: bool = False, headless_mode: bool = True, + github_user_id: str | None = None, ): super().__init__( config, @@ -55,6 +56,7 @@ class RemoteRuntime(ActionExecutionClient): status_callback, attach_to_existing, headless_mode, + github_user_id, ) if self.config.sandbox.api_key is None: raise ValueError( @@ -291,7 +293,8 @@ class RemoteRuntime(ActionExecutionClient): stop=tenacity.stop_after_delay( self.config.sandbox.remote_runtime_init_timeout ) - | stop_if_should_exit() | self._stop_if_closed, + | stop_if_should_exit() + | self._stop_if_closed, reraise=True, retry=tenacity.retry_if_exception_type(AgentRuntimeNotReadyError), wait=tenacity.wait_fixed(2), @@ -394,10 +397,14 @@ class RemoteRuntime(ActionExecutionClient): retry_decorator = tenacity.retry( retry=tenacity.retry_if_exception_type(ConnectionError), - stop=tenacity.stop_after_attempt(3) | stop_if_should_exit() | self._stop_if_closed, + stop=tenacity.stop_after_attempt(3) + | stop_if_should_exit() + | self._stop_if_closed, wait=tenacity.wait_exponential(multiplier=1, min=4, max=60), ) - return retry_decorator(self._send_action_server_request_impl)(method, url, **kwargs) + return retry_decorator(self._send_action_server_request_impl)( + method, url, **kwargs + ) def _send_action_server_request_impl(self, method, url, **kwargs): try: @@ -430,6 +437,6 @@ class RemoteRuntime(ActionExecutionClient): ) from e else: raise e - + def _stop_if_closed(self, retry_state: tenacity.RetryCallState) -> bool: return self._runtime_closed diff --git a/openhands/server/config/server_config.py b/openhands/server/config/server_config.py index 32b2deab4e..ae86d8d43a 100644 --- a/openhands/server/config/server_config.py +++ b/openhands/server/config/server_config.py @@ -18,8 +18,6 @@ class ServerConfig(ServerConfigInterface): ) conversation_manager_class: str = 'openhands.server.conversation_manager.standalone_conversation_manager.StandaloneConversationManager' - github_service_class: str = 'openhands.services.github.github_service.GitHubService' - def verify_config(self): if self.config_cls: raise ValueError('Unexpected config path provided') diff --git a/openhands/server/conversation_manager/standalone_conversation_manager.py b/openhands/server/conversation_manager/standalone_conversation_manager.py index a1748038d6..2f49de63dc 100644 --- a/openhands/server/conversation_manager/standalone_conversation_manager.py +++ b/openhands/server/conversation_manager/standalone_conversation_manager.py @@ -149,8 +149,8 @@ class StandaloneConversationManager(ConversationManager): self._close_session(sid) for sid in self._local_agent_loops_by_sid ) return - except Exception as e: - logger.error(f'error_cleaning_stale') + except Exception: + logger.error('error_cleaning_stale') await asyncio.sleep(_CLEANUP_INTERVAL) async def get_running_agent_loops( diff --git a/openhands/server/routes/github.py b/openhands/server/routes/github.py index c50c1b2ca7..51013beff7 100644 --- a/openhands/server/routes/github.py +++ b/openhands/server/routes/github.py @@ -2,20 +2,17 @@ from fastapi import APIRouter, Depends from fastapi.responses import JSONResponse from pydantic import SecretStr -from openhands.server.auth import get_github_token, get_user_id -from openhands.server.shared import server_config -from openhands.services.github.github_service import ( +from openhands.integrations.github.github_service import GithubServiceImpl +from openhands.integrations.github.github_types import ( GhAuthenticationError, GHUnknownException, - GitHubService, + GitHubRepository, + GitHubUser, ) -from openhands.services.github.github_types import GitHubRepository, GitHubUser -from openhands.utils.import_utils import get_impl +from openhands.server.auth import get_github_token, get_user_id app = APIRouter(prefix='/api/github') -GithubServiceImpl = get_impl(GitHubService, server_config.github_service_class) - @app.get('/repositories') async def get_github_repositories( diff --git a/openhands/server/routes/manage_conversations.py b/openhands/server/routes/manage_conversations.py index 41c79a2d25..4edfb47c51 100644 --- a/openhands/server/routes/manage_conversations.py +++ b/openhands/server/routes/manage_conversations.py @@ -9,9 +9,9 @@ from pydantic import BaseModel, SecretStr from openhands.core.logger import openhands_logger as logger from openhands.events.action.message import MessageAction from openhands.events.stream import EventStreamSubscriber +from openhands.integrations.github.github_service import GithubServiceImpl from openhands.runtime import get_runtime_cls from openhands.server.auth import get_github_token, get_user_id -from openhands.server.routes.github import GithubServiceImpl from openhands.server.session.conversation_init_data import ConversationInitData from openhands.server.shared import ( ConversationStoreImpl, @@ -131,8 +131,8 @@ async def new_conversation(request: Request, data: InitSessionRequest): """ logger.info('Initializing new conversation') user_id = get_user_id(request) - github_service = GithubServiceImpl(user_id=user_id, token=get_github_token(request)) - github_token = await github_service.get_latest_token() + gh_client = GithubServiceImpl(user_id=user_id, token=get_github_token(request)) + github_token = await gh_client.get_latest_token() selected_repository = data.selected_repository initial_user_msg = data.initial_user_msg diff --git a/openhands/server/routes/settings.py b/openhands/server/routes/settings.py index 8fb7b5f06c..66ed76a23e 100644 --- a/openhands/server/routes/settings.py +++ b/openhands/server/routes/settings.py @@ -3,10 +3,10 @@ from fastapi.responses import JSONResponse from pydantic import SecretStr from openhands.core.logger import openhands_logger as logger +from openhands.integrations.github.github_service import GithubServiceImpl from openhands.server.auth import get_github_token, get_user_id from openhands.server.settings import GETSettingsModel, POSTSettingsModel, Settings from openhands.server.shared import SettingsStoreImpl, config -from openhands.services.github.github_service import GitHubService app = APIRouter(prefix='/api') @@ -51,7 +51,9 @@ async def store_settings( try: # We check if the token is valid by getting the user # If the token is invalid, this will raise an exception - github = GitHubService(user_id=None, token=SecretStr(settings.github_token)) + github = GithubServiceImpl( + user_id=None, token=SecretStr(settings.github_token) + ) await github.get_user() except Exception as e: diff --git a/openhands/server/session/agent_session.py b/openhands/server/session/agent_session.py index 37202e2bcb..298474b884 100644 --- a/openhands/server/session/agent_session.py +++ b/openhands/server/session/agent_session.py @@ -17,6 +17,7 @@ from openhands.events.stream import EventStream from openhands.microagent import BaseMicroAgent from openhands.runtime import get_runtime_cls from openhands.runtime.base import Runtime +from openhands.runtime.impl.remote.remote_runtime import RemoteRuntime from openhands.security import SecurityAnalyzer, options from openhands.storage.files import FileStore from openhands.utils.async_utils import call_sync_from_async @@ -49,6 +50,7 @@ class AgentSession: sid: str, file_store: FileStore, status_callback: Optional[Callable] = None, + github_user_id: str | None = None, ): """Initializes a new instance of the Session class @@ -61,6 +63,7 @@ class AgentSession: self.event_stream = EventStream(sid, file_store) self.file_store = file_store self._status_callback = status_callback + self.github_user_id = github_user_id async def start( self, @@ -202,6 +205,11 @@ class AgentSession: if github_token else None ) + + kwargs = {} + if runtime_cls == RemoteRuntime: + kwargs['github_user_id'] = self.github_user_id + self.runtime = runtime_cls( config=config, event_stream=self.event_stream, @@ -210,6 +218,7 @@ class AgentSession: status_callback=self._status_callback, headless_mode=False, env_vars=env_vars, + **kwargs, ) # FIXME: this sleep is a terrible hack. diff --git a/openhands/server/session/session.py b/openhands/server/session/session.py index 848b51efc4..5d34baf4f6 100644 --- a/openhands/server/session/session.py +++ b/openhands/server/session/session.py @@ -55,7 +55,10 @@ class Session: self.last_active_ts = int(time.time()) self.file_store = file_store self.agent_session = AgentSession( - sid, file_store, status_callback=self.queue_status_message + sid, + file_store, + status_callback=self.queue_status_message, + github_user_id=user_id, ) self.agent_session.event_stream.subscribe( EventStreamSubscriber.SERVER, self.on_event, self.sid diff --git a/tests/unit/test_github_service.py b/tests/unit/test_github_service.py index e24faf6cb1..222be16767 100644 --- a/tests/unit/test_github_service.py +++ b/tests/unit/test_github_service.py @@ -4,8 +4,8 @@ import httpx import pytest from pydantic import SecretStr -from openhands.services.github.github_service import GitHubService -from openhands.services.github.github_types import GhAuthenticationError +from openhands.integrations.github.github_service import GitHubService +from openhands.integrations.github.github_types import GhAuthenticationError @pytest.mark.asyncio