From 9bc09aa754c544a6ce12b8305966ce275b6cb93e Mon Sep 17 00:00:00 2001 From: Chris Volzer Date: Wed, 24 Jun 2026 15:51:44 -0400 Subject: [PATCH 1/4] fix(ui): reset MCP connect loading state --- .../mcp-servers/hooks/useMcpServers.test.tsx | 112 ++++++++++++++++++ .../mcp-servers/hooks/useMcpServers.ts | 10 +- 2 files changed, 117 insertions(+), 5 deletions(-) create mode 100644 packages/ui/src/features/mcp-servers/hooks/useMcpServers.test.tsx diff --git a/packages/ui/src/features/mcp-servers/hooks/useMcpServers.test.tsx b/packages/ui/src/features/mcp-servers/hooks/useMcpServers.test.tsx new file mode 100644 index 0000000000..a61175365f --- /dev/null +++ b/packages/ui/src/features/mcp-servers/hooks/useMcpServers.test.tsx @@ -0,0 +1,112 @@ +import type { McpRecommendedServer } from "@posthog/api-client/posthog-client"; +import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; +import { act, renderHook, waitFor } from "@testing-library/react"; +import type { ReactNode } from "react"; +import { beforeEach, describe, expect, it, vi } from "vitest"; + +const mockClient = vi.hoisted(() => ({ + getMcpServerInstallations: vi.fn(), + getMcpServers: vi.fn(), + installMcpTemplate: vi.fn(), + installCustomMcpServer: vi.fn(), + uninstallMcpServer: vi.fn(), + updateMcpServerInstallation: vi.fn(), + authorizeMcpInstallation: vi.fn(), +})); + +const mockTrpcClient = vi.hoisted(() => ({ + mcpCallback: { + getCallbackUrl: { query: vi.fn() }, + openAndWaitForCallback: { mutate: vi.fn() }, + }, +})); + +const mockTrpc = vi.hoisted(() => ({ + mcpCallback: { + onOAuthComplete: { + subscriptionOptions: vi.fn(() => ({})), + }, + }, +})); + +vi.mock("@posthog/ui/features/auth/authClient", () => ({ + useOptionalAuthenticatedClient: () => mockClient, +})); + +vi.mock("@posthog/host-router/react", () => ({ + useHostTRPC: () => mockTrpc, + useHostTRPCClient: () => mockTrpcClient, +})); + +vi.mock("@trpc/tanstack-react-query", () => ({ + useSubscription: vi.fn(), +})); + +vi.mock("sonner", () => ({ + toast: { + error: vi.fn(), + success: vi.fn(), + }, +})); + +import { useMcpServers } from "./useMcpServers"; + +function wrapper({ children }: { children: ReactNode }) { + const queryClient = new QueryClient({ + defaultOptions: { + queries: { retry: false }, + mutations: { retry: false }, + }, + }); + return ( + {children} + ); +} + +const template = { + id: "granola", + name: "Granola", + auth_type: "oauth", +} as McpRecommendedServer; + +describe("useMcpServers", () => { + beforeEach(() => { + vi.clearAllMocks(); + mockClient.getMcpServerInstallations.mockResolvedValue([]); + mockClient.getMcpServers.mockResolvedValue([]); + mockTrpcClient.mcpCallback.getCallbackUrl.query.mockResolvedValue({ + callbackUrl: "posthog-code://mcp-oauth-complete", + }); + }); + + it("reverts template connect loading state after a failed install", async () => { + let rejectInstall!: (error: Error) => void; + mockClient.installMcpTemplate.mockReturnValue( + new Promise((_resolve, reject) => { + rejectInstall = reject; + }), + ); + + const { result } = renderHook(() => useMcpServers(), { wrapper }); + + act(() => { + result.current.installTemplate(template); + }); + + await waitFor(() => expect(result.current.installingId).toBe("granola")); + await waitFor(() => + expect(mockClient.installMcpTemplate).toHaveBeenCalledWith({ + template_id: "granola", + install_source: "posthog-code", + posthog_code_callback_url: "posthog-code://mcp-oauth-complete", + api_key: undefined, + }), + ); + + await act(async () => { + rejectInstall(new Error("Connection failed")); + }); + + await waitFor(() => expect(result.current.installingId).toBeNull()); + }); +}); diff --git a/packages/ui/src/features/mcp-servers/hooks/useMcpServers.ts b/packages/ui/src/features/mcp-servers/hooks/useMcpServers.ts index 50ca0dac80..758651b0e1 100644 --- a/packages/ui/src/features/mcp-servers/hooks/useMcpServers.ts +++ b/packages/ui/src/features/mcp-servers/hooks/useMcpServers.ts @@ -14,7 +14,7 @@ import { useAuthenticatedMutation } from "@posthog/ui/hooks/useAuthenticatedMuta import { useAuthenticatedQuery } from "@posthog/ui/hooks/useAuthenticatedQuery"; import { useQueryClient } from "@tanstack/react-query"; import { useSubscription } from "@trpc/tanstack-react-query"; -import { useCallback, useMemo, useState } from "react"; +import { useCallback, useMemo } from "react"; import { toast } from "sonner"; export const mcpKeys = { @@ -38,7 +38,6 @@ export function useMcpServers() { const trpc = useHostTRPC(); const trpcClient = useHostTRPCClient(); const oauth = useMemo(() => createOAuthCallback(trpcClient), [trpcClient]); - const [installingId, setInstallingId] = useState(null); const queryClient = useQueryClient(); const { data: installations, isLoading: installationsLoading } = @@ -120,18 +119,15 @@ export function useMcpServers() { toast.error(data.error); } invalidateInstallations(); - setInstallingId(null); }, onError: (error: Error) => { toast.error(error.message || "Failed to connect server"); - setInstallingId(null); }, }, ); const installTemplate = useCallback( (template: McpRecommendedServer, opts?: { api_key?: string }) => { - setInstallingId(template.id); installTemplateMutation.mutate({ template_id: template.id, api_key: opts?.api_key, @@ -140,6 +136,10 @@ export function useMcpServers() { [installTemplateMutation], ); + const installingId = installTemplateMutation.isPending + ? (installTemplateMutation.variables?.template_id ?? null) + : null; + const installCustomMutation = useAuthenticatedMutation( ( client, From 9e594d32189629abb10cad396e01de609b834ad6 Mon Sep 17 00:00:00 2001 From: Chris Volzer Date: Wed, 24 Jun 2026 15:57:19 -0400 Subject: [PATCH 2/4] fix(api-client): handle MCP authorize redirects --- packages/api-client/src/fetcher.test.ts | 31 +++++++++++++ packages/api-client/src/fetcher.ts | 15 +++++- .../api-client/src/posthog-client.test.ts | 46 +++++++++++++++++++ packages/api-client/src/posthog-client.ts | 11 +++++ 4 files changed, 101 insertions(+), 2 deletions(-) diff --git a/packages/api-client/src/fetcher.test.ts b/packages/api-client/src/fetcher.test.ts index b75e4bdff2..fe5f44677c 100644 --- a/packages/api-client/src/fetcher.test.ts +++ b/packages/api-client/src/fetcher.test.ts @@ -26,6 +26,16 @@ describe("buildApiFetcher", () => { }; return response; }; + const redirect = (location: string) => + ({ + ok: false, + status: 302, + headers: new Headers({ location }), + json: () => Promise.reject(new Error("redirect has no JSON body")), + clone: () => ({ + text: () => Promise.resolve(""), + }), + }) as Response; beforeEach(() => { vi.resetAllMocks(); @@ -128,6 +138,27 @@ describe("buildApiFetcher", () => { expect(refreshAccessToken).not.toHaveBeenCalled(); }); + it("returns manual redirect responses without treating them as errors", async () => { + const refreshAccessToken = vi.fn().mockResolvedValue("new-token"); + mockFetch.mockResolvedValueOnce(redirect("https://auth.example.com")); + + const fetcher = buildApiFetcher({ + getAccessToken: vi.fn().mockResolvedValue("token"), + refreshAccessToken, + appVersion: "test", + }); + + const response = await fetcher.fetch({ + ...mockInput, + overrides: { redirect: "manual" }, + }); + + expect(response.status).toBe(302); + expect(response.headers.get("location")).toBe("https://auth.example.com"); + expect(refreshAccessToken).not.toHaveBeenCalled(); + expect(mockFetch.mock.calls[0][1].redirect).toBe("manual"); + }); + it("throws when the retry still returns 401", async () => { const getAccessToken = vi.fn().mockResolvedValue("token-1"); const refreshAccessToken = vi.fn().mockResolvedValue("token-2"); diff --git a/packages/api-client/src/fetcher.ts b/packages/api-client/src/fetcher.ts index b9a7a4ed50..99f2874644 100644 --- a/packages/api-client/src/fetcher.ts +++ b/packages/api-client/src/fetcher.ts @@ -10,6 +10,13 @@ export const buildApiFetcher: ( config: ApiFetcherConfig, ) => Parameters[0] = (config) => { const userAgent = `posthog/desktop.hog.dev; version: ${config.appVersion}`; + const isManualRedirectResponse = ( + input: Parameters[0]["fetch"]>[0], + response: Response, + ): boolean => + input.overrides?.redirect === "manual" && + response.status >= 300 && + response.status < 400; const makeRequest = async ( input: Parameters[0]["fetch"]>[0], @@ -78,7 +85,11 @@ export const buildApiFetcher: ( fetch: async (input) => { let response = await makeRequest(input, await config.getAccessToken()); - if (!response.ok && (await isAuthFailure(response))) { + if ( + !response.ok && + !isManualRedirectResponse(input, response) && + (await isAuthFailure(response)) + ) { try { response = await makeRequest( input, @@ -97,7 +108,7 @@ export const buildApiFetcher: ( } } - if (!response.ok) { + if (!response.ok && !isManualRedirectResponse(input, response)) { const cloned = response.clone(); const errorResponse = await response .json() diff --git a/packages/api-client/src/posthog-client.test.ts b/packages/api-client/src/posthog-client.test.ts index cd660438df..a05a5b7109 100644 --- a/packages/api-client/src/posthog-client.test.ts +++ b/packages/api-client/src/posthog-client.test.ts @@ -219,6 +219,52 @@ describe("PostHogAPIClient", () => { ); }); + it("returns the redirect location when authorizing an MCP installation", async () => { + const fetch = vi.fn().mockResolvedValue({ + ok: false, + status: 302, + headers: new Headers({ + location: "https://auth.example.com/authorize?state=abc", + }), + json: async () => { + throw new Error("redirect response should not be parsed as JSON"); + }, + }); + const client = new PostHogAPIClient( + "http://localhost:8000", + async () => "token", + async () => "token", + 123, + ); + + ( + client as unknown as { + api: { baseUrl: string; fetcher: { fetch: typeof fetch } }; + } + ).api = { + baseUrl: "http://localhost:8000", + fetcher: { fetch }, + }; + + await expect( + client.authorizeMcpInstallation({ + installation_id: "inst-123", + install_source: "posthog-code", + posthog_code_callback_url: "posthog-code://mcp-oauth-complete", + }), + ).resolves.toEqual({ + redirect_url: "https://auth.example.com/authorize?state=abc", + }); + + expect(fetch).toHaveBeenCalledWith( + expect.objectContaining({ + method: "get", + path: "/api/environments/123/mcp_server_installations/authorize/", + overrides: { redirect: "manual" }, + }), + ); + }); + describe("warmTask", () => { function makeClient(fetch: ReturnType) { const client = new PostHogAPIClient( diff --git a/packages/api-client/src/posthog-client.ts b/packages/api-client/src/posthog-client.ts index eed207fd55..0a7a0b942a 100644 --- a/packages/api-client/src/posthog-client.ts +++ b/packages/api-client/src/posthog-client.ts @@ -3790,8 +3790,19 @@ export class PostHogAPIClient { method: "get", url, path, + overrides: { redirect: "manual" }, }); + if (response.status >= 300 && response.status < 400) { + const redirectUrl = response.headers.get("location"); + if (!redirectUrl) { + throw new Error( + "Failed to authorize MCP installation: missing redirect location", + ); + } + return { redirect_url: redirectUrl }; + } + if (!response.ok) { const errorData = await response.json().catch(() => ({})); throw new Error( From e138df29ca651d1bc52eebd1e5b1deceabb08506 Mon Sep 17 00:00:00 2001 From: Chris Volzer Date: Wed, 24 Jun 2026 16:05:14 -0400 Subject: [PATCH 3/4] fix(core): tolerate transient preemptive token refresh failures --- packages/core/src/auth/auth.test.ts | 66 +++++++++++++++++++++++++++++ packages/core/src/auth/auth.ts | 29 +++++++++++-- 2 files changed, 91 insertions(+), 4 deletions(-) diff --git a/packages/core/src/auth/auth.test.ts b/packages/core/src/auth/auth.test.ts index 1b9448e04a..80384696ee 100644 --- a/packages/core/src/auth/auth.test.ts +++ b/packages/core/src/auth/auth.test.ts @@ -708,6 +708,72 @@ describe("AuthService", () => { expect(service.getState().status).toBe("restoring"); expect(oauthFlow.refreshToken).toHaveBeenCalledTimes(3); }); + + it("uses the current access token when a preemptive refresh fails before expiry", async () => { + vi.useFakeTimers(); + try { + oauthFlow.startFlow.mockResolvedValue( + mockTokenResponse({ + accessToken: "current-access-token", + refreshToken: "current-refresh-token", + }), + ); + stubAuthFetch(); + + await service.initialize(); + await service.login("us"); + + oauthFlow.refreshToken.mockReset(); + oauthFlow.refreshToken.mockResolvedValue({ + success: false, + error: "Token refresh failed: 500 Internal Server Error", + errorCode: "server_error", + }); + + await vi.advanceTimersByTimeAsync(3_599_500); + + await expect(service.getValidAccessToken()).resolves.toMatchObject({ + accessToken: "current-access-token", + }); + expect(oauthFlow.refreshToken).toHaveBeenCalledTimes(3); + expect(service.getState().status).toBe("authenticated"); + } finally { + vi.useRealTimers(); + } + }); + + it("does not use the current access token when refresh token auth fails", async () => { + vi.useFakeTimers(); + try { + oauthFlow.startFlow.mockResolvedValue( + mockTokenResponse({ + accessToken: "current-access-token", + refreshToken: "current-refresh-token", + }), + ); + stubAuthFetch(); + + await service.initialize(); + await service.login("us"); + + oauthFlow.refreshToken.mockReset(); + oauthFlow.refreshToken.mockResolvedValue({ + success: false, + error: "Token revoked", + errorCode: "auth_error", + }); + + await vi.advanceTimersByTimeAsync(3_599_500); + + await expect(service.getValidAccessToken()).rejects.toThrow( + "Token revoked", + ); + expect(service.getState().status).toBe("anonymous"); + expect(sessionPort.getCurrent()).toBeNull(); + } finally { + vi.useRealTimers(); + } + }); }); describe("transient org fetch failures", () => { diff --git a/packages/core/src/auth/auth.ts b/packages/core/src/auth/auth.ts index b7b47b424c..77e5111734 100644 --- a/packages/core/src/auth/auth.ts +++ b/packages/core/src/auth/auth.ts @@ -487,12 +487,13 @@ export class AuthService extends TypedEventEmitter { private async ensureValidSession( forceRefresh = false, ): Promise { + const currentSession = this.session; if ( - this.session && + currentSession && !forceRefresh && - !this.isSessionExpiring(this.session) + !this.isSessionExpiring(currentSession) ) { - return this.session; + return currentSession; } if (this.refreshPromise) { @@ -502,7 +503,24 @@ export class AuthService extends TypedEventEmitter { const sessionInput = this.getSessionInputForRefresh(); const refreshAndSync = async (): Promise => { - const session = await this.refreshSession(sessionInput); + let session: InMemorySession; + try { + session = await this.refreshSession(sessionInput); + } catch (error) { + if ( + currentSession && + this.session === currentSession && + !forceRefresh && + !this.isSessionExpired(currentSession) + ) { + this.logger.warn( + "Preemptive session refresh failed; using current access token", + { error }, + ); + return currentSession; + } + throw error; + } await this.syncAuthenticatedSession(session); return session; }; @@ -833,6 +851,9 @@ export class AuthService extends TypedEventEmitter { private isSessionExpiring(session: InMemorySession): boolean { return session.accessTokenExpiresAt - Date.now() <= TOKEN_EXPIRY_SKEW_MS; } + private isSessionExpired(session: InMemorySession): boolean { + return session.accessTokenExpiresAt <= Date.now(); + } private async fetchUserContext( accessToken: string, cloudRegion: CloudRegion, From 54ea3d7c5264ec45eb11017178d2df5487184d89 Mon Sep 17 00:00:00 2001 From: Chris Volzer Date: Wed, 24 Jun 2026 16:25:25 -0400 Subject: [PATCH 4/4] fix(api-client): expect JSON MCP authorize redirects --- packages/api-client/src/fetcher.test.ts | 32 ------------------- packages/api-client/src/fetcher.ts | 16 ++-------- .../api-client/src/posthog-client.test.ts | 15 ++++----- packages/api-client/src/posthog-client.ts | 11 ------- 4 files changed, 8 insertions(+), 66 deletions(-) diff --git a/packages/api-client/src/fetcher.test.ts b/packages/api-client/src/fetcher.test.ts index fe5f44677c..6c0352e30f 100644 --- a/packages/api-client/src/fetcher.test.ts +++ b/packages/api-client/src/fetcher.test.ts @@ -26,17 +26,6 @@ describe("buildApiFetcher", () => { }; return response; }; - const redirect = (location: string) => - ({ - ok: false, - status: 302, - headers: new Headers({ location }), - json: () => Promise.reject(new Error("redirect has no JSON body")), - clone: () => ({ - text: () => Promise.resolve(""), - }), - }) as Response; - beforeEach(() => { vi.resetAllMocks(); vi.stubGlobal("fetch", mockFetch); @@ -138,27 +127,6 @@ describe("buildApiFetcher", () => { expect(refreshAccessToken).not.toHaveBeenCalled(); }); - it("returns manual redirect responses without treating them as errors", async () => { - const refreshAccessToken = vi.fn().mockResolvedValue("new-token"); - mockFetch.mockResolvedValueOnce(redirect("https://auth.example.com")); - - const fetcher = buildApiFetcher({ - getAccessToken: vi.fn().mockResolvedValue("token"), - refreshAccessToken, - appVersion: "test", - }); - - const response = await fetcher.fetch({ - ...mockInput, - overrides: { redirect: "manual" }, - }); - - expect(response.status).toBe(302); - expect(response.headers.get("location")).toBe("https://auth.example.com"); - expect(refreshAccessToken).not.toHaveBeenCalled(); - expect(mockFetch.mock.calls[0][1].redirect).toBe("manual"); - }); - it("throws when the retry still returns 401", async () => { const getAccessToken = vi.fn().mockResolvedValue("token-1"); const refreshAccessToken = vi.fn().mockResolvedValue("token-2"); diff --git a/packages/api-client/src/fetcher.ts b/packages/api-client/src/fetcher.ts index 99f2874644..802260973f 100644 --- a/packages/api-client/src/fetcher.ts +++ b/packages/api-client/src/fetcher.ts @@ -10,14 +10,6 @@ export const buildApiFetcher: ( config: ApiFetcherConfig, ) => Parameters[0] = (config) => { const userAgent = `posthog/desktop.hog.dev; version: ${config.appVersion}`; - const isManualRedirectResponse = ( - input: Parameters[0]["fetch"]>[0], - response: Response, - ): boolean => - input.overrides?.redirect === "manual" && - response.status >= 300 && - response.status < 400; - const makeRequest = async ( input: Parameters[0]["fetch"]>[0], token: string, @@ -85,11 +77,7 @@ export const buildApiFetcher: ( fetch: async (input) => { let response = await makeRequest(input, await config.getAccessToken()); - if ( - !response.ok && - !isManualRedirectResponse(input, response) && - (await isAuthFailure(response)) - ) { + if (!response.ok && (await isAuthFailure(response))) { try { response = await makeRequest( input, @@ -108,7 +96,7 @@ export const buildApiFetcher: ( } } - if (!response.ok && !isManualRedirectResponse(input, response)) { + if (!response.ok) { const cloned = response.clone(); const errorResponse = await response .json() diff --git a/packages/api-client/src/posthog-client.test.ts b/packages/api-client/src/posthog-client.test.ts index a05a5b7109..22041f0d4d 100644 --- a/packages/api-client/src/posthog-client.test.ts +++ b/packages/api-client/src/posthog-client.test.ts @@ -219,16 +219,13 @@ describe("PostHogAPIClient", () => { ); }); - it("returns the redirect location when authorizing an MCP installation", async () => { + it("returns the redirect URL when authorizing an MCP installation", async () => { const fetch = vi.fn().mockResolvedValue({ - ok: false, - status: 302, - headers: new Headers({ - location: "https://auth.example.com/authorize?state=abc", + ok: true, + status: 200, + json: async () => ({ + redirect_url: "https://auth.example.com/authorize?state=abc", }), - json: async () => { - throw new Error("redirect response should not be parsed as JSON"); - }, }); const client = new PostHogAPIClient( "http://localhost:8000", @@ -260,9 +257,9 @@ describe("PostHogAPIClient", () => { expect.objectContaining({ method: "get", path: "/api/environments/123/mcp_server_installations/authorize/", - overrides: { redirect: "manual" }, }), ); + expect(fetch.mock.calls[0][0]).not.toHaveProperty("overrides"); }); describe("warmTask", () => { diff --git a/packages/api-client/src/posthog-client.ts b/packages/api-client/src/posthog-client.ts index 0a7a0b942a..eed207fd55 100644 --- a/packages/api-client/src/posthog-client.ts +++ b/packages/api-client/src/posthog-client.ts @@ -3790,19 +3790,8 @@ export class PostHogAPIClient { method: "get", url, path, - overrides: { redirect: "manual" }, }); - if (response.status >= 300 && response.status < 400) { - const redirectUrl = response.headers.get("location"); - if (!redirectUrl) { - throw new Error( - "Failed to authorize MCP installation: missing redirect location", - ); - } - return { redirect_url: redirectUrl }; - } - if (!response.ok) { const errorData = await response.json().catch(() => ({})); throw new Error(