feat: persist current workspace in user sessions

crates/crank-registry/src/postgres.rs

@@ -257,6 +257,7 @@ impl PostgresRegistry {
         &self,
         session_id: &UserSessionId,
         user_id: &UserId,
+        current_workspace_id: Option<&WorkspaceId>,
         secret_hash: &str,
         expires_at: &str,
     ) -> Result<(), RegistryError> {
@@ -264,6 +265,7 @@ impl PostgresRegistry {
             "insert into user_sessions (
                 id,
                 user_id,
+                current_workspace_id,
                 secret_hash,
                 status,
                 expires_at,
@@ -273,14 +275,16 @@ impl PostgresRegistry {
                 $1,
                 $2,
                 $3,
+                $4,
                 'active',
-                $4::timestamptz,
+                $5::timestamptz,
                 now(),
                 now()
              )",
         )
         .bind(session_id.as_str())
         .bind(user_id.as_str())
+        .bind(current_workspace_id.map(|id| id.as_str()))
         .bind(secret_hash)
         .bind(expires_at)
         .execute(&self.pool)
@@ -298,6 +302,7 @@ impl PostgresRegistry {
             "select
                 s.id,
                 s.user_id,
+                s.current_workspace_id,
                 u.email,
                 u.display_name,
                 u.status,
@@ -328,11 +333,26 @@ impl PostgresRegistry {
             created_at: row.try_get("created_at")?,
         };
         let memberships = self.list_workspaces_for_user(&user_id).await?;
+        let stored_workspace_id = row
+            .try_get::<Option<String>, _>("current_workspace_id")?
+            .map(WorkspaceId::new);
+        let current_workspace_id = stored_workspace_id
+            .filter(|workspace_id| {
+                memberships
+                    .iter()
+                    .any(|membership| membership.workspace.id == *workspace_id)
+            })
+            .or_else(|| {
+                memberships
+                    .first()
+                    .map(|membership| membership.workspace.id.clone())
+            });
 
         Ok(Some(SessionRecord {
             session_id: UserSessionId::new(row.try_get::<String, _>("id")?),
             user,
             memberships,
+            current_workspace_id,
         }))
     }
 
@@ -368,6 +388,24 @@ impl PostgresRegistry {
         Ok(())
     }
 
+    pub async fn set_user_session_current_workspace(
+        &self,
+        session_id: &UserSessionId,
+        workspace_id: &WorkspaceId,
+    ) -> Result<(), RegistryError> {
+        sqlx::query(
+            "update user_sessions
+             set current_workspace_id = $2
+             where id = $1",
+        )
+        .bind(session_id.as_str())
+        .bind(workspace_id.as_str())
+        .execute(&self.pool)
+        .await?;
+
+        Ok(())
+    }
+
     pub async fn user_has_workspace_access(
         &self,
         user_id: &UserId,