Skip to content

Commit ac66d2d

Browse files
cvolzer3cvolzer3
authored
fix(mcp-store): fix MCP connect button loading state (#2917)
1 parent d3ffcd3 commit ac66d2d

5 files changed

Lines changed: 251 additions & 9 deletions

File tree

‎packages/api-client/src/posthog-client.test.ts‎

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -418,6 +418,49 @@ describe("PostHogAPIClient", () => {
418418
);
419419
});
420420

421+
it("returns the redirect URL when authorizing an MCP installation", async () => {
422+
const fetch = vi.fn().mockResolvedValue({
423+
ok: true,
424+
status: 200,
425+
json: async () => ({
426+
redirect_url: "https://auth.example.com/authorize?state=abc",
427+
}),
428+
});
429+
const client = new PostHogAPIClient(
430+
"http://localhost:8000",
431+
async () => "token",
432+
async () => "token",
433+
123,
434+
);
435+
436+
(
437+
client as unknown as {
438+
api: { baseUrl: string; fetcher: { fetch: typeof fetch } };
439+
}
440+
).api = {
441+
baseUrl: "http://localhost:8000",
442+
fetcher: { fetch },
443+
};
444+
445+
await expect(
446+
client.authorizeMcpInstallation({
447+
installation_id: "inst-123",
448+
install_source: "posthog-code",
449+
posthog_code_callback_url: "posthog-code://mcp-oauth-complete",
450+
}),
451+
).resolves.toEqual({
452+
redirect_url: "https://auth.example.com/authorize?state=abc",
453+
});
454+
455+
expect(fetch).toHaveBeenCalledWith(
456+
expect.objectContaining({
457+
method: "get",
458+
path: "/api/environments/123/mcp_server_installations/authorize/",
459+
}),
460+
);
461+
expect(fetch.mock.calls[0][0]).not.toHaveProperty("overrides");
462+
});
463+
421464
describe("warmTask", () => {
422465
function makeClient(fetch: ReturnType<typeof vi.fn>) {
423466
const client = new PostHogAPIClient(

‎packages/core/src/auth/auth.test.ts‎

Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -715,6 +715,72 @@ describe("AuthService", () => {
715715
expect(service.getState().status).toBe("restoring");
716716
expect(oauthFlow.refreshToken).toHaveBeenCalledTimes(3);
717717
});
718+
719+
it("uses the current access token when a preemptive refresh fails before expiry", async () => {
720+
vi.useFakeTimers();
721+
try {
722+
oauthFlow.startFlow.mockResolvedValue(
723+
mockTokenResponse({
724+
accessToken: "current-access-token",
725+
refreshToken: "current-refresh-token",
726+
}),
727+
);
728+
stubAuthFetch();
729+
730+
await service.initialize();
731+
await service.login("us");
732+
733+
oauthFlow.refreshToken.mockReset();
734+
oauthFlow.refreshToken.mockResolvedValue({
735+
success: false,
736+
error: "Token refresh failed: 500 Internal Server Error",
737+
errorCode: "server_error",
738+
});
739+
740+
await vi.advanceTimersByTimeAsync(3_599_500);
741+
742+
await expect(service.getValidAccessToken()).resolves.toMatchObject({
743+
accessToken: "current-access-token",
744+
});
745+
expect(oauthFlow.refreshToken).toHaveBeenCalledTimes(3);
746+
expect(service.getState().status).toBe("authenticated");
747+
} finally {
748+
vi.useRealTimers();
749+
}
750+
});
751+
752+
it("does not use the current access token when refresh token auth fails", async () => {
753+
vi.useFakeTimers();
754+
try {
755+
oauthFlow.startFlow.mockResolvedValue(
756+
mockTokenResponse({
757+
accessToken: "current-access-token",
758+
refreshToken: "current-refresh-token",
759+
}),
760+
);
761+
stubAuthFetch();
762+
763+
await service.initialize();
764+
await service.login("us");
765+
766+
oauthFlow.refreshToken.mockReset();
767+
oauthFlow.refreshToken.mockResolvedValue({
768+
success: false,
769+
error: "Token revoked",
770+
errorCode: "auth_error",
771+
});
772+
773+
await vi.advanceTimersByTimeAsync(3_599_500);
774+
775+
await expect(service.getValidAccessToken()).rejects.toThrow(
776+
"Token revoked",
777+
);
778+
expect(service.getState().status).toBe("anonymous");
779+
expect(sessionPort.getCurrent()).toBeNull();
780+
} finally {
781+
vi.useRealTimers();
782+
}
783+
});
718784
});
719785

720786
describe("transient org fetch failures", () => {

‎packages/core/src/auth/auth.ts‎

Lines changed: 25 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -487,12 +487,13 @@ export class AuthService extends TypedEventEmitter<AuthServiceEvents> {
487487
private async ensureValidSession(
488488
forceRefresh = false,
489489
): Promise<InMemorySession> {
490+
const currentSession = this.session;
490491
if (
491-
this.session &&
492+
currentSession &&
492493
!forceRefresh &&
493-
!this.isSessionExpiring(this.session)
494+
!this.isSessionExpiring(currentSession)
494495
) {
495-
return this.session;
496+
return currentSession;
496497
}
497498

498499
if (this.refreshPromise) {
@@ -502,7 +503,24 @@ export class AuthService extends TypedEventEmitter<AuthServiceEvents> {
502503
const sessionInput = this.getSessionInputForRefresh();
503504

504505
const refreshAndSync = async (): Promise<InMemorySession> => {
505-
const session = await this.refreshSession(sessionInput);
506+
let session: InMemorySession;
507+
try {
508+
session = await this.refreshSession(sessionInput);
509+
} catch (error) {
510+
if (
511+
currentSession &&
512+
this.session === currentSession &&
513+
!forceRefresh &&
514+
!this.isSessionExpired(currentSession)
515+
) {
516+
this.logger.warn(
517+
"Preemptive session refresh failed; using current access token",
518+
{ error },
519+
);
520+
return currentSession;
521+
}
522+
throw error;
523+
}
506524
await this.syncAuthenticatedSession(session);
507525
return session;
508526
};
@@ -843,6 +861,9 @@ export class AuthService extends TypedEventEmitter<AuthServiceEvents> {
843861
private isSessionExpiring(session: InMemorySession): boolean {
844862
return session.accessTokenExpiresAt - Date.now() <= TOKEN_EXPIRY_SKEW_MS;
845863
}
864+
private isSessionExpired(session: InMemorySession): boolean {
865+
return session.accessTokenExpiresAt <= Date.now();
866+
}
846867
private async fetchUserContext(
847868
accessToken: string,
848869
cloudRegion: CloudRegion,

‎packages/ui/src/features/mcp-servers/hooks/useMcpServers.test.tsx‎

Lines changed: 112 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,112 @@
1+
import type { McpRecommendedServer } from "@posthog/api-client/posthog-client";
2+
import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
3+
import { act, renderHook, waitFor } from "@testing-library/react";
4+
import type { ReactNode } from "react";
5+
import { beforeEach, describe, expect, it, vi } from "vitest";
6+
7+
const mockClient = vi.hoisted(() => ({
8+
getMcpServerInstallations: vi.fn(),
9+
getMcpServers: vi.fn(),
10+
installMcpTemplate: vi.fn(),
11+
installCustomMcpServer: vi.fn(),
12+
uninstallMcpServer: vi.fn(),
13+
updateMcpServerInstallation: vi.fn(),
14+
authorizeMcpInstallation: vi.fn(),
15+
}));
16+
17+
const mockTrpcClient = vi.hoisted(() => ({
18+
mcpCallback: {
19+
getCallbackUrl: { query: vi.fn() },
20+
openAndWaitForCallback: { mutate: vi.fn() },
21+
},
22+
}));
23+
24+
const mockTrpc = vi.hoisted(() => ({
25+
mcpCallback: {
26+
onOAuthComplete: {
27+
subscriptionOptions: vi.fn(() => ({})),
28+
},
29+
},
30+
}));
31+
32+
vi.mock("@posthog/ui/features/auth/authClient", () => ({
33+
useOptionalAuthenticatedClient: () => mockClient,
34+
}));
35+
36+
vi.mock("@posthog/host-router/react", () => ({
37+
useHostTRPC: () => mockTrpc,
38+
useHostTRPCClient: () => mockTrpcClient,
39+
}));
40+
41+
vi.mock("@trpc/tanstack-react-query", () => ({
42+
useSubscription: vi.fn(),
43+
}));
44+
45+
vi.mock("@posthog/ui/primitives/toast", () => ({
46+
toast: {
47+
error: vi.fn(),
48+
success: vi.fn(),
49+
},
50+
}));
51+
52+
import { useMcpServers } from "./useMcpServers";
53+
54+
function wrapper({ children }: { children: ReactNode }) {
55+
const queryClient = new QueryClient({
56+
defaultOptions: {
57+
queries: { retry: false },
58+
mutations: { retry: false },
59+
},
60+
});
61+
return (
62+
<QueryClientProvider client={queryClient}>{children}</QueryClientProvider>
63+
);
64+
}
65+
66+
const template = {
67+
id: "granola",
68+
name: "Granola",
69+
auth_type: "oauth",
70+
} as McpRecommendedServer;
71+
72+
describe("useMcpServers", () => {
73+
beforeEach(() => {
74+
vi.clearAllMocks();
75+
mockClient.getMcpServerInstallations.mockResolvedValue([]);
76+
mockClient.getMcpServers.mockResolvedValue([]);
77+
mockTrpcClient.mcpCallback.getCallbackUrl.query.mockResolvedValue({
78+
callbackUrl: "posthog-code://mcp-oauth-complete",
79+
});
80+
});
81+
82+
it("reverts template connect loading state after a failed install", async () => {
83+
let rejectInstall!: (error: Error) => void;
84+
mockClient.installMcpTemplate.mockReturnValue(
85+
new Promise((_resolve, reject) => {
86+
rejectInstall = reject;
87+
}),
88+
);
89+
90+
const { result } = renderHook(() => useMcpServers(), { wrapper });
91+
92+
act(() => {
93+
result.current.installTemplate(template);
94+
});
95+
96+
await waitFor(() => expect(result.current.installingId).toBe("granola"));
97+
await waitFor(() =>
98+
expect(mockClient.installMcpTemplate).toHaveBeenCalledWith({
99+
template_id: "granola",
100+
install_source: "posthog-code",
101+
posthog_code_callback_url: "posthog-code://mcp-oauth-complete",
102+
api_key: undefined,
103+
}),
104+
);
105+
106+
await act(async () => {
107+
rejectInstall(new Error("Connection failed"));
108+
});
109+
110+
await waitFor(() => expect(result.current.installingId).toBeNull());
111+
});
112+
});

‎packages/ui/src/features/mcp-servers/hooks/useMcpServers.ts‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@ import { useAuthenticatedQuery } from "@posthog/ui/hooks/useAuthenticatedQuery";
1818
import { toast } from "@posthog/ui/primitives/toast";
1919
import { useQueryClient } from "@tanstack/react-query";
2020
import { useSubscription } from "@trpc/tanstack-react-query";
21-
import { useCallback, useMemo, useState } from "react";
21+
import { useCallback, useMemo } from "react";
2222

2323
// `mcpKeys` + `createOAuthCallback` now live in the shared mcp-server-manager
2424
// module (also used by the agent-applications builder). Re-exported here so
@@ -29,7 +29,6 @@ export function useMcpServers() {
2929
const trpc = useHostTRPC();
3030
const trpcClient = useHostTRPCClient();
3131
const oauth = useMemo(() => createOAuthCallback(trpcClient), [trpcClient]);
32-
const [installingId, setInstallingId] = useState<string | null>(null);
3332
const queryClient = useQueryClient();
3433

3534
const { data: installations, isLoading: installationsLoading } =
@@ -111,18 +110,15 @@ export function useMcpServers() {
111110
toast.error(data.error);
112111
}
113112
invalidateInstallations();
114-
setInstallingId(null);
115113
},
116114
onError: (error: Error) => {
117115
toast.error(error.message || "Failed to connect server");
118-
setInstallingId(null);
119116
},
120117
},
121118
);
122119

123120
const installTemplate = useCallback(
124121
(template: McpRecommendedServer, opts?: { api_key?: string }) => {
125-
setInstallingId(template.id);
126122
installTemplateMutation.mutate({
127123
template_id: template.id,
128124
api_key: opts?.api_key,
@@ -131,6 +127,10 @@ export function useMcpServers() {
131127
[installTemplateMutation],
132128
);
133129

130+
const installingId = installTemplateMutation.isPending
131+
? (installTemplateMutation.variables?.template_id ?? null)
132+
: null;
133+
134134
const installCustomMutation = useAuthenticatedMutation(
135135
(
136136
client,

0 commit comments

Comments
 (0)