diff --git a/packages/core/src/auth/auth.test.ts b/packages/core/src/auth/auth.test.ts index 38be8ae05e..b65d8b6fce 100644 --- a/packages/core/src/auth/auth.test.ts +++ b/packages/core/src/auth/auth.test.ts @@ -79,16 +79,21 @@ const mockLogger: RootLogger = { function mockTokenResponse( overrides: { accessToken?: string; - refreshToken?: string; + refreshToken?: string | null; + expiresIn?: number; scopedOrgs?: string[]; } = {}, ) { + const refreshToken = + overrides.refreshToken === undefined + ? "refresh-token" + : overrides.refreshToken; return { success: true as const, data: { access_token: overrides.accessToken ?? "access-token", - refresh_token: overrides.refreshToken ?? "refresh-token", - expires_in: 3600, + ...(refreshToken ? { refresh_token: refreshToken } : {}), + expires_in: overrides.expiresIn ?? 3600, token_type: "Bearer", scope: "", scoped_organizations: overrides.scopedOrgs ?? ["org-1"], @@ -242,9 +247,157 @@ describe("AuthService", () => { currentProjectId: null, hasCodeAccess: null, needsScopeReauth: false, + sessionType: null, + sessionExpiresAt: null, + sessionEndReason: null, }); }); + it("uses an impersonated session without persisting it", async () => { + oauthFlow.startFlow.mockResolvedValue( + mockTokenResponse({ refreshToken: null }), + ); + stubAuthFetch(); + + await service.initialize(); + await service.login("us"); + + expect(service.getState()).toMatchObject({ + status: "authenticated", + sessionType: "impersonated", + sessionExpiresAt: expect.any(Number), + }); + await expect(service.getValidAccessToken()).resolves.toMatchObject({ + accessToken: "access-token", + }); + expect(sessionPort.getCurrent()).toBeNull(); + }); + + it("signs out an impersonated session when a refresh is required", async () => { + oauthFlow.startFlow.mockResolvedValue( + mockTokenResponse({ refreshToken: null }), + ); + stubAuthFetch(); + + await service.initialize(); + await service.login("us"); + + await expect(service.refreshAccessToken()).rejects.toThrow( + "Your impersonated session has expired", + ); + expect(service.getState()).toMatchObject({ + status: "anonymous", + sessionType: null, + sessionExpiresAt: null, + sessionEndReason: "impersonation_expired", + }); + expect(oauthFlow.refreshToken).not.toHaveBeenCalled(); + }); + + it("signs out when an impersonated access token expires", async () => { + vi.useFakeTimers(); + try { + oauthFlow.startFlow.mockResolvedValue( + mockTokenResponse({ refreshToken: null, expiresIn: 61 }), + ); + stubAuthFetch(); + + await service.initialize(); + await service.login("us"); + await vi.advanceTimersByTimeAsync(60_000); + + await expect(service.getValidAccessToken()).resolves.toMatchObject({ + accessToken: "access-token", + }); + + await vi.advanceTimersByTimeAsync(1_001); + + expect(service.getState()).toMatchObject({ + status: "anonymous", + sessionEndReason: "impersonation_expired", + }); + } finally { + vi.useRealTimers(); + } + }); + + it("does not persist or extend impersonated credentials after a project change", async () => { + vi.useFakeTimers(); + try { + oauthFlow.startFlow.mockResolvedValue( + mockTokenResponse({ refreshToken: null, expiresIn: 61 }), + ); + stubAuthFetch({ + orgs: { + "org-1": { + name: "Org 1", + projects: [ + { id: 42, name: "Project 42" }, + { id: 84, name: "Project 84" }, + ], + }, + }, + }); + + await service.initialize(); + await service.login("us"); + await service.selectProject(84); + + expect(service.getState().currentProjectId).toBe(84); + expect(sessionPort.getCurrent()).toBeNull(); + + await vi.advanceTimersByTimeAsync(61_001); + + expect(service.getState()).toMatchObject({ + status: "anonymous", + sessionEndReason: "impersonation_expired", + }); + } finally { + vi.useRealTimers(); + } + }); + + it("does not restore an impersonated session after restart", async () => { + oauthFlow.startFlow.mockResolvedValue( + mockTokenResponse({ refreshToken: null }), + ); + stubAuthFetch(); + + await service.initialize(); + await service.login("us"); + service.shutdown(); + + service = createService(); + service.init(); + await service.initialize(); + + expect(service.getState()).toMatchObject({ + status: "anonymous", + sessionType: null, + }); + expect(oauthFlow.refreshToken).not.toHaveBeenCalled(); + }); + + it("returns permission-denied responses without ending impersonation", async () => { + oauthFlow.startFlow.mockResolvedValue( + mockTokenResponse({ refreshToken: null }), + ); + stubAuthFetch(); + + await service.initialize(); + await service.login("us"); + vi.mocked(fetch).mockResolvedValueOnce(new Response(null, { status: 403 })); + + const response = await service.authenticatedFetch( + fetch, + "https://us.posthog.com/api/restricted/", + ); + + expect(response.status).toBe(403); + expect(service.getState().status).toBe("authenticated"); + expect(oauthFlow.refreshToken).not.toHaveBeenCalled(); + }); + it("requires scope reauthentication when the stored scope version is stale", async () => { seedStoredSession({ refreshToken: "refresh-token", @@ -263,6 +416,9 @@ describe("AuthService", () => { currentProjectId: 123, hasCodeAccess: null, needsScopeReauth: true, + sessionType: null, + sessionExpiresAt: null, + sessionEndReason: null, }); }); @@ -312,6 +468,24 @@ describe("AuthService", () => { ); }); + it("keeps the existing refresh token when the server does not rotate it", async () => { + seedStoredSession({ refreshToken: "existing-refresh-token" }); + oauthFlow.refreshToken.mockResolvedValue( + mockTokenResponse({ refreshToken: null }), + ); + stubAuthFetch(); + + await service.initialize(); + + expect(service.getState()).toMatchObject({ + status: "authenticated", + sessionType: "persistent", + }); + expect(sessionPort.getCurrent()?.refreshTokenEncrypted).toBe( + "existing-refresh-token", + ); + }); + it("completes bootstrap but stays restoring when the stored-session restore hangs", async () => { vi.useFakeTimers(); try { diff --git a/packages/core/src/auth/auth.ts b/packages/core/src/auth/auth.ts index ec0512ed4e..d325c0d808 100644 --- a/packages/core/src/auth/auth.ts +++ b/packages/core/src/auth/auth.ts @@ -52,7 +52,8 @@ interface InMemorySession { accountKey: string | null; accessToken: string; accessTokenExpiresAt: number; - refreshToken: string; + refreshToken: string | null; + sessionType: "persistent" | "impersonated"; cloudRegion: CloudRegion; orgProjectsMap: OrgProjectsMap; currentOrgId: string | null; @@ -69,6 +70,7 @@ interface StoredSessionInput { interface TokenResponseOptions { cloudRegion: CloudRegion; selectedProjectId: number | null; + fallbackRefreshToken?: string; } @injectable() @@ -82,10 +84,14 @@ export class AuthService extends TypedEventEmitter { currentProjectId: null, hasCodeAccess: null, needsScopeReauth: false, + sessionType: null, + sessionExpiresAt: null, + sessionEndReason: null, }; private session: InMemorySession | null = null; private initializePromise: Promise | null = null; private refreshPromise: Promise | null = null; + private impersonationExpiryTimer: ReturnType | null = null; // Serializes session-state commits so overlapping selections can't // interleave across async encryption (see commitSessionState). private commitChain: Promise = Promise.resolve(); @@ -157,7 +163,7 @@ export class AuthService extends TypedEventEmitter { } async getOAuthCredentials(): Promise<{ access: string; - refresh: string; + refresh: string | null; expires: number; region: CloudRegion; } | null> { @@ -216,6 +222,13 @@ export class AuthService extends TypedEventEmitter { initialAuth.accessToken, ); + if ( + response.status === 403 && + this.session?.sessionType === "impersonated" + ) { + return response; + } + if (response.status === 401 || response.status === 403) { const refreshedAuth = await this.refreshAccessToken(); response = await this.executeAuthenticatedFetch( @@ -383,11 +396,13 @@ export class AuthService extends TypedEventEmitter { // encryption may reject). Mutate this.session, the preference, and // published state only after it resolves, so a rejection leaves every // layer on the prior session. - await this.persistSession({ - refreshToken: nextSession.refreshToken, - cloudRegion: nextSession.cloudRegion, - selectedProjectId: next.currentProjectId, - }); + if (nextSession.refreshToken) { + await this.persistSession({ + refreshToken: nextSession.refreshToken, + cloudRegion: nextSession.cloudRegion, + selectedProjectId: next.currentProjectId, + }); + } this.session = nextSession; this.persistProjectPreference(nextSession); @@ -417,6 +432,7 @@ export class AuthService extends TypedEventEmitter { const { cloudRegion, currentProjectId } = this.state; this.authSession.clearCurrent(); + this.clearImpersonationExpiryTimer(); this.session = null; this.setAnonymousState({ cloudRegion, currentProjectId }); return this.getState(); @@ -538,6 +554,16 @@ export class AuthService extends TypedEventEmitter { return currentSession; } + if (currentSession && !currentSession.refreshToken) { + if (!forceRefresh && !this.isSessionExpired(currentSession)) { + return currentSession; + } + this.endImpersonatedSession(currentSession); + throw new NotAuthenticatedError( + "Your impersonated session has expired. Impersonate the user again to continue.", + ); + } + if (this.refreshPromise) { return this.refreshPromise; } @@ -579,6 +605,11 @@ export class AuthService extends TypedEventEmitter { private async getSessionInputForRefresh(): Promise { if (this.session) { + if (!this.session.refreshToken) { + throw new NotAuthenticatedError( + "Your impersonated session has expired. Impersonate the user again to continue.", + ); + } return { refreshToken: this.session.refreshToken, cloudRegion: this.session.cloudRegion, @@ -613,7 +644,10 @@ export class AuthService extends TypedEventEmitter { ); if (result.success && result.data) { - return await this.createSessionFromTokenResponse(result.data, input); + return await this.createSessionFromTokenResponse(result.data, { + ...input, + fallbackRefreshToken: input.refreshToken, + }); } lastError = result.error || "Token refresh failed"; @@ -686,11 +720,14 @@ export class AuthService extends TypedEventEmitter { lastSelectedOrgId: lastPrefs?.lastSelectedOrgId ?? null, }); + const refreshToken = + tokenResponse.refresh_token ?? options.fallbackRefreshToken ?? null; const session: InMemorySession = { accountKey, accessToken: tokenResponse.access_token, accessTokenExpiresAt: Date.now() + tokenResponse.expires_in * 1000, - refreshToken: tokenResponse.refresh_token, + refreshToken, + sessionType: refreshToken ? "persistent" : "impersonated", cloudRegion: options.cloudRegion, orgProjectsMap, currentOrgId, @@ -839,13 +876,18 @@ export class AuthService extends TypedEventEmitter { session: InMemorySession, ): Promise { this.persistProjectPreference(session); - await this.persistSession({ - refreshToken: session.refreshToken, - cloudRegion: session.cloudRegion, - selectedProjectId: session.currentProjectId, - }); + if (session.refreshToken) { + await this.persistSession({ + refreshToken: session.refreshToken, + cloudRegion: session.cloudRegion, + selectedProjectId: session.currentProjectId, + }); + } else { + this.authSession.clearCurrent(); + } this.session = session; + this.scheduleImpersonationExpiry(session); this.updateState({ status: "authenticated", bootstrapComplete: true, @@ -854,6 +896,9 @@ export class AuthService extends TypedEventEmitter { currentOrgId: session.currentOrgId, currentProjectId: session.currentProjectId, needsScopeReauth: false, + sessionType: session.sessionType, + sessionExpiresAt: session.accessTokenExpiresAt, + sessionEndReason: null, }); await this.updateCodeAccessFromSession(); @@ -967,6 +1012,7 @@ export class AuthService extends TypedEventEmitter { | "cloudRegion" | "currentProjectId" | "needsScopeReauth" + | "sessionEndReason" > = {}, ): void { this.updateState({ @@ -978,6 +1024,9 @@ export class AuthService extends TypedEventEmitter { currentProjectId: partial.currentProjectId ?? null, hasCodeAccess: null, needsScopeReauth: partial.needsScopeReauth ?? false, + sessionType: null, + sessionExpiresAt: null, + sessionEndReason: partial.sessionEndReason ?? null, }); } private async updateCodeAccessFromSession(): Promise { @@ -1088,11 +1137,46 @@ export class AuthService extends TypedEventEmitter { } @preDestroy() shutdown(): void { + this.clearImpersonationExpiryTimer(); this.connectivityUnsubscribe?.(); this.connectivityUnsubscribe = null; this.resumeUnsubscribe?.(); this.resumeUnsubscribe = null; } + + private scheduleImpersonationExpiry(session: InMemorySession): void { + this.clearImpersonationExpiryTimer(); + if (session.sessionType !== "impersonated") return; + + const delayMs = Math.max(0, session.accessTokenExpiresAt - Date.now()); + this.impersonationExpiryTimer = setTimeout(() => { + this.impersonationExpiryTimer = null; + const currentSession = this.session; + if ( + currentSession?.sessionType === "impersonated" && + this.isSessionExpired(currentSession) + ) { + this.endImpersonatedSession(currentSession); + } + }, delayMs); + } + + private clearImpersonationExpiryTimer(): void { + if (this.impersonationExpiryTimer) { + clearTimeout(this.impersonationExpiryTimer); + this.impersonationExpiryTimer = null; + } + } + + private endImpersonatedSession(session: InMemorySession): void { + this.clearImpersonationExpiryTimer(); + this.session = null; + this.setAnonymousState({ + cloudRegion: session.cloudRegion, + currentProjectId: session.currentProjectId, + sessionEndReason: "impersonation_expired", + }); + } private handleResume = (): void => { this.attemptSessionRecovery(); }; diff --git a/packages/core/src/auth/oauth.schemas.ts b/packages/core/src/auth/oauth.schemas.ts index c8d51cc73d..81d7ffe1b2 100644 --- a/packages/core/src/auth/oauth.schemas.ts +++ b/packages/core/src/auth/oauth.schemas.ts @@ -23,7 +23,7 @@ export const oAuthTokenResponse = z.object({ expires_in: z.number(), token_type: z.string(), scope: z.string().optional().default(""), - refresh_token: z.string(), + refresh_token: z.string().optional(), scoped_organizations: z.array(z.string()).optional(), }); export type OAuthTokenResponse = z.infer; diff --git a/packages/core/src/auth/schemas.ts b/packages/core/src/auth/schemas.ts index 3fa2d74ad9..5318136b0f 100644 --- a/packages/core/src/auth/schemas.ts +++ b/packages/core/src/auth/schemas.ts @@ -80,6 +80,9 @@ export const authStateSchema = z.object({ currentProjectId: z.number().nullable(), hasCodeAccess: z.boolean().nullable(), needsScopeReauth: z.boolean(), + sessionType: z.enum(["persistent", "impersonated"]).nullable(), + sessionExpiresAt: z.number().nullable(), + sessionEndReason: z.enum(["impersonation_expired"]).nullable().optional(), }); export type AuthState = z.infer; diff --git a/packages/core/src/oauth/oauth.test.ts b/packages/core/src/oauth/oauth.test.ts index bf680e42f3..7e0d22d7de 100644 --- a/packages/core/src/oauth/oauth.test.ts +++ b/packages/core/src/oauth/oauth.test.ts @@ -227,4 +227,29 @@ describe("OAuthService deep-link callback handler", () => { expect(mainWindow.restore).toHaveBeenCalled(); expect(mainWindow.focus).toHaveBeenCalled(); }); + + it("accepts a short-lived impersonated session without a refresh token", async () => { + const { service, getCallbackHandler } = createDeps(); + fetchMock.mockResolvedValue( + jsonResponse({ + access_token: "at", + expires_in: 1800, + token_type: "Bearer", + scope: "", + }), + ); + + const login = service.startFlow("us"); + getCallbackHandler()?.("callback", new URLSearchParams("code=abc")); + + await expect(login).resolves.toEqual({ + success: true, + data: { + access_token: "at", + expires_in: 1800, + token_type: "Bearer", + scope: "", + }, + }); + }); }); diff --git a/packages/core/src/oauth/oauth.ts b/packages/core/src/oauth/oauth.ts index 5149630174..215bf4bb02 100644 --- a/packages/core/src/oauth/oauth.ts +++ b/packages/core/src/oauth/oauth.ts @@ -28,10 +28,10 @@ import { OAUTH_HOST, type OAuthHost } from "./identifiers"; import type { CancelFlowOutput, CloudRegion, - OAuthTokenResponse, RefreshTokenOutput, StartFlowOutput, } from "./schemas"; +import { type OAuthTokenResponse, oAuthTokenResponse } from "./schemas"; const OAUTH_TIMEOUT_MS = 180_000; // 3 minutes const TOKEN_FETCH_TIMEOUT_MS = 30_000; @@ -39,7 +39,6 @@ const DEV_CALLBACK_PORT = 8237; const NETWORK_ERROR_MESSAGE = "Could not connect to PostHog. Please check your internet connection and try again."; - const TOKEN_FETCH_MAX_ATTEMPTS = 3; const TOKEN_FETCH_BACKOFF: BackoffOptions = { initialDelayMs: 1_000, @@ -403,7 +402,15 @@ export class OAuthService { } if (response.ok) { - return (await response.json()) as OAuthTokenResponse; + const tokenResponse = oAuthTokenResponse.safeParse( + await response.json(), + ); + if (!tokenResponse.success) { + throw new Error( + "PostHog returned an invalid authentication response. Please try again.", + ); + } + return tokenResponse.data; } lastError = `Token exchange failed: ${response.status} ${response.statusText}`; diff --git a/packages/ui/src/features/auth/components/AuthScreen.tsx b/packages/ui/src/features/auth/components/AuthScreen.tsx index 1f6dc856f7..1d3d8c75cc 100644 --- a/packages/ui/src/features/auth/components/AuthScreen.tsx +++ b/packages/ui/src/features/auth/components/AuthScreen.tsx @@ -1,9 +1,11 @@ import { happyHog } from "@posthog/ui/assets/hedgehogs"; import { FullScreenLayout } from "@posthog/ui/primitives/FullScreenLayout"; -import { Flex } from "@radix-ui/themes"; +import { Callout, Flex } from "@radix-ui/themes"; +import { useAuthStateValue } from "../store"; import { SignInCard } from "./SignInCard"; export function AuthScreen() { + const sessionEndReason = useAuthStateValue((state) => state.sessionEndReason); return ( @@ -20,6 +22,14 @@ export function AuthScreen() { > + {sessionEndReason === "impersonation_expired" && ( + + + Your impersonated session ended. Impersonate the user + again, then sign in to continue. + + + )} state.currentOrgId); + const sessionType = useAuthStateValue((state) => state.sessionType); + const sessionExpiresAt = useAuthStateValue((state) => state.sessionExpiresAt); const client = useOptionalAuthenticatedClient(); const { data: currentUser } = useCurrentUser({ client }); const selectProjectMutation = useSelectProjectMutation(); @@ -82,6 +84,13 @@ export function ProjectSwitcher() { currentOrgGroup?.orgName ?? currentProject?.organization.name ?? "No organization"; + const impersonationExpiry = + sessionType === "impersonated" && sessionExpiresAt + ? new Date(sessionExpiresAt).toLocaleTimeString([], { + hour: "numeric", + minute: "2-digit", + }) + : null; const projectItems = useMemo( () => (currentOrgGroup?.projects ?? []).map((project) => ({ @@ -187,7 +196,9 @@ export function ProjectSwitcher() { {currentProject?.name ?? "No project selected"} - {currentUser?.email ?? "No email"} + {impersonationExpiry + ? `Impersonating until ${impersonationExpiry}` + : (currentUser?.email ?? "No email")} @@ -219,6 +230,11 @@ export function ProjectSwitcher() { {currentUser.email} + {impersonationExpiry && ( + + Impersonated session ends at {impersonationExpiry} + + )} ) : ( diff --git a/packages/workspace-server/src/services/agent/auth-adapter.test.ts b/packages/workspace-server/src/services/agent/auth-adapter.test.ts index ddff8be973..6028511bf5 100644 --- a/packages/workspace-server/src/services/agent/auth-adapter.test.ts +++ b/packages/workspace-server/src/services/agent/auth-adapter.test.ts @@ -26,9 +26,15 @@ function createDependencies() { accessToken: "fresh-access-token", apiHost: "https://app.posthog.com", }), - getState: vi.fn((): { currentProjectId: number | null } => ({ - currentProjectId: 1, - })), + getState: vi.fn( + (): { + currentProjectId: number | null; + sessionType?: "persistent" | "impersonated" | null; + } => ({ + currentProjectId: 1, + sessionType: "persistent", + }), + ), authenticatedFetch: vi .fn() .mockImplementation( @@ -83,6 +89,8 @@ describe("AgentAuthAdapter", () => { afterEach(() => { vi.restoreAllMocks(); + delete process.env.POSTHOG_API_KEY; + delete process.env.POSTHOG_AUTH_HEADER; }); describe("getCurrentCredentials", () => { @@ -266,6 +274,24 @@ describe("AgentAuthAdapter", () => { expect(process.env.PATH).toBe(pathBefore); }); + it("does not export impersonated credentials to the process environment", async () => { + process.env.POSTHOG_API_KEY = "stale-token"; + process.env.POSTHOG_AUTH_HEADER = "Bearer stale-token"; + deps.authService.getState.mockReturnValue({ + currentProjectId: 1, + sessionType: "impersonated", + }); + + await adapter.configureProcessEnv({ + credentials: baseCredentials, + proxyUrl: "http://127.0.0.1:9999", + claudeCliPath: "/mock/claude-cli.js", + }); + + expect(process.env.POSTHOG_API_KEY).toBeUndefined(); + expect(process.env.POSTHOG_AUTH_HEADER).toBeUndefined(); + }); + it.each([ { rtkEnabled: false, expected: "0" }, { rtkEnabled: true, expected: undefined }, diff --git a/packages/workspace-server/src/services/agent/auth-adapter.ts b/packages/workspace-server/src/services/agent/auth-adapter.ts index 2ddf7c8cd1..2a33faf29d 100644 --- a/packages/workspace-server/src/services/agent/auth-adapter.ts +++ b/packages/workspace-server/src/services/agent/auth-adapter.ts @@ -185,6 +185,11 @@ export class AgentAuthAdapter { } private syncTokenEnvironment(token: string): void { + if (this.authService.getState().sessionType === "impersonated") { + delete process.env.POSTHOG_API_KEY; + delete process.env.POSTHOG_AUTH_HEADER; + return; + } process.env.POSTHOG_API_KEY = token; process.env.POSTHOG_AUTH_HEADER = `Bearer ${token}`; } diff --git a/packages/workspace-server/src/services/agent/ports.ts b/packages/workspace-server/src/services/agent/ports.ts index 76d87364a3..9ed476dd35 100644 --- a/packages/workspace-server/src/services/agent/ports.ts +++ b/packages/workspace-server/src/services/agent/ports.ts @@ -61,12 +61,15 @@ export interface AgentAuth { getValidAccessToken(): Promise<{ accessToken: string; apiHost: string }>; getOAuthCredentials(): Promise<{ access: string; - refresh: string; + refresh: string | null; expires: number; region: CloudRegion; } | null>; refreshAccessToken(): Promise<{ accessToken: string; apiHost: string }>; - getState(): { currentProjectId: number | null }; + getState(): { + currentProjectId: number | null; + sessionType?: "persistent" | "impersonated" | null; + }; authenticatedFetch( fetchImpl: AgentFetchLike, input: string | Request,