@workos-inc/authkit-nextjs 3.0.0-beta.1 → 3.0.1

package/src/components/impersonation.spec.tsx

@@ -1,3 +1,4 @@
+import type { Mock } from 'vitest';
 import { render, act, screen } from '@testing-library/react';
 import '@testing-library/jest-dom';
 import { Impersonation } from './impersonation.js';
@@ -7,39 +8,26 @@ import * as React from 'react';
 import { handleSignOutAction } from '../actions.js';
 
 // Mock the useAuth hook
-jest.mock('./authkit-provider', () => ({
-  useAuth: jest.fn(),
+vi.mock('./authkit-provider', () => ({
+  useAuth: vi.fn(),
 }));
 
 // Mock the getOrganizationAction
-jest.mock('../actions', () => ({
-  getOrganizationAction: jest.fn(),
-  handleSignOutAction: jest.fn(),
+vi.mock('../actions', () => ({
+  getOrganizationAction: vi.fn(),
+  handleSignOutAction: vi.fn(),
 }));
 
 describe('Impersonation', () => {
   beforeEach(() => {
-    jest.clearAllMocks();
+    vi.clearAllMocks();
   });
 
   it('should return null if not impersonating', () => {
-    (useAuth as jest.Mock).mockReturnValue({
+    (useAuth as Mock).mockReturnValue({
       impersonator: null,
       user: { id: '123', email: 'user@example.com' },
       organizationId: null,
-      loading: false,
-    });
-
-    const { container } = render(<Impersonation />);
-    expect(container).toBeEmptyDOMElement();
-  });
-
-  it('should return null if loading', () => {
-    (useAuth as jest.Mock).mockReturnValue({
-      impersonator: { email: 'admin@example.com' },
-      user: { id: '123', email: 'user@example.com' },
-      organizationId: null,
-      loading: true,
     });
 
     const { container } = render(<Impersonation />);
@@ -47,11 +35,10 @@ describe('Impersonation', () => {
   });
 
   it('should render impersonation banner when impersonating', () => {
-    (useAuth as jest.Mock).mockReturnValue({
+    (useAuth as Mock).mockReturnValue({
       impersonator: { email: 'admin@example.com' },
       user: { id: '123', email: 'user@example.com' },
       organizationId: null,
-      loading: false,
     });
 
     const { container } = render(<Impersonation />);
@@ -59,14 +46,13 @@ describe('Impersonation', () => {
   });
 
   it('should render with organization info when organizationId is provided', async () => {
-    (useAuth as jest.Mock).mockReturnValue({
+    (useAuth as Mock).mockReturnValue({
       impersonator: { email: 'admin@example.com' },
       user: { id: '123', email: 'user@example.com' },
       organizationId: 'org_123',
-      loading: false,
     });
 
-    (getOrganizationAction as jest.Mock).mockResolvedValue({
+    (getOrganizationAction as Mock).mockResolvedValue({
       id: 'org_123',
       name: 'Test Org',
     });
@@ -79,11 +65,10 @@ describe('Impersonation', () => {
   });
 
   it('should render at the bottom by default', () => {
-    (useAuth as jest.Mock).mockReturnValue({
+    (useAuth as Mock).mockReturnValue({
       impersonator: { email: 'admin@example.com' },
       user: { id: '123', email: 'user@example.com' },
       organizationId: null,
-      loading: false,
     });
 
     const { container } = render(<Impersonation />);
@@ -92,11 +77,10 @@ describe('Impersonation', () => {
   });
 
   it('should render at the top when side prop is "top"', () => {
-    (useAuth as jest.Mock).mockReturnValue({
+    (useAuth as Mock).mockReturnValue({
       impersonator: { email: 'admin@example.com' },
       user: { id: '123', email: 'user@example.com' },
       organizationId: null,
-      loading: false,
     });
 
     const { container } = render(<Impersonation side="top" />);
@@ -105,30 +89,163 @@ describe('Impersonation', () => {
   });
 
   it('should merge custom styles with default styles', () => {
-    (useAuth as jest.Mock).mockReturnValue({
+    (useAuth as Mock).mockReturnValue({
       impersonator: { email: 'admin@example.com' },
       user: { id: '123', email: 'user@example.com' },
       organizationId: null,
-      loading: false,
     });
 
     const customStyle = { backgroundColor: 'red' };
     const { container } = render(<Impersonation style={customStyle} />);
     const root = container.querySelector('[data-workos-impersonation-root]');
-    expect(root).toHaveStyle({ backgroundColor: 'red' });
+    expect((root as HTMLElement).style.backgroundColor).toBe('red');
   });
 
   it('should should sign out when the Stop button is called', async () => {
-    (useAuth as jest.Mock).mockReturnValue({
+    (useAuth as Mock).mockReturnValue({
       impersonator: { email: 'admin@example.com' },
       user: { id: '123', email: 'user@example.com' },
       organizationId: null,
-      loading: false,
     });
 
     render(<Impersonation />);
     const stopButton = await screen.findByText('Stop');
     stopButton.click();
-    expect(handleSignOutAction).toHaveBeenCalled();
+    expect(handleSignOutAction).toHaveBeenCalledWith({});
+  });
+
+  it('should pass returnTo prop to handleSignOutAction when provided', async () => {
+    (useAuth as Mock).mockReturnValue({
+      impersonator: { email: 'admin@example.com' },
+      user: { id: '123', email: 'user@example.com' },
+      organizationId: null,
+      loading: false,
+    });
+
+    const returnTo = '/dashboard';
+    render(<Impersonation returnTo={returnTo} />);
+    const stopButton = await screen.findByText('Stop');
+    stopButton.click();
+    expect(handleSignOutAction).toHaveBeenCalledWith({ returnTo });
+  });
+
+  it('should not call getOrganizationAction when organizationId is not provided', () => {
+    (useAuth as Mock).mockReturnValue({
+      impersonator: { email: 'admin@example.com' },
+      user: { id: '123', email: 'user@example.com' },
+      organizationId: null,
+    });
+
+    render(<Impersonation />);
+    expect(getOrganizationAction).not.toHaveBeenCalled();
+  });
+
+  it('should not call getOrganizationAction when impersonator is not present', () => {
+    (useAuth as Mock).mockReturnValue({
+      impersonator: null,
+      user: { id: '123', email: 'user@example.com' },
+      organizationId: 'org_123',
+    });
+
+    render(<Impersonation />);
+    expect(getOrganizationAction).not.toHaveBeenCalled();
+  });
+
+  it('should not call getOrganizationAction when user is not present', () => {
+    (useAuth as Mock).mockReturnValue({
+      impersonator: { email: 'admin@example.com' },
+      user: null,
+      organizationId: 'org_123',
+    });
+
+    render(<Impersonation />);
+    expect(getOrganizationAction).not.toHaveBeenCalled();
+  });
+
+  it('should not call getOrganizationAction again when organization is already loaded with same ID', async () => {
+    const mockOrg = {
+      id: 'org_123',
+      name: 'Test Org',
+    };
+
+    (getOrganizationAction as Mock).mockResolvedValue(mockOrg);
+
+    const { rerender } = await act(async () => {
+      (useAuth as Mock).mockReturnValue({
+        impersonator: { email: 'admin@example.com' },
+        user: { id: '123', email: 'user@example.com' },
+        organizationId: 'org_123',
+      });
+
+      return render(<Impersonation />);
+    });
+
+    // Wait for the initial call to complete
+    await act(async () => {
+      await new Promise((resolve) => setTimeout(resolve, 0));
+    });
+
+    expect(getOrganizationAction).toHaveBeenCalledTimes(1);
+
+    // Rerender with the same organizationId
+    await act(async () => {
+      (useAuth as Mock).mockReturnValue({
+        impersonator: { email: 'admin@example.com' },
+        user: { id: '123', email: 'user@example.com' },
+        organizationId: 'org_123',
+      });
+
+      rerender(<Impersonation />);
+    });
+
+    // Should still be called only once
+    expect(getOrganizationAction).toHaveBeenCalledTimes(1);
+  });
+
+  it('should call getOrganizationAction again when organizationId changes', async () => {
+    const mockOrg1 = {
+      id: 'org_123',
+      name: 'Test Org 1',
+    };
+
+    const mockOrg2 = {
+      id: 'org_456',
+      name: 'Test Org 2',
+    };
+
+    (getOrganizationAction as Mock).mockResolvedValueOnce(mockOrg1).mockResolvedValueOnce(mockOrg2);
+
+    const { rerender } = await act(async () => {
+      (useAuth as Mock).mockReturnValue({
+        impersonator: { email: 'admin@example.com' },
+        user: { id: '123', email: 'user@example.com' },
+        organizationId: 'org_123',
+      });
+
+      return render(<Impersonation />);
+    });
+
+    // Wait for the initial call to complete
+    await act(async () => {
+      await new Promise((resolve) => setTimeout(resolve, 0));
+    });
+
+    expect(getOrganizationAction).toHaveBeenCalledTimes(1);
+    expect(getOrganizationAction).toHaveBeenCalledWith('org_123');
+
+    // Rerender with a different organizationId
+    await act(async () => {
+      (useAuth as Mock).mockReturnValue({
+        impersonator: { email: 'admin@example.com' },
+        user: { id: '123', email: 'user@example.com' },
+        organizationId: 'org_456',
+      });
+
+      rerender(<Impersonation />);
+    });
+
+    // Should be called again with the new ID
+    expect(getOrganizationAction).toHaveBeenCalledTimes(2);
+    expect(getOrganizationAction).toHaveBeenCalledWith('org_456');
   });
 });

package/src/components/impersonation.tsx

@@ -9,53 +9,60 @@ import { useAuth } from './authkit-provider.js';
 
 interface ImpersonationProps extends React.ComponentPropsWithoutRef<'div'> {
   side?: 'top' | 'bottom';
+  returnTo?: string;
 }
 
-export function Impersonation({ side = 'bottom', ...props }: ImpersonationProps) {
-  const { user, impersonator, organizationId, loading } = useAuth();
+export function Impersonation({ side = 'bottom', returnTo, ...props }: ImpersonationProps) {
+  const { user, impersonator, organizationId } = useAuth();
 
   const [organization, setOrganization] = React.useState<Organization | null>(null);
 
   React.useEffect(() => {
-    if (!organizationId) return;
+    if (!organizationId || !impersonator || !user) return;
+    if (organization && organization.id === organizationId) return;
     getOrganizationAction(organizationId).then(setOrganization);
-  }, [organizationId]);
+  }, [organizationId, impersonator, user]);
 
-  if (loading || !impersonator || !user) return null;
+  if (!impersonator || !user) return null;
 
   return (
     <div
       {...props}
       data-workos-impersonation-root=""
-      style={{
-        'position': 'fixed',
-        'inset': 0,
-        'pointerEvents': 'none',
-        'zIndex': 9999,
-
-        // short properties with defaults for authoring convenience
-        '--wi-minimized': '0',
-        '--wi-s': 'min(max(var(--workos-impersonation-size, 4px), 2px), 15px)',
-        '--wi-bgc': 'var(--workos-impersonation-background-color, #fce654)',
-        '--wi-c': 'var(--workos-impersonation-color, #1a1600)',
-        '--wi-bc': 'var(--workos-impersonation-border-color, #e0c36c)',
-        '--wi-bw': 'var(--workos-impersonation-border-width, 1px)',
-
-        ...props.style,
-      }}
+      style={
+        {
+          position: 'fixed',
+          inset: 0,
+          pointerEvents: 'none',
+          zIndex: 9999,
+
+          // short properties with defaults for authoring convenience
+          '--wi-minimized': '0',
+          '--wi-s': 'min(max(var(--workos-impersonation-size, 4px), 2px), 15px)',
+          '--wi-bgc': 'var(--workos-impersonation-background-color, #fce654)',
+          '--wi-c': 'var(--workos-impersonation-color, #1a1600)',
+          '--wi-bc': 'var(--workos-impersonation-border-color, #e0c36c)',
+          '--wi-bw': 'var(--workos-impersonation-border-width, 1px)',
+
+          ...props.style,
+        } as React.CSSProperties
+      }
     >
       <div
-        style={{
-          '--wi-frame-size': 'calc(var(--wi-s) * (1 - var(--wi-minimized)) + var(--wi-minimized) * var(--wi-bw) * -1)',
-          'position': 'absolute',
-          'inset': 'calc(var(--wi-frame-size) * -1)',
-          'borderRadius': 'calc(var(--wi-frame-size) * 3)',
-          'boxShadow': `
+        style={
+          {
+            '--wi-frame-size':
+              'calc(var(--wi-s) * (1 - var(--wi-minimized)) + var(--wi-minimized) * var(--wi-bw) * -1)',
+            position: 'absolute',
+            inset: 'calc(var(--wi-frame-size) * -1)',
+            borderRadius: 'calc(var(--wi-frame-size) * 3)',
+            boxShadow: `
 						inset 0 0 0 calc(var(--wi-frame-size) * 2) var(--wi-bgc),
 						inset 0 0 0 calc(var(--wi-frame-size) * 2 + var(--wi-bw)) var(--wi-bc)
 					`,
-          'transition': 'all 500ms cubic-bezier(0.16, 1, 0.3, 1)',
-        }}
+            transition: 'all 500ms cubic-bezier(0.16, 1, 0.3, 1)',
+          } as React.CSSProperties
+        }
       />
 
       <div
@@ -78,7 +85,7 @@ export function Impersonation({ side = 'bottom', ...props }: ImpersonationProps)
         <form
           onSubmit={async (event) => {
             event.preventDefault();
-            await handleSignOutAction();
+            await handleSignOutAction({ returnTo });
           }}
           style={{
             display: 'flex',

package/src/components/min-max-button.spec.tsx

@@ -14,6 +14,7 @@ describe('MinMaxButton', () => {
   afterEach(() => {
     // Clean up after each test
     document.body.innerHTML = '';
+    vi.restoreAllMocks();
   });
 
   it('sets minimized value when clicked', () => {
@@ -33,7 +34,7 @@ describe('MinMaxButton', () => {
     const root = document.querySelector('[data-workos-impersonation-root]');
 
     // Mock querySelector to return null for this test
-    jest.spyOn(document, 'querySelector').mockReturnValue(null);
+    vi.spyOn(document, 'querySelector').mockReturnValue(null);
 
     act(() => {
       getByRole('button').click();

package/src/components/tokenStore.spec.ts

@@ -1,20 +1,21 @@
+import type { Mock } from 'vitest';
 import { tokenStore, TokenStore } from './tokenStore.js';
 import { getAccessTokenAction, refreshAccessTokenAction } from '../actions.js';
 
-jest.mock('../actions.js', () => ({
-  getAccessTokenAction: jest.fn(),
-  refreshAccessTokenAction: jest.fn(),
+vi.mock('../actions.js', () => ({
+  getAccessTokenAction: vi.fn(),
+  refreshAccessTokenAction: vi.fn(),
 }));
 
-const mockGetAccessTokenAction = getAccessTokenAction as jest.Mock;
-const mockRefreshAccessTokenAction = refreshAccessTokenAction as jest.Mock;
+const mockGetAccessTokenAction = getAccessTokenAction as Mock;
+const mockRefreshAccessTokenAction = refreshAccessTokenAction as Mock;
 // eslint-disable-next-line @typescript-eslint/no-explicit-any
 const _global = global as any;
 
 describe('tokenStore', () => {
   beforeEach(() => {
-    jest.useFakeTimers();
-    jest.resetAllMocks();
+    vi.useFakeTimers();
+    vi.resetAllMocks();
     tokenStore.reset();
 
     // Clean up DOM globals
@@ -23,10 +24,10 @@ describe('tokenStore', () => {
   });
 
   afterEach(() => {
-    jest.clearAllTimers();
-    jest.useRealTimers();
+    vi.clearAllTimers();
+    vi.useRealTimers();
     tokenStore.reset();
-    jest.restoreAllMocks();
+    vi.restoreAllMocks();
 
     // Clean up DOM globals
     delete _global.document;
@@ -54,7 +55,7 @@ describe('tokenStore', () => {
         resolvePromise = resolve;
       });
 
-      mockRefreshAccessTokenAction.mockReturnValue(slowPromise);
+      mockRefreshAccessTokenAction.mockReturnValue(slowPromise.then((t) => ({ accessToken: t })));
 
       expect(tokenStore.isRefreshing()).toBe(false);
 
@@ -124,18 +125,25 @@ describe('tokenStore', () => {
       const refreshedToken =
         'eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiJyZWZyZXNoZWQiLCJzaWQiOiJzZXNzaW9uXzEyMyIsImV4cCI6OTk5OTk5OTk5OX0.mock-signature-2';
 
-      // Set expiring token first
+      // Set expiring token first — also set refresh mock since getAccessTokenSilently
+      // will trigger refresh for expiring tokens
       mockGetAccessTokenAction.mockResolvedValue(expiringToken);
+      mockRefreshAccessTokenAction.mockResolvedValueOnce({ accessToken: expiringToken });
       await tokenStore.getAccessTokenSilently();
 
-      // Setup refresh mock
-      mockRefreshAccessTokenAction.mockResolvedValue(refreshedToken);
+      // Clear mocks to track subsequent calls
+      mockGetAccessTokenAction.mockClear();
+      mockRefreshAccessTokenAction.mockClear();
 
-      // Now call getAccessToken - should trigger refresh
+      // Setup refresh to return new token
+      mockRefreshAccessTokenAction.mockResolvedValue({ accessToken: refreshedToken });
+
+      // Now call getAccessToken - should trigger refresh due to expiring token
       const token = await tokenStore.getAccessToken();
 
-      expect(token).toBe(refreshedToken);
+      // Should have called refresh since existing token was expiring
       expect(mockRefreshAccessTokenAction).toHaveBeenCalled();
+      expect(token).toBe(refreshedToken);
     });
 
     it('should refresh when no token exists', async () => {
@@ -192,7 +200,7 @@ describe('tokenStore', () => {
         'eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiJyZWZyZXNoZWQiLCJzaWQiOiJzZXNzaW9uXzEyMyIsImV4cCI6OTk5OTk5OTk5OX0.mock-signature-2';
 
       mockGetAccessTokenAction.mockResolvedValue(expiredToken);
-      mockRefreshAccessTokenAction.mockResolvedValue(refreshedToken);
+      mockRefreshAccessTokenAction.mockResolvedValue({ accessToken: refreshedToken });
 
       const token = await tokenStore.getAccessTokenSilently();
 
@@ -213,7 +221,7 @@ describe('tokenStore', () => {
       };
       const validToken = `eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.${btoa(JSON.stringify(validPayload))}.mock-signature`;
 
-      const setTimeoutSpy = jest.spyOn(global, 'setTimeout');
+      const setTimeoutSpy = vi.spyOn(global, 'setTimeout');
       mockGetAccessTokenAction.mockResolvedValue(validToken);
 
       await tokenStore.getAccessTokenSilently();
@@ -227,7 +235,7 @@ describe('tokenStore', () => {
 
   describe('subscriber management', () => {
     it('should notify subscribers when state changes', () => {
-      const listener = jest.fn();
+      const listener = vi.fn();
       const unsubscribe = tokenStore.subscribe(listener);
 
       // Trigger a state change
@@ -256,14 +264,14 @@ describe('tokenStore', () => {
       mockGetAccessTokenAction.mockResolvedValue(validToken);
 
       // Subscribe to create a listener
-      const listener = jest.fn();
+      const listener = vi.fn();
       const unsubscribe = tokenStore.subscribe(listener);
 
       // Get token to schedule a refresh
       await tokenStore.getAccessTokenSilently();
 
       // Spy on clearTimeout
-      const clearTimeoutSpy = jest.spyOn(global, 'clearTimeout');
+      const clearTimeoutSpy = vi.spyOn(global, 'clearTimeout');
 
       // Unsubscribe the last (only) subscriber - should clear timeout
       unsubscribe();
@@ -287,15 +295,18 @@ describe('tokenStore', () => {
       const refreshedToken =
         'eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiJyZWZyZXNoZWQiLCJzaWQiOiJzZXNzaW9uXzEyMyIsImV4cCI6OTk5OTk5OTk5OX0.mock-signature-2';
 
-      // First set an expiring token
+      // First set an expiring token — also set refresh mock since getAccessTokenSilently
+      // will trigger refresh for expiring tokens
       mockGetAccessTokenAction.mockResolvedValue(expiringToken);
+      mockRefreshAccessTokenAction.mockResolvedValueOnce({ accessToken: expiringToken });
       await tokenStore.getAccessTokenSilently();
 
       // Clear mocks
       mockGetAccessTokenAction.mockClear();
+      mockRefreshAccessTokenAction.mockClear();
 
       // Setup refresh to return new token
-      mockRefreshAccessTokenAction.mockResolvedValue(refreshedToken);
+      mockRefreshAccessTokenAction.mockResolvedValue({ accessToken: refreshedToken });
 
       // Call getAccessToken again - should trigger refresh due to expiring token
       const token = await tokenStore.getAccessToken();
@@ -354,7 +365,7 @@ describe('tokenStore', () => {
 
     it('should consume eager auth cookie on first getAccessToken call', async () => {
       const eagerToken = 'eager-auth-token';
-      const mockCookieSetter = jest.fn();
+      const mockCookieSetter = vi.fn();
 
       // Mock document.cookie with both getter and setter
       let cookieValue = `workos-access-token=${eagerToken};`;
@@ -404,7 +415,7 @@ describe('tokenStore', () => {
         iat: now - 40,
       };
       const fastToken = `eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.${btoa(JSON.stringify(fastPayload))}.mock-signature`;
-      const mockCookieSetter = jest.fn();
+      const mockCookieSetter = vi.fn();
 
       let cookieValue = `workos-access-token=${fastToken};`;
 
@@ -433,7 +444,7 @@ describe('tokenStore', () => {
         configurable: true,
       });
 
-      const setTimeoutSpy = jest.spyOn(global, 'setTimeout');
+      const setTimeoutSpy = vi.spyOn(global, 'setTimeout');
 
       // Call getAccessTokenSilently to trigger fast cookie consumption and refresh scheduling
       const token = await tokenStore.getAccessTokenSilently();
@@ -480,7 +491,7 @@ describe('tokenStore', () => {
 
     it('should handle HTTP protocol for cookie deletion', async () => {
       const eagerToken = 'http-token';
-      const mockCookieSetter = jest.fn();
+      const mockCookieSetter = vi.fn();
 
       let cookieValue = `workos-access-token=${eagerToken};`;
 
@@ -526,11 +537,14 @@ describe('tokenStore', () => {
       await tokenStore.getAccessTokenSilently();
 
       // Now simulate network error during refresh
-      mockRefreshAccessTokenAction.mockRejectedValue(new Error('Network error'));
+      mockRefreshAccessTokenAction.mockResolvedValue({
+        accessToken: undefined,
+        error: 'Failed to refresh access token',
+      });
 
       try {
         await tokenStore.refreshToken();
-      } catch (e) {
+      } catch {
         // Expected to throw
       }
 
@@ -543,11 +557,11 @@ describe('tokenStore', () => {
     it('should convert non-Error objects to Error instances', async () => {
       const errorString = 'network timeout';
 
-      mockRefreshAccessTokenAction.mockRejectedValue(errorString);
+      mockRefreshAccessTokenAction.mockResolvedValue({ accessToken: undefined, error: errorString });
 
       try {
         await tokenStore.refreshToken();
-      } catch (e) {
+      } catch {
         // Expected to throw
       }
 
@@ -624,7 +638,7 @@ describe('tokenStore', () => {
       await tokenStore.getAccessTokenSilently();
 
       // Spy on clearTimeout
-      const clearTimeoutSpy = jest.spyOn(global, 'clearTimeout');
+      const clearTimeoutSpy = vi.spyOn(global, 'clearTimeout');
 
       // Clear token should clear the refresh timeout
       tokenStore.clearToken();
@@ -649,7 +663,7 @@ describe('tokenStore', () => {
 
       mockRefreshAccessTokenAction.mockImplementation(() => {
         callCount++;
-        return slowPromise;
+        return slowPromise.then((t) => ({ accessToken: t }));
       });
 
       // Clear any existing refresh promise
@@ -691,21 +705,22 @@ describe('tokenStore', () => {
   });
 
   describe('refresh state management', () => {
-    it('should preserve Error instances without conversion', async () => {
-      const errorInstance = new Error('actual error instance');
+    it('should create Error from server action error string', async () => {
+      const errorMessage = 'actual error instance';
 
-      // Mock refresh to throw an Error instance
-      mockRefreshAccessTokenAction.mockRejectedValue(errorInstance);
+      // Mock refresh to return an error result (as server actions do)
+      mockRefreshAccessTokenAction.mockResolvedValue({ accessToken: undefined, error: errorMessage });
 
       try {
         await tokenStore.refreshToken();
-      } catch (e) {
+      } catch {
         // Expected to throw
       }
 
-      // Verify the Error instance was preserved without conversion
+      // Verify an Error was created from the error string
       const state = tokenStore.getSnapshot();
-      expect(state.error).toBe(errorInstance); // Same instance, not a new one
+      expect(state.error).toBeInstanceOf(Error);
+      expect(state.error?.message).toBe(errorMessage);
     });
 
     it('should update state for manual refresh', async () => {
@@ -717,7 +732,7 @@ describe('tokenStore', () => {
       await tokenStore.getAccessTokenSilently();
 
       // Mock refresh to return new token
-      mockRefreshAccessTokenAction.mockResolvedValue(newToken);
+      mockRefreshAccessTokenAction.mockResolvedValue({ accessToken: newToken });
 
       // Call manual refresh which should update state
       const result = await tokenStore.refreshToken();
@@ -736,9 +751,9 @@ describe('tokenStore', () => {
 
       // Clear mocks and set up spy on setState
       mockGetAccessTokenAction.mockClear();
-      mockRefreshAccessTokenAction.mockResolvedValue(existingToken); // Same token
+      mockRefreshAccessTokenAction.mockResolvedValue({ accessToken: existingToken }); // Same token
 
-      const listener = jest.fn();
+      const listener = vi.fn();
       tokenStore.subscribe(listener);
 
       // Force a silent refresh that returns the same token
@@ -752,7 +767,7 @@ describe('tokenStore', () => {
 
   describe('TokenStore constructor', () => {
     const setupMockEnv = (cookieValue = '', protocol = 'https:') => {
-      const mockCookieSetter = jest.fn();
+      const mockCookieSetter = vi.fn();
 
       Object.defineProperty(_global, 'document', {
         value: { cookie: cookieValue },