feat: persist current workspace in user sessions

apps/admin-api/src/app.rs

@@ -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() {

apps/admin-api/src/auth.rs

@@ -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)]

apps/admin-api/src/routes/auth.rs

@@ -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)))
 }

apps/admin-api/src/service.rs

@@ -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(