From eb1b8693645f946d48c1689930ea4fca8d1556a3 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Daniel=20M=C3=BCller?= <daniel.mueller@surnet.ch>
Date: Wed, 10 Dec 2025 13:18:10 +0100
Subject: [PATCH] Add optional function to modify headers

---
 .changeset/heavy-apples-hug.md                    |  5 +++++
 packages/sandbox/src/EmbeddedSandbox.ts           |  7 +++++++
 .../src/helpers/postMessageRelayHelpers.ts        |  4 ++++
 packages/sandbox/src/setupSandboxEmbedRelay.ts    | 15 +++++++++++++--
 4 files changed, 29 insertions(+), 2 deletions(-)
 create mode 100644 .changeset/heavy-apples-hug.md

diff --git a/.changeset/heavy-apples-hug.md b/.changeset/heavy-apples-hug.md
new file mode 100644
index 0000000..d03268f
--- /dev/null
+++ b/.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
diff --git a/packages/sandbox/src/EmbeddedSandbox.ts b/packages/sandbox/src/EmbeddedSandbox.ts
index c6fb866..e164565 100644
--- a/packages/sandbox/src/EmbeddedSandbox.ts
+++ b/packages/sandbox/src/EmbeddedSandbox.ts
@@ -6,6 +6,7 @@ import { defaultHandleRequest } from './helpers/defaultHandleRequest';
 import type {
   DisposableResource,
   HandleRequest,
+  ModifyHeaders,
 } from './helpers/postMessageRelayHelpers';
 import { setupSandboxEmbedRelay } from './setupSandboxEmbedRelay';
 import packageJSON from '../package.json';
@@ -59,6 +60,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.
@@ -128,6 +134,7 @@ export class EmbeddedSandbox {
     this.disposable = setupSandboxEmbedRelay({
       embeddedSandboxIFrameElement: this.embeddedSandboxIFrameElement,
       handleRequest: this.handleRequest,
+      modifyHeaders: this.options.modifyHeaders,
       __testLocal__: !!this.__testLocal__,
     });
   }
diff --git a/packages/sandbox/src/helpers/postMessageRelayHelpers.ts b/packages/sandbox/src/helpers/postMessageRelayHelpers.ts
index 1774de0..abc0803 100644
--- a/packages/sandbox/src/helpers/postMessageRelayHelpers.ts
+++ b/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';
 
diff --git a/packages/sandbox/src/setupSandboxEmbedRelay.ts b/packages/sandbox/src/setupSandboxEmbedRelay.ts
index b625521..ded23d3 100644
--- a/packages/sandbox/src/setupSandboxEmbedRelay.ts
+++ b/packages/sandbox/src/setupSandboxEmbedRelay.ts
@@ -14,22 +14,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;
 }): DisposableResource {
   const embedUrlOrigin = EMBEDDABLE_SANDBOX_URL_ORIGIN(__testLocal__);
   // Callback definition
-  const onPostMessageReceived = (event: IncomingEmbedMessage) => {
+  const onPostMessageReceived = async (event: IncomingEmbedMessage) => {
     handleAuthenticationPostMessage({
       event,
       embedUrlOrigin,
@@ -87,7 +90,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) {
@@ -97,6 +100,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,
@@ -114,6 +121,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,