INN-2882 Add selective header forwarding when sending events

packages/inngest/src/components/Inngest.ts

@@ -350,6 +350,20 @@ export class Inngest<TClientOpts extends ClientOptions = ClientOptions> {
   public async send<Payload extends SendEventPayload<GetEvents<this>>>(
     payload: Payload
   ): Promise<SendEventOutput<TClientOpts>> {
+    return this._send({ payload });
+  }
+
+  /**
+   * Internal method for sending an event, used to allow Inngest internals to
+   * further customize the request sent to an Inngest Server.
+   */
+  private async _send<Payload extends SendEventPayload<GetEvents<this>>>({
+    payload,
+    headers,
+  }: {
+    payload: Payload;
+    headers?: Record<string, string>;
+  }): Promise<SendEventOutput<TClientOpts>> {
     const hooks = await getHookStack(
       this.middleware,
       "onSendEvent",
@@ -462,7 +476,7 @@ export class Inngest<TClientOpts extends ClientOptions = ClientOptions> {
     const response = await this.fetch(url, {
       method: "POST",
       body: stringify(payloads),
-      headers: { ...this.headers },
+      headers: { ...this.headers, ...headers },
     });
 
     let body: SendEventResponse | undefined;

packages/inngest/src/components/InngestCommHandler.ts

@@ -722,12 +722,35 @@ export class InngestCommHandler<
           (await getQuerystring("processing run request", queryKeys.StepId)) ||
           null;
 
+        const headersToForward = [
+          headerKeys.TraceParent,
+          headerKeys.TraceState,
+        ];
+
+        const headers = await headersToForward.reduce<
+          Promise<Record<string, string>>
+        >(async (acc, header) => {
+          const value = await actions.headers(
+            `fetching ${header} for forwarding`,
+            header
+          );
+          if (value) {
+            return {
+              ...(await acc),
+              [header]: value,
+            };
+          }
+
+          return acc;
+        }, Promise.resolve({}));
+
         const { version, result } = this.runStep({
           functionId: fnId,
           data: body,
           stepId,
           timer,
           reqArgs,
+          headers,
         });
         const stepOutput = await result;
 
@@ -901,12 +924,14 @@ export class InngestCommHandler<
     data,
     timer,
     reqArgs,
+    headers,
   }: {
     functionId: string;
     stepId: string | null;
     data: unknown;
     timer: ServerTiming;
     reqArgs: unknown[];
+    headers: Record<string, string>;
   }): { version: ExecutionVersion; result: Promise<ExecutionResult> } {
     const fn = this.fns[functionId];
     if (!fn) {
@@ -972,6 +997,7 @@ export class InngestCommHandler<
               isFailureHandler: fn.onFailure,
               stepCompletionOrder: ctx?.stack?.stack ?? [],
               reqArgs,
+              headers,
             },
           };
         },
@@ -1007,6 +1033,7 @@ export class InngestCommHandler<
               disableImmediateExecution: ctx?.disable_immediate_execution,
               stepCompletionOrder: ctx?.stack?.stack ?? [],
               reqArgs,
+              headers,
             },
           };
         },

packages/inngest/src/components/InngestFunction.test.ts

@@ -156,6 +156,7 @@ describe("runFn", () => {
                 stepState: {},
                 stepCompletionOrder: [],
                 reqArgs: [],
+                headers: {},
               },
             });
 
@@ -200,6 +201,7 @@ describe("runFn", () => {
                 runId: "run",
                 stepCompletionOrder: [],
                 reqArgs: [],
+                headers: {},
               },
             });
 

packages/inngest/src/components/InngestStepTools.ts

@@ -30,6 +30,7 @@ import {
 } from "./Inngest";
 import { InngestFunction } from "./InngestFunction";
 import { InngestFunctionReference } from "./InngestFunctionReference";
+import { type InngestExecution } from "./execution/InngestExecution";
 
 export interface FoundStep extends HashedOp {
   hashedId: string;
@@ -123,6 +124,7 @@ export const createStepTools = <
   TTriggers extends TriggersFromClient<TClient> = TriggersFromClient<TClient>,
 >(
   client: TClient,
+  execution: InngestExecution,
   stepHandler: StepHandler
 ) => {
   /**
@@ -201,7 +203,10 @@ export const createStepTools = <
       {
         nonStepExecuteInline: true,
         fn: (idOrOptions, payload) => {
-          return client.send(payload);
+          return client["_send"]({
+            payload,
+            headers: execution["options"]["headers"],
+          });
         },
       }
     ),

packages/inngest/src/components/execution/InngestExecution.ts

@@ -54,6 +54,11 @@ export interface InngestExecutionOptions {
   data: Omit<Context.Any, "step">;
   stepState: Record<string, MemoizedOp>;
   stepCompletionOrder: string[];
+
+  /**
+   * Headers to be sent with any request to Inngest during this execution.
+   */
+  headers: Record<string, string>;
   requestedRunStep?: string;
   timer?: ServerTiming;
   isFailureHandler?: boolean;
@@ -68,7 +73,6 @@ export class InngestExecution {
   protected debug: Debugger;
 
   constructor(protected options: InngestExecutionOptions) {
-    this.options = options;
     this.debug = Debug(debugPrefix).extend(this.options.runId);
   }
 }

packages/inngest/src/components/execution/v0.ts

@@ -417,7 +417,7 @@ export class V0InngestExecution
       });
     };
 
-    const step = createStepTools(this.options.client, stepHandler);
+    const step = createStepTools(this.options.client, this, stepHandler);
 
     let fnArg = {
       ...(this.options.data as { event: EventPayload }),

packages/inngest/src/components/execution/v1.ts

@@ -859,7 +859,7 @@ class V1InngestExecution extends InngestExecution implements IInngestExecution {
       return promise;
     };
 
-    return createStepTools(this.options.client, stepHandler);
+    return createStepTools(this.options.client, this, stepHandler);
   }
 
   private getUserFnToRun(): Handler.Any {

packages/inngest/src/helpers/consts.ts

@@ -122,6 +122,8 @@ export enum headerKeys {
   RetryAfter = "retry-after",
   InngestServerKind = "x-inngest-server-kind",
   InngestExpectedServerKind = "x-inngest-expected-server-kind",
+  TraceParent = "traceparent",
+  TraceState = "tracestate",
 }
 
 export const defaultInngestApiBaseUrl = "https://api.inngest.com/";

packages/inngest/src/test/helpers.ts

@@ -11,6 +11,8 @@ import {
 } from "@local/components/InngestStepTools";
 import {
   ExecutionVersion,
+  IInngestExecution,
+  InngestExecution,
   InngestExecutionOptions,
   PREFERRED_EXECUTION_VERSION,
 } from "@local/components/execution/InngestExecution";
@@ -64,9 +66,31 @@ export const createClient = <T extends ConstructorParameters<typeof Inngest>>(
 export const testClientId = "__test_client__";
 
 export const getStepTools = (
-  client: Inngest.Any = createClient({ id: testClientId })
+  client: Inngest.Any = createClient({ id: testClientId }),
+  executionOptions: Partial<InngestExecutionOptions> = {}
 ) => {
-  const step = createStepTools(client, ({ args, matchOp }) => {
+  const execution = client
+    .createFunction({ id: "test" }, { event: "test" }, () => undefined)
+    ["createExecution"]({
+      version: PREFERRED_EXECUTION_VERSION,
+      partialOptions: {
+        data: fromPartial({
+          event: { name: "foo", data: {} },
+        }),
+        runId: "run",
+        stepState: {},
+        stepCompletionOrder: [],
+        isFailureHandler: false,
+        requestedRunStep: undefined,
+        timer: new ServerTiming(),
+        disableImmediateExecution: false,
+        reqArgs: [],
+        headers: {},
+        ...executionOptions,
+      },
+    }) as IInngestExecution & InngestExecution;
+
+  const step = createStepTools(client, execution, ({ args, matchOp }) => {
     const stepOptions = getStepOptions(args[0]);
     return Promise.resolve(matchOp(stepOptions, ...args.slice(1)));
   });
@@ -106,6 +130,7 @@ export const runFnWithStack = (
       timer: new ServerTiming(),
       disableImmediateExecution: opts?.disableImmediateExecution,
       reqArgs: [],
+      headers: {},
     },
   });