Prepare cli auth & access token retrieval for api key hashing

src/app/(auth)/auth/cli/page.tsx

@@ -2,14 +2,10 @@ import { Suspense } from 'react'
 import { redirect } from 'next/navigation'
 import { CloudIcon, LaptopIcon, Link2Icon } from 'lucide-react'
 import { createClient } from '@/lib/clients/supabase/server'
-import { logger } from '@/lib/clients/logger'
+import { logError } from '@/lib/clients/logger'
 import { ERROR_CODES } from '@/configs/logs'
 import { AUTH_URLS, PROTECTED_URLS } from '@/configs/urls'
-import {
-  bailOutFromPPR,
-  getTeamApiKey,
-  getUserAccessToken,
-} from '@/lib/utils/server'
+import { bailOutFromPPR, generateE2BUserAccessToken } from '@/lib/utils/server'
 import { encodedRedirect } from '@/lib/utils/auth'
 import { Alert, AlertDescription, AlertTitle } from '@/ui/primitives/alert'
 import { getDefaultTeamRelation } from '@/server/auth/get-default-team'
@@ -26,22 +22,26 @@ type CLISearchParams = Promise<{
 
 // Server Actions
 
-async function handleCLIAuth(next: string, userId: string, userEmail: string) {
+async function handleCLIAuth(
+  next: string,
+  userId: string,
+  userEmail: string,
+  supabaseAccessToken: string
+) {
   if (!next?.startsWith('http://localhost')) {
     throw new Error('Invalid redirect URL')
   }
 
   try {
     const defaultTeam = await getDefaultTeamRelation(userId)
-    const [apiKey, accessToken] = await Promise.all([
-      getTeamApiKey(userId, defaultTeam.team_id),
-      getUserAccessToken(userId),
-    ])
+    const e2bAccessToken = await generateE2BUserAccessToken(
+      supabaseAccessToken,
+      userId
+    )
 
     const searchParams = new URLSearchParams({
       email: userEmail,
-      defaultTeamApiKey: apiKey,
-      accessToken,
+      accessToken: e2bAccessToken.token,
       defaultTeamId: defaultTeam.team_id,
     })
 
@@ -107,7 +107,7 @@ export default async function CLIAuthPage({
 
   // Validate redirect URL
   if (!next?.startsWith('http://localhost')) {
-    logger.error(ERROR_CODES.CLI_AUTH, 'Invalid redirect URL')
+    logError(ERROR_CODES.CLI_AUTH, 'Invalid redirect URL')
     redirect(PROTECTED_URLS.DASHBOARD)
   }
 
@@ -122,13 +122,26 @@ export default async function CLIAuthPage({
   // Handle CLI callback if authenticated
   if (!error && next && user) {
     try {
-      return await handleCLIAuth(next, user.id, user.email!)
+      const {
+        data: { session },
+      } = await supabase.auth.getSession()
+
+      if (!session?.access_token) {
+        throw new Error('No provider access token found')
+      }
+
+      return await handleCLIAuth(
+        next,
+        user.id,
+        user.email!,
+        session.access_token
+      )
     } catch (err) {
       if (err instanceof Error && err.message.includes('NEXT_REDIRECT')) {
         throw err
       }
 
-      logger.error(ERROR_CODES.CLI_AUTH, err)
+      logError(ERROR_CODES.CLI_AUTH, err)
 
       return encodedRedirect('error', '/auth/cli', (err as Error).message, {
         next,
@@ -142,7 +155,7 @@ export default async function CLIAuthPage({
       <h2 className="mt-6 text-base leading-7">
         Linking CLI with your account
       </h2>
-      <div className="mt-12 leading-8 text-fg-500">
+      <div className="text-fg-500 mt-12 leading-8">
         <Suspense fallback={<div>Loading...</div>}>
           {error ? (
             <ErrorAlert message={decodeURIComponent(error)} />

src/app/api/team/state/route.ts

@@ -0,0 +1,35 @@
+import { cookies } from 'next/headers'
+import { COOKIE_KEYS } from '@/configs/keys'
+import { z } from 'zod'
+import { ResponseCookie } from 'next/dist/compiled/@edge-runtime/cookies'
+
+const TeamStateSchema = z.object({
+  teamId: z.string(),
+  teamSlug: z.string(),
+})
+
+export async function POST(request: Request) {
+  try {
+    const body = TeamStateSchema.parse(await request.json())
+
+    const COOKIE_SETTINGS: Partial<ResponseCookie> = {
+      path: '/',
+      maxAge: 60 * 60 * 24 * 7, // 7 days
+      sameSite: 'lax',
+      secure: process.env.NODE_ENV === 'production',
+    }
+
+    const cookieStore = await cookies()
+
+    cookieStore.set(COOKIE_KEYS.SELECTED_TEAM_ID, body.teamId, COOKIE_SETTINGS)
+    cookieStore.set(
+      COOKIE_KEYS.SELECTED_TEAM_SLUG,
+      body.teamSlug,
+      COOKIE_SETTINGS
+    )
+
+    return Response.json({ success: true })
+  } catch (error) {
+    return Response.json({ error: 'Invalid request' }, { status: 400 })
+  }
+}

src/features/dashboard/account/user-access-token.tsx

@@ -3,10 +3,8 @@
 import { useState } from 'react'
 import { Button } from '@/ui/primitives/button'
 import { Input } from '@/ui/primitives/input'
-import HelpTooltip from '@/ui/help-tooltip'
 import { useToast } from '@/lib/hooks/use-toast'
 import { Eye, EyeOff } from 'lucide-react'
-import { Label } from '@/ui/primitives/label'
 import { Loader } from '@/ui/loader'
 import { getUserAccessTokenAction } from '@/server/user/user-actions'
 import CopyButton from '@/ui/copy-button'
@@ -27,7 +25,7 @@ export default function UserAccessToken({ className }: UserAccessTokenProps) {
     {
       onSuccess: (result) => {
         if (result.data) {
-          setToken(result.data.accessToken)
+          setToken(result.data.token)
           setIsVisible(true)
         }
       },

src/features/dashboard/sidebar/menu.tsx

@@ -40,16 +40,22 @@ export default function DashboardSidebarMenu({
   const [createTeamOpen, setCreateTeamOpen] = useState(false)
   const pathname = usePathname()
 
-  const handleTeamChange = (teamId: string) => {
+  const handleTeamChange = async (teamId: string) => {
     const team = teams.find((t) => t.id === teamId)
-    if (team && selectedTeam) {
-      router.push(
-        pathname
-          .replace(selectedTeam.slug, team.slug)
-          .replace(selectedTeam.id, team.id)
-      )
-      router.refresh()
-    }
+
+    if (!team || !selectedTeam) return
+
+    await fetch('/api/team/state', {
+      method: 'POST',
+      body: JSON.stringify({ teamId: team.id, teamSlug: team.slug }),
+    })
+
+    router.push(
+      pathname
+        .replace(selectedTeam.slug, team.slug)
+        .replace(selectedTeam.id, team.id)
+    )
+    router.refresh()
   }
 
   const handleLogout = () => {

src/lib/utils/server.ts

@@ -12,11 +12,13 @@ import { z } from 'zod'
 import { cookies } from 'next/headers'
 import { unstable_noStore } from 'next/cache'
 import { COOKIE_KEYS } from '@/configs/keys'
-import { logger } from '../clients/logger'
+import { logError, logger } from '../clients/logger'
 import { kv } from '@/lib/clients/kv'
 import { KV_KEYS } from '@/configs/keys'
 import { ERROR_CODES, INFO_CODES } from '@/configs/logs'
 import { getEncryptedCookie } from './cookies'
+import { SUPABASE_AUTH_HEADERS } from '@/configs/constants'
+import { CreatedAccessToken } from '@/types/api'
 
 /*
  *  This function checks if the user is authenticated and returns the user and the supabase client.
@@ -93,22 +95,35 @@ export async function getTeamApiKey(userId: string, teamId: string) {
 }
 
 /*
- *  This function fetches a user access token for a given user.
- *  If the user does not have an active access token, it throws an error.
+ *  This function generates an e2b user access token for a given user.
  */
-export async function getUserAccessToken(userId: string) {
-  const { data: userAccessTokenData, error: userAccessTokenError } =
-    await supabaseAdmin.from('access_tokens').select('*').eq('user_id', userId)
+export async function generateE2BUserAccessToken(
+  supabaseAccessToken: string,
+  userId: string
+) {
+  const TOKEN_NAME = 'e2b_generated_access_token'
 
-  if (userAccessTokenError) {
-    throw userAccessTokenError
-  }
+  const apiUrl = await getApiUrl()
+
+  const response = await fetch(`${apiUrl.url}/access-tokens`, {
+    method: 'POST',
+    headers: {
+      'Content-Type': 'application/json',
+      ...SUPABASE_AUTH_HEADERS(supabaseAccessToken),
+    },
+    body: JSON.stringify({ name: TOKEN_NAME }),
+  })
 
-  if (!userAccessTokenData || userAccessTokenData.length === 0) {
-    throw new Error(`No user access token found for user (user: ${userId})`)
+  if (!response.ok) {
+    const text = await response.text()
+    throw new Error(
+      `Failed to generate e2b user access token for user (${userId}): ${text ?? response.statusText}`
+    )
   }
 
-  return userAccessTokenData[0].access_token
+  const data: CreatedAccessToken = await response.json()
+
+  return data
 }
 
 // TODO: we should probably add some team permission system here

src/server/user/user-actions.ts

@@ -2,7 +2,7 @@
 
 import { authActionClient } from '@/lib/clients/action'
 import { supabaseAdmin } from '@/lib/clients/supabase/admin'
-import { getUserAccessToken } from '@/lib/utils/server'
+import { generateE2BUserAccessToken } from '@/lib/utils/server'
 import { z } from 'zod'
 import { headers } from 'next/headers'
 import { returnValidationErrors } from 'next-safe-action'
@@ -99,11 +99,12 @@ export const deleteAccountAction = authActionClient
 export const getUserAccessTokenAction = authActionClient
   .metadata({ actionName: 'getUserAccessToken' })
   .action(async ({ ctx }) => {
-    const { user } = ctx
+    const { user, session } = ctx
 
-    const accessToken = await getUserAccessToken(user.id)
+    const token = await generateE2BUserAccessToken(
+      session.access_token,
+      user.id
+    )
 
-    return {
-      accessToken,
-    }
+    return token
   })

src/types/api.d.ts

@@ -38,4 +38,18 @@ interface SandboxMetrics {
   timestamp: string
 }
 
-export type { Sandbox, Template, SandboxMetrics, DefaultTemplate }
+interface CreatedAccessToken {
+  id: string
+  name: string
+  token: string
+  tokenMask: string
+  createdAt: string
+}
+
+export type {
+  Sandbox,
+  Template,
+  SandboxMetrics,
+  DefaultTemplate,
+  CreatedAccessToken,
+}