From 1a7a52a63ae51b5ab04d0d73bcae604f2d7d647d Mon Sep 17 00:00:00 2001
From: Kyle Graham Matzen <kgrahammatzen@gmail.com>
Date: Fri, 27 Feb 2026 14:03:07 -0500
Subject: [PATCH] feat: add key param to useCustomer and useEntity for
 declarative revalidation

---
 package/src/libraries/react/hooks/useCustomerBase.tsx |  4 +++-
 package/src/libraries/react/hooks/useEntity.tsx       |  8 ++++++--
 package/src/libraries/react/hooks/useEntityBase.tsx   |  5 ++++-
 package/src/next/client/hooks/useEntity.tsx           | 11 +++++++++--
 4 files changed, 22 insertions(+), 6 deletions(-)

diff --git a/package/src/libraries/react/hooks/useCustomerBase.tsx b/package/src/libraries/react/hooks/useCustomerBase.tsx
index 9a9f8d1b..2c9b01bf 100644
--- a/package/src/libraries/react/hooks/useCustomerBase.tsx
+++ b/package/src/libraries/react/hooks/useCustomerBase.tsx
@@ -130,6 +130,8 @@ export interface UseCustomerParams<
 	errorOnNotFound?: boolean;
 	expand?: T;
 	swrConfig?: SWRConfiguration;
+	/** Extra key(s) appended to the SWR query key that trigger a refetch when any value changes. */
+	extraQueryKeys?: (string | null | undefined)[];
 }
 
 export const useCustomerBase = <
@@ -158,7 +160,7 @@ export const useCustomerBase = <
 	}
 
 	const baseUrl = client?.backendUrl || "";
-	const queryKey = ["customer", baseUrl, params?.expand];
+	const queryKey = ["customer", baseUrl, params?.expand, ...(params?.extraQueryKeys ?? [])];
 
 	const fetchCustomer = async () => {
 		const { data, error } = await client!.createCustomer<T>({
diff --git a/package/src/libraries/react/hooks/useEntity.tsx b/package/src/libraries/react/hooks/useEntity.tsx
index 790b7a57..c8b7e5c7 100644
--- a/package/src/libraries/react/hooks/useEntity.tsx
+++ b/package/src/libraries/react/hooks/useEntity.tsx
@@ -4,7 +4,11 @@ import { useEntityBase } from "./useEntityBase";
 
 export const useEntity = (
   entityId: string | null,
-  params?: GetEntityParams
+  params?: GetEntityParams & {
+    /** Extra key(s) appended to the SWR query key that trigger a refetch when any value changes. */
+    extraQueryKeys?: (string | null | undefined)[];
+  }
 ) => {
-  return useEntityBase({ AutumnContext, entityId, params });
+  const { extraQueryKeys, ...restParams } = params ?? {};
+  return useEntityBase({ AutumnContext, entityId, params: restParams, extraQueryKeys });
 };
diff --git a/package/src/libraries/react/hooks/useEntityBase.tsx b/package/src/libraries/react/hooks/useEntityBase.tsx
index ec0d4bd4..6a36d169 100644
--- a/package/src/libraries/react/hooks/useEntityBase.tsx
+++ b/package/src/libraries/react/hooks/useEntityBase.tsx
@@ -15,10 +15,13 @@ import { AutumnError, CheckResult, Entity } from "@sdk";
 export const useEntityBase = ({
   entityId,
   params,
+  extraQueryKeys,
   AutumnContext,
 }: {
   entityId: string | null;
   params?: GetEntityParams;
+  /** Extra key(s) appended to the SWR query key that trigger a refetch when any value changes. */
+  extraQueryKeys?: (string | null | undefined)[];
   AutumnContext: React.Context<AutumnContextParams>;
 }): {
   entity: Entity | null;
@@ -31,7 +34,7 @@ export const useEntityBase = ({
   track: (params: TrackParams) => void;
 } => {
   const { client } = useContext(AutumnContext);
-  const queryKey = ["entity", entityId, params?.expand];
+  const queryKey = ["entity", entityId, params?.expand, ...(extraQueryKeys ?? [])];
 
   const context = useAutumnContext({
     AutumnContext,
diff --git a/package/src/next/client/hooks/useEntity.tsx b/package/src/next/client/hooks/useEntity.tsx
index 49951fbb..2d732c34 100644
--- a/package/src/next/client/hooks/useEntity.tsx
+++ b/package/src/next/client/hooks/useEntity.tsx
@@ -2,6 +2,13 @@ import { GetEntityParams } from "../../../sdk/customers/entities/entTypes";
 import { useEntityBase } from "../../../libraries/react/hooks/useEntityBase";
 import { AutumnContext } from "../../../libraries/react/AutumnContext";
 
-export const useEntity = (entityId: string | null, params?: GetEntityParams) => {
-  return useEntityBase({ AutumnContext, entityId, params });
+export const useEntity = (
+  entityId: string | null,
+  params?: GetEntityParams & {
+    /** Extra key(s) appended to the SWR query key that trigger a refetch when any value changes. */
+    extraQueryKeys?: (string | null | undefined)[];
+  }
+) => {
+  const { extraQueryKeys, ...restParams } = params ?? {};
+  return useEntityBase({ AutumnContext, entityId, params: restParams, extraQueryKeys });
 };