feat(customer-portal): let customers update subscription units

clients/apps/web/src/components/CustomerPortal/CustomerPortalSubscription.tsx

@@ -30,6 +30,7 @@ import CustomerCancellationModal from './CustomerCancellationModal'
 import CustomerPauseSubscriptionModal from './CustomerPauseSubscriptionModal'
 import { SubscriptionStatusLabel } from '../Subscriptions/utils'
 import { CustomerPortalGrants } from './CustomerPortalGrants'
+import { CustomerUnitQuantityManager } from './CustomerUnitQuantityManager'
 import { SeatManagementTable } from './SeatManagementTable'
 
 const CustomerPortalSubscription = ({
@@ -93,6 +94,11 @@ const CustomerPortalSubscription = ({
     (price) => price.amount_type === 'seat_based',
   )
 
+  const unitPrice = subscription.prices.find(
+    (price): price is schemas['ProductPriceUnitBased'] =>
+      price.amount_type === 'unit_based',
+  )
+
   // Check customer portal settings for seat management visibility
   const portalSettings =
     subscription.product.organization.customer_portal_settings
@@ -269,6 +275,12 @@ const CustomerPortalSubscription = ({
                 value={`${subscription.seats} -> ${pendingUpdate.seats}`}
               />
             )}
+            {pendingUpdate.units !== null && (
+              <DetailRow
+                label="Units"
+                value={`${subscription.units} -> ${pendingUpdate.units}`}
+              />
+            )}
             <DetailRow
               label="Update in effect from"
               value={
@@ -342,6 +354,21 @@ const CustomerPortalSubscription = ({
         />
       )}
 
+      {unitPrice && canManageBilling && !isCancelled && (
+        <div className="flex flex-col gap-y-2">
+          <h3 className="text-lg">Units</h3>
+          <CustomerUnitQuantityManager
+            api={api}
+            subscription={subscription}
+            unitPrice={unitPrice}
+            prorationBehavior={
+              subscription.product.organization.proration_behavior
+            }
+            onUpdate={() => router.refresh()}
+          />
+        </div>
+      )}
+
       <CustomerPortalGrants api={api} subscriptionId={subscription.id} />
 
       <div className="flex w-full flex-col gap-4">

clients/apps/web/src/components/CustomerPortal/CustomerUnitQuantityManager.tsx

@@ -0,0 +1,231 @@
+'use client'
+
+import { useCustomerUpdateSubscription } from '@/hooks/queries/customerPortal'
+import { setValidationErrors } from '@/utils/api/errors'
+import { getUnitLabels } from '@/utils/product'
+import { Client, isValidationError, schemas } from '@polar-sh/client'
+import { Button, Input } from '@polar-sh/orbit'
+import { MinusIcon, PlusIcon } from 'lucide-react'
+import { useCallback, useMemo, useState } from 'react'
+import { useForm } from 'react-hook-form'
+import { toast } from '../Toast/use-toast'
+
+interface CustomerUnitQuantityManagerProps {
+  api: Client
+  subscription: schemas['CustomerSubscription']
+  unitPrice: schemas['ProductPriceUnitBased']
+  prorationBehavior?: schemas['CustomerOrganization']['proration_behavior']
+  onUpdate?: () => void
+}
+
+const getUnitBounds = (
+  unitPrice: schemas['ProductPriceUnitBased'],
+): { min: number; max: number | null } => {
+  const min = Math.max(1, unitPrice.minimum_units ?? 0)
+  if (unitPrice.maximum_units != null) {
+    return { min, max: unitPrice.maximum_units }
+  }
+  const lastBound = unitPrice.tiers.tiers.at(-1)?.bound ?? null
+  return { min, max: lastBound }
+}
+
+export const CustomerUnitQuantityManager = ({
+  api,
+  subscription,
+  unitPrice,
+  prorationBehavior,
+  onUpdate,
+}: CustomerUnitQuantityManagerProps) => {
+  const updateSubscription = useCustomerUpdateSubscription(api)
+
+  const currentUnits = subscription.units ?? 1
+  const { min, max } = useMemo(() => getUnitBounds(unitPrice), [unitPrice])
+  const { unitLabel, unitLabelPlural } = getUnitLabels(unitPrice)
+
+  const { handleSubmit, watch, setValue, setError } = useForm<{
+    units: number
+  }>({
+    values: {
+      units: currentUnits,
+    },
+  })
+
+  // eslint-disable-next-line react-hooks/incompatible-library
+  const units = watch('units')
+  const canDecrease = units > min
+  const canIncrease = max == null || units < max
+  const hasChanges = units !== currentUnits
+
+  const [draft, setDraft] = useState(String(currentUnits))
+
+  const setUnits = (value: number) => {
+    setValue('units', value, { shouldValidate: true })
+    setDraft(String(value))
+  }
+
+  const invoicingMessage = useMemo(() => {
+    if (!prorationBehavior) return null
+    switch (prorationBehavior) {
+      case 'invoice':
+        return "I'll be charged immediately with a proration for the current month."
+      case 'prorate':
+        return 'Your next invoice will include the updated units plus the proration for the current month.'
+      case 'next_period':
+        return 'The unit update will be applied on your next billing cycle.'
+    }
+  }, [prorationBehavior])
+
+  const onSubmit = useCallback(
+    async (data: { units: number }) => {
+      try {
+        const result = await updateSubscription.mutateAsync({
+          id: subscription.id,
+          body: {
+            units: data.units,
+          },
+        })
+
+        if (result.error) {
+          const errorMessage =
+            typeof result.error.detail === 'string'
+              ? result.error.detail
+              : 'Failed to update units'
+          toast({
+            title: 'Error updating units',
+            description: errorMessage,
+            variant: 'error',
+          })
+        } else {
+          const noun = data.units === 1 ? unitLabel : unitLabelPlural
+          const description = `Subscription now has ${data.units} ${noun}.`
+          toast({
+            title: 'Units updated',
+            description,
+          })
+          onUpdate?.()
+        }
+      } catch (error) {
+        if (isValidationError(error)) {
+          setValidationErrors(error, setError)
+        } else {
+          toast({
+            title: 'Error updating units',
+            description:
+              error instanceof Error
+                ? error.message
+                : 'An unexpected error occurred',
+            variant: 'error',
+          })
+        }
+      }
+    },
+    [
+      updateSubscription,
+      subscription.id,
+      unitLabel,
+      unitLabelPlural,
+      onUpdate,
+      setError,
+    ],
+  )
+
+  const handleIncrement = () => {
+    if (canIncrease) {
+      setUnits(units + 1)
+    }
+  }
+
+  const handleDecrement = () => {
+    if (canDecrease) {
+      setUnits(units - 1)
+    }
+  }
+
+  const clampUnits = (value: number): number => {
+    const floored = Math.max(min, Math.floor(value))
+    return max != null ? Math.min(max, floored) : floored
+  }
+
+  const handleInputChange = (event: React.ChangeEvent<HTMLInputElement>) => {
+    const raw = event.target.value
+    if (!/^\d*$/.test(raw)) {
+      return
+    }
+    setDraft(raw)
+    if (raw !== '') {
+      setValue('units', Number.parseInt(raw, 10), { shouldValidate: true })
+    }
+  }
+
+  const handleInputBlur = () => {
+    const parsed = draft === '' ? min : Number.parseInt(draft, 10)
+    setUnits(clampUnits(parsed))
+  }
+
+  return (
+    <form onSubmit={handleSubmit(onSubmit)} className="rounded-2xl border p-4">
+      <div className="flex flex-col gap-2 text-sm">
+        <div className="flex items-center justify-between">
+          <div>
+            <span className="font-medium">
+              {units === 1
+                ? unitLabel.charAt(0).toUpperCase() + unitLabel.slice(1)
+                : unitLabelPlural.charAt(0).toUpperCase() +
+                  unitLabelPlural.slice(1)}
+            </span>
+          </div>
+          <div className="flex items-center gap-1">
+            <Button
+              type="button"
+              variant="secondary"
+              size="icon"
+              onClick={handleDecrement}
+              disabled={!canDecrease || updateSubscription.isPending}
+            >
+              <MinusIcon className="h-4 w-4" />
+            </Button>
+
+            <Input
+              type="text"
+              inputMode="numeric"
+              value={draft}
+              onChange={handleInputChange}
+              onBlur={handleInputBlur}
+              disabled={updateSubscription.isPending}
+              className="w-20 text-center font-medium"
+            />
+
+            <Button
+              type="button"
+              variant="secondary"
+              size="icon"
+              onClick={handleIncrement}
+              disabled={!canIncrease || updateSubscription.isPending}
+            >
+              <PlusIcon className="h-4 w-4" />
+            </Button>
+          </div>
+        </div>
+      </div>
+
+      {hasChanges && (
+        <div className="mt-4 flex flex-col gap-3">
+          {invoicingMessage && (
+            <span className="dark:text-polar-500 text-sm text-gray-500">
+              {invoicingMessage}
+            </span>
+          )}
+          <div className="flex flex-row-reverse gap-3">
+            <Button
+              loading={updateSubscription.isPending}
+              onClick={handleSubmit(onSubmit)}
+              className="w-full"
+            >
+              Update units
+            </Button>
+          </div>
+        </div>
+      )}
+    </form>
+  )
+}

server/polar/customer_portal/endpoints/subscription.py

@@ -36,6 +36,7 @@
     RevokeNotAllowed,
     UpdateSubscriptionPlanNotAllowed,
     UpdateSubscriptionSeatsNotAllowed,
+    UpdateSubscriptionUnitsNotAllowed,
 )
 from ..service.subscription import (
     customer_subscription as customer_subscription_service,
@@ -185,7 +186,8 @@ async def get_cancel_preview(
             "description": "Previewing this change is not allowed.",
             "model": AlreadyCanceledSubscription.schema()
             | UpdateSubscriptionPlanNotAllowed.schema()
-            | UpdateSubscriptionSeatsNotAllowed.schema(),
+            | UpdateSubscriptionSeatsNotAllowed.schema()
+            | UpdateSubscriptionUnitsNotAllowed.schema(),
         },
         404: SubscriptionNotFound,
     },
@@ -228,7 +230,9 @@ async def preview_change(
                 "or pausing/resuming is not enabled for the organization."
             ),
             "model": AlreadyCanceledSubscription.schema()
-            | PauseResumeNotAllowed.schema(),
+            | PauseResumeNotAllowed.schema()
+            | UpdateSubscriptionSeatsNotAllowed.schema()
+            | UpdateSubscriptionUnitsNotAllowed.schema(),
         },
         404: SubscriptionNotFound,
         409: {