fix(org): show impersonation session controls (#7675)

Author: waleedlatif1

apps/sim/app/o/[organizationId]/layout.test.tsx

@@ -7,12 +7,26 @@ import { authMockFns } from '@sim/testing'
 import { renderToStaticMarkup } from 'react-dom/server'
 import { beforeEach, describe, expect, it, vi } from 'vitest'
 
-const { mockGetOrganizationSurfaceContext, mockWorkspaceChrome, mockPrefetchUserProfile } =
-  vi.hoisted(() => ({
-    mockGetOrganizationSurfaceContext: vi.fn(),
-    mockWorkspaceChrome: vi.fn(({ children }: { children: ReactNode }) => children),
-    mockPrefetchUserProfile: vi.fn(async () => undefined),
-  }))
+const {
+  mockGetOrganizationSurfaceContext,
+  mockWorkspaceChrome,
+  mockPrefetchUserProfile,
+  mockUseSession,
+} = vi.hoisted(() => ({
+  mockGetOrganizationSurfaceContext: vi.fn(),
+  mockWorkspaceChrome: vi.fn(({ children }: { children: ReactNode }) => children),
+  mockPrefetchUserProfile: vi.fn(async () => undefined),
+  mockUseSession: vi.fn(),
+}))
+
+vi.mock('@/lib/auth/auth-client', () => ({ useSession: mockUseSession }))
+vi.mock('@/hooks/queries/admin-users', () => ({
+  useStopImpersonating: () => ({ mutate: vi.fn(), isPending: false }),
+}))
+vi.mock('@/stores', () => ({ clearUserData: vi.fn() }))
+vi.mock('@/lib/auth/stale-session-recovery', () => ({
+  recoverFromStaleSession: vi.fn(),
+}))
 
 vi.mock('@tanstack/react-query', () => ({
   HydrationBoundary: ({ children }: { children: ReactNode }) => children,
@@ -67,6 +81,7 @@ describe('OrganizationLayout', () => {
   beforeEach(() => {
     vi.clearAllMocks()
     mockGetSession.mockResolvedValue({ user: { id: 'viewer-1' } })
+    mockUseSession.mockReturnValue({ data: { user: { id: 'viewer-1' } }, isPending: false })
   })
 
   it('returns signed-out visitors to the organization entry after sign-in', async () => {
@@ -93,12 +108,58 @@ describe('OrganizationLayout', () => {
     expect(mockGetOrganizationSurfaceContext).toHaveBeenCalledWith('org-1', 'viewer-1')
     expect(mockPrefetchUserProfile).toHaveBeenCalledWith({}, 'viewer-1')
     expect(html).toContain('Organization child')
+    expect(html).not.toContain('Stop impersonating')
     expect(mockWorkspaceChrome).toHaveBeenCalledWith(
       expect.objectContaining({ initialSidebarCollapsed: true }),
       undefined
     )
   })
 
+  it('shows the shared impersonation banner above organization content', async () => {
+    const session = {
+      user: { id: 'viewer-1', name: 'QA Member', email: 'member@example.com' },
+      session: { impersonatedBy: 'platform-admin' },
+    }
+    mockGetSession.mockResolvedValue(session)
+    mockUseSession.mockReturnValue({ data: session, isPending: false })
+    mockGetOrganizationSurfaceContext.mockResolvedValue(SURFACE_CONTEXT)
+
+    const html = renderToStaticMarkup(
+      await OrganizationLayout({
+        children: <div>Organization child</div>,
+        params: Promise.resolve({ organizationId: 'org-1' }),
+      })
+    )
+
+    expect(mockGetOrganizationSurfaceContext).toHaveBeenCalledWith('org-1', 'viewer-1')
+    expect(html).toContain('Impersonating QA Member (member@example.com)')
+    expect(html).toContain('Stop impersonating')
+    expect(html.indexOf('Stop impersonating')).toBeLessThan(html.indexOf('Organization child'))
+  })
+
+  it('does not use the impersonating admin to enter an organization outside the rollout', async () => {
+    mockGetSession.mockResolvedValue({
+      user: { id: 'customer-member' },
+      session: { impersonatedBy: 'platform-admin' },
+    })
+    mockGetOrganizationSurfaceContext.mockResolvedValue({
+      ...SURFACE_CONTEXT,
+      searchAccess: { memberScoped: false, sourceMirrored: false },
+    })
+
+    await expect(
+      OrganizationLayout({
+        children: <div>Organization child</div>,
+        params: Promise.resolve({ organizationId: 'customer-org' }),
+      })
+    ).rejects.toThrow('redirect:/workspace?redirect=settings')
+    expect(mockGetOrganizationSurfaceContext).toHaveBeenCalledWith(
+      'customer-org',
+      'customer-member'
+    )
+    expect(mockWorkspaceChrome).not.toHaveBeenCalled()
+  })
+
   it('renders an explicit denial for a non-member without the surface', async () => {
     mockGetOrganizationSurfaceContext.mockResolvedValue(null)
 

apps/sim/app/o/[organizationId]/layout.tsx

@@ -10,6 +10,8 @@ import { buildAuthCrossLink } from '@/app/(auth)/auth-redirect'
 import { OrganizationAccessDenied } from '@/app/o/[organizationId]/components/organization-access-denied'
 import { OrganizationSidebar } from '@/app/o/[organizationId]/components/organization-sidebar'
 import { OrganizationProvider } from '@/app/o/[organizationId]/providers/organization-provider'
+import { ImpersonationBanner } from '@/app/workspace/[workspaceId]/components/impersonation-banner'
+import { SessionExpired } from '@/app/workspace/[workspaceId]/components/session-expired'
 import { WorkspaceChrome } from '@/app/workspace/[workspaceId]/components/workspace-chrome'
 import { GlobalCommandsProvider } from '@/app/workspace/[workspaceId]/providers/global-commands-provider'
 
@@ -58,6 +60,8 @@ export default async function OrganizationLayout({
       <OrganizationProvider context={context}>
         <GlobalCommandsProvider>
           <div className='workspace-root flex h-screen w-full flex-col overflow-hidden bg-[var(--surface-1)]'>
+            <ImpersonationBanner />
+            <SessionExpired />
             <WorkspaceChrome
               sidebar={<OrganizationSidebar />}
               initialSidebarCollapsed={initialSidebarCollapsed}

apps/sim/lib/navigation/organization-rollout.test.ts

@@ -0,0 +1,82 @@
+/** @vitest-environment node */
+import { resetEnvFlagsMock, setEnvFlags } from '@sim/testing'
+import { afterAll, beforeEach, describe, expect, it, vi } from 'vitest'
+import type { FeatureFlagsConfig } from '@/lib/core/config/feature-flags'
+
+const mocks = vi.hoisted(() => ({
+  appConfig: vi.fn(),
+  landing: vi.fn(),
+  platformAdmin: vi.fn(),
+}))
+
+vi.mock('@/lib/core/config/appconfig', () => ({ fetchAppConfigProfile: mocks.appConfig }))
+vi.mock('@/lib/permissions/super-user', () => ({ isPlatformAdmin: mocks.platformAdmin }))
+vi.mock('@/lib/organizations/surface', () => ({ resolveOrganizationLanding: mocks.landing }))
+vi.mock('@/lib/billing/core/subscription', () => ({
+  isOrganizationOnEnterprisePlan: vi.fn().mockResolvedValue(true),
+  getOrganizationSubscriptionUsable: vi
+    .fn()
+    .mockResolvedValue({ plan: 'enterprise', status: 'active' }),
+}))
+vi.mock('@/lib/billing/core/access', () => ({
+  isOrganizationBillingBlocked: vi.fn().mockResolvedValue(false),
+}))
+
+import { requireOrganizationSearchAvailable } from '@/lib/knowledge/access/availability'
+import { resolveAppEntryPath } from '@/lib/navigation/resolve-app-entry'
+
+afterAll(resetEnvFlagsMock)
+
+describe('organization rollout during impersonation', () => {
+  beforeEach(() => {
+    vi.clearAllMocks()
+    setEnvFlags({ isAppConfigEnabled: true, isHosted: true })
+    mocks.landing.mockImplementation(async (userId: string) =>
+      userId === 'platform-admin' ? 'admin-org' : 'customer-org'
+    )
+    mocks.platformAdmin.mockImplementation(async (userId: string) => userId === 'platform-admin')
+  })
+
+  it.each([
+    { knowledge: false, groups: false },
+    { knowledge: false, groups: true },
+    { knowledge: true, groups: false },
+    { knowledge: true, groups: true },
+  ])('uses the customer organization for both gates: %j', async ({ knowledge, groups }) => {
+    const flags: FeatureFlagsConfig = {
+      'knowledge-member-access': {
+        orgIds: ['admin-org', ...(knowledge ? ['customer-org'] : [])],
+        userIds: ['platform-admin'],
+        adminEnabled: true,
+      },
+      'credential-groups': {
+        orgIds: ['admin-org', ...(groups ? ['customer-org'] : [])],
+        userIds: ['platform-admin'],
+        adminEnabled: true,
+      },
+    }
+    mocks.appConfig.mockResolvedValue(flags)
+
+    await expect(resolveAppEntryPath({ user: { id: 'platform-admin' } })).resolves.toBe(
+      '/o/admin-org/home'
+    )
+
+    const impersonatedSession = {
+      user: { id: 'customer-member' },
+      session: { impersonatedBy: 'platform-admin', activeOrganizationId: 'customer-org' },
+    }
+    await expect(resolveAppEntryPath(impersonatedSession)).resolves.toBe(
+      knowledge && groups ? '/o/customer-org/home' : '/workspace?redirect=settings'
+    )
+    expect(mocks.landing).toHaveBeenLastCalledWith('customer-member', 'customer-org')
+    expect(mocks.platformAdmin).not.toHaveBeenCalled()
+
+    if (knowledge && groups) {
+      await expect(requireOrganizationSearchAvailable('customer-org')).resolves.toBeUndefined()
+    } else {
+      await expect(requireOrganizationSearchAvailable('customer-org')).rejects.toMatchObject({
+        code: 'forbidden',
+      })
+    }
+  })
+})