From 764ad7f72b1dfa3b6c3c90ac730926e84fa7e759 Mon Sep 17 00:00:00 2001
From: unnoq <contact@unnoq.com>
Date: Mon, 21 Jul 2025 09:23:44 +0700
Subject: [PATCH] fix(tanstack-query): client context should include inside
 query/mutation key

Queries have different client context can behave differently, so we should treat them as separate queries
---
 packages/tanstack-query/src/key.test.ts       |  2 ++
 packages/tanstack-query/src/key.ts            | 13 ++++---
 .../src/procedure-utils.test.ts               | 16 +++++----
 .../tanstack-query/src/procedure-utils.ts     | 35 ++++++++++++------
 packages/tanstack-query/src/types.ts          | 36 ++++++++++++++-----
 5 files changed, 71 insertions(+), 31 deletions(-)

diff --git a/packages/tanstack-query/src/key.test.ts b/packages/tanstack-query/src/key.test.ts
index be885ad2c..4ce15d5fa 100644
--- a/packages/tanstack-query/src/key.test.ts
+++ b/packages/tanstack-query/src/key.test.ts
@@ -8,4 +8,6 @@ it('generateOperationKey', () => {
     .toEqual([['planet', 'find'], { type: 'query', input: { id: 1 } }])
   expect(generateOperationKey(['planet', 'stream'], { type: 'streamed', input: { cursor: 0 }, fnOptions: { refetchMode: 'append' } }))
     .toEqual([['planet', 'stream'], { type: 'streamed', input: { cursor: 0 }, fnOptions: { refetchMode: 'append' } }])
+  expect(generateOperationKey(['planet', 'find'], { type: 'query', input: { id: 1 }, context: { cache: true } }))
+    .toEqual([['planet', 'find'], { type: 'query', input: { id: 1 }, context: { cache: true } }])
 })
diff --git a/packages/tanstack-query/src/key.ts b/packages/tanstack-query/src/key.ts
index b1abe8de6..71555a520 100644
--- a/packages/tanstack-query/src/key.ts
+++ b/packages/tanstack-query/src/key.ts
@@ -1,12 +1,17 @@
+import type { ClientContext } from '@orpc/client'
 import type { OperationKey, OperationKeyOptions, OperationType } from './types'
 
-export function generateOperationKey<TType extends OperationType, TInput>(
+/**
+ * @todo move TClientContext to second position + remove default in next major version
+ */
+export function generateOperationKey<TType extends OperationType, TInput, TClientContext extends ClientContext = ClientContext>(
   path: readonly string[],
-  state: OperationKeyOptions<TType, TInput> = {},
-): OperationKey<TType, TInput> {
+  state: OperationKeyOptions<TType, TInput, TClientContext> = {},
+): OperationKey<TType, TInput, TClientContext> {
   return [path, {
-    ...state.input !== undefined ? { input: state.input } : {},
     ...state.type !== undefined ? { type: state.type } : {},
+    ...state.context !== undefined ? { context: state.context } : {},
     ...state.fnOptions !== undefined ? { fnOptions: state.fnOptions } : {},
+    ...state.input !== undefined ? { input: state.input } : {},
   } as any]
 }
diff --git a/packages/tanstack-query/src/procedure-utils.test.ts b/packages/tanstack-query/src/procedure-utils.test.ts
index 60ecf7f4f..6f52700f4 100644
--- a/packages/tanstack-query/src/procedure-utils.test.ts
+++ b/packages/tanstack-query/src/procedure-utils.test.ts
@@ -50,7 +50,7 @@ describe('createProcedureUtils', () => {
 
       expect(options.queryKey).toBe(generateOperationKeySpy.mock.results[0]!.value)
       expect(generateOperationKeySpy).toHaveBeenCalledTimes(1)
-      expect(generateOperationKeySpy).toHaveBeenCalledWith(['ping'], { type: 'query', input: { search: '__search__' } })
+      expect(generateOperationKeySpy).toHaveBeenCalledWith(['ping'], { type: 'query', context: { batch: '__batch__' }, input: { search: '__search__' } })
 
       client.mockResolvedValueOnce('__output__')
       await expect(options.queryFn!({ signal } as any)).resolves.toEqual('__output__')
@@ -71,7 +71,7 @@ describe('createProcedureUtils', () => {
 
       expect(options.queryKey).toBe(generateOperationKeySpy.mock.results[0]!.value)
       expect(generateOperationKeySpy).toHaveBeenCalledTimes(1)
-      expect(generateOperationKeySpy).toHaveBeenCalledWith(['ping'], { type: 'query', input: skipToken })
+      expect(generateOperationKeySpy).toHaveBeenCalledWith(['ping'], { type: 'query', context: { batch: '__batch__' }, input: skipToken })
 
       expect(() => options.queryFn!({ signal } as any)).toThrow('queryFn should not be called with skipToken used as input')
       expect(client).toHaveBeenCalledTimes(0)
@@ -103,6 +103,7 @@ describe('createProcedureUtils', () => {
       expect(generateOperationKeySpy).toHaveBeenCalledTimes(1)
       expect(generateOperationKeySpy).toHaveBeenCalledWith(['ping'], {
         type: 'streamed',
+        context: { batch: '__batch__' },
         input: { search: '__search__' },
         fnOptions: {
           refetchMode: 'replace',
@@ -144,7 +145,7 @@ describe('createProcedureUtils', () => {
 
       expect(options.queryKey).toBe(generateOperationKeySpy.mock.results[0]!.value)
       expect(generateOperationKeySpy).toHaveBeenCalledTimes(1)
-      expect(generateOperationKeySpy).toHaveBeenCalledWith(['ping'], { type: 'streamed', input: skipToken })
+      expect(generateOperationKeySpy).toHaveBeenCalledWith(['ping'], { type: 'streamed', context: { batch: '__batch__' }, input: skipToken })
 
       await expect(options.queryFn!({ signal, client: queryClient } as any)).rejects.toThrow('queryFn should not be called with skipToken used as input')
       expect(client).toHaveBeenCalledTimes(0)
@@ -188,6 +189,7 @@ describe('createProcedureUtils', () => {
       expect(generateOperationKeySpy).toHaveBeenCalledTimes(1)
       expect(generateOperationKeySpy).toHaveBeenCalledWith(['ping'], {
         type: 'live',
+        context: { batch: '__batch__' },
         input: { search: '__search__' },
       })
 
@@ -223,7 +225,7 @@ describe('createProcedureUtils', () => {
 
       expect(options.queryKey).toBe(generateOperationKeySpy.mock.results[0]!.value)
       expect(generateOperationKeySpy).toHaveBeenCalledTimes(1)
-      expect(generateOperationKeySpy).toHaveBeenCalledWith(['ping'], { type: 'live', input: skipToken })
+      expect(generateOperationKeySpy).toHaveBeenCalledWith(['ping'], { type: 'live', context: { batch: '__batch__' }, input: skipToken })
 
       await expect(options.queryFn!({ signal, client: queryClient } as any)).rejects.toThrow('queryFn should not be called with skipToken used as input')
       expect(client).toHaveBeenCalledTimes(0)
@@ -272,7 +274,7 @@ describe('createProcedureUtils', () => {
 
       expect(options.queryKey).toBe(generateOperationKeySpy.mock.results[0]!.value)
       expect(generateOperationKeySpy).toHaveBeenCalledTimes(1)
-      expect(generateOperationKeySpy).toHaveBeenCalledWith(['ping'], { type: 'infinite', input: { search: '__search__', pageParam: '__initialPageParam__' } })
+      expect(generateOperationKeySpy).toHaveBeenCalledWith(['ping'], { type: 'infinite', context: { batch: '__batch__' }, input: { search: '__search__', pageParam: '__initialPageParam__' } })
 
       expect(options.initialPageParam).toEqual('__initialPageParam__')
       expect(options.getNextPageParam).toBe(getNextPageParam)
@@ -309,7 +311,7 @@ describe('createProcedureUtils', () => {
 
       expect(options.queryKey).toBe(generateOperationKeySpy.mock.results[0]!.value)
       expect(generateOperationKeySpy).toHaveBeenCalledTimes(1)
-      expect(generateOperationKeySpy).toHaveBeenCalledWith(['ping'], { type: 'infinite', input: skipToken })
+      expect(generateOperationKeySpy).toHaveBeenCalledWith(['ping'], { type: 'infinite', context: { batch: '__batch__' }, input: skipToken })
 
       expect(options.initialPageParam).toEqual('__initialPageParam__')
       expect(options.getNextPageParam).toBe(getNextPageParam)
@@ -335,7 +337,7 @@ describe('createProcedureUtils', () => {
 
     expect(options.mutationKey).toBe(generateOperationKeySpy.mock.results[0]!.value)
     expect(generateOperationKeySpy).toHaveBeenCalledTimes(1)
-    expect(generateOperationKeySpy).toHaveBeenCalledWith(['ping'], { type: 'mutation' })
+    expect(generateOperationKeySpy).toHaveBeenCalledWith(['ping'], { type: 'mutation', context: { batch: '__batch__' } })
 
     client.mockResolvedValueOnce('__output__')
     await expect(options.mutationFn!('__input__')).resolves.toEqual('__output__')
diff --git a/packages/tanstack-query/src/procedure-utils.ts b/packages/tanstack-query/src/procedure-utils.ts
index be8bae01e..848aa9089 100644
--- a/packages/tanstack-query/src/procedure-utils.ts
+++ b/packages/tanstack-query/src/procedure-utils.ts
@@ -7,6 +7,7 @@ import type {
   experimental_StreamedQueryOutput,
   InfiniteOptionsBase,
   InfiniteOptionsIn,
+  MutationKeyOptions,
   MutationOptions,
   MutationOptionsIn,
   OperationContext,
@@ -39,7 +40,7 @@ export interface ProcedureUtils<TClientContext extends ClientContext, TInput, TO
    */
   queryKey(
     ...rest: MaybeOptionalOptions<
-      QueryKeyOptions<TInput>
+      QueryKeyOptions<TInput, TClientContext>
     >
   ): DataTag<QueryKey, TOutput, TError>
 
@@ -61,7 +62,7 @@ export interface ProcedureUtils<TClientContext extends ClientContext, TInput, TO
    */
   experimental_streamedKey(
     ...rest: MaybeOptionalOptions<
-      experimental_StreamedKeyOptions<TInput>
+      experimental_StreamedKeyOptions<TInput, TClientContext>
     >
   ): DataTag<QueryKey, experimental_StreamedQueryOutput<TOutput>, TError>
 
@@ -85,7 +86,7 @@ export interface ProcedureUtils<TClientContext extends ClientContext, TInput, TO
    */
   experimental_liveKey(
     ...rest: MaybeOptionalOptions<
-      QueryKeyOptions<TInput>
+      QueryKeyOptions<TInput, TClientContext>
     >
   ): DataTag<QueryKey, experimental_LiveQueryOutput<TOutput>, TError>
 
@@ -110,7 +111,7 @@ export interface ProcedureUtils<TClientContext extends ClientContext, TInput, TO
   infiniteKey<UPageParam>(
     options: Pick<
       InfiniteOptionsIn<TClientContext, TInput, TOutput, TError, InfiniteData<TOutput, UPageParam>, UPageParam>,
-      'input' | 'initialPageParam' | 'queryKey'
+      'context' | 'initialPageParam' | 'input' | 'queryKey'
     >
   ): DataTag<QueryKey, InfiniteData<TOutput, UPageParam>, TError>
 
@@ -129,9 +130,8 @@ export interface ProcedureUtils<TClientContext extends ClientContext, TInput, TO
    * @see {@link https://orpc.unnoq.com/docs/integrations/tanstack-query#query-mutation-key Tanstack Query/Mutation Key Docs}
    */
   mutationKey(
-    options?: Pick<
-      MutationOptionsIn<TClientContext, TInput, TOutput, TError, any>,
-      'mutationKey'
+    ...rest: MaybeOptionalOptions<
+      MutationKeyOptions<TClientContext>
     >
   ): DataTag<QueryKey, TOutput, TError>
 
@@ -159,7 +159,10 @@ export function createProcedureUtils<TClientContext extends ClientContext, TInpu
     call: client,
 
     queryKey(...[optionsIn = {} as any]) {
-      const queryKey = optionsIn.queryKey ?? generateOperationKey(options.path, { type: 'query', input: optionsIn.input })
+      const queryKey = optionsIn.queryKey ?? generateOperationKey(
+        options.path,
+        { type: 'query', input: optionsIn.input, context: optionsIn.context },
+      )
 
       return queryKey
     },
@@ -191,7 +194,10 @@ export function createProcedureUtils<TClientContext extends ClientContext, TInpu
     },
 
     experimental_streamedKey(...[optionsIn = {} as any]) {
-      const queryKey = optionsIn.queryKey ?? generateOperationKey(options.path, { type: 'streamed', input: optionsIn.input, fnOptions: optionsIn.queryFnOptions })
+      const queryKey = optionsIn.queryKey ?? generateOperationKey(
+        options.path,
+        { type: 'streamed', input: optionsIn.input, fnOptions: optionsIn.queryFnOptions, context: optionsIn.context },
+      )
 
       return queryKey
     },
@@ -232,7 +238,10 @@ export function createProcedureUtils<TClientContext extends ClientContext, TInpu
     },
 
     experimental_liveKey(...[optionsIn = {} as any]) {
-      const queryKey = optionsIn.queryKey ?? generateOperationKey(options.path, { type: 'live', input: optionsIn.input })
+      const queryKey = optionsIn.queryKey ?? generateOperationKey(
+        options.path,
+        { type: 'live', input: optionsIn.input, context: optionsIn.context },
+      )
 
       return queryKey
     },
@@ -273,6 +282,7 @@ export function createProcedureUtils<TClientContext extends ClientContext, TInpu
       const queryKey = optionsIn.queryKey ?? generateOperationKey(options.path, {
         type: 'infinite',
         input: optionsIn.input === skipToken ? skipToken : optionsIn.input(optionsIn.initialPageParam) as any,
+        context: optionsIn.context,
       })
 
       return queryKey as any
@@ -305,7 +315,10 @@ export function createProcedureUtils<TClientContext extends ClientContext, TInpu
     },
 
     mutationKey(...[optionsIn = {} as any]) {
-      const mutationKey = optionsIn.mutationKey ?? generateOperationKey(options.path, { type: 'mutation' })
+      const mutationKey = optionsIn.mutationKey ?? generateOperationKey(
+        options.path,
+        { type: 'mutation', context: optionsIn.context },
+      )
 
       return mutationKey
     },
diff --git a/packages/tanstack-query/src/types.ts b/packages/tanstack-query/src/types.ts
index bbc7d2714..c7884c7e3 100644
--- a/packages/tanstack-query/src/types.ts
+++ b/packages/tanstack-query/src/types.ts
@@ -5,6 +5,7 @@ import type {
   experimental_streamedQuery,
   InfiniteData,
   InfiniteQueryObserverOptions,
+  MutationKey,
   MutationObserverOptions,
   QueryFunction,
   QueryKey,
@@ -19,30 +20,40 @@ type experimental_StreamedQueryOptions = Omit<Parameters<typeof experimental_str
 
 export type OperationType = 'query' | 'streamed' | 'live' | 'infinite' | 'mutation'
 
-export type OperationKeyOptions<TType extends OperationType, TInput> = {
+/**
+ * @todo move TClientContext to second position + remove default in next major version
+ */
+export type OperationKeyOptions<TType extends OperationType, TInput, TClientContext extends ClientContext = ClientContext> = {
   type?: TType
-  input?: TType extends 'mutation' ? never : PartialDeep<TInput>
+  context?: TClientContext
   fnOptions?: TType extends 'streamed' ? experimental_StreamedQueryOptions : never
+  input?: TType extends 'mutation' ? never : PartialDeep<TInput>
 }
 
-export type OperationKey<TType extends OperationType, TInput> = [path: readonly string[], options: OperationKeyOptions<TType, TInput>]
+/**
+ * @todo move TClientContext to second position + remove default in next major version
+ */
+export type OperationKey<TType extends OperationType, TInput, TClientContext extends ClientContext = ClientContext> = [path: readonly string[], options: OperationKeyOptions<TType, TInput, TClientContext>]
 
 export const OPERATION_CONTEXT_SYMBOL: unique symbol = Symbol('ORPC_OPERATION_CONTEXT')
 
 export interface OperationContext {
   [OPERATION_CONTEXT_SYMBOL]?: {
-    key: QueryKey
     type: OperationType
+    key: QueryKey
   }
 }
 
-export type QueryKeyOptions<TInput>
+/**
+ * @todo move TClientContext to first position + remove default in next major version
+ */
+export type QueryKeyOptions<TInput, TClientContext extends ClientContext = ClientContext>
   = & (undefined extends TInput ? { input?: TInput | SkipToken } : { input: TInput | SkipToken })
+    & (Record<never, never> extends TClientContext ? { context?: TClientContext } : { context: TClientContext })
     & { queryKey?: QueryKey }
 
 export type QueryOptionsIn<TClientContext extends ClientContext, TInput, TOutput, TError, TSelectData>
-  = & QueryKeyOptions<TInput>
-    & (Record<never, never> extends TClientContext ? { context?: TClientContext } : { context: TClientContext })
+  = & QueryKeyOptions<TInput, TClientContext>
     & Omit<QueryObserverOptions<TOutput, TError, TSelectData>, 'queryKey'>
 
 export interface QueryOptionsBase<TOutput, TError> {
@@ -53,8 +64,11 @@ export interface QueryOptionsBase<TOutput, TError> {
   enabled: boolean
 }
 
-export type experimental_StreamedKeyOptions<TInput>
-  = & QueryKeyOptions<TInput>
+/**
+ * @todo move TClientContext to first position + remove default in next major version
+ */
+export type experimental_StreamedKeyOptions<TInput, TClientContext extends ClientContext = ClientContext>
+  = & QueryKeyOptions<TInput, TClientContext>
     & { queryFnOptions?: experimental_StreamedQueryOptions }
 
 export type experimental_StreamedOptionsIn<TClientContext extends ClientContext, TInput, TOutput, TError, TSelectData>
@@ -78,6 +92,10 @@ export interface InfiniteOptionsBase<TOutput, TError, TPageParam> {
   enabled: boolean
 }
 
+export type MutationKeyOptions<TClientContext extends ClientContext>
+    = & (Record<never, never> extends TClientContext ? { context?: TClientContext } : { context: TClientContext })
+      & { mutationKey?: MutationKey }
+
 export type MutationOptionsIn<TClientContext extends ClientContext, TInput, TOutput, TError, TMutationContext>
     = & (Record<never, never> extends TClientContext ? { context?: TClientContext } : { context: TClientContext })
       & MutationOptions<TInput, TOutput, TError, TMutationContext>