From 81ced64019c72335ee6c22d808cee91b4304ca8d Mon Sep 17 00:00:00 2001
From: Wei Ouyang <oeway007@gmail.com>
Date: Sat, 28 Sep 2024 17:04:22 -0700
Subject: [PATCH] Support current workspace in http proxy (#691)

* support _token and remove prefix in login

* allow pass workspace when login

* Support current workspace for http endpoint

* fix login optional

* increase page size for list workspaces

* Bump version for hypha-rpc 0.20.38

* Update change logs and login instructions
---
 CHANGELOG.md                                  |  8 ++
 docs/getting-started.md                       | 26 ++++++-
 docs/migration-guide.md                       | 10 +--
 docs/service-type-annotation.md               |  2 +-
 helm-charts/aks-hypha.md                      |  2 +-
 helm-charts/hypha-server/Chart.yaml           |  2 +-
 helm-charts/hypha-server/values.yaml          |  2 +-
 hypha/VERSION                                 |  2 +-
 hypha/apps.py                                 |  6 +-
 hypha/core/__init__.py                        |  2 +
 hypha/core/auth.py                            | 65 ++++++++++++----
 hypha/core/store.py                           | 28 ++++---
 hypha/core/workspace.py                       |  4 +-
 hypha/http.py                                 | 16 ++--
 hypha/templates/apps/login_template.html      | 77 +++++++++++--------
 .../hypha-core-app/hypha-app-webpython.js     |  2 +-
 hypha/templates/ws/index.html                 |  2 +-
 hypha/websocket.py                            |  9 +++
 requirements.txt                              |  2 +-
 setup.py                                      |  2 +-
 20 files changed, 183 insertions(+), 86 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index d5fba2c7..ff7ed89e 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,5 +1,13 @@
 # Hypha Change Log
 
+### 0.20.38
+ - Allow passing workspace and expires_in to the `login` function to generate workspace specific token.
+ - When using http endpoint to access the service, you can now pass workspace specific token to the http header `Authorization` to access the service. (Previously, all the services are assumed to be accessed from the same service provider workspace)
+
+### 0.20.37
+ - Add s3-proxy to allow accessing s3 presigned url in case the s3 server is not directly accessible. Use `--enable-s3-proxy` to enable the s3 proxy when starting Hypha.
+ - Add `artifact-manager` service to provide comprehensive artifact management, used for creating gallery-like service portal. The artifact manager service is backed by s3 storage and supports presigned url for direct access to the artifacts. This is a replacement of the previous `card` service.
+
 ### 0.20.36
 
  - Upgrade hypha-rpc to support updating reconnection token (otherwise it generate token expired error after some time)
diff --git a/docs/getting-started.md b/docs/getting-started.md
index 02a45dd7..96bba80e 100644
--- a/docs/getting-started.md
+++ b/docs/getting-started.md
@@ -220,7 +220,7 @@ svc = await get_remote_service("http://localhost:9527/ws-user-scintillating-lawy
 Include the following script in your HTML file to load the `hypha-rpc` client:
 
 ```html
-<script src="https://cdn.jsdelivr.net/npm/hypha-rpc@0.20.37/dist/hypha-rpc-websocket.min.js"></script>
+<script src="https://cdn.jsdelivr.net/npm/hypha-rpc@0.20.38/dist/hypha-rpc-websocket.min.js"></script>
 ```
 
 Use the following code in JavaScript to connect to the server and access an existing service:
@@ -250,7 +250,7 @@ By default all the clients connected to Hypha server communicate via the websock
 
 ### User Login and Token-Based Authentication
 
-To access the full features of the Hypha server, users need to log in and obtain a token for authentication. The new `login()` function provides a convenient way to display a login URL, once the user click it and login, it can then return the token for connecting to the server.
+To access the full features of the Hypha server, users need to log in and obtain a token for authentication. The `login()` function provides a convenient way to display a login URL, once the user click it and login, it can then return the token for connecting to the server.
 
 Here is an example of how the login process works using the `login()` function:
 
@@ -294,7 +294,27 @@ The output will provide a URL for the user to open in their browser and
 
  perform the login process. Once the user clicks the link and successfully logs in, the `login()` function will return, providing the token.
 
-The `login()` function also supports additional arguments:
+#### Additional Arguments for Login
+
+You can specify the `workspace` and `expires_in` arguments in the `login()` function so that the token is generated for a specific workspace and expires after a certain period of time.
+
+```python
+token = await login(
+    {
+        "server_url": SERVER_URL,
+        "workspace": "my-workspace",
+        "expires_in": 3600,
+    }
+)
+```
+
+If a token is generated for a specific workspace, when calling `connect_to_server`, you need to specify the same workspace as well:
+
+```python
+server = await connect_to_server({"server_url": "https://ai.imjoy.io", "token": token, "workspace": "my-workspace"})
+```
+
+The `login()` function also supports other additional arguments:
 
 ```python
 token = await login(
diff --git a/docs/migration-guide.md b/docs/migration-guide.md
index f8dba6d3..979e3be7 100644
--- a/docs/migration-guide.md
+++ b/docs/migration-guide.md
@@ -15,7 +15,7 @@ To connect to the server, instead of installing the `imjoy-rpc` module, you will
 pip install -U hypha-rpc # new install
 ```
 
-We also changed our versioning strategy, we use the same version number for the server and client, so it's easier to match the client and server versions. For example, `hypha-rpc` version `0.20.37` is compatible with Hypha server version `0.20.37`.
+We also changed our versioning strategy, we use the same version number for the server and client, so it's easier to match the client and server versions. For example, `hypha-rpc` version `0.20.38` is compatible with Hypha server version `0.20.38`.
 
 #### 2. Change the imports to use `hypha-rpc`
 
@@ -128,10 +128,10 @@ loop.run_forever()
 To connect to the server, instead of using the `imjoy-rpc` module, you will need to use the `hypha-rpc` module. The `hypha-rpc` module is a standalone module that provides the RPC connection to the Hypha server. You can include it in your HTML using a script tag:
 
 ```html
-<script src="https://cdn.jsdelivr.net/npm/hypha-rpc@0.20.37/dist/hypha-rpc-websocket.min.js"></script>
+<script src="https://cdn.jsdelivr.net/npm/hypha-rpc@0.20.38/dist/hypha-rpc-websocket.min.js"></script>
 ```
 
-We also changed our versioning strategy, we use the same version number for the server and client, so it's easier to match the client and server versions. For example, `hypha-rpc` version `0.20.37` is compatible with Hypha server version `0.20.37`.
+We also changed our versioning strategy, we use the same version number for the server and client, so it's easier to match the client and server versions. For example, `hypha-rpc` version `0.20.38` is compatible with Hypha server version `0.20.38`.
 
 #### 2. Change the connection method and use camelCase for service function names
 
@@ -149,7 +149,7 @@ Here is a suggested list of search and replace operations to update your code:
 Here is an example of how the updated code might look:
 
 ```html
-<script src="https://cdn.jsdelivr.net/npm/hypha-rpc@0.20.37/dist/hypha-rpc-websocket.min.js"></script>
+<script src="https://cdn.jsdelivr.net/npm/hypha-rpc@0.20.38/dist/hypha-rpc-websocket.min.js"></script>
 <script>
 async function main(){
     const server = await hyphaWebsocketClient.connectToServer({"server_url": "https://hypha.amun.ai"});
@@ -197,7 +197,7 @@ We created a tutorial to introduce this new feature: [service type annotation](.
 Here is a quick example in JavaScript:
 
 ```html
-<script src="https://cdn.jsdelivr.net/npm/hypha-rpc@0.20.37/dist/hypha-rpc-websocket.min.js"></script>
+<script src="https://cdn.jsdelivr.net/npm/hypha-rpc@0.20.38/dist/hypha-rpc-websocket.min.js"></script>
 
 <script>
 async function main(){
diff --git a/docs/service-type-annotation.md b/docs/service-type-annotation.md
index bb8ea0e1..d5df4806 100644
--- a/docs/service-type-annotation.md
+++ b/docs/service-type-annotation.md
@@ -134,7 +134,7 @@ if __name__ == "__main__":
 **JavaScript Client: Service Usage**
 
 ```html
-<script src="https://cdn.jsdelivr.net/npm/hypha-rpc@0.20.37/dist/hypha-rpc-websocket.min.js"></script>
+<script src="https://cdn.jsdelivr.net/npm/hypha-rpc@0.20.38/dist/hypha-rpc-websocket.min.js"></script>
 <script>
 async function main() {
     const server = await hyphaWebsocketClient.connectToServer({"server_url": "https://hypha.amun.ai"});
diff --git a/helm-charts/aks-hypha.md b/helm-charts/aks-hypha.md
index 08611056..7aebdfac 100644
--- a/helm-charts/aks-hypha.md
+++ b/helm-charts/aks-hypha.md
@@ -146,7 +146,7 @@ replicaCount: 1
 image:
   repository: ghcr.io/amun-ai/hypha
   pullPolicy: IfNotPresent
-  tag: "0.20.37"
+  tag: "0.20.38"
 
 serviceAccount:
   create: true
diff --git a/helm-charts/hypha-server/Chart.yaml b/helm-charts/hypha-server/Chart.yaml
index f0ca269f..e11fd706 100644
--- a/helm-charts/hypha-server/Chart.yaml
+++ b/helm-charts/hypha-server/Chart.yaml
@@ -15,7 +15,7 @@ type: application
 # This is the chart version. This version number should be incremented each time you make changes
 # to the chart and its templates, including the app version.
 # Versions are expected to follow Semantic Versioning (https://semver.org/)
-version: 0.20.37
+version: 0.20.38
 
 # This is the version number of the application being deployed. This version number should be
 # incremented each time you make changes to the application. Versions are not expected to
diff --git a/helm-charts/hypha-server/values.yaml b/helm-charts/hypha-server/values.yaml
index 5384ee5b..8bd42ed1 100644
--- a/helm-charts/hypha-server/values.yaml
+++ b/helm-charts/hypha-server/values.yaml
@@ -8,7 +8,7 @@ image:
   repository: ghcr.io/amun-ai/hypha
   pullPolicy: IfNotPresent
   # Overrides the image tag whose default is the chart appVersion.
-  tag: "0.20.37"
+  tag: "0.20.38"
 
 imagePullSecrets: []
 nameOverride: ""
diff --git a/hypha/VERSION b/hypha/VERSION
index 14f65389..6e281531 100644
--- a/hypha/VERSION
+++ b/hypha/VERSION
@@ -1,3 +1,3 @@
 {
-    "version": "0.20.37.post3"
+    "version": "0.20.38"
 }
diff --git a/hypha/apps.py b/hypha/apps.py
index 85bab035..5038140a 100644
--- a/hypha/apps.py
+++ b/hypha/apps.py
@@ -360,7 +360,7 @@ async def start(
         workspace = context["ws"]
         user_info = UserInfo.model_validate(context["user"])
 
-        async with self.store.get_workspace_interface(workspace, user_info) as ws:
+        async with self.store.get_workspace_interface(user_info, workspace) as ws:
             token = await ws.generate_token()
 
         if not user_info.check_permission(workspace, UserPermission.read):
@@ -520,7 +520,9 @@ async def list_running(self, context: Optional[dict] = None) -> List[str]:
     async def list_apps(self, context: Optional[dict] = None):
         """List applications in the workspace."""
         try:
-            apps = await self.artifact_manager.read(prefix="applications", context=context)
+            apps = await self.artifact_manager.read(
+                prefix="applications", context=context
+            )
             return apps["collection"]
         except KeyError:
             return []
diff --git a/hypha/core/__init__.py b/hypha/core/__init__.py
index 1d373ac3..83c3e2b7 100644
--- a/hypha/core/__init__.py
+++ b/hypha/core/__init__.py
@@ -115,6 +115,8 @@ class UserTokenInfo(BaseModel):
     """Represent user profile."""
 
     token: constr(max_length=1024)  # type: ignore
+    workspace: Optional[str] = None
+    expires_in: Optional[int] = None
     email: Optional[EmailStr] = None
     email_verified: Optional[bool] = None
     name: Optional[constr(max_length=64)] = None  # type: ignore
diff --git a/hypha/core/auth.py b/hypha/core/auth.py
index d478af45..b240d580 100644
--- a/hypha/core/auth.py
+++ b/hypha/core/auth.py
@@ -36,6 +36,7 @@
 AUTH0_NAMESPACE = env.get("AUTH0_NAMESPACE", "https://amun.ai/")
 JWT_SECRET = env.get("JWT_SECRET")
 LOGIN_SERVICE_URL = "/public/services/hypha-login"
+LOGIN_KEY_PREFIX = "login_key:"
 
 if not JWT_SECRET:
     logger.warning(
@@ -67,6 +68,9 @@ def get_user_info(credentials):
         scope=scope,
         expires_at=expires_at,
     )
+    # make sure the user has admin permission to their own workspace
+    user_workspace = info.get_workspace()
+    info.scope.workspaces[user_workspace] = UserPermission.admin
     return info
 
 
@@ -292,7 +296,7 @@ def create_scope(
 
 
 def update_user_scope(
-    user_info: UserInfo, workspace_info: WorkspaceInfo, client_id: str
+    user_info: UserInfo, workspace_info: WorkspaceInfo, client_id: str = None
 ):
     """Update the user scope for a workspace."""
     user_info.scope = user_info.scope or ScopeInfo()
@@ -355,13 +359,20 @@ def create_login_service(store):
         auth0_issuer=AUTH0_ISSUER,
     )
 
-    async def start_login():
+    async def start_login(workspace: str = None, expires_in: int = None):
         """Start the login process."""
-        key = "login_key:" + str(random_id(readable=False))
+        key = str(random_id(readable=False))
         # set the key and with expire time
-        await redis.setex(key, MAXIMUM_LOGIN_TIME, "")
+        await redis.setex(LOGIN_KEY_PREFIX + key, MAXIMUM_LOGIN_TIME, "")
         return {
-            "login_url": f"{login_service_url.replace('/services/', '/apps/')}/?key={key}",
+            "login_url": f"{login_service_url.replace('/services/', '/apps/')}/?key={key}"
+            + (
+                f"&workspace={workspace}"
+                if workspace
+                else "" + f"&expires_in={expires_in}"
+                if expires_in
+                else ""
+            ),
             "key": key,
             "report_url": f"{login_service_url}/report",
             "check_url": f"{login_service_url}/check",
@@ -378,15 +389,17 @@ async def index(event):
 
     async def check_login(key, timeout=MAXIMUM_LOGIN_TIME, profile=False):
         """Check the status of a login session."""
-        assert await redis.exists(key), "Invalid key, key does not exist"
+        assert await redis.exists(
+            LOGIN_KEY_PREFIX + key
+        ), "Invalid key, key does not exist"
         if timeout <= 0:
-            user_info = await redis.get(key)
+            user_info = await redis.get(LOGIN_KEY_PREFIX + key)
             if user_info == b"":
                 return None
             user_info = json.loads(user_info)
             user_info = UserTokenInfo.model_validate(user_info)
             if user_info:
-                await redis.delete(key)
+                await redis.delete(LOGIN_KEY_PREFIX + key)
             return (
                 user_info.model_dump(mode="json")
                 if profile
@@ -394,7 +407,7 @@ async def check_login(key, timeout=MAXIMUM_LOGIN_TIME, profile=False):
             )
         count = 0
         while True:
-            user_info = await redis.get(key)
+            user_info = await redis.get(LOGIN_KEY_PREFIX + key)
             if user_info != b"":
                 user_info = json.loads(user_info)
                 user_info = UserTokenInfo.model_validate(user_info)
@@ -403,7 +416,7 @@ async def check_login(key, timeout=MAXIMUM_LOGIN_TIME, profile=False):
                         f"Login session expired, the maximum login time is {MAXIMUM_LOGIN_TIME} seconds"
                     )
                 if user_info:
-                    await redis.delete(key)
+                    await redis.delete(LOGIN_KEY_PREFIX + key)
                     return (
                         user_info.model_dump(mode="json")
                         if profile
@@ -417,6 +430,8 @@ async def check_login(key, timeout=MAXIMUM_LOGIN_TIME, profile=False):
     async def report_login(
         key,
         token,
+        workspace=None,
+        expires_in=None,
         email=None,
         email_verified=None,
         name=None,
@@ -425,9 +440,14 @@ async def report_login(
         picture=None,
     ):
         """Report a token associated with a login session."""
-        assert await redis.exists(key), "Invalid key, key does not exist or expired"
+        assert await redis.exists(
+            LOGIN_KEY_PREFIX + key
+        ), "Invalid key, key does not exist or expired"
+        # workspace = workspace or ("ws-user-" + user_id)
         kwargs = {
             "token": token,
+            "workspace": workspace,
+            "expires_in": expires_in or None,
             "email": email,
             "email_verified": email_verified,
             "name": name,
@@ -435,9 +455,26 @@ async def report_login(
             "user_id": user_id,
             "picture": picture,
         }
-        user_info = UserTokenInfo.model_validate(kwargs)
-        user_info = user_info.model_dump(mode="json")
-        await redis.setex(key, MAXIMUM_LOGIN_TIME, json.dumps(user_info))
+
+        user_token_info = UserTokenInfo.model_validate(kwargs)
+        if workspace:
+            user_info = parse_token(token)
+            # based on the user token, create a scoped token
+            workspace = workspace or user_info.get_workspace()
+            # generate scoped token
+            workspace_info = await store.load_or_create_workspace(user_info, workspace)
+            user_info.scope = update_user_scope(user_info, workspace_info)
+            if not user_info.check_permission(workspace, UserPermission.read):
+                raise Exception(f"Invalid permission for the workspace {workspace}")
+
+            token = generate_presigned_token(user_info, int(expires_in or 3600))
+            # replace the token
+            user_token_info.token = token
+        await redis.setex(
+            LOGIN_KEY_PREFIX + key,
+            MAXIMUM_LOGIN_TIME,
+            user_token_info.model_dump_json(),
+        )
 
     logger.info(
         f"To preview the login page, visit: {login_service_url.replace('/services/', '/apps/')}"
diff --git a/hypha/core/store.py b/hypha/core/store.py
index a008ee23..1ffcf950 100644
--- a/hypha/core/store.py
+++ b/hypha/core/store.py
@@ -6,7 +6,7 @@
 import datetime
 from typing import List, Union
 from pydantic import BaseModel
-from fastapi import Header, Cookie
+from fastapi import Header, Cookie, Query
 
 from hypha_rpc import RPC
 from hypha_rpc.utils.schema import schema_method
@@ -440,8 +440,8 @@ async def init(self, reset_redis, startup_functions=None):
     async def _register_root_services(self):
         """Register root services."""
         self._root_workspace_interface = await self.get_workspace_interface(
-            self._root_user.get_workspace(),
             self._root_user,
+            self._root_user.get_workspace(),
             client_id=self._server_id,
             silent=False,
         )
@@ -476,7 +476,7 @@ async def get_public_api(self):
         """Get the public API."""
         if self._public_workspace_interface is None:
             self._public_workspace_interface = await self.get_workspace_interface(
-                "public", self._root_user, client_id=self._server_id, silent=False
+                self._root_user, "public", client_id=self._server_id, silent=False
             )
         return self._public_workspace_interface
 
@@ -527,17 +527,14 @@ async def parse_user_token(self, token):
         key = "revoked_token:" + token
         if await self._redis.exists(key):
             raise Exception("Token has been revoked")
-        # automatically add user's own workspace to the scope
-        if not user_info.scope.workspaces:
-            user_info.scope.workspaces = {
-                user_info.get_workspace(): UserPermission.admin
-            }
         if "admin" in user_info.roles:
             user_info.scope.workspaces["*"] = UserPermission.admin
         return user_info
 
     async def login_optional(
-        self, authorization: str = Header(None), access_token: str = Cookie(None)
+        self,
+        authorization: str = Header(None),
+        access_token: str = Cookie(None),
     ):
         """Return user info or create an anonymouse user.
 
@@ -547,9 +544,16 @@ async def login_optional(
         token = authorization or access_token
         if token:
             user_info = await self.parse_user_token(token)
+            if user_info.scope.current_workspace is None:
+                user_info.scope.current_workspace = user_info.get_workspace()
             return user_info
         else:
-            return generate_anonymous_user()
+            user_info = generate_anonymous_user()
+            user_workspace = user_info.get_workspace()
+            user_info.scope = create_scope(
+                f"{user_workspace}#a", current_workspace=user_workspace
+            )
+            return user_info
 
     async def get_all_workspace(self):
         """Get all workspaces."""
@@ -588,7 +592,7 @@ def connect_to_workspace(
             client_id = self._server_id + "-" + random_id(readable=False)
         user_info = user_info or self._root_user
         return self.get_workspace_interface(
-            workspace, user_info, client_id=client_id, timeout=timeout, silent=silent
+            user_info, workspace, client_id=client_id, timeout=timeout, silent=silent
         )
 
     def get_manager_id(self):
@@ -614,8 +618,8 @@ def get_server_info(self):
 
     def get_workspace_interface(
         self,
-        workspace: str,
         user_info: UserInfo,
+        workspace: str,
         client_id=None,
         timeout=10,
         silent=True,
diff --git a/hypha/core/workspace.py b/hypha/core/workspace.py
index 2266d657..5846fba9 100644
--- a/hypha/core/workspace.py
+++ b/hypha/core/workspace.py
@@ -1213,7 +1213,7 @@ async def list_workspaces(
         self,
         match: dict = Field(None, description="Match pattern for filtering workspaces"),
         page: int = Field(1, description="Page number for pagination"),
-        page_size: int = Field(10, description="Number of items per page"),
+        page_size: int = Field(256, description="Number of items per page"),
         context=None,
     ) -> List[Dict[str, Any]]:
         """Get all workspaces with pagination."""
@@ -1347,7 +1347,7 @@ async def unload_if_empty(self, context=None):
             await self.unload(context=context)
         else:
             logger.warning(
-                f"Workspace {context['ws']} is not empty, remaining clients: {client_keys[:10]}..."  # only list the first 10 maximum
+                f"Skip unloading workspace {context['ws']} because it is not empty, remaining clients: {client_keys[:10]}..."  # only list the first 10 maximum
             )
 
     async def unload(self, context=None):
diff --git a/hypha/http.py b/hypha/http.py
index b22562ce..db3ae21b 100644
--- a/hypha/http.py
+++ b/hypha/http.py
@@ -240,11 +240,11 @@ async def __call__(self, scope, receive, send):
                     )
 
                     async with self.store.get_workspace_interface(
-                        workspace, user_info
+                        user_info, user_info.scope.current_workspace
                     ) as api:
                         # Call get_service_type_id to check if it's an ASGI service
                         service_info = await api.get_service_info(
-                            service_id, {"mode": _mode}
+                            workspace + "/" + service_id, {"mode": _mode}
                         )
                         service = await api.get_service(service_info.id)
                         # intercept the request if it's an ASGI service
@@ -521,7 +521,7 @@ async def get_workspace_services(
             """Route for get services under a workspace."""
             try:
                 async with self.store.get_workspace_interface(
-                    workspace, user_info
+                    user_info, workspace
                 ) as manager:
                     services = await manager.list_services()
                 info = serialize(services)
@@ -548,7 +548,7 @@ async def get_service_info(
             """Route for checking details of a service."""
             try:
                 async with self.store.get_workspace_interface(
-                    workspace, user_info
+                    user_info, workspace
                 ) as api:
                     if service_id == "ws":
                         return serialize(api)
@@ -584,7 +584,7 @@ async def get_workspace_apps(
             """Route for get apps under a workspace."""
             try:
                 async with self.store.get_workspace_interface(
-                    workspace, user_info
+                    user_info, user_info.scope.current_workspace
                 ) as manager:
                     try:
                         controller = await manager.get_service("public/server-apps")
@@ -776,12 +776,14 @@ async def service_function(
             try:
                 workspace, service_id, function_key = function_info
                 async with self.store.get_workspace_interface(
-                    workspace, user_info
+                    user_info, user_info.scope.current_workspace
                 ) as api:
                     if service_id == "ws":
                         service = api
                     else:
-                        info = await api.get_service_info(service_id, {"mode": _mode})
+                        info = await api.get_service_info(
+                            workspace + "/" + service_id, {"mode": _mode}
+                        )
                         service = await api.get_service(info.id)
                     func = get_value(function_key, service)
                     if not func:
diff --git a/hypha/templates/apps/login_template.html b/hypha/templates/apps/login_template.html
index 3f8113af..dfb3b0f9 100644
--- a/hypha/templates/apps/login_template.html
+++ b/hypha/templates/apps/login_template.html
@@ -102,39 +102,47 @@
             }
         });
 
+        const _generateToken = async (workspace, expires_in) =>{
+            const generate_token_url = "{{ generate_token_url }}";
+            if(!generate_token_url.startswith("http")){
+                throw new Error("Invalid generate token URL");
+            }
+            // make a post request with {config: {expires_in, workspace}}
+            // add auth barier token to the request
+            const response = await fetch(generate_token_url, {
+                method: "POST",
+                headers: {
+                    "Content-Type": "application/json",
+                    "Authorization": `Bearer ${accessToken}`,
+                },
+                body: JSON.stringify({ config: { expires_in, workspace } }),
+            });
+            if (response.ok) {
+                const token = await response.json();
+                return token;
+            }
+            else{
+                throw new Error(`Failed to generate token. Please try again later. (${response.status}: ${response.statusText})`);
+            }
+        }
+
         window.generateToken = async () => {
             const expires_in = document.getElementById("expiry-time").value || 3600; // Default to 1 hour if not specified
             const workspace = document.getElementById("workspace").value || "";
             const accessToken = await auth0.getTokenSilently({ audience: AUTH0_AUDIENCE });
-
-            const generate_token_url = "{{ generate_token_url }}";
-            if (generate_token_url.startsWith('http')) {
-                // make a post request with {config: {expires_in, workspace}}
-                // add auth barier token to the request
-                const response = await fetch(generate_token_url, {
-                    method: "POST",
-                    headers: {
-                        "Content-Type": "application/json",
-                        "Authorization": `Bearer ${accessToken}`,
-                    },
-                    body: JSON.stringify({ config: { expires_in, workspace } }),
-                });
-                if (response.ok) {
-                    const token = await response.json();
-                    document.getElementById("access-token").innerHTML = token;
-                    document.getElementById("info-panel").style.display = "block";
-                    document.getElementById("toast").innerHTML = "New access token generated successfully!";
-                    document.getElementById("toast").style.display = "block";
-                    setTimeout(() => {
-                        document.getElementById("toast").style.display = "none";
-                    }, 2000);
-                } else {
-                    console.error("Failed to generate token: ", response.status, response.statusText);
-                    alert(`Failed to generate token. Please try again later. (${response.status}: ${response.statusText})`);
-                }
-            } else {
-                console.error("Invalid generate_token_url: ", generate_token_url);
-                alert("Invalid generate_token_url. Please contact the site administrator.");
+            try {
+                const token = await _generateToken(workspace, expires_in);
+                document.getElementById("access-token").innerHTML = token;
+                document.getElementById("info-panel").style.display = "block";
+                document.getElementById("toast").innerHTML = "New access token generated successfully!";
+                document.getElementById("toast").style.display = "block";
+                setTimeout(() => {
+                    document.getElementById("toast").style.display = "none";
+                }, 2000);
+            }
+            catch(e){
+                console.error("Failed to generate token: ", response.status, response.statusText);
+                alert(`Failed to generate token. Please try again later. (${response.status}: ${response.statusText})`);
             }
         };
 
@@ -175,12 +183,17 @@
             document.getElementById("user-id").innerHTML = user.sub;
             document.getElementById("access-token").innerHTML = accessToken;
             document.getElementById("info-panel").style.display = "block";
-            document.getElementById("workspace").value = user.sub;
+            document.getElementById("workspace").value = "ws-user-" + user.sub;
 
-            const key = window.location.search.split('key=')[1];
+            // extract key from search
+            const urlParams = new URLSearchParams(window.location.search);
+            const key = urlParams.get('key');
+            const workspace = urlParams.get('workspace') || ""
+            const expires_in = urlParams.get('expires_in') || ""
             const report_url = "{{ login_service_url }}/report";
+
             if (key && report_url.startsWith('http')) {
-                const url = `${report_url}?key=${key}&token=${accessToken}&email=${user.email}&email_verified=${user.email_verified}&name=${user.name}&nickname=${user.nickname}&user_id=${user.sub}&picture=${encodeURIComponent(user.picture)}`;
+                const url = `${report_url}?key=${key}&token=${accessToken}&workspace=${workspace}&expires_in=${expires_in}&email=${user.email}&email_verified=${user.email_verified}&name=${user.name}&nickname=${user.nickname}&user_id=${user.sub}&picture=${encodeURIComponent(user.picture)}`;
                 const response = await fetch(url);
                 if (!response.ok) {
                     console.error("Failed to report token: ", response.status, response.statusText);
diff --git a/hypha/templates/hypha-core-app/hypha-app-webpython.js b/hypha/templates/hypha-core-app/hypha-app-webpython.js
index 9f9bcb76..27daf5d5 100644
--- a/hypha/templates/hypha-core-app/hypha-app-webpython.js
+++ b/hypha/templates/hypha-core-app/hypha-app-webpython.js
@@ -27,7 +27,7 @@ loadPyodide().then(async (pyodide) => {
     pyodide.setStderr({ batched: (msg) => console.error(msg) });
     await pyodide.loadPackage("micropip");
     const micropip = pyodide.pyimport("micropip");
-    await micropip.install('hypha-rpc==0.20.37');
+    await micropip.install('hypha-rpc==0.20.38');
     const isWindow = typeof window !== "undefined";
     
     setTimeout(() => {
diff --git a/hypha/templates/ws/index.html b/hypha/templates/ws/index.html
index 7dfb7bee..65edebed 100644
--- a/hypha/templates/ws/index.html
+++ b/hypha/templates/ws/index.html
@@ -160,7 +160,7 @@
 <body class="bg-black text-white font-poppins">
     <div id="app"></div>
     <script type="module">
-        import { HyphaCore } from "https://cdn.jsdelivr.net/npm/hypha-core@0.20.37/dist/hypha-core.mjs";
+        import { HyphaCore } from "https://cdn.jsdelivr.net/npm/hypha-core@0.20.38/dist/hypha-core.mjs";
 
         const defaultService = {
             getServerConfig(context){
diff --git a/hypha/websocket.py b/hypha/websocket.py
index 347afae5..9a4b252b 100644
--- a/hypha/websocket.py
+++ b/hypha/websocket.py
@@ -223,6 +223,12 @@ async def authenticate_user(
                 user_info = await self.store.parse_user_token(token)
                 # user token doesn't have client id, so we add that
                 user_info.scope.client_id = client_id
+                if user_info.scope.current_workspace and workspace:
+                    assert (
+                        workspace == user_info.scope.current_workspace
+                    ), f"Current workspace encoded in the token ({scope.current_workspace}) does not the specified workspace ({workspace})"
+                if not workspace and user_info.scope.current_workspace:
+                    workspace = user_info.scope.current_workspace
             else:
                 raise RuntimeError("No authentication information provided")
             return user_info, workspace
@@ -248,6 +254,9 @@ async def force_disconnect(_):
             )
 
         event_bus = self.store.get_event_bus()
+        assert (
+            user_info.scope.current_workspace == workspace
+        ), f"Workspace mismatch: {workspace} != {user_info.current_workspace}"
         conn = RedisRPCConnection(event_bus, workspace, client_id, user_info, None)
         self._websockets[f"{workspace}/{client_id}"] = websocket
         try:
diff --git a/requirements.txt b/requirements.txt
index 0535aed9..4cbb60cf 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -2,7 +2,7 @@ aioboto3==13.1.1
 aiofiles==23.2.1
 base58==2.1.1
 fastapi==0.106.0
-hypha-rpc==0.20.37
+hypha-rpc==0.20.38
 jinja2==3.1.4
 lxml==4.9.3
 msgpack==1.0.8
diff --git a/setup.py b/setup.py
index cc251845..e0e48fb6 100644
--- a/setup.py
+++ b/setup.py
@@ -12,7 +12,7 @@
 REQUIREMENTS = [
     "aiofiles",
     "fastapi>=0.70.0,<=0.106.0",
-    "hypha-rpc>=0.20.37",
+    "hypha-rpc>=0.20.38",
     "msgpack>=1.0.2",
     "numpy",
     "pydantic[email]>=2.6.1",