fix workspace switcher links: current workspace & profile

GooeyAI · Dec 16, 2024 · 7ec2d44 · 7ec2d44
1 parent c28aefc
commit 7ec2d44
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 6 deletions.
diff --git a/routers/account.py b/routers/account.py
@@ -24,7 +24,7 @@
 from routers.root import explore_page, page_wrapper, get_og_url_path
 from workspaces.models import Workspace, WorkspaceInvite
 from workspaces.views import invitation_page, workspaces_page
-from workspaces.widgets import get_current_workspace
+from workspaces.widgets import get_current_workspace, SWITCH_WORKSPACE_KEY
 
 if typing.TYPE_CHECKING:
     from app_users.models import AppUser
@@ -111,9 +111,16 @@ def account_route(request: Request):
 
 @gui.route(app, "/account/profile/")
 def profile_route(request: Request):
+    is_switching_workspace = gui.session_state.get(SWITCH_WORKSPACE_KEY)
     with account_page_wrapper(request, AccountTabs.profile) as current_workspace:
         if not current_workspace.is_personal:
-            raise gui.RedirectException(get_route_path(members_route))
+            if is_switching_workspace:
+                raise gui.RedirectException(get_route_path(members_route))
+            else:
+                gui.session_state[SWITCH_WORKSPACE_KEY] = str(
+                    request.user.get_or_create_personal_workspace()[0].id
+                )
+                gui.rerun()
         profile_tab(request)
     url = get_og_url_path(request)
     return dict(

diff --git a/workspaces/widgets.py b/workspaces/widgets.py
@@ -8,19 +8,25 @@
 
 
 SESSION_SELECTED_WORKSPACE = "selected-workspace-id"
+SWITCH_WORKSPACE_KEY = "--switch-workspace"
 
 
 def global_workspace_selector(user: AppUser, session: dict):
     from daras_ai_v2.base import BasePage
-    from routers.account import members_route
+    from routers.account import members_route, profile_route
 
     try:
         del user.cached_workspaces  # invalidate cache on every re-render
     except AttributeError:
         pass
     workspaces = user.cached_workspaces
 
-    if switch_workspace_id := gui.session_state.pop("--switch-workspace", None):
+    if switch_workspace_id := gui.session_state.pop(SWITCH_WORKSPACE_KEY, None):
+        try:
+            if str(session[SESSION_SELECTED_WORKSPACE]) == switch_workspace_id:
+                raise gui.RedirectException(get_route_path(members_route))
+        except KeyError:
+            pass
         set_current_workspace(session, int(switch_workspace_id))
 
     try:
@@ -59,7 +65,7 @@ def global_workspace_selector(user: AppUser, session: dict):
             with gui.tag(
                 "button",
                 className="bg-transparent border-0 text-start bg-hover-light px-3 my-1",
-                name="--switch-workspace",
+                name=SWITCH_WORKSPACE_KEY,
                 type="submit",
                 value=str(workspace.id),
                 style=dict(height=row_height),
@@ -119,7 +125,7 @@ def global_workspace_selector(user: AppUser, session: dict):
         gui.html('<hr class="my-1"/>')
 
         with gui.link(
-            to="/account/profile/",
+            to=get_route_path(profile_route),
             className="text-decoration-none d-block bg-hover-light align-items-center px-3 my-1 py-1",
             style=dict(height=row_height),
         ):