Implement ChannelFetch util to relay fetch requests over the channel,…

… and have URQL use it

src/Panel.tsx

@@ -28,9 +28,10 @@ import { ControlsProvider } from "./screens/VisualTests/ControlsContext";
 import { RunBuildProvider } from "./screens/VisualTests/RunBuildContext";
 import { VisualTests } from "./screens/VisualTests/VisualTests";
 import { APIInfoPayload, GitInfoPayload, LocalBuildProgress, UpdateStatusFunction } from "./types";
-import { client, Provider, useAccessToken } from "./utils/graphQLClient";
+import { createClient, GraphQLClientProvider, useAccessToken } from "./utils/graphQLClient";
 import { TelemetryProvider } from "./utils/TelemetryContext";
 import { useBuildEvents } from "./utils/useBuildEvents";
+import { useChannelFetch } from "./utils/useChannelFetch";
 import { useProjectId } from "./utils/useProjectId";
 import { clearSessionState, useSessionState } from "./utils/useSessionState";
 import { useSharedState } from "./utils/useSharedState";
@@ -81,8 +82,9 @@ export const Panel = ({ active, api }: PanelProps) => {
   const trackEvent = useCallback((data: any) => emit(TELEMETRY, data), [emit]);
   const { isRunning, startBuild, stopBuild } = useBuildEvents({ localBuildProgress, accessToken });
 
+  const fetch = useChannelFetch();
   const withProviders = (children: React.ReactNode) => (
-    <Provider key={PANEL_ID} value={client}>
+    <GraphQLClientProvider key={PANEL_ID} value={createClient({ fetch })}>
       <TelemetryProvider value={trackEvent}>
         <AuthProvider value={{ accessToken, setAccessToken }}>
           <UninstallProvider
@@ -99,7 +101,7 @@ export const Panel = ({ active, api }: PanelProps) => {
           </UninstallProvider>
         </AuthProvider>
       </TelemetryProvider>
-    </Provider>
+    </GraphQLClientProvider>
   );
 
   if (!active) {

src/preset.ts

@@ -35,6 +35,7 @@ import {
   LocalBuildProgress,
   ProjectInfoPayload,
 } from "./types";
+import { ChannelFetch } from "./utils/ChannelFetch";
 import { SharedState } from "./utils/SharedState";
 import { updateChromaticConfig } from "./utils/updateChromaticConfig";
 
@@ -193,6 +194,9 @@ const watchConfigFile = async (
 async function serverChannel(channel: Channel, options: Options & { configFile?: string }) {
   const { configFile, presets } = options;
 
+  // Handle relayed fetch requests from the client
+  ChannelFetch.subscribe(ADDON_ID, channel);
+
   // Lazy load these APIs since we don't need them right away
   const apiPromise = presets.apply<any>("experimental_serverAPI");
   const corePromise = presets.apply("core");

src/utils/ChannelFetch.ts

@@ -0,0 +1,48 @@
+import type { Channel } from "@storybook/channels";
+
+export const FETCH_ABORTED = "ChannelFetch/aborted";
+export const FETCH_REQUEST = "ChannelFetch/request";
+export const FETCH_RESPONSE = "ChannelFetch/response";
+
+type ChannelLike = Pick<Channel, "emit" | "on" | "off">;
+
+const instances = new Map<string, ChannelFetch>();
+
+export class ChannelFetch {
+  channel: ChannelLike;
+
+  abortControllers: Map<string, AbortController>;
+
+  constructor(channel: ChannelLike) {
+    this.channel = channel;
+    this.abortControllers = new Map<string, AbortController>();
+
+    this.channel.on(FETCH_ABORTED, ({ requestId }) => {
+      this.abortControllers.get(requestId)?.abort();
+      this.abortControllers.delete(requestId);
+    });
+
+    this.channel.on(FETCH_REQUEST, async ({ requestId, input, init }) => {
+      const controller = new AbortController();
+      this.abortControllers.set(requestId, controller);
+
+      try {
+        const res = await fetch(input as RequestInfo, { ...init, signal: controller.signal });
+        const body = await res.text();
+        const headers = Array.from(res.headers as any);
+        const response = { body, headers, status: res.status, statusText: res.statusText };
+        this.channel.emit(FETCH_RESPONSE, { requestId, response });
+      } catch (error) {
+        this.channel.emit(FETCH_RESPONSE, { requestId, error });
+      } finally {
+        this.abortControllers.delete(requestId);
+      }
+    });
+  }
+
+  static subscribe<T>(key: string, channel: ChannelLike) {
+    const instance = instances.get(key) || new ChannelFetch(channel);
+    if (!instances.has(key)) instances.set(key, instance);
+    return instance;
+  }
+}

src/utils/graphQLClient.tsx

@@ -1,13 +1,11 @@
 import { useAddonState } from "@storybook/manager-api";
 import { authExchange } from "@urql/exchange-auth";
 import React from "react";
-import { Client, fetchExchange, mapExchange, Provider } from "urql";
+import { Client, ClientOptions, fetchExchange, mapExchange, Provider } from "urql";
 import { v4 as uuid } from "uuid";
 
 import { ACCESS_TOKEN_KEY, ADDON_ID, CHROMATIC_API_URL } from "../constants";
 
-export { Provider };
-
 let currentToken: string | null;
 let currentTokenExpiration: number | null;
 const setCurrentToken = (token: string | null) => {
@@ -56,56 +54,69 @@ export const getFetchOptions = (token?: string) => ({
   },
 });
 
-export const client = new Client({
-  url: CHROMATIC_API_URL,
-  exchanges: [
-    // We don't use cacheExchange, because it would inadvertently share data between stories.
-    mapExchange({
-      onResult(result) {
-        // Not all queries contain the `viewer` field, in which case it will be `undefined`.
-        // When we do retrieve the field but the token is invalid, it will be `null`.
-        if (result.data?.viewer === null) setCurrentToken(null);
-      },
-    }),
-    authExchange(async (utils) => {
-      return {
-        addAuthToOperation(operation) {
-          if (!currentToken) return operation;
-          return utils.appendHeaders(operation, { Authorization: `Bearer ${currentToken}` });
+export const createClient = (options?: Partial<ClientOptions>) =>
+  new Client({
+    url: CHROMATIC_API_URL,
+    exchanges: [
+      // We don't use cacheExchange, because it would inadvertently share data between stories.
+      mapExchange({
+        onResult(result) {
+          // Not all queries contain the `viewer` field, in which case it will be `undefined`.
+          // When we do retrieve the field but the token is invalid, it will be `null`.
+          if (result.data?.viewer === null) setCurrentToken(null);
         },
+      }),
+      authExchange(async (utils) => {
+        return {
+          addAuthToOperation(operation) {
+            if (!currentToken) return operation;
+            return utils.appendHeaders(operation, { Authorization: `Bearer ${currentToken}` });
+          },
 
-        // Determine if the current error is an authentication error.
-        didAuthError: (error) =>
-          error.response.status === 401 ||
-          error.graphQLErrors.some((e) => e.message.includes("Must login")),
+          // Determine if the current error is an authentication error.
+          didAuthError: (error) =>
+            error.response.status === 401 ||
+            error.graphQLErrors.some((e) => e.message.includes("Must login")),
 
-        // If didAuthError returns true, clear the token. Ideally we should refresh the token here.
-        // The operation will be retried automatically.
-        async refreshAuth() {
-          setCurrentToken(null);
-        },
+          // If didAuthError returns true, clear the token. Ideally we should refresh the token here.
+          // The operation will be retried automatically.
+          async refreshAuth() {
+            setCurrentToken(null);
+          },
 
-        // Prevent making a request if we know the token is missing, invalid or expired.
-        // This handler is called repeatedly so we avoid parsing the token each time.
-        willAuthError() {
-          if (!currentToken) return true;
-          try {
-            if (!currentTokenExpiration) {
-              const { exp } = JSON.parse(atob(currentToken.split(".")[1]));
-              currentTokenExpiration = exp;
+          // Prevent making a request if we know the token is missing, invalid or expired.
+          // This handler is called repeatedly so we avoid parsing the token each time.
+          willAuthError() {
+            if (!currentToken) return true;
+            try {
+              if (!currentTokenExpiration) {
+                const { exp } = JSON.parse(atob(currentToken.split(".")[1]));
+                currentTokenExpiration = exp;
+              }
+              return Date.now() / 1000 > (currentTokenExpiration || 0);
+            } catch (e) {
+              return true;
             }
-            return Date.now() / 1000 > (currentTokenExpiration || 0);
-          } catch (e) {
-            return true;
-          }
-        },
-      };
-    }),
-    fetchExchange,
-  ],
-  fetchOptions: getFetchOptions(), // Auth header (token) is handled by authExchange
-});
+          },
+        };
+      }),
+      fetchExchange,
+    ],
+    fetchOptions: getFetchOptions(), // Auth header (token) is handled by authExchange
+    ...options,
+  });
 
-export const GraphQLClientProvider = ({ children }: { children: React.ReactNode }) => {
-  return <Provider value={client}>{children}</Provider>;
+export const GraphQLClientProvider = ({
+  children,
+  value = createClient(),
+  ...rest
+}: {
+  children: React.ReactNode;
+  value?: Client;
+}) => {
+  return (
+    <Provider value={value} {...rest}>
+      {children}
+    </Provider>
+  );
 };

src/utils/useChannelFetch.ts

@@ -0,0 +1,54 @@
+import { useChannel } from "@storybook/manager-api";
+
+const FETCH_ABORTED = "ChannelFetch/aborted";
+const FETCH_REQUEST = "ChannelFetch/request";
+const FETCH_RESPONSE = "ChannelFetch/response";
+
+type SerializedResponse = {
+  status: number;
+  statusText: string;
+  headers: [string, string][];
+  body: string;
+};
+
+const pendingRequests = new Map<
+  string,
+  { resolve: (value: Response) => void; reject: (reason?: any) => void }
+>();
+
+export const useChannelFetch: () => typeof fetch = () => {
+  const emit = useChannel({
+    [FETCH_RESPONSE]: (
+      data:
+        | { requestId: string; response: SerializedResponse }
+        | { requestId: string; error: string }
+    ) => {
+      const request = pendingRequests.get(data.requestId);
+      if (!request) return;
+
+      pendingRequests.delete(data.requestId);
+      if ("error" in data) {
+        request.reject(new Error(data.error));
+      } else {
+        const { body, headers, status, statusText } = data.response;
+        const res = new Response(body, { headers, status, statusText });
+        request.resolve(res);
+      }
+    },
+  });
+
+  return async (input: string | URL | Request, { signal, ...init }: RequestInit = {}) => {
+    const requestId = Math.random().toString(36).slice(2);
+    emit(FETCH_REQUEST, { requestId, input, init });
+
+    signal?.addEventListener("abort", () => emit(FETCH_ABORTED, { requestId }));
+
+    return new Promise((resolve, reject) => {
+      pendingRequests.set(requestId, { resolve, reject });
+      setTimeout(() => {
+        reject(new Error("Request timed out"));
+        pendingRequests.delete(requestId);
+      }, 30000);
+    });
+  };
+};