From ed974ca8ed9e7c032b08d8dd8c581ad7840fae43 Mon Sep 17 00:00:00 2001 From: waleed Date: Tue, 13 Jan 2026 20:02:50 -0800 Subject: [PATCH] fix(invitations): preserve tokens after error --- .../invitations/[invitationId]/route.test.ts | 206 ++++++++++++++++-- .../invitations/[invitationId]/route.ts | 18 +- apps/sim/app/invite/[id]/invite.tsx | 43 ++-- 3 files changed, 219 insertions(+), 48 deletions(-) diff --git a/apps/sim/app/api/workspaces/invitations/[invitationId]/route.test.ts b/apps/sim/app/api/workspaces/invitations/[invitationId]/route.test.ts index 12833c9695..389de676cd 100644 --- a/apps/sim/app/api/workspaces/invitations/[invitationId]/route.test.ts +++ b/apps/sim/app/api/workspaces/invitations/[invitationId]/route.test.ts @@ -2,13 +2,6 @@ import { createSession, createWorkspaceRecord, loggerMock } from '@sim/testing' import { NextRequest } from 'next/server' import { beforeEach, describe, expect, it, vi } from 'vitest' -/** - * Tests for workspace invitation by ID API route - * Tests GET (details + token acceptance), DELETE (cancellation) - * - * @vitest-environment node - */ - const mockGetSession = vi.fn() const mockHasWorkspaceAdminAccess = vi.fn() @@ -227,7 +220,7 @@ describe('Workspace Invitation [invitationId] API Route', () => { expect(response.headers.get('location')).toBe('https://test.sim.ai/workspace/workspace-456/w') }) - it('should redirect to error page when invitation expired', async () => { + it('should redirect to error page with token preserved when invitation expired', async () => { const session = createSession({ userId: mockUser.id, email: 'invited@example.com', @@ -250,12 +243,13 @@ describe('Workspace Invitation [invitationId] API Route', () => { const response = await GET(request, { params }) expect(response.status).toBe(307) - expect(response.headers.get('location')).toBe( - 'https://test.sim.ai/invite/invitation-789?error=expired' + const location = response.headers.get('location') + expect(location).toBe( + 'https://test.sim.ai/invite/invitation-789?error=expired&token=token-abc123' ) }) - it('should redirect to error page when email mismatch', async () => { + it('should redirect to error page with token preserved when email mismatch', async () => { const session = createSession({ userId: mockUser.id, email: 'wrong@example.com', @@ -277,12 +271,13 @@ describe('Workspace Invitation [invitationId] API Route', () => { const response = await GET(request, { params }) expect(response.status).toBe(307) - expect(response.headers.get('location')).toBe( - 'https://test.sim.ai/invite/invitation-789?error=email-mismatch' + const location = response.headers.get('location') + expect(location).toBe( + 'https://test.sim.ai/invite/invitation-789?error=email-mismatch&token=token-abc123' ) }) - it('should return 404 when invitation not found', async () => { + it('should return 404 when invitation not found (without token)', async () => { const session = createSession({ userId: mockUser.id, email: mockUser.email }) mockGetSession.mockResolvedValue(session) dbSelectResults = [[]] @@ -296,6 +291,189 @@ describe('Workspace Invitation [invitationId] API Route', () => { expect(response.status).toBe(404) expect(data).toEqual({ error: 'Invitation not found or has expired' }) }) + + it('should redirect to error page with token preserved when invitation not found (with token)', async () => { + const session = createSession({ userId: mockUser.id, email: mockUser.email }) + mockGetSession.mockResolvedValue(session) + dbSelectResults = [[]] + + const request = new NextRequest( + 'http://localhost/api/workspaces/invitations/non-existent?token=some-invalid-token' + ) + const params = Promise.resolve({ invitationId: 'non-existent' }) + + const response = await GET(request, { params }) + + expect(response.status).toBe(307) + const location = response.headers.get('location') + expect(location).toBe( + 'https://test.sim.ai/invite/non-existent?error=invalid-token&token=some-invalid-token' + ) + }) + + it('should redirect to error page with token preserved when invitation already processed', async () => { + const session = createSession({ + userId: mockUser.id, + email: 'invited@example.com', + name: mockUser.name, + }) + mockGetSession.mockResolvedValue(session) + + const acceptedInvitation = { + ...mockInvitation, + status: 'accepted', + } + + dbSelectResults = [[acceptedInvitation], [mockWorkspace]] + + const request = new NextRequest( + 'http://localhost/api/workspaces/invitations/token-abc123?token=token-abc123' + ) + const params = Promise.resolve({ invitationId: 'token-abc123' }) + + const response = await GET(request, { params }) + + expect(response.status).toBe(307) + const location = response.headers.get('location') + expect(location).toBe( + 'https://test.sim.ai/invite/invitation-789?error=already-processed&token=token-abc123' + ) + }) + + it('should redirect to error page with token preserved when workspace not found', async () => { + const session = createSession({ + userId: mockUser.id, + email: 'invited@example.com', + name: mockUser.name, + }) + mockGetSession.mockResolvedValue(session) + + dbSelectResults = [[mockInvitation], []] + + const request = new NextRequest( + 'http://localhost/api/workspaces/invitations/token-abc123?token=token-abc123' + ) + const params = Promise.resolve({ invitationId: 'token-abc123' }) + + const response = await GET(request, { params }) + + expect(response.status).toBe(307) + const location = response.headers.get('location') + expect(location).toBe( + 'https://test.sim.ai/invite/invitation-789?error=workspace-not-found&token=token-abc123' + ) + }) + + it('should redirect to error page with token preserved when user not found', async () => { + const session = createSession({ + userId: mockUser.id, + email: 'invited@example.com', + name: mockUser.name, + }) + mockGetSession.mockResolvedValue(session) + + dbSelectResults = [[mockInvitation], [mockWorkspace], []] + + const request = new NextRequest( + 'http://localhost/api/workspaces/invitations/token-abc123?token=token-abc123' + ) + const params = Promise.resolve({ invitationId: 'token-abc123' }) + + const response = await GET(request, { params }) + + expect(response.status).toBe(307) + const location = response.headers.get('location') + expect(location).toBe( + 'https://test.sim.ai/invite/invitation-789?error=user-not-found&token=token-abc123' + ) + }) + + it('should URL encode special characters in token when preserving in error redirects', async () => { + const session = createSession({ + userId: mockUser.id, + email: 'wrong@example.com', + name: mockUser.name, + }) + mockGetSession.mockResolvedValue(session) + + dbSelectResults = [ + [mockInvitation], + [mockWorkspace], + [{ ...mockUser, email: 'wrong@example.com' }], + ] + + const specialToken = 'token+with/special=chars&more' + const request = new NextRequest( + `http://localhost/api/workspaces/invitations/token-abc123?token=${encodeURIComponent(specialToken)}` + ) + const params = Promise.resolve({ invitationId: 'token-abc123' }) + + const response = await GET(request, { params }) + + expect(response.status).toBe(307) + const location = response.headers.get('location') + expect(location).toContain('error=email-mismatch') + expect(location).toContain(`token=${encodeURIComponent(specialToken)}`) + }) + }) + + describe('Token Preservation - Full Flow Scenario', () => { + it('should preserve token through email mismatch so user can retry with correct account', async () => { + const wrongSession = createSession({ + userId: 'wrong-user', + email: 'wrong@example.com', + name: 'Wrong User', + }) + mockGetSession.mockResolvedValue(wrongSession) + + dbSelectResults = [ + [mockInvitation], + [mockWorkspace], + [{ id: 'wrong-user', email: 'wrong@example.com' }], + ] + + const request1 = new NextRequest( + 'http://localhost/api/workspaces/invitations/token-abc123?token=token-abc123' + ) + const params1 = Promise.resolve({ invitationId: 'token-abc123' }) + + const response1 = await GET(request1, { params: params1 }) + + expect(response1.status).toBe(307) + const location1 = response1.headers.get('location') + expect(location1).toBe( + 'https://test.sim.ai/invite/invitation-789?error=email-mismatch&token=token-abc123' + ) + + vi.clearAllMocks() + dbSelectCallIndex = 0 + + const correctSession = createSession({ + userId: mockUser.id, + email: 'invited@example.com', + name: mockUser.name, + }) + mockGetSession.mockResolvedValue(correctSession) + + dbSelectResults = [ + [mockInvitation], + [mockWorkspace], + [{ ...mockUser, email: 'invited@example.com' }], + [], + ] + + const request2 = new NextRequest( + 'http://localhost/api/workspaces/invitations/token-abc123?token=token-abc123' + ) + const params2 = Promise.resolve({ invitationId: 'token-abc123' }) + + const response2 = await GET(request2, { params: params2 }) + + expect(response2.status).toBe(307) + expect(response2.headers.get('location')).toBe( + 'https://test.sim.ai/workspace/workspace-456/w' + ) + }) }) describe('DELETE /api/workspaces/invitations/[invitationId]', () => { diff --git a/apps/sim/app/api/workspaces/invitations/[invitationId]/route.ts b/apps/sim/app/api/workspaces/invitations/[invitationId]/route.ts index 4d4ac7928f..c7574a61e2 100644 --- a/apps/sim/app/api/workspaces/invitations/[invitationId]/route.ts +++ b/apps/sim/app/api/workspaces/invitations/[invitationId]/route.ts @@ -31,7 +31,6 @@ export async function GET( const isAcceptFlow = !!token // If token is provided, this is an acceptance flow if (!session?.user?.id) { - // For token-based acceptance flows, redirect to login if (isAcceptFlow) { return NextResponse.redirect(new URL(`/invite/${invitationId}?token=${token}`, getBaseUrl())) } @@ -51,8 +50,9 @@ export async function GET( if (!invitation) { if (isAcceptFlow) { + const tokenParam = token ? `&token=${encodeURIComponent(token)}` : '' return NextResponse.redirect( - new URL(`/invite/${invitationId}?error=invalid-token`, getBaseUrl()) + new URL(`/invite/${invitationId}?error=invalid-token${tokenParam}`, getBaseUrl()) ) } return NextResponse.json({ error: 'Invitation not found or has expired' }, { status: 404 }) @@ -60,8 +60,9 @@ export async function GET( if (new Date() > new Date(invitation.expiresAt)) { if (isAcceptFlow) { + const tokenParam = token ? `&token=${encodeURIComponent(token)}` : '' return NextResponse.redirect( - new URL(`/invite/${invitation.id}?error=expired`, getBaseUrl()) + new URL(`/invite/${invitation.id}?error=expired${tokenParam}`, getBaseUrl()) ) } return NextResponse.json({ error: 'Invitation has expired' }, { status: 400 }) @@ -75,17 +76,20 @@ export async function GET( if (!workspaceDetails) { if (isAcceptFlow) { + const tokenParam = token ? `&token=${encodeURIComponent(token)}` : '' return NextResponse.redirect( - new URL(`/invite/${invitation.id}?error=workspace-not-found`, getBaseUrl()) + new URL(`/invite/${invitation.id}?error=workspace-not-found${tokenParam}`, getBaseUrl()) ) } return NextResponse.json({ error: 'Workspace not found' }, { status: 404 }) } if (isAcceptFlow) { + const tokenParam = token ? `&token=${encodeURIComponent(token)}` : '' + if (invitation.status !== ('pending' as WorkspaceInvitationStatus)) { return NextResponse.redirect( - new URL(`/invite/${invitation.id}?error=already-processed`, getBaseUrl()) + new URL(`/invite/${invitation.id}?error=already-processed${tokenParam}`, getBaseUrl()) ) } @@ -100,7 +104,7 @@ export async function GET( if (!userData) { return NextResponse.redirect( - new URL(`/invite/${invitation.id}?error=user-not-found`, getBaseUrl()) + new URL(`/invite/${invitation.id}?error=user-not-found${tokenParam}`, getBaseUrl()) ) } @@ -108,7 +112,7 @@ export async function GET( if (!isValidMatch) { return NextResponse.redirect( - new URL(`/invite/${invitation.id}?error=email-mismatch`, getBaseUrl()) + new URL(`/invite/${invitation.id}?error=email-mismatch${tokenParam}`, getBaseUrl()) ) } diff --git a/apps/sim/app/invite/[id]/invite.tsx b/apps/sim/app/invite/[id]/invite.tsx index 41a6c20588..caa2659d47 100644 --- a/apps/sim/app/invite/[id]/invite.tsx +++ b/apps/sim/app/invite/[id]/invite.tsx @@ -178,22 +178,24 @@ export default function Invite() { useEffect(() => { const errorReason = searchParams.get('error') - - if (errorReason) { - setError(getInviteError(errorReason)) - setIsLoading(false) - return - } - const isNew = searchParams.get('new') === 'true' setIsNewUser(isNew) const tokenFromQuery = searchParams.get('token') - const effectiveToken = tokenFromQuery || inviteId + if (tokenFromQuery) { + setToken(tokenFromQuery) + sessionStorage.setItem('inviteToken', tokenFromQuery) + } else { + const storedToken = sessionStorage.getItem('inviteToken') + if (storedToken && storedToken !== inviteId) { + setToken(storedToken) + } + } - if (effectiveToken) { - setToken(effectiveToken) - sessionStorage.setItem('inviteToken', effectiveToken) + if (errorReason) { + setError(getInviteError(errorReason)) + setIsLoading(false) + return } }, [searchParams, inviteId]) @@ -203,7 +205,6 @@ export default function Invite() { async function fetchInvitationDetails() { setIsLoading(true) try { - // Fetch invitation details using the invitation ID from the URL path const workspaceInviteResponse = await fetch(`/api/workspaces/invitations/${inviteId}`, { method: 'GET', }) @@ -220,7 +221,6 @@ export default function Invite() { return } - // Handle workspace invitation errors with specific status codes if (!workspaceInviteResponse.ok && workspaceInviteResponse.status !== 404) { const errorCode = parseApiError(null, workspaceInviteResponse.status) const errorData = await workspaceInviteResponse.json().catch(() => ({})) @@ -229,7 +229,6 @@ export default function Invite() { error: errorData, }) - // Refine error code based on response body if available if (errorData.error) { const refinedCode = parseApiError(errorData.error, workspaceInviteResponse.status) setError(getInviteError(refinedCode)) @@ -254,13 +253,11 @@ export default function Invite() { if (data) { setInvitationType('organization') - // Check if user is already in an organization BEFORE showing the invitation const activeOrgResponse = await client.organization .getFullOrganization() .catch(() => ({ data: null })) if (activeOrgResponse?.data) { - // User is already in an organization setCurrentOrgName(activeOrgResponse.data.name) setError(getInviteError('already-in-organization')) setIsLoading(false) @@ -289,7 +286,6 @@ export default function Invite() { throw { code: 'invalid-invitation' } } } catch (orgErr: any) { - // If this is our structured error, use it directly if (orgErr.code) { throw orgErr } @@ -316,7 +312,6 @@ export default function Invite() { window.location.href = `/api/workspaces/invitations/${encodeURIComponent(inviteId)}?token=${encodeURIComponent(token || '')}` } else { try { - // Get the organizationId from invitation details const orgId = invitationDetails?.data?.organizationId if (!orgId) { @@ -325,7 +320,6 @@ export default function Invite() { return } - // Use our custom API endpoint that handles Pro usage snapshot const response = await fetch(`/api/organizations/${orgId}/invitations/${inviteId}`, { method: 'PUT', headers: { @@ -347,7 +341,6 @@ export default function Invite() { return } - // Set the organization as active await client.organization.setActive({ organizationId: orgId, }) @@ -360,7 +353,6 @@ export default function Invite() { } catch (err: any) { logger.error('Error accepting invitation:', err) - // Reset accepted state on error setAccepted(false) const errorCode = parseApiError(err) @@ -371,7 +363,9 @@ export default function Invite() { } const getCallbackUrl = () => { - return `/invite/${inviteId}${token && token !== inviteId ? `?token=${token}` : ''}` + const effectiveToken = + token || sessionStorage.getItem('inviteToken') || searchParams.get('token') + return `/invite/${inviteId}${effectiveToken && effectiveToken !== inviteId ? `?token=${effectiveToken}` : ''}` } if (!session?.user && !isPending) { @@ -435,7 +429,6 @@ export default function Invite() { if (error) { const callbackUrl = encodeURIComponent(getCallbackUrl()) - // Special handling for already in organization if (error.code === 'already-in-organization') { return ( @@ -463,7 +456,6 @@ export default function Invite() { ) } - // Handle email mismatch - user needs to sign in with a different account if (error.code === 'email-mismatch') { return ( @@ -490,7 +482,6 @@ export default function Invite() { ) } - // Handle auth-related errors - prompt user to sign in if (error.requiresAuth) { return ( @@ -518,7 +509,6 @@ export default function Invite() { ) } - // Handle retryable errors const actions: Array<{ label: string onClick: () => void @@ -550,7 +540,6 @@ export default function Invite() { ) } - // Show success only if accepted AND no error if (accepted && !error) { return (