feat: Add user team selection per session

Author: LubomirGeorgiev

src/actions/session.action.ts

@@ -0,0 +1,68 @@
+"use server";
+
+import { z } from "zod";
+import { ZSAError, createServerAction } from "zsa";
+import { getSessionFromCookie } from "@/utils/auth";
+import { updateKVSessionSelectedTeam } from "@/utils/kv-session";
+
+const updateSelectedTeamSchema = z.object({
+  selectedTeam: z.string().optional(),
+});
+
+/**
+ * Update the selected team for the current user's session
+ */
+export const updateSelectedTeamAction = createServerAction()
+  .input(updateSelectedTeamSchema)
+  .handler(async ({ input }) => {
+    try {
+      const session = await getSessionFromCookie();
+
+      if (!session) {
+        throw new ZSAError(
+          "FORBIDDEN",
+          "You must be logged in to update your selected team"
+        );
+      }
+
+      // Validate that the selected team exists in the user's teams (if provided)
+      if (input.selectedTeam && session.teams) {
+        const teamExists = session.teams.some(team => team.id === input.selectedTeam);
+        if (!teamExists) {
+          throw new ZSAError(
+            "FORBIDDEN",
+            "Team not found or you are not a member"
+          );
+        }
+      }
+
+      const updatedSession = await updateKVSessionSelectedTeam(
+        session.id,
+        session.userId,
+        input.selectedTeam
+      );
+
+      if (!updatedSession) {
+        throw new ZSAError(
+          "INTERNAL_SERVER_ERROR",
+          "Failed to update selected team"
+        );
+      }
+
+      return {
+        success: true,
+        selectedTeam: updatedSession.selectedTeam
+      };
+    } catch (error) {
+      console.error("Failed to update selected team:", error);
+
+      if (error instanceof ZSAError) {
+        throw error;
+      }
+
+      throw new ZSAError(
+        "INTERNAL_SERVER_ERROR",
+        "Failed to update selected team"
+      );
+    }
+  });

src/app/(dashboard)/dashboard/marketplace/components-catalog.tsx

@@ -5,6 +5,7 @@ import SeparatorWithText from "@/components/separator-with-text"
 import { NavUser } from "@/components/nav-user"
 import { Button } from "@/components/ui/button"
 import { PageHeader } from "@/components/page-header"
+import { ComponentProps } from "react"
 
 interface MarketplaceComponent {
   id: string
@@ -15,22 +16,20 @@ interface MarketplaceComponent {
   preview: () => React.ReactNode
 }
 
-interface Team {
-  name: string
-  iconName: string
-  plan: string
-}
+type Team = ComponentProps<typeof TeamSwitcher>['teams'][number];
 
 const demoTeams: Team[] = [
   {
+    id: "acme-inc",
     name: "Acme Inc",
-    iconName: "boxes",
-    plan: "Pro Plan",
+    logo: Boxes,
+    role: "admin",
   },
   {
+    id: "monsters-inc",
     name: "Monsters Inc",
-    iconName: "boxes",
-    plan: "Free Plan",
+    logo: Boxes,
+    role: "member",
   },
 ]
 

src/components/app-sidebar.tsx

@@ -46,9 +46,10 @@ type Data = {
     email: string
   }
   teams: {
+    id: string
     name: string
     logo: ComponentType
-    plan: string
+    role: string
   }[]
   navMain: NavMainItem[]
   projects: NavItem[]
@@ -64,11 +65,11 @@ export function AppSidebar({ ...props }: React.ComponentProps<typeof Sidebar>) {
       // Map teams from session to the format expected by TeamSwitcher
       const teamData = session.teams.map(team => {
         return {
+          id: team.id,
           name: team.name,
           // TODO Get the actual logo when we implement team avatars
           logo: Building2,
-          // Default plan - you might want to add plan data to your team structure
-          plan: team.role.name || "Member"
+          role: team.role.name || "Member",
         };
       });
 

src/components/team-switcher.tsx

@@ -1,6 +1,6 @@
 "use client"
 
-import * as React from "react"
+import { useCallback, useMemo, type ElementType } from "react"
 import { Building2, ChevronsUpDown, Plus } from "lucide-react"
 
 import {
@@ -19,28 +19,54 @@ import {
   useSidebar,
 } from "@/components/ui/sidebar"
 import Link from "next/link"
+import { useSessionStore } from "@/state/session"
+import { useServerAction } from "zsa-react"
+import { updateSelectedTeamAction } from "@/actions/session.action"
+import { toast } from "sonner"
 
 export function TeamSwitcher({
   teams,
 }: {
   teams: {
+    id: string
     name: string
-    logo: React.ElementType
-    plan: string
+    logo: ElementType
+    role: string
   }[]
 }) {
   const { isMobile, setOpenMobile } = useSidebar()
-  const [activeTeam, setActiveTeam] = React.useState<typeof teams[0] | null>(null)
+  const session = useSessionStore()
+  const selectedTeamId = session.selectedTeam()
+  const setSelectedTeam = session.setSelectedTeam
 
-  // Update activeTeam when teams change or on initial render
-  React.useEffect(() => {
-    if (teams.length > 0 && (!activeTeam || !teams.find(t => t.name === activeTeam.name))) {
-      setActiveTeam(teams[0])
-    }
-  }, [teams, activeTeam])
+  const { execute: updateSelectedTeam, isPending } = useServerAction(updateSelectedTeamAction)
+
+  // Find the active team based on selectedTeamId from session
+  const activeTeam = useMemo(() => {
+    if (!selectedTeamId) return teams[0] || null
+
+    // Find the matching team in the props by ID
+    return teams.find(t => t.id === selectedTeamId) || teams[0] || null
+  }, [selectedTeamId, teams])
 
   const LogoComponent = activeTeam?.logo || Building2
 
+  const handleTeamChange = useCallback(async (team: typeof teams[0]) => {
+    // Optimistically update local state
+    setSelectedTeam(team.id)
+
+    // Call server action to persist
+    // eslint-disable-next-line @typescript-eslint/no-unused-vars
+    const [_data, error] = await updateSelectedTeam({ selectedTeam: team.id })
+
+    if (error) {
+      console.error("Failed to update selected team:", error)
+      // Revert optimistic update
+      setSelectedTeam(selectedTeamId)
+      toast.error("Failed to update selected team")
+    }
+  }, [selectedTeamId, setSelectedTeam, updateSelectedTeam])
+
   return (
     <SidebarMenu>
       <SidebarMenuItem>
@@ -57,7 +83,7 @@ export function TeamSwitcher({
                 <span className="truncate font-semibold">
                   {activeTeam?.name || "No Team"}
                 </span>
-                <span className="truncate text-xs">{activeTeam?.plan || "Select a team"}</span>
+                <span className="truncate text-xs capitalize">{activeTeam?.role || "Member"}</span>
               </div>
               <ChevronsUpDown className="ml-auto" />
             </SidebarMenuButton>
@@ -74,9 +100,10 @@ export function TeamSwitcher({
             {teams.length > 0 ? (
               teams.map((team, index) => (
                 <DropdownMenuItem
-                  key={team.name}
-                  onClick={() => setActiveTeam(team)}
+                  key={team.id}
+                  onClick={() => handleTeamChange(team)}
                   className="gap-2 p-2"
+                  disabled={isPending}
                 >
                   <div className="flex size-6 items-center justify-center rounded-sm border">
                     <team.logo className="size-4 shrink-0" />

src/state/session.ts

@@ -17,9 +17,11 @@ interface SessionActions {
   setSession: (session: SessionValidationResult) => void;
   clearSession: () => void;
   refetchSession: () => void;
+  setSelectedTeam: (teamId: string | undefined) => void;
 
   // Team related selectors
   teams: () => KVSession['teams'] | undefined;
+  selectedTeam: () => string | undefined;
   isTeamMember: (teamId: string) => boolean;
   hasTeamRole: (teamId: string, roleId: string, isSystemRole?: boolean) => boolean;
   hasTeamPermission: (teamId: string, permission: string) => boolean;
@@ -39,9 +41,23 @@ export const useSessionStore = create(
       clearSession: () => set({ session: null, isLoading: false, lastFetched: null }),
       refetchSession: () => set({ isLoading: true }),
 
+      setSelectedTeam: (teamId: string | undefined) => {
+        const currentSession = get().session;
+        if (currentSession) {
+          set({
+            session: {
+              ...currentSession,
+              selectedTeam: teamId,
+            }
+          });
+        }
+      },
+
       // Team related selectors
       teams: () => get().session?.teams,
 
+      selectedTeam: () => get().session?.selectedTeam,
+
       isTeamMember: (teamId: string) => {
         return !!get().session?.teams?.some(team => team.id === teamId);
       },

src/utils/auth.ts

@@ -169,7 +169,8 @@ export async function createSession({
     user,
     authenticationType,
     passkeyCredentialId,
-    teams: teamsWithPermissions
+    teams: teamsWithPermissions,
+    selectedTeam: teamsWithPermissions?.length > 0 ? teamsWithPermissions?.[0]?.id : undefined
   });
 }
 

src/utils/kv-session.ts

@@ -44,6 +44,10 @@ export interface KVSession {
     };
     permissions: string[];
   }[];
+  /**
+   * The ID of the currently selected team for this session
+   */
+  selectedTeam?: string;
   /**
    *  !!!!!!!!!!!!!!!!!!!!!
    *  !!!   IMPORTANT   !!!
@@ -63,16 +67,17 @@ export interface KVSession {
  * IF YOU MAKE ANY CHANGES TO THE KVSESSION TYPE ABOVE, YOU NEED TO INCREMENT THIS VERSION.
  * THIS IS HOW WE TRACK WHEN WE NEED TO UPDATE THE SESSIONS IN THE KV STORE.
  */
-export const CURRENT_SESSION_VERSION = 3;
+export const CURRENT_SESSION_VERSION = 4;
 
 export async function getKV() {
   const { env } = getCloudflareContext();
   return env.NEXT_INC_CACHE_KV;
 }
 
-export interface CreateKVSessionParams extends Omit<KVSession, "id" | "createdAt" | "expiresAt"> {
+export interface CreateKVSessionParams extends Omit<KVSession, "id" | "createdAt" | "expiresAt" | "selectedTeam"> {
   sessionId: string;
   expiresAt: Date;
+  selectedTeam?: string;
 }
 
 export async function createKVSession({
@@ -82,7 +87,8 @@ export async function createKVSession({
   user,
   authenticationType,
   passkeyCredentialId,
-  teams
+  teams,
+  selectedTeam
 }: CreateKVSessionParams): Promise<KVSession> {
   const { cf } = getCloudflareContext();
   const headersList = await headers();
@@ -106,6 +112,7 @@ export async function createKVSession({
     authenticationType,
     passkeyCredentialId,
     teams,
+    selectedTeam,
     version: CURRENT_SESSION_VERSION
   };
 
@@ -222,6 +229,41 @@ export async function deleteKVSession(sessionId: string, userId: string): Promis
   await kv.delete(getSessionKey(userId, sessionId));
 }
 
+export async function updateKVSessionSelectedTeam(sessionId: string, userId: string, selectedTeam: string | undefined): Promise<KVSession | null> {
+  const session = await getKVSession(sessionId, userId);
+  if (!session) return null;
+
+  const updatedSession: KVSession = {
+    ...session,
+    selectedTeam,
+  };
+
+  const kv = await getKV();
+
+  if (!kv) {
+    throw new Error("Can't connect to KV store");
+  }
+
+  // Calculate the remaining TTL from the existing expiration
+  const remainingTtl = Math.floor((session.expiresAt - Date.now()) / 1000);
+
+  // Only update if the session hasn't expired
+  if (remainingTtl > 0) {
+    await kv.put(
+      getSessionKey(userId, sessionId),
+      JSON.stringify(updatedSession),
+      {
+        expirationTtl: remainingTtl
+      }
+    );
+  } else {
+    // Session has expired, return null
+    return null;
+  }
+
+  return updatedSession;
+}
+
 export async function getAllSessionIdsOfUser(userId: string) {
   const kv = await getKV();