Add optional function to modify headers

.changeset/heavy-apples-hug.md

@@ -0,0 +1,5 @@
+---
+'@apollo/sandbox': minor
+---
+
+Adding support for optional `modifyHeaders` function which can modify the headers before sending the request

packages/sandbox/src/EmbeddedSandbox.ts

@@ -1,6 +1,6 @@
 import { EMBEDDABLE_SANDBOX_URL, IFRAME_DOM_ID } from './helpers/constants';
 import { defaultHandleRequest } from './helpers/defaultHandleRequest';
-import type { HandleRequest } from './helpers/postMessageRelayHelpers';
+import type { HandleRequest, ModifyHeaders } from './helpers/postMessageRelayHelpers';
 import { setupSandboxEmbedRelay } from './setupSandboxEmbedRelay';
 import packageJSON from '../package.json';
 import type { JSONObject } from './helpers/types';
@@ -53,6 +53,11 @@ export interface EmbeddableSandboxOptions {
    */
   handleRequest?: HandleRequest;
 
+   /**
+    * optional. Function that accepts the original headers and returns the modified headers.
+    */
+  modifyHeaders?: ModifyHeaders;
+
   /**
    * optional. If this is passed, its value will take precedence over your sandbox connection settings `includeCookies` value.
    * If you pass `handleRequest`, that will override this value and its behavior.
@@ -122,6 +127,7 @@ export class EmbeddedSandbox {
     this.disposable = setupSandboxEmbedRelay({
       embeddedSandboxIFrameElement: this.embeddedSandboxIFrameElement,
       handleRequest: this.handleRequest,
+      modifyHeaders: this.options.modifyHeaders,
       __testLocal__: !!this.__testLocal__,
     });
   }

packages/sandbox/src/helpers/postMessageRelayHelpers.ts

@@ -41,6 +41,10 @@ export type HandleRequest = (
   endpointUrl: string,
   options: Omit<RequestInit, 'headers'> & { headers: Record<string, string> }
 ) => Promise<Response>;
+export type ModifyHeaders = (
+  endpointUrl: string,
+  headers: Record<string, string> | undefined
+) => Promise<Record<string, string> | undefined>;
 
 export type SocketStatus = 'disconnected' | 'connecting' | 'connected';
 

packages/sandbox/src/setupSandboxEmbedRelay.ts

@@ -12,22 +12,25 @@ import {
   handleAuthenticationPostMessage,
   HandleRequest,
   IncomingEmbedMessage,
+  ModifyHeaders,
   sendPostMessageToEmbed,
 } from './helpers/postMessageRelayHelpers';
 import { executeSubscription } from './helpers/subscriptionPostMessageRelayHelpers';
 
 export function setupSandboxEmbedRelay({
   handleRequest,
+  modifyHeaders,
   embeddedSandboxIFrameElement,
   __testLocal__,
 }: {
   handleRequest: HandleRequest;
+  modifyHeaders?: ModifyHeaders;
   embeddedSandboxIFrameElement: HTMLIFrameElement;
   __testLocal__: boolean;
 }) {
   const embedUrl = EMBEDDABLE_SANDBOX_URL(__testLocal__);
   // Callback definition
-  const onPostMessageReceived = (event: IncomingEmbedMessage) => {
+  const onPostMessageReceived = async (event: IncomingEmbedMessage) => {
     handleAuthenticationPostMessage({
       event,
       embedUrl,
@@ -85,7 +88,7 @@ export function setupSandboxEmbedRelay({
         data.operationId
       ) {
         // Extract the operation details from the event.data object
-        const { operation, variables, operationName, operationId, headers } =
+        const { operation, variables, operationName, operationId, headers: originalHeaders } =
           data;
 
         if (isQueryOrMutation) {
@@ -95,6 +98,10 @@ export function setupSandboxEmbedRelay({
               'Something went wrong, we should not have gotten here. The sandbox endpoint url was not sent.'
             );
           }
+          // If the user has provided a function to modify headers, call it
+          const headers = modifyHeaders
+            ? await modifyHeaders(endpointUrl, originalHeaders)
+            : originalHeaders;
           executeOperation({
             endpointUrl,
             handleRequest,
@@ -112,6 +119,10 @@ export function setupSandboxEmbedRelay({
           });
         } else if (isSubscription) {
           const { httpMultipartParams } = data;
+          // If the user has provided a function to modify headers, call it
+          const headers = modifyHeaders
+            ? await modifyHeaders(data.subscriptionUrl, originalHeaders)
+            : originalHeaders;
           executeSubscription({
             operation,
             operationName,