Add support for headers

packages/connect-query-core/src/call-unary-method.ts

@@ -34,13 +34,14 @@ export async function callUnaryMethod<
   input: MessageInitShape<I> | undefined,
   options?: {
     signal?: AbortSignal;
+    headers?: HeadersInit;
   },
 ): Promise<MessageShape<O>> {
   const result = await transport.unary(
     schema,
     options?.signal,
     undefined,
-    undefined,
+    options?.headers,
     input ?? create(schema.input),
     undefined,
   );

packages/connect-query-core/src/connect-query-key.test.ts

@@ -178,6 +178,77 @@ describe("createConnectQueryKey", () => {
     });
   });
 
+  describe("headers", () => {
+    it("allows headers to be passed as an object", () => {
+      const key = createConnectQueryKey({
+        schema: ElizaService.method.say,
+        input: create(SayRequestSchema, { sentence: "hi" }),
+        cardinality: "finite",
+        headers: {
+          "x-custom-header": "custom-value",
+        },
+      });
+      expect(key[1].headers).toEqual({
+        "x-custom-header": "custom-value",
+      });
+    });
+    it("allows headers to be passed as a tuple", () => {
+      const key = createConnectQueryKey({
+        schema: ElizaService.method.say,
+        input: create(SayRequestSchema, { sentence: "hi" }),
+        cardinality: "finite",
+        headers: [["x-custom-header", "custom-value"]],
+      });
+      expect(key[1].headers).toEqual({
+        "x-custom-header": "custom-value",
+      });
+    });
+    it("allows headers to be passed as a HeadersInit", () => {
+      const key = createConnectQueryKey({
+        schema: ElizaService.method.say,
+        input: create(SayRequestSchema, { sentence: "hi" }),
+        cardinality: "finite",
+        headers: new Headers({
+          "x-custom-header": "custom-value",
+        }),
+      });
+      expect(key[1].headers).toEqual({
+        "x-custom-header": "custom-value",
+      });
+    });
+    it("normalizes header values", () => {
+      const keyA = createConnectQueryKey({
+        schema: ElizaService.method.say,
+        input: create(SayRequestSchema, { sentence: "hi" }),
+        cardinality: "finite",
+        headers: {
+          foo: "a",
+          Foo: "b",
+        },
+      });
+      const keyB = createConnectQueryKey({
+        schema: ElizaService.method.say,
+        input: create(SayRequestSchema, { sentence: "hi" }),
+        cardinality: "finite",
+        headers: {
+          foo: "a, b",
+        },
+      });
+      const keyC = createConnectQueryKey({
+        schema: ElizaService.method.say,
+        input: create(SayRequestSchema, { sentence: "hi" }),
+        cardinality: "finite",
+        headers: [
+          ["foo", "a"],
+          ["foo", "b"],
+        ],
+      });
+
+      expect(keyA[1].headers).toEqual(keyB[1].headers);
+      expect(keyA[1].headers).toEqual(keyC[1].headers);
+    });
+  });
+
   describe("infinite queries", () => {
     it("contains type hints to indicate the output type", () => {
       const sampleQueryClient = new QueryClient();

packages/connect-query-core/src/connect-query-key.ts

@@ -44,6 +44,11 @@ type SharedConnectQueryOptions = {
    * or "skipped".
    */
   input?: Record<string, unknown> | "skipped";
+  /**
+   * Headers to be sent with the request.
+   * Note that invalid HTTP header names will raise a TypeError, and that the Set-Cookie header is not supported.
+   */
+  headers?: Record<string, string>;
 };
 
 type InfiniteConnectQueryKey<OutputMessage extends DescMessage = DescMessage> =
@@ -125,6 +130,11 @@ type KeyParamsForMethod<Desc extends DescMethod> = {
    * If omit the field with this name from the key for infinite queries.
    */
   pageParamKey?: keyof MessageInitShape<Desc["input"]>;
+  /**
+   * Set `headers` in the key.
+   * Note that invalid HTTP header names will raise a TypeError, and that the Set-Cookie header is not supported.
+   */
+  headers?: HeadersInit;
 };
 
 type KeyParamsForService<Desc extends DescService> = {
@@ -220,6 +230,7 @@ export function createConnectQueryKey<
     transport?: string;
     cardinality?: "finite" | "infinite";
     input?: "skipped" | Record<string, unknown>;
+    headers?: Record<string, string>;
   } =
     params.schema.kind == "rpc"
       ? {
@@ -246,5 +257,25 @@ export function createConnectQueryKey<
       );
     }
   }
+  if (
+    params.schema.kind === "rpc" &&
+    "headers" in params &&
+    params.headers !== undefined
+  ) {
+    props.headers = createHeadersKey(params.headers);
+  }
   return ["connect-query", props] as ConnectQueryKey<O>;
 }
+
+/**
+ * Creates a record of headers from a HeadersInit object.
+ *
+ */
+function createHeadersKey(headers: HeadersInit): Record<string, string> {
+  const result: Record<string, string> = {};
+
+  for (const [key, value] of new Headers(headers)) {
+    result[key] = value;
+  }
+  return result;
+}

packages/connect-query-core/src/create-infinite-query-options.ts

@@ -49,6 +49,7 @@ export interface ConnectInfiniteQueryOptions<
     MessageInitShape<I>[ParamKey],
     MessageShape<O>
   >;
+  headers?: HeadersInit;
 }
 
 // eslint-disable-next-line @typescript-eslint/max-params -- we have 4 required arguments
@@ -79,6 +80,7 @@ function createUnaryInfiniteQueryFn<
     };
     return callUnaryMethod(transport, schema, inputCombinedWithPageParam, {
       signal: context.signal,
+      headers: context.queryKey[1].headers,
     });
   };
 }
@@ -97,6 +99,7 @@ export function createInfiniteQueryOptions<
     transport,
     getNextPageParam,
     pageParamKey,
+    headers,
   }: ConnectInfiniteQueryOptions<I, O, ParamKey> & { transport: Transport },
 ): {
   getNextPageParam: ConnectInfiniteQueryOptions<
@@ -124,6 +127,7 @@ export function createInfiniteQueryOptions<
     transport,
     getNextPageParam,
     pageParamKey,
+    headers,
   }: ConnectInfiniteQueryOptions<I, O, ParamKey> & { transport: Transport },
 ): {
   getNextPageParam: ConnectInfiniteQueryOptions<
@@ -149,6 +153,7 @@ export function createInfiniteQueryOptions<
     transport,
     getNextPageParam,
     pageParamKey,
+    headers,
   }: ConnectInfiniteQueryOptions<I, O, ParamKey> & { transport: Transport },
 ): {
   getNextPageParam: ConnectInfiniteQueryOptions<
@@ -180,6 +185,7 @@ export function createInfiniteQueryOptions<
     transport,
     getNextPageParam,
     pageParamKey,
+    headers,
   }: ConnectInfiniteQueryOptions<I, O, ParamKey> & { transport: Transport },
 ): {
   getNextPageParam: ConnectInfiniteQueryOptions<
@@ -203,6 +209,7 @@ export function createInfiniteQueryOptions<
     schema,
     transport,
     input,
+    headers,
   });
   const structuralSharing = createStructuralSharing(schema.output);
   const queryFn =

packages/connect-query-core/src/create-query-options.test.ts

@@ -48,12 +48,18 @@ describe("createQueryOptions", () => {
       input: { sentence: "hi" },
       transport: mockedElizaTransport,
       cardinality: "finite",
+      headers: {
+        "x-custom-header": "custom-value",
+      },
     });
     const opt = createQueryOptions(
       sayMethodDescriptor,
       { sentence: "hi" },
       {
         transport: mockedElizaTransport,
+        headers: {
+          "x-custom-header": "custom-value",
+        },
       },
     );
     expect(opt.queryKey).toStrictEqual(want);

packages/connect-query-core/src/create-query-options.ts

@@ -36,6 +36,7 @@ function createUnaryQueryFn<I extends DescMessage, O extends DescMessage>(
   return async (context) => {
     return callUnaryMethod(transport, schema, input, {
       signal: context.signal,
+      headers: context.queryKey[1].headers,
     });
   };
 }
@@ -51,8 +52,10 @@ export function createQueryOptions<
   input: MessageInitShape<I> | undefined,
   {
     transport,
+    headers,
   }: {
     transport: Transport;
+    headers?: HeadersInit;
   },
 ): {
   queryKey: ConnectQueryKey<O>;
@@ -67,8 +70,10 @@ export function createQueryOptions<
   input: SkipToken,
   {
     transport,
+    headers,
   }: {
     transport: Transport;
+    headers?: HeadersInit;
   },
 ): {
   queryKey: ConnectQueryKey<O>;
@@ -83,8 +88,10 @@ export function createQueryOptions<
   input: SkipToken | MessageInitShape<I> | undefined,
   {
     transport,
+    headers,
   }: {
     transport: Transport;
+    headers?: HeadersInit;
   },
 ): {
   queryKey: ConnectQueryKey<O>;
@@ -99,8 +106,10 @@ export function createQueryOptions<
   input: SkipToken | MessageInitShape<I> | undefined,
   {
     transport,
+    headers,
   }: {
     transport: Transport;
+    headers?: HeadersInit;
   },
 ): {
   queryKey: ConnectQueryKey<O>;
@@ -112,6 +121,7 @@ export function createQueryOptions<
     input: input ?? create(schema.input),
     transport,
     cardinality: "finite",
+    headers,
   });
   const structuralSharing = createStructuralSharing(schema.output);
   const queryFn =

packages/connect-query/src/call-unary-method.test.ts

@@ -72,4 +72,42 @@ describe("callUnaryMethod", () => {
     });
     expect(result.current.query1.data?.sentence).toEqual("Response 1");
   });
+  it("can pass headers through", async () => {
+    let resolve: () => void;
+    const promise = new Promise<void>((res) => {
+      resolve = res;
+    });
+    const transport = mockEliza(
+      {
+        sentence: "Response 1",
+      },
+      false,
+      {
+        router: {
+          interceptors: [
+            (next) => (req) => {
+              expect(req.header.get("x-custom-header")).toEqual("custom-value");
+              resolve();
+              return next(req);
+            },
+          ],
+        },
+      },
+    );
+    const input: SayRequest = create(SayRequestSchema, {
+      sentence: "query 1",
+    });
+    const res = await callUnaryMethod(
+      transport,
+      ElizaService.method.say,
+      input,
+      {
+        headers: {
+          "x-custom-header": "custom-value",
+        },
+      },
+    );
+    await promise;
+    expect(res.sentence).toEqual("Response 1");
+  });
 });

packages/connect-query/src/use-infinite-query.test.ts

@@ -396,4 +396,53 @@ describe("useSuspenseInfiniteQuery", () => {
       wrapper({}, mockedPaginatedTransport),
     );
   });
+
+  it("can pass headers through", async () => {
+    let resolve: () => void;
+    const promise = new Promise<void>((res) => {
+      resolve = res;
+    });
+    const transport = mockPaginatedTransport(
+      {
+        items: ["Intercepted!"],
+        page: 0n,
+      },
+      false,
+      {
+        router: {
+          interceptors: [
+            (next) => (req) => {
+              expect(req.header.get("x-custom-header")).toEqual("custom-value");
+              resolve();
+              return next(req);
+            },
+          ],
+        },
+      },
+    );
+    const { result } = renderHook(() => {
+      return useSuspenseInfiniteQuery(
+        methodDescriptor,
+        {
+          page: 0n,
+        },
+        {
+          getNextPageParam: (lastPage) => lastPage.page + 1n,
+          pageParamKey: "page",
+          transport,
+          headers: {
+            "x-custom-header": "custom-value",
+          },
+        },
+      );
+    }, wrapper({}));
+
+    await waitFor(() => {
+      expect(result.current.isSuccess).toBeTruthy();
+    });
+
+    await promise;
+
+    expect(result.current.data.pages[0].items).toEqual(["Intercepted!"]);
+  });
 });

packages/connect-query/src/use-infinite-query.ts

@@ -128,6 +128,7 @@ export function useSuspenseInfiniteQuery<
     transport,
     pageParamKey,
     getNextPageParam,
+    headers,
     ...queryOptions
   }: UseSuspenseInfiniteQueryOptions<I, O, ParamKey>,
 ): UseSuspenseInfiniteQueryResult<InfiniteData<MessageShape<O>>, ConnectError> {
@@ -136,6 +137,7 @@ export function useSuspenseInfiniteQuery<
     transport: transport ?? transportFromCtx,
     getNextPageParam,
     pageParamKey,
+    headers,
   });
   return tsUseSuspenseInfiniteQuery({
     ...baseOptions,