From 9d8a855a4e2295e5174e8fad531180225077c9e7 Mon Sep 17 00:00:00 2001
From: ross-rl <ross@runloop.ai>
Date: Wed, 27 Nov 2024 09:30:23 -0600
Subject: [PATCH] Add vscode url to runloop runtime

---
 .../runtime/impl/runloop/runloop_runtime.py   | 50 ++++++++++++++++++-
 1 file changed, 48 insertions(+), 2 deletions(-)

diff --git a/openhands/runtime/impl/runloop/runloop_runtime.py b/openhands/runtime/impl/runloop/runloop_runtime.py
index 20f0ea46a1b8..76f9b254fdcf 100644
--- a/openhands/runtime/impl/runloop/runloop_runtime.py
+++ b/openhands/runtime/impl/runloop/runloop_runtime.py
@@ -99,6 +99,7 @@ class RunloopRuntime(EventStreamRuntime):
     """The RunloopRuntime class is an EventStreamRuntime that utilizes Runloop Devbox as a runtime environment."""
 
     _sandbox_port: int = 4444
+    _vscode_port: int = 4445
 
     def __init__(
         self,
@@ -109,6 +110,7 @@ def __init__(
         env_vars: dict[str, str] | None = None,
         status_callback: Callable | None = None,
         attach_to_existing: bool = False,
+        headless_mode: bool = True,
     ):
         assert config.runloop_api_key is not None, 'Runloop API key is required'
         self.devbox: DevboxView | None = None
@@ -127,9 +129,11 @@ def __init__(
             env_vars,
             status_callback,
             attach_to_existing,
+            headless_mode,
         )
         # Buffer for container logs
         self.log_buffer: LogBuffer | None = None
+        self._vscode_url: str | None = None
 
     @tenacity.retry(
         stop=tenacity.stop_after_attempt(120),
@@ -192,7 +196,7 @@ def _create_new_devbox(self) -> DevboxView:
             environment_variables={'DEBUG': 'true'} if self.config.debug else {},
             prebuilt='openhands',
             launch_parameters=LaunchParameters(
-                available_ports=[self._sandbox_port],
+                available_ports=[self._sandbox_port, self._vscode_port],
                 resource_size_request='LARGE',
             ),
             metadata={'container-name': self.container_name},
@@ -221,7 +225,7 @@ async def connect(self):
 
         # Hook up logs
         self.log_buffer = RunloopLogBuffer(self.runloop_api_client, self.devbox.id)
-        self.api_url = f'https://{tunnel.url}'
+        self.api_url = tunnel.url
         logger.info(f'Container started. Server url: {self.api_url}')
 
         # End Runloop connect
@@ -273,3 +277,45 @@ def close(self, rm_all_containers: bool | None = True):
 
         if self.devbox:
             self.runloop_api_client.devboxes.shutdown(self.devbox.id)
+
+    @property
+    def vscode_url(self) -> str | None:
+        if self.vscode_enabled and self.devbox and self.devbox.status == 'running':
+            if self._vscode_url is not None:
+                return self._vscode_url
+
+            try:
+                with send_request(
+                    self.session,
+                    'GET',
+                    f'{self.api_url}/vscode/connection_token',
+                    timeout=10,
+                ) as response:
+                    response_json = response.json()
+                    assert isinstance(response_json, dict)
+                    if response_json['token'] is None:
+                        return None
+                    token = response_json['token']
+
+                self._vscode_url = (
+                    self.runloop_api_client.devboxes.create_tunnel(
+                        id=self.devbox.id,
+                        port=self._vscode_port,
+                    ).url
+                    + f'/?tkn={token}&folder={self.config.workspace_mount_path_in_sandbox}'
+                )
+
+                self.log(
+                    'debug',
+                    f'VSCode URL: {self._vscode_url}',
+                )
+
+                return self._vscode_url
+            except Exception as e:
+                self.log(
+                    'error',
+                    f'Failed to create vscode tunnel {e}',
+                )
+                return None
+        else:
+            return None