@assistant-ui/react-mcp 0.1.15 → 0.1.17

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (45) hide show
  1. package/dist/auth/createOAuthProvider.d.ts +19 -2
  2. package/dist/auth/createOAuthProvider.d.ts.map +1 -1
  3. package/dist/auth/createOAuthProvider.js +197 -35
  4. package/dist/auth/createOAuthProvider.js.map +1 -1
  5. package/dist/auth/types.d.ts +6 -1
  6. package/dist/auth/types.d.ts.map +1 -1
  7. package/dist/hooks/useMcpOAuthCallback.d.ts.map +1 -1
  8. package/dist/hooks/useMcpOAuthCallback.js +5 -6
  9. package/dist/hooks/useMcpOAuthCallback.js.map +1 -1
  10. package/dist/resources/McpManagerResource.d.ts.map +1 -1
  11. package/dist/resources/McpManagerResource.js +2 -1
  12. package/dist/resources/McpManagerResource.js.map +1 -1
  13. package/dist/resources/McpServerResource.d.ts +2 -1
  14. package/dist/resources/McpServerResource.d.ts.map +1 -1
  15. package/dist/resources/McpServerResource.js +79 -19
  16. package/dist/resources/McpServerResource.js.map +1 -1
  17. package/dist/resources/storage/McpLocalStorage.d.ts +7 -0
  18. package/dist/resources/storage/McpLocalStorage.d.ts.map +1 -1
  19. package/dist/resources/storage/McpLocalStorage.js +146 -37
  20. package/dist/resources/storage/McpLocalStorage.js.map +1 -1
  21. package/dist/resources/storage/McpMemoryStorage.d.ts.map +1 -1
  22. package/dist/resources/storage/McpMemoryStorage.js +9 -6
  23. package/dist/resources/storage/McpMemoryStorage.js.map +1 -1
  24. package/dist/resources/storage/types.d.ts +12 -0
  25. package/dist/resources/storage/types.d.ts.map +1 -1
  26. package/dist/utils/serverUrl.d.ts +8 -0
  27. package/dist/utils/serverUrl.d.ts.map +1 -0
  28. package/dist/utils/serverUrl.js +15 -0
  29. package/dist/utils/serverUrl.js.map +1 -0
  30. package/package.json +7 -7
  31. package/src/auth/createOAuthProvider.test.ts +919 -5
  32. package/src/auth/createOAuthProvider.ts +328 -42
  33. package/src/auth/types.ts +6 -1
  34. package/src/hooks/useMcpOAuthCallback.test.ts +74 -1
  35. package/src/hooks/useMcpOAuthCallback.tsx +11 -8
  36. package/src/resources/McpManagerResource.ts +5 -1
  37. package/src/resources/McpServerResource.test.ts +612 -16
  38. package/src/resources/McpServerResource.ts +95 -23
  39. package/src/resources/storage/McpLocalStorage.test.ts +97 -1
  40. package/src/resources/storage/McpLocalStorage.ts +90 -47
  41. package/src/resources/storage/McpMemoryStorage.test.ts +99 -0
  42. package/src/resources/storage/McpMemoryStorage.ts +23 -17
  43. package/src/resources/storage/types.ts +12 -0
  44. package/src/utils/serverUrl.test.ts +66 -0
  45. package/src/utils/serverUrl.ts +23 -0
@@ -3,7 +3,9 @@ import type { ClientOutput } from "@assistant-ui/store";
3
3
  import { useEffect, useState } from "react";
4
4
  import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
5
5
  import type { MCPAuthConfig } from "../mcp-scope";
6
+ import type { MCPPersistedAuthState } from "../auth/types";
6
7
  import type { MCPStorage } from "./storage/types";
8
+ import type { McpServerResourceProps } from "./McpServerResource";
7
9
 
8
10
  const mocks = vi.hoisted(() => {
9
11
  const clients: any[] = [];
@@ -70,7 +72,8 @@ vi.mock("@modelcontextprotocol/client", async (importOriginal) => ({
70
72
  StreamableHTTPClientTransport: mocks.StreamableHTTPClientTransport,
71
73
  }));
72
74
 
73
- const { McpServerResource } = await import("./McpServerResource");
75
+ const { McpServerResource, getConnectionDependencies } =
76
+ await import("./McpServerResource");
74
77
 
75
78
  const never = <T>() => new Promise<T>(() => {});
76
79
 
@@ -187,9 +190,191 @@ const mount = (
187
190
  });
188
191
  };
189
192
 
193
+ const unboundAuthMessage =
194
+ 'MCP server "docs" has saved authentication for a different URL. Authenticate again to connect to https://example.com/mcp.';
195
+
190
196
  describe("McpServerResource automatic authentication", () => {
191
197
  beforeEach(resetMocks);
192
198
 
199
+ it("does not auto-connect with authentication from another server URL", async () => {
200
+ const storage = createStorage();
201
+ vi.mocked(storage.loadAuthState).mockResolvedValue({
202
+ serverUrl: "https://other.example.com/mcp",
203
+ token: "secret",
204
+ });
205
+ const root = mount({
206
+ auth: { type: "bearer" },
207
+ storage,
208
+ autoConnect: true,
209
+ });
210
+
211
+ try {
212
+ await waitForResourceUpdate(
213
+ () => root.getValue().getState().lastError !== null,
214
+ );
215
+
216
+ expect(root.getValue().getState()).toMatchObject({
217
+ connectionState: "disconnected",
218
+ lastError: { message: unboundAuthMessage },
219
+ });
220
+ expect(mocks.StreamableHTTPClientTransport).not.toHaveBeenCalled();
221
+ } finally {
222
+ root.unmount();
223
+ }
224
+ });
225
+
226
+ it("does not auto-connect with an unbound legacy bearer token", async () => {
227
+ const storage = createStorage();
228
+ vi.mocked(storage.loadAuthState).mockResolvedValue({ token: "secret" });
229
+ const root = mount({
230
+ auth: { type: "bearer" },
231
+ storage,
232
+ autoConnect: true,
233
+ });
234
+
235
+ try {
236
+ await waitForResourceUpdate(
237
+ () => root.getValue().getState().lastError !== null,
238
+ );
239
+
240
+ expect(root.getValue().getState()).toMatchObject({
241
+ connectionState: "disconnected",
242
+ lastError: { message: unboundAuthMessage },
243
+ });
244
+ expect(mocks.StreamableHTTPClientTransport).not.toHaveBeenCalled();
245
+ } finally {
246
+ root.unmount();
247
+ }
248
+ });
249
+
250
+ it("does not auto-connect with unbound legacy OAuth tokens", async () => {
251
+ const storage = createStorage();
252
+ vi.mocked(storage.loadAuthState).mockResolvedValue({
253
+ tokens: { access_token: "secret", token_type: "bearer" },
254
+ });
255
+ const root = mount({
256
+ auth: { type: "oauth" },
257
+ storage,
258
+ autoConnect: true,
259
+ });
260
+
261
+ try {
262
+ await waitForResourceUpdate(
263
+ () => root.getValue().getState().lastError !== null,
264
+ );
265
+
266
+ expect(root.getValue().getState()).toMatchObject({
267
+ connectionState: "disconnected",
268
+ lastError: { message: unboundAuthMessage },
269
+ });
270
+ expect(mocks.StreamableHTTPClientTransport).not.toHaveBeenCalled();
271
+ } finally {
272
+ root.unmount();
273
+ }
274
+ });
275
+
276
+ it("does not auto-connect OAuth tokens for a different client", async () => {
277
+ const storage = createStorage();
278
+ vi.mocked(storage.loadAuthState).mockResolvedValue({
279
+ serverUrl: "https://example.com/mcp",
280
+ clientInformation: {
281
+ client_id: "client-a",
282
+ redirect_uris: ["https://example.com/callback"],
283
+ },
284
+ clientInformationSource: "registered",
285
+ tokens: { access_token: "secret", token_type: "bearer" },
286
+ tokensClientId: "client-a",
287
+ });
288
+ const root = mount({
289
+ auth: { type: "oauth", clientId: "client-b" },
290
+ storage,
291
+ autoConnect: true,
292
+ });
293
+
294
+ try {
295
+ await waitFor(() => storage.loadAuthState.mock.calls.length > 0);
296
+ await flushMacrotask();
297
+
298
+ expect(root.getValue().getState()).toMatchObject({
299
+ connectionState: "disconnected",
300
+ lastError: null,
301
+ });
302
+ expect(mocks.StreamableHTTPClientTransport).not.toHaveBeenCalled();
303
+ } finally {
304
+ root.unmount();
305
+ }
306
+ });
307
+
308
+ it("keeps a static bearer token usable when the saved record is unbound", async () => {
309
+ const storage = createStorage();
310
+ vi.mocked(storage.loadAuthState).mockResolvedValue({ token: "stale" });
311
+ const root = mount({
312
+ auth: { type: "bearer", token: "static" },
313
+ storage,
314
+ autoConnect: true,
315
+ });
316
+
317
+ try {
318
+ await waitFor(() => mocks.transports.length > 0);
319
+ await flushMacrotask();
320
+
321
+ expect(mocks.StreamableHTTPClientTransport).toHaveBeenCalledWith(
322
+ new URL("https://example.com/mcp"),
323
+ { requestInit: { headers: { Authorization: "Bearer static" } } },
324
+ );
325
+ expect(root.getValue().getState().lastError).toBeNull();
326
+ } finally {
327
+ root.unmount();
328
+ }
329
+ });
330
+
331
+ it("does not report a record that holds no credentials", async () => {
332
+ const storage = createStorage();
333
+ vi.mocked(storage.loadAuthState).mockResolvedValue({});
334
+ const root = mount({
335
+ auth: { type: "bearer" },
336
+ storage,
337
+ autoConnect: true,
338
+ });
339
+
340
+ try {
341
+ await waitFor(() => storage.loadAuthState.mock.calls.length > 0);
342
+ await flushMacrotask();
343
+
344
+ expect(root.getValue().getState()).toMatchObject({
345
+ connectionState: "disconnected",
346
+ lastError: null,
347
+ });
348
+ expect(mocks.StreamableHTTPClientTransport).not.toHaveBeenCalled();
349
+ } finally {
350
+ root.unmount();
351
+ }
352
+ });
353
+
354
+ it("reports an unbound bearer record on a manual connect", async () => {
355
+ const storage = createStorage();
356
+ vi.mocked(storage.loadAuthState).mockResolvedValue({
357
+ serverUrl: "https://other.example.com/mcp",
358
+ token: "secret",
359
+ });
360
+ const root = mount({ auth: { type: "bearer" }, storage });
361
+
362
+ try {
363
+ await expect(root.getValue().connect()).resolves.toBeUndefined();
364
+ await waitForResourceUpdate(
365
+ () => root.getValue().getState().connectionState === "error",
366
+ );
367
+
368
+ expect(root.getValue().getState()).toMatchObject({
369
+ connectionState: "error",
370
+ lastError: { message: unboundAuthMessage },
371
+ });
372
+ expect(mocks.StreamableHTTPClientTransport).not.toHaveBeenCalled();
373
+ } finally {
374
+ root.unmount();
375
+ }
376
+ });
377
+
193
378
  it("reports auth storage load failures", async () => {
194
379
  const storage = createStorage();
195
380
  vi.mocked(storage.loadAuthState).mockRejectedValue(
@@ -559,14 +744,130 @@ describe("McpServerResource connection lifecycle", () => {
559
744
  describe("McpServerResource completeAuth", () => {
560
745
  beforeEach(resetMocks);
561
746
 
747
+ it("lets callback validation win over mount-time auto-connect", async () => {
748
+ const pendingLoads: Array<(value: MCPPersistedAuthState | null) => void> =
749
+ [];
750
+ const storage = createStorage();
751
+ vi.mocked(storage.loadAuthState).mockImplementation(
752
+ () =>
753
+ new Promise((resolve) => {
754
+ pendingLoads.push(resolve);
755
+ }),
756
+ );
757
+ const root = mount({
758
+ auth: { type: "oauth" },
759
+ storage,
760
+ autoConnect: true,
761
+ });
762
+
763
+ try {
764
+ await waitFor(() => pendingLoads.length > 1);
765
+ const callbackLoadIndex = pendingLoads.length;
766
+ const completeAuth = root
767
+ .getValue()
768
+ .completeAuth("https://example.com/callback?code=abc&state=expected");
769
+ await waitFor(() => pendingLoads.length > callbackLoadIndex);
770
+
771
+ for (const resolve of pendingLoads.slice(0, callbackLoadIndex)) {
772
+ resolve({
773
+ serverUrl: "https://example.com/mcp",
774
+ clientInformation: {
775
+ client_id: "registered-client",
776
+ redirect_uris: ["https://example.com/callback"],
777
+ },
778
+ clientInformationSource: "registered",
779
+ tokens: { access_token: "persisted", token_type: "bearer" },
780
+ tokensClientId: "registered-client",
781
+ });
782
+ }
783
+ await flushMacrotask();
784
+
785
+ expect(mocks.StreamableHTTPClientTransport).not.toHaveBeenCalled();
786
+
787
+ pendingLoads[callbackLoadIndex]!({
788
+ serverUrl: "https://example.com/mcp",
789
+ state: "expected",
790
+ });
791
+ await expect(completeAuth).resolves.toBeUndefined();
792
+ await flushMacrotask();
793
+
794
+ expect(mocks.transports).toHaveLength(1);
795
+ expect(mocks.transports[0].finishAuth).toHaveBeenCalledTimes(1);
796
+ expect(root.getValue().getState().connectionState).toBe("connected");
797
+ } finally {
798
+ root.unmount();
799
+ }
800
+ });
801
+
802
+ it("resumes auto-connect when callback validation fails", async () => {
803
+ const pendingLoads: Array<(value: MCPPersistedAuthState | null) => void> =
804
+ [];
805
+ const storage = createStorage();
806
+ vi.mocked(storage.loadAuthState).mockImplementation(
807
+ () =>
808
+ new Promise((resolve) => {
809
+ pendingLoads.push(resolve);
810
+ }),
811
+ );
812
+ const root = mount({
813
+ auth: { type: "oauth" },
814
+ storage,
815
+ autoConnect: true,
816
+ });
817
+
818
+ try {
819
+ await waitFor(() => pendingLoads.length > 1);
820
+ const callbackLoadIndex = pendingLoads.length;
821
+ const completeAuth = root
822
+ .getValue()
823
+ .completeAuth("https://example.com/callback?code=abc&state=expected");
824
+ await waitFor(() => pendingLoads.length > callbackLoadIndex);
825
+
826
+ for (const resolve of pendingLoads.slice(0, callbackLoadIndex)) {
827
+ resolve({
828
+ serverUrl: "https://example.com/mcp",
829
+ clientInformation: {
830
+ client_id: "registered-client",
831
+ redirect_uris: ["https://example.com/callback"],
832
+ },
833
+ clientInformationSource: "registered",
834
+ tokens: { access_token: "persisted", token_type: "bearer" },
835
+ tokensClientId: "registered-client",
836
+ });
837
+ }
838
+ await flushMacrotask();
839
+
840
+ expect(mocks.StreamableHTTPClientTransport).not.toHaveBeenCalled();
841
+
842
+ pendingLoads[callbackLoadIndex]!({
843
+ serverUrl: "https://example.com/mcp",
844
+ state: "different",
845
+ });
846
+ await expect(completeAuth).rejects.toThrow(
847
+ "OAuth state does not match the authorization request",
848
+ );
849
+ await waitFor(() => mocks.transports.length === 1);
850
+ await waitForResourceUpdate(
851
+ () => root.getValue().getState().connectionState === "connected",
852
+ );
853
+ } finally {
854
+ root.unmount();
855
+ }
856
+ });
857
+
562
858
  it("completes auth across the StrictMode effect replay", async () => {
859
+ const storage = createStorage();
860
+ vi.mocked(storage.loadAuthState).mockResolvedValue({
861
+ serverUrl: "https://example.com/mcp",
862
+ state: "aui-mcp:ZG9jcw.nonce",
863
+ });
563
864
  let completeAuth: Promise<void> | undefined;
564
865
  let started = false;
565
- const root = mount({ auth: { type: "oauth" } }, (server) => {
866
+ const root = mount({ auth: { type: "oauth" }, storage }, (server) => {
566
867
  if (started) return;
567
868
  started = true;
568
869
  completeAuth = server.completeAuth(
569
- "https://example.com/callback?code=abc",
870
+ "https://example.com/callback?code=abc&state=aui-mcp%3AZG9jcw.nonce&iss=https%3A%2F%2Fauth.example.com",
570
871
  );
571
872
  });
572
873
 
@@ -575,26 +876,162 @@ describe("McpServerResource completeAuth", () => {
575
876
  await flushMacrotask();
576
877
 
577
878
  expect(mocks.transports[0].finishAuth).toHaveBeenCalledTimes(1);
879
+ const params = mocks.transports[0].finishAuth.mock.calls[0][0];
880
+ expect(params).toBeInstanceOf(URLSearchParams);
881
+ expect(params.get("code")).toBe("abc");
882
+ expect(params.get("iss")).toBe("https://auth.example.com");
578
883
  expect(root.getValue().getState().connectionState).toBe("connected");
579
884
  } finally {
580
885
  root.unmount();
581
886
  }
582
887
  });
583
888
 
584
- it("rejects when the callback URL has no authorization code", async () => {
585
- const root = mount();
889
+ it("rejects before transport setup without a usable authorization code", async () => {
890
+ const storage = createStorage();
891
+ vi.mocked(storage.loadAuthState).mockResolvedValue({
892
+ serverUrl: "https://example.com/mcp",
893
+ state: "abc",
894
+ });
895
+ const root = mount({ auth: { type: "oauth" }, storage });
586
896
 
587
897
  try {
588
898
  await expect(
589
- root.getValue().completeAuth("https://example.com/callback?state=abc"),
899
+ root
900
+ .getValue()
901
+ .completeAuth("https://example.com/callback?state=abc&code="),
590
902
  ).rejects.toThrow("missing authorization code in callback URL");
591
903
  await flushMacrotask();
592
904
 
905
+ expect(root.getValue().getState()).toMatchObject({
906
+ connectionState: "disconnected",
907
+ lastError: null,
908
+ });
909
+ expect(mocks.transports).toHaveLength(0);
910
+ } finally {
911
+ root.unmount();
912
+ }
913
+ });
914
+
915
+ it("forwards OAuth error callbacks to the transport", async () => {
916
+ const storage = createStorage();
917
+ vi.mocked(storage.loadAuthState).mockResolvedValue({
918
+ serverUrl: "https://example.com/mcp",
919
+ state: "expected",
920
+ });
921
+ mocks.finishAuthResults.push(() =>
922
+ Promise.reject(new Error("access_denied: Denied")),
923
+ );
924
+ const root = mount({ auth: { type: "oauth" }, storage });
925
+
926
+ try {
927
+ await expect(
928
+ root
929
+ .getValue()
930
+ .completeAuth(
931
+ "https://example.com/callback?error=access_denied&error_description=Denied&state=expected&iss=https%3A%2F%2Fauth.example.com",
932
+ ),
933
+ ).rejects.toThrow("access_denied: Denied");
934
+ await flushMacrotask();
935
+
936
+ const params = mocks.transports[0].finishAuth.mock.calls[0][0];
937
+ expect(params).toBeInstanceOf(URLSearchParams);
938
+ expect(params.get("error")).toBe("access_denied");
939
+ expect(params.get("error_description")).toBe("Denied");
940
+ expect(params.get("iss")).toBe("https://auth.example.com");
593
941
  expect(root.getValue().getState()).toMatchObject({
594
942
  connectionState: "error",
595
- lastError: {
596
- message: "missing authorization code in callback URL",
597
- },
943
+ lastError: { message: "access_denied: Denied" },
944
+ });
945
+ } finally {
946
+ root.unmount();
947
+ }
948
+ });
949
+
950
+ it("rejects callbacks without a pending authorization request", async () => {
951
+ const storage = createStorage();
952
+ vi.mocked(storage.loadAuthState).mockResolvedValue(null);
953
+ const root = mount({ auth: { type: "oauth" }, storage });
954
+
955
+ try {
956
+ await expect(
957
+ root
958
+ .getValue()
959
+ .completeAuth("https://example.com/callback?code=abc&state=expected"),
960
+ ).rejects.toThrow(
961
+ "no pending OAuth authorization request for this server",
962
+ );
963
+
964
+ expect(mocks.transports).toHaveLength(0);
965
+ } finally {
966
+ root.unmount();
967
+ }
968
+ });
969
+
970
+ it("rejects callbacks whose state does not match the authorization request", async () => {
971
+ const storage = createStorage();
972
+ vi.mocked(storage.loadAuthState).mockResolvedValue({
973
+ serverUrl: "https://example.com/mcp",
974
+ state: "aui-mcp:ZG9jcw.expected",
975
+ });
976
+ const root = mount({ auth: { type: "oauth" }, storage });
977
+
978
+ try {
979
+ await root.getValue().connect();
980
+ const transport = mocks.transports[0];
981
+
982
+ await expect(
983
+ root
984
+ .getValue()
985
+ .completeAuth(
986
+ "https://example.com/callback?code=abc&state=aui-mcp%3AZG9jcw.forged",
987
+ ),
988
+ ).rejects.toThrow("OAuth state does not match the authorization request");
989
+ await flushMacrotask();
990
+
991
+ expect(mocks.transports).toHaveLength(1);
992
+ expect(transport.close).not.toHaveBeenCalled();
993
+ expect(root.getValue().getState()).toMatchObject({
994
+ connectionState: "connected",
995
+ lastError: null,
996
+ });
997
+ } finally {
998
+ root.unmount();
999
+ }
1000
+ });
1001
+
1002
+ it("does not resume authorization after disconnecting during validation", async () => {
1003
+ let resolveAuthState!: (state: {
1004
+ serverUrl: string;
1005
+ state: string;
1006
+ }) => void;
1007
+ const storage = createStorage();
1008
+ vi.mocked(storage.loadAuthState).mockImplementation(
1009
+ () =>
1010
+ new Promise((resolve) => {
1011
+ resolveAuthState = resolve;
1012
+ }),
1013
+ );
1014
+ const root = mount({ auth: { type: "oauth" }, storage });
1015
+
1016
+ try {
1017
+ const completeAuth = root
1018
+ .getValue()
1019
+ .completeAuth("https://example.com/callback?code=abc&state=expected");
1020
+ await waitFor(() => resolveAuthState !== undefined);
1021
+
1022
+ await root.getValue().disconnect();
1023
+ resolveAuthState({
1024
+ serverUrl: "https://example.com/mcp",
1025
+ state: "expected",
1026
+ });
1027
+
1028
+ await expect(completeAuth).rejects.toThrow(
1029
+ 'MCP server "docs" authorization was interrupted before completion.',
1030
+ );
1031
+ expect(mocks.transports).toHaveLength(0);
1032
+ expect(root.getValue().getState()).toMatchObject({
1033
+ connectionState: "disconnected",
1034
+ lastError: null,
598
1035
  });
599
1036
  } finally {
600
1037
  root.unmount();
@@ -605,11 +1042,18 @@ describe("McpServerResource completeAuth", () => {
605
1042
  mocks.finishAuthResults.push(() =>
606
1043
  Promise.reject(new Error("invalid_grant")),
607
1044
  );
608
- const root = mount({ auth: { type: "oauth" } });
1045
+ const storage = createStorage();
1046
+ vi.mocked(storage.loadAuthState).mockResolvedValue({
1047
+ serverUrl: "https://example.com/mcp",
1048
+ state: "expected",
1049
+ });
1050
+ const root = mount({ auth: { type: "oauth" }, storage });
609
1051
 
610
1052
  try {
611
1053
  await expect(
612
- root.getValue().completeAuth("https://example.com/callback?code=abc"),
1054
+ root
1055
+ .getValue()
1056
+ .completeAuth("https://example.com/callback?code=abc&state=expected"),
613
1057
  ).rejects.toThrow("invalid_grant");
614
1058
  await flushMacrotask();
615
1059
 
@@ -619,7 +1063,9 @@ describe("McpServerResource completeAuth", () => {
619
1063
  message: "invalid_grant",
620
1064
  },
621
1065
  });
622
- expect(mocks.transports[0].finishAuth).toHaveBeenCalledWith("abc");
1066
+ expect(mocks.transports[0].finishAuth).toHaveBeenCalledWith(
1067
+ expect.any(URLSearchParams),
1068
+ );
623
1069
  expect(mocks.transports[0].close).toHaveBeenCalledTimes(1);
624
1070
  } finally {
625
1071
  root.unmount();
@@ -634,13 +1080,18 @@ describe("McpServerResource completeAuth", () => {
634
1080
  resolveFinishAuth = resolve;
635
1081
  }),
636
1082
  );
637
- const root = mount({ auth: { type: "oauth" } });
1083
+ const storage = createStorage();
1084
+ vi.mocked(storage.loadAuthState).mockResolvedValue({
1085
+ serverUrl: "https://example.com/mcp",
1086
+ state: "expected",
1087
+ });
1088
+ const root = mount({ auth: { type: "oauth" }, storage });
638
1089
  let didUnmount = false;
639
1090
 
640
1091
  try {
641
1092
  const completeAuth = root
642
1093
  .getValue()
643
- .completeAuth("https://example.com/callback?code=abc");
1094
+ .completeAuth("https://example.com/callback?code=abc&state=expected");
644
1095
  await waitFor(
645
1096
  () => mocks.transports[0]?.finishAuth.mock.calls.length === 1,
646
1097
  );
@@ -665,13 +1116,18 @@ describe("McpServerResource completeAuth", () => {
665
1116
  rejectFinishAuth = reject;
666
1117
  }),
667
1118
  );
668
- const root = mount({ auth: { type: "oauth" } });
1119
+ const storage = createStorage();
1120
+ vi.mocked(storage.loadAuthState).mockResolvedValue({
1121
+ serverUrl: "https://example.com/mcp",
1122
+ state: "expected",
1123
+ });
1124
+ const root = mount({ auth: { type: "oauth" }, storage });
669
1125
  let didUnmount = false;
670
1126
 
671
1127
  try {
672
1128
  const completeAuth = root
673
1129
  .getValue()
674
- .completeAuth("https://example.com/callback?code=abc");
1130
+ .completeAuth("https://example.com/callback?code=abc&state=expected");
675
1131
  await waitFor(
676
1132
  () => mocks.transports[0]?.finishAuth.mock.calls.length === 1,
677
1133
  );
@@ -1481,3 +1937,143 @@ describe("McpServerResource resource methods", () => {
1481
1937
  }
1482
1938
  });
1483
1939
  });
1940
+
1941
+ describe("getConnectionDependencies storage scope", () => {
1942
+ const propsWith = (storage: MCPStorage): McpServerResourceProps => ({
1943
+ id: "docs",
1944
+ kind: "connector",
1945
+ name: "Docs",
1946
+ url: "https://example.com/mcp",
1947
+ auth: { type: "oauth" },
1948
+ storage,
1949
+ redirectUri: "https://example.com/callback",
1950
+ autoConnect: false,
1951
+ onRemove: async () => {},
1952
+ });
1953
+
1954
+ it("keys the connection on a declared storage scopeId", () => {
1955
+ const a = { ...createStorage(), scopeId: "local-storage:a" };
1956
+ const b = { ...createStorage(), scopeId: "local-storage:b" };
1957
+
1958
+ expect(getConnectionDependencies(propsWith(a))).not.toEqual(
1959
+ getConnectionDependencies(propsWith(b)),
1960
+ );
1961
+ });
1962
+
1963
+ it("treats storages sharing a scopeId as the same connection target", () => {
1964
+ const a = { ...createStorage(), scopeId: "local-storage:same" };
1965
+ const b = { ...createStorage(), scopeId: "local-storage:same" };
1966
+
1967
+ expect(getConnectionDependencies(propsWith(a))).toEqual(
1968
+ getConnectionDependencies(propsWith(b)),
1969
+ );
1970
+ });
1971
+
1972
+ it("does not key the connection on storage identity when no scopeId is declared", () => {
1973
+ expect(getConnectionDependencies(propsWith(createStorage()))).toEqual(
1974
+ getConnectionDependencies(propsWith(createStorage())),
1975
+ );
1976
+ });
1977
+
1978
+ it("ignores the storage scope for none-auth servers", () => {
1979
+ const a = { ...createStorage(), scopeId: "local-storage:a" };
1980
+ const b = { ...createStorage(), scopeId: "local-storage:b" };
1981
+ const noneProps = (storage: MCPStorage): McpServerResourceProps => ({
1982
+ ...propsWith(storage),
1983
+ auth: { type: "none" },
1984
+ });
1985
+
1986
+ expect(getConnectionDependencies(noneProps(a))).toEqual(
1987
+ getConnectionDependencies(noneProps(b)),
1988
+ );
1989
+ });
1990
+ });
1991
+
1992
+ describe("McpServerResource oauth storage swap", () => {
1993
+ beforeEach(resetMocks);
1994
+
1995
+ it("reconnects onto the replacement storage when the scope changes", async () => {
1996
+ const persisted = {
1997
+ serverUrl: "https://example.com/mcp",
1998
+ clientInformation: {
1999
+ client_id: "registered-client",
2000
+ redirect_uris: ["https://example.com/callback"],
2001
+ },
2002
+ clientInformationSource: "registered" as const,
2003
+ tokens: { access_token: "tok", token_type: "bearer" },
2004
+ tokensClientId: "registered-client",
2005
+ };
2006
+ const storageA = {
2007
+ ...createStorage(),
2008
+ scopeId: "scope:a",
2009
+ loadAuthState: vi.fn(async () => persisted),
2010
+ };
2011
+ const storageB = {
2012
+ ...createStorage(),
2013
+ scopeId: "scope:b",
2014
+ loadAuthState: vi.fn(async () => persisted),
2015
+ };
2016
+ let setStorage!: (s: MCPStorage) => void;
2017
+
2018
+ const Host = resource(function useHost() {
2019
+ const [storage, set] = useState<MCPStorage>(storageA);
2020
+ setStorage = set;
2021
+ return useResource(
2022
+ McpServerResource({
2023
+ id: "docs",
2024
+ kind: "connector",
2025
+ name: "Docs",
2026
+ url: "https://example.com/mcp",
2027
+ auth: { type: "oauth" },
2028
+ storage,
2029
+ redirectUri: "https://example.com/callback",
2030
+ autoConnect: true,
2031
+ connectionTimeout: 10_000,
2032
+ onRemove: vi.fn(async () => {}),
2033
+ }),
2034
+ );
2035
+ });
2036
+
2037
+ const root = createTapRoot(function SwapRoot() {
2038
+ return useResource(Host());
2039
+ });
2040
+
2041
+ try {
2042
+ await waitFor(() => mocks.transports.length === 1);
2043
+ const authProviderA =
2044
+ mocks.StreamableHTTPClientTransport.mock.calls[0]?.[1]?.authProvider;
2045
+ await authProviderA.tokens();
2046
+ expect(storageA.loadAuthState).toHaveBeenCalledWith("docs");
2047
+ expect(storageB.loadAuthState).not.toHaveBeenCalled();
2048
+
2049
+ setStorage(storageB);
2050
+ await waitForResourceUpdate(() => mocks.transports.length === 2);
2051
+ await waitForResourceUpdate(
2052
+ () => vi.mocked(mocks.transports[0]!.close).mock.calls.length > 0,
2053
+ );
2054
+
2055
+ const authProviderB =
2056
+ mocks.StreamableHTTPClientTransport.mock.calls[1]?.[1]?.authProvider;
2057
+ const callsBefore = vi.mocked(storageA.loadAuthState).mock.calls.length;
2058
+ await authProviderB.tokens();
2059
+ expect(storageB.loadAuthState).toHaveBeenCalledWith("docs");
2060
+ expect(vi.mocked(storageA.loadAuthState).mock.calls.length).toBe(
2061
+ callsBefore,
2062
+ );
2063
+
2064
+ await authProviderB.saveTokens({
2065
+ access_token: "fresh",
2066
+ token_type: "bearer",
2067
+ });
2068
+ expect(storageB.saveAuthState).toHaveBeenCalledWith(
2069
+ "docs",
2070
+ expect.objectContaining({
2071
+ tokens: expect.objectContaining({ access_token: "fresh" }),
2072
+ }),
2073
+ );
2074
+ expect(storageA.saveAuthState).not.toHaveBeenCalled();
2075
+ } finally {
2076
+ root.unmount();
2077
+ }
2078
+ });
2079
+ });