feat: persist current workspace in user sessions

2026-03-31 15:17:07 +03:00
parent 84f4437ce0
commit 86b61523bd
15 changed files with 331 additions and 57 deletions
@@ -15,7 +15,10 @@ use crate::{
            create_agent, delete_agent, get_agent, get_agent_version, list_agents, publish_agent,
            save_agent_bindings, update_agent,
        },
-        auth::{change_password, get_profile, get_session, login, logout, update_profile},
+        auth::{
+            change_password, get_profile, get_session, login, logout, update_current_workspace,
+            update_profile,
+        },
        auth_profiles::{create_auth_profile, get_auth_profile, list_auth_profiles},
        observability::{get_agent_usage, get_log, get_operation_usage, get_usage, list_logs},
        operations::{
@@ -149,6 +152,7 @@ pub fn build_app(state: AppState) -> Router {
        .route("/logout", post(logout))
        .route("/session", get(get_session))
        .route("/profile", get(get_profile).patch(update_profile))
+        .route("/current-workspace", post(update_current_workspace))
        .route("/password", post(change_password))
        .layer(middleware::from_fn_with_state(
            state.clone(),
@@ -991,6 +995,57 @@ mod tests {
        assert_eq!(relogin_status, reqwest::StatusCode::OK);
    }

+    #[tokio::test(flavor = "multi_thread")]
+    #[serial]
+    async fn switches_current_workspace_in_session() {
+        let registry = test_registry().await;
+        let storage_root = test_storage_root("session_workspace");
+        let base_url = spawn_admin_api(build_test_app(registry, storage_root)).await;
+        let root_url = base_url
+            .as_ref()
+            .split("/api/admin/workspaces/")
+            .next()
+            .unwrap()
+            .to_owned();
+        let client = authorized_client(&base_url).await;
+
+        let created_workspace = assert_success_json(
+            client
+                .post(format!("{root_url}/api/admin/workspaces"))
+                .json(&json!({
+                    "slug": "growth-lab",
+                    "display_name": "Growth Lab",
+                    "settings": {}
+                }))
+                .send()
+                .await
+                .unwrap(),
+        )
+        .await;
+        let workspace_id = created_workspace["workspace"]["id"].as_str().unwrap();
+
+        let switched = assert_success_json(
+            client
+                .post(format!("{root_url}/api/auth/current-workspace"))
+                .json(&json!({ "workspace_id": workspace_id }))
+                .send()
+                .await
+                .unwrap(),
+        )
+        .await;
+        assert_eq!(switched["current_workspace_id"], workspace_id);
+
+        let session = assert_success_json(
+            client
+                .get(format!("{root_url}/api/auth/session"))
+                .send()
+                .await
+                .unwrap(),
+        )
+        .await;
+        assert_eq!(session["current_workspace_id"], workspace_id);
+    }
+
    #[tokio::test(flavor = "multi_thread")]
    #[serial]
    async fn exposes_logs_and_usage_from_real_test_runs() {
@@ -38,8 +38,10 @@ pub struct AuthSettings {

 #[derive(Clone, Debug, Serialize)]
 pub struct AuthenticatedSession {
+    pub session_id: UserSessionId,
    pub user: User,
    pub memberships: Vec<WorkspaceMembershipRecord>,
+    pub current_workspace_id: Option<WorkspaceId>,
 }

 #[derive(Clone)]
@@ -5,7 +5,9 @@ use serde_json::json;
 use crate::{
    auth::{AuthenticatedSession, cleared_session_cookie, extract_session_token, session_cookie},
    error::ApiError,
-    service::{ChangePasswordPayload, LoginPayload, UpdateProfilePayload},
+    service::{
+        ChangePasswordPayload, LoginPayload, UpdateCurrentWorkspacePayload, UpdateProfilePayload,
+    },
    state::AppState,
 };

@@ -55,9 +57,11 @@ pub async fn get_session(
 pub async fn get_profile(
    Extension(session): Extension<AuthenticatedSession>,
 ) -> Result<Json<serde_json::Value>, ApiError> {
-    Ok(Json(
-        json!({ "user": session.user, "memberships": session.memberships }),
-    ))
+    Ok(Json(json!({
+        "user": session.user,
+        "memberships": session.memberships,
+        "current_workspace_id": session.current_workspace_id
+    })))
 }

 pub async fn update_profile(
@@ -67,7 +71,27 @@ pub async fn update_profile(
 ) -> Result<Json<serde_json::Value>, ApiError> {
    let updated = state
        .service
-        .update_profile(&session.user.id, payload)
+        .update_profile(
+            &session.user.id,
+            session.current_workspace_id.as_ref(),
+            payload,
+        )
+        .await?;
+    Ok(Json(json!(updated)))
+}
+
+pub async fn update_current_workspace(
+    State(state): State<AppState>,
+    Extension(session): Extension<AuthenticatedSession>,
+    Json(payload): Json<UpdateCurrentWorkspacePayload>,
+) -> Result<Json<serde_json::Value>, ApiError> {
+    let updated = state
+        .service
+        .set_current_workspace(
+            &session.session_id,
+            &session.user.id,
+            &payload.workspace_id.as_str().into(),
+        )
        .await?;
    Ok(Json(json!(updated)))
 }
@@ -64,6 +64,7 @@ pub struct LoginPayload {
 pub struct SessionResponse {
    pub user: crank_core::User,
    pub memberships: Vec<WorkspaceMembershipRecord>,
+    pub current_workspace_id: Option<String>,
 }

 #[derive(Clone, Debug, Deserialize)]
@@ -78,6 +79,11 @@ pub struct ChangePasswordPayload {
    pub new_password: String,
 }

+#[derive(Clone, Debug, Deserialize)]
+pub struct UpdateCurrentWorkspacePayload {
+    pub workspace_id: String,
+}
+
 #[derive(Clone, Debug, Deserialize, Serialize)]
 pub struct OperationPayload {
    pub name: String,
@@ -546,8 +552,10 @@ impl AdminService {
            .get_user_session(session_id, &secret_hash)
            .await?
            .map(|record| AuthenticatedSession {
+                session_id: record.session_id,
                user: record.user,
                memberships: record.memberships,
+                current_workspace_id: record.current_workspace_id,
            });

        Ok(session)
@@ -582,24 +590,31 @@ impl AdminService {
            &session_cookie.value,
            &self.auth_settings.session_secret,
        );
+        let memberships = self
+            .registry
+            .list_workspaces_for_user(&user.user.id)
+            .await?;
+        let current_workspace_id = memberships
+            .first()
+            .map(|membership| membership.workspace.id.as_str().to_owned());
        self.registry
            .create_user_session(
                &session_cookie.session_id,
                &user.user.id,
+                memberships
+                    .first()
+                    .map(|membership| &membership.workspace.id),
                &secret_hash,
                &session_cookie.expires_at,
            )
            .await?;
-        let memberships = self
-            .registry
-            .list_workspaces_for_user(&user.user.id)
-            .await?;

        Ok((
            session_cookie,
            SessionResponse {
                user: user.user,
                memberships,
+                current_workspace_id,
            },
        ))
    }
@@ -672,12 +687,16 @@ impl AdminService {
            .map(|session| SessionResponse {
                user: session.user,
                memberships: session.memberships,
+                current_workspace_id: session
+                    .current_workspace_id
+                    .map(|id| id.as_str().to_owned()),
            }))
    }

    pub async fn update_profile(
        &self,
        user_id: &crank_core::UserId,
+        current_workspace_id: Option<&WorkspaceId>,
        payload: UpdateProfilePayload,
    ) -> Result<SessionResponse, ApiError> {
        let display_name = payload.display_name.trim();
@@ -696,7 +715,43 @@ impl AdminService {
            .await?;
        let memberships = self.registry.list_workspaces_for_user(user_id).await?;

-        Ok(SessionResponse { user, memberships })
+        Ok(SessionResponse {
+            user,
+            memberships,
+            current_workspace_id: current_workspace_id.map(|id| id.as_str().to_owned()),
+        })
+    }
+
+    pub async fn set_current_workspace(
+        &self,
+        session_id: &UserSessionId,
+        user_id: &crank_core::UserId,
+        workspace_id: &WorkspaceId,
+    ) -> Result<SessionResponse, ApiError> {
+        if !self
+            .user_has_workspace_access(user_id, workspace_id)
+            .await?
+        {
+            return Err(ApiError::forbidden("workspace access denied"));
+        }
+
+        self.registry
+            .set_user_session_current_workspace(session_id, workspace_id)
+            .await?;
+
+        let user = self
+            .registry
+            .get_auth_user_by_id(user_id)
+            .await?
+            .ok_or_else(|| ApiError::not_found(format!("user {} was not found", user_id.as_str())))?
+            .user;
+        let memberships = self.registry.list_workspaces_for_user(user_id).await?;
+
+        Ok(SessionResponse {
+            user,
+            memberships,
+            current_workspace_id: Some(workspace_id.as_str().to_owned()),
+        })
    }

    pub async fn change_password(
@@ -162,6 +162,13 @@
        body: JSON.stringify(payload),
      });
    },
+    setCurrentWorkspace: function(workspaceId) {
+      return request(AUTH_BASE + '/current-workspace', {
+        method: 'POST',
+        headers: headers({ 'Content-Type': 'application/json' }),
+        body: JSON.stringify({ workspace_id: workspaceId }),
+      });
+    },
    listWorkspaces: function() {
      return get('/workspaces');
    },
@@ -24,9 +24,29 @@
  }

  function primaryMembership(session) {
-    return session && session.memberships && session.memberships.length
-      ? session.memberships[0]
-      : null;
+    if (!(session && session.memberships && session.memberships.length)) {
+      return null;
+    }
+    if (session.current_workspace_id) {
+      var currentMembership = session.memberships.find(function(item) {
+        return item.workspace && item.workspace.id === session.current_workspace_id;
+      });
+      if (currentMembership) {
+        return currentMembership;
+      }
+    }
+    if (window.getCurrentWorkspace) {
+      var currentWorkspace = window.getCurrentWorkspace();
+      if (currentWorkspace) {
+        var matched = session.memberships.find(function(item) {
+          return item.workspace && item.workspace.id === currentWorkspace.id;
+        });
+        if (matched) {
+          return matched;
+        }
+      }
+    }
+    return session.memberships[0];
  }

  function initials(displayName, email) {
@@ -138,6 +158,7 @@
  window.CrankAuth = {
    fetchSession: fetchSession,
    replaceSession: replaceSession,
+    getCachedSession: function() { return sessionCache; },
    guardProtectedPage: guardProtectedPage,
    guardLoginPage: guardLoginPage,
    login: login,
@@ -249,7 +249,7 @@ async function loadWorkspaceSettings() {
          var workspaces = await window.refreshWorkspaces();
          var updated = workspaces.find(function (entry) { return entry.id === item.id; });
          if (updated && window.setCurrentWorkspace) {
-            window.setCurrentWorkspace(updated);
+            await window.setCurrentWorkspace(updated);
          }
        }

@@ -324,7 +324,7 @@ async function submitForm() {
      settings: payload.settings,
    };

-    window.setCurrentWorkspace(mapped);
+    await window.setCurrentWorkspace(mapped);
    workspaceFormState.workspaceId = workspace.id;
    workspaceFormState.workspaceRecord = { workspace: workspace };
    setFormDirty(false);
@@ -4,6 +4,7 @@ var WS_LIST = [

 var workspaceLoadPromise = null;
 var workspaceLoadFailed = false;
+var lastWorkspaceId = null;

 function workspaceColor(index) {
  return ['#0d9488', '#7c3aed', '#0891b2', '#f59e0b', '#2563eb'][index % 5];
@@ -13,11 +14,14 @@ function mapWorkspace(record, index) {
  var workspace = record && record.workspace ? record.workspace : record;
  var displayName = workspace.display_name || workspace.slug || workspace.id;
  var settings = workspace.settings || {};
+  var role = record && record.role
+    ? String(record.role).replace(/^\w/, function(char) { return char.toUpperCase(); })
+    : 'Owner';
  return {
    id: workspace.id,
    slug: workspace.slug,
    name: displayName,
-    role: 'Owner',
+    role: role,
    letter: displayName.charAt(0).toUpperCase(),
    color: settings.color || workspaceColor(index),
    status: workspace.status,
@@ -25,29 +29,60 @@ function mapWorkspace(record, index) {
  };
 }

-function selectedWorkspaceId() {
-  return localStorage.getItem('crank_workspace_id');
+function cachedWorkspaceId() {
+  try {
+    return localStorage.getItem('crank_workspace_id');
+  } catch (_error) {
+    return null;
+  }
 }

-function selectedWorkspaceSlug() {
-  return localStorage.getItem('crank_workspace_slug');
+function cachedWorkspaceSlug() {
+  try {
+    return localStorage.getItem('crank_workspace_slug');
+  } catch (_error) {
+    return null;
+  }
 }

-function persistCurrentWorkspace(workspace) {
+function sessionWorkspaceId() {
+  if (!window.CrankAuth || typeof window.CrankAuth.getCachedSession !== 'function') {
+    return null;
+  }
+  var session = window.CrankAuth.getCachedSession();
+  return session && session.current_workspace_id ? session.current_workspace_id : null;
+}
+
+function cacheCurrentWorkspace(workspace) {
  if (!workspace) return;
-  localStorage.setItem('crank_workspace_id', workspace.id);
-  localStorage.setItem('crank_workspace_slug', workspace.slug);
+  try {
+    localStorage.setItem('crank_workspace_id', workspace.id);
+    localStorage.setItem('crank_workspace_slug', workspace.slug);
+  } catch (_error) {}
+}
+
+function resolveCurrentWorkspace() {
+  var workspaceId = sessionWorkspaceId() || cachedWorkspaceId();
+  var workspaceSlug = cachedWorkspaceSlug();
+  return WS_LIST.find(function(item) { return item.id === workspaceId; })
+    || WS_LIST.find(function(item) { return item.slug === workspaceSlug; })
+    || WS_LIST[0]
+    || null;
+}
+
+function emitWorkspaceChange(workspace) {
+  if (!workspace) {
+    return;
+  }
+  lastWorkspaceId = workspace.id;
+  window.dispatchEvent(new CustomEvent('crank:workspacechange', { detail: workspace }));
 }

 function getCurrentWs() {
-  var workspaceId = selectedWorkspaceId();
-  var workspaceSlug = selectedWorkspaceSlug();
-  var workspace = WS_LIST.find(function(item) { return item.id === workspaceId; })
-    || WS_LIST.find(function(item) { return item.slug === workspaceSlug; })
-    || WS_LIST[0];
+  var workspace = resolveCurrentWorkspace();

  if (workspace) {
-    persistCurrentWorkspace(workspace);
+    cacheCurrentWorkspace(workspace);
  }

  return workspace;
@@ -134,17 +169,39 @@ function toggleWsSwitcher(e) {
 function switchWorkspace(workspaceId) {
  var workspace = WS_LIST.find(function(item) { return item.id === workspaceId; });
  if (!workspace) return;
-  persistCurrentWorkspace(workspace);
+
  var dd = document.getElementById('ws-dropdown');
  if (dd) dd.style.display = 'none';
-  renderWorkspaceList();
-  window.dispatchEvent(new CustomEvent('crank:workspacechange', { detail: workspace }));
+
+  if (!window.CrankApi || !window.CrankAuth || typeof window.CrankAuth.replaceSession !== 'function') {
+    cacheCurrentWorkspace(workspace);
+    renderWorkspaceList();
+    emitWorkspaceChange(workspace);
+    return;
+  }
+
+  return window.CrankApi
+    .setCurrentWorkspace(workspace.id)
+    .then(function(session) {
+      window.CrankAuth.replaceSession(session);
+      cacheCurrentWorkspace(workspace);
+      renderWorkspaceList();
+      emitWorkspaceChange(workspace);
+      return workspace;
+    })
+    .catch(function(error) {
+      if (window.CrankUi) {
+        window.CrankUi.error(error.message || 'Failed to switch workspace', 'Workspace switch failed');
+      }
+      throw error;
+    });
 }

 function setCurrentWorkspace(workspace) {
-  persistCurrentWorkspace(workspace);
-  renderWorkspaceList();
-  window.dispatchEvent(new CustomEvent('crank:workspacechange', { detail: workspace }));
+  if (!workspace) {
+    return Promise.resolve(null);
+  }
+  return Promise.resolve(switchWorkspace(workspace.id));
 }

 window.getCurrentWorkspace = getCurrentWs;
@@ -155,6 +212,14 @@ window.hasWorkspaceApiFailure = function() { return workspaceLoadFailed; };

 document.addEventListener('DOMContentLoaded', initWorkspaceSwitcher);

+window.addEventListener('crank:sessionchange', function() {
+  var workspace = getCurrentWs();
+  renderWorkspaceList();
+  if (workspace && workspace.id !== lastWorkspaceId) {
+    emitWorkspaceChange(workspace);
+  }
+});
+
 document.addEventListener('click', function(e) {
  if (!e.target.closest('#ws-switcher')) {
    var dd = document.getElementById('ws-dropdown');