fix(auth): give context on the login page and tolerate transient 401s

frontend/src/lib/logic/apiStatusLogic.test.ts

@@ -78,6 +78,36 @@ describe('apiStatusLogic', () => {
             logoutSpy.mockRestore()
             submitSpy.mockRestore()
         })
+
+        it('tolerates a transient 401 that clears on recheck', async () => {
+            let callCount = 0
+            useMocks({
+                get: {
+                    '/api/users/@me/': () => {
+                        callCount += 1
+                        // First recheck still 401 (the blip), second recheck succeeds
+                        return callCount === 1 ? [401, {}] : [200, MOCK_DEFAULT_USER]
+                    },
+                },
+            })
+            initKeaTests()
+            userLogic.mount()
+            userLogic.actions.loadUserSuccess(MOCK_DEFAULT_USER)
+
+            logic = apiStatusLogic()
+            logic.mount()
+
+            const logoutSpy = jest.spyOn(userLogic.actions, 'logout')
+
+            const mockResponse = { status: 401, ok: false } as Response
+
+            await expectLogic(logic, () => {
+                logic.actions.onApiResponse(mockResponse)
+            }).toFinishAllListeners()
+
+            expect(logoutSpy).not.toHaveBeenCalled()
+            logoutSpy.mockRestore()
+        })
     })
 
     describe('read-only impersonation 403 handling', () => {

frontend/src/lib/logic/apiStatusLogic.ts

@@ -6,6 +6,9 @@ import api from 'lib/api'
 import { twoFactorLogic } from 'scenes/authentication/two-factor-setup/twoFactorLogic'
 import { userLogic } from 'scenes/userLogic'
 
+// How long to wait before rechecking a 401 that could be a transient blip rather than a truly expired session.
+const TRANSIENT_401_RECHECK_DELAY_MS = 2000
+
 // Generated by kea-typegen. Update if you're an agent, ignore if you're human.
 export interface apiStatusLogicValues {
     internetConnectionIssue: boolean
@@ -148,11 +151,29 @@ export const apiStatusLogic = kea<apiStatusLogicType>([
                 if (now - 10000 > (cache.lastUnauthorizedCheck ?? 0)) {
                     cache.lastUnauthorizedCheck = Date.now()
 
-                    await api.get('api/users/@me/').catch((error: any) => {
-                        if (error.status === 401) {
+                    const isStillUnauthorized = async (): Promise<boolean> => {
+                        try {
+                            await api.get('api/users/@me/')
+                            return false
+                        } catch (error: any) {
+                            return error.status === 401
+                        }
+                    }
+
+                    if (await isStillUnauthorized()) {
+                        // A single 401 can be a transient blip (e.g. a load balancer hiccup or a race
+                        // during token rotation) rather than a truly expired session. Give it a moment
+                        // and check again before ejecting a live session.
+                        //
+                        // A plain timer (not kea's `breakpoint`) is used deliberately: the recheck
+                        // itself calls the API, which re-triggers this same `onApiResponse` listener
+                        // re-entrantly (see api.ts's own onApiResponse call), which would otherwise
+                        // invalidate `breakpoint`'s cancellation counter and abort this check.
+                        await new Promise<void>((resolve) => setTimeout(resolve, TRANSIENT_401_RECHECK_DELAY_MS))
+                        if (await isStillUnauthorized()) {
                             userLogic.findMounted()?.actions.logout(true)
                         }
-                    })
+                    }
                 }
             }
         },

frontend/src/scenes/authentication/login/Login.stories.tsx

@@ -157,6 +157,20 @@ export const SSOError: Story = {
     },
 }
 
+export const SessionExpired: Story = {
+    render: () => {
+        useStorybookMocks({
+            get: {
+                '/_preflight': preflightJson,
+            },
+        })
+
+        useDelayedOnMountEffect(() => router.actions.push(`${urls.login()}?next=/settings/user-api-keys`))
+
+        return <Login />
+    },
+}
+
 export const SecondFactor: Story = {
     render: () => {
         useDelayedOnMountEffect(() => router.actions.push(urls.login2FA()))

frontend/src/scenes/authentication/login/SessionRiskBanner.tsx

@@ -5,15 +5,24 @@ import { LemonBanner } from 'lib/lemon-ui/LemonBanner'
 import { loginLogic } from './loginLogic'
 
 export function SessionRiskBanner({ className }: { className?: string }): JSX.Element | null {
-    const { wasSignedOutForSessionRisk } = useValues(loginLogic)
+    const { wasSignedOutForSessionRisk, sessionExpiredRedirectPath } = useValues(loginLogic)
 
-    if (!wasSignedOutForSessionRisk) {
-        return null
+    if (wasSignedOutForSessionRisk) {
+        return (
+            <LemonBanner type="warning" className={className}>
+                For your security, we signed you out because this session showed unusual activity. Sign back in to
+                continue.
+            </LemonBanner>
+        )
     }
 
-    return (
-        <LemonBanner type="warning" className={className}>
-            For your security, we signed you out because this session showed unusual activity. Sign back in to continue.
-        </LemonBanner>
-    )
+    if (sessionExpiredRedirectPath) {
+        return (
+            <LemonBanner type="info" className={className}>
+                Your session expired. Sign in again to get back to {sessionExpiredRedirectPath}.
+            </LemonBanner>
+        )
+    }
+
+    return null
 }

frontend/src/scenes/authentication/login/loginLogic.test.ts

@@ -56,6 +56,33 @@ describe('loginLogic', () => {
         }
     })
 
+    describe('sessionExpiredRedirectPath', () => {
+        let logic: ReturnType<typeof loginLogic.build>
+
+        beforeEach(() => {
+            initKeaTests()
+            logic = loginLogic()
+            logic.mount()
+        })
+
+        const cases: [string, string | null][] = [
+            ['/login?next=/settings/user-api-keys', '/settings/user-api-keys'],
+            // The session-risk banner already explains the redirect, so don't show both
+            ['/login?next=/settings/user-api-keys&reason=session_risk', null],
+            // No next param means the user didn't get here via a login_required bounce
+            ['/login', null],
+            // Sanitized away by getRelativeNextPath - must not leak into the banner
+            ['/login?next=//evil.com', null],
+        ]
+
+        for (const [url, expected] of cases) {
+            it(`for "${url}" it returns ${expected}`, () => {
+                router.actions.push(url)
+                expect(logic.values.sessionExpiredRedirectPath).toEqual(expected)
+            })
+        }
+    })
+
     describe('parseLoginRedirectURL', () => {
         let logic: ReturnType<typeof loginLogic.build>
 

frontend/src/scenes/authentication/login/loginLogic.ts

@@ -132,6 +132,7 @@ export interface loginLogicValues {
         success: boolean
     } | null
     resendResponseLoading: boolean
+    sessionExpiredRedirectPath: string | null
     showCodeVerificationErrors: boolean
     showLoginErrors: boolean
     signupUrl: string
@@ -279,6 +280,7 @@ export interface loginLogicMeta {
     __keaTypeGenInternalSelectorTypes: {
         signupUrl: (searchParams: Record<string, any>) => string
         wasSignedOutForSessionRisk: (searchParams: Record<string, any>) => boolean
+        sessionExpiredRedirectPath: (searchParams: Record<string, any>) => string | null
     }
 }
 
@@ -386,6 +388,25 @@ export const loginLogic = kea<loginLogicType>([
             () => [router.selectors.searchParams],
             (searchParams: Record<string, string>): boolean => searchParams['reason'] === 'session_risk',
         ],
+        sessionExpiredRedirectPath: [
+            () => [router.selectors.searchParams],
+            (searchParams: Record<string, string>): string | null => {
+                // The session-risk banner already explains why the user landed here, so don't
+                // also show a generic "session expired" message in that case.
+                if (searchParams['reason'] === 'session_risk') {
+                    return null
+                }
+                const nextPath = getRelativeNextPath(searchParams['next'], location)
+                if (!nextPath) {
+                    return null
+                }
+                try {
+                    return new URL(location.origin + nextPath).pathname
+                } catch {
+                    return null
+                }
+            },
+        ],
     })),
     forms(({ actions, values }) => ({
         login: {