Repository: All-Hands-AI/OpenHands
Pull request #7670: [Fix]: Ensure refresh logic works for restarted conversations in cloud openhands
Author: login: malhotra5
Created: 2025-04-02T18:13:57Z | Merged: 2025-04-03T15:58:01Z

- [ ] This change is worth documenting at https://docs.all-hands.dev/
- [ ] Include this change in the Release Notes. If checked, you **must** provide an **end-user friendly** description for your change below

**End-user friendly description of the problem this fixes or functionality that this introduces.**

- This fixes refresh logic for multiple providers in cloud openhands; this, however, is just scaffolding changes and won't have any effect in the upcoming release

---
**Give a summary of what the PR does, explaining any non-trivial design decisions.**

#7532 included a hotfix where GitHub token refresh failed because provider information was lost for conversations that were restarted. Since the refresh logic only exports the tokens for the providers that are currently set, it wouldn't export any tokens due to the lost provider information. 

This PR ensures that when conversations are joined, the information about which providers are set is available in the runtime. This ensures that refresh logic takes place as expected  

---
**Link of any specific issues this addresses.**

---

To run this PR locally, use the following command:
```
docker run -it --rm   -p 3000:3000   -v /var/run/docker.sock:/var/run/docker.sock   --add-host host.docker.internal:host-gateway   -e SANDBOX_RUNTIME_CONTAINER_IMAGE=docker.all-hands.dev/all-hands-ai/runtime:7e367a1-nikolaik   --name openhands-app-7e367a1   docker.all-hands.dev/all-hands-ai/openhands:7e367a1
```


Diff:
diff --git a/frontend/__tests__/context/ws-client-provider.test.tsx b/frontend/__tests__/context/ws-client-provider.test.tsx
index 777cbd1225c5..3d41a2eab32b 100644
--- a/frontend/__tests__/context/ws-client-provider.test.tsx
+++ b/frontend/__tests__/context/ws-client-provider.test.tsx
@@ -8,6 +8,7 @@ import {
   WsClientProvider,
   useWsClient,
 } from "#/context/ws-client-provider";
+import { AuthProvider } from "#/context/auth-context";
 
 describe("Propagate error message", () => {
   it("should do nothing when no message was passed from server", () => {
@@ -90,9 +91,11 @@ describe("WsClientProvider", () => {
     const { getByText } = render(<TestComponent />, {
       wrapper: ({ children }) => (
         <QueryClientProvider client={new QueryClient()}>
-          <WsClientProvider conversationId="test-conversation-id">
-            {children}
-          </WsClientProvider>
+          <AuthProvider initialProviderTokens={[]}>
+            <WsClientProvider conversationId="test-conversation-id">
+              {children}
+            </WsClientProvider>
+          </AuthProvider>
         </QueryClientProvider>
       ),
     });
diff --git a/frontend/src/context/ws-client-provider.tsx b/frontend/src/context/ws-client-provider.tsx
index 15209549c1ff..ad7965cdcee6 100644
--- a/frontend/src/context/ws-client-provider.tsx
+++ b/frontend/src/context/ws-client-provider.tsx
@@ -9,6 +9,7 @@ import {
   AssistantMessageAction,
   UserMessageAction,
 } from "#/types/core/actions";
+import { useAuth } from "./auth-context";
 
 const isOpenHandsEvent = (event: unknown): event is OpenHandsParsedEvent =>
   typeof event === "object" &&
@@ -110,6 +111,7 @@ export function WsClientProvider({
   );
   const [events, setEvents] = React.useState<Record<string, unknown>[]>([]);
   const lastEventRef = React.useRef<Record<string, unknown> | null>(null);
+  const { providerTokensSet } = useAuth();
 
   const messageRateHandler = useRate({ threshold: 250 });
 
@@ -168,6 +170,7 @@ export function WsClientProvider({
     const query = {
       latest_event_id: lastEvent?.id ?? -1,
       conversation_id: conversationId,
+      providers_set: providerTokensSet,
     };
 
     const baseUrl =
diff --git a/openhands/integrations/provider.py b/openhands/integrations/provider.py
index 2843dfc4278a..2aeb4f6099d3 100644
--- a/openhands/integrations/provider.py
+++ b/openhands/integrations/provider.py
@@ -14,7 +14,6 @@
 )
 from pydantic.json import pydantic_encoder
 
-from openhands.core.logger import openhands_logger as logger
 from openhands.events.action.action import Action
 from openhands.events.action.commands import CmdRunAction
 from openhands.events.stream import EventStream
@@ -319,9 +318,7 @@ async def get_env_vars(
             get_latest: Get the latest working token for the providers if True, otherwise get the existing ones
         """
 
-        # TODO: We should remove `not get_latest` in the future. More
-        # details about the error this fixes is in the next comment below
-        if not self.provider_tokens and not get_latest:
+        if not self.provider_tokens:
             return {}
 
         env_vars: dict[ProviderType, SecretStr] = {}
@@ -342,20 +339,6 @@ async def get_env_vars(
                 if token:
                     env_vars[provider] = token
 
-        # TODO: we have an error where reinitializing the runtime doesn't happen with
-        # the provider tokens; thus the code above believes that github isn't a provider
-        # when it really is. We need to share information about current providers set
-        # for the user when the socket event for connect is sent
-        if ProviderType.GITHUB not in env_vars and get_latest:
-            logger.info(
-                f'Force refresh runtime token for user: {self.external_auth_id}'
-            )
-            service = GithubServiceImpl(
-                external_auth_id=self.external_auth_id,
-                external_token_manager=self.external_token_manager,
-            )
-            env_vars[ProviderType.GITHUB] = await service.get_latest_token()
-
         if not expose_secrets:
             return env_vars
 
diff --git a/openhands/server/listen_socket.py b/openhands/server/listen_socket.py
index c432c418b575..37a1f47e84f4 100644
--- a/openhands/server/listen_socket.py
+++ b/openhands/server/listen_socket.py
@@ -1,3 +1,4 @@
+from types import MappingProxyType
 from urllib.parse import parse_qs
 
 from socketio.exceptions import ConnectionRefusedError
@@ -16,6 +17,9 @@
     RecallObservation,
 )
 from openhands.events.serialization import event_to_dict
+from openhands.integrations.provider import PROVIDER_TOKEN_TYPE, ProviderToken
+from openhands.integrations.service_types import ProviderType
+from openhands.server.session.conversation_init_data import ConversationInitData
 from openhands.server.shared import (
     SettingsStoreImpl,
     config,
@@ -27,12 +31,26 @@
 )
 
 
+def create_provider_tokens_object(
+    providers_set: list[ProviderType],
+) -> PROVIDER_TOKEN_TYPE:
+    provider_information = {}
+
+    for provider in providers_set:
+        provider_information[provider] = ProviderToken(token=None, user_id=None)
+
+    return MappingProxyType(provider_information)
+
+
 @sio.event
 async def connect(connection_id: str, environ):
     logger.info(f'sio:connect: {connection_id}')
     query_params = parse_qs(environ.get('QUERY_STRING', ''))
     latest_event_id = int(query_params.get('latest_event_id', [-1])[0])
     conversation_id = query_params.get('conversation_id', [None])[0]
+    providers_raw: list[str] = query_params.get('providers_set', [])
+    providers_set: list[ProviderType] = [ProviderType(p) for p in providers_raw]
+
     if not conversation_id:
         logger.error('No conversation_id in query params')
         raise ConnectionRefusedError('No conversation_id in query params')
@@ -50,9 +68,17 @@ async def connect(connection_id: str, environ):
         raise ConnectionRefusedError(
             'Settings not found', {'msg_id': 'CONFIGURATION$SETTINGS_NOT_FOUND'}
         )
+    session_init_args: dict = {}
+    if settings:
+        session_init_args = {**settings.__dict__, **session_init_args}
+
+    session_init_args['git_provider_tokens'] = create_provider_tokens_object(
+        providers_set
+    )
+    conversation_init_data = ConversationInitData(**session_init_args)
 
     event_stream = await conversation_manager.join_conversation(
-        conversation_id, connection_id, settings, user_id, github_user_id
+        conversation_id, connection_id, conversation_init_data, user_id, github_user_id
     )
     logger.info(
         f'Connected to conversation {conversation_id} with connection_id {connection_id}. Replaying event stream...'


Reviews:
user: login: rbren
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-04-02T19:22:56Z
user: login: Prateikx
state: COMMENTED
author_association: NONE
submitted_at: 2025-04-02T20:12:13Z
user: login: malhotra5
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-03T02:00:00Z

Review comments:
diff_hunk: @@ -27,12 +31,26 @@
 )
 
 
+def create_provider_tokens_object(
+    providers_set: list[ProviderType],
+) -> PROVIDER_TOKEN_TYPE:
+    provider_information = {}
+
+    for provider in providers_set:
+        provider_information[provider] = ProviderToken(token=None, user_id=None)
+
+    return MappingProxyType(provider_information)
+
+
 @sio.event
 async def connect(connection_id: str, environ):
     logger.info(f'sio:connect: {connection_id}')
     query_params = parse_qs(environ.get('QUERY_STRING', ''))
     latest_event_id = int(query_params.get('latest_event_id', [-1])[0])
     conversation_id = query_params.get('conversation_id', [None])[0]
+    providers_raw: list[str] = query_params.get('providers_set', [])
path: openhands/server/listen_socket.py
user: login: Prateikx
body: **✨ AI Suggestion:** The `providers_set` variable is created by iterating through `providers_raw` and converting each element to `ProviderType`. If any element in `providers_raw` is not a valid `ProviderType`, a `ValueError` will be raised, potentially causing the connection to fail. It's important to handle this potential error to prevent unexpected behavior.
```suggestion
providers_raw: list[str] = query_params.get('providers_set', [])
providers_set: list[ProviderType] = []
for p in providers_raw:
    try:
        providers_set.append(ProviderType(p))
    except ValueError:
        logger.warning(f"Invalid provider type received: {p}")
        # Handle the error appropriately, e.g., skip the invalid provider or disconnect
```
created_at: 2025-04-02T20:12:13Z
author_association: NONE
diff_hunk: @@ -27,12 +31,26 @@
 )
 
 
+def create_provider_tokens_object(
+    providers_set: list[ProviderType],
+) -> PROVIDER_TOKEN_TYPE:
+    provider_information = {}
+
+    for provider in providers_set:
+        provider_information[provider] = ProviderToken(token=None, user_id=None)
+
+    return MappingProxyType(provider_information)
+
+
 @sio.event
 async def connect(connection_id: str, environ):
     logger.info(f'sio:connect: {connection_id}')
     query_params = parse_qs(environ.get('QUERY_STRING', ''))
     latest_event_id = int(query_params.get('latest_event_id', [-1])[0])
     conversation_id = query_params.get('conversation_id', [None])[0]
+    providers_raw: list[str] = query_params.get('providers_set', [])
path: openhands/server/listen_socket.py
user: login: malhotra5
body: We want the application to fail hard is we ever come across this error.

The provider values should be consistent globally :)
created_at: 2025-04-03T02:00:00Z
author_association: COLLABORATOR

PR comments:
user: login: malhotra5
created_at: 2025-04-02T18:18:10Z
body: @openhands please fix the failing frontend issues
author_association: COLLABORATOR
user: login: openhands-ai[bot]
created_at: 2025-04-02T18:18:18Z
body: I'm on it! @malhotra5 can [track my progress at all-hands.dev](https://app.all-hands.dev/conversations/e00b8c19973f47a0abe1c9ce332da8c3)
author_association: NONE
user: login: malhotra5
created_at: 2025-04-02T18:18:30Z
body: @openhands please fix the failing frontend tests
author_association: COLLABORATOR
user: login: openhands-ai[bot]
created_at: 2025-04-02T18:18:39Z
body: I'm on it! @malhotra5 can [track my progress at all-hands.dev](https://app.all-hands.dev/conversations/d43320fdac2842e1aa0960a23783b1f4)
author_association: NONE
user: login: openhands-ai[bot]
created_at: 2025-04-02T18:19:12Z
body: I haven't been able to complete the task yet. Here's a summary of what I've found and what needs to be done:

## Current Status
- I've examined the PR #7670 on the cloud-refresh-fix branch
- The PR appears to be removing a temporary refresh hack in the codebase
- Key changes include:
  - Downgrading vite from 6.2.4 to 6.2.3 in package-lock.json
  - Adding a `providerTokensSet` parameter to WebSocket connections in the frontend
  - Removing a GitHub token refresh workaround in the backend
  - Downgrading OpenVSCode Server from v1.98.2 to v1.94.2
  - Modifying the socket connection handling to use provider tokens

## Issues Identified
- I attempted to run the frontend linting but encountered an error because the dependencies weren't installed
- When I tried to install the dependencies, that command also failed

## Next Steps
To fix the failing frontend issues, I would need to:
1. Successfully install the frontend dependencies
2. Run the linting and tests to identify specific failures
3. Fix any TypeScript/React issues related to the changes in the WebSocket client provider
4. Ensure the build process completes successfully

I wasn't able to complete these steps due to environment issues with installing the dependencies.
author_association: NONE
user: login: openhands-ai[bot]
created_at: 2025-04-02T18:23:36Z
body: I've successfully fixed the failing frontend tests in the cloud-refresh-fix branch. Here's a summary of what I did:

## Issue Identified
The failing test was in `__tests__/context/ws-client-provider.test.tsx` with the error: "useAuth must be used within a AuthProvider". This occurred because recent changes to the `ws-client-provider.tsx` file added a dependency on the Auth context (`const { providerTokensSet } = useAuth();`), but the test wasn't properly set up to provide this context.

## Solution Implemented
1. I imported the `AuthProvider` component from `"#/context/auth-context"` in the test file
2. Updated the test wrapper to include the `AuthProvider` component:
   ```jsx
   <QueryClientProvider client={new QueryClient()}>
     <AuthProvider initialProviderTokens={[]}>
       <WsClientProvider conversationId="test-conversation-id">
         {children}
       </WsClientProvider>
     </AuthProvider>
   </QueryClientProvider>
   ```
3. Committed the changes with a descriptive message: "Fix ws-client-provider test by adding AuthProvider to test wrapper"
4. Pushed the changes to the cloud-refresh-fix branch

## Results
All frontend tests are now passing. The PR should now be ready for review and merging.
author_association: NONE