fix(web): roll back rejected onboarding updates

.github/workflows/ci.yml

@@ -29,5 +29,8 @@ jobs:
       - name: Run TypeScript type checking
         run: bunx turbo run check-types --filter='@supermemory/ai-sdk' --filter='@supermemory/memory-graph'
 
+      - name: Run shared library unit tests
+        run: bun test packages/hooks packages/lib
+
       - name: Run Biome CI (format & lint on changed files)
         run: bunx biome ci --changed --since=origin/main --no-errors-on-unmatched

packages/hooks/use-org-onboarding.test.ts

@@ -0,0 +1,145 @@
+import {
+	afterEach,
+	beforeAll,
+	beforeEach,
+	describe,
+	expect,
+	it,
+	mock,
+	spyOn,
+} from "bun:test"
+import { createElement } from "react"
+import { renderToStaticMarkup } from "react-dom/server"
+import { mergeOrganizationMetadata } from "../lib/organization-metadata"
+
+interface TestOrganization {
+	id: string
+	metadata: Record<string, unknown>
+}
+
+type UpdateResult = {
+	error: { message: string } | null
+}
+
+let activeOrganization: TestOrganization | null = null
+let persistOrganization: () => Promise<UpdateResult>
+
+const updateOrgMetadata = mock(
+	(organizationId: string, partial: Record<string, unknown>) => {
+		activeOrganization = mergeOrganizationMetadata(
+			activeOrganization,
+			organizationId,
+			partial,
+		)
+	},
+)
+
+mock.module("@lib/auth-context", () => ({
+	useAuth: () => ({
+		org: activeOrganization,
+		updateOrgMetadata,
+	}),
+}))
+
+mock.module("@lib/auth", () => ({
+	authClient: {
+		organization: {
+			update: () => persistOrganization(),
+		},
+	},
+}))
+
+let useOrgOnboarding: typeof import("./use-org-onboarding").useOrgOnboarding
+
+beforeAll(async () => {
+	;({ useOrgOnboarding } = await import("./use-org-onboarding"))
+})
+
+beforeEach(() => {
+	activeOrganization = {
+		id: "org-a",
+		metadata: { isOnboarded: false },
+	}
+	persistOrganization = async () => ({ error: null })
+	updateOrgMetadata.mockClear()
+	spyOn(console, "error").mockImplementation(() => {})
+})
+
+afterEach(() => {
+	mock.restore()
+})
+
+function renderOnboardingHook() {
+	let hook: ReturnType<typeof useOrgOnboarding> | undefined
+
+	renderToStaticMarkup(
+		createElement(() => {
+			hook = useOrgOnboarding()
+			return null
+		}),
+	)
+
+	if (!hook) throw new Error("Onboarding hook did not render")
+	return hook
+}
+
+async function flushUpdate() {
+	await Promise.resolve()
+	await Promise.resolve()
+	await Promise.resolve()
+}
+
+describe("useOrgOnboarding", () => {
+	it("rolls back a rejected mark update in the same organization", async () => {
+		persistOrganization = async () => ({
+			error: { message: "Update rejected" },
+		})
+
+		renderOnboardingHook().markOrgOnboarded()
+		expect(activeOrganization?.metadata.isOnboarded).toBe(true)
+
+		await flushUpdate()
+
+		expect(activeOrganization?.metadata.isOnboarded).toBe(false)
+	})
+
+	it("rolls back a rejected reset update in the same organization", async () => {
+		activeOrganization = {
+			id: "org-a",
+			metadata: { isOnboarded: true },
+		}
+		persistOrganization = async () => ({
+			error: { message: "Update rejected" },
+		})
+
+		renderOnboardingHook().resetOrgOnboarded()
+		expect(activeOrganization?.metadata.isOnboarded).toBe(false)
+
+		await flushUpdate()
+
+		expect(activeOrganization?.metadata.isOnboarded).toBe(true)
+	})
+
+	it("does not roll back a different organization", async () => {
+		let resolveUpdate: (result: UpdateResult) => void = () => {}
+		persistOrganization = () =>
+			new Promise((resolve) => {
+				resolveUpdate = resolve
+			})
+
+		renderOnboardingHook().markOrgOnboarded()
+		expect(activeOrganization?.metadata.isOnboarded).toBe(true)
+
+		activeOrganization = {
+			id: "org-b",
+			metadata: { isOnboarded: true },
+		}
+		resolveUpdate({ error: { message: "Update rejected" } })
+		await flushUpdate()
+
+		expect(activeOrganization).toEqual({
+			id: "org-b",
+			metadata: { isOnboarded: true },
+		})
+	})
+})

packages/hooks/use-org-onboarding.ts

@@ -27,7 +27,7 @@ export function useOrgOnboarding() {
 		}
 
 		// Optimistic update: update in-memory state immediately
-		updateOrgMetadata({ isOnboarded: true })
+		updateOrgMetadata(org.id, { isOnboarded: true })
 
 		authClient.organization
 			.update({
@@ -39,9 +39,16 @@ export function useOrgOnboarding() {
 					},
 				},
 			})
+			.then((result) => {
+				if (result.error) {
+					throw new Error(
+						result.error.message ?? "Failed to mark organization as onboarded",
+					)
+				}
+			})
 			.catch((error) => {
 				console.error("Failed to mark organization as onboarded:", error)
-				updateOrgMetadata({ isOnboarded: false })
+				updateOrgMetadata(org.id, { isOnboarded: false })
 			})
 	}, [org, updateOrgMetadata])
 
@@ -52,7 +59,7 @@ export function useOrgOnboarding() {
 		}
 
 		// Optimistic update: update in-memory state immediately
-		updateOrgMetadata({ isOnboarded: false })
+		updateOrgMetadata(org.id, { isOnboarded: false })
 
 		authClient.organization
 			.update({
@@ -64,9 +71,16 @@ export function useOrgOnboarding() {
 					},
 				},
 			})
+			.then((result) => {
+				if (result.error) {
+					throw new Error(
+						result.error.message ?? "Failed to reset organization onboarding",
+					)
+				}
+			})
 			.catch((error) => {
 				console.error("Failed to reset organization onboarding:", error)
-				updateOrgMetadata({ isOnboarded: true })
+				updateOrgMetadata(org.id, { isOnboarded: true })
 			})
 	}, [org, updateOrgMetadata])
 

packages/lib/auth-context.tsx

@@ -9,6 +9,7 @@ import {
 	useState,
 } from "react"
 import { authClient, useSession } from "./auth"
+import { mergeOrganizationMetadata } from "./organization-metadata"
 
 type Organization = typeof authClient.$Infer.ActiveOrganization
 type SessionData = NonNullable<ReturnType<typeof useSession>["data"]>
@@ -44,7 +45,10 @@ interface AuthContextType {
 	isSessionPending: boolean
 	setActiveOrg: (orgSlug: string) => Promise<void>
 	clearActiveOrg: () => Promise<void>
-	updateOrgMetadata: (partial: Record<string, unknown>) => void
+	updateOrgMetadata: (
+		organizationId: string,
+		partial: Record<string, unknown>,
+	) => void
 	refetchActiveOrg: () => Promise<Organization | null>
 	refetchOrganizations: () => Promise<unknown>
 }
@@ -89,18 +93,12 @@ export function AuthProvider({ children }: { children: ReactNode }) {
 		} catch {}
 	}, [])
 
-	const updateOrgMetadata = useCallback((partial: Record<string, unknown>) => {
-		setOrg((prev) => {
-			if (!prev) return prev
-			return {
-				...prev,
-				metadata: {
-					...prev.metadata,
-					...partial,
-				},
-			}
-		})
-	}, [])
+	const updateOrgMetadata = useCallback(
+		(organizationId: string, partial: Record<string, unknown>) => {
+			setOrg((prev) => mergeOrganizationMetadata(prev, organizationId, partial))
+		},
+		[],
+	)
 
 	const refetchActiveOrg = useCallback(async () => {
 		const full = await authClient.organization.getFullOrganization()

packages/lib/organization-metadata.test.ts

@@ -0,0 +1,41 @@
+import { describe, expect, it } from "bun:test"
+import { mergeOrganizationMetadata } from "./organization-metadata"
+
+describe("mergeOrganizationMetadata", () => {
+	it("merges metadata for the expected organization", () => {
+		expect(
+			mergeOrganizationMetadata(
+				{
+					id: "org-a",
+					name: "Organization A",
+					metadata: { plan: "pro", isOnboarded: false },
+				},
+				"org-a",
+				{ isOnboarded: true },
+			),
+		).toEqual({
+			id: "org-a",
+			name: "Organization A",
+			metadata: { plan: "pro", isOnboarded: true },
+		})
+	})
+
+	it("leaves a different current organization untouched", () => {
+		const current = {
+			id: "org-b",
+			metadata: { isOnboarded: true },
+		}
+
+		expect(
+			mergeOrganizationMetadata(current, "org-a", {
+				isOnboarded: false,
+			}),
+		).toBe(current)
+	})
+
+	it("leaves an empty organization state untouched", () => {
+		expect(
+			mergeOrganizationMetadata(null, "org-a", { isOnboarded: true }),
+		).toBeNull()
+	})
+})

packages/lib/organization-metadata.ts

@@ -0,0 +1,22 @@
+interface OrganizationWithMetadata {
+	id: string
+	metadata?: unknown
+}
+
+export function mergeOrganizationMetadata<
+	Organization extends OrganizationWithMetadata,
+>(
+	current: Organization | null,
+	organizationId: string,
+	partial: Record<string, unknown>,
+): Organization | null {
+	if (!current || current.id !== organizationId) return current
+
+	return {
+		...current,
+		metadata: {
+			...(current.metadata as Record<string, unknown> | null),
+			...partial,
+		},
+	} as Organization
+}