fix: onboarding flow and versioning

server/src/internal/customers/attach/checkout/previewToCheckoutRes.ts

@@ -21,7 +21,7 @@ import {
   getProductItemResponse,
   getProductResponse,
 } from "@/internal/products/productUtils/productResponseUtils/getProductResponse.js";
-import { toProductItem } from "@/internal/products/product-items/mapToItem.js";
+import { toProductItem } from "@autumn/shared";
 import { getPriceEntitlement } from "@/internal/products/prices/priceUtils.js";
 import { formatUnixToDateTime, notNullish } from "@/utils/genUtils.js";
 import { isPriceItem } from "@/internal/products/product-items/productItemUtils/getItemType.js";

server/src/internal/customers/getFullCusQuery.ts

@@ -377,7 +377,7 @@ export const getPaginatedFullCusQuery = ({
       WHERE c.org_id = ${orgId}
         AND c.env = ${env}
       ${
-        internalCustomerIds
+        internalCustomerIds && internalCustomerIds.length > 0
           ? sql`AND c.internal_id IN (${sql.join(
               internalCustomerIds.map((id) => sql`${id}`),
               sql`, `

server/src/internal/customers/internalCusRouter.ts

@@ -7,24 +7,18 @@ import {
   CusProductStatus,
   ErrCode,
   FullCusProduct,
-  FullCustomerEntitlement,
-  FullCustomerPrice,
   productToCusProduct,
 } from "@autumn/shared";
 
 import RecaseError, { handleFrontendReqError } from "@/utils/errorUtils.js";
-import { RewardService } from "../rewards/RewardService.js";
 import { EventService } from "../api/events/EventService.js";
 import { createStripeCli } from "@/external/stripe/utils.js";
-import { getLatestProducts } from "../products/productUtils.js";
-import { getProductVersionCounts } from "../products/productUtils.js";
-import { notNullish, nullish } from "@/utils/genUtils.js";
 import { mapToProductV2 } from "../products/productV2Utils.js";
 import { RewardRedemptionService } from "../rewards/RewardRedemptionService.js";
 import { CusReadService } from "./CusReadService.js";
 import { StatusCodes } from "http-status-codes";
 import { cusProductToProduct } from "@autumn/shared";
-import { createOrgResponse, isStripeConnected } from "../orgs/orgUtils.js";
+import { isStripeConnected } from "../orgs/orgUtils.js";
 import { routeHandler } from "@/utils/routerUtils.js";
 import { CusSearchService } from "./CusSearchService.js";
 import { CusBatchService } from "../api/batch/CusBatchService.js";

server/src/internal/invoices/previewItemUtils/getItemsForNewProduct.ts

@@ -33,19 +33,12 @@ import {
 import { newPriceToInvoiceDescription } from "../invoiceFormatUtils.js";
 import { calculateProrationAmount } from "../prorationUtils.js";
 import { getPricecnPrice } from "../../products/pricecn/pricecnUtils.js";
-import { toProductItem } from "../../products/product-items/mapToItem.js";
+import { toProductItem } from "@autumn/shared";
 import { formatAmount } from "@/utils/formatUtils.js";
-import {
-  formatUnixToDate,
-  formatUnixToDateTime,
-  formatUnixToUTCDateTime,
-  notNullish,
-} from "@/utils/genUtils.js";
+import { formatUnixToDate, notNullish } from "@/utils/genUtils.js";
 import {
   addBillingIntervalUnix,
   addIntervalForProration,
-  getAlignedIntervalUnix,
-  subtractBillingIntervalUnix,
   subtractFromUnixTillAligned,
   subtractIntervalForProration,
 } from "../../products/prices/billingIntervalUtils.js";

server/src/internal/products/entitlements/entitlementUtils.ts

@@ -1,5 +1,5 @@
 import RecaseError from "@/utils/errorUtils.js";
-import { generateId, notNullOrUndefined } from "@/utils/genUtils.js";
+
 import {
   EntInterval,
   FreeTrial,
@@ -7,21 +7,16 @@ import {
   AllowanceType,
   EntitlementWithFeature,
   FeatureType,
-  CreateEntitlement,
-  CreateEntitlementSchema,
   Feature,
   ErrCode,
   UsagePriceConfig,
   PriceType,
   Price,
   FullProduct,
-  BillingType,
   FullEntitlement,
 } from "@autumn/shared";
 
 import { addDays } from "date-fns";
-import { getBillingType } from "@/internal/products/prices/priceUtils.js";
-import { features } from "process";
 
 export const entIntervalToTrialDuration = ({
   interval,

server/src/internal/products/internalProductRouter.ts

@@ -24,6 +24,7 @@ import { sortFullProducts } from "./productUtils/sortProductUtils.js";
 
 export const productRouter: Router = Router({ mergeParams: true });
 
+// Get list of products
 productRouter.get("/products", async (req: any, res) => {
   try {
     const { db } = req;
@@ -49,6 +50,7 @@ productRouter.get("/products", async (req: any, res) => {
   }
 });
 
+// Get counts for all products
 productRouter.get("/product_counts", async (req: any, res) => {
   try {
     let { db } = req;
@@ -90,6 +92,7 @@ productRouter.get("/product_counts", async (req: any, res) => {
   }
 });
 
+// Get list of features
 productRouter.get("/features", async (req: any, res) => {
   try {
     res.status(200).json({ features: req.features });
@@ -99,6 +102,7 @@ productRouter.get("/features", async (req: any, res) => {
   }
 });
 
+// Get list of rewards
 productRouter.get("/rewards", async (req: any, res) => {
   try {
     const { db, orgId, env } = req;
@@ -119,19 +123,28 @@ productRouter.get("/rewards", async (req: any, res) => {
   }
 });
 
+// Get single product data
 productRouter.get("/:productId/data2", async (req: any, res) => {
   try {
     const { productId } = req.params;
     const { version } = req.query;
     const { db, orgId, env } = req;
 
-    const product = await ProductService.getFull({
-      db,
-      idOrInternalId: productId,
-      orgId,
-      env,
-      version: version ? parseInt(version) : undefined,
-    });
+    const [product, latestProduct] = await Promise.all([
+      ProductService.getFull({
+        db,
+        idOrInternalId: productId,
+        orgId,
+        env,
+        version: version ? parseInt(version) : undefined,
+      }),
+      ProductService.getFull({
+        db,
+        idOrInternalId: productId,
+        orgId,
+        env,
+      }),
+    ]);
 
     if (!product) {
       throw new RecaseError({
@@ -148,13 +161,77 @@ productRouter.get("/:productId/data2", async (req: any, res) => {
       features: req.features,
     });
 
-    res.status(200).json({ product: productV2 });
+    res
+      .status(200)
+      .json({ product: productV2, numVersions: latestProduct.version });
   } catch (error) {
     console.error("Failed to get product", error);
     res.status(500).send(error);
   }
 });
 
+// Get counts for a single product
+productRouter.get("/:productId/count", async (req: any, res) => {
+  try {
+    const { db, orgId, env } = req;
+    const { productId } = req.params;
+    const { version } = req.query;
+
+    const product = await ProductService.get({
+      db,
+      id: productId,
+      orgId,
+      env,
+      version: version ? parseInt(version) : undefined,
+    });
+
+    if (!product) {
+      throw new RecaseError({
+        message: `Product ${productId} ${
+          version ? `(v${version})` : ""
+        } not found`,
+        code: ErrCode.ProductNotFound,
+        statusCode: StatusCodes.NOT_FOUND,
+      });
+    }
+
+    // Get counts from postgres
+    const counts = await CusProdReadService.getCounts({
+      db,
+      internalProductId: product.internal_id,
+    });
+
+    res.status(200).send(counts);
+  } catch (error) {
+    handleFrontendReqError({
+      error,
+      req,
+      res,
+      action: "Get product counts (internal)",
+    });
+  }
+});
+
+// Get list of migrations
+productRouter.get("/migrations", async (req: any, res) => {
+  try {
+    const { db, orgId, env } = req;
+    const migrations = await MigrationService.getExistingJobs({
+      db,
+      orgId,
+      env,
+    });
+    res.status(200).send({ migrations });
+  } catch (error) {
+    handleFrontendReqError({
+      error,
+      req,
+      res,
+      action: "Get migrations",
+    });
+  }
+});
+
 productRouter.get("/data", async (req: any, res) => {
   try {
     let { db } = req;
@@ -395,47 +472,6 @@ productRouter.get("/:productId/data", async (req: any, res) => {
   }
 });
 
-productRouter.get("/:productId/count", async (req: any, res) => {
-  try {
-    const { db, orgId, env } = req;
-    const { productId } = req.params;
-    const { version } = req.query;
-
-    const product = await ProductService.get({
-      db,
-      id: productId,
-      orgId,
-      env,
-      version: version ? parseInt(version) : undefined,
-    });
-
-    if (!product) {
-      throw new RecaseError({
-        message: `Product ${productId} ${
-          version ? `(v${version})` : ""
-        } not found`,
-        code: ErrCode.ProductNotFound,
-        statusCode: StatusCodes.NOT_FOUND,
-      });
-    }
-
-    // Get counts from postgres
-    const counts = await CusProdReadService.getCounts({
-      db,
-      internalProductId: product.internal_id,
-    });
-
-    res.status(200).send(counts);
-  } catch (error) {
-    handleFrontendReqError({
-      error,
-      req,
-      res,
-      action: "Get product counts (internal)",
-    });
-  }
-});
-
 productRouter.post("/product_options", async (req: any, res: any) => {
   try {
     const { items } = req.body;

server/src/internal/products/product-items/productItemUtils.ts

@@ -74,27 +74,6 @@ export const itemToPriceOrTiers = ({
     }
   }
 };
-export const getItemFeatureType = ({
-  item,
-  features,
-}: {
-  item: ProductItem;
-  features: Feature[];
-}) => {
-  let feature = features.find((f) => f.id == item.feature_id);
-
-  if (feature) {
-    if (feature.type == FeatureType.Boolean) {
-      return ProductItemFeatureType.Static;
-    } else if (feature.type == FeatureType.CreditSystem) {
-      return ProductItemFeatureType.SingleUse;
-    } else {
-      return feature.config?.usage_type;
-    }
-  }
-
-  return undefined;
-};
 
 // FOR TESTS?
 export const constructFeatureItem = ({

server/src/internal/products/productV2Utils.ts

@@ -1,23 +1,15 @@
 import {
   EntitlementWithFeature,
   Feature,
-  FeatureType,
-  FreeTrial,
-  FreeTrialResponseSchema,
   FullProduct,
   Price,
   ProductItem,
-  ProductItemInterval,
-  ProductItemResponseSchema,
-  ProductResponseSchema,
   ProductV2,
 } from "@autumn/shared";
 import { getEntRelatedPrice } from "./entitlements/entitlementUtils.js";
 import { getPriceEntitlement } from "./prices/priceUtils.js";
-import { toProductItem } from "./product-items/mapToItem.js";
-import { getItemFeatureType } from "./product-items/productItemUtils.js";
-import { sortProductItems } from "./pricecn/pricecnUtils.js";
-import { getItemType } from "./product-items/productItemUtils/getItemType.js";
+import { toProductItem } from "@autumn/shared";
+import { getItemFeatureType } from "@autumn/shared";
 
 export const mapToProductItems = ({
   prices,
@@ -42,7 +34,7 @@ export const mapToProductItems = ({
     let relatedEnt = getPriceEntitlement(
       price,
       entitlements,
-      allowFeatureMatch,
+      allowFeatureMatch
     );
 
     if (!relatedEnt) {

shared/utils/index.ts

@@ -7,3 +7,11 @@ export * from "./cusProductUtils/convertCusProduct.js";
 export * from "./cusProductUtils/productIdToCusProduct.js";
 export * from "./cusProductUtils/cusProductConstants.js";
 export * from "./utils.js";
+
+// Item utils
+export * from "./productV2Utils/productItemUtils/mapToItem.js";
+export * from "./productV2Utils/productItemUtils/productItemUtils.js";
+export * from "./productV2Utils/mapToProductV2.js";
+
+// Product utils
+export * from "./productUtils/convertUtils.js";

shared/utils/productUtils/convertUtils.ts

@@ -0,0 +1,53 @@
+import { Entitlement } from "../../models/productModels/entModels/entModels.js";
+import { Price } from "../../models/productModels/priceModels/priceModels.js";
+
+// export const getEntRelatedPrice = (
+//   entitlement: Entitlement,
+//   prices: Price[],
+//   allowFeatureMatch = false
+// ) => {
+//   return prices.find((price) => {
+//     if (price.config?.type === PriceType.Fixed) {
+//       return false;
+//     }
+
+//     let config = price.config as UsagePriceConfig;
+
+//     if (allowFeatureMatch) {
+//       return entitlement.internal_feature_id == config.internal_feature_id;
+//     }
+
+//     let entIdMatch = entitlement.id == price.entitlement_id;
+//     let productIdMatch =
+//       entitlement.internal_product_id == price.internal_product_id;
+//     return entIdMatch && productIdMatch;
+//   });
+// };
+
+export const entToPrice = ({
+  ent,
+  prices,
+}: {
+  ent: Entitlement;
+  prices: Price[];
+}) => {
+  return prices.find(
+    (price) =>
+      price.entitlement_id === ent.id &&
+      price.internal_product_id === ent.internal_product_id
+  );
+};
+
+export const priceToEnt = ({
+  price,
+  entitlements,
+}: {
+  price: Price;
+  entitlements: Entitlement[];
+}) => {
+  return entitlements.find(
+    (ent) =>
+      ent.id === price.entitlement_id &&
+      ent.internal_product_id === price.internal_product_id
+  );
+};

shared/utils/productV2Utils/mapToProductV2.ts

@@ -0,0 +1,57 @@
+import { Feature } from "../../models/featureModels/featureModels.js";
+import { FullProduct } from "../../models/productModels/productModels.js";
+import { ProductItem } from "../../models/productV2Models/productItemModels/productItemModels.js";
+import { ProductV2 } from "../../models/productV2Models/productV2Models.js";
+import { entToPrice, priceToEnt } from "../productUtils/convertUtils.js";
+import { toProductItem } from "./productItemUtils/mapToItem.js";
+import { getItemFeatureType } from "./productItemUtils/productItemUtils.js";
+
+export const mapToProductV2 = ({
+  product,
+  features,
+}: {
+  product: FullProduct;
+  features?: Feature[];
+}): ProductV2 => {
+  let items: ProductItem[] = [];
+
+  for (const ent of product.entitlements) {
+    let relatedPrice = entToPrice({ ent, prices: product.prices });
+    items.push(toProductItem({ ent, price: relatedPrice }));
+  }
+
+  for (const price of product.prices) {
+    let relatedEnt = priceToEnt({ price, entitlements: product.entitlements });
+
+    if (!relatedEnt) {
+      items.push(toProductItem({ price }));
+    }
+  }
+
+  if (!features) {
+    features = product.entitlements.map((ent) => ent.feature);
+  }
+
+  for (const item of items) {
+    item.feature_type = getItemFeatureType({ item, features });
+  }
+
+  let productV2: ProductV2 = {
+    internal_id: product.internal_id,
+
+    id: product.id,
+    name: product.name,
+    is_add_on: product.is_add_on,
+    is_default: product.is_default,
+    version: product.version,
+    group: product.group,
+    free_trial: product.free_trial,
+    created_at: product.created_at,
+
+    items: items,
+    stripe_id: product.processor?.id || null,
+    archived: product.archived || false,
+  };
+
+  return productV2;
+};

shared/utils/productV2Utils/productItemUtils/mapToItem.ts

@@ -1,23 +1,26 @@
+import { FeatureType } from "../../../models/featureModels/featureEnums.js";
 import {
   AllowanceType,
-  BillWhen,
   EntitlementWithFeature,
-  FeatureType,
-  FixedPriceConfig,
-  Infinite,
-  Price,
-  ProductItem,
-  UsageModel,
-  TierInfinite,
+} from "../../../models/productModels/entModels/entModels.js";
+import { FixedPriceConfig } from "../../../models/productModels/priceModels/priceConfig/fixedPriceConfig.js";
+import {
+  BillWhen,
   UsagePriceConfig,
+} from "../../../models/productModels/priceModels/priceConfig/usagePriceConfig.js";
+import { Price } from "../../../models/productModels/priceModels/priceModels.js";
+import { Infinite } from "../../../models/productModels/productEnums.js";
+import {
+  ProductItem,
   ProductItemFeatureType,
-} from "@autumn/shared";
-
-import { nullish } from "@/utils/genUtils.js";
+  TierInfinite,
+  UsageModel,
+} from "../../../models/productV2Models/productItemModels/productItemModels.js";
+import { nullish } from "../../utils.js";
 import {
   billingToItemInterval,
   entToItemInterval,
-} from "./itemIntervalUtils.js";
+} from "./productItemUtils.js";
 
 export const toProductItem = ({
   ent,

shared/utils/productV2Utils/productItemUtils/productItemUtils.ts

@@ -0,0 +1,46 @@
+import { FeatureType } from "../../../models/featureModels/featureEnums.js";
+import { Feature } from "../../../models/featureModels/featureModels.js";
+import { EntInterval } from "../../../models/productModels/entModels/entEnums.js";
+import { BillingInterval } from "../../../models/productModels/priceModels/priceEnums.js";
+import {
+  ProductItem,
+  ProductItemFeatureType,
+  ProductItemInterval,
+} from "../../../models/productV2Models/productItemModels/productItemModels.js";
+
+export const entToItemInterval = (entInterval: EntInterval) => {
+  if (entInterval == EntInterval.Lifetime) {
+    return null;
+  }
+  return entInterval as unknown as ProductItemInterval;
+};
+
+export const billingToItemInterval = (billingInterval: BillingInterval) => {
+  if (billingInterval == BillingInterval.OneOff) {
+    return null;
+  }
+
+  return billingInterval as unknown as ProductItemInterval;
+};
+
+export const getItemFeatureType = ({
+  item,
+  features,
+}: {
+  item: ProductItem;
+  features: Feature[];
+}) => {
+  let feature = features.find((f) => f.id == item.feature_id);
+
+  if (feature) {
+    if (feature.type == FeatureType.Boolean) {
+      return ProductItemFeatureType.Static;
+    } else if (feature.type == FeatureType.CreditSystem) {
+      return ProductItemFeatureType.SingleUse;
+    } else {
+      return feature.config?.usage_type;
+    }
+  }
+
+  return undefined;
+};

vite/src/components/autumn/checkout-dialog.tsx

@@ -50,8 +50,7 @@ export default function CheckoutDialog(params: CheckoutDialogProps) {
   const [checkoutResult, setCheckoutResult] = useState<
     CheckoutResult | undefined
   >(params?.checkoutResult);
-
-  const { mutateAutumnProducts } = useModelPricingContext();
+  const { refetch } = useModelPricingContext();
 
   useEffect(() => {
     if (params.checkoutResult) {
@@ -102,7 +101,7 @@ export default function CheckoutDialog(params: CheckoutDialogProps) {
                 options,
                 openInNewTab: true,
               });
-              await mutateAutumnProducts();
+              await refetch();
               setOpen(false);
               setLoading(false);
             }}

vite/src/components/autumn/pricing-table.tsx

@@ -1,4 +1,5 @@
 import React from "react";
+import { Loader2 } from "lucide-react";
 
 import { createContext, useContext, useState } from "react";
 import { cn } from "@/lib/utils";
@@ -7,21 +8,18 @@ import { Button } from "@/components/ui/button";
 import CheckoutDialog from "@/components/autumn/checkout-dialog";
 import { getPricingTableContent } from "@/lib/autumn/pricing-table-content";
 import type { Product, ProductItem } from "autumn-js";
-import { Loader2 } from "lucide-react";
-import { toast } from "sonner";
+
 import { useCustomer } from "autumn-js/react";
-import { useModelPricingContext } from "@/views/onboarding2/model-pricing/ModelPricingContext";
+import { useOrg } from "@/hooks/common/useOrg";
 
 export default function PricingTable({
   products,
-  stripeConnected,
   setConnectStripeOpen,
 }: {
   products?: Product[];
-  stripeConnected: boolean;
   setConnectStripeOpen: (open: boolean) => void;
 }) {
-  const { mutateAutumnProducts } = useModelPricingContext();
+  const { org } = useOrg();
   const { checkout } = useCustomer();
   const [isAnnual, setIsAnnual] = useState(false);
 
@@ -70,8 +68,7 @@ export default function PricingTable({
                   product.scenario === "scheduled",
 
                 onClick: async () => {
-                  console.log("Inside onClick function");
-                  if (!stripeConnected) {
+                  if (!org.stripe_connected) {
                     setConnectStripeOpen(true);
                     return;
                   }

vite/src/utils/product/product-item/validateProductItem.ts

@@ -144,8 +144,6 @@ export const validateProductItem = ({
     item.tiers = finalTiers;
   }
 
-  console.log("Final tiers:", item.tiers);
-
   // Validate billing units
   if (item.billing_units && invalidNumber(item.billing_units)) {
     toast.error("Please enter valid billing units");

vite/src/views/customers/customer/customer-product-list/AttachProductDropdown.tsx

@@ -9,7 +9,7 @@ import {
 import React, { useState } from "react";
 import { useCustomerContext } from "../CustomerContext";
 import { useAxiosInstance } from "@/services/useAxiosInstance";
-import { AddProductContext } from "./CreateCheckoutContext";
+import { AddProductContext } from "../add-product/CreateCheckoutContext";
 import { Input } from "@/components/ui/input";
 
 import { useNavigate } from "react-router";
@@ -23,7 +23,7 @@ import { useOrg } from "@/hooks/common/useOrg";
 import { useCustomer } from "autumn-js/react";
 import { useCusQuery } from "../hooks/useCusQuery";
 
-function AddProduct({
+function AttachProductDropdown({
   setMultiAttachOpen,
 }: {
   setMultiAttachOpen: (open: boolean) => void;
@@ -143,7 +143,7 @@ function AddProduct({
   );
 }
 
-export default AddProduct;
+export default AttachProductDropdown;
 
 const DropdownProductItem = ({
   product,
@@ -159,9 +159,8 @@ const DropdownProductItem = ({
       onClick={async (e) => {
         e.stopPropagation();
         e.preventDefault();
-        setIsLoading(true);
+
         await handleAddProduct(product.id, setIsLoading);
-        // setIsLoading(false);
       }}
     >
       {isLoading && <SmallSpinner />}

vite/src/views/customers/customer/customer-product-list/CustomerProductList.tsx

@@ -14,7 +14,6 @@ import { useState } from "react";
 import { Badge } from "@/components/ui/badge";
 
 import { AdminHover } from "@/components/general/AdminHover";
-import AddProduct from "../add-product/NewProductDropdown";
 import { Item, Row } from "@/components/general/TableGrid";
 import { cn } from "@/lib/utils";
 
@@ -25,6 +24,7 @@ import { MultiAttachDialog } from "../product/multi-attach/MultiAttachDialog";
 import { useCusQuery } from "../hooks/useCusQuery";
 import { useProductsQuery } from "@/hooks/queries/useProductsQuery";
 import { getVersionCounts } from "@/utils/productUtils";
+import AttachProductDropdown from "./AttachProductDropdown";
 
 export const CustomerProductList = () => {
   const navigate = useNavigate();
@@ -125,7 +125,7 @@ export const CustomerProductList = () => {
                 open={multiAttachOpen}
                 setOpen={setMultiAttachOpen}
               />
-              <AddProduct setMultiAttachOpen={setMultiAttachOpen} />
+              <AttachProductDropdown setMultiAttachOpen={setMultiAttachOpen} />
             </div>
           </div>
         </div>

vite/src/views/customers/customer/product/CustomerProductView.tsx

@@ -62,8 +62,6 @@ export default function CustomerProductView() {
   const env = useEnv();
   const initialProductRef = useRef<ProductV2 | null>(null);
 
-  const [product, setProduct] = useState<ProductV2 | null>(null);
-
   const [options, setOptions] = useState<OptionValue[]>([]);
   const [entityId, setEntityId] = useState<string | null>(entityIdParam);
   const [entityFeatureIds, setEntityFeatureIds] = useState<string[]>([]);
@@ -77,6 +75,9 @@ export default function CustomerProductView() {
     isLoading,
     error,
   } = useCusProductQuery();
+  const [product, setProduct] = useState<ProductV2 | null>(
+    originalProduct ?? null
+  );
 
   const { isLoading: cusLoading } = useCusQuery();
 
@@ -133,7 +134,7 @@ export default function CustomerProductView() {
     );
   }
 
-  if (isLoading || !product || cusLoading || orgLoading || featuresLoading)
+  if (isLoading || cusLoading || orgLoading || featuresLoading || !product)
     return <LoadingScreen />;
 
   if (!customer_id || !product_id) {

vite/src/views/customers/customer/product/hooks/useAttachState.tsx

@@ -52,7 +52,7 @@ export const useAttachState = ({
   product: ProductV2 | null;
   setProduct: (product: ProductV2) => void;
 
-  cusProduct: FullCusProduct | undefined;
+  cusProduct: FullCusProduct | undefined | null;
   initialProductRef: React.RefObject<ProductV2>;
 }) => {
   const [preview, setPreview] = useState<AttachPreview | null>(null);

vite/src/views/customers/customer/product/hooks/useCusProductCache.tsx

@@ -1,9 +1,13 @@
 import { useQueryClient } from "@tanstack/react-query";
 import {
   ACTIVE_STATUSES,
+  cusProductToProduct,
   FullCusProduct,
   FullCustomer,
+  mapToProductV2,
+  notNullish,
   productToCusProduct,
+  ProductV2,
 } from "@autumn/shared";
 
 export const useCusProductCache = ({
@@ -21,7 +25,30 @@ export const useCusProductCache = ({
 }) => {
   const queryClient = useQueryClient();
 
-  const getCachedCusProduct = (): FullCusProduct | null => {
+  const getSortedProductsQueries = () => {
+    const queryCache = queryClient.getQueryCache();
+    const productsQuery = queryCache.findAll({
+      queryKey: ["products"],
+    });
+
+    // Sort by most recently updated first to get the freshest data
+    const sortedQueries = productsQuery.sort((a, b) => {
+      const aTime = a.state.dataUpdatedAt || 0;
+      const bTime = b.state.dataUpdatedAt || 0;
+      return bTime - aTime;
+    });
+
+    return sortedQueries
+      .map((query) => query.state.data)
+      .filter(notNullish) as {
+      products: ProductV2[];
+    }[];
+  };
+
+  const getCachedCusProduct = (): {
+    cusProduct: FullCusProduct | null;
+    product: ProductV2;
+  } | null => {
     if (!customerId || !productId) return null;
 
     // Check all cached full customers queries
@@ -65,26 +92,33 @@ export const useCusProductCache = ({
             // inStatuses: ACTIVE_STATUSES,
           });
 
-          console.log("Cached cus product:", cusProduct);
+          if (cusProduct) {
+            const fullProduct = cusProductToProduct({ cusProduct });
+            const productV2 = mapToProductV2({ product: fullProduct });
+            return {
+              cusProduct,
+              product: productV2,
+            };
+          }
         }
-
-        // if (cachedData?.fullCustomers) {
-        //   const cachedCustomer = cachedData.fullCustomers.find(
-        //     (cusProduct) =>
-        //       cusProduct.id === customerId ||
-        //       cusProduct.internal_customer_id === customerId
-        //   );
-
-        //   if (cachedCustomer) {
-        //     return cachedCustomer;
-        //   }
-        // }
       }
     }
 
-    // const productQueries = queryCache.findAll({
-    //   queryKey: ["products"],
-    // });
+    const sortedProductsQueries = getSortedProductsQueries();
+
+    for (const query of sortedProductsQueries) {
+      const product = query.products.find(
+        (p) =>
+          p.id === productId &&
+          (queryStates.version ? p.version === queryStates.version : true)
+      );
+      if (product) {
+        return {
+          cusProduct: null,
+          product,
+        };
+      }
+    }
 
     return null;
   };

vite/src/views/customers/customer/product/hooks/useCusProductQuery.tsx

@@ -35,6 +35,9 @@ export const useCusProductQuery = () => {
     };
 
     try {
+      console.log(
+        `Fetching customer product ${product_id} with version ${queryStates.version}`
+      );
       const { data } = await axiosInstance.get(
         `/customers/${customer_id}/product/${product_id}`,
         { params: queryParams }
@@ -49,14 +52,24 @@ export const useCusProductQuery = () => {
     cusProduct: FullCusProduct;
     product: ProductV2;
   }>({
-    queryKey: ["customer_product", customer_id, product_id],
+    queryKey: [
+      "customer_product",
+      customer_id,
+      product_id,
+      queryStates.version,
+      queryStates.customer_product_id,
+      queryStates.entity_id,
+    ],
     queryFn: fetcher,
   });
 
+  const finalData = data || cachedCusProduct;
+  const isLoadingWithCache = cachedCusProduct ? false : isLoading;
+
   return {
-    cusProduct: data?.cusProduct,
-    product: data?.product,
-    isLoading,
+    cusProduct: finalData?.cusProduct,
+    product: finalData?.product,
+    isLoading: isLoadingWithCache,
     error,
     refetch,
   };

vite/src/views/developer/DevView.tsx

@@ -12,6 +12,7 @@ import { PublishableKeySection } from "./publishable-key";
 import { useAutumnFlags } from "@/hooks/common/useAutumnFlags";
 import { ConfigureStripe } from "./configure-stripe/ConfigureStripe";
 import { useSecondaryTab } from "@/hooks/common/useSecondaryTab";
+import { useAppQueryStates } from "@/hooks/common/useAppQueryStates";
 
 export default function DevScreen({ env }: { env: AppEnv }) {
   const { data, isLoading, mutate } = useAxiosSWR({
@@ -20,7 +21,9 @@ export default function DevScreen({ env }: { env: AppEnv }) {
     withAuth: true,
   });
 
-  const secondaryTab = useSecondaryTab({ defaultTab: "api_keys" });
+  // const secondaryTab = useSecondaryTab({ defaultTab: "api_keys" });
+  const { queryStates } = useAppQueryStates({ defaultTab: "api_keys" });
+  const tab = queryStates.tab;
   const { pkey, webhooks } = useAutumnFlags();
 
   const apiKeys = data?.api_keys || [];
@@ -32,15 +35,15 @@ export default function DevScreen({ env }: { env: AppEnv }) {
       <div className="flex flex-col gap-4 h-fit relative w-full text-sm">
         <h1 className="text-xl font-medium shrink-0 pt-6 pl-10">Developer</h1>
 
-        {(secondaryTab === "api_keys" || !secondaryTab) && (
+        {(tab === "api_keys" || !tab) && (
           <div className="flex flex-col gap-16">
             <ApiKeysView apiKeys={apiKeys} />
             {pkey && <PublishableKeySection org={data.org} />}
           </div>
         )}
 
-        {secondaryTab === "stripe" && <ConfigureStripe />}
-        {secondaryTab === "webhooks" && webhooks && (
+        {tab === "stripe" && <ConfigureStripe />}
+        {tab === "webhooks" && webhooks && (
           <ConfigureWebhookSection dashboardUrl={data.svix_dashboard_url} />
         )}
       </div>

vite/src/views/onboarding/onboarding-steps/ProductList.tsx

@@ -26,7 +26,7 @@ import {
   TooltipContent,
   TooltipTrigger,
 } from "@/components/ui/tooltip";
-import CreateProduct from "@/views/products/products/components/CreateProduct";
+import CreateProduct from "@/views/products/products/components/CreateProductDialog";
 import { ProductsTable } from "@/views/products/products/components/ProductsTable";
 
 export const ProductList = ({

vite/src/views/onboarding2/ConnectStripeDialog.tsx

@@ -1,11 +1,6 @@
-import {
-  Dialog,
-  DialogContent,
-  DialogHeader,
-  DialogTitle,
-} from "@/components/ui/dialog";
+import { Dialog, DialogHeader, DialogTitle } from "@/components/ui/dialog";
 import { useModelPricingContext } from "./model-pricing/ModelPricingContext";
-import { ConnectStripeStep } from "./integrate/ConnectStripeStep";
+
 import {
   CustomDialogBody,
   CustomDialogContent,
@@ -16,6 +11,7 @@ import { useState } from "react";
 import { Input } from "@/components/ui/input";
 import { connectStripe } from "./utils/connectStripe";
 import { useAxiosInstance } from "@/services/useAxiosInstance";
+import { useOrg } from "@/hooks/common/useOrg";
 
 export default function ConnectStripeDialog({
   open,
@@ -24,15 +20,15 @@ export default function ConnectStripeDialog({
   open: boolean;
   setOpen: (open: boolean) => void;
 }) {
-  const { stripeConnected, mutate, productDataState, data } =
-    useModelPricingContext();
+  const { org, mutate: mutateOrg } = useOrg();
+
   const axiosInstance = useAxiosInstance();
 
   const [testApiKey, setTestApiKey] = useState("");
   const [loading, setLoading] = useState(false);
   const handleConnectStripe = async () => {
     setLoading(true);
-    await connectStripe({ testApiKey, axiosInstance, mutate });
+    await connectStripe({ testApiKey, axiosInstance, mutate: mutateOrg });
     setOpen(false);
     setLoading(false);
   };
@@ -59,9 +55,9 @@ export default function ConnectStripeDialog({
           <Input
             className="w-full"
             placeholder="Stripe secret key (sk_test_...)"
-            value={stripeConnected ? "Stripe connected  ✅ " : testApiKey}
+            value={org?.stripe_connected ? "Stripe connected  ✅ " : testApiKey}
             onChange={(e) => setTestApiKey(e.target.value)}
-            disabled={stripeConnected}
+            disabled={org?.stripe_connected}
           />
         </CustomDialogBody>
         <CustomDialogFooter>

vite/src/views/onboarding2/OnboardingView2.tsx

@@ -8,25 +8,29 @@ import { useEffect, useRef, useState } from "react";
 import { useAxiosInstance } from "@/services/useAxiosInstance";
 import { useSearchParams } from "react-router";
 import { useSession } from "@/lib/auth-client";
+import { useProductsQuery } from "@/hooks/queries/useProductsQuery";
+import { useOnboardingQueryState } from "./hooks/useOnboardingQueryState";
+import { useOrg } from "@/hooks/common/useOrg";
 
 export default function OnboardingView2() {
-  const [queryStates, setQueryStates] = useQueryStates(
-    {
-      page: parseAsString.withDefault("pricing"),
-      reactTypescript: parseAsBoolean.withDefault(true),
-      frontend: parseAsString.withDefault(""),
-      backend: parseAsString.withDefault(""),
-      auth: parseAsString.withDefault(""),
-      customerType: parseAsString.withDefault("user"),
-      productId: parseAsString.withDefault(""),
-    },
-    {
-      history: "push",
-    }
-  );
+  // const [queryStates, setQueryStates] = useQueryStates(
+  //   {
+  //     page: parseAsString.withDefault("pricing"),
+  //     reactTypescript: parseAsBoolean.withDefault(true),
+  //     frontend: parseAsString.withDefault(""),
+  //     backend: parseAsString.withDefault(""),
+  //     auth: parseAsString.withDefault(""),
+  //     customerType: parseAsString.withDefault("user"),
+  //     productId: parseAsString.withDefault(""),
+  //     token: parseAsString.withDefault(""),
+  //   },
+  //   {
+  //     history: "push",
+  //   }
+  // );
+
+  const { queryStates, setQueryStates } = useOnboardingQueryState();
 
-  const [searchParams] = useSearchParams();
-  const token = searchParams.get("token");
   const [loading, setLoading] = useState(true);
   const axiosInstance = useAxiosInstance();
   const hasHandledToken = useRef(false);
@@ -36,28 +40,21 @@ export default function OnboardingView2() {
   const {
     products: autumnProducts,
     isLoading: isAutumnLoading,
-    mutate: mutateAutumnProducts,
+    mutate: refetchAutumnProducts,
   } = useListProducts({ customerId: "onboarding_demo_user" });
 
-  const {
-    data: productsData,
-    mutate: productMutate,
-    isLoading,
-  } = useAxiosSWR({ url: `/products/data` });
-
-  const { data: productCounts, mutate: mutateCounts } = useAxiosSWR({
-    url: `/products/counts?latest_version=true`,
-  });
+  const { isLoading: productsLoading } = useProductsQuery();
+  const { isLoading: orgLoading } = useOrg();
 
   useEffect(() => {
     const handleToken = async () => {
       try {
         await axiosInstance.post("/onboarding", {
-          token,
+          token: queryStates.token,
         });
 
-        await productMutate();
-        await mutateAutumnProducts();
+        // await productMutate();
+        await refetchAutumnProducts();
       } catch (error) {
         console.error(error);
       } finally {
@@ -65,42 +62,44 @@ export default function OnboardingView2() {
       }
     };
 
-    if (token && !hasHandledToken.current) {
+    if (queryStates.token && !hasHandledToken.current) {
       hasHandledToken.current = true;
       handleToken();
     }
-  }, [searchParams, token, axiosInstance, productMutate]);
+  }, [queryStates.token, axiosInstance]);
 
   useEffect(() => {
-    if (orgId && !token) {
+    if (orgId && !queryStates.token) {
       setLoading(false);
     }
-  }, [orgId, token]);
+  }, [orgId, queryStates.token]);
 
-  if (isLoading || isAutumnLoading || loading) return <LoadingScreen />;
+  if (isAutumnLoading || loading || productsLoading || orgLoading)
+    return <LoadingScreen />;
 
   return (
     <>
       {queryStates.page === "integrate" ? (
         <IntegrateAutumn
-          data={productsData}
-          mutate={productMutate}
-          queryStates={queryStates}
-          setQueryStates={setQueryStates}
+        // data={productsData}
+        // mutate={productMutate}
+        // queryStates={queryStates}
+        // setQueryStates={setQueryStates}
         />
       ) : (
         <ModelPricing
-          data={productsData}
-          mutate={async () => {
-            await productMutate();
-            await mutateAutumnProducts();
-          }}
-          mutateAutumnProducts={mutateAutumnProducts}
+          // data={productsData}
+          // mutate={async () => {
+          //   await productMutate();
+          //   await mutateAutumnProducts();
+          // }}
+          // mutateAutumnProducts={mutateAutumnProducts}
           autumnProducts={autumnProducts}
-          productCounts={productCounts}
-          mutateCounts={mutateCounts}
-          queryStates={queryStates}
-          setQueryStates={setQueryStates}
+          refetchAutumnProducts={refetchAutumnProducts}
+          // productCounts={productCounts}
+          // mutateCounts={mutateCounts}
+          // queryStates={queryStates}
+          // setQueryStates={setQueryStates}
         />
       )}
     </>

vite/src/views/onboarding2/SampleApp.tsx

@@ -1,7 +1,4 @@
 import { Button } from "@/components/ui/button";
-import Step from "@/components/general/OnboardingStep";
-import CheckDialog from "@/components/autumn/paywall-dialog";
-
 import { useEnv } from "@/utils/envUtils";
 
 import {
@@ -10,7 +7,7 @@ import {
   DialogTrigger,
   DialogTitle,
 } from "@/components/ui/dialog";
-import { useEffect, useState } from "react";
+import { useState } from "react";
 import { Dialog } from "@/components/ui/dialog";
 
 import { toast } from "sonner";
@@ -30,12 +27,10 @@ import {
 } from "lucide-react";
 import { cn } from "@/lib/utils";
 import CodeBlock from "@/views/onboarding/components/CodeBlock";
-import { useListProducts } from "./model-pricing/usePricingTable";
+import { useFeaturesQuery } from "@/hooks/queries/useFeaturesQuery";
 
-export const SampleApp = ({ data }: { data: any }) => {
-  const env = useEnv();
+export const SampleApp = () => {
   const [searchParams] = useSearchParams();
-  const token = searchParams.get("token");
   const [open, setOpen] = useState(false);
   const [checkData, setCheckData] = useState<any>(null);
   const [trackData, setTrackData] = useState<any>(null);
@@ -43,14 +38,17 @@ export const SampleApp = ({ data }: { data: any }) => {
   const [showCheckSnippet, setShowCheckSnippet] = useState(true);
   const [showTrackSnippet, setShowTrackSnippet] = useState(true);
   const [showCustomerSnippet, setShowCustomerSnippet] = useState(true);
+
+  const { features, isLoading } = useFeaturesQuery();
+
   const [lastUsedFeature, setLastUsedFeature] = useState<any>({
-    featureId: data.features?.[0]?.id,
+    featureId: features?.[0]?.id,
     value: 1,
   });
 
   const { customer, openBillingPortal } = useCustomer();
-  const { products } = useListProducts({ customerId: "onboarding_demo_user" });
 
+  if (isLoading) return null;
   return (
     <Dialog open={open} onOpenChange={setOpen}>
       <DialogTrigger asChild>
@@ -97,10 +95,10 @@ export const SampleApp = ({ data }: { data: any }) => {
                 <h3 className="text-md font-medium">Available Features</h3>
               </div>
               <div className="flex flex-wrap gap-2 items-center">
-                {data.features
+                {features
                   ?.filter((feature: any) => customer?.features?.[feature.id])
                   .concat(
-                    data.features?.filter(
+                    features?.filter(
                       (feature: any) => !customer?.features?.[feature.id]
                     ) || []
                   )
@@ -183,7 +181,7 @@ const { allowed } = useCustomer();
 const handleCheckFeature = async () => {
 if ( !allowed({ featureId: '${
                                     lastUsedFeature?.featureId ||
-                                    data.features?.[0]?.id ||
+                                    features?.[0]?.id ||
                                     "feature-id"
                                   }' }) ) {
     alert('Feature not allowed');
@@ -207,7 +205,7 @@ const autumn = new Autumn({
 const { data } = await autumn.check({
   customerId: 'user_123',
   featureId: '${
-    lastUsedFeature?.featureId || data.features?.[0]?.id || "feature-id"
+    lastUsedFeature?.featureId || features?.[0]?.id || "feature-id"
   }'
 });
 `
@@ -261,7 +259,7 @@ const { track } = useCustomer();
 const handleTrackUsage = async () => {
   await track({ 
     featureId: '${
-      lastUsedFeature?.featureId || data.features?.[0]?.id || "feature-id"
+      lastUsedFeature?.featureId || features?.[0]?.id || "feature-id"
     }',
     value: ${lastUsedFeature?.value || 1}
   });
@@ -283,7 +281,7 @@ const autumn = new Autumn({
 
 const response = await autumn.track({
   featureId: '${
-    lastUsedFeature?.featureId || data.features?.[0]?.id || "feature-id"
+    lastUsedFeature?.featureId || features?.[0]?.id || "feature-id"
   }',
   value: ${lastUsedFeature?.value || 1}
 });
@@ -337,7 +335,7 @@ const { customer, refetch } = useCustomer();
 
 console.log('Customer balance:', customer?.features.${
                               lastUsedFeature?.featureId ||
-                              data.features?.[0]?.id ||
+                              features?.[0]?.id ||
                               "feature-id"
                             }?.balance);
 

vite/src/views/onboarding2/hooks/useOnboardingQueryState.tsx

@@ -0,0 +1,21 @@
+import { parseAsBoolean, parseAsString, useQueryStates } from "nuqs";
+
+export const useOnboardingQueryState = () => {
+  const [queryStates, setQueryStates] = useQueryStates(
+    {
+      page: parseAsString.withDefault("pricing"),
+      reactTypescript: parseAsBoolean.withDefault(true),
+      frontend: parseAsString.withDefault(""),
+      backend: parseAsString.withDefault(""),
+      auth: parseAsString.withDefault(""),
+      customerType: parseAsString.withDefault("user"),
+      productId: parseAsString.withDefault(""),
+      token: parseAsString.withDefault(""),
+    },
+    {
+      history: "push",
+    }
+  );
+
+  return { queryStates, setQueryStates };
+};

vite/src/views/onboarding2/integrate/ConnectStripeStep.tsx

@@ -7,14 +7,11 @@ import { AppEnv } from "@autumn/shared";
 import { toast } from "sonner";
 import { getBackendErr } from "@/utils/genUtils";
 import { StepHeader } from "./StepHeader";
+import { useOrg } from "@/hooks/common/useOrg";
+
+export const ConnectStripeStep = () => {
+  const { org, mutate: mutateOrg } = useOrg();
 
-export const ConnectStripeStep = ({
-  mutate,
-  productData,
-}: {
-  mutate: () => Promise<void>;
-  productData: any;
-}) => {
   const [testApiKey, setTestApiKey] = useState("");
 
   const [loading, setLoading] = useState(false);
@@ -29,7 +26,7 @@ export const ConnectStripeStep = ({
       });
 
       toast.success("Successfully connected to Stripe");
-      await mutate();
+      await mutateOrg();
     } catch (error) {
       console.log("Failed to connect Stripe", error);
       toast.error(getBackendErr(error, "Failed to connect Stripe"));
@@ -39,7 +36,7 @@ export const ConnectStripeStep = ({
   };
 
   // console.log("productData", productData);
-  const stripeConnected = productData?.org.stripe_connected;
+  const stripeConnected = org?.stripe_connected;
   return (
     <div className="w-full flex flex-col gap-4">
       <StepHeader number={2} title="Connect Stripe" />

vite/src/views/onboarding2/integrate/IntegrateAutumn.tsx

@@ -14,25 +14,15 @@ import { Button } from "@/components/ui/button";
 import { NextSteps } from "./NextSteps";
 import { AutumnProvider } from "autumn-js/react";
 import { ConnectStripeStep } from "./ConnectStripeStep";
+import { useOnboardingQueryState } from "../hooks/useOnboardingQueryState";
 
-export default function IntegrateAutumn({
-  data,
-  mutate,
-  queryStates,
-  setQueryStates,
-}: {
-  data: any;
-  mutate: any;
-  queryStates: any;
-  setQueryStates: any;
-}) {
+export default function IntegrateAutumn() {
+  const { queryStates, setQueryStates } = useOnboardingQueryState();
   const stackSelected =
     queryStates.frontend && queryStates.backend && queryStates.auth;
 
   return (
-    <IntegrateContext.Provider
-      value={{ queryStates, setQueryStates, data, mutate }}
-    >
+    <IntegrateContext.Provider value={{ queryStates, setQueryStates }}>
       <div className="w-full h-full p-10 flex flex-col items-center justify-start overflow-y-scroll">
         <div className="max-w-[600px] w-full flex flex-col gap-6">
           <div className="flex flex-col gap-2">
@@ -60,7 +50,7 @@ export default function IntegrateAutumn({
             <SelectStack />
             {stackSelected && queryStates.reactTypescript && (
               <>
-                <ConnectStripeStep mutate={mutate} productData={data} />
+                <ConnectStripeStep />
                 <EnvStep />
                 <Install />
                 <AutumnHandler />

vite/src/views/onboarding2/model-pricing/AddTrialButton.tsx

@@ -2,14 +2,14 @@ import { Button } from "@/components/ui/button";
 import { CreateFreeTrial } from "@/views/products/product/free-trial/CreateFreeTrial";
 import { useState } from "react";
 import { useProductContext } from "@/views/products/product/ProductContext";
-import { PlusIcon, Trash, X } from "lucide-react";
+import { PlusIcon, X } from "lucide-react";
 import { handleAutoSave } from "./model-pricing-utils/modelPricingUtils";
 import { useAxiosInstance } from "@/services/useAxiosInstance";
 
 export const AddTrialButton = () => {
-  const { product, setProduct, mutate, autoSave } = useProductContext();
-  const [open, setOpen] = useState(false);
   const axiosInstance = useAxiosInstance();
+  const { product, setProduct, autoSave, refetch } = useProductContext();
+  const [open, setOpen] = useState(false);
 
   return (
     <>
@@ -40,12 +40,12 @@ export const AddTrialButton = () => {
                   free_trial: null,
                 });
 
-                if (!autoSave) {
+                if (autoSave) {
                   handleAutoSave({
                     axiosInstance,
                     productId: product.id,
                     product: { ...product, free_trial: null },
-                    mutate,
+                    refetch,
                   });
                 }
               }}

vite/src/views/onboarding2/model-pricing/EditProduct.tsx

@@ -18,19 +18,25 @@ import { updateProduct } from "@/views/products/product/utils/updateProduct";
 import { getBackendErr } from "@/utils/genUtils";
 import { EditProductDetails } from "./edit-product/EditProductDetails";
 import { ToggleDefaultProduct } from "@/views/products/product/product-sidebar/ToggleDefaultProduct";
+import { useFeaturesQuery } from "@/hooks/queries/useFeaturesQuery";
+import { useProductsQuery } from "@/hooks/queries/useProductsQuery";
 
-export const EditProduct = ({ mutate }: { mutate: any }) => {
-  const [freeTrialModalOpen, setFreeTrialModalOpen] = useState(false);
-  const { data, productCount, productDataState, mutateCounts } =
-    useModelPricingContext();
+export const EditProduct = ({
+  refetchAutumnProducts,
+}: {
+  refetchAutumnProducts: () => Promise<void>;
+}) => {
+  const { refetch, counts, products } = useProductsQuery();
+  const { productDataState } = useModelPricingContext();
 
   const [showNewVersionDialog, setShowNewVersionDialog] = useState(false);
+  const [freeTrialModalOpen, setFreeTrialModalOpen] = useState(false);
+
+  const productCount = counts?.[productDataState.product?.id];
 
   const {
     product,
     setProduct,
-    features,
-    setFeatures,
     entityFeatureIds,
     setEntityFeatureIds,
     actionState,
@@ -42,24 +48,16 @@ export const EditProduct = ({ mutate }: { mutate: any }) => {
   });
 
   const axiosInstance = useAxiosInstance();
-  const env = useEnv();
-
   const [saveLoading, setSaveLoading] = useState(false);
 
   const runUpdateProduct = async () => {
     setSaveLoading(true);
-    try {
-      await updateProduct({
-        axiosInstance,
-        product,
-        mutate,
-        mutateCount: mutateCounts,
-      });
-    } catch (error) {
-      toast.error(getBackendErr(error, "Failed to update product"));
-    } finally {
-      setSaveLoading(false);
-    }
+    await updateProduct({
+      axiosInstance,
+      product,
+      onSuccess: refetch,
+    });
+    setSaveLoading(false);
   };
 
   const handleSaveClicked = async () => {
@@ -73,9 +71,7 @@ export const EditProduct = ({ mutate }: { mutate: any }) => {
   const hasItems = product.items.length > 0;
   const hasCustomers = productCount?.all > 0;
   const showSaveButton = hasCustomers || product.version > 1;
-  const firstProductCreated = data.products.length > 0;
-
-  const autoSave = !showSaveButton && firstProductCreated;
+  const firstProductCreated = products.length > 0;
 
   const handleToggleSettings = async (key: string) => {
     if (!product) return;
@@ -96,7 +92,7 @@ export const EditProduct = ({ mutate }: { mutate: any }) => {
         axiosInstance,
         productId: product.id ? product.id : details.id,
         product: { ...product, [key]: !curValue },
-        mutate,
+        refetch,
       });
     }
   };
@@ -106,13 +102,13 @@ export const EditProduct = ({ mutate }: { mutate: any }) => {
       <div className="flex gap-4 transition-all duration-500 ease-in-out">
         <ProductContext.Provider
           value={{
-            groupDefaults: data.groupToDefaults?.[product?.group || ""],
+            // groupDefaults: data.groupToDefaults?.[product?.group || ""],
             product,
             setProduct,
-            mutate,
-            env,
-            features,
-            setFeatures,
+            refetch: async () => {
+              console.log("Refetching edit product...");
+              await Promise.all([refetch(), refetchAutumnProducts()]);
+            },
             entityFeatureIds,
             setEntityFeatureIds,
             isOnboarding: true,
@@ -122,7 +118,6 @@ export const EditProduct = ({ mutate }: { mutate: any }) => {
           <ConfirmNewVersionDialog
             open={showNewVersionDialog}
             setOpen={setShowNewVersionDialog}
-            createProduct={runUpdateProduct}
           />
           <CreateFreeTrial
             open={freeTrialModalOpen}
@@ -182,13 +177,6 @@ export const EditProduct = ({ mutate }: { mutate: any }) => {
           >
             <div className="flex flex-col gap-4" style={{ width: "320px" }}>
               <div>
-                {/* <ToggleButton
-                    disabled={product?.is_add_on}
-                    buttonText="Default Product"
-                    value={product?.is_default}
-                    className="text-t2 font-medium h-fit mb-2"
-                    setValue={() => handleToggleSettings("is_default")}
-                  /> */}
                 <div className="flex items-center text-sm text-t2 gap-2">
                   <p className="text-t2 font-medium">Default Product</p>
                   <ToggleDefaultProduct toggleKey="is_default" />
@@ -199,13 +187,6 @@ export const EditProduct = ({ mutate }: { mutate: any }) => {
                 </div>
               </div>
               <div className="">
-                {/* <ToggleButton
-                    disabled={product?.is_default}
-                    buttonText="Add-on Product"
-                    className="text-t2 font-medium h-fit mb-2"
-                    value={product?.is_add_on}
-                    setValue={() => handleToggleSettings("is_add_on")}
-                  /> */}
                 <div className="flex items-center text-sm text-t2 gap-2">
                   <p className="text-t2 font-medium">Add On Product</p>
                   <ToggleDefaultProduct toggleKey="is_add_on" />

vite/src/views/onboarding2/model-pricing/EditProductOld.tsx

@@ -1,182 +0,0 @@
-import { DialogFooter } from "@/components/ui/dialog";
-import { ProductService } from "@/services/products/ProductService";
-import { useAxiosInstance } from "@/services/useAxiosInstance";
-import { useEnv } from "@/utils/envUtils";
-import { getBackendErr } from "@/utils/genUtils";
-import { CreateFreeTrial } from "@/views/products/product/free-trial/CreateFreeTrial";
-import { ManageProduct } from "@/views/products/product/ManageProduct";
-import { ProductContext } from "@/views/products/product/ProductContext";
-import { Dialog, DialogContent, DialogTitle } from "@/components/ui/dialog";
-import { X } from "lucide-react";
-import { useState } from "react";
-import { Button } from "@/components/ui/button";
-import { toast } from "sonner";
-import { ToggleButton } from "@/components/general/ToggleButton";
-
-export const EditProductDialog = ({
-  product,
-  features,
-  setProduct,
-  setFeatures,
-  mutate,
-  open,
-  setOpen,
-  originalProduct,
-  entityFeatureIds,
-  setEntityFeatureIds,
-}: {
-  product: any;
-  setProduct: (product: any) => void;
-  features: any[];
-  setFeatures: (features: any[]) => void;
-  mutate: () => Promise<void>;
-  open: boolean;
-  setOpen: (open: boolean) => void;
-  originalProduct: any;
-  entityFeatureIds: string[];
-  setEntityFeatureIds: (entityFeatureIds: string[]) => void;
-}) => {
-  const env = useEnv();
-  const axiosInstance = useAxiosInstance();
-  const [createProductLoading, setCreateProductLoading] = useState(false);
-  const [freeTrialModalOpen, setFreeTrialModalOpen] = useState(false);
-
-  // Store the original product state when modal opens
-  const handleOpenChange = async (newOpen: boolean) => {
-    if (!newOpen && open && product?.id) {
-      // Modal is being closed, check if there are changes
-      const hasChanges =
-        originalProduct &&
-        JSON.stringify(product) !== JSON.stringify(originalProduct);
-
-      if (hasChanges) {
-        // Only update if there are changes
-        updateProduct();
-      }
-    }
-    setOpen(newOpen);
-  };
-
-  const updateProduct = async () => {
-    setCreateProductLoading(true);
-    try {
-      const res = await ProductService.updateProduct(
-        axiosInstance,
-        product.id,
-        product
-      );
-      toast.success("Product updated successfully");
-      await mutate();
-      setOpen(false);
-    } catch (error) {
-      toast.error(getBackendErr(error, "Failed to update product"));
-    }
-    setCreateProductLoading(false);
-  };
-
-  const handleFreeTrialClick = () => {
-    if (product?.free_trial) {
-      // Delete the free trial
-      setProduct({ ...product, free_trial: null });
-    } else {
-      // Open the free trial modal
-      setFreeTrialModalOpen(true);
-    }
-  };
-
-  return (
-    <Dialog open={open} onOpenChange={handleOpenChange}>
-      <DialogContent className="p-0 py-8 min-w-[500px] min-h-[300px] flex flex-col justify-between">
-        <DialogTitle className="text-t2 font-semibold px-10 hidden">
-          {/* Edit Product */}
-        </DialogTitle>
-        <div>
-          <ProductContext.Provider
-            value={{
-              product,
-              setProduct,
-              mutate,
-              env,
-              features,
-              setFeatures,
-              entityFeatureIds,
-              setEntityFeatureIds,
-            }}
-          >
-            <CreateFreeTrial
-              open={freeTrialModalOpen}
-              setOpen={setFreeTrialModalOpen}
-            />
-            <ManageProduct hideAdminHover={true} />
-          </ProductContext.Provider>
-        </div>
-        <DialogFooter>
-          <div className="flex justify-between items-center gap-2 px-10 w-full mt-6">
-            <div className="flex gap-6">
-              <ToggleButton
-                disabled={product?.is_add_on}
-                buttonText="Default"
-                infoContent="This product is enabled by default for all new users, typically used for your free plan"
-                value={product?.is_default}
-                setValue={() =>
-                  setProduct({
-                    ...product,
-                    is_default: !product?.is_default,
-                  })
-                }
-              />
-              <ToggleButton
-                disabled={product?.is_default}
-                buttonText="Add-on"
-                infoContent="This product is an add-on that can be bought together with your base products (eg, for top ups)"
-                value={product?.is_add_on}
-                setValue={() =>
-                  setProduct({ ...product, is_add_on: !product?.is_add_on })
-                }
-              />
-              <Button
-                variant="outline"
-                onClick={handleFreeTrialClick}
-                className={`min-w-32 flex items-center gap-2`}
-              >
-                {product?.free_trial ? (
-                  <div className="flex items-center gap-2 justify-between w-full">
-                    <p>
-                      {product?.free_trial?.length}{" "}
-                      {product.free_trial?.duration} trial
-                    </p>
-                    <Button
-                      variant="ghost"
-                      size="icon"
-                      isIcon
-                      onClick={() => {
-                        setProduct({ ...product, free_trial: null });
-                      }}
-                      className="hover:bg-zinc-300 !h-4 !w-4 text-t3"
-                    >
-                      <X size={12} />
-                    </Button>
-                  </div>
-                ) : (
-                  // <div className="w-3 h-3 bg-lime-500 rounded-full flex items-center justify-center">
-
-                  //   {/* <Check className="w-2 h-2 text-white" /> */}
-                  // </div>
-                  <p>Add Free Trial</p>
-                )}
-              </Button>
-            </div>
-            <Button
-              isLoading={createProductLoading}
-              variant="gradientPrimary"
-              onClick={updateProduct}
-              className="min-w-44 w-44 max-w-44"
-            >
-              Update Product
-            </Button>
-          </div>
-        </DialogFooter>
-      </DialogContent>
-    </Dialog>
-  );
-};

vite/src/views/onboarding2/model-pricing/ModelPricing.tsx

@@ -26,6 +26,11 @@ import { ConnectStripeStep } from "../integrate/ConnectStripeStep";
 import { AutumnProvider } from "autumn-js/react";
 import { useProductData } from "@/views/products/product/hooks/useProductData";
 import ConnectStripeDialog from "../ConnectStripeDialog";
+import { useOnboardingQueryState } from "../hooks/useOnboardingQueryState";
+import { useProductsQuery } from "@/hooks/queries/useProductsQuery";
+import { ProductV2 } from "@autumn/shared";
+import { useOrg } from "@/hooks/common/useOrg";
+import { useProductContext } from "@/views/products/product/ProductContext";
 
 const defaultProduct = {
   id: "",
@@ -37,46 +42,44 @@ const defaultProduct = {
 };
 
 export const ModelPricing = ({
-  data,
-  mutate,
-  mutateAutumnProducts,
+  // data,
+  // mutate,
+  // mutateAutumnProducts,
+  refetchAutumnProducts,
   autumnProducts,
-  productCounts,
-  mutateCounts,
-  queryStates,
-  setQueryStates,
+  // productCounts,
+  // mutateCounts,
+  // queryStates,
+  // setQueryStates,
 }: {
-  data: any;
-  mutate: any;
-  mutateAutumnProducts: any;
+  // data: any;
+  // mutate: any;
+  // mutateAutumnProducts: any;
+  refetchAutumnProducts: any;
   autumnProducts: Product[];
-  productCounts: any;
-  mutateCounts: any;
-  queryStates: any;
-  setQueryStates: any;
+  // productCounts: any;
+  // mutateCounts: any;
+  // queryStates: any;
+  // setQueryStates: any;
 }) => {
+  const { queryStates, setQueryStates } = useOnboardingQueryState();
+  const { products, refetch } = useProductsQuery();
+
   const getCurProduct = () => {
     if (queryStates.productId) {
-      const prod = data.products.find(
-        (p: Product) => p.id === queryStates.productId
+      const prod = products.find(
+        (p: ProductV2) => p.id === queryStates.productId
       );
-      if (prod) {
-        return prod;
-      }
+
+      if (prod) return prod;
     }
 
-    if (data.products.length > 0) {
-      return data.products[0];
-    }
+    if (products.length > 0) return products[0];
 
     return defaultProduct;
   };
 
   const curProduct = getCurProduct();
-  // // console.log("curProduct:", curProduct);
-  // const [firstItemCreated, setFirstItemCreated] = useState(
-  //   autumnProducts.some((p: Product) => p.items.length > 0)
-  // );
 
   const firstItemCreated = autumnProducts.some(
     (p: Product) => p.items.length > 0
@@ -90,14 +93,19 @@ export const ModelPricing = ({
 
   const productDataState = useProductData({
     originalProduct: curProduct as any,
-    originalFeatures: data.features as any,
   });
 
   const { product } = productDataState;
 
-  if (!product) return null;
+  useEffect(() => {
+    if (curProduct && !queryStates.productId) {
+      setQueryStates({
+        productId: curProduct.id,
+      });
+    }
+  }, [curProduct]);
 
-  const stripeConnected = data?.org.stripe_connected;
+  if (!product) return null;
 
   return (
     <ModelPricingContext.Provider
@@ -106,16 +114,24 @@ export const ModelPricing = ({
         // setFirstItemCreated,
         editingNewProduct,
         setEditingNewProduct,
+
+        // For <EditProduct />
+        productDataState,
+
+        // For <CheckoutDialog /> & <NewProductPopover />
+        refetch: async () => {
+          await Promise.all([refetch(), refetchAutumnProducts()]);
+        },
+
         // product,
         // setProduct,
-        productDataState,
-        mutate,
-        data,
-        productCount: productCounts?.[product?.id ?? ""],
-        queryStates,
-        setQueryStates,
-        mutateAutumnProducts,
-        mutateCounts,
+        // mutate,
+        // data,
+        // productCount: productCounts?.[product?.id ?? ""],
+        // queryStates,
+        // setQueryStates,
+        // mutateAutumnProducts,
+        // mutateCounts,
         setConnectStripeOpen,
       }}
     >
@@ -123,77 +139,76 @@ export const ModelPricing = ({
         open={connectStripeOpen}
         setOpen={setConnectStripeOpen}
       />
-      <ProductsContext.Provider value={{ productCounts, mutate }}>
-        <div className="flex flex-col w-full h-full items-center justify-between overflow-y-auto overflow-x-hidden">
-          <div className="w-full p-10 flex flex-col gap-4 justify-center items-center">
-            <div className="max-w-[800px] w-full">
-              <div className="flex gap-4 items-center justify-between mb-6">
-                <div className="flex flex-col gap-1">
-                  <p className="text-xl font-medium">Create your products</p>
-                  <p className="text-t2 text-sm w-md">
-                    To start, model your app's pricing by creating a product for
-                    your free plans, paid plans and any add-ons or top-ups.
-                  </p>
-                </div>
-                {firstItemCreated && (
-                  <div className="flex gap-0 items-center">
-                    <SelectEditProduct />
-                    {data.products.length > 0 && <NewProductPopover />}
-                  </div>
-                )}
-              </div>
-              <div className="flex flex-col gap-4 w-full">
-                <EditProduct mutate={mutate} />
-              </div>
-            </div>
-          </div>
 
-          <div className="w-full">
-            <div
-              className={cn(
-                "w-full px-10 flex flex-col gap-4 items-center ",
-                "transition-all duration-800 overflow-hidden py-10 pb-4",
-                !firstItemCreated && "hidden"
-                // firstItemCreated ? "max-h-[800px] py-10" : "max-h-0 py-0"
+      <div className="flex flex-col w-full h-full items-center justify-between overflow-y-auto overflow-x-hidden">
+        <div className="w-full p-10 flex flex-col gap-4 justify-center items-center">
+          <div className="max-w-[800px] w-full">
+            <div className="flex gap-4 items-center justify-between mb-6">
+              <div className="flex flex-col gap-1">
+                <p className="text-xl font-medium">Create your products</p>
+                <p className="text-t2 text-sm w-md">
+                  To start, model your app's pricing by creating a product for
+                  your free plans, paid plans and any add-ons or top-ups.
+                </p>
+              </div>
+              {firstItemCreated && (
+                <div className="flex gap-0 items-center">
+                  <SelectEditProduct />
+                  {products.length > 0 && <NewProductPopover />}
+                </div>
               )}
-            >
-              <AutumnProvider
-                backendUrl={`${import.meta.env.VITE_BACKEND_URL}/demo`}
-                includeCredentials={true}
-              >
-                <div className="gap-8 flex justify-center max-w-[800px] w-full flex-col">
-                  <PricingTable
-                    products={autumnProducts}
-                    stripeConnected={stripeConnected}
-                    setConnectStripeOpen={setConnectStripeOpen}
-                  />
-                </div>
-              </AutumnProvider>
-
-              <div className="w-full flex justify-end translate-x-6">
-                <Button
-                  onClick={() => {
-                    setQueryStates({
-                      page: "integrate",
-                    });
-                  }}
-                  className="bg-zinc-800 w-fit hover:bg-zinc-700"
-                >
-                  Next: Integrate Autumn <ArrowRight className="w-4 h-4" />
-                </Button>
-              </div>
+            </div>
+            <div className="flex flex-col gap-4 w-full">
+              <EditProduct refetchAutumnProducts={refetchAutumnProducts} />
             </div>
           </div>
         </div>
-      </ProductsContext.Provider>
+
+        <div className="w-full">
+          <div
+            className={cn(
+              "w-full px-10 flex flex-col gap-4 items-center ",
+              "transition-all duration-800 overflow-hidden py-10 pb-4",
+              !firstItemCreated && "hidden"
+              // firstItemCreated ? "max-h-[800px] py-10" : "max-h-0 py-0"
+            )}
+          >
+            <AutumnProvider
+              backendUrl={`${import.meta.env.VITE_BACKEND_URL}/demo`}
+              includeCredentials={true}
+            >
+              <div className="gap-8 flex justify-center max-w-[800px] w-full flex-col">
+                <PricingTable
+                  products={autumnProducts}
+                  setConnectStripeOpen={setConnectStripeOpen}
+                />
+              </div>
+            </AutumnProvider>
+
+            <div className="w-full flex justify-end translate-x-6">
+              <Button
+                onClick={() => {
+                  setQueryStates({
+                    page: "integrate",
+                  });
+                }}
+                className="bg-zinc-800 w-fit hover:bg-zinc-700"
+              >
+                Next: Integrate Autumn <ArrowRight className="w-4 h-4" />
+              </Button>
+            </div>
+          </div>
+        </div>
+      </div>
     </ModelPricingContext.Provider>
   );
 };
 
 const NewProductPopover = () => {
   const [open, setOpen] = useState(false);
+
   const {
-    mutate,
+    refetch,
     productDataState: { setProduct },
   } = useModelPricingContext();
 
@@ -212,7 +227,7 @@ const NewProductPopover = () => {
         name: details.name,
         id: details.id,
       });
-      await mutate();
+      await refetch();
       const newProduct = {
         ...defaultProduct,
         name: details.name,

vite/src/views/onboarding2/model-pricing/ModelPricingContext.tsx

@@ -5,11 +5,11 @@ export const ModelPricingContext = createContext<any>(null);
 export const useModelPricingContext = () => {
   const context = useContext(ModelPricingContext);
 
-  if (context === undefined) {
-    throw new Error(
-      "useProductContext must be used within a ProductContextProvider"
-    );
-  }
+  // if (context === undefined) {
+  //   throw new Error(
+  //     "useProductContext must be used within a ProductContextProvider"
+  //   );
+  // }
 
   return context;
 };

vite/src/views/onboarding2/model-pricing/SelectEditProduct.tsx

@@ -11,23 +11,27 @@ import { ChevronDownIcon, Trash } from "lucide-react";
 import { DeleteProductDialog } from "@/views/products/products/product-row-toolbar/DeleteProductDialog";
 import { useState } from "react";
 import { cn } from "@/lib/utils";
+import { useProductsQuery } from "@/hooks/queries/useProductsQuery";
+import { useOnboardingQueryState } from "../hooks/useOnboardingQueryState";
 
 export const SelectEditProduct = () => {
   const {
-    data,
     productDataState: { product, setProduct },
-    queryStates,
-    setQueryStates,
   } = useModelPricingContext();
+
+  const { products } = useProductsQuery();
+  const { setQueryStates } = useOnboardingQueryState();
+
   const [deleteProductOpen, setDeleteProductOpen] = useState(false);
   const [deleteProductSelected, setDeleteProductSelected] = useState(product);
 
   const selectedClassName = "!bg-zinc-100 h-7 border";
-  if (data.products.length > 3) {
+  if (products.length > 3) {
     return (
       <>
         <DeleteProductDialog
           product={deleteProductSelected}
+          dropdownOpen={deleteProductOpen}
           open={deleteProductOpen}
           setOpen={setDeleteProductOpen}
         />
@@ -42,7 +46,7 @@ export const SelectEditProduct = () => {
             </Button>
           </DropdownMenuTrigger>
           <DropdownMenuContent align="start" className="max-h-[600px]">
-            {data.products.map((p: any) => {
+            {products.map((p: any) => {
               if (!p.name) {
                 return null;
               }
@@ -67,6 +71,7 @@ export const SelectEditProduct = () => {
                       e.preventDefault();
                       setDeleteProductSelected(p);
                       setDeleteProductOpen(true);
+                      console.log("Delete product selected:", p);
                     }}
                   >
                     <Trash size={12} />
@@ -89,10 +94,11 @@ export const SelectEditProduct = () => {
         product={product}
         open={deleteProductOpen}
         setOpen={setDeleteProductOpen}
-      ></DeleteProductDialog>
+        dropdownOpen={deleteProductOpen}
+      />
       <Tabs className="" value={product.id}>
         <TabsList className="gap-1 mr-1">
-          {data.products.map((p: any) => {
+          {products.map((p: any) => {
             if (!p.name) {
               return null;
             }

vite/src/views/onboarding2/model-pricing/edit-product/EditProductDetails.tsx

@@ -10,15 +10,15 @@ import { useProductContext } from "@/views/products/product/ProductContext";
 import { InfoTooltip } from "@/components/general/modal-components/InfoTooltip";
 import { getBackendErr } from "@/utils/genUtils";
 import { toast } from "sonner";
+import { useProductsQuery } from "@/hooks/queries/useProductsQuery";
 
 export const EditProductDetails = () => {
-  const {
-    productDataState: { product, setProduct },
-    data,
-    mutate,
-  } = useModelPricingContext();
+  const { productDataState } = useModelPricingContext();
+  const { product, setProduct } = productDataState;
 
-  const allowCreate = data.products.length === 0;
+  const { products, refetch } = useProductsQuery();
+
+  const allowCreate = products.length === 0;
   const { autoSave } = useProductContext();
   const [createLoading, setCreateLoading] = useState(false);
 
@@ -44,7 +44,7 @@ export const EditProductDetails = () => {
         name: details.name,
         id: details.id,
       });
-      await mutate();
+      await refetch();
     } catch (error) {
       toast.error(getBackendErr(error, "Failed to create product"));
     } finally {
@@ -67,7 +67,7 @@ export const EditProductDetails = () => {
                   name: details.name,
                   id: details.id,
                 },
-                mutate,
+                refetch,
               });
             }
             setProduct({

vite/src/views/onboarding2/model-pricing/model-pricing-utils/modelPricingUtils.ts

@@ -7,12 +7,12 @@ export const handleAutoSave = async ({
   axiosInstance,
   productId,
   product,
-  mutate,
+  refetch,
 }: {
   axiosInstance: AxiosInstance;
   productId: string;
   product: ProductV2;
-  mutate: any;
+  refetch?: any;
 }) => {
   if (!productId || !product.id) return;
   try {
@@ -23,7 +23,9 @@ export const handleAutoSave = async ({
         group: notNullish(product.group) ? product.group : undefined,
       }
     );
-    await mutate();
+    if (refetch) {
+      await refetch();
+    }
   } catch (error) {
     console.log(error);
     toast.error(getBackendErr(error, "Failed to auto save product"));

vite/src/views/products/product/ProductProps.tsx

@@ -26,16 +26,17 @@ import { ToggleButton } from "@/components/general/ToggleButton";
 import { InfoTooltip } from "@/components/general/modal-components/InfoTooltip";
 import { ToggleDefaultProduct } from "./product-sidebar/ToggleDefaultProduct";
 import { getBackendErr } from "@/utils/genUtils";
+import { useProductCountsQuery } from "./hooks/queries/useProductCountsQuery";
+import { useProductQuery } from "./hooks/useProductQuery";
 
 export const ProductProps = () => {
-  const { product, setProduct, counts, mutate } = useProductContext();
   const axiosInstance = useAxiosInstance();
-  const [defaultOpen, setDefaultOpen] = React.useState(false);
-  const [defaultTrialOpen, setDefaultTrialOpen] = React.useState(false);
-  const [addOnOpen, setAddOnOpen] = React.useState(false);
+  const { product, setProduct } = useProductContext();
+  const { refetch } = useProductQuery();
+  const { counts } = useProductCountsQuery();
+
   const [groupModalOpen, setGroupModalOpen] = React.useState(false);
   const [tempGroup, setTempGroup] = React.useState(product.group || "");
-  const [archivedOpen, setArchivedOpen] = React.useState(false);
 
   return (
     <>
@@ -184,7 +185,7 @@ export const ProductProps = () => {
                       { archived: value },
                       product.version
                     );
-                    await mutate();
+                    await refetch();
                     toast.success(
                       value
                         ? "Product archived successfully"

vite/src/views/products/product/ProductSidebar.tsx

@@ -55,11 +55,11 @@ export default function ProductSidebar() {
             <ProductProps />
           </SideAccordion>
         </div>
-        {/* <div className="flex w-full border-b p-4">
+        <div className="flex w-full border-b p-4">
           <SideAccordion title="Versions" value="versions">
             <ProductVersions />
           </SideAccordion>
-        </div> */}
+        </div>
         <div className="flex w-full border-b p-4 relative">
           <SideAccordion
             title="Entities"

vite/src/views/products/product/ProductVersions.tsx

@@ -11,11 +11,12 @@ import {
 } from "@/components/ui/select";
 import { CountAndMigrate } from "./versioning/CountAndMigrate";
 import { formatUnixToDate } from "@/utils/formatUtils/formatDateUtils";
+import { useProductQuery, useProductQueryState } from "./hooks/useProductQuery";
 
 export const ProductVersions = () => {
-  const { product, numVersions, customer, version } = useProductContext();
-  const navigate = useNavigate();
-  const env = useEnv();
+  const { setQueryStates } = useProductQueryState();
+  const { product, isCusProductView } = useProductContext();
+  const { numVersions } = useProductQuery();
 
   return (
     <div className="flex justify-between gap-4 w-full text-xs">
@@ -34,16 +35,20 @@ export const ProductVersions = () => {
           <p className=" text-t3 font-medium text-center">Version History </p>
           {numVersions > 1 ? (
             <Select
-              value={version ? version.toString() : product.version.toString()}
+              value={product.version.toString()}
               onValueChange={async (value) => {
-                navigate(
-                  getRedirectUrl(
-                    `${customer ? `/customers/${customer.id}` : "/products"}/${
-                      product.id
-                    }?version=${value}`,
-                    env
-                  )
-                );
+                setQueryStates({
+                  version: parseInt(value),
+                });
+                // Set query states...
+                // navigate(
+                //   getRedirectUrl(
+                //     `${customer ? `/customers/${customer.id}` : "/products"}/${
+                //       product.id
+                //     }?version=${value}`,
+                //     env
+                //   )
+                // );
               }}
             >
               <SelectTrigger
@@ -61,16 +66,9 @@ export const ProductVersions = () => {
                       value={version.toString()}
                       className="px-3 py-2 hover:bg-gray-100 cursor-pointer w-full"
                       onClick={() => {
-                        navigate(
-                          getRedirectUrl(
-                            `${
-                              customer
-                                ? `/customers/${customer.id}`
-                                : "/products"
-                            }/${product.id}?version=${version}`,
-                            env
-                          )
-                        );
+                        setQueryStates({
+                          version: version,
+                        });
                       }}
                     >
                       v{version}
@@ -82,7 +80,7 @@ export const ProductVersions = () => {
             <p className="text-sm text-t3 pr-2">None</p>
           )}
         </div>
-        {!customer && <CountAndMigrate />}
+        {!isCusProductView && <CountAndMigrate />}
       </div>
     </div>
   );

vite/src/views/products/product/ProductView.tsx

@@ -4,36 +4,22 @@ import ErrorScreen from "@/views/general/ErrorScreen";
 import LoadingScreen from "@/views/general/LoadingScreen";
 import ProductSidebar from "./ProductSidebar";
 import ProductViewBreadcrumbs from "./components/ProductViewBreadcrumbs";
-import ConfirmNewVersionDialog from "./versioning/ConfirmNewVersionDialog";
 
-import { toast } from "sonner";
 import { useState } from "react";
-import { useAxiosSWR } from "@/services/useAxiosSwr";
 import { ProductContext } from "./ProductContext";
 import { useParams, useSearchParams } from "react-router";
 import { useAxiosInstance } from "@/services/useAxiosInstance";
 import { ManageProduct } from "./ManageProduct";
-import { AppEnv, UpdateProductSchema } from "@autumn/shared";
-import { ProductService } from "@/services/products/ProductService";
-import { getBackendErr } from "@/utils/genUtils";
+import { AppEnv } from "@autumn/shared";
 import { useProductChangedAlert } from "./hooks/useProductChangedAlert";
 import { useProductData } from "./hooks/useProductData";
 import { UpdateProductButton } from "./components/UpdateProductButton";
 import { useProductQuery } from "./hooks/useProductQuery";
+import ConfirmNewVersionDialog from "./versioning/ConfirmNewVersionDialog";
 
-function ProductView({ env }: { env: AppEnv }) {
-  const axiosInstance = useAxiosInstance();
+function ProductView() {
   const { product_id } = useParams();
-  const [searchParams] = useSearchParams();
-  const version = searchParams.get("version");
 
-  // const url = `/products/${product_id}/data?version=${version}`;
-  // const { data, isLoading, mutate } = useAxiosSWR({ url });
-
-  // const countUrl = `/products/${product_id}/count?version=${version}`;
-  // const { data: counts, mutate: mutateCount } = useAxiosSWR({ url: countUrl });
-
-  const [showNewVersionDialog, setShowNewVersionDialog] = useState(false);
   const { product: originalProduct, isLoading, error } = useProductQuery();
 
   const {
@@ -43,10 +29,10 @@ function ProductView({ env }: { env: AppEnv }) {
     entityFeatureIds,
     setEntityFeatureIds,
     actionState,
-    isNewProduct,
   } = useProductData({ originalProduct });
 
   const { modal } = useProductChangedAlert({ hasChanges });
+  const [showNewVersionDialog, setShowNewVersionDialog] = useState(false);
 
   if (isLoading) return <LoadingScreen />;
   if (error) {
@@ -113,6 +99,7 @@ function ProductView({ env }: { env: AppEnv }) {
         // mutate,
         // env,
 
+        setShowNewVersionDialog,
         product,
         setProduct,
 
@@ -124,11 +111,10 @@ function ProductView({ env }: { env: AppEnv }) {
         hasChanges,
       }}
     >
-      {/* <ConfirmNewVersionDialog
+      <ConfirmNewVersionDialog
         open={showNewVersionDialog}
         setOpen={setShowNewVersionDialog}
-        createProduct={updateProduct}
-      /> */}
+      />
       <div className="flex w-full">
         <div className="flex flex-col gap-4 w-full">
           <ProductViewBreadcrumbs />

vite/src/views/products/product/components/UpdateProductButton.tsx

@@ -4,51 +4,42 @@ import { useProductContext } from "@/views/products/product/ProductContext";
 import { Upload } from "lucide-react";
 import { useProductQuery } from "../hooks/useProductQuery";
 import { useAxiosInstance } from "@/services/useAxiosInstance";
-import { UpdateProductSchema } from "@autumn/shared";
 import { toast } from "sonner";
-import { getBackendErr } from "@/utils/genUtils";
+import { useProductCountsQuery } from "../hooks/queries/useProductCountsQuery";
+import { updateProduct } from "../utils/updateProduct";
 
 export const UpdateProductButton = () => {
   const [buttonLoading, setButtonLoading] = useState(false);
   const axiosInstance = useAxiosInstance();
 
-  const { actionState, product } = useProductContext();
+  const { actionState, product, setShowNewVersionDialog } = useProductContext();
+  const { counts, isLoading } = useProductCountsQuery();
 
   const { refetch } = useProductQuery();
 
-  const updateProduct = async () => {
-    try {
-      await axiosInstance.post(`/v1/products/${product.id}`, {
-        ...UpdateProductSchema.parse(product),
-        items: product.items,
-        free_trial: product.free_trial,
-      });
-      // await ProductService.updateProduct(axiosInstance, product.id, {
-      //   ...UpdateProductSchema.parse(product),
-      //   items: product.items,
-      //   free_trial: product.free_trial,
-      // });
+  const handleUpdateClicked = async () => {
+    if (isLoading) toast.error("Product counts are loading");
 
-      toast.success("Product updated successfully");
-      // if (isNewProduct) {
-      //   toast.success("Product created successfully");
-      // } else {
-      // }
-
-      await refetch();
-      // await mutateCount();
-    } catch (error) {
-      toast.error(getBackendErr(error, "Failed to update product"));
+    if (counts?.all > 0) {
+      setShowNewVersionDialog(true);
+      return;
     }
+
+    setButtonLoading(true);
+    const success = await updateProduct({
+      axiosInstance,
+      product,
+      onSuccess: async () => {
+        await refetch();
+      },
+    });
+
+    setButtonLoading(false);
   };
 
   return (
     <Button
-      onClick={async () => {
-        setButtonLoading(true);
-        await updateProduct();
-        setButtonLoading(false);
-      }}
+      onClick={handleUpdateClicked}
       variant="gradientPrimary"
       className="w-full gap-2"
       isLoading={buttonLoading}

vite/src/views/products/product/free-trial/CreateFreeTrial.tsx

@@ -20,8 +20,8 @@ export const CreateFreeTrial = ({
   open: boolean;
   setOpen: (open: boolean) => void;
 }) => {
-  const [loading, setLoading] = useState(false);
-  const { product, setProduct, autoSave, mutate } = useProductContext();
+  const [loading] = useState(false);
+  const { product, setProduct, autoSave, refetch } = useProductContext();
 
   const axiosInstance = useAxiosInstance();
 
@@ -48,19 +48,22 @@ export const CreateFreeTrial = ({
         card_required: freeTrial.card_required,
       },
     });
-    // if (autoSave) {
-    //   handleAutoSave({
-    //     axiosInstance,
-    //     productId: product.id,
-    //     product: { ...product, free_trial: {
-    //       length: lengthInt,
-    //       unique_fingerprint: freeTrial.unique_fingerprint,
-    //       duration: freeTrial.duration,
-    //       card_required: freeTrial.card_required,
-    //     } },
-    //     mutate,
-    //   });
-    // }
+    if (autoSave) {
+      handleAutoSave({
+        axiosInstance,
+        productId: product.id,
+        product: {
+          ...product,
+          free_trial: {
+            length: lengthInt,
+            unique_fingerprint: freeTrial.unique_fingerprint,
+            duration: freeTrial.duration,
+            card_required: freeTrial.card_required,
+          },
+        },
+        refetch,
+      });
+    }
     setOpen(false);
   };
 

vite/src/views/products/product/hooks/queries/useMigrationsQuery.tsx.tsx

@@ -0,0 +1,18 @@
+import { useAxiosInstance } from "@/services/useAxiosInstance";
+import { useQuery } from "@tanstack/react-query";
+
+export const useMigrationsQuery = () => {
+  const axiosInstance = useAxiosInstance();
+
+  const fetchProductMigrations = async () => {
+    const { data } = await axiosInstance.get("/products/migrations");
+    return data;
+  };
+
+  const { data, isLoading, error, refetch } = useQuery({
+    queryKey: ["migrations"],
+    queryFn: fetchProductMigrations,
+  });
+
+  return { migrations: data?.migrations || [], isLoading, error, refetch };
+};

vite/src/views/products/product/hooks/queries/useProductCountsQuery.tsx

@@ -0,0 +1,29 @@
+import { useAxiosInstance } from "@/services/useAxiosInstance";
+import { useQuery } from "@tanstack/react-query";
+import { useParams } from "react-router-dom";
+import { useProductQueryState } from "../useProductQuery";
+
+export const useProductCountsQuery = () => {
+  const axiosInstance = useAxiosInstance();
+  const { product_id } = useParams();
+  const { queryStates } = useProductQueryState();
+
+  const productId = queryStates.productId || product_id;
+
+  const fetchProductCounts = async () => {
+    if (!productId) return null;
+    const { data } = await axiosInstance.get(`/products/${productId}/count`, {
+      params: {
+        version: queryStates.version,
+      },
+    });
+    return data;
+  };
+
+  const { data, isLoading, error, refetch } = useQuery({
+    queryKey: ["product_counts", productId, queryStates.version],
+    queryFn: fetchProductCounts,
+  });
+
+  return { counts: data, isLoading, error, refetch };
+};

vite/src/views/products/product/hooks/useProductQuery.tsx

@@ -5,15 +5,18 @@ import { useQueryStates } from "nuqs";
 import { useParams, useSearchParams } from "react-router";
 import { useCachedProduct } from "./getCachedProduct";
 import { useMemo } from "react";
+import { useProductCountsQuery } from "./queries/useProductCountsQuery";
+import { useMigrationsQuery } from "./queries/useMigrationsQuery.tsx";
 
 // Product query state...
 export const useProductQueryState = () => {
   const [queryStates, setQueryStates] = useQueryStates(
     {
       version: parseAsInteger,
+      productId: parseAsString,
     },
     {
-      history: "replace",
+      history: "push",
     }
   );
 
@@ -23,14 +26,16 @@ export const useProductQueryState = () => {
 export const useProductQuery = () => {
   const { product_id } = useParams();
   const { queryStates } = useProductQueryState();
+  const productId = queryStates.productId || product_id;
 
   const axiosInstance = useAxiosInstance();
-  const { getCachedProduct } = useCachedProduct({ productId: product_id });
+  const { getCachedProduct } = useCachedProduct({ productId: productId });
 
   const cachedProduct = useMemo(getCachedProduct, [getCachedProduct]);
 
   const fetcher = async () => {
-    const url = `/products/${product_id}/data2`;
+    if (!productId) return null;
+    const url = `/products/${productId}/data2`;
     const queryParams = {
       version: queryStates.version,
     };
@@ -40,15 +45,24 @@ export const useProductQuery = () => {
   };
 
   const { data, isLoading, refetch, error } = useQuery({
-    queryKey: ["product", product_id, queryStates.version],
+    queryKey: ["product", productId, queryStates.version],
     queryFn: fetcher,
   });
 
+  const { refetch: refetchCounts } = useProductCountsQuery();
+  const { refetch: refetchMigrations } = useMigrationsQuery();
+
   const product = data?.product || cachedProduct;
-  console.log("Cached product:", cachedProduct);
-  console.log("Error:", error);
-  console.log("Is loading:", isLoading);
   const isLoadingWithCache = cachedProduct ? false : isLoading;
 
-  return { product, isLoading: isLoadingWithCache, refetch, error };
+  return {
+    product,
+    numVersions: data?.numVersions || cachedProduct?.version || 1,
+    isLoading: isLoadingWithCache,
+    refetch: async () => {
+      await refetch();
+      await Promise.all([refetchMigrations(), refetchCounts()]);
+    },
+    error,
+  };
 };

vite/src/views/products/product/product-item/CreateProductItem2.tsx

@@ -14,7 +14,6 @@ import { useProductContext } from "../ProductContext";
 import { validateProductItem } from "@/utils/product/product-item/validateProductItem";
 import { PlusIcon } from "lucide-react";
 import { CreateItemDialogContent } from "./create-product-item/CreateItemDialogContent";
-import { useModelPricingContext } from "@/views/onboarding2/model-pricing/ModelPricingContext";
 import { useSteps } from "./useSteps";
 import { CreateItemStep } from "./utils/CreateItemStep";
 import { cn } from "@/lib/utils";

vite/src/views/products/product/product-item/product-item-config/item-config-footer/ItemConfigFooter.tsx

@@ -15,7 +15,14 @@ export const ItemConfigFooter = ({
   handleBack?: () => void;
 }) => {
   const axiosInstance = useAxiosInstance();
-  const { entityFeatureIds, product, mutate, autoSave } = useProductContext();
+  const {
+    entityFeatureIds,
+    product,
+    autoSave,
+    // This refetch is used in the onboarding flow
+    refetch,
+  } = useProductContext();
+
   const {
     item,
     handleCreateProductItem,
@@ -58,7 +65,7 @@ export const ItemConfigFooter = ({
                   axiosInstance,
                   productId: product.id,
                   product,
-                  mutate,
+                  refetch,
                 });
               }
             }}
@@ -77,7 +84,7 @@ export const ItemConfigFooter = ({
                   axiosInstance,
                   productId: product.id,
                   product,
-                  mutate,
+                  refetch,
                 });
               }
             }}
@@ -99,7 +106,7 @@ export const ItemConfigFooter = ({
                   axiosInstance,
                   productId: product.id,
                   product: newProduct,
-                  mutate,
+                  refetch,
                 });
               }
             }}
@@ -112,12 +119,3 @@ export const ItemConfigFooter = ({
     </CustomDialogFooter>
   );
 };
-
-// {handleCreateProductItem &&
-//   show.feature &&
-//   item.feature_id &&
-//   !entityFeatureIds.includes(item.feature_id) &&
-//   entityFeatureIds.length > 0 ? (
-
-//     </>
-//   )

vite/src/views/products/product/utils/updateProduct.ts

@@ -8,13 +8,15 @@ import { isFreeProduct } from "@/utils/product/priceUtils";
 export const updateProduct = async ({
   axiosInstance,
   product,
-  mutate,
-  mutateCount,
+  onSuccess,
+  // mutate,
+  // mutateCount,
 }: {
   axiosInstance: AxiosInstance;
   product: ProductV2;
-  mutate: () => void;
-  mutateCount: () => void;
+  onSuccess: () => Promise<void>;
+  // mutate: () => void;
+  // mutateCount: () => void;
 }) => {
   try {
     await ProductService.updateProduct(axiosInstance, product.id, {
@@ -25,9 +27,11 @@ export const updateProduct = async ({
 
     toast.success("Product updated successfully");
 
-    await mutate();
-    await mutateCount();
+    await onSuccess();
+    return true;
   } catch (error) {
+    console.error(error);
     toast.error(getBackendErr(error, "Failed to update product"));
+    return false;
   }
 };

vite/src/views/products/product/versioning/ConfirmNewVersionDialog.tsx

@@ -12,19 +12,23 @@ import { Input } from "@/components/ui/input";
 import { useProductContext } from "../ProductContext";
 import { toast } from "sonner";
 import { useState } from "react";
+import { updateProduct } from "../utils/updateProduct";
+import { useAxiosInstance } from "@/services/useAxiosInstance";
+import { useProductQuery } from "../hooks/useProductQuery";
 
 export default function ConfirmNewVersionDialog({
   open,
   setOpen,
-  createProduct,
 }: {
   open: boolean;
   setOpen: (open: boolean) => void;
-  createProduct: () => Promise<void>;
 }) {
-  const { product, version } = useProductContext();
-  let [confirmText, setConfirmText] = useState("");
-  let [isLoading, setIsLoading] = useState(false);
+  const axiosInstance = useAxiosInstance();
+  const { product } = useProductContext();
+  const { refetch } = useProductQuery();
+
+  const [confirmText, setConfirmText] = useState("");
+  const [isLoading, setIsLoading] = useState(false);
 
   const onClick = async () => {
     if (confirmText !== product.id) {
@@ -33,16 +37,21 @@ export default function ConfirmNewVersionDialog({
     }
 
     setIsLoading(true);
-    await createProduct();
+    await updateProduct({
+      axiosInstance,
+      product,
+      onSuccess: async () => {
+        await refetch();
+      },
+    });
     setIsLoading(false);
     setOpen(false);
+    // toast.success("New version created successfully");
   };
 
   return (
     <Dialog open={open} onOpenChange={setOpen}>
-      <DialogTrigger asChild>
-        {/* <Button>Confirm New Version</Button> */}
-      </DialogTrigger>
+      <DialogTrigger asChild></DialogTrigger>
       <DialogContent>
         <DialogHeader>
           <DialogTitle>Create new version?</DialogTitle>

vite/src/views/products/product/versioning/CountAndMigrate.tsx

@@ -1,6 +1,5 @@
 import SmallSpinner from "@/components/general/SmallSpinner";
 import { useAxiosInstance } from "@/services/useAxiosInstance";
-import { useEnv } from "@/utils/envUtils";
 import { getBackendErr } from "@/utils/genUtils";
 import { isOneOffProduct } from "@/utils/product/priceUtils";
 import { MigrationJob, MigrationJobStep } from "@autumn/shared";
@@ -15,17 +14,24 @@ import { Button } from "@/components/ui/button";
 import { toast } from "sonner";
 import { useProductContext } from "../ProductContext";
 import ConfirmMigrateDialog from "./ConfirmMigrateDialog";
+import { useProductQuery } from "../hooks/useProductQuery";
+import { useProductCountsQuery } from "../hooks/queries/useProductCountsQuery";
+import { useMigrationsQuery } from "../hooks/queries/useMigrationsQuery.tsx";
 
 export const CountAndMigrate = () => {
-  const {
-    product,
-    counts,
-    numVersions,
-    version,
-    existingMigrations,
-    mutate,
-    mutateCount,
-  } = useProductContext();
+  // const {
+  //   product,
+  //   counts,
+  //   numVersions,
+  //   version,
+  //   existingMigrations,
+  //   mutate,
+  //   mutateCount,
+  // } = useProductContext();
+
+  const { product, numVersions, refetch } = useProductQuery();
+  const { counts, refetch: refetchCounts } = useProductCountsQuery();
+  const { migrations, refetch: refetchMigrations } = useMigrationsQuery();
 
   const axiosInstance = useAxiosInstance();
   const [loading, setLoading] = useState(false);
@@ -36,11 +42,12 @@ export const CountAndMigrate = () => {
     try {
       const { data } = await axiosInstance.post("/v1/migrations", {
         from_product_id: product.id,
-        from_version: version,
+        from_version: product.version,
         to_product_id: product.id,
         to_version: numVersions,
       });
-      await mutate();
+
+      await refetchMigrations();
 
       toast.success(`Migration started. ID: ${data.id}`);
     } catch (error) {
@@ -54,22 +61,22 @@ export const CountAndMigrate = () => {
   };
 
   useEffect(() => {
-    if (existingMigrations.length > 0) {
+    if (migrations.length > 0) {
       // Run poll job on mutate
       const pollInterval = setInterval(() => {
-        mutate();
-        mutateCount();
+        refetchCounts();
+        refetchMigrations();
       }, 5000);
       return () => clearInterval(pollInterval);
     }
-  }, [existingMigrations]);
+  }, [migrations]);
 
   if (!counts) {
     return <></>;
   }
 
   const renderCurrentMigration = () => {
-    const migration: MigrationJob = existingMigrations[0];
+    const migration: MigrationJob = migrations[0];
 
     const getCusDetails = migration.step_details[MigrationJobStep.GetCustomers];
     const migrateDetails =
@@ -85,7 +92,7 @@ export const CountAndMigrate = () => {
             "
             >
               <span className="flex items-center gap-2 justify-between w-full pr-2">
-                <span className="text-sm font-medium text-t3 text-xs">
+                <span className="font-medium text-t3 text-xs">
                   Migration in progress
                 </span>
                 <SmallSpinner size={16} />
@@ -129,8 +136,8 @@ export const CountAndMigrate = () => {
   };
 
   const fromIsOneOff = isOneOffProduct(product.items);
-
   const migrateCount = counts?.active - counts?.canceled - counts?.custom;
+  const version = product.version;
 
   const canMigrate =
     counts &&
@@ -146,7 +153,7 @@ export const CountAndMigrate = () => {
         setOpen={setConfirmMigrateOpen}
         startMigration={migrateCustomers}
       />
-      {!canMigrate ? null : existingMigrations.length > 0 ? (
+      {!canMigrate ? null : migrations.length > 0 ? (
         renderCurrentMigration()
       ) : (
         <Button

vite/src/views/products/products/ProductsPage.tsx

@@ -1,4 +1,4 @@
-import CreateProduct from "./components/CreateProduct";
+import CreateProduct from "./components/CreateProductDialog";
 import { PageSectionHeader } from "@/components/general/PageSectionHeader";
 import { Badge } from "@/components/ui/badge";
 import { ProductsTable } from "./components/ProductsTable";

vite/src/views/products/products/components/CreateProductDialog.tsx

@@ -37,7 +37,7 @@ function CreateProduct({
   const [product, setProduct] = useState(defaultProduct);
   const [open, setOpen] = useState(false);
 
-  const { groupToDefaults, refetch } = useProductsQuery();
+  const { groupToDefaults } = useProductsQuery();
 
   const axiosInstance = useAxiosInstance();
   const navigate = useNavigate();
@@ -50,8 +50,6 @@ function CreateProduct({
         product
       );
 
-      await refetch();
-
       if (onSuccess) {
         await onSuccess(newProduct);
       } else {

vite/src/views/products/products/product-row-toolbar/DeleteProductDialog.tsx

@@ -20,28 +20,27 @@ import {
 } from "@/components/ui/select";
 import { getBackendErr } from "@/utils/genUtils";
 import { toast } from "sonner";
-import { useProductInfoQuery } from "./hooks/useProductInfoQuery";
 import { useProductsQuery } from "@/hooks/queries/useProductsQuery";
 import { useGeneralQuery } from "@/hooks/queries/useGeneralQuery";
+import { useModelPricingContext } from "@/views/onboarding2/model-pricing/ModelPricingContext";
 
 export const DeleteProductDialog = ({
   product,
   dropdownOpen,
   open,
   setOpen,
-  productCounts,
 }: {
   product: ProductV2;
   dropdownOpen: boolean;
   open: boolean;
   setOpen: (open: boolean) => void;
-  productCounts?: ProductCounts;
 }) => {
   const [deleteLoading, setDeleteLoading] = useState(false);
   const [archiveLoading, setArchiveLoading] = useState(false);
   const axiosInstance = useAxiosInstance();
 
-  const { refetch } = useProductsQuery();
+  const modelPricingContext = useModelPricingContext();
+  const { refetch: refetchProducts } = useProductsQuery();
 
   const { data: productInfo, isLoading } = useGeneralQuery({
     url: `/products/${product.id}/info`,
@@ -59,7 +58,12 @@ export const DeleteProductDialog = ({
         product.id,
         deleteAllVersions
       );
-      await refetch();
+
+      if (modelPricingContext) {
+        await modelPricingContext.refetch();
+      } else {
+        await refetchProducts();
+      }
       setOpen(false);
     } catch (error) {
       console.error("Error deleting product:", error);
@@ -112,7 +116,7 @@ export const DeleteProductDialog = ({
         }
         await Promise.all(updatePromises);
       }
-      await refetch();
+      await refetchProducts();
       toast.success(
         `Product ${product.name} ${newArchivedState ? "archived" : "unarchived"} successfully`
       );