Add support for context values

README.md

@@ -193,6 +193,7 @@ Use this helper to get the default transport that's currently attached to the Re
 > [!TIP]
 >
 > All hooks accept a `transport` in the options. You can use the Transport from the context, or create one dynamically. If you create a Transport dynamically, make sure to memoize it, because it is taken into consideration when building query keys.
+> All hooks also accept `contextValues`, which can be used to pass additional context on each call to the transport and any interceptors.
 
 ### `useQuery`
 
@@ -204,7 +205,11 @@ function useQuery<
 >(
   schema: DescMethodUnary<I, O>,
   input?: SkipToken | MessageInitShape<I>,
-  { transport, ...queryOptions }: UseQueryOptions<I, O, SelectOutData> = {},
+  {
+    transport,
+    contextValues,
+    ...queryOptions
+  }: UseQueryOptions<I, O, SelectOutData> = {},
 ): UseQueryResult<SelectOutData, ConnectError>;
 ```
 
@@ -232,6 +237,7 @@ function useInfiniteQuery<
     transport,
     pageParamKey,
     getNextPageParam,
+    contextValues,
     ...queryOptions
   }: UseInfiniteQueryOptions<I, O, ParamKey>,
 ): UseInfiniteQueryResult<InfiniteData<MessageShape<O>>, ConnectError>;
@@ -250,7 +256,11 @@ Identical to useInfiniteQuery but mapping to the `useSuspenseInfiniteQuery` hook
 ```ts
 function useMutation<I extends DescMessage, O extends DescMessage>(
   schema: DescMethodUnary<I, O>,
-  { transport, ...queryOptions }: UseMutationOptions<I, O, Ctx> = {},
+  {
+    transport,
+    contextValues,
+    ...queryOptions
+  }: UseMutationOptions<I, O, Ctx> = {},
 ): UseMutationResult<MessageShape<O>, ConnectError, PartialMessage<I>>;
 ```
 
@@ -349,6 +359,7 @@ function callUnaryMethod<I extends DescMessage, O extends DescMessage>(
   input: MessageInitShape<I> | undefined,
   options?: {
     signal?: AbortSignal;
+    contextValues?: SerializableContextValues;
   },
 ): Promise<O>;
 ```
@@ -395,8 +406,10 @@ function createQueryOptions<I extends DescMessage, O extends DescMessage>(
   input: SkipToken | PartialMessage<I> | undefined,
   {
     transport,
+    contextValues,
   }: {
     transport: Transport;
+    contextValues?: SerializableContextValues;
   },
 ): {
   queryKey: ConnectQueryKey;
@@ -508,6 +521,10 @@ type ConnectQueryKey = [
      * Whether this is an infinite query, or a regular one.
      */
     cardinality?: "infinite" | "finite";
+    /**
+     * The stringified version of contextValues, if present.
+     */
+    contextValues?: string;
   },
 ];
 ```

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

@@ -20,6 +20,7 @@ import type {
 } from "@bufbuild/protobuf";
 import { create } from "@bufbuild/protobuf";
 import type { Transport } from "@connectrpc/connect";
+import type { SerializableContextValues } from "./serializable-context-values.js";
 
 /**
  * Call a unary method given its signature and input.
@@ -34,6 +35,7 @@ export async function callUnaryMethod<
   input: MessageInitShape<I> | undefined,
   options?: {
     signal?: AbortSignal;
+    contextValues?: SerializableContextValues;
   },
 ): Promise<MessageShape<O>> {
   const result = await transport.unary(
@@ -42,7 +44,7 @@ export async function callUnaryMethod<
     undefined,
     undefined,
     input ?? create(schema.input),
-    undefined,
+    options?.contextValues,
   );
   return result.message;
 }

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

@@ -13,13 +13,13 @@
 // limitations under the License.
 
 import { create } from "@bufbuild/protobuf";
-import type { Transport } from "@connectrpc/connect";
+import { createContextValues, type Transport } from "@connectrpc/connect";
 import { ElizaService, SayRequestSchema } from "test-utils/gen/eliza_pb.js";
 import { ListRequestSchema, ListService } from "test-utils/gen/list_pb.js";
 import { describe, expect, it } from "vitest";
 
 import { createConnectQueryKey } from "./connect-query-key.js";
-import { skipToken } from "./index.js";
+import { skipToken, type SerializableContextValues } from "./index.js";
 import { createMessageKey } from "./message-key.js";
 import { createTransportKey } from "./transport-key.js";
 
@@ -133,4 +133,21 @@ describe("createConnectQueryKey", () => {
       cardinality: undefined,
     });
   });
+
+  it("allows to set contextValues", () => {
+    const baseContextValues = createContextValues();
+    const fakeContextValues: SerializableContextValues = {
+      ...baseContextValues,
+      toString() {
+        return "serialized";
+      },
+    };
+    const key = createConnectQueryKey({
+      schema: ElizaService.method.say,
+      input: skipToken,
+      cardinality: "finite",
+      contextValues: fakeContextValues,
+    });
+    expect(key[1].contextValues).toEqual("serialized");
+  });
 });

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

@@ -24,6 +24,7 @@ import type { SkipToken } from "@tanstack/query-core";
 
 import { createMessageKey } from "./message-key.js";
 import { createTransportKey } from "./transport-key.js";
+import type { SerializableContextValues } from "./serializable-context-values.js";
 
 /**
  * TanStack Query manages query caching for you based on query keys. `QueryKey`s in TanStack Query are arrays with arbitrary JSON-serializable data - typically handwritten for each endpoint.
@@ -71,6 +72,10 @@ export type ConnectQueryKey = [
      * Whether this is an infinite query, or a regular one.
      */
     cardinality?: "infinite" | "finite" | undefined;
+    /**
+     * The stringified version of contextValues, if present.
+     */
+    contextValues?: string;
   },
 ];
 
@@ -98,6 +103,10 @@ type KeyParamsForMethod<Desc extends DescMethod> = {
    * If omit the field with this name from the key for infinite queries.
    */
   pageParamKey?: keyof MessageInitShape<Desc["input"]>;
+  /**
+   * Any contextValues that need to be passed to the query.
+   */
+  contextValues?: SerializableContextValues;
 };
 
 type KeyParamsForService<Desc extends DescService> = {
@@ -168,6 +177,10 @@ export function createConnectQueryKey<
       : {
           serviceName: params.schema.typeName,
         };
+
+  if ("contextValues" in params && params.contextValues !== undefined) {
+    props.contextValues = params.contextValues.toString();
+  }
   if (params.transport !== undefined) {
     props.transport = createTransportKey(params.transport);
   }

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

@@ -33,6 +33,7 @@ import {
 } from "./connect-query-key.js";
 import { createStructuralSharing } from "./structural-sharing.js";
 import { assert } from "./utils.js";
+import type { SerializableContextValues } from "./serializable-context-values.js";
 
 /**
  * Options specific to connect-query
@@ -49,6 +50,7 @@ export interface ConnectInfiniteQueryOptions<
     MessageInitShape<I>[ParamKey],
     MessageShape<O>
   >;
+  contextValues?: SerializableContextValues;
 }
 
 // eslint-disable-next-line @typescript-eslint/max-params -- we have 4 required arguments
@@ -62,8 +64,10 @@ function createUnaryInfiniteQueryFn<
   input: MessageInitShape<I>,
   {
     pageParamKey,
+    contextValues,
   }: {
     pageParamKey: ParamKey;
+    contextValues?: SerializableContextValues;
   },
 ): QueryFunction<
   MessageShape<O>,
@@ -79,6 +83,7 @@ function createUnaryInfiniteQueryFn<
     };
     return callUnaryMethod(transport, schema, inputCombinedWithPageParam, {
       signal: context.signal,
+      contextValues: contextValues,
     });
   };
 }
@@ -97,6 +102,7 @@ export function createInfiniteQueryOptions<
     transport,
     getNextPageParam,
     pageParamKey,
+    contextValues,
   }: ConnectInfiniteQueryOptions<I, O, ParamKey> & { transport: Transport },
 ): {
   getNextPageParam: ConnectInfiniteQueryOptions<
@@ -149,6 +155,7 @@ export function createInfiniteQueryOptions<
     transport,
     getNextPageParam,
     pageParamKey,
+    contextValues,
   }: ConnectInfiniteQueryOptions<I, O, ParamKey> & { transport: Transport },
 ): {
   getNextPageParam: ConnectInfiniteQueryOptions<
@@ -180,6 +187,7 @@ export function createInfiniteQueryOptions<
     transport,
     getNextPageParam,
     pageParamKey,
+    contextValues,
   }: ConnectInfiniteQueryOptions<I, O, ParamKey> & { transport: Transport },
 ): {
   getNextPageParam: ConnectInfiniteQueryOptions<
@@ -203,13 +211,15 @@ export function createInfiniteQueryOptions<
     schema,
     transport,
     input,
+    contextValues,
   });
   const structuralSharing = createStructuralSharing(schema.output);
   const queryFn =
     input === skipToken
       ? skipToken
       : createUnaryInfiniteQueryFn(transport, schema, input, {
           pageParamKey,
+          contextValues,
         });
   return {
     getNextPageParam,

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

@@ -27,15 +27,18 @@ import { callUnaryMethod } from "./call-unary-method.js";
 import type { ConnectQueryKey } from "./connect-query-key.js";
 import { createConnectQueryKey } from "./connect-query-key.js";
 import { createStructuralSharing } from "./structural-sharing.js";
+import type { SerializableContextValues } from "./serializable-context-values.js";
 
 function createUnaryQueryFn<I extends DescMessage, O extends DescMessage>(
   transport: Transport,
   schema: DescMethodUnary<I, O>,
   input: MessageInitShape<I> | undefined,
+  contextValues?: SerializableContextValues,
 ): QueryFunction<MessageShape<O>, ConnectQueryKey> {
   return async (context) => {
     return callUnaryMethod(transport, schema, input, {
       signal: context.signal,
+      contextValues,
     });
   };
 }
@@ -51,8 +54,10 @@ export function createQueryOptions<
   input: MessageInitShape<I> | undefined,
   {
     transport,
+    contextValues,
   }: {
     transport: Transport;
+    contextValues?: SerializableContextValues;
   },
 ): {
   queryKey: ConnectQueryKey;
@@ -67,8 +72,10 @@ export function createQueryOptions<
   input: SkipToken,
   {
     transport,
+    contextValues,
   }: {
     transport: Transport;
+    contextValues?: SerializableContextValues;
   },
 ): {
   queryKey: ConnectQueryKey;
@@ -83,8 +90,10 @@ export function createQueryOptions<
   input: SkipToken | MessageInitShape<I> | undefined,
   {
     transport,
+    contextValues,
   }: {
     transport: Transport;
+    contextValues?: SerializableContextValues;
   },
 ): {
   queryKey: ConnectQueryKey;
@@ -99,8 +108,10 @@ export function createQueryOptions<
   input: SkipToken | MessageInitShape<I> | undefined,
   {
     transport,
+    contextValues,
   }: {
     transport: Transport;
+    contextValues?: SerializableContextValues;
   },
 ): {
   queryKey: ConnectQueryKey;
@@ -112,12 +123,13 @@ export function createQueryOptions<
     input: input ?? create(schema.input),
     transport,
     cardinality: "finite",
+    contextValues,
   });
   const structuralSharing = createStructuralSharing(schema.output);
   const queryFn =
     input === skipToken
       ? skipToken
-      : createUnaryQueryFn(transport, schema, input);
+      : createUnaryQueryFn(transport, schema, input, contextValues);
   return {
     queryKey,
     queryFn,

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

@@ -23,3 +23,4 @@ export { createQueryOptions } from "./create-query-options.js";
 export { addStaticKeyToTransport } from "./transport-key.js";
 export type { SkipToken } from "@tanstack/query-core";
 export { skipToken } from "@tanstack/query-core";
+export type { SerializableContextValues } from "./serializable-context-values.js";

packages/connect-query-core/src/serializable-context-values.ts

@@ -0,0 +1,24 @@
+// Copyright 2021-2023 The Connect Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+import { type ContextValues } from "@connectrpc/connect";
+
+/**
+ * A version of ContextValues that can be serialized to a string.
+ * This string will be used to detect if the context has changed
+ * and the query needs to be invalidated.
+ */
+export interface SerializableContextValues extends ContextValues {
+  toString(): string;
+}

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

@@ -78,6 +78,7 @@ export function useInfiniteQuery<
     transport,
     pageParamKey,
     getNextPageParam,
+    contextValues,
     ...queryOptions
   }: UseInfiniteQueryOptions<I, O, ParamKey>,
 ): UseInfiniteQueryResult<InfiniteData<MessageShape<O>>, ConnectError> {
@@ -86,6 +87,7 @@ export function useInfiniteQuery<
     transport: transport ?? transportFromCtx,
     getNextPageParam,
     pageParamKey,
+    contextValues,
   });
   return tsUseInfiniteQuery({
     ...baseOptions,