From 1bac3af6aad85ccb8a40a828c71b4978eec02f41 Mon Sep 17 00:00:00 2001
From: Paul Sachs <psachs@buf.build>
Date: Tue, 29 Apr 2025 12:02:49 -0400
Subject: [PATCH 1/9] Add datatag to connect key so we can infer type info

Signed-off-by: Paul Sachs <psachs@buf.build>
---
 .../src/connect-query-key.test.ts               | 17 +++++++++++++++--
 .../connect-query-core/src/connect-query-key.ts | 16 +++++++++-------
 2 files changed, 24 insertions(+), 9 deletions(-)

diff --git a/packages/connect-query-core/src/connect-query-key.test.ts b/packages/connect-query-core/src/connect-query-key.test.ts
index 19379650..4330628a 100644
--- a/packages/connect-query-core/src/connect-query-key.test.ts
+++ b/packages/connect-query-core/src/connect-query-key.test.ts
@@ -14,14 +14,15 @@
 
 import { create } from "@bufbuild/protobuf";
 import type { Transport } from "@connectrpc/connect";
-import { ElizaService, SayRequestSchema } from "test-utils/gen/eliza_pb.js";
+import { ElizaService, SayRequestSchema, type SayResponse } from "test-utils/gen/eliza_pb.js";
 import { ListRequestSchema, ListService } from "test-utils/gen/list_pb.js";
-import { describe, expect, it } from "vitest";
+import { describe, expect, expectTypeOf, it } from "vitest";
 
 import { createConnectQueryKey } from "./connect-query-key.js";
 import { skipToken } from "./index.js";
 import { createMessageKey } from "./message-key.js";
 import { createTransportKey } from "./transport-key.js";
+import { QueryClient } from "@tanstack/query-core";
 
 describe("createConnectQueryKey", () => {
   const fakeTransport: Transport = {
@@ -133,4 +134,16 @@ describe("createConnectQueryKey", () => {
       cardinality: undefined,
     });
   });
+
+  it("contains type hints to indicate the output type", () => {
+    const sampleQueryClient = new QueryClient();
+    const key = createConnectQueryKey({
+      schema: ElizaService.method.say,
+      input: create(SayRequestSchema, { sentence: "hi" }),
+      cardinality: "finite",
+    });
+    const data = sampleQueryClient.getQueryData(key);
+    
+    expectTypeOf(data).toEqualTypeOf<SayResponse | undefined>();
+  })
 });
diff --git a/packages/connect-query-core/src/connect-query-key.ts b/packages/connect-query-core/src/connect-query-key.ts
index c7e3d58f..8986e6a9 100644
--- a/packages/connect-query-core/src/connect-query-key.ts
+++ b/packages/connect-query-core/src/connect-query-key.ts
@@ -18,9 +18,10 @@ import type {
   DescMethodUnary,
   DescService,
   MessageInitShape,
+  MessageShape,
 } from "@bufbuild/protobuf";
-import type { Transport } from "@connectrpc/connect";
-import type { SkipToken } from "@tanstack/query-core";
+import type { ConnectError, Transport } from "@connectrpc/connect";
+import type { DataTag, SkipToken } from "@tanstack/query-core";
 
 import { createMessageKey } from "./message-key.js";
 import { createTransportKey } from "./transport-key.js";
@@ -44,7 +45,7 @@ import { createTransportKey } from "./transport-key.js";
  *   }
  * ]
  */
-export type ConnectQueryKey = [
+export type ConnectQueryKey<OutputMessage extends DescMessage> = DataTag<[
   /**
    * To distinguish Connect query keys from other query keys, they always start with the string "connect-query".
    */
@@ -72,7 +73,7 @@ export type ConnectQueryKey = [
      */
     cardinality?: "infinite" | "finite" | undefined;
   },
-];
+], MessageShape<OutputMessage>, ConnectError>;
 
 type KeyParamsForMethod<Desc extends DescMethod> = {
   /**
@@ -158,8 +159,9 @@ export function createConnectQueryKey<
   Desc extends DescService,
 >(
   params: KeyParamsForMethod<DescMethodUnary<I, O>> | KeyParamsForService<Desc>,
-): ConnectQueryKey {
-  const props: ConnectQueryKey[1] =
+): ConnectQueryKey<O>
+ {
+  const props: ConnectQueryKey<O>[1] =
     params.schema.kind == "rpc"
       ? {
           serviceName: params.schema.parent.typeName,
@@ -185,5 +187,5 @@ export function createConnectQueryKey<
       );
     }
   }
-  return ["connect-query", props];
+  return ["connect-query", props] as ConnectQueryKey<O>;
 }

From 80f175efd117f56999b57ec1ade9799dd4f413bb Mon Sep 17 00:00:00 2001
From: Paul Sachs <psachs@buf.build>
Date: Tue, 29 Apr 2025 14:04:27 -0400
Subject: [PATCH 2/9] Fix some connectkey type references to include param

Signed-off-by: Paul Sachs <psachs@buf.build>
---
 .../src/connect-query-key.test.ts             | 10 ++-
 .../src/connect-query-key.ts                  | 66 ++++++++++---------
 .../src/create-infinite-query-options.ts      | 16 ++---
 .../src/create-query-options.ts               | 16 ++---
 .../connect-query/src/use-infinite-query.ts   |  4 +-
 packages/connect-query/src/use-query.ts       |  4 +-
 6 files changed, 62 insertions(+), 54 deletions(-)

diff --git a/packages/connect-query-core/src/connect-query-key.test.ts b/packages/connect-query-core/src/connect-query-key.test.ts
index 4330628a..e18d3bdb 100644
--- a/packages/connect-query-core/src/connect-query-key.test.ts
+++ b/packages/connect-query-core/src/connect-query-key.test.ts
@@ -14,7 +14,11 @@
 
 import { create } from "@bufbuild/protobuf";
 import type { Transport } from "@connectrpc/connect";
-import { ElizaService, SayRequestSchema, type SayResponse } from "test-utils/gen/eliza_pb.js";
+import {
+  ElizaService,
+  SayRequestSchema,
+  type SayResponse,
+} from "test-utils/gen/eliza_pb.js";
 import { ListRequestSchema, ListService } from "test-utils/gen/list_pb.js";
 import { describe, expect, expectTypeOf, it } from "vitest";
 
@@ -143,7 +147,7 @@ describe("createConnectQueryKey", () => {
       cardinality: "finite",
     });
     const data = sampleQueryClient.getQueryData(key);
-    
+
     expectTypeOf(data).toEqualTypeOf<SayResponse | undefined>();
-  })
+  });
 });
diff --git a/packages/connect-query-core/src/connect-query-key.ts b/packages/connect-query-core/src/connect-query-key.ts
index 8986e6a9..72191670 100644
--- a/packages/connect-query-core/src/connect-query-key.ts
+++ b/packages/connect-query-core/src/connect-query-key.ts
@@ -45,35 +45,40 @@ import { createTransportKey } from "./transport-key.js";
  *   }
  * ]
  */
-export type ConnectQueryKey<OutputMessage extends DescMessage> = DataTag<[
-  /**
-   * To distinguish Connect query keys from other query keys, they always start with the string "connect-query".
-   */
-  "connect-query",
-  {
-    /**
-     * A key for a Transport reference, created with createTransportKey().
-     */
-    transport?: string;
-    /**
-     * The name of the service, e.g. connectrpc.eliza.v1.ElizaService
-     */
-    serviceName: string;
-    /**
-     * The name of the method, e.g. Say.
-     */
-    methodName?: string;
-    /**
-     * A key for the request message, created with createMessageKey(),
-     * or "skipped".
-     */
-    input?: Record<string, unknown> | "skipped";
-    /**
-     * Whether this is an infinite query, or a regular one.
-     */
-    cardinality?: "infinite" | "finite" | undefined;
-  },
-], MessageShape<OutputMessage>, ConnectError>;
+export type ConnectQueryKey<OutputMessage extends DescMessage = DescMessage> =
+  DataTag<
+    [
+      /**
+       * To distinguish Connect query keys from other query keys, they always start with the string "connect-query".
+       */
+      "connect-query",
+      {
+        /**
+         * A key for a Transport reference, created with createTransportKey().
+         */
+        transport?: string;
+        /**
+         * The name of the service, e.g. connectrpc.eliza.v1.ElizaService
+         */
+        serviceName: string;
+        /**
+         * The name of the method, e.g. Say.
+         */
+        methodName?: string;
+        /**
+         * A key for the request message, created with createMessageKey(),
+         * or "skipped".
+         */
+        input?: Record<string, unknown> | "skipped";
+        /**
+         * Whether this is an infinite query, or a regular one.
+         */
+        cardinality?: "infinite" | "finite" | undefined;
+      },
+    ],
+    MessageShape<OutputMessage>,
+    ConnectError
+  >;
 
 type KeyParamsForMethod<Desc extends DescMethod> = {
   /**
@@ -159,8 +164,7 @@ export function createConnectQueryKey<
   Desc extends DescService,
 >(
   params: KeyParamsForMethod<DescMethodUnary<I, O>> | KeyParamsForService<Desc>,
-): ConnectQueryKey<O>
- {
+): ConnectQueryKey<O> {
   const props: ConnectQueryKey<O>[1] =
     params.schema.kind == "rpc"
       ? {
diff --git a/packages/connect-query-core/src/create-infinite-query-options.ts b/packages/connect-query-core/src/create-infinite-query-options.ts
index f4aa2d5f..176ef108 100644
--- a/packages/connect-query-core/src/create-infinite-query-options.ts
+++ b/packages/connect-query-core/src/create-infinite-query-options.ts
@@ -67,7 +67,7 @@ function createUnaryInfiniteQueryFn<
   },
 ): QueryFunction<
   MessageShape<O>,
-  ConnectQueryKey,
+  ConnectQueryKey<O>,
   MessageInitShape<I>[ParamKey]
 > {
   return async (context) => {
@@ -104,10 +104,10 @@ export function createInfiniteQueryOptions<
     O,
     ParamKey
   >["getNextPageParam"];
-  queryKey: ConnectQueryKey;
+  queryKey: ConnectQueryKey<O>;
   queryFn: QueryFunction<
     MessageShape<O>,
-    ConnectQueryKey,
+    ConnectQueryKey<O>,
     MessageInitShape<I>[ParamKey]
   >;
   structuralSharing: (oldData: unknown, newData: unknown) => unknown;
@@ -131,7 +131,7 @@ export function createInfiniteQueryOptions<
     O,
     ParamKey
   >["getNextPageParam"];
-  queryKey: ConnectQueryKey;
+  queryKey: ConnectQueryKey<O>;
   queryFn: SkipToken;
   structuralSharing: (oldData: unknown, newData: unknown) => unknown;
   initialPageParam: MessageInitShape<I>[ParamKey];
@@ -156,11 +156,11 @@ export function createInfiniteQueryOptions<
     O,
     ParamKey
   >["getNextPageParam"];
-  queryKey: ConnectQueryKey;
+  queryKey: ConnectQueryKey<O>;
   queryFn:
     | QueryFunction<
         MessageShape<O>,
-        ConnectQueryKey,
+        ConnectQueryKey<O>,
         MessageInitShape<I>[ParamKey]
       >
     | SkipToken;
@@ -187,11 +187,11 @@ export function createInfiniteQueryOptions<
     O,
     ParamKey
   >["getNextPageParam"];
-  queryKey: ConnectQueryKey;
+  queryKey: ConnectQueryKey<O>;
   queryFn:
     | QueryFunction<
         MessageShape<O>,
-        ConnectQueryKey,
+        ConnectQueryKey<O>,
         MessageInitShape<I>[ParamKey]
       >
     | SkipToken;
diff --git a/packages/connect-query-core/src/create-query-options.ts b/packages/connect-query-core/src/create-query-options.ts
index 3d7069b8..6b8b0e7e 100644
--- a/packages/connect-query-core/src/create-query-options.ts
+++ b/packages/connect-query-core/src/create-query-options.ts
@@ -32,7 +32,7 @@ function createUnaryQueryFn<I extends DescMessage, O extends DescMessage>(
   transport: Transport,
   schema: DescMethodUnary<I, O>,
   input: MessageInitShape<I> | undefined,
-): QueryFunction<MessageShape<O>, ConnectQueryKey> {
+): QueryFunction<MessageShape<O>, ConnectQueryKey<O>> {
   return async (context) => {
     return callUnaryMethod(transport, schema, input, {
       signal: context.signal,
@@ -55,8 +55,8 @@ export function createQueryOptions<
     transport: Transport;
   },
 ): {
-  queryKey: ConnectQueryKey;
-  queryFn: QueryFunction<MessageShape<O>, ConnectQueryKey>;
+  queryKey: ConnectQueryKey<O>;
+  queryFn: QueryFunction<MessageShape<O>, ConnectQueryKey<O>>;
   structuralSharing: (oldData: unknown, newData: unknown) => unknown;
 };
 export function createQueryOptions<
@@ -71,7 +71,7 @@ export function createQueryOptions<
     transport: Transport;
   },
 ): {
-  queryKey: ConnectQueryKey;
+  queryKey: ConnectQueryKey<O>;
   queryFn: SkipToken;
   structuralSharing: (oldData: unknown, newData: unknown) => unknown;
 };
@@ -87,8 +87,8 @@ export function createQueryOptions<
     transport: Transport;
   },
 ): {
-  queryKey: ConnectQueryKey;
-  queryFn: QueryFunction<MessageShape<O>, ConnectQueryKey> | SkipToken;
+  queryKey: ConnectQueryKey<O>;
+  queryFn: QueryFunction<MessageShape<O>, ConnectQueryKey<O>> | SkipToken;
   structuralSharing: (oldData: unknown, newData: unknown) => unknown;
 };
 export function createQueryOptions<
@@ -103,8 +103,8 @@ export function createQueryOptions<
     transport: Transport;
   },
 ): {
-  queryKey: ConnectQueryKey;
-  queryFn: QueryFunction<MessageShape<O>, ConnectQueryKey> | SkipToken;
+  queryKey: ConnectQueryKey<O>;
+  queryFn: QueryFunction<MessageShape<O>, ConnectQueryKey<O>> | SkipToken;
   structuralSharing: (oldData: unknown, newData: unknown) => unknown;
 } {
   const queryKey = createConnectQueryKey({
diff --git a/packages/connect-query/src/use-infinite-query.ts b/packages/connect-query/src/use-infinite-query.ts
index f3d71ce8..afa84fb3 100644
--- a/packages/connect-query/src/use-infinite-query.ts
+++ b/packages/connect-query/src/use-infinite-query.ts
@@ -52,7 +52,7 @@ export type UseInfiniteQueryOptions<
     ConnectError,
     InfiniteData<MessageShape<O>>,
     MessageShape<O>,
-    ConnectQueryKey,
+    ConnectQueryKey<O>,
     MessageInitShape<I>[ParamKey]
   >,
   "getNextPageParam" | "initialPageParam" | "queryFn" | "queryKey"
@@ -106,7 +106,7 @@ export type UseSuspenseInfiniteQueryOptions<
     ConnectError,
     InfiniteData<MessageShape<O>>,
     MessageShape<O>,
-    ConnectQueryKey,
+    ConnectQueryKey<O>,
     MessageInitShape<I>[ParamKey]
   >,
   "getNextPageParam" | "initialPageParam" | "queryFn" | "queryKey"
diff --git a/packages/connect-query/src/use-query.ts b/packages/connect-query/src/use-query.ts
index 6db782dc..3675af22 100644
--- a/packages/connect-query/src/use-query.ts
+++ b/packages/connect-query/src/use-query.ts
@@ -48,7 +48,7 @@ export type UseQueryOptions<
     MessageShape<O>,
     ConnectError,
     SelectOutData,
-    ConnectQueryKey
+    ConnectQueryKey<O>
   >,
   "queryFn" | "queryKey"
 > & {
@@ -89,7 +89,7 @@ export type UseSuspenseQueryOptions<
     MessageShape<O>,
     ConnectError,
     SelectOutData,
-    ConnectQueryKey
+    ConnectQueryKey<O>
   >,
   "queryFn" | "queryKey"
 > & {

From 58e75aadaad808c3408df07bf8d11e45062ef680 Mon Sep 17 00:00:00 2001
From: Paul Sachs <psachs@buf.build>
Date: Tue, 29 Apr 2025 16:18:23 -0400
Subject: [PATCH 3/9] Add tests and support for infinite data

Signed-off-by: Paul Sachs <psachs@buf.build>
---
 .../src/connect-query-key.test.ts             |  67 +++++++++-
 .../src/connect-query-key.ts                  | 124 ++++++++++++------
 2 files changed, 152 insertions(+), 39 deletions(-)

diff --git a/packages/connect-query-core/src/connect-query-key.test.ts b/packages/connect-query-core/src/connect-query-key.test.ts
index e18d3bdb..cee37a48 100644
--- a/packages/connect-query-core/src/connect-query-key.test.ts
+++ b/packages/connect-query-core/src/connect-query-key.test.ts
@@ -17,6 +17,7 @@ import type { Transport } from "@connectrpc/connect";
 import {
   ElizaService,
   SayRequestSchema,
+  SayResponseSchema,
   type SayResponse,
 } from "test-utils/gen/eliza_pb.js";
 import { ListRequestSchema, ListService } from "test-utils/gen/list_pb.js";
@@ -26,7 +27,7 @@ import { createConnectQueryKey } from "./connect-query-key.js";
 import { skipToken } from "./index.js";
 import { createMessageKey } from "./message-key.js";
 import { createTransportKey } from "./transport-key.js";
-import { QueryClient } from "@tanstack/query-core";
+import { type InfiniteData, QueryClient } from "@tanstack/query-core";
 
 describe("createConnectQueryKey", () => {
   const fakeTransport: Transport = {
@@ -116,6 +117,7 @@ describe("createConnectQueryKey", () => {
       schema: ElizaService.method.say,
       cardinality: undefined,
     });
+    // @ts-expect-error(2322) cardinality is type undefined as well as actually undefined
     expect(key[1].cardinality).toBeUndefined();
   });
 
@@ -130,9 +132,9 @@ describe("createConnectQueryKey", () => {
 
   it("cannot except invalid input", () => {
     createConnectQueryKey({
+      // @ts-expect-error(2322) cannot create a key with invalid input
       schema: ElizaService.method.say,
       input: {
-        // @ts-expect-error(2322) cannot create a key with invalid input
         sentence: 1,
       },
       cardinality: undefined,
@@ -150,4 +152,65 @@ describe("createConnectQueryKey", () => {
 
     expectTypeOf(data).toEqualTypeOf<SayResponse | undefined>();
   });
+
+  it("supports typesafe data updaters", () => {
+    const sampleQueryClient = new QueryClient();
+    const key = createConnectQueryKey({
+      schema: ElizaService.method.say,
+      input: create(SayRequestSchema, { sentence: "hi" }),
+      cardinality: "finite",
+    });
+    // @ts-expect-error(2345) this is a test to check if the type is correct
+    sampleQueryClient.setQueryData(key, { sentence: 1 });
+    // @ts-expect-error(2345) $typename is required
+    sampleQueryClient.setQueryData(key, {
+      sentence: "a proper value but missing $typename",
+    });
+    sampleQueryClient.setQueryData(
+      key,
+      create(SayResponseSchema, { sentence: "a proper value" })
+    );
+  });
+
+  describe("infinite queries", () => {
+    it("contains type hints to indicate the output type", () => {
+      const sampleQueryClient = new QueryClient();
+      const key = createConnectQueryKey({
+        schema: ElizaService.method.say,
+        input: create(SayRequestSchema, { sentence: "hi" }),
+        cardinality: "infinite",
+      });
+      const data = sampleQueryClient.getQueryData(key);
+
+      expectTypeOf(data).toEqualTypeOf<InfiniteData<SayResponse> | undefined>();
+    });
+
+    it("supports typesafe data updaters", () => {
+      const sampleQueryClient = new QueryClient();
+      const key = createConnectQueryKey({
+        schema: ElizaService.method.say,
+        input: create(SayRequestSchema, { sentence: "hi" }),
+        cardinality: "infinite",
+      });
+      sampleQueryClient.setQueryData(key, {
+        pages: [
+          // @ts-expect-error(2345) make sure the shape is as expected
+          { sentence: 1 },
+        ],
+      });
+      sampleQueryClient.setQueryData(key, {
+        // @ts-expect-error(2345) $typename is required
+        pages: [{ sentence: "a proper value but missing $typename" }],
+      });
+      sampleQueryClient.setQueryData(
+        key,
+        {
+          pageParams: [0],
+          pages: [
+            create(SayResponseSchema, { sentence: "a proper value" }),
+          ]
+        }
+      );
+    });
+  });
 });
diff --git a/packages/connect-query-core/src/connect-query-key.ts b/packages/connect-query-core/src/connect-query-key.ts
index 72191670..253440aa 100644
--- a/packages/connect-query-core/src/connect-query-key.ts
+++ b/packages/connect-query-core/src/connect-query-key.ts
@@ -21,11 +21,57 @@ import type {
   MessageShape,
 } from "@bufbuild/protobuf";
 import type { ConnectError, Transport } from "@connectrpc/connect";
-import type { DataTag, SkipToken } from "@tanstack/query-core";
+import type { DataTag, InfiniteData, SkipToken } from "@tanstack/query-core";
 
 import { createMessageKey } from "./message-key.js";
 import { createTransportKey } from "./transport-key.js";
 
+type SharedConnectQueryOptions = {
+  /**
+   * A key for a Transport reference, created with createTransportKey().
+   */
+  transport?: string;
+  /**
+   * The name of the service, e.g. connectrpc.eliza.v1.ElizaService
+   */
+  serviceName: string;
+  /**
+   * The name of the method, e.g. Say.
+   */
+  methodName?: string;
+  /**
+   * A key for the request message, created with createMessageKey(),
+   * or "skipped".
+   */
+  input?: Record<string, unknown> | "skipped";
+};
+
+type InfiniteConnectQueryKey<OutputMessage extends DescMessage = DescMessage> =
+  DataTag<
+    [
+      "connect-query",
+      SharedConnectQueryOptions & {
+        /** This data represents a infinite, paged result */
+        cardinality: "infinite";
+      },
+    ],
+    InfiniteData<MessageShape<OutputMessage>>,
+    ConnectError
+  >;
+
+type FiniteConnectQueryKey<OutputMessage extends DescMessage = DescMessage> =
+  DataTag<
+    [
+      "connect-query",
+      SharedConnectQueryOptions & {
+        /** This data represents a finite result */
+        cardinality: "finite";
+      },
+    ],
+    MessageShape<OutputMessage>,
+    ConnectError
+  >;
+
 /**
  * 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.
  *
@@ -46,39 +92,9 @@ import { createTransportKey } from "./transport-key.js";
  * ]
  */
 export type ConnectQueryKey<OutputMessage extends DescMessage = DescMessage> =
-  DataTag<
-    [
-      /**
-       * To distinguish Connect query keys from other query keys, they always start with the string "connect-query".
-       */
-      "connect-query",
-      {
-        /**
-         * A key for a Transport reference, created with createTransportKey().
-         */
-        transport?: string;
-        /**
-         * The name of the service, e.g. connectrpc.eliza.v1.ElizaService
-         */
-        serviceName: string;
-        /**
-         * The name of the method, e.g. Say.
-         */
-        methodName?: string;
-        /**
-         * A key for the request message, created with createMessageKey(),
-         * or "skipped".
-         */
-        input?: Record<string, unknown> | "skipped";
-        /**
-         * Whether this is an infinite query, or a regular one.
-         */
-        cardinality?: "infinite" | "finite" | undefined;
-      },
-    ],
-    MessageShape<OutputMessage>,
-    ConnectError
-  >;
+  | InfiniteConnectQueryKey<OutputMessage>
+  | FiniteConnectQueryKey<OutputMessage>
+  | ["connect-query", SharedConnectQueryOptions];
 
 type KeyParamsForMethod<Desc extends DescMethod> = {
   /**
@@ -158,14 +174,48 @@ type KeyParamsForService<Desc extends DescService> = {
  *
  * @see ConnectQueryKey for information on the components of Connect-Query's keys.
  */
+export function createConnectQueryKey<
+  I extends DescMessage,
+  O extends DescMessage,
+>(
+  params: KeyParamsForMethod<DescMethodUnary<I, O>> & {
+    cardinality: "finite";
+  }
+): FiniteConnectQueryKey<O>;
+export function createConnectQueryKey<
+  I extends DescMessage,
+  O extends DescMessage,
+>(
+  params: KeyParamsForMethod<DescMethodUnary<I, O>> & {
+    cardinality: "infinite";
+  }
+): InfiniteConnectQueryKey<O>;
+export function createConnectQueryKey<
+  I extends DescMessage,
+  O extends DescMessage,
+>(
+  params: KeyParamsForMethod<DescMethodUnary<I, O>> & {
+    cardinality: undefined;
+  }
+): ConnectQueryKey<O>;
+export function createConnectQueryKey<
+  O extends DescMessage,
+  Desc extends DescService,
+>(params: KeyParamsForService<Desc>): ConnectQueryKey<O>;
 export function createConnectQueryKey<
   I extends DescMessage,
   O extends DescMessage,
   Desc extends DescService,
 >(
-  params: KeyParamsForMethod<DescMethodUnary<I, O>> | KeyParamsForService<Desc>,
+  params: KeyParamsForMethod<DescMethodUnary<I, O>> | KeyParamsForService<Desc>
 ): ConnectQueryKey<O> {
-  const props: ConnectQueryKey<O>[1] =
+  const props: {
+    serviceName: string;
+    methodName?: string;
+    transport?: string;
+    cardinality?: "finite" | "infinite";
+    input?: "skipped" | Record<string, unknown>;
+  } =
     params.schema.kind == "rpc"
       ? {
           serviceName: params.schema.parent.typeName,
@@ -187,7 +237,7 @@ export function createConnectQueryKey<
       props.input = createMessageKey(
         params.schema.input,
         params.input,
-        params.pageParamKey,
+        params.pageParamKey
       );
     }
   }

From 1d61a0daf427a70a48d01582a48e73952f03b14b Mon Sep 17 00:00:00 2001
From: Paul Sachs <psachs@buf.build>
Date: Tue, 29 Apr 2025 16:22:44 -0400
Subject: [PATCH 4/9] Ran formatter

Signed-off-by: Paul Sachs <psachs@buf.build>
---
 .../src/connect-query-key.test.ts                 | 15 +++++----------
 .../connect-query-core/src/connect-query-key.ts   | 10 +++++-----
 2 files changed, 10 insertions(+), 15 deletions(-)

diff --git a/packages/connect-query-core/src/connect-query-key.test.ts b/packages/connect-query-core/src/connect-query-key.test.ts
index cee37a48..d6596200 100644
--- a/packages/connect-query-core/src/connect-query-key.test.ts
+++ b/packages/connect-query-core/src/connect-query-key.test.ts
@@ -168,7 +168,7 @@ describe("createConnectQueryKey", () => {
     });
     sampleQueryClient.setQueryData(
       key,
-      create(SayResponseSchema, { sentence: "a proper value" })
+      create(SayResponseSchema, { sentence: "a proper value" }),
     );
   });
 
@@ -202,15 +202,10 @@ describe("createConnectQueryKey", () => {
         // @ts-expect-error(2345) $typename is required
         pages: [{ sentence: "a proper value but missing $typename" }],
       });
-      sampleQueryClient.setQueryData(
-        key,
-        {
-          pageParams: [0],
-          pages: [
-            create(SayResponseSchema, { sentence: "a proper value" }),
-          ]
-        }
-      );
+      sampleQueryClient.setQueryData(key, {
+        pageParams: [0],
+        pages: [create(SayResponseSchema, { sentence: "a proper value" })],
+      });
     });
   });
 });
diff --git a/packages/connect-query-core/src/connect-query-key.ts b/packages/connect-query-core/src/connect-query-key.ts
index 253440aa..d02381ff 100644
--- a/packages/connect-query-core/src/connect-query-key.ts
+++ b/packages/connect-query-core/src/connect-query-key.ts
@@ -180,7 +180,7 @@ export function createConnectQueryKey<
 >(
   params: KeyParamsForMethod<DescMethodUnary<I, O>> & {
     cardinality: "finite";
-  }
+  },
 ): FiniteConnectQueryKey<O>;
 export function createConnectQueryKey<
   I extends DescMessage,
@@ -188,7 +188,7 @@ export function createConnectQueryKey<
 >(
   params: KeyParamsForMethod<DescMethodUnary<I, O>> & {
     cardinality: "infinite";
-  }
+  },
 ): InfiniteConnectQueryKey<O>;
 export function createConnectQueryKey<
   I extends DescMessage,
@@ -196,7 +196,7 @@ export function createConnectQueryKey<
 >(
   params: KeyParamsForMethod<DescMethodUnary<I, O>> & {
     cardinality: undefined;
-  }
+  },
 ): ConnectQueryKey<O>;
 export function createConnectQueryKey<
   O extends DescMessage,
@@ -207,7 +207,7 @@ export function createConnectQueryKey<
   O extends DescMessage,
   Desc extends DescService,
 >(
-  params: KeyParamsForMethod<DescMethodUnary<I, O>> | KeyParamsForService<Desc>
+  params: KeyParamsForMethod<DescMethodUnary<I, O>> | KeyParamsForService<Desc>,
 ): ConnectQueryKey<O> {
   const props: {
     serviceName: string;
@@ -237,7 +237,7 @@ export function createConnectQueryKey<
       props.input = createMessageKey(
         params.schema.input,
         params.input,
-        params.pageParamKey
+        params.pageParamKey,
       );
     }
   }

From 71ca6234c62576135ca288896d777f44937e66d3 Mon Sep 17 00:00:00 2001
From: Paul Sachs <psachs@buf.build>
Date: Tue, 29 Apr 2025 16:26:59 -0400
Subject: [PATCH 5/9] Add ignore for vitest temp file

Signed-off-by: Paul Sachs <psachs@buf.build>
---
 .gitignore | 1 +
 1 file changed, 1 insertion(+)

diff --git a/.gitignore b/.gitignore
index a77f46a2..2e42cf58 100644
--- a/.gitignore
+++ b/.gitignore
@@ -3,3 +3,4 @@
 node_modules
 /packages/*/dist
 /packages/*/coverage
+tsconfig.vitest-temp.json
\ No newline at end of file

From 8bcf4a1dad701a1bc0a6c2a64bd34edf9f4fbbda Mon Sep 17 00:00:00 2001
From: Paul Sachs <psachs@buf.build>
Date: Tue, 29 Apr 2025 16:33:06 -0400
Subject: [PATCH 6/9] Add tests with functional callback formats

Signed-off-by: Paul Sachs <psachs@buf.build>
---
 .../src/connect-query-key.test.ts                | 16 ++++++++++++++++
 1 file changed, 16 insertions(+)

diff --git a/packages/connect-query-core/src/connect-query-key.test.ts b/packages/connect-query-core/src/connect-query-key.test.ts
index d6596200..27d26046 100644
--- a/packages/connect-query-core/src/connect-query-key.test.ts
+++ b/packages/connect-query-core/src/connect-query-key.test.ts
@@ -170,6 +170,13 @@ describe("createConnectQueryKey", () => {
       key,
       create(SayResponseSchema, { sentence: "a proper value" }),
     );
+
+    sampleQueryClient.setQueryData(key, (prev) => {
+      expectTypeOf(prev).toEqualTypeOf<SayResponse | undefined>();
+      return create(SayResponseSchema, {
+        sentence: "a proper value",
+      });
+    });
   });
 
   describe("infinite queries", () => {
@@ -206,6 +213,15 @@ describe("createConnectQueryKey", () => {
         pageParams: [0],
         pages: [create(SayResponseSchema, { sentence: "a proper value" })],
       });
+      sampleQueryClient.setQueryData(key, (prev) => {
+        expectTypeOf(prev).toEqualTypeOf<
+          InfiniteData<SayResponse> | undefined
+        >();
+        return {
+          pageParams: [0],
+          pages: [create(SayResponseSchema, { sentence: "a proper value" })],
+        };
+      });
     });
   });
 });

From 995218f7f6988a6799e3cf8838872cc3fb2314b0 Mon Sep 17 00:00:00 2001
From: Paul Sachs <psachs@buf.build>
Date: Tue, 20 May 2025 10:29:32 -0400
Subject: [PATCH 7/9] Add cardinality to empty connect query key

Signed-off-by: Paul Sachs <psachs@buf.build>
---
 packages/connect-query-core/src/connect-query-key.test.ts | 1 -
 packages/connect-query-core/src/connect-query-key.ts      | 7 ++++++-
 2 files changed, 6 insertions(+), 2 deletions(-)

diff --git a/packages/connect-query-core/src/connect-query-key.test.ts b/packages/connect-query-core/src/connect-query-key.test.ts
index 27d26046..857e9a06 100644
--- a/packages/connect-query-core/src/connect-query-key.test.ts
+++ b/packages/connect-query-core/src/connect-query-key.test.ts
@@ -117,7 +117,6 @@ describe("createConnectQueryKey", () => {
       schema: ElizaService.method.say,
       cardinality: undefined,
     });
-    // @ts-expect-error(2322) cardinality is type undefined as well as actually undefined
     expect(key[1].cardinality).toBeUndefined();
   });
 
diff --git a/packages/connect-query-core/src/connect-query-key.ts b/packages/connect-query-core/src/connect-query-key.ts
index d02381ff..9f7623f3 100644
--- a/packages/connect-query-core/src/connect-query-key.ts
+++ b/packages/connect-query-core/src/connect-query-key.ts
@@ -94,7 +94,12 @@ type FiniteConnectQueryKey<OutputMessage extends DescMessage = DescMessage> =
 export type ConnectQueryKey<OutputMessage extends DescMessage = DescMessage> =
   | InfiniteConnectQueryKey<OutputMessage>
   | FiniteConnectQueryKey<OutputMessage>
-  | ["connect-query", SharedConnectQueryOptions];
+  | [
+      "connect-query",
+      SharedConnectQueryOptions & {
+        cardinality: undefined;
+      },
+    ];
 
 type KeyParamsForMethod<Desc extends DescMethod> = {
   /**

From 1fdf1f39ff75eeffa7f9599fcf1c6953b5a7feb4 Mon Sep 17 00:00:00 2001
From: Paul Sachs <psachs@buf.build>
Date: Tue, 20 May 2025 11:24:09 -0400
Subject: [PATCH 8/9] Updated peer dep and readme with clarifying details

Signed-off-by: Paul Sachs <psachs@buf.build>
---
 README.md                                | 30 ++++++++++--------------
 package-lock.json                        |  4 ++--
 packages/connect-query-core/package.json |  2 +-
 packages/connect-query-core/src/utils.ts |  2 ++
 packages/connect-query/package.json      |  2 +-
 5 files changed, 18 insertions(+), 22 deletions(-)

diff --git a/README.md b/README.md
index 0174f92c..ee5163d5 100644
--- a/README.md
+++ b/README.md
@@ -260,19 +260,17 @@ Any additional `options` you pass to `useMutation` will be merged with the optio
 
 ### `createConnectQueryKey`
 
-```ts
-function createConnectQueryKey<Desc extends DescMethod | DescService>(
-  params: KeyParams<Desc>,
-): ConnectQueryKey;
-```
-
-This function is used under the hood of `useQuery` and other hooks to compute a [`queryKey`](https://tanstack.com/query/v4/docs/react/guides/query-keys) for TanStack Query. You can use it to create (partial) keys yourself to filter queries.
+This function is used under the hood of `useQuery` and other hooks to compute a [`queryKey`](https://tanstack.com/query/v4/docs/react/guides/query-keys) for TanStack Query. You can use it to create keys yourself to filter queries.
 
 `useQuery` creates a query key with the following parameters:
 
 1. The qualified name of the RPC.
 2. The transport being used.
 3. The request message.
+4. The cardinality of the RPC (either "finite" or "infinite").
+5. Adds a DataTag which brands the key with the associated data type of the response.
+
+The DataTag type allows @tanstack/react-query functions to properly infer the type of the data returned by the query. This is useful for things like `QueryClient.setQueryData` and `QueryClient.getQueryData`.
 
 To create the same key manually, you simply provide the same parameters:
 
@@ -355,7 +353,7 @@ function callUnaryMethod<I extends DescMessage, O extends DescMessage>(
 
 This API allows you to directly call the method using the provided transport. Use this if you need to manually call a method outside of the context of a React component, or need to call it where you can't use hooks.
 
-### `createProtobufSafeUpdater`
+### `createProtobufSafeUpdater` (deprecated)
 
 Creates a typesafe updater that can be used to update data in a query cache. Used in combination with a queryClient.
 
@@ -387,6 +385,8 @@ queryClient.setQueryData(
 
 ```
 
+** Note: This API is deprecated and will be removed in a future version. `ConnectQueryKey` now contains type information to make it safer to use `setQueryData` directly. **
+
 ### `createQueryOptions`
 
 ```ts
@@ -599,21 +599,15 @@ Connect-Query does require React, but the core (`createConnectQueryKey` and `cal
 
 ### How do I do Prefetching?
 
-When you might not have access to React context, you can use the `create` series of functions and provide a transport directly. For example:
+When you might not have access to React context, you can use `createQueryOptions` and provide a transport directly. For example:
 
 ```ts
 import { say } from "./gen/eliza-ElizaService_connectquery";
 
 function prefetch() {
-  return queryClient.prefetchQuery({
-    queryKey: createConnectQueryKey({
-      schema: say,
-      transport: myTransport,
-      input: { sentence: "Hello" },
-      cardinality: "finite",
-    }),
-    queryFn: () => callUnaryMethod(myTransport, say, { sentence: "Hello" }),
-  });
+  return queryClient.prefetchQuery(
+    createQueryOptions(say, { sentence: "Hello" }, { transport: myTransport })
+  );
 }
 ```
 
diff --git a/package-lock.json b/package-lock.json
index 80906232..85218402 100644
--- a/package-lock.json
+++ b/package-lock.json
@@ -9536,7 +9536,7 @@
       "peerDependencies": {
         "@bufbuild/protobuf": "2.x",
         "@connectrpc/connect": "^2.0.1",
-        "@tanstack/react-query": "5.x",
+        "@tanstack/react-query": ">=5.62.0",
         "react": "^18 || ^19",
         "react-dom": "^18 || ^19"
       }
@@ -9560,7 +9560,7 @@
       "peerDependencies": {
         "@bufbuild/protobuf": "2.x",
         "@connectrpc/connect": "^2.0.1",
-        "@tanstack/query-core": "5.x"
+        "@tanstack/query-core": ">=5.62.0"
       }
     },
     "packages/examples/react/basic": {
diff --git a/packages/connect-query-core/package.json b/packages/connect-query-core/package.json
index f9e53d73..e9f7c80a 100644
--- a/packages/connect-query-core/package.json
+++ b/packages/connect-query-core/package.json
@@ -44,7 +44,7 @@
   "peerDependencies": {
     "@bufbuild/protobuf": "2.x",
     "@connectrpc/connect": "^2.0.1",
-    "@tanstack/query-core": "5.x"
+    "@tanstack/query-core": ">=5.62.0"
   },
   "files": [
     "dist/**"
diff --git a/packages/connect-query-core/src/utils.ts b/packages/connect-query-core/src/utils.ts
index b51b6a89..1f6e498b 100644
--- a/packages/connect-query-core/src/utils.ts
+++ b/packages/connect-query-core/src/utils.ts
@@ -60,6 +60,8 @@ export type ConnectUpdater<O extends DescMessage> =
 
 /**
  * This helper makes sure that the type for the original response message is returned.
+ * 
+ * @deprecated the ConnectQueryKey type now links to the return data type so `setQueryData` can be called safely without this helper.
  */
 export const createProtobufSafeUpdater =
   <O extends DescMessage>(
diff --git a/packages/connect-query/package.json b/packages/connect-query/package.json
index f610582e..6433705f 100644
--- a/packages/connect-query/package.json
+++ b/packages/connect-query/package.json
@@ -52,7 +52,7 @@
   "peerDependencies": {
     "@bufbuild/protobuf": "2.x",
     "@connectrpc/connect": "^2.0.1",
-    "@tanstack/react-query": "5.x",
+    "@tanstack/react-query": ">=5.62.0",
     "react": "^18 || ^19",
     "react-dom": "^18 || ^19"
   },

From b59790e4294bcf52bfc8b501bcfad42ed013331a Mon Sep 17 00:00:00 2001
From: Paul Sachs <psachs@buf.build>
Date: Tue, 20 May 2025 12:34:17 -0400
Subject: [PATCH 9/9] Format

Signed-off-by: Paul Sachs <psachs@buf.build>
---
 README.md                                | 2 +-
 packages/connect-query-core/src/utils.ts | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/README.md b/README.md
index ee5163d5..7589d550 100644
--- a/README.md
+++ b/README.md
@@ -606,7 +606,7 @@ import { say } from "./gen/eliza-ElizaService_connectquery";
 
 function prefetch() {
   return queryClient.prefetchQuery(
-    createQueryOptions(say, { sentence: "Hello" }, { transport: myTransport })
+    createQueryOptions(say, { sentence: "Hello" }, { transport: myTransport }),
   );
 }
 ```
diff --git a/packages/connect-query-core/src/utils.ts b/packages/connect-query-core/src/utils.ts
index 1f6e498b..ffd7d7f5 100644
--- a/packages/connect-query-core/src/utils.ts
+++ b/packages/connect-query-core/src/utils.ts
@@ -60,7 +60,7 @@ export type ConnectUpdater<O extends DescMessage> =
 
 /**
  * This helper makes sure that the type for the original response message is returned.
- * 
+ *
  * @deprecated the ConnectQueryKey type now links to the return data type so `setQueryData` can be called safely without this helper.
  */
 export const createProtobufSafeUpdater =