"use client"; import { useAtom } from "jotai"; import { atomWithStorage } from "jotai/utils"; import { usePathname } from "next/navigation"; import React, { createContext, useCallback, useContext, useEffect, useMemo, useRef, useState } from "react"; import { useSocketContext } from "../../../contexts/SocketContext"; import { Modules, rehydrate } from "../../../core"; import { Action, checkPermissions, ModuleWithPermissions } from "../../../permissions"; import { getRoleId } from "../../../roles"; import { getTokenHandler } from "../../auth/config"; import { CompanyInterface } from "../../company/data/company.interface"; import { FeatureInterface } from "../../feature"; import { RoleInterface } from "../../role"; import { UserInterface, UserService } from "../data"; import { CURRENT_USER_STORAGE_KEY } from "../data/user.storage"; export interface CurrentUserContextType { currentUser: T | null; company: CompanyInterface | null; setUser: (user?: T) => void; hasPermissionToModule: (params: { module: M; action: Action; data?: any; }) => boolean; hasPermissionToModules: (params: { modules: M[]; action: Action; data?: any; }) => boolean; hasPermissionToPath: (path: string) => boolean; hasAccesToFeature: (featureIdentifier: string) => boolean; matchUrlToModule: (prarms?: { path: string }) => ModuleWithPermissions | undefined; hasRole: (roleId: string) => boolean; refreshUser: (options?: { skipCookieUpdate?: boolean }) => Promise; isRefreshing: boolean; } const userAtom = atomWithStorage(CURRENT_USER_STORAGE_KEY, null); const CurrentUserContext = createContext(undefined); export const CurrentUserProvider = ({ children }: { children: React.ReactNode }) => { const path = usePathname(); const [dehydratedUser, setDehydratedUser] = useAtom(userAtom); const matchUrlToModule = (_params?: { path: string }): ModuleWithPermissions | undefined => { const moduleKeys = Object.getOwnPropertyNames(Modules).filter( (key) => key !== "prototype" && key !== "_factory" && key !== "length" && key !== "name", ); const matchedModuleKey = moduleKeys.find((key) => { const descriptor = Object.getOwnPropertyDescriptor(Modules, key); if (!descriptor?.get) return false; const selectedModule = descriptor.get.call(Modules); return path.toLowerCase().startsWith(selectedModule.pageUrl?.toLowerCase()); }); if (!matchedModuleKey) return undefined; const descriptor = Object.getOwnPropertyDescriptor(Modules, matchedModuleKey); return descriptor?.get?.call(Modules); }; const currentUser = useMemo( () => (dehydratedUser ? rehydrate(Modules.User, dehydratedUser) : null), [dehydratedUser], ); // Balances pushed by company:credits_updated. Transient: a real user refresh is // authoritative and clears this. const [creditOverride, setCreditOverride] = useState<{ availableMonthlyCredits: number; availableExtraCredits: number; } | null>(null); const company = useMemo(() => { const c = currentUser?.company ?? null; if (!c || !creditOverride) return c; return Object.assign(Object.create(Object.getPrototypeOf(c)), c, { _availableMonthlyCredits: creditOverride.availableMonthlyCredits, _availableExtraCredits: creditOverride.availableExtraCredits, }); }, [currentUser, creditOverride]); const setUser = (user?: UserInterface): void => { if (user) setDehydratedUser(user.dehydrate() as any); else setDehydratedUser(null); setCreditOverride(null); }; const hasRole = (roleId: string): boolean => { if (!currentUser) return false; return !!currentUser.roles?.some((userRole: RoleInterface) => userRole.id === roleId); }; const hasAccesToFeature = (featureIdentifier: string): boolean => { if (hasRole(getRoleId().Administrator)) return true; if (!company) return false; return company.features.some((feature: FeatureInterface) => feature.id === featureIdentifier); }; function hasPermissionToModule(params: { module: M; action: Action; data?: M extends ModuleWithPermissions ? any : never; }): boolean { if (!currentUser) return false; if (!!params.module.feature && !hasAccesToFeature(params.module.feature)) return false; return checkPermissions({ module: params.module, action: params.action, data: params.data, user: currentUser }); } function hasPermissionToModules(params: { modules: M[]; action: Action; data?: M extends ModuleWithPermissions ? any : never; }): boolean { if (!currentUser) return false; if (!params.modules.every((module) => !module.feature || hasAccesToFeature(module.feature))) return false; return params.modules.every((module) => checkPermissions({ module: module, action: params.action, data: params.data, user: currentUser }), ); } function hasPermissionToPath(path: string): boolean { if (!currentUser) return false; if (path === "#" || path === "/") return true; const selectedModule = matchUrlToModule({ path: path }); if (!selectedModule) return true; const response = hasPermissionToModule({ module: selectedModule, action: Action.Read }); return response; } // State for tracking refresh status const [isRefreshing, setIsRefreshing] = useState(false); // Function to refresh user data from the API // skipCookieUpdate: When true, only updates React state without calling the Server Action // This prevents page reloads when refresh is triggered by WebSocket events const refreshUser = useCallback( async (options?: { skipCookieUpdate?: boolean }): Promise => { if (isRefreshing) { return; } setIsRefreshing(true); try { const fullUser = await UserService.findFullUser(); if (fullUser) { // Update authentication cookies with fresh user data BEFORE writing the atom, // so downstream observers see cookies-then-user, never user-then-cookies. // Skip when triggered by WebSocket to prevent page reload (Server Actions modify cookies). if (!options?.skipCookieUpdate) { await getTokenHandler()?.updateToken({ userId: fullUser.id, companyId: fullUser.company?.id, roles: fullUser.roles.map((role) => role.id), features: fullUser.company?.features?.map((feature) => feature.id) ?? [], modules: fullUser.modules.map((module) => ({ id: module.id, permissions: module.permissions, })), }); } setDehydratedUser(fullUser.dehydrate() as any); setCreditOverride(null); } } catch (error) { console.error("Failed to refresh user data:", error); } finally { setIsRefreshing(false); } }, [isRefreshing, setDehydratedUser], ); // WebSocket integration for real-time token updates const { socket, isConnected } = useSocketContext(); // Use ref for stable refreshUser reference to avoid effect re-runs const refreshUserRef = useRef(refreshUser); refreshUserRef.current = refreshUser; // Track refresh in progress to prevent duplicate API calls const isRefreshingRef = useRef(false); // Listen for company WebSocket events that trigger user refresh useEffect(() => { if (!socket || !isConnected || !currentUser?.company?.id) { return; } // Hot path: fires once per LLM call. Fully described by its payload — must // never trigger an HTTP request. const handleCreditsUpdated = (data: { companyId: string; availableMonthlyCredits: number; availableExtraCredits: number; }) => { if (data.companyId !== currentUser.company?.id) return; setCreditOverride({ availableMonthlyCredits: data.availableMonthlyCredits, availableExtraCredits: data.availableExtraCredits, }); }; // Rare (Stripe webhooks only) and genuinely changes features, modules and // permissions — a payload cannot economically carry that, so refetch. const handleSubscriptionUpdated = (data: { companyId: string; type: string }) => { if (data.companyId === currentUser.company?.id && !isRefreshingRef.current) { isRefreshingRef.current = true; // Skip cookie update to prevent page reload - only update React state refreshUserRef.current({ skipCookieUpdate: true }).finally(() => { isRefreshingRef.current = false; }); } }; socket.on("company:credits_updated", handleCreditsUpdated); socket.on("company:subscription_updated", handleSubscriptionUpdated); return () => { socket.off("company:credits_updated", handleCreditsUpdated); socket.off("company:subscription_updated", handleSubscriptionUpdated); }; }, [socket, isConnected, currentUser?.company?.id]); return ( {children} ); }; export function useCurrentUserContext(): CurrentUserContextType { const context = useContext(CurrentUserContext); if (context === undefined) { throw new Error("useCurrentUserContext must be used within a UserProvider"); } return context as unknown as CurrentUserContextType; }