From 6760ea46f9fec622314f56fd3b2eade5a1e054bf Mon Sep 17 00:00:00 2001
From: yohlo <yohlerkyle@gmail.com>
Date: Wed, 24 Sep 2025 08:04:09 -0500
Subject: [PATCH] update query integration

---
 package.json                                  |  1 +
 src/app/router.tsx                            | 40 ++++++++++---------
 .../hooks/use-optimistic-mutation.ts          |  8 ++--
 .../hooks/use-server-mutation.ts              |  8 ++--
 4 files changed, 31 insertions(+), 26 deletions(-)

diff --git a/package.json b/package.json
index a1caf07..8a9ada5 100644
--- a/package.json
+++ b/package.json
@@ -52,6 +52,7 @@
     "zustand": "^5.0.7"
   },
   "devDependencies": {
+    "@tanstack/react-router-ssr-query": "^1.132.2",
     "@tanstack/router-plugin": "^1.132.2",
     "@types/node": "^22.5.4",
     "@types/pg": "^8.15.5",
diff --git a/src/app/router.tsx b/src/app/router.tsx
index 0138f21..c70f40b 100644
--- a/src/app/router.tsx
+++ b/src/app/router.tsx
@@ -1,6 +1,6 @@
 import { QueryClient } from "@tanstack/react-query";
 import { createRouter as createTanStackRouter } from "@tanstack/react-router";
-import { routerWithQueryClient } from "@tanstack/react-router-with-query";
+import { setupRouterSsrQueryIntegration } from "@tanstack/react-router-ssr-query";
 import { routeTree } from "./routeTree.gen";
 import { DefaultCatchBoundary } from "../components/DefaultCatchBoundary";
 import { defaultHeaderConfig } from "@/features/core/hooks/use-router-config";
@@ -18,24 +18,28 @@ export function getRouter() {
     },
   });
 
-  return routerWithQueryClient(
-    createTanStackRouter({
-      routeTree,
-      context: {
-        queryClient,
-        auth: undefined!,
-        header: defaultHeaderConfig,
-        refresh: [],
-        withPadding: true,
-        fullWidth: false,
-      },
-      defaultPreload: "intent",
-      defaultErrorComponent: DefaultCatchBoundary,
-      scrollRestoration: true,
-      defaultViewTransition: false,
-    }),
+  const router = createTanStackRouter({
+    routeTree,
+    context: {
+      queryClient,
+      auth: undefined!,
+      header: defaultHeaderConfig,
+      refresh: [],
+      withPadding: true,
+      fullWidth: false,
+    },
+    defaultPreload: "intent",
+    defaultErrorComponent: DefaultCatchBoundary,
+    scrollRestoration: true,
+    defaultViewTransition: false,
+  });
+
+  setupRouterSsrQueryIntegration({
+    router,
     queryClient
-  );
+  })
+
+  return router;
 }
 
 declare module "@tanstack/react-router" {
diff --git a/src/lib/tanstack-query/hooks/use-optimistic-mutation.ts b/src/lib/tanstack-query/hooks/use-optimistic-mutation.ts
index f46f7a3..c14c8cf 100644
--- a/src/lib/tanstack-query/hooks/use-optimistic-mutation.ts
+++ b/src/lib/tanstack-query/hooks/use-optimistic-mutation.ts
@@ -23,15 +23,15 @@ export function useOptimisticMutation<TData, TVariables = unknown>(
 
       return { previousData };
     },
-    onError: (error, variables, context) => {
+    onError: (error, variables, onMutateResult, context) => {
       if (context && typeof context === 'object' && 'previousData' in context && context.previousData) {
         queryClient.setQueryData(queryKey, context.previousData);
       }
-      mutationOptions.onError?.(error, variables, context);
+      mutationOptions.onError?.(error, variables, onMutateResult, context);
     },
-    onSettled: (data, error, variables, context) => {
+    onSettled: (data, error, variables, onMutateResult, context) => {
       queryClient.invalidateQueries({ queryKey });
-      mutationOptions.onSettled?.(data, error, variables, context);
+      mutationOptions.onSettled?.(data, error, variables, onMutateResult, context);
     }
   });
 }
diff --git a/src/lib/tanstack-query/hooks/use-server-mutation.ts b/src/lib/tanstack-query/hooks/use-server-mutation.ts
index a6e3ccf..c6de452 100644
--- a/src/lib/tanstack-query/hooks/use-server-mutation.ts
+++ b/src/lib/tanstack-query/hooks/use-server-mutation.ts
@@ -52,14 +52,14 @@ export function useServerMutation<TData, TVariables = unknown>(
         throw error;
       }
     },
-    onSuccess: (data, variables, context) => {
+    onSuccess: (data, variables, onMutateResult, context) => {
       if (showSuccessToast && successMessage) {
         toast.success(successMessage);
       }
-      onSuccess?.(data, variables, context);
+      onSuccess?.(data, variables, onMutateResult, context);
     },
-    onError: (error, variables, context) => {
-      onError?.(error, variables, context);
+    onError: (error, variables, onMutateResult, context) => {
+      onError?.(error, variables, onMutateResult, context);
     }
   });
 }