From cbb22550242972aee645263a6cb8ead91d280329 Mon Sep 17 00:00:00 2001
From: Usman Baig <mubaig@pm.me>
Date: Thu, 22 Jan 2026 01:36:56 +0100
Subject: [PATCH] fix: handle access_token only response from switchContext

---
 app/actions/auth.ts              | 23 ++++++++++++++---------
 app/layout-content.tsx           |  4 ++--
 components/WorkspaceSwitcher.tsx |  5 +++--
 lib/api/organization.ts          |  4 ++--
 lib/auth/context.tsx             |  4 ++--
 5 files changed, 23 insertions(+), 17 deletions(-)

diff --git a/app/actions/auth.ts b/app/actions/auth.ts
index fd3ab9b..e93d239 100644
--- a/app/actions/auth.ts
+++ b/app/actions/auth.ts
@@ -97,9 +97,11 @@ export async function exchangeAuthCode(code: string, codeVerifier: string, redir
   }
 }
 
-export async function setSessionAction(accessToken: string, refreshToken: string) {
+export async function setSessionAction(accessToken: string, refreshToken?: string) {
     try {
         console.log('[setSessionAction] Decoding token...')
+        if (!accessToken) throw new Error('Access token is missing')
+        
         const payloadPart = accessToken.split('.')[1]
         const payload: UserPayload = JSON.parse(Buffer.from(payloadPart, 'base64').toString())
         
@@ -119,14 +121,17 @@ export async function setSessionAction(accessToken: string, refreshToken: string
             maxAge: 60 * 15
         })
 
-        cookieStore.set('refresh_token', refreshToken, {
-            httpOnly: true,
-            secure: process.env.NODE_ENV === 'production',
-            sameSite: 'lax',
-            path: '/',
-            domain: cookieDomain,
-            maxAge: 60 * 60 * 24 * 30
-        })
+        // * Only update refresh token if provided
+        if (refreshToken) {
+            cookieStore.set('refresh_token', refreshToken, {
+                httpOnly: true,
+                secure: process.env.NODE_ENV === 'production',
+                sameSite: 'lax',
+                path: '/',
+                domain: cookieDomain,
+                maxAge: 60 * 60 * 24 * 30
+            })
+        }
         
         console.log('[setSessionAction] Cookies set successfully')
 
diff --git a/app/layout-content.tsx b/app/layout-content.tsx
index 36c00fd..7e9c372 100644
--- a/app/layout-content.tsx
+++ b/app/layout-content.tsx
@@ -25,8 +25,8 @@ export default function LayoutContent({ children }: { children: React.ReactNode
 
   const handleSwitchWorkspace = async (orgId: string) => {
     try {
-      const { token, refresh_token } = await switchContext(orgId)
-      await setSessionAction(token, refresh_token)
+      const { access_token } = await switchContext(orgId)
+      await setSessionAction(access_token)
       window.location.reload()
     } catch (err) {
       console.error('Failed to switch workspace', err)
diff --git a/components/WorkspaceSwitcher.tsx b/components/WorkspaceSwitcher.tsx
index 286dbba..29c2466 100644
--- a/components/WorkspaceSwitcher.tsx
+++ b/components/WorkspaceSwitcher.tsx
@@ -28,10 +28,11 @@ export default function WorkspaceSwitcher({ orgs, activeOrgId }: { orgs: Organiz
           return
       }
 
-      const { token, refresh_token } = await switchContext(orgId)
+      const { access_token } = await switchContext(orgId)
       
       // * Update session cookie via server action
-      await setSessionAction(token, refresh_token)
+      // * Note: switchContext only returns access_token, we keep existing refresh token
+      await setSessionAction(access_token)
       
       // Force reload to pick up new permissions
       window.location.reload() 
diff --git a/lib/api/organization.ts b/lib/api/organization.ts
index 7d03722..126cee9 100644
--- a/lib/api/organization.ts
+++ b/lib/api/organization.ts
@@ -37,9 +37,9 @@ export async function createOrganization(name: string, slug: string): Promise<Or
 }
 
 // * Switch context to organization (returns new token)
-export async function switchContext(organizationId: string): Promise<{ token: string, refresh_token: string }> {
+export async function switchContext(organizationId: string): Promise<{ access_token: string }> {
   // * Route in main.go is /api/v1/auth/switch-context
-  return apiRequest<{ token: string, refresh_token: string }>('/auth/switch-context', {
+  return apiRequest<{ access_token: string }>('/auth/switch-context', {
     method: 'POST',
     body: JSON.stringify({ organization_id: organizationId }),
   })
diff --git a/lib/auth/context.tsx b/lib/auth/context.tsx
index 327696b..eff2360 100644
--- a/lib/auth/context.tsx
+++ b/lib/auth/context.tsx
@@ -123,10 +123,10 @@ export function AuthProvider({ children }: { children: React.ReactNode }) {
              console.log('Auto-switching to organization:', firstOrg.organization_name)
              
              try {
-                 const { token, refresh_token } = await switchContext(firstOrg.organization_id)
+                 const { access_token } = await switchContext(firstOrg.organization_id)
                  
                  // * Update session cookie
-                 const result = await setSessionAction(token, refresh_token)
+                 const result = await setSessionAction(access_token)
                  if (result.success && result.user) {
                      setUser(result.user)
                      localStorage.setItem('user', JSON.stringify(result.user))