[Enhancement]: Handle GH token refresh inside runtime (#6632)

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)

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
             )

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

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

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')

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(

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(

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

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:

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.

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

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