[dev] [Marfuen] mariano/fix-vendor-risks-links (#1832)

* refactor(risk): update getRisks and getAssignees functions to accept orgId

* chore(policy-editor): gate policy ai assistant behind feature flag

---------

Co-authored-by: Mariano Fuentes <[email protected]>

Author: github-actions[bot]

apps/app/src/app/(app)/[orgId]/policies/[policyId]/editor/components/PolicyDetails.tsx

@@ -13,6 +13,7 @@ import type { JSONContent } from '@tiptap/react';
 import { structuredPatch } from 'diff';
 import { CheckCircle, Loader2, Sparkles, X } from 'lucide-react';
 import { useAction } from 'next-safe-action/hooks';
+import { useFeatureFlagEnabled } from 'posthog-js/react';
 import { useMemo, useState } from 'react';
 import { toast } from 'sonner';
 import { switchPolicyDisplayFormatAction } from '../../actions/switch-policy-display-format';
@@ -47,6 +48,7 @@ export function PolicyContentManager({
 
   const [proposedPolicyMarkdown, setProposedPolicyMarkdown] = useState<string | null>(null);
   const [isApplying, setIsApplying] = useState(false);
+  const isAiPolicyAssistantEnabled = useFeatureFlagEnabled('is-ai-policy-assistant-enabled');
 
   const switchFormat = useAction(switchPolicyDisplayFormatAction, {
     onError: () => toast.error('Failed to switch view.'),
@@ -103,17 +105,19 @@ export function PolicyContentManager({
                       PDF View
                     </TabsTrigger>
                   </TabsList>
-                  {!isPendingApproval && displayFormat === 'EDITOR' && (
-                    <Button
-                      variant={showAiAssistant ? 'default' : 'outline'}
-                      size="sm"
-                      onClick={() => setShowAiAssistant((prev) => !prev)}
-                      className="gap-2"
-                    >
-                      <Sparkles className="h-4 w-4" />
-                      AI Assistant
-                    </Button>
-                  )}
+                  {!isPendingApproval &&
+                    displayFormat === 'EDITOR' &&
+                    isAiPolicyAssistantEnabled && (
+                      <Button
+                        variant={showAiAssistant ? 'default' : 'outline'}
+                        size="sm"
+                        onClick={() => setShowAiAssistant((prev) => !prev)}
+                        className="gap-2"
+                      >
+                        <Sparkles className="h-4 w-4" />
+                        AI Assistant
+                      </Button>
+                    )}
                 </div>
                 <TabsContent value="EDITOR" className="mt-4">
                   <PolicyEditorWrapper
@@ -134,7 +138,7 @@ export function PolicyContentManager({
               </Tabs>
             </div>
 
-            {showAiAssistant && (
+            {showAiAssistant && isAiPolicyAssistantEnabled && (
               <div className="w-80 shrink-0 min-h-[400px] self-stretch flex flex-col">
                 <PolicyAiAssistant
                   policyId={policyId}

apps/app/src/app/(app)/[orgId]/risk/(overview)/RisksTable.tsx

@@ -5,7 +5,6 @@ import { DataTableToolbar } from '@/components/data-table/data-table-toolbar';
 import { CreateRiskSheet } from '@/components/sheets/create-risk-sheet';
 import { useDataTable } from '@/hooks/use-data-table';
 import { getFiltersStateParser, getSortingStateParser } from '@/lib/parsers';
-import { useSession } from '@/utils/auth-client';
 import type { Member, Risk, User } from '@db';
 import { Risk as RiskType } from '@db';
 import { ColumnDef } from '@tanstack/react-table';
@@ -42,15 +41,15 @@ export const RisksTable = ({
   pageCount: initialPageCount,
   onboardingRunId,
   searchParams: initialSearchParams,
+  orgId,
 }: {
   risks: RiskRow[];
   assignees: (Member & { user: User })[];
   pageCount: number;
   onboardingRunId?: string | null;
   searchParams: GetRiskSchema;
+  orgId: string;
 }) => {
-  const session = useSession();
-  const orgId = session?.data?.session?.activeOrganizationId;
   const [_, setOpenSheet] = useQueryState('create-risk-sheet');
 
   const { itemStatuses, progress, itemsInfo, isActive, isLoading } = useOnboardingStatus(
@@ -99,8 +98,9 @@ export const RisksTable = ({
 
   // Fetcher function for SWR
   const fetcher = useCallback(async () => {
-    return await getRisksAction(currentSearchParams);
-  }, [currentSearchParams]);
+    if (!orgId) return { data: [], pageCount: 0 };
+    return await getRisksAction({ orgId, searchParams: currentSearchParams });
+  }, [orgId, currentSearchParams]);
 
   // Use SWR to fetch risks with polling when onboarding is active
   const { data: risksData } = useSWR(swrKey, fetcher, {
@@ -202,7 +202,7 @@ export const RisksTable = ({
           residualImpact: 'insignificant' as const,
           treatmentStrategy: 'accept' as const,
           treatmentStrategyDescription: null,
-          organizationId: orgId || '',
+          organizationId: orgId,
           assigneeId: null,
           assignee: null,
           createdAt: new Date(),
@@ -229,7 +229,7 @@ export const RisksTable = ({
           residualImpact: 'insignificant' as const,
           treatmentStrategy: 'accept' as const,
           treatmentStrategyDescription: null,
-          organizationId: orgId || '',
+          organizationId: orgId,
           assigneeId: null,
           assignee: null,
           createdAt: new Date(),
@@ -241,7 +241,7 @@ export const RisksTable = ({
     return [...risksWithStatus, ...pendingRisks, ...tempRisks];
   }, [risks, itemsInfo, itemStatuses, orgId, isActive, onboardingRunId]);
 
-  const columns = useMemo<ColumnDef<RiskRow>[]>(() => getColumns(orgId ?? ''), [orgId]);
+  const columns = useMemo<ColumnDef<RiskRow>[]>(() => getColumns(orgId), [orgId]);
 
   const { table } = useDataTable({
     data: mergedRisks,

apps/app/src/app/(app)/[orgId]/risk/(overview)/actions/get-risks-action.ts

@@ -3,6 +3,11 @@
 import { getRisks } from '../data/getRisks';
 import type { GetRiskSchema } from '../data/validations';
 
-export async function getRisksAction(input: GetRiskSchema) {
-  return await getRisks(input);
+type GetRisksActionInput = {
+  orgId: string;
+  searchParams: GetRiskSchema;
+};
+
+export async function getRisksAction({ orgId, searchParams }: GetRisksActionInput) {
+  return await getRisks({ orgId, searchParams });
 }

apps/app/src/app/(app)/[orgId]/risk/(overview)/data/getRisks.ts

@@ -1,11 +1,14 @@
 import 'server-only';
 
-import { auth } from '@/utils/auth';
 import { db, Prisma, type User } from '@db';
-import { headers } from 'next/headers';
 import type { GetRiskSchema } from './validations';
 
-export async function getRisks(input: GetRiskSchema): Promise<{
+export type GetRisksInput = {
+  orgId: string;
+  searchParams: GetRiskSchema;
+};
+
+export async function getRisks({ orgId, searchParams }: GetRisksInput): Promise<{
   data: (Omit<
     Prisma.RiskGetPayload<{
       include: { assignee: { include: { user: true } } };
@@ -14,16 +17,11 @@ export async function getRisks(input: GetRiskSchema): Promise<{
   > & { assignee: User | null })[];
   pageCount: number;
 }> {
-  const session = await auth.api.getSession({
-    headers: await headers(),
-  });
-
-  if (!session?.session.activeOrganizationId) {
-    // In case of unauthorized or missing org, return empty data and 0 pageCount
+  if (!orgId) {
     return { data: [], pageCount: 0 };
   }
 
-  const { title, page, perPage, sort, filters, joinOperator } = input;
+  const { title, page, perPage, sort, filters, joinOperator } = searchParams;
 
   const orderBy = sort.map((s) => ({
     [s.id]: s.desc ? 'desc' : 'asc',
@@ -37,7 +35,7 @@ export async function getRisks(input: GetRiskSchema): Promise<{
   });
 
   const where: Prisma.RiskWhereInput = {
-    organizationId: session.session.activeOrganizationId,
+    organizationId: orgId,
     ...(title && {
       title: {
         contains: title,

apps/app/src/app/(app)/[orgId]/risk/(overview)/page.tsx

@@ -2,10 +2,8 @@ import { AppOnboarding } from '@/components/app-onboarding';
 import PageWithBreadcrumb from '@/components/pages/PageWithBreadcrumb';
 import { CreateRiskSheet } from '@/components/sheets/create-risk-sheet';
 import { getValidFilters } from '@/lib/data-table';
-import { auth } from '@/utils/auth';
 import { db } from '@db';
 import type { Metadata } from 'next';
-import { headers } from 'next/headers';
 import { cache } from 'react';
 import { RisksTable } from './RisksTable';
 import { getRisks } from './data/getRisks';
@@ -35,8 +33,8 @@ export default async function RiskRegisterPage(props: {
   };
 
   const [risksResult, assignees, onboarding] = await Promise.all([
-    getRisks(searchParamsForTable),
-    getAssignees(),
+    getRisks({ orgId, searchParams: searchParamsForTable }),
+    getAssignees(orgId),
     db.onboarding.findFirst({
       where: { organizationId: orgId },
       select: { triggerJobId: true },
@@ -92,6 +90,7 @@ export default async function RiskRegisterPage(props: {
         assignees={assignees}
         onboardingRunId={onboarding?.triggerJobId ?? null}
         searchParams={searchParamsForTable}
+        orgId={orgId}
       />
     </PageWithBreadcrumb>
   );
@@ -103,18 +102,14 @@ export async function generateMetadata(): Promise<Metadata> {
   };
 }
 
-const getAssignees = cache(async () => {
-  const session = await auth.api.getSession({
-    headers: await headers(),
-  });
-
-  if (!session || !session.session.activeOrganizationId) {
+const getAssignees = cache(async (orgId: string) => {
+  if (!orgId) {
     return [];
   }
 
   return await db.member.findMany({
     where: {
-      organizationId: session.session.activeOrganizationId,
+      organizationId: orgId,
       isActive: true,
       deactivated: false,
       role: {

apps/app/src/app/(app)/[orgId]/vendors/(overview)/actions/get-vendors-action.ts

@@ -1,19 +1,17 @@
 'use server';
 
-import { auth } from '@/utils/auth';
-import { headers } from 'next/headers';
 import { getVendors } from '../data/queries';
 import type { GetVendorsSchema } from '../data/validations';
 
-export async function getVendorsAction(input: GetVendorsSchema) {
-  const session = await auth.api.getSession({
-    headers: await headers(),
-  });
+export type GetVendorsActionInput = {
+  orgId: string;
+  searchParams: GetVendorsSchema;
+};
 
-  if (!session?.session.activeOrganizationId) {
+export async function getVendorsAction({ orgId, searchParams }: GetVendorsActionInput) {
+  if (!orgId) {
     return { data: [], pageCount: 0 };
   }
 
-  return await getVendors(session.session.activeOrganizationId, input);
+  return await getVendors(orgId, searchParams);
 }
-

apps/app/src/app/(app)/[orgId]/vendors/(overview)/components/VendorsTable.tsx

@@ -5,15 +5,14 @@ import { DataTableToolbar } from '@/components/data-table/data-table-toolbar';
 import { OnboardingLoadingAnimation } from '@/components/onboarding-loading-animation';
 import { useDataTable } from '@/hooks/use-data-table';
 import { getFiltersStateParser, getSortingStateParser } from '@/lib/parsers';
-import { useSession } from '@/utils/auth-client';
 import { Departments, Vendor } from '@db';
 import { ColumnDef } from '@tanstack/react-table';
 import { Loader2 } from 'lucide-react';
 import { parseAsInteger, parseAsString, parseAsStringEnum, useQueryState } from 'nuqs';
 import { useCallback, useMemo } from 'react';
 import useSWR from 'swr';
 import { CreateVendorSheet } from '../../components/create-vendor-sheet';
-import { getVendorsAction } from '../actions/get-vendors-action';
+import { getVendorsAction, type GetVendorsActionInput } from '../actions/get-vendors-action';
 import type { GetAssigneesResult, GetVendorsResult } from '../data/queries';
 import type { GetVendorsSchema } from '../data/validations';
 import { useOnboardingStatus } from '../hooks/use-onboarding-status';
@@ -25,6 +24,10 @@ export type VendorRow = GetVendorsResult['data'][number] & {
   isAssessing?: boolean;
 };
 
+const callGetVendorsAction = getVendorsAction as unknown as (
+  input: GetVendorsActionInput,
+) => Promise<GetVendorsResult>;
+
 const ACTIVE_STATUSES: Array<'pending' | 'processing' | 'created' | 'assessing'> = [
   'pending',
   'processing',
@@ -38,6 +41,7 @@ interface VendorsTableProps {
   assignees: GetAssigneesResult;
   onboardingRunId?: string | null;
   searchParams: GetVendorsSchema;
+  orgId: string;
 }
 
 export function VendorsTable({
@@ -46,10 +50,8 @@ export function VendorsTable({
   assignees,
   onboardingRunId,
   searchParams: initialSearchParams,
+  orgId,
 }: VendorsTableProps) {
-  const session = useSession();
-  const orgId = session?.data?.session?.activeOrganizationId;
-
   const { itemStatuses, progress, itemsInfo, isActive, isLoading } = useOnboardingStatus(
     onboardingRunId,
     'vendors',
@@ -103,8 +105,9 @@ export function VendorsTable({
 
   // Fetcher function for SWR
   const fetcher = useCallback(async () => {
-    return await getVendorsAction(currentSearchParams);
-  }, [currentSearchParams]);
+    if (!orgId) return { data: [], pageCount: 0 };
+    return await callGetVendorsAction({ orgId, searchParams: currentSearchParams });
+  }, [orgId, currentSearchParams]);
 
   // Use SWR to fetch vendors with polling when onboarding is active
   const { data: vendorsData } = useSWR(swrKey, fetcher, {
@@ -203,7 +206,7 @@ export function VendorsTable({
           residualProbability: 'very_unlikely' as const,
           residualImpact: 'insignificant' as const,
           website: null,
-          organizationId: orgId || '',
+          organizationId: orgId,
           assigneeId: null,
           assignee: null,
           createdAt: new Date(),
@@ -228,7 +231,7 @@ export function VendorsTable({
           residualProbability: 'very_unlikely' as const,
           residualImpact: 'insignificant' as const,
           website: null,
-          organizationId: orgId || '',
+          organizationId: orgId,
           assigneeId: null,
           assignee: null,
           createdAt: new Date(),
@@ -240,7 +243,7 @@ export function VendorsTable({
     return [...vendorsWithStatus, ...pendingVendors, ...tempVendors];
   }, [vendors, itemsInfo, itemStatuses, orgId, isActive, onboardingRunId]);
 
-  const columns = useMemo<ColumnDef<VendorRow>[]>(() => getColumns(orgId ?? ''), [orgId]);
+  const columns = useMemo<ColumnDef<VendorRow>[]>(() => getColumns(orgId), [orgId]);
 
   const { table } = useDataTable({
     data: mergedVendors,

apps/app/src/app/(app)/[orgId]/vendors/(overview)/page.tsx

@@ -90,6 +90,7 @@ export default async function Page({
         assignees={assignees}
         onboardingRunId={onboarding?.triggerJobId ?? null}
         searchParams={parsedSearchParams}
+        orgId={orgId}
       />
     </PageWithBreadcrumb>
   );