fix(web): attach Amplitude user ID before firing registration event (#37091)

Co-authored-by: CodingOnStar <hanxujiang@dify.com>
Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>

web/app/components/__tests__/oauth-registration-analytics.spec.tsx

@@ -4,9 +4,9 @@ import { beforeEach, describe, expect, it, vi } from 'vitest'
 import { useSearchParams } from '@/next/navigation'
 import { OAuthRegistrationAnalytics } from '../oauth-registration-analytics'
 
-const { mockSendGAEvent, mockTrackEvent } = vi.hoisted(() => ({
+const { mockSendGAEvent, mockRememberRegistrationSuccess } = vi.hoisted(() => ({
   mockSendGAEvent: vi.fn(),
-  mockTrackEvent: vi.fn(),
+  mockRememberRegistrationSuccess: vi.fn(),
 }))
 
 vi.mock('@/utils/gtag', () => ({
@@ -17,8 +17,8 @@ vi.mock('@/next/navigation', () => ({
   useSearchParams: vi.fn(),
 }))
 
-vi.mock('../base/amplitude', () => ({
-  trackEvent: (...args: unknown[]) => mockTrackEvent(...args),
+vi.mock('../base/amplitude/registration-tracking', () => ({
+  rememberRegistrationSuccess: (...args: unknown[]) => mockRememberRegistrationSuccess(...args),
 }))
 
 const mockUseSearchParams = vi.mocked(useSearchParams)
@@ -48,10 +48,9 @@ describe('OAuthRegistrationAnalytics', () => {
     render(<OAuthRegistrationAnalytics />)
 
     await waitFor(() => {
-      expect(mockTrackEvent).toHaveBeenCalledWith('user_registration_success_with_utm', {
+      expect(mockRememberRegistrationSuccess).toHaveBeenCalledWith({
         method: 'oauth',
-        utm_source: 'linkedin',
-        slug: 'agent-launch',
+        utmInfo: { utm_source: 'linkedin', slug: 'agent-launch' },
       })
     })
     expect(mockSendGAEvent).toHaveBeenCalledWith('user_registration_success_with_utm', {
@@ -73,8 +72,9 @@ describe('OAuthRegistrationAnalytics', () => {
     render(<OAuthRegistrationAnalytics />)
 
     await waitFor(() => {
-      expect(mockTrackEvent).toHaveBeenCalledWith('user_registration_success', {
+      expect(mockRememberRegistrationSuccess).toHaveBeenCalledWith({
         method: 'oauth',
+        utmInfo: null,
       })
     })
     expect(mockSendGAEvent).toHaveBeenCalledWith('user_registration_success', {
@@ -87,7 +87,7 @@ describe('OAuthRegistrationAnalytics', () => {
   it('should do nothing without the oauth registration query flag', () => {
     render(<OAuthRegistrationAnalytics />)
 
-    expect(mockTrackEvent).not.toHaveBeenCalled()
+    expect(mockRememberRegistrationSuccess).not.toHaveBeenCalled()
     expect(mockSendGAEvent).not.toHaveBeenCalled()
   })
 
@@ -100,7 +100,7 @@ describe('OAuthRegistrationAnalytics', () => {
     await waitFor(() => {
       expect(replaceStateSpy).toHaveBeenCalledWith(null, '', '/signin')
     })
-    expect(mockTrackEvent).not.toHaveBeenCalled()
+    expect(mockRememberRegistrationSuccess).not.toHaveBeenCalled()
     expect(mockSendGAEvent).not.toHaveBeenCalled()
   })
 })

web/app/components/base/amplitude/__tests__/registration-tracking.spec.ts

@@ -0,0 +1,191 @@
+import {
+  flushRegistrationSuccess,
+  REGISTRATION_SUCCESS_STORAGE_KEY,
+  rememberRegistrationSuccess,
+} from '../registration-tracking'
+
+const mockTrackEvent = vi.hoisted(() => vi.fn())
+
+vi.mock('../utils', () => ({
+  trackEvent: (...args: unknown[]) => mockTrackEvent(...args),
+}))
+
+describe('registration tracking', () => {
+  beforeEach(() => {
+    vi.clearAllMocks()
+    vi.unstubAllGlobals()
+    window.sessionStorage.clear()
+  })
+
+  // Captures the registration event for a later flush instead of firing it right away.
+  describe('rememberRegistrationSuccess', () => {
+    it('should store the base event and not track immediately when there is no utm info', () => {
+      rememberRegistrationSuccess({ method: 'email' })
+
+      expect(JSON.parse(window.sessionStorage.getItem(REGISTRATION_SUCCESS_STORAGE_KEY)!)).toEqual({
+        eventName: 'user_registration_success',
+        properties: { method: 'email' },
+      })
+      expect(mockTrackEvent).not.toHaveBeenCalled()
+    })
+
+    it('should store the utm event and merge utm info into properties when utm info is present', () => {
+      rememberRegistrationSuccess({
+        method: 'oauth',
+        utmInfo: { utm_source: 'linkedin', slug: 'agent-launch' },
+      })
+
+      expect(JSON.parse(window.sessionStorage.getItem(REGISTRATION_SUCCESS_STORAGE_KEY)!)).toEqual({
+        eventName: 'user_registration_success_with_utm',
+        properties: { method: 'oauth', utm_source: 'linkedin', slug: 'agent-launch' },
+      })
+    })
+
+    it('should swallow errors when writing to sessionStorage fails', () => {
+      vi.stubGlobal('window', {
+        sessionStorage: {
+          getItem: vi.fn(() => null),
+          setItem: () => {
+            throw new Error('quota exceeded')
+          },
+          removeItem: vi.fn(),
+        },
+      })
+
+      try {
+        expect(() => rememberRegistrationSuccess({ method: 'email' })).not.toThrow()
+      }
+      finally {
+        vi.unstubAllGlobals()
+      }
+    })
+  })
+
+  // Replays the remembered event exactly once, after the user ID has been attached.
+  describe('flushRegistrationSuccess', () => {
+    it('should track the remembered event and clear it from storage', () => {
+      rememberRegistrationSuccess({ method: 'email', utmInfo: { utm_source: 'blog' } })
+
+      flushRegistrationSuccess()
+
+      expect(mockTrackEvent).toHaveBeenCalledTimes(1)
+      expect(mockTrackEvent).toHaveBeenCalledWith('user_registration_success_with_utm', {
+        method: 'email',
+        utm_source: 'blog',
+      })
+      expect(window.sessionStorage.getItem(REGISTRATION_SUCCESS_STORAGE_KEY)).toBeNull()
+    })
+
+    it('should do nothing when there is no pending event', () => {
+      flushRegistrationSuccess()
+
+      expect(mockTrackEvent).not.toHaveBeenCalled()
+    })
+
+    it('should fire the event at most once across repeated flushes', () => {
+      rememberRegistrationSuccess({ method: 'oauth' })
+
+      flushRegistrationSuccess()
+      flushRegistrationSuccess()
+
+      expect(mockTrackEvent).toHaveBeenCalledTimes(1)
+    })
+
+    it('should clear malformed pending data without tracking', () => {
+      window.sessionStorage.setItem(REGISTRATION_SUCCESS_STORAGE_KEY, '{not-json')
+
+      flushRegistrationSuccess()
+
+      expect(mockTrackEvent).not.toHaveBeenCalled()
+      expect(window.sessionStorage.getItem(REGISTRATION_SUCCESS_STORAGE_KEY)).toBeNull()
+    })
+
+    it('should clear the pending entry without tracking when it has no event name', () => {
+      window.sessionStorage.setItem(
+        REGISTRATION_SUCCESS_STORAGE_KEY,
+        JSON.stringify({ properties: { method: 'email' } }),
+      )
+
+      flushRegistrationSuccess()
+
+      expect(mockTrackEvent).not.toHaveBeenCalled()
+      expect(window.sessionStorage.getItem(REGISTRATION_SUCCESS_STORAGE_KEY)).toBeNull()
+    })
+
+    it('should stop without tracking when reading from sessionStorage throws', () => {
+      vi.stubGlobal('window', {
+        sessionStorage: {
+          getItem: () => {
+            throw new Error('read failed')
+          },
+          setItem: vi.fn(),
+          removeItem: vi.fn(),
+        },
+      })
+
+      try {
+        expect(() => flushRegistrationSuccess()).not.toThrow()
+        expect(mockTrackEvent).not.toHaveBeenCalled()
+      }
+      finally {
+        vi.unstubAllGlobals()
+      }
+    })
+
+    it('should still track when clearing the pending entry fails', () => {
+      const pending = { eventName: 'user_registration_success', properties: { method: 'email' } }
+      vi.stubGlobal('window', {
+        sessionStorage: {
+          getItem: () => JSON.stringify(pending),
+          setItem: vi.fn(),
+          removeItem: () => {
+            throw new Error('remove failed')
+          },
+        },
+      })
+
+      try {
+        flushRegistrationSuccess()
+
+        expect(mockTrackEvent).toHaveBeenCalledWith('user_registration_success', { method: 'email' })
+      }
+      finally {
+        vi.unstubAllGlobals()
+      }
+    })
+  })
+
+  // Both producers and the consumer must degrade gracefully when sessionStorage is
+  // missing (SSR) or blocked (privacy mode / disabled storage).
+  describe('when sessionStorage is unavailable', () => {
+    it('should no-op without throwing when window is undefined', () => {
+      vi.stubGlobal('window', undefined)
+
+      try {
+        expect(() => rememberRegistrationSuccess({ method: 'email' })).not.toThrow()
+        expect(() => flushRegistrationSuccess()).not.toThrow()
+        expect(mockTrackEvent).not.toHaveBeenCalled()
+      }
+      finally {
+        vi.unstubAllGlobals()
+      }
+    })
+
+    it('should no-op without throwing when accessing sessionStorage throws', () => {
+      vi.stubGlobal('window', {
+        get sessionStorage() {
+          throw new Error('storage disabled')
+        },
+      })
+
+      try {
+        expect(() => rememberRegistrationSuccess({ method: 'oauth' })).not.toThrow()
+        expect(() => flushRegistrationSuccess()).not.toThrow()
+        expect(mockTrackEvent).not.toHaveBeenCalled()
+      }
+      finally {
+        vi.unstubAllGlobals()
+      }
+    })
+  })
+})

web/app/components/base/amplitude/registration-tracking.ts

@@ -0,0 +1,88 @@
+import { trackEvent } from './utils'
+
+/**
+ * Storage key for a registration success event that is waiting to be sent to
+ * Amplitude until a user ID has been attached.
+ */
+export const REGISTRATION_SUCCESS_STORAGE_KEY = 'pending_registration_success_event'
+
+type RegistrationMethod = 'email' | 'oauth'
+
+type PendingRegistrationSuccessEvent = {
+  eventName: string
+  properties: Record<string, unknown>
+}
+
+const getSessionStorage = (): Storage | null => {
+  try {
+    if (typeof window === 'undefined')
+      return null
+    return window.sessionStorage
+  }
+  catch {
+    return null
+  }
+}
+
+/**
+ * Remember a registration success event so it can be sent to Amplitude *after* the
+ * user ID is attached (see `flushRegistrationSuccess`).
+ *
+ * Amplitude attributes events to whatever identity is active when `track` runs. At
+ * registration time the client does not yet know the user ID, so firing the event
+ * immediately records it under an anonymous profile. We persist the event here and
+ * replay it once `setUserId` runs in the app context provider after the redirect.
+ */
+export const rememberRegistrationSuccess = (
+  { method, utmInfo }: { method: RegistrationMethod, utmInfo?: Record<string, unknown> | null },
+) => {
+  const storage = getSessionStorage()
+  if (!storage)
+    return
+
+  const pending: PendingRegistrationSuccessEvent = {
+    eventName: utmInfo ? 'user_registration_success_with_utm' : 'user_registration_success',
+    properties: { method, ...utmInfo },
+  }
+
+  try {
+    storage.setItem(REGISTRATION_SUCCESS_STORAGE_KEY, JSON.stringify(pending))
+  }
+  catch {}
+}
+
+/**
+ * Send a previously remembered registration success event to Amplitude.
+ *
+ * MUST be called after `setUserId` so the event lands on the identified user profile.
+ * No-op when nothing is pending. The pending entry is removed before tracking so the
+ * event fires at most once even if this runs multiple times.
+ */
+export const flushRegistrationSuccess = () => {
+  const storage = getSessionStorage()
+  if (!storage)
+    return
+
+  let raw: string | null = null
+  try {
+    raw = storage.getItem(REGISTRATION_SUCCESS_STORAGE_KEY)
+  }
+  catch {
+    return
+  }
+
+  if (!raw)
+    return
+
+  try {
+    storage.removeItem(REGISTRATION_SUCCESS_STORAGE_KEY)
+  }
+  catch {}
+
+  try {
+    const pending = JSON.parse(raw) as PendingRegistrationSuccessEvent
+    if (pending?.eventName)
+      trackEvent(pending.eventName, pending.properties)
+  }
+  catch {}
+}

web/app/components/oauth-registration-analytics.tsx

@@ -4,7 +4,7 @@ import Cookies from 'js-cookie'
 import { useEffect, useRef } from 'react'
 import { useSearchParams } from '@/next/navigation'
 import { sendGAEvent } from '@/utils/gtag'
-import { trackEvent } from './base/amplitude'
+import { rememberRegistrationSuccess } from './base/amplitude/registration-tracking'
 
 const OAUTH_NEW_USER_PARAM = 'oauth_new_user'
 
@@ -48,10 +48,10 @@ export function OAuthRegistrationAnalytics() {
 
     const eventName = utmInfo ? 'user_registration_success_with_utm' : 'user_registration_success'
 
-    trackEvent(eventName, {
-      method: 'oauth',
-      ...utmInfo,
-    })
+    // Defer the Amplitude event until the user ID is attached. It is flushed in
+    // AppContextProvider after setUserId runs. Firing it here would record it under an
+    // anonymous Amplitude profile (no user ID set yet).
+    rememberRegistrationSuccess({ method: 'oauth', utmInfo })
 
     sendGAEvent(eventName, {
       method: 'oauth',

web/app/signup/set-password/__tests__/page.spec.tsx

@@ -2,6 +2,7 @@ import type { ReactElement } from 'react'
 import type { MockedFunction } from 'vitest'
 import { QueryClient, QueryClientProvider } from '@tanstack/react-query'
 import { fireEvent, render, screen, waitFor } from '@testing-library/react'
+import Cookies from 'js-cookie'
 import { beforeEach, describe, expect, it, vi } from 'vitest'
 import { useLocale } from '@/context/i18n'
 import { useRouter, useSearchParams } from '@/next/navigation'
@@ -9,6 +10,11 @@ import { useMailRegister } from '@/service/use-common'
 import { getBrowserTimezone } from '@/utils/timezone'
 import ChangePasswordForm from '../page'
 
+const { mockRememberRegistrationSuccess, mockSendGAEvent } = vi.hoisted(() => ({
+  mockRememberRegistrationSuccess: vi.fn(),
+  mockSendGAEvent: vi.fn(),
+}))
+
 vi.mock('@/context/i18n', () => ({
   useLocale: vi.fn(),
 }))
@@ -27,11 +33,11 @@ vi.mock('@/utils/timezone', () => ({
 }))
 
 vi.mock('@/utils/gtag', () => ({
-  sendGAEvent: vi.fn(),
+  sendGAEvent: (...args: unknown[]) => mockSendGAEvent(...args),
 }))
 
-vi.mock('@/app/components/base/amplitude', () => ({
-  trackEvent: vi.fn(),
+vi.mock('@/app/components/base/amplitude/registration-tracking', () => ({
+  rememberRegistrationSuccess: (...args: unknown[]) => mockRememberRegistrationSuccess(...args),
 }))
 
 vi.mock('@/utils/create-app-tracking', () => ({
@@ -64,6 +70,7 @@ const renderWithQueryClient = (ui: ReactElement) => {
 describe('Signup Set Password Page', () => {
   beforeEach(() => {
     vi.clearAllMocks()
+    Cookies.remove('utm_info')
     mockUseLocale.mockReturnValue('zh-Hans')
     mockUseSearchParams.mockReturnValue(new URLSearchParams('token=register-token') as unknown as ReturnType<typeof useSearchParams>)
     mockUseRouter.mockReturnValue({ replace: mockReplace } as unknown as ReturnType<typeof useRouter>)
@@ -98,4 +105,56 @@ describe('Signup Set Password Page', () => {
       })
     })
   })
+
+  // On successful registration the Amplitude event is deferred (remembered) so it can
+  // fire after the user ID is attached, while the GA event still fires immediately.
+  describe('Registration success tracking', () => {
+    const fillAndSubmit = () => {
+      fireEvent.change(screen.getByLabelText('common.account.newPassword'), {
+        target: { value: 'ValidPass123!' },
+      })
+      fireEvent.change(screen.getByLabelText('common.account.confirmPassword'), {
+        target: { value: 'ValidPass123!' },
+      })
+      fireEvent.click(screen.getByRole('button', { name: 'login.changePasswordBtn' }))
+    }
+
+    it('should defer the amplitude event and fire GA immediately when registration succeeds', async () => {
+      mockRegister.mockResolvedValue({ result: 'success', data: {} })
+
+      renderWithQueryClient(<ChangePasswordForm />)
+      fillAndSubmit()
+
+      await waitFor(() => {
+        expect(mockRememberRegistrationSuccess).toHaveBeenCalledWith({
+          method: 'email',
+          utmInfo: null,
+        })
+      })
+      expect(mockSendGAEvent).toHaveBeenCalledWith('user_registration_success', {
+        method: 'email',
+      })
+      expect(mockReplace).toHaveBeenCalledWith('/')
+    })
+
+    it('should remember the utm event and clear the utm cookie when a utm_info cookie is present', async () => {
+      Cookies.set('utm_info', JSON.stringify({ utm_source: 'twitter' }))
+      mockRegister.mockResolvedValue({ result: 'success', data: {} })
+
+      renderWithQueryClient(<ChangePasswordForm />)
+      fillAndSubmit()
+
+      await waitFor(() => {
+        expect(mockRememberRegistrationSuccess).toHaveBeenCalledWith({
+          method: 'email',
+          utmInfo: { utm_source: 'twitter' },
+        })
+      })
+      expect(mockSendGAEvent).toHaveBeenCalledWith('user_registration_success_with_utm', {
+        method: 'email',
+        utm_source: 'twitter',
+      })
+      expect(Cookies.get('utm_info')).toBeUndefined()
+    })
+  })
 })

web/app/signup/set-password/page.tsx

@@ -7,7 +7,7 @@ import { useQueryClient } from '@tanstack/react-query'
 import Cookies from 'js-cookie'
 import { useCallback, useState } from 'react'
 import { useTranslation } from 'react-i18next'
-import { trackEvent } from '@/app/components/base/amplitude'
+import { rememberRegistrationSuccess } from '@/app/components/base/amplitude/registration-tracking'
 import Input from '@/app/components/base/input'
 import { validPassword } from '@/config'
 import { useLocale } from '@/context/i18n'
@@ -78,10 +78,10 @@ const ChangePasswordForm = () => {
       if (result === 'success') {
         const utmInfo = parseUtmInfo()
         rememberCreateAppExternalAttribution({ utmInfo })
-        trackEvent(utmInfo ? 'user_registration_success_with_utm' : 'user_registration_success', {
-          method: 'email',
-          ...utmInfo,
-        })
+        // Defer the Amplitude event until the user ID is attached. It is flushed in
+        // AppContextProvider after setUserId runs once the redirect lands on /apps.
+        // Firing it here would record it under an anonymous Amplitude profile.
+        rememberRegistrationSuccess({ method: 'email', utmInfo })
 
         sendGAEvent(utmInfo ? 'user_registration_success_with_utm' : 'user_registration_success', {
           method: 'email',

web/context/app-context-provider.tsx

@@ -7,6 +7,7 @@ import type { ICurrentWorkspace, LangGeniusVersionResponse } from '@/models/comm
 import { useQuery, useQueryClient, useSuspenseQuery } from '@tanstack/react-query'
 import { useCallback, useEffect, useMemo } from 'react'
 import { setUserId, setUserProperties } from '@/app/components/base/amplitude'
+import { flushRegistrationSuccess } from '@/app/components/base/amplitude/registration-tracking'
 import { setZendeskConversationFields } from '@/app/components/base/zendesk/utils'
 import MaintenanceNotice from '@/app/components/header/maintenance-notice'
 import { ZENDESK_FIELD_IDS } from '@/config'
@@ -160,6 +161,11 @@ export const AppContextProvider: FC<AppContextProviderProps> = ({ children }) =>
       }
 
       setUserProperties(properties)
+
+      // The user ID is now attached, so replay any registration success event captured
+      // at signup time. This makes it land on the identified Amplitude profile instead
+      // of an anonymous one (no-op when nothing was deferred).
+      flushRegistrationSuccess()
     }
   }, [userProfile, currentWorkspace])