feat(laboratory/preflight-script): set request headers

cypress/e2e/laboratory-preflight-script.cy.ts

@@ -1,5 +1,16 @@
 import { dedent } from '../support/testkit';
 
+const selectors = {
+  buttonModal: 'preflight-script-modal-button',
+  buttonToggle: 'toggle-preflight-script',
+  graphiql: {
+    buttonExecute: '.graphiql-execute-button',
+  },
+  modal: {
+    buttonSubmit: 'preflight-script-modal-submit',
+  },
+};
+
 beforeEach(() => {
   cy.clearLocalStorage().then(async () => {
     cy.task('seedTarget').then(({ slug, refreshToken }: any) => {
@@ -186,6 +197,66 @@ throw new TypeError('Test')`,
 });
 
 describe('Execution', () => {
+  it('result.request.headers are added to the graphiql request base headers', () => {
+    const headers = {
+      foo: { name: 'foo', value: 'bar' },
+    };
+    cy.dataCy(selectors.buttonToggle).click();
+    cy.dataCy(selectors.buttonModal).click();
+    setEditorScript(`lab.request.headers.append('${headers.foo.name}', '${headers.foo.value}')`);
+    cy.dataCy(selectors.modal.buttonSubmit).click();
+    cy.intercept({
+      method: 'POST',
+      headers: {
+        [headers.foo.name]: headers.foo.value,
+      },
+    }).as('post');
+    cy.get(selectors.graphiql.buttonExecute).click();
+    cy.wait('@post');
+  });
+
+  it('result.request.headers take precedence over graphiql request base headers', () => {
+    // --- Pre Assert Integrity Check: make sure the header we think we're overriding is actually there.
+    const baseHeaders = {
+      accept: {
+        name: 'accept',
+        value: 'application/json, multipart/mixed',
+      },
+    };
+    cy.intercept({
+      method: 'POST',
+      headers: {
+        [baseHeaders.accept.name]: baseHeaders.accept.value,
+      },
+    }).as('integrityCheck');
+    cy.get(selectors.graphiql.buttonExecute).click();
+    cy.wait('@integrityCheck');
+    // ---
+
+    const preflightHeaders = {
+      accept: {
+        name: 'accept',
+        value: 'application/graphql-response+json; charset=utf-8, application/json; charset=utf-8',
+      },
+    };
+    cy.dataCy(selectors.buttonToggle).click();
+    cy.dataCy(selectors.buttonModal).click();
+    setEditorScript(
+      `lab.request.headers.append('${preflightHeaders.accept.name}', '${preflightHeaders.accept.value}')`,
+    );
+    cy.dataCy(selectors.modal.buttonSubmit).click();
+    cy.intercept({
+      method: 'POST',
+      headers: {
+        [preflightHeaders.accept.name]: preflightHeaders.accept.value,
+      },
+    }).as('post');
+    cy.get(selectors.graphiql.buttonExecute).click();
+    cy.wait('@post');
+  });
+
+  it('result.request.headers do not receive placeholder substitution');
+
   it('header placeholders are substituted with environment variables', () => {
     cy.dataCy('toggle-preflight-script').click();
     cy.get('[data-name="headers"]').click();

packages/web/app/src/lib/kit/helpers.ts

@@ -0,0 +1,10 @@
+export const tryOr = <$PrimaryResult, $FallbackResult>(
+  fn: () => $PrimaryResult,
+  fallback: () => $FallbackResult,
+): $PrimaryResult | $FallbackResult => {
+  try {
+    return fn();
+  } catch {
+    return fallback();
+  }
+};

packages/web/app/src/lib/kit/index.ts

@@ -0,0 +1,6 @@
+// eslint-disable-next-line import/no-self-import
+export * as Kit from './index';
+
+export * from './never';
+export * from './types/headers';
+export * from './helpers';

packages/web/app/src/lib/kit/never.ts

@@ -0,0 +1,15 @@
+/**
+ * This case is impossible.
+ * If it happens, then that means there is a bug in our code.
+ */
+export const neverCase = (value: never): never => {
+  never({ type: 'case', value });
+};
+
+/**
+ * This code cannot be reached.
+ * If it is reached, then that means there is a bug in our code.
+ */
+export const never: (context?: object) => never = context => {
+  throw new Error(`Something that should be impossible happened`, { cause: context });
+};

packages/web/app/src/lib/kit/types/headers.ts

@@ -0,0 +1,4 @@
+// eslint-disable-next-line @typescript-eslint/no-namespace
+export namespace Headers {
+  export type Encoded = [name: string, value: string][];
+}

packages/web/app/src/lib/preflight-sandbox/graphiql-plugin.tsx

@@ -32,10 +32,16 @@ import { Editor as MonacoEditor, OnMount, type Monaco } from '@monaco-editor/rea
 import { Cross2Icon, InfoCircledIcon, Pencil1Icon, TriangleRightIcon } from '@radix-ui/react-icons';
 import { captureException } from '@sentry/react';
 import { useParams } from '@tanstack/react-router';
+import { Kit } from '../kit';
 import { cn } from '../utils';
 import labApiDefinitionRaw from './lab-api-declaration?raw';
 import { IFrameEvents, LogMessage } from './shared-types';
 
+export type PreflightScriptResultData = Omit<
+  IFrameEvents.Outgoing.EventData.Result,
+  'type' | 'runId'
+>;
+
 export const preflightScriptPlugin: GraphiQLPlugin = {
   icon: () => (
     <svg
@@ -135,14 +141,6 @@ const PreflightScript_TargetFragment = graphql(`
 
 export type LogRecord = LogMessage | { type: 'separator' };
 
-function safeParseJSON(str: string): Record<string, unknown> | null {
-  try {
-    return JSON.parse(str);
-  } catch {
-    return null;
-  }
-}
-
 export const enum PreflightWorkerState {
   running,
   ready,
@@ -173,9 +171,24 @@ export function usePreflightScript(args: {
 
   const currentRun = useRef<null | Function>(null);
 
-  async function execute(script = target?.preflightScript?.sourceCode ?? '', isPreview = false) {
+  async function execute(
+    script = target?.preflightScript?.sourceCode ?? '',
+    isPreview = false,
+  ): Promise<PreflightScriptResultData> {
+    const resultEnvironmentVariablesDecoded: PreflightScriptResultData['environmentVariables'] =
+      Kit.tryOr(
+        () => JSON.parse(latestEnvironmentVariablesRef.current),
+        () => ({}),
+      );
+    const result: PreflightScriptResultData = {
+      request: {
+        headers: [],
+      },
+      environmentVariables: resultEnvironmentVariablesDecoded,
+    };
+
     if (isPreview === false && !isPreflightScriptEnabled) {
-      return safeParseJSON(latestEnvironmentVariablesRef.current);
+      return result;
     }
 
     const id = crypto.randomUUID();
@@ -201,7 +214,8 @@ export function usePreflightScript(args: {
           type: IFrameEvents.Incoming.Event.run,
           id,
           script,
-          environmentVariables: (environmentVariables && safeParseJSON(environmentVariables)) || {},
+          // Preflight Script has read/write relationship with environment variables.
+          environmentVariables: result.environmentVariables,
         } satisfies IFrameEvents.Incoming.EventData,
         '*',
       );
@@ -257,16 +271,23 @@ export function usePreflightScript(args: {
         }
 
         if (ev.data.type === IFrameEvents.Outgoing.Event.result) {
-          const mergedEnvironmentVariables = JSON.stringify(
-            {
-              ...safeParseJSON(latestEnvironmentVariablesRef.current),
-              ...ev.data.environmentVariables,
-            },
+          const mergedEnvironmentVariables = {
+            ...result.environmentVariables,
+            ...ev.data.environmentVariables,
+          };
+          result.environmentVariables = mergedEnvironmentVariables;
+          result.request.headers = ev.data.request.headers;
+
+          // Cause the new state of environment variables to be
+          // written back to local storage.
+          const mergedEnvironmentVariablesEncoded = JSON.stringify(
+            result.environmentVariables,
             null,
             2,
           );
-          setEnvironmentVariables(mergedEnvironmentVariables);
-          latestEnvironmentVariablesRef.current = mergedEnvironmentVariables;
+          setEnvironmentVariables(mergedEnvironmentVariablesEncoded);
+          latestEnvironmentVariablesRef.current = mergedEnvironmentVariablesEncoded;
+
           setLogs(logs => [
             ...logs,
             {
@@ -301,7 +322,6 @@ export function usePreflightScript(args: {
           ]);
           setFinished();
           closedOpenedPrompts();
-
           return;
         }
 
@@ -310,6 +330,27 @@ export function usePreflightScript(args: {
           setLogs(logs => [...logs, log]);
           return;
         }
+
+        if (ev.data.type === IFrameEvents.Outgoing.Event.ready) {
+          console.debug('preflight sandbox graphiql plugin: noop iframe event:', ev.data);
+          return;
+        }
+
+        if (ev.data.type === IFrameEvents.Outgoing.Event.start) {
+          console.debug('preflight sandbox graphiql plugin: noop iframe event:', ev.data);
+          return;
+        }
+
+        // Window message events can be emitted from unknowable sources.
+        // For example when our e2e tests runs within Cypress GUI, we see a `MessageEvent` with `.data` of `{ vscodeScheduleAsyncWork: 3 }`.
+        // Since we cannot know if the event source is Preflight Script, we cannot perform an exhaustive check.
+        //
+        // Kit.neverCase(ev.data);
+        //
+        console.debug(
+          'preflight sandbox graphiql plugin: An unknown window message event received. Ignoring.',
+          ev,
+        );
       }
 
       window.addEventListener('message', eventHandler);
@@ -328,7 +369,8 @@ export function usePreflightScript(args: {
       window.removeEventListener('message', eventHandler);
 
       setState(PreflightWorkerState.ready);
-      return safeParseJSON(latestEnvironmentVariablesRef.current);
+
+      return result;
     } catch (err) {
       if (err instanceof Error) {
         setLogs(prev => [
@@ -346,7 +388,7 @@ export function usePreflightScript(args: {
           },
         ]);
         setState(PreflightWorkerState.ready);
-        return safeParseJSON(latestEnvironmentVariablesRef.current);
+        return result;
       }
       throw err;
     }

packages/web/app/src/lib/preflight-sandbox/lab-api-declaration.ts

@@ -8,6 +8,22 @@
 // and use Prettier to format it, have syntax highlighting, etc.
 
 interface LabAPI {
+  /**
+   * Contains aspects of the request that you can manipulate before it is sent.
+   */
+  request: {
+    /**
+     * Headers that will be added to the request. They are merged
+     * using the following rules:
+     *
+     * 1. Do *not* interpolate environment variables.
+     *
+     * 2. Upon a collision with a base header, this header takes precedence.
+     *    This means that if the base headers contain "foo: bar" and you've added
+     *    "foo: qux" here, the final headers become "foo: qux" (*not* "foo: bar, qux").
+     */
+    headers: Headers;
+  };
   /**
    * [CryptoJS](https://cryptojs.gitbook.io/docs) library.
    */

packages/web/app/src/lib/preflight-sandbox/preflight-script-worker.ts

@@ -1,9 +1,16 @@
 import CryptoJS from 'crypto-js';
 import CryptoJSPackageJson from 'crypto-js/package.json';
 import { ALLOWED_GLOBALS } from './allowed-globals';
-import { isJSONPrimitive } from './json';
+import { isJSONPrimitive, JSONPrimitive } from './json';
 import { LogMessage, WorkerEvents } from './shared-types';
 
+interface WorkerData {
+  request: {
+    headers: Headers;
+  };
+  environmentVariables: Record<string, JSONPrimitive>;
+}
+
 /**
  * Unique id for each prompt request.
  * Incremented each time a prompt is requested.
@@ -45,11 +52,16 @@ async function execute(args: WorkerEvents.Incoming.EventData): Promise<void> {
     return;
   }
 
-  const { environmentVariables, script } = args;
+  const { script } = args;
 
-  // When running in worker `environmentVariables` will not be a reference to the main thread value
-  // but sometimes this will be tested outside the worker, so we don't want to mutate the input in that case
-  const workingEnvironmentVariables = { ...environmentVariables };
+  const workerData: WorkerData = {
+    request: {
+      headers: new Headers(),
+    },
+    // When running in worker `environmentVariables` will not be a reference to the main thread value
+    // but sometimes this will be tested outside the worker, so we don't want to mutate the input in that case
+    environmentVariables: { ...args.environmentVariables },
+  };
 
   // generate list of all in scope variables, we do getOwnPropertyNames and `for in` because each contain slightly different sets of keys
   const allGlobalKeys = Object.getOwnPropertyNames(globalThis);
@@ -123,17 +135,20 @@ async function execute(args: WorkerEvents.Incoming.EventData): Promise<void> {
     },
     environment: {
       get(key: string) {
-        return Object.freeze(workingEnvironmentVariables[key]);
+        return Object.freeze(workerData.environmentVariables[key]);
       },
       set(key: string, value: unknown) {
         const validValue = getValidEnvVariable(value);
         if (validValue === undefined) {
-          delete workingEnvironmentVariables[key];
+          delete workerData.environmentVariables[key];
         } else {
-          workingEnvironmentVariables[key] = validValue;
+          workerData.environmentVariables[key] = validValue;
         }
       },
     },
+    request: {
+      headers: workerData.request.headers,
+    },
     /**
      * Mimics the `prompt` function in the browser, by sending a message to the main thread
      * and waiting for a response.
@@ -179,9 +194,13 @@ ${script}})()`;
     });
     return;
   }
+
   sendMessage({
     type: WorkerEvents.Outgoing.Event.result,
-    environmentVariables: workingEnvironmentVariables,
+    environmentVariables: workerData.environmentVariables,
+    request: {
+      headers: Array.from(workerData.request.headers.entries()),
+    },
   });
 }